From 802f6dd89e46698e76adbd998c85f3654d84bfb0 Mon Sep 17 00:00:00 2001 From: Florian Hines Date: Tue, 8 Sep 2026 11:53:45 -0500 Subject: [PATCH 1/9] feat(auth): extract shared token foundation --- ENVIRONMENT.md | 10 + ...resource-delegation.servicecontrol.test.ts | 128 +++++ .../src/lib/auth/resource-delegation.test.ts | 398 +++++++++++++++ apps/web/src/lib/auth/resource-delegation.ts | 465 ++++++++++++++++++ .../lib/config.server.resource-tokens.test.ts | 107 ++++ apps/web/src/lib/config.server.ts | 28 ++ apps/web/src/lib/token.test.ts | 15 + apps/web/src/lib/tokens.ts | 127 ++++- apps/web/src/lib/user/server.test.ts | 84 ++++ apps/web/src/lib/user/server.ts | 39 +- apps/web/src/routers/test-utils.ts | 7 +- packages/worker-utils/package.json | 3 + .../src/kilo-auth-middleware.test.ts | 51 +- .../worker-utils/src/kilo-auth-middleware.ts | 22 +- .../src/kilo-token-policy.test.ts | 50 +- .../worker-utils/src/kilo-token-policy.ts | 56 ++- packages/worker-utils/src/kilo-token.test.ts | 22 + packages/worker-utils/src/kilo-token.ts | 21 +- .../src/runtime-authorization-contract.ts | 82 +++ .../src/runtime-authorization.test.ts | 346 +++++++++++++ .../worker-utils/src/runtime-authorization.ts | 368 ++++++++++++++ .../src/runtime-proxy-attestation.test.ts | 65 +++ .../src/runtime-proxy-attestation.ts | 125 +++++ .../src/middleware/kilo-jwt-auth.test.ts | 88 +++- .../src/middleware/kilo-jwt-auth.ts | 37 +- 25 files changed, 2706 insertions(+), 38 deletions(-) create mode 100644 apps/web/src/lib/auth/resource-delegation.servicecontrol.test.ts create mode 100644 apps/web/src/lib/auth/resource-delegation.test.ts create mode 100644 apps/web/src/lib/auth/resource-delegation.ts create mode 100644 apps/web/src/lib/config.server.resource-tokens.test.ts create mode 100644 packages/worker-utils/src/runtime-authorization-contract.ts create mode 100644 packages/worker-utils/src/runtime-authorization.test.ts create mode 100644 packages/worker-utils/src/runtime-authorization.ts create mode 100644 packages/worker-utils/src/runtime-proxy-attestation.test.ts create mode 100644 packages/worker-utils/src/runtime-proxy-attestation.ts diff --git a/ENVIRONMENT.md b/ENVIRONMENT.md index 20529b19ce..c7b7aa3e73 100644 --- a/ENVIRONMENT.md +++ b/ENVIRONMENT.md @@ -81,6 +81,16 @@ Manage shared web env var additions and rotations with `pnpm web:env set { + shared.enabled = true; + shared.family = ''; +}); +jest.mock('@/lib/config.server', () => ({ + NEXTAUTH_SECRET: 'service-control-test-secret', + isResourceTokenIssuanceEnabled: (family: string) => + shared.enabled && (!shared.family || shared.family === family), +})); +jest.mock('@/lib/user/server', () => ({ getUserFromSessionForCredentialIssuance: jest.fn() })); + +import { generateCloudAgentWorkflowToken, generateWorkflowGatewayToken } from '@/lib/tokens'; +import { defineTestUser } from '@/tests/helpers/user.helper'; + +describe('workflow service control tokens', () => { + test('uses bounded modern gateway workflow owner claims', () => { + const user = defineTestUser({ api_token_pepper: 'workflow-pepper' }); + const token = generateWorkflowGatewayToken(user, { + organizationId: 'organization-id', + tokenSource: 'reviewer', + }); + const claims = jwt.verify(token, 'service-control-test-secret') as jwt.JwtPayload; + + expect(claims).toMatchObject({ + aud: 'kilo-gateway', + kiloUserId: user.id, + apiTokenPepper: 'workflow-pepper', + organizationId: 'organization-id', + tokenSource: 'reviewer', + tokenPurpose: 'delegated-workload', + credentialExchange: false, + }); + expect(claims).not.toHaveProperty('runtimeAdmission'); + expect(claims.exp! - claims.iat!).toBe(60 * 60); + }); + + test('uses bounded modern automation admission', () => { + const user = defineTestUser({ api_token_pepper: 'workflow-pepper' }); + const token = generateCloudAgentWorkflowToken(user, { + expiresIn: 300, + tokenSource: 'reviewer', + botId: 'reviewer', + }); + const claims = jwt.verify(token, 'service-control-test-secret') as jwt.JwtPayload; + expect(claims).toMatchObject({ + aud: 'cloud-agent-next', + tokenPurpose: 'internal-service', + credentialExchange: false, + runtimeAdmission: { + source: 'automation', + authorizationUserId: user.id, + authorizationPepper: 'workflow-pepper', + }, + }); + expect(claims.exp! - claims.iat!).toBe(300); + }); + + test('caps modern workflow admission to one hour', () => { + const user = defineTestUser({ api_token_pepper: 'workflow-pepper' }); + const token = generateCloudAgentWorkflowToken(user, { + expiresIn: 5 * 365 * 24 * 60 * 60, + tokenSource: 'reviewer', + }); + const claims = jwt.decode(token) as jwt.JwtPayload; + + expect(claims.exp! - claims.iat!).toBe(60 * 60); + }); + + test('requires an authorization pepper for modern workflow admission', () => { + const user = defineTestUser({ api_token_pepper: 'workflow-pepper' }); + const authorizationUser = defineTestUser({ api_token_pepper: null }); + + expect(() => + generateCloudAgentWorkflowToken(user, { + expiresIn: 300, + tokenSource: 'reviewer', + authorizationUser, + }) + ).toThrow('current authorization pepper'); + }); + + test('preserves the legacy workflow token shape when shared issuance is disabled', () => { + shared.enabled = false; + const user = defineTestUser({ api_token_pepper: 'workflow-pepper' }); + const token = generateCloudAgentWorkflowToken(user, { + expiresIn: 300, + tokenSource: 'reviewer', + botId: 'reviewer', + }); + const claims = jwt.verify(token, 'service-control-test-secret') as jwt.JwtPayload; + expect(claims).toMatchObject({ + kiloUserId: user.id, + tokenSource: 'reviewer', + botId: 'reviewer', + }); + expect(claims).not.toHaveProperty('tokenPurpose'); + expect(claims.exp! - claims.iat!).toBe(300); + }); +}); + +test.each(['cloud-agent-next', 'workflow-gateway'])( + 'only enables the selected workflow family %s', + family => { + shared.family = family; + const user = defineTestUser({ api_token_pepper: 'workflow-pepper' }); + const cloud = jwt.decode( + generateCloudAgentWorkflowToken(user, { + expiresIn: 300, + tokenSource: 'reviewer', + }) + ) as jwt.JwtPayload; + const gateway = jwt.decode( + generateWorkflowGatewayToken(user, { + tokenSource: 'reviewer', + }) + ) as jwt.JwtPayload; + expect(cloud.tokenPurpose).toBe(family === 'cloud-agent-next' ? 'internal-service' : undefined); + expect(gateway.tokenPurpose).toBe( + family === 'workflow-gateway' ? 'delegated-workload' : undefined + ); + expect(cloud.exp! - cloud.iat!).toBe(300); + if (family === 'cloud-agent-next') expect(gateway).not.toHaveProperty('aud'); + } +); diff --git a/apps/web/src/lib/auth/resource-delegation.test.ts b/apps/web/src/lib/auth/resource-delegation.test.ts new file mode 100644 index 0000000000..74a90296cd --- /dev/null +++ b/apps/web/src/lib/auth/resource-delegation.test.ts @@ -0,0 +1,398 @@ +import { isResourceTokenIssuanceEnabled } from '@/lib/config.server'; +import { afterEach, describe, expect, test } from '@jest/globals'; +import { + device_sessions, + kilocode_users, + organization_memberships, + organizations, +} from '@kilocode/db/schema'; +import { eq, inArray } from 'drizzle-orm'; +import jwt from 'jsonwebtoken'; +import { buildModernKiloTokenPayload } from '@kilocode/worker-utils/kilo-token-policy'; + +const shared = { enabled: true, family: '' }; +jest.mock('@/lib/config.server', () => ({ + NEXTAUTH_SECRET: 'resource-delegation-test-secret', + isResourceTokenIssuanceEnabled: jest.fn( + (family: string) => shared.enabled && (!shared.family || shared.family === family) + ), +})); +jest.mock('@/lib/user/server', () => ({ + getUserFromSessionForCredentialIssuance: jest.fn(), +})); + +import { + canIssueLegacyOrganizationToken, + createControlTokenForRequest, + createDelegatedResourceToken, + getResourceDelegationAuthority, +} from './resource-delegation'; +import { db } from '@/lib/drizzle'; +import { getUserFromSessionForCredentialIssuance } from '@/lib/user/server'; +import { insertTestUser } from '@/tests/helpers/user.helper'; + +const secret = 'resource-delegation-test-secret'; +const cleanups: string[] = []; + +afterEach(async () => { + if (cleanups.length) { + await db.delete(kilocode_users).where(inArray(kilocode_users.id, cleanups)); + cleanups.length = 0; + } + shared.enabled = true; + shared.family = ''; + jest.clearAllMocks(); +}); + +async function user() { + const row = await insertTestUser({ api_token_pepper: crypto.randomUUID() }); + cleanups.push(row.id); + return row; +} + +async function organizationFor(userId: string) { + const [organization] = await db + .insert(organizations) + .values({ + name: `Resource delegation ${crypto.randomUUID()}`, + created_by_kilo_user_id: userId, + require_seats: false, + }) + .returning(); + return organization; +} + +function bearer(token: string) { + return new Headers({ authorization: `Bearer ${token}` }); +} + +describe('legacy organization token compatibility', () => { + test('allows only requests without an Authorization header', () => { + expect(canIssueLegacyOrganizationToken(new Headers())).toBe(true); + expect(canIssueLegacyOrganizationToken(bearer('restricted-token'))).toBe(false); + expect(canIssueLegacyOrganizationToken(new Headers({ authorization: '' }))).toBe(false); + }); +}); + +function modernToken( + user: { id: string; api_token_pepper: string | null }, + options?: { + purpose?: 'human-api' | 'device-access'; + exchange?: boolean; + deviceSessionId?: string; + } +) { + const now = Math.floor(Date.now() / 1000); + const base = { + userId: user.id, + pepper: user.api_token_pepper, + env: process.env.NODE_ENV, + audience: 'kilo-api', + issuedAt: now, + expiresAt: now + 1800, + extra: options?.deviceSessionId ? { deviceSessionId: options.deviceSessionId } : undefined, + }; + const payload = + options?.purpose === 'device-access' + ? buildModernKiloTokenPayload({ + ...base, + tokenPurpose: 'device-access', + credentialExchange: false, + }) + : buildModernKiloTokenPayload({ + ...base, + tokenPurpose: 'human-api', + credentialExchange: options?.exchange ?? true, + }); + return jwt.sign(payload, secret, { algorithm: 'HS256' }); +} + +describe('resource delegation authority', () => { + test.each([true, false])( + 'accepts direct billing_manager membership when shared issuance is %s', + async enabled => { + shared.enabled = enabled; + const current = await user(); + const organization = await organizationFor(current.id); + await db.insert(organization_memberships).values({ + organization_id: organization.id, + kilo_user_id: current.id, + role: 'billing_manager', + }); + const headers = enabled ? bearer(modernToken(current)) : new Headers(); + if (!enabled) { + jest.mocked(getUserFromSessionForCredentialIssuance).mockResolvedValue({ + user: current, + authFailedResponse: null, + }); + } + + const result = await createControlTokenForRequest(current, 'cloud-agent-next', { + headers, + organizationId: organization.id, + }); + const claims = jwt.decode(result.token) as jwt.JwtPayload; + expect(claims).not.toHaveProperty('organizationRole'); + if (enabled) { + expect(claims.organizationId).toBe(organization.id); + } else { + expect(claims).not.toHaveProperty('organizationId'); + } + } + ); + + test.each([ + [true, 'removed'], + [false, 'removed'], + [true, 'deleted'], + [false, 'deleted'], + ] as const)( + 'denies direct organization access when shared issuance is %s and state is %s', + async (enabled, state) => { + shared.enabled = enabled; + const current = await user(); + const organization = await organizationFor(current.id); + await db.insert(organization_memberships).values({ + organization_id: organization.id, + kilo_user_id: current.id, + role: 'billing_manager', + }); + if (state === 'removed') { + await db + .delete(organization_memberships) + .where(eq(organization_memberships.organization_id, organization.id)); + } else { + await db + .update(organizations) + .set({ deleted_at: new Date().toISOString() }) + .where(eq(organizations.id, organization.id)); + } + + const headers = enabled ? bearer(modernToken(current)) : new Headers(); + if (!enabled) { + jest.mocked(getUserFromSessionForCredentialIssuance).mockResolvedValue({ + user: current, + authFailedResponse: null, + }); + } + + await expect( + createControlTokenForRequest(current, 'cloud-agent-next', { + headers, + organizationId: organization.id, + }) + ).rejects.toThrow('Unauthorized resource delegation request'); + } + ); + + test('accepts an exchangeable modern human API credential and preserves its provenance', async () => { + const current = await user(); + const authority = await getResourceDelegationAuthority(current, { + headers: bearer(modernToken(current)), + }); + + expect(authority).toMatchObject({ + credentialKind: 'human-api', + isModern: true, + runtimeAdmission: { source: 'user', authorizationUserId: current.id }, + }); + }); + + test('requires an active owned device session for a modern device credential', async () => { + const current = await user(); + const [session] = await db + .insert(device_sessions) + .values({ kilo_user_id: current.id, user_agent: 'resource-delegation-test' }) + .returning({ id: device_sessions.id }); + + const authority = await getResourceDelegationAuthority(current, { + headers: bearer( + modernToken(current, { + purpose: 'device-access', + exchange: false, + deviceSessionId: session.id, + }) + ), + }); + expect(authority.credentialKind).toBe('device-access'); + }); + + test('rejects restricted modern principals rather than projecting away their source claim', async () => { + const current = await user(); + const claims = jwt.decode(modernToken(current)); + if (!claims || typeof claims === 'string') throw new Error('Expected JWT claims'); + const token = jwt.sign({ ...claims, tokenSource: 'cloud-agent' }, secret, { + algorithm: 'HS256', + }); + await expect( + getResourceDelegationAuthority(current, { headers: bearer(token) }) + ).rejects.toThrow('Unauthorized resource delegation request'); + }); + + test('does not fall back to ambient cookies for malformed supplied authorization', async () => { + const current = await user(); + jest.mocked(getUserFromSessionForCredentialIssuance).mockResolvedValue({ + user: current, + authFailedResponse: null, + }); + await expect( + getResourceDelegationAuthority(current, { + headers: new Headers({ authorization: 'Basic bad' }), + }) + ).rejects.toThrow('Unauthorized resource delegation request'); + }); + + test('mints a bounded single-audience control token from a verified authority', async () => { + const current = await user(); + const result = await createControlTokenForRequest(current, 'gastown', { + headers: bearer(modernToken(current)), + expiresIn: 7200, + }); + const claims = jwt.verify(result.token, secret) as jwt.JwtPayload; + expect(claims).toMatchObject({ + aud: 'gastown', + tokenPurpose: 'human-api', + credentialExchange: false, + runtimeAdmission: { source: 'user', authorizationUserId: current.id }, + }); + expect(claims.exp! - claims.iat!).toBeLessThanOrEqual(3600); + }); + + test('keeps a legitimate legacy session request on the legacy contract when disabled', async () => { + shared.enabled = false; + const current = await user(); + jest.mocked(getUserFromSessionForCredentialIssuance).mockResolvedValue({ + user: current, + authFailedResponse: null, + }); + const result = await createControlTokenForRequest(current, 'gastown', { + headers: new Headers(), + legacyExpiresIn: 60, + }); + const claims = jwt.verify(result.token, secret) as jwt.JwtPayload; + expect(claims.aud).toBeUndefined(); + expect(claims.exp! - claims.iat!).toBe(60); + }); + + test('returns migration unavailable for a verified modern user principal when disabled', async () => { + shared.enabled = false; + const current = await user(); + await expect( + createControlTokenForRequest(current, 'wasteland', { headers: bearer(modernToken(current)) }) + ).rejects.toMatchObject({ status: 503, delegationCode: 'MIGRATION_UNAVAILABLE' }); + }); + + test('continues bounded control issuance for an active modern device after shared rollout rollback', async () => { + const current = await user(); + const [session] = await db + .insert(device_sessions) + .values({ kilo_user_id: current.id, user_agent: 'resource-delegation-test' }) + .returning({ id: device_sessions.id }); + const headers = bearer( + modernToken(current, { + purpose: 'device-access', + exchange: false, + deviceSessionId: session.id, + }) + ); + + const active = await createControlTokenForRequest(current, 'cloud-agent-next', { headers }); + expect(jwt.verify(active.token, secret)).toMatchObject({ + aud: 'cloud-agent-next', + tokenPurpose: 'device-access', + credentialExchange: false, + deviceSessionId: session.id, + runtimeAdmission: { source: 'user', authorizationUserId: current.id }, + }); + + shared.enabled = false; + const rollback = await createControlTokenForRequest(current, 'gastown', { headers }); + const claims = jwt.verify(rollback.token, secret) as jwt.JwtPayload; + expect(claims).toMatchObject({ + aud: 'gastown', + tokenPurpose: 'device-access', + credentialExchange: false, + deviceSessionId: session.id, + runtimeAdmission: { source: 'user', authorizationUserId: current.id }, + }); + expect(claims.exp! - claims.iat!).toBeLessThanOrEqual(1800); + + const rotatedPepper = crypto.randomUUID(); + await db + .update(kilocode_users) + .set({ api_token_pepper: rotatedPepper }) + .where(eq(kilocode_users.id, current.id)); + await expect(createControlTokenForRequest(current, 'wasteland', { headers })).rejects.toThrow( + 'Unauthorized resource delegation request' + ); + await db + .update(kilocode_users) + .set({ api_token_pepper: current.api_token_pepper }) + .where(eq(kilocode_users.id, current.id)); + + await db + .update(device_sessions) + .set({ revoked_at: new Date().toISOString(), revoked_reason: 'test' }) + .where(eq(device_sessions.id, session.id)); + await expect(createControlTokenForRequest(current, 'wasteland', { headers })).rejects.toThrow( + 'Unauthorized resource delegation request' + ); + }); + + test('mints a non-exchangeable delegated token for exactly the selected audience', async () => { + const current = await user(); + const result = await createDelegatedResourceToken(current, 'gateway', { + headers: bearer(modernToken(current)), + }); + const claims = jwt.verify(result.token, secret) as jwt.JwtPayload; + expect(claims).toMatchObject({ + aud: 'kilo-gateway', + tokenPurpose: 'delegated-workload', + credentialExchange: false, + }); + expect(claims.exp! - claims.iat!).toBeLessThanOrEqual(15 * 60); + }); + + test('caps delegated tokens to the parent credential expiry', async () => { + const current = await user(); + const now = Math.floor(Date.now() / 1000); + const claims = buildModernKiloTokenPayload({ + userId: current.id, + pepper: current.api_token_pepper, + env: process.env.NODE_ENV, + audience: 'kilo-api', + issuedAt: now - 1, + expiresAt: now + 30, + tokenPurpose: 'human-api', + credentialExchange: true, + }); + const parent = jwt.sign(claims, secret, { algorithm: 'HS256' }); + const result = await createDelegatedResourceToken(current, 'api', { headers: bearer(parent) }); + const delegated = jwt.verify(result.token, secret) as jwt.JwtPayload; + expect(delegated.exp! - delegated.iat!).toBeLessThanOrEqual(30); + }); +}); + +test.each(['cloud-agent-next', 'gastown', 'wasteland'] as const)( + 'isolates request control family %s', + async family => { + shared.family = family; + const current = await user(); + jest + .mocked(getUserFromSessionForCredentialIssuance) + .mockResolvedValue({ user: current, authFailedResponse: null }); + for (const resource of ['cloud-agent-next', 'gastown', 'wasteland'] as const) { + const result = await createControlTokenForRequest(current, resource, { + headers: new Headers(), + }); + const claims = jwt.decode(result.token) as jwt.JwtPayload; + expect(isResourceTokenIssuanceEnabled).toHaveBeenLastCalledWith(resource); + expect(claims.aud).toBe(resource === family ? resource : undefined); + expect(claims.tokenPurpose).toBe(resource === family ? 'human-api' : undefined); + } + await expect( + createDelegatedResourceToken(current, 'gateway', { headers: new Headers() }) + ).rejects.toMatchObject({ status: 503, delegationCode: 'MIGRATION_UNAVAILABLE' }); + expect(isResourceTokenIssuanceEnabled).toHaveBeenLastCalledWith('delegated-resource'); + } +); diff --git a/apps/web/src/lib/auth/resource-delegation.ts b/apps/web/src/lib/auth/resource-delegation.ts new file mode 100644 index 0000000000..8833e69f80 --- /dev/null +++ b/apps/web/src/lib/auth/resource-delegation.ts @@ -0,0 +1,465 @@ +import 'server-only'; + +import type { User } from '@kilocode/db/schema'; +import { + device_sessions, + kilocode_users, + organization_memberships, + organizations, +} from '@kilocode/db/schema'; +import { and, eq, isNull, ne } from 'drizzle-orm'; +import jwt from 'jsonwebtoken'; +import { TRPCError } from '@trpc/server'; +import { headers as nextHeaders } from 'next/headers'; +import { + AI_ATTRIBUTION_AUDIENCE, + CLOUD_AGENT_NEXT_AUDIENCE, + GASTOWN_AUDIENCE, + HTML_DEPLOY_AUDIENCE, + KILO_API_AUDIENCE, + KILO_GATEWAY_AUDIENCE, + WASTELAND_AUDIENCE, +} from '@kilocode/worker-utils/internal-service-token-audiences'; +import { + buildModernKiloTokenPayload, + isKiloCredentialExchangeEligible, + verifyKiloTokenForPolicy, +} from '@kilocode/worker-utils/kilo-token-policy'; +import type { RuntimeAdmission } from '@kilocode/worker-utils/runtime-authorization-contract'; +import { isResourceTokenIssuanceEnabled, NEXTAUTH_SECRET } from '@/lib/config.server'; +import { db } from '@/lib/drizzle'; +import { generateApiToken, TOKEN_EXPIRY } from '@/lib/tokens'; +import { getUserFromSessionForCredentialIssuance } from '@/lib/user/server'; + +const ONE_HOUR_SECONDS = 60 * 60; +const LEGACY_DEVICE_SESSION_SECONDS = ONE_HOUR_SECONDS; + +export type DelegableResource = 'api' | 'gateway' | 'attribution' | 'html-deploy'; +export type ControlResource = 'cloud-agent-next' | 'gastown' | 'wasteland'; +type Resource = DelegableResource | ControlResource; +type CredentialKind = 'human-api' | 'device-access'; + +export type ResourceDelegationAuthority = { + user: User; + credentialKind: CredentialKind; + expiresAt?: number; + organizationId?: string; + audience?: string | string[]; + tokenSource?: string; + deviceSessionId?: string; + isModern: boolean; + runtimeAdmission: RuntimeAdmission; +}; + +export class TypedResourceDelegationError extends TRPCError { + constructor( + public readonly status: 401 | 403 | 503, + public readonly delegationCode: 'UNAUTHORIZED' | 'FORBIDDEN' | 'MIGRATION_UNAVAILABLE', + message: string + ) { + super({ + code: status === 401 ? 'UNAUTHORIZED' : status === 503 ? 'SERVICE_UNAVAILABLE' : 'FORBIDDEN', + message, + }); + } +} + +function unauthorized(message = 'Unauthorized resource delegation request'): never { + throw new TypedResourceDelegationError(401, 'UNAUTHORIZED', message); +} + +function forbidden(message: string): never { + throw new TypedResourceDelegationError(403, 'FORBIDDEN', message); +} + +export type CreateControlTokenOptions = { + headers?: Headers; + organizationId?: string; + tokenSource?: string; + expiresIn?: number; + legacyExpiresIn?: number; + extra?: { + botId?: string; + createdOnPlatform?: string; + isAdmin?: boolean; + gastownAccess?: boolean; + orgMemberships?: Array<{ + orgId: string; + role: 'owner' | 'admin' | 'member' | 'billing_manager'; + }>; + }; +}; + +export type CreateDelegatedResourceTokenOptions = { + headers?: Headers; + organizationId?: string; + organizationRole?: 'owner' | 'admin' | 'member'; + tokenSource?: string; + expiresIn?: number; +}; + +function resourceAudience(resource: Resource): string { + return { + api: KILO_API_AUDIENCE, + gateway: KILO_GATEWAY_AUDIENCE, + attribution: AI_ATTRIBUTION_AUDIENCE, + 'html-deploy': HTML_DEPLOY_AUDIENCE, + 'cloud-agent-next': CLOUD_AGENT_NEXT_AUDIENCE, + gastown: GASTOWN_AUDIENCE, + wasteland: WASTELAND_AUDIENCE, + }[resource]; +} + +export function isDelegableResource(value: unknown): value is DelegableResource { + return ( + value === 'api' || value === 'gateway' || value === 'attribution' || value === 'html-deploy' + ); +} + +export function canIssueLegacyOrganizationToken(requestHeaders: Headers): boolean { + return !requestHeaders.has('authorization'); +} + +function tokenFromHeaders(requestHeaders: Headers): string | null { + const value = requestHeaders.get('authorization'); + if (!value || !value.toLowerCase().startsWith('bearer ')) return null; + const token = value.slice(7).trim(); + return token || null; +} + +function hasUnsafeLegacyClaims(claimNames: readonly string[]): boolean { + const allowed = new Set([ + 'version', + 'kiloUserId', + 'apiTokenPepper', + 'env', + 'iat', + 'exp', + 'deviceAuthRequestCode', + 'deviceSessionId', + ]); + return claimNames.some(name => !allowed.has(name)); +} + +async function currentUser(expectedUser: User): Promise { + const user = await db.query.kilocode_users.findFirst({ + where: eq(kilocode_users.id, expectedUser.id), + }); + if (!user || user.blocked_at || user.blocked_reason) unauthorized(); + return user; +} + +async function assertActiveDeviceSession(id: string, userId: string): Promise { + const session = await db.query.device_sessions.findFirst({ + where: and( + eq(device_sessions.id, id), + eq(device_sessions.kilo_user_id, userId), + isNull(device_sessions.revoked_at) + ), + }); + if (!session) unauthorized(); +} + +async function membershipsFor(userId: string) { + return db + .select({ + orgId: organization_memberships.organization_id, + role: organization_memberships.role, + }) + .from(organization_memberships) + .where( + and( + eq(organization_memberships.kilo_user_id, userId), + ne(organization_memberships.role, 'billing_manager') + ) + ); +} + +async function assertCurrentOrganizationMembership( + userId: string, + organizationId?: string +): Promise { + if (!organizationId) return; + const membership = await db + .select({ id: organization_memberships.id }) + .from(organization_memberships) + .innerJoin(organizations, eq(organizations.id, organization_memberships.organization_id)) + .where( + and( + eq(organization_memberships.kilo_user_id, userId), + eq(organization_memberships.organization_id, organizationId), + isNull(organizations.deleted_at) + ) + ) + .limit(1); + if (!membership[0]) unauthorized(); +} + +function hasRestrictedPrincipalClaims(claimNames: readonly string[]): boolean { + const allowed = new Set([ + 'version', + 'kiloUserId', + 'apiTokenPepper', + 'env', + 'aud', + 'iat', + 'exp', + 'tokenPurpose', + 'credentialExchange', + 'deviceSessionId', + ]); + return claimNames.some(name => !allowed.has(name)); +} + +export async function getResourceDelegationAuthority( + expectedUser: User, + options?: { headers?: Headers; organizationId?: string } +): Promise { + const requestHeaders = options?.headers ?? (await nextHeaders()); + const bearer = tokenFromHeaders(requestHeaders); + const authorizationPresent = requestHeaders.has('authorization'); + const user = await currentUser(expectedUser); + await assertCurrentOrganizationMembership(user.id, options?.organizationId); + + if (!bearer) { + if (authorizationPresent) { + unauthorized(); + } + const session = await getUserFromSessionForCredentialIssuance(); + if (!session.user || session.user.id !== expectedUser.id) { + unauthorized(); + } + return { + user, + credentialKind: 'human-api', + isModern: false, + runtimeAdmission: { + source: 'user', + authorizationUserId: user.id, + authorizationPepper: user.api_token_pepper, + }, + }; + } + + let verified: Awaited>; + try { + verified = await verifyKiloTokenForPolicy(bearer, NEXTAUTH_SECRET, { + audience: KILO_API_AUDIENCE, + mode: 'allow-legacy', + }); + } catch { + unauthorized(); + } + const claims = verified.claims; + if ( + verified.userId !== expectedUser.id || + claims.env !== process.env.NODE_ENV || + claims.apiTokenPepper !== user.api_token_pepper + ) { + unauthorized(); + } + + if (claims.tokenPurpose !== undefined) { + if (claims.tokenPurpose !== 'human-api' && claims.tokenPurpose !== 'device-access') { + forbidden('Delegated credentials cannot mint resource control tokens'); + } + if (claims.aud !== KILO_API_AUDIENCE || hasRestrictedPrincipalClaims(verified.claimNames)) { + unauthorized(); + } + if (claims.tokenPurpose === 'device-access') { + if (!claims.deviceSessionId) unauthorized(); + await assertActiveDeviceSession(claims.deviceSessionId, user.id); + } + return { + user, + credentialKind: claims.tokenPurpose, + expiresAt: claims.exp, + organizationId: claims.organizationId, + audience: claims.aud, + tokenSource: claims.tokenSource, + deviceSessionId: claims.deviceSessionId, + isModern: true, + runtimeAdmission: { + source: 'user', + authorizationUserId: user.id, + authorizationPepper: user.api_token_pepper, + }, + }; + } + + if (hasUnsafeLegacyClaims(verified.claimNames)) { + forbidden('Unsupported legacy credential context'); + } + const isLegacyDevice = claims.deviceSessionId !== undefined; + if (isLegacyDevice) { + if (claims.exp - claims.iat !== LEGACY_DEVICE_SESSION_SECONDS) { + forbidden('Unsupported legacy credential context'); + } + await assertActiveDeviceSession(claims.deviceSessionId, user.id); + } else if (!isKiloCredentialExchangeEligible(verified, { legacy: 'five-year-api' })) { + forbidden('Unsupported legacy credential context'); + } + return { + user, + credentialKind: isLegacyDevice ? 'device-access' : 'human-api', + expiresAt: claims.exp, + isModern: false, + runtimeAdmission: { + source: 'user', + authorizationUserId: user.id, + authorizationPepper: user.api_token_pepper, + }, + }; +} + +export async function isModernResourceDelegationRequest(requestHeaders: Headers): Promise { + const bearer = tokenFromHeaders(requestHeaders); + if (!bearer) return false; + const verified = await verifyKiloTokenForPolicy(bearer, NEXTAUTH_SECRET, { + audience: KILO_API_AUDIENCE, + mode: 'allow-legacy', + }); + return verified.claims.tokenPurpose !== undefined; +} + +async function createModernControlToken( + authority: ResourceDelegationAuthority, + resource: ControlResource, + options?: CreateControlTokenOptions +): Promise<{ token: string; expiresAt: string; user: User; tokenSource?: string }> { + const now = Math.floor(Date.now() / 1000); + const requested = options?.expiresIn ?? ONE_HOUR_SECONDS; + const expiresIn = Math.min( + requested, + ONE_HOUR_SECONDS, + authority.expiresAt ? authority.expiresAt - now : requested + ); + if (expiresIn <= 0) unauthorized('Resource delegation authority has expired'); + const memberships = + resource === 'cloud-agent-next' ? undefined : await membershipsFor(authority.user.id); + const payload = buildModernKiloTokenPayload({ + userId: authority.user.id, + pepper: authority.user.api_token_pepper, + env: process.env.NODE_ENV, + audience: resourceAudience(resource), + issuedAt: now, + expiresAt: now + expiresIn, + tokenPurpose: authority.credentialKind, + credentialExchange: false, + extra: { + organizationId: options?.organizationId ?? authority.organizationId, + tokenSource: options?.tokenSource ?? authority.tokenSource, + deviceSessionId: authority.deviceSessionId, + botId: options?.extra?.botId, + createdOnPlatform: options?.extra?.createdOnPlatform, + isAdmin: options?.extra?.isAdmin === true && authority.user.is_admin, + gastownAccess: options?.extra?.gastownAccess, + orgMemberships: memberships, + ...(resource === 'wasteland' ? {} : { runtimeAdmission: authority.runtimeAdmission }), + }, + }); + const token = jwt.sign(payload, NEXTAUTH_SECRET, { algorithm: 'HS256' }); + return { + token, + expiresAt: new Date((now + expiresIn) * 1000).toISOString(), + user: authority.user, + tokenSource: options?.tokenSource ?? authority.tokenSource, + }; +} + +export async function createControlTokenForRequest( + user: User, + resource: ControlResource, + options?: CreateControlTokenOptions +): Promise<{ token: string; expiresAt: string; user: User; tokenSource?: string }> { + const authority = await getResourceDelegationAuthority(user, options); + if (authority.organizationId) { + forbidden('Organization-scoped credentials cannot mint resource control tokens'); + } + if (!isResourceTokenIssuanceEnabled(resource)) { + if (authority.isModern) { + if (authority.credentialKind === 'device-access' && authority.deviceSessionId) { + return await createModernControlToken(authority, resource, options); + } + throw new TypedResourceDelegationError( + 503, + 'MIGRATION_UNAVAILABLE', + 'Shared resource token migration is unavailable' + ); + } + const expiresIn = + options?.legacyExpiresIn ?? + (resource === 'cloud-agent-next' ? TOKEN_EXPIRY.default : ONE_HOUR_SECONDS); + const token = generateApiToken( + authority.user, + { + ...options?.extra, + tokenSource: options?.tokenSource, + ...(resource === 'cloud-agent-next' + ? {} + : { + isAdmin: options?.extra?.isAdmin === true && authority.user.is_admin, + orgMemberships: await membershipsFor(authority.user.id), + }), + }, + { expiresIn } + ); + return { + token, + expiresAt: new Date(Date.now() + expiresIn * 1000).toISOString(), + user: authority.user, + tokenSource: options?.tokenSource, + }; + } + return await createModernControlToken(authority, resource, options); +} + +export async function createDelegatedResourceToken( + user: User, + resource: DelegableResource, + options?: CreateDelegatedResourceTokenOptions +): Promise<{ token: string; expiresAt: string; user: User; tokenSource?: string }> { + const authority = await getResourceDelegationAuthority(user, options); + if (authority.organizationId && authority.organizationId !== options?.organizationId) { + forbidden('Scoped credentials cannot mint tokens for another organization'); + } + if (authority.organizationId && authority.audience !== resourceAudience(resource)) { + forbidden('Scoped credentials cannot broaden their resource audience'); + } + if (!isResourceTokenIssuanceEnabled('delegated-resource')) { + throw new TypedResourceDelegationError( + 503, + 'MIGRATION_UNAVAILABLE', + 'Shared resource token migration is unavailable' + ); + } + const now = Math.floor(Date.now() / 1000); + const requested = options?.expiresIn ?? 15 * 60; + const expiresIn = Math.min( + requested, + 15 * 60, + authority.expiresAt ? authority.expiresAt - now : requested + ); + if (expiresIn <= 0) unauthorized('Resource delegation authority has expired'); + const payload = buildModernKiloTokenPayload({ + userId: authority.user.id, + pepper: authority.user.api_token_pepper, + env: process.env.NODE_ENV, + audience: resourceAudience(resource), + issuedAt: now, + expiresAt: now + expiresIn, + tokenPurpose: 'delegated-workload', + credentialExchange: false, + extra: { + organizationId: options?.organizationId ?? authority.organizationId, + organizationRole: options?.organizationRole, + tokenSource: options?.tokenSource ?? authority.tokenSource, + }, + }); + return { + token: jwt.sign(payload, NEXTAUTH_SECRET, { algorithm: 'HS256' }), + expiresAt: new Date((now + expiresIn) * 1000).toISOString(), + user: authority.user, + tokenSource: options?.tokenSource ?? authority.tokenSource, + }; +} diff --git a/apps/web/src/lib/config.server.resource-tokens.test.ts b/apps/web/src/lib/config.server.resource-tokens.test.ts new file mode 100644 index 0000000000..77d340a9ec --- /dev/null +++ b/apps/web/src/lib/config.server.resource-tokens.test.ts @@ -0,0 +1,107 @@ +import jwt from 'jsonwebtoken'; +import type { User } from '@kilocode/db/schema'; +import { generateCloudAgentWorkflowToken, generateWorkflowGatewayToken } from './tokens'; +import { + isNativeResourceCredentialIssuanceEnabled, + isResourceTokenIssuanceEnabled, + type ResourceTokenFamily, +} from './config.server'; + +const families = { + 'cloud-agent-next': 'CLOUD_AGENT_RESOURCE_TOKENS_ENABLED', + gastown: 'GASTOWN_RESOURCE_TOKENS_ENABLED', + wasteland: 'WASTELAND_RESOURCE_TOKENS_ENABLED', + chat: 'CHAT_RESOURCE_TOKENS_ENABLED', + 'delegated-resource': 'DELEGATED_RESOURCE_TOKENS_ENABLED', + 'workflow-gateway': 'WORKFLOW_GATEWAY_RESOURCE_TOKENS_ENABLED', + benchmark: 'BENCHMARK_RESOURCE_TOKENS_ENABLED', +} satisfies Record; +const entries = Object.entries(families) as [ResourceTokenFamily, string][]; +const master = 'SHARED_RESOURCE_TOKENS_ENABLED'; +const native = 'NATIVE_RESOURCE_TOKENS_ENABLED'; +const keys = [master, native, ...Object.values(families)]; +const saved = new Map(keys.map(key => [key, process.env[key]])); + +beforeEach(() => { + for (const key of keys) delete process.env[key]; +}); +afterEach(() => { + for (const [key, value] of saved) { + if (value === undefined) delete process.env[key]; + else process.env[key] = value; + } +}); + +it('defaults every family and native adoption off', () => { + for (const [family] of entries) expect(isResourceTokenIssuanceEnabled(family)).toBe(false); + expect(isNativeResourceCredentialIssuanceEnabled()).toBe(false); +}); + +it('keeps all families off when only the master is enabled', () => { + process.env[master] = 'true'; + for (const [family] of entries) expect(isResourceTokenIssuanceEnabled(family)).toBe(false); +}); + +it.each([undefined, '', 'false', 'TRUE', '1', ' true', 'true '])( + 'requires exact true on the master (%s)', + value => { + if (value !== undefined) process.env[master] = value; + for (const [, key] of entries) process.env[key] = 'true'; + process.env[native] = 'true'; + for (const [family] of entries) expect(isResourceTokenIssuanceEnabled(family)).toBe(false); + expect(isNativeResourceCredentialIssuanceEnabled()).toBe(false); + } +); + +it.each(entries)('enables only %s and leaves native independent', (family, key) => { + process.env[master] = 'true'; + process.env[key] = 'true'; + for (const [candidate] of entries) { + expect(isResourceTokenIssuanceEnabled(candidate)).toBe(candidate === family); + } + expect(isNativeResourceCredentialIssuanceEnabled()).toBe(false); +}); + +it.each(['', 'false', 'TRUE', '1', ' true', 'true '])( + 'requires exact true on each family and native flag (%s)', + value => { + process.env[master] = 'true'; + for (const [, key] of entries) process.env[key] = value; + process.env[native] = value; + for (const [family] of entries) expect(isResourceTokenIssuanceEnabled(family)).toBe(false); + expect(isNativeResourceCredentialIssuanceEnabled()).toBe(false); + } +); + +it('native adoption needs no producer family enabled and enables none', () => { + process.env[master] = 'true'; + process.env[native] = 'true'; + expect(isNativeResourceCredentialIssuanceEnabled()).toBe(true); + for (const [family] of entries) expect(isResourceTokenIssuanceEnabled(family)).toBe(false); +}); + +it.each([ + ['CLOUD_AGENT_RESOURCE_TOKENS_ENABLED', true, false], + ['WORKFLOW_GATEWAY_RESOURCE_TOKENS_ENABLED', false, true], +] as const)('real issuer integration isolates %s', (flag, cloudModern, gatewayModern) => { + process.env[master] = 'true'; + process.env[flag] = 'true'; + const user = { id: 'oauth/family-test', api_token_pepper: 'family-test-pepper' } as User; + const cloud = jwt.decode( + generateCloudAgentWorkflowToken(user, { + expiresIn: 7200, + tokenSource: 'reviewer', + }) + ) as jwt.JwtPayload; + const gateway = jwt.decode( + generateWorkflowGatewayToken(user, { + tokenSource: 'reviewer', + }) + ) as jwt.JwtPayload; + expect(cloud.aud).toBe(cloudModern ? 'cloud-agent-next' : undefined); + expect(cloud.tokenPurpose).toBe(cloudModern ? 'internal-service' : undefined); + expect(cloud.exp! - cloud.iat!).toBe(cloudModern ? 3600 : 7200); + expect(gateway.aud).toBe(gatewayModern ? 'kilo-gateway' : undefined); + expect(gateway.tokenPurpose).toBe(gatewayModern ? 'delegated-workload' : undefined); + expect(gateway.exp! - gateway.iat!).toBe(gatewayModern ? 3600 : 5 * 365 * 24 * 3600); +}); diff --git a/apps/web/src/lib/config.server.ts b/apps/web/src/lib/config.server.ts index 27a9a111cc..22f1d97a16 100644 --- a/apps/web/src/lib/config.server.ts +++ b/apps/web/src/lib/config.server.ts @@ -63,6 +63,34 @@ export const INTERNAL_API_SECRET = getEnvVariable('INTERNAL_API_SECRET'); export function isBoundedInternalServiceTokenIssuanceEnabled(): boolean { return getEnvVariable('BOUNDED_INTERNAL_SERVICE_TOKENS_ENABLED') === 'true'; } +export function isNativeResourceCredentialIssuanceEnabled(): boolean { + return ( + getEnvVariable('NATIVE_RESOURCE_TOKENS_ENABLED') === 'true' && + isSharedResourceTokenIssuanceEnabled() + ); +} +export function isSharedResourceTokenIssuanceEnabled(): boolean { + return getEnvVariable('SHARED_RESOURCE_TOKENS_ENABLED') === 'true'; +} +const resourceTokenFamilyFlags = { + 'cloud-agent-next': 'CLOUD_AGENT_RESOURCE_TOKENS_ENABLED', + gastown: 'GASTOWN_RESOURCE_TOKENS_ENABLED', + wasteland: 'WASTELAND_RESOURCE_TOKENS_ENABLED', + chat: 'CHAT_RESOURCE_TOKENS_ENABLED', + 'delegated-resource': 'DELEGATED_RESOURCE_TOKENS_ENABLED', + 'workflow-gateway': 'WORKFLOW_GATEWAY_RESOURCE_TOKENS_ENABLED', + benchmark: 'BENCHMARK_RESOURCE_TOKENS_ENABLED', +} as const; + +export type ResourceTokenFamily = keyof typeof resourceTokenFamilyFlags; + +export function isResourceTokenIssuanceEnabled(family: ResourceTokenFamily): boolean { + return ( + isSharedResourceTokenIssuanceEnabled() && + getEnvVariable(resourceTokenFamilyFlags[family]) === 'true' + ); +} + export const USER_DATA_EXPORT_WORKER_URL = getEnvVariable('USER_DATA_EXPORT_WORKER_URL') || (process.env.NODE_ENV === 'development' ? 'http://127.0.0.1:8818' : ''); diff --git a/apps/web/src/lib/token.test.ts b/apps/web/src/lib/token.test.ts index fa0e7d37b2..aeff0fa39d 100644 --- a/apps/web/src/lib/token.test.ts +++ b/apps/web/src/lib/token.test.ts @@ -313,6 +313,21 @@ describe('Token Functions', () => { // validateAuthorizationHeader should NOT return organizationId - that's handled at a higher level }); + it('fails closed when a cloud-agent-next runtime bearer is used directly', () => { + const token = signedToken({ + aud: KILO_API_AUDIENCE, + runtimeAuthorization: { + id: '11111111-1111-4111-8111-111111111111', + resourceKind: 'cloud-agent-next', + resourceId: 'agent_123', + }, + }); + + expect( + validateAuthorizationHeader(new Headers({ authorization: `Bearer ${token}` })).error + ).toMatch(/^Invalid token \([a-f0-9-]+\)$/); + }); + it('should return error when authorization header is missing', () => { const headers = new Headers(); diff --git a/apps/web/src/lib/tokens.ts b/apps/web/src/lib/tokens.ts index a331acf211..e5d031c0d5 100644 --- a/apps/web/src/lib/tokens.ts +++ b/apps/web/src/lib/tokens.ts @@ -8,6 +8,7 @@ import { GITHUB_USER_AUTHORIZATION_DISCONNECT_AUDIENCE, GITLAB_CREDENTIAL_BROKER_AUDIENCE, KILO_API_AUDIENCE, + KILO_GATEWAY_AUDIENCE, SESSION_INGEST_AUDIENCE, SESSION_INGEST_USER_DELETION_AUDIENCE, USER_DATA_EXPORT_AUDIENCE, @@ -16,10 +17,15 @@ import { buildModernKiloTokenPayload, isKiloResourceAudienceAllowed, } from '@kilocode/worker-utils/kilo-token-policy'; +import { CloudAgentNextRuntimeAuthorizationClaimSchema } from '@kilocode/worker-utils/runtime-proxy-attestation'; import type { OrganizationRole } from '@/lib/organizations/organization-types'; import jwt from 'jsonwebtoken'; import { warnExceptInTest } from '@/lib/utils.server'; -import { isBoundedInternalServiceTokenIssuanceEnabled, NEXTAUTH_SECRET } from '@/lib/config.server'; +import { + isBoundedInternalServiceTokenIssuanceEnabled, + isResourceTokenIssuanceEnabled, + NEXTAUTH_SECRET, +} from '@/lib/config.server'; export { BITBUCKET_REPOSITORY_LIST_AUDIENCE } from '@kilocode/worker-utils/internal-service-token-audiences'; @@ -193,6 +199,7 @@ export type JWTTokenPayload = { kiloUserId: string; version: number; apiTokenPepper?: string; + runtimeAuthorization?: unknown; } & JWTTokenExtraPayload; function tryJwtVerify(token: string) { @@ -209,7 +216,7 @@ function tryJwtVerify(token: string) { export function validateAuthorizationHeader( headers: Headers, - options?: { expectedAudience?: string } + options?: { expectedAudience?: string; runtimeProxyAttestationVerified?: boolean } ) { const traceability_logging_id = crypto.randomUUID(); const authHeader = headers.get('authorization'); @@ -244,6 +251,21 @@ export function validateAuthorizationHeader( return { error: `Token version outdated, please re-authenticate (${traceability_logging_id})` }; } + if ( + typeof payload.runtimeAuthorization === 'object' && + payload.runtimeAuthorization !== null && + 'resourceKind' in payload.runtimeAuthorization && + payload.runtimeAuthorization.resourceKind === 'cloud-agent-next' + ) { + const runtimeAuthorization = CloudAgentNextRuntimeAuthorizationClaimSchema.safeParse( + payload.runtimeAuthorization + ); + if (!runtimeAuthorization.success || !options?.runtimeProxyAttestationVerified) { + warnExceptInTest(`Invalid token (${traceability_logging_id})`); + return { error: `Invalid token (${traceability_logging_id})` }; + } + } + return { kiloUserId: payload.kiloUserId, apiTokenPepper: payload.apiTokenPepper, @@ -260,3 +282,104 @@ export function validateAuthorizationHeader( export function generateCloudAgentToken(user: User) { return generateApiToken(user, { tokenSource: 'cloud-agent' }); } + +export function generateCloudAgentWorkflowToken( + user: User, + options: { + organizationId?: string; + tokenSource: string; + botId?: string; + createdOnPlatform?: string; + expiresIn: number; + authorizationUser?: User; + } +): string { + if (!isResourceTokenIssuanceEnabled('cloud-agent-next')) { + return generateApiToken( + user, + { + organizationId: options.organizationId, + tokenSource: options.tokenSource, + botId: options.botId, + createdOnPlatform: options.createdOnPlatform, + }, + { expiresIn: options.expiresIn } + ); + } + if (!user.api_token_pepper) { + throw new Error('Workflow control tokens require a current user pepper'); + } + const expiresIn = Math.min(options.expiresIn, 60 * 60); + if (expiresIn <= 0) { + throw new Error('Workflow control token expiry must be positive'); + } + const authorizationUser = options.authorizationUser ?? user; + if (!authorizationUser.api_token_pepper) { + throw new Error('Workflow control tokens require a current authorization pepper'); + } + const issuedAt = Math.floor(Date.now() / 1000); + const payload = buildModernKiloTokenPayload({ + userId: user.id, + pepper: user.api_token_pepper, + env: process.env.NODE_ENV, + audience: 'cloud-agent-next', + issuedAt, + expiresAt: issuedAt + expiresIn, + tokenPurpose: 'internal-service', + credentialExchange: false, + extra: { + organizationId: options.organizationId, + tokenSource: options.tokenSource, + botId: options.botId, + createdOnPlatform: options.createdOnPlatform, + runtimeAdmission: { + source: 'automation', + authorizationUserId: authorizationUser.id, + authorizationPepper: authorizationUser.api_token_pepper, + }, + }, + }); + return jwt.sign(payload, NEXTAUTH_SECRET, { algorithm: jwtSigningAlgorithm }); +} + +/** + * Generates a gateway credential for a server-side workflow. Unlike resource + * delegation from an HTTP request, queued workflows have no browser bearer + * token to delegate from. The workflow owner is therefore the authority and + * the token is deliberately limited to the gateway and one hour. + */ +export function generateWorkflowGatewayToken( + user: User, + options: { + organizationId?: string; + tokenSource: string; + expiresIn?: number; + } +): string { + if (!isResourceTokenIssuanceEnabled('workflow-gateway')) { + return generateApiToken(user, { tokenSource: options.tokenSource }); + } + if (!user.api_token_pepper) { + throw new Error('Workflow gateway tokens require a current user pepper'); + } + const expiresIn = Math.min(options.expiresIn ?? ONE_HOUR_IN_SECONDS, ONE_HOUR_IN_SECONDS); + if (expiresIn <= 0) { + throw new Error('Workflow gateway token expiry must be positive'); + } + const issuedAt = Math.floor(Date.now() / 1000); + const payload = buildModernKiloTokenPayload({ + userId: user.id, + pepper: user.api_token_pepper, + env: process.env.NODE_ENV, + audience: KILO_GATEWAY_AUDIENCE, + issuedAt, + expiresAt: issuedAt + expiresIn, + tokenPurpose: 'delegated-workload', + credentialExchange: false, + extra: { + organizationId: options.organizationId, + tokenSource: options.tokenSource, + }, + }); + return jwt.sign(payload, NEXTAUTH_SECRET, { algorithm: jwtSigningAlgorithm }); +} diff --git a/apps/web/src/lib/user/server.test.ts b/apps/web/src/lib/user/server.test.ts index c0da277cf9..80f6351b40 100644 --- a/apps/web/src/lib/user/server.test.ts +++ b/apps/web/src/lib/user/server.test.ts @@ -51,6 +51,7 @@ import { createTestOrganization } from '@/tests/helpers/organization.helper'; import { insertTestUser } from '@/tests/helpers/user.helper'; import { createCallerForUser } from '@/routers/test-utils'; import { generateApiToken, JWT_TOKEN_VERSION } from '@/lib/tokens'; +import { ORGANIZATION_ID_HEADER } from '@/lib/constants'; import { eq } from 'drizzle-orm'; import { v5 as uuidv5 } from 'uuid'; import jwt from 'jsonwebtoken'; @@ -611,6 +612,89 @@ describe('getUserFromAuth', () => { expect(apiResult.authFailedResponse?.status).toBe(401); }); + test('uses the signed organization instead of a conflicting request header', async () => { + const user = await insertTestUser({ api_token_pepper: 'signed-organization-pepper' }); + const signedOrganization = await createTestOrganization( + `Signed Organization ${crypto.randomUUID()}`, + user.id, + 0 + ); + const headerOrganization = await createTestOrganization( + `Header Organization ${crypto.randomUUID()}`, + user.id, + 0 + ); + const token = signPolicyClaims({ + version: JWT_TOKEN_VERSION, + kiloUserId: user.id, + apiTokenPepper: user.api_token_pepper, + env: process.env.NODE_ENV, + aud: KILO_GATEWAY_AUDIENCE, + organizationId: signedOrganization.id, + }); + mockHeaders.mockResolvedValue( + new Headers({ + Authorization: `Bearer ${token}`, + [ORGANIZATION_ID_HEADER]: headerOrganization.id, + }) + ); + + const result = await getUserFromAuth({ + adminOnly: false, + expectedAudience: KILO_GATEWAY_AUDIENCE, + }); + + expect(result.authFailedResponse).toBeNull(); + expect(result.organizationId).toBe(signedOrganization.id); + }); + + test('uses the signed organization when no request header is supplied', async () => { + const user = await insertTestUser({ api_token_pepper: 'signed-organization-no-header-pepper' }); + const organization = await createTestOrganization( + `Signed Organization ${crypto.randomUUID()}`, + user.id, + 0 + ); + const token = signPolicyClaims({ + version: JWT_TOKEN_VERSION, + kiloUserId: user.id, + apiTokenPepper: user.api_token_pepper, + env: process.env.NODE_ENV, + aud: KILO_GATEWAY_AUDIENCE, + organizationId: organization.id, + }); + mockHeaders.mockResolvedValue(new Headers({ Authorization: `Bearer ${token}` })); + + const result = await getUserFromAuth({ + adminOnly: false, + expectedAudience: KILO_GATEWAY_AUDIENCE, + }); + + expect(result.authFailedResponse).toBeNull(); + expect(result.organizationId).toBe(organization.id); + }); + + test('uses the request header for a token without a signed organization', async () => { + const user = await insertTestUser({ api_token_pepper: 'unbound-organization-pepper' }); + const organization = await createTestOrganization( + `Header Organization ${crypto.randomUUID()}`, + user.id, + 0 + ); + const token = generateApiToken(user); + mockHeaders.mockResolvedValue( + new Headers({ + Authorization: `Bearer ${token}`, + [ORGANIZATION_ID_HEADER]: organization.id, + }) + ); + + const result = await getUserFromAuth({ adminOnly: false }); + + expect(result.authFailedResponse).toBeNull(); + expect(result.organizationId).toBe(organization.id); + }); + test('allows API-token authentication for users from SSO-protected domains', async () => { const ssoDomain = `${crypto.randomUUID()}.example.com`; const user = await insertTestUser({ diff --git a/apps/web/src/lib/user/server.ts b/apps/web/src/lib/user/server.ts index e78aefcb0e..640d68c76f 100644 --- a/apps/web/src/lib/user/server.ts +++ b/apps/web/src/lib/user/server.ts @@ -1,6 +1,11 @@ import { getEnvVariable } from '@/lib/dotenvx'; import 'server-only'; import { validateAuthorizationHeader, JWT_TOKEN_VERSION } from '@/lib/tokens'; +import { + CloudAgentNextRuntimeAuthorizationClaimSchema, + RuntimeProxyAttestationAudienceSchema, + verifyRuntimeProxyAttestation, +} from '@kilocode/worker-utils/runtime-proxy-attestation'; import { NextResponse } from 'next/server'; import { cookies, headers } from 'next/headers'; @@ -1219,8 +1224,35 @@ async function resolveUserFromAuth( // all calls from the extension including the openrouter proxy call use this auth method // also val.town and other blessed API users who are given their own custom JWTs use this path if (headersList.get('Authorization')) { + const rawAuthorization = headersList.get('Authorization'); + const bearer = rawAuthorization?.match(/^Bearer (.+)$/i)?.[1]; + const decoded = bearer ? jwt.decode(bearer) : null; + const decodedPayload = decoded !== null && typeof decoded !== 'string' ? decoded : null; + const decodedRuntimeAuthorization = decodedPayload?.runtimeAuthorization; + const runtimeAuthorization = CloudAgentNextRuntimeAuthorizationClaimSchema.safeParse( + decodedRuntimeAuthorization + ); + const attestationAudience = RuntimeProxyAttestationAudienceSchema.safeParse( + opts.expectedAudience ?? KILO_API_AUDIENCE + ); + const runtimeProxyAttestationVerified = + bearer !== undefined && + runtimeAuthorization.success && + attestationAudience.success && + typeof decodedPayload?.kiloUserId === 'string' + ? await verifyRuntimeProxyAttestation({ + value: headersList.get('X-Kilo-Runtime-Proxy-Attestation'), + secret: NEXTAUTH_SECRET, + audience: attestationAudience.data, + userId: decodedPayload.kiloUserId, + authorizationId: runtimeAuthorization.data.id, + resourceId: runtimeAuthorization.data.resourceId, + bearer, + }) + : false; const authorizationValidationResult = validateAuthorizationHeader(headersList, { expectedAudience: opts.expectedAudience, + runtimeProxyAttestationVerified, }); if (authorizationValidationResult.error != undefined) { return authError(401, authorizationValidationResult.error, '?'); @@ -1234,7 +1266,12 @@ async function resolveUserFromAuth( ) { return authError(401, 'Invalid API token', user.id); } - const organizationId = headersList.get(ORGANIZATION_ID_HEADER) || undefined; + // A token-bound organization is signed; the request header is mutable. + // Legacy and personal tokens intentionally continue to use the header. + const organizationId = + authorizationValidationResult.organizationId ?? + headersList.get(ORGANIZATION_ID_HEADER) ?? + undefined; const internalApiUse = authorizationValidationResult.internalApiUse; const botId = authorizationValidationResult.botId; const tokenSource = authorizationValidationResult.tokenSource; diff --git a/apps/web/src/routers/test-utils.ts b/apps/web/src/routers/test-utils.ts index 3e90d770a7..e9bc3134a7 100644 --- a/apps/web/src/routers/test-utils.ts +++ b/apps/web/src/routers/test-utils.ts @@ -1,6 +1,7 @@ import { createCallerFactory } from '@/lib/trpc/init'; import { findUserById } from '@/lib/user'; import { rootRouter } from '@/routers/root-router'; +import { generateApiToken } from '@/lib/tokens'; const createCaller = createCallerFactory(rootRouter); @@ -10,5 +11,9 @@ export async function createCallerForUser(userId: string, opts?: { deviceSession if (!user) { throw new Error(`Test user not found: ${userId}`); } - return createCaller({ user, deviceSessionId: opts?.deviceSessionId }); + return createCaller({ + user, + deviceSessionId: opts?.deviceSessionId, + headersList: new Headers({ Authorization: `Bearer ${generateApiToken(user)}` }), + }); } diff --git a/packages/worker-utils/package.json b/packages/worker-utils/package.json index ab2332bf97..ef2afd39f7 100644 --- a/packages/worker-utils/package.json +++ b/packages/worker-utils/package.json @@ -13,6 +13,9 @@ "./kilo-token-auth": "./src/kilo-token-auth.ts", "./kilo-token": "./src/kilo-token.ts", "./kilo-token-policy": "./src/kilo-token-policy.ts", + "./runtime-authorization": "./src/runtime-authorization.ts", + "./runtime-authorization-contract": "./src/runtime-authorization-contract.ts", + "./runtime-proxy-attestation": "./src/runtime-proxy-attestation.ts", "./kilo-auth-middleware": "./src/kilo-auth-middleware.ts", "./sandbox-id": "./src/sandbox-id.ts", "./hostname-label": "./src/hostname-label.ts", diff --git a/packages/worker-utils/src/kilo-auth-middleware.test.ts b/packages/worker-utils/src/kilo-auth-middleware.test.ts index 9ecedb49ab..5116ae3f8b 100644 --- a/packages/worker-utils/src/kilo-auth-middleware.test.ts +++ b/packages/worker-utils/src/kilo-auth-middleware.test.ts @@ -2,7 +2,7 @@ import { describe, expect, it } from 'vitest'; import { Hono, type Context } from 'hono'; import { SignJWT } from 'jose'; import { GASTOWN_AUDIENCE } from './internal-service-token-audiences'; -import { createKiloAuthMiddleware } from './kilo-auth-middleware'; +import { createKiloAuthMiddleware, type KiloAuthVariables } from './kilo-auth-middleware'; const TEST_SECRET = 'test-secret-that-is-long-enough-for-hs256'; @@ -11,17 +11,12 @@ const resolveSecret = async (binding: { get(): Promise } | string) => type TestEnv = { Bindings: { NEXTAUTH_SECRET?: string }; - Variables: { - kiloUserId: string; - kiloIsAdmin: boolean; - kiloApiTokenPepper: string | null; - kiloGastownAccess: boolean; - kiloOrgMemberships: { orgId: string; role: 'owner' | 'member' | 'billing_manager' }[]; - }; + Variables: KiloAuthVariables; }; function createApp() { let downstreamCalls = 0; + const authenticated: KiloAuthVariables[] = []; const app = new Hono(); const kiloAuthMiddleware = createKiloAuthMiddleware({ resolveSecret, @@ -31,6 +26,7 @@ function createApp() { app.use('/trpc/*', kiloAuthMiddleware); const handler = (c: Context) => { downstreamCalls += 1; + authenticated.push(c.var); return c.json({ kiloUserId: c.get('kiloUserId'), isAdmin: c.get('kiloIsAdmin'), @@ -41,7 +37,7 @@ function createApp() { }; app.get('/api/whoami', handler); app.get('/trpc/whoami', handler); - return { app, downstreamCalls: () => downstreamCalls }; + return { app, downstreamCalls: () => downstreamCalls, authenticated }; } async function signToken( @@ -95,6 +91,43 @@ describe('createKiloAuthMiddleware', () => { expect(downstreamCalls()).toBe(0); }); + it('captures verified raw claims and control token for authorization', async () => { + const { app, authenticated } = createApp(); + const token = await signToken({ + aud: GASTOWN_AUDIENCE, + tokenPurpose: 'human-api', + credentialExchange: false, + apiTokenPepper: 'current-pepper', + customRestriction: 'retain-for-authorization', + orgMemberships: [{ orgId: 'org-a', role: 'admin' }], + }); + const responses = await request(app, token); + expect(responses.map(response => response.status)).toEqual([200, 200]); + expect(authenticated).toHaveLength(2); + for (const variables of authenticated) { + expect(variables.kiloControlToken).toBe(token); + expect(variables.kiloUsesModernToken).toBe(true); + expect(variables.kiloTokenClaims).toMatchObject({ + tokenPurpose: 'human-api', + customRestriction: 'retain-for-authorization', + }); + expect(variables.kiloOrgMemberships).toEqual([{ orgId: 'org-a', role: 'admin' }]); + } + }); + + it.each([ + { tokenPurpose: 'human-api', credentialExchange: false }, + { aud: GASTOWN_AUDIENCE, tokenPurpose: 'human-api' }, + { aud: GASTOWN_AUDIENCE, tokenPurpose: 'unknown', credentialExchange: false }, + { aud: GASTOWN_AUDIENCE, tokenPurpose: 'delegated-workload', credentialExchange: true }, + ])('rejects invalid modern claims before authorization: %j', async claims => { + const { app, downstreamCalls, authenticated } = createApp(); + const responses = await request(app, await signToken(claims)); + expect(responses.map(response => response.status)).toEqual([401, 401]); + expect(downstreamCalls()).toBe(0); + expect(authenticated).toEqual([]); + }); + it('accepts legacy tokens without an audience or dates', async () => { const { app, downstreamCalls } = createApp(); const responses = await request(app, await signToken({}, { dates: false })); diff --git a/packages/worker-utils/src/kilo-auth-middleware.ts b/packages/worker-utils/src/kilo-auth-middleware.ts index 0362e83d06..6eddfa0eeb 100644 --- a/packages/worker-utils/src/kilo-auth-middleware.ts +++ b/packages/worker-utils/src/kilo-auth-middleware.ts @@ -1,9 +1,10 @@ +import { decodeJwt, type JWTPayload } from 'jose'; import { createMiddleware } from 'hono/factory'; import type { MiddlewareHandler } from 'hono'; import { extractBearerToken } from './extract-bearer-token.js'; import type { KiloTokenPayload } from './kilo-token.js'; import type { KiloResourceAudiencePolicy } from './kilo-token-policy.js'; -import { verifyKiloTokenForResource } from './kilo-token-policy.js'; +import { verifyKiloTokenForPolicy, verifyKiloTokenForResource } from './kilo-token-policy.js'; import { resError } from './res.js'; /** @@ -13,12 +14,12 @@ import { resError } from './res.js'; */ export type SecretBinding = { get(): Promise } | string; -export type KiloAuthOrgMembership = { - orgId: string; - role: 'owner' | 'member' | 'billing_manager'; -}; +export type KiloAuthOrgMembership = NonNullable[number]; export type KiloAuthVariables = { + kiloControlToken: string; + kiloUsesModernToken: boolean; + kiloTokenClaims: JWTPayload; kiloUserId: string; kiloIsAdmin: boolean; kiloApiTokenPepper: string | null; @@ -74,6 +75,17 @@ export function createKiloAuthMiddleware( try { const payload = await verifyKiloTokenForResource(token, secret, audiencePolicy); + // Preserve claims stripped by the legacy payload schema for downstream authorization. + // Decode only after signature and resource audience verification succeeds. + const claims = decodeJwt(token); + const usesModernToken = + claims.tokenPurpose !== undefined || claims.credentialExchange !== undefined; + if (usesModernToken) { + await verifyKiloTokenForPolicy(token, secret, audiencePolicy); + } + c.set('kiloControlToken', token); + c.set('kiloUsesModernToken', usesModernToken); + c.set('kiloTokenClaims', claims); c.set('kiloUserId', payload.kiloUserId); c.set('kiloIsAdmin', payload.isAdmin === true); c.set('kiloApiTokenPepper', payload.apiTokenPepper ?? null); diff --git a/packages/worker-utils/src/kilo-token-policy.test.ts b/packages/worker-utils/src/kilo-token-policy.test.ts index d0b8a9d032..0a8d623b31 100644 --- a/packages/worker-utils/src/kilo-token-policy.test.ts +++ b/packages/worker-utils/src/kilo-token-policy.test.ts @@ -22,6 +22,8 @@ import { GITHUB_USER_ACCESS_TOKEN_AUDIENCE, USER_DATA_EXPORT_AUDIENCE, SESSION_INGEST_USER_DELETION_AUDIENCE, + KILO_API_AUDIENCE, + KILO_GATEWAY_AUDIENCE, } from './internal-service-token-audiences.js'; const SECRET = 'synthetic-policy-test-secret-at-least-32-chars'; @@ -855,7 +857,7 @@ describe('buildModernKiloTokenPayload and compatibility', () => { }; type ExpectedReadonlyOrganizationMemberships = readonly { readonly orgId: string; - readonly role: 'owner' | 'member' | 'billing_manager'; + readonly role: 'owner' | 'admin' | 'member' | 'billing_manager'; }[]; expectTypeOf().toEqualTypeOf(); @@ -881,8 +883,52 @@ describe('buildModernKiloTokenPayload and compatibility', () => { ).resolves.toMatchObject({ userId: 'synthetic-user' }); }); + it('supports unique audience arrays only for non-exchangeable modern tokens', () => { + expect( + buildModernKiloTokenPayload({ + userId: 'synthetic-user', + audience: ['kilo-api', 'kilo-gateway'], + issuedAt: NOW_SECONDS, + expiresAt: NOW_SECONDS + 60, + tokenPurpose: 'delegated-workload', + credentialExchange: false, + }).aud + ).toEqual(['kilo-api', 'kilo-gateway']); + expect(() => + buildModernKiloTokenPayload({ + userId: 'synthetic-user', + pepper: 'synthetic-pepper', + audience: ['kilo-api', 'kilo-gateway'], + issuedAt: NOW_SECONDS, + expiresAt: NOW_SECONDS + 60, + tokenPurpose: 'human-api', + credentialExchange: true, + }) + ).toThrow(); + }); + + it('preserves the ordered API and gateway audience set for delegated workloads', () => { + const payload = buildModernKiloTokenPayload({ + userId: 'synthetic-user', + pepper: 'synthetic-pepper', + audience: [KILO_API_AUDIENCE, KILO_GATEWAY_AUDIENCE], + issuedAt: NOW_SECONDS, + expiresAt: NOW_SECONDS + 60, + tokenPurpose: 'delegated-workload', + credentialExchange: false, + }); + + expect(payload.aud).toEqual([KILO_API_AUDIENCE, KILO_GATEWAY_AUDIENCE]); + expect(isKiloResourceAudienceAllowed(payload.aud, API_POLICY)).toBe(true); + expect( + isKiloResourceAudienceAllowed(payload.aud, { + audience: 'unrelated-audience', + mode: 'required', + }) + ).toBe(false); + }); + it.each([ - { audience: ['kilo-api'] }, { issuedAt: -1 }, { expiresAt: NOW_SECONDS }, { tokenPurpose: 'device-access', credentialExchange: true }, diff --git a/packages/worker-utils/src/kilo-token-policy.ts b/packages/worker-utils/src/kilo-token-policy.ts index 34079040b2..4af9b45672 100644 --- a/packages/worker-utils/src/kilo-token-policy.ts +++ b/packages/worker-utils/src/kilo-token-policy.ts @@ -47,14 +47,41 @@ const policyClaims = kiloTokenPayload message: 'Modern tokens require audience, purpose and exchange eligibility', }); } - if (claims.credentialExchange === true && claims.tokenPurpose !== 'human-api') { - ctx.addIssue({ - code: 'custom', - message: 'Only human API tokens may permit credential exchange', - path: ['credentialExchange'], - }); + } + if (claims.runtimeAdmission !== undefined && claims.runtimeAuthorization !== undefined) { + ctx.addIssue({ + code: 'custom', + message: 'Runtime admission and runtime authorization cannot coexist', + }); + } + if (claims.runtimeAdmission !== undefined) { + const permittedPurpose = + (claims.runtimeAdmission.source === 'user' && + (claims.tokenPurpose === 'human-api' || claims.tokenPurpose === 'device-access')) || + (claims.runtimeAdmission.source === 'automation' && + claims.tokenPurpose === 'internal-service'); + if ( + !permittedPurpose || + claims.credentialExchange !== false || + claims.apiTokenPepper === undefined || + claims.tokenPurpose === 'delegated-workload' + ) { + ctx.addIssue({ code: 'custom', message: 'Invalid runtime admission' }); } } + if ( + claims.runtimeAuthorization !== undefined && + (claims.tokenPurpose !== 'delegated-workload' || claims.credentialExchange !== false) + ) { + ctx.addIssue({ code: 'custom', message: 'Invalid runtime authorization reference' }); + } + if (claims.credentialExchange === true && claims.tokenPurpose !== 'human-api') { + ctx.addIssue({ + code: 'custom', + message: 'Only human API tokens may permit credential exchange', + path: ['credentialExchange'], + }); + } }); export type KiloTokenPolicyClaims = z.infer; @@ -242,7 +269,7 @@ export type ModernKiloTokenPurpose = z.infer; const modernClaims = policyClaims .safeExtend({ - aud: audienceName, + aud: audienceClaim, tokenPurpose: purposeClaim, credentialExchange: z.boolean(), }) @@ -253,7 +280,9 @@ const modernClaims = policyClaims .filter(([, value]) => value !== undefined) .map(([name]) => name); if ( - claims.aud !== KILO_API_AUDIENCE || + (typeof claims.aud === 'string' + ? claims.aud !== KILO_API_AUDIENCE + : claims.aud.length !== 1 || claims.aud[0] !== KILO_API_AUDIENCE) || claims.apiTokenPepper === undefined || !hasOnlyExchangeSafeClaims(serializedClaimNames) ) { @@ -282,8 +311,9 @@ export type SignModernKiloTokenParams = { expiresInSeconds: number; pepper?: string | null; env?: string; - audience: string; + audience: string | string[]; extra?: z.infer; + now?: Date; } & ModernKiloTokenSigningPurpose; export function buildModernKiloTokenPayload( @@ -291,14 +321,14 @@ export function buildModernKiloTokenPayload( userId: string; pepper?: string | null; env?: string; - audience: string; + audience: string | string[]; issuedAt: number; expiresAt: number; extra?: z.infer; } & ModernKiloTokenPurpose ): ModernKiloTokenClaims { const extra = modernTokenExtra.parse(params.extra ?? {}); - const audience = audienceName.parse(params.audience); + const audience = audienceClaim.parse(params.audience); const claims = { ...extra, version: KILO_TOKEN_VERSION, @@ -318,7 +348,9 @@ export async function signModernKiloToken( params: SignModernKiloTokenParams ): Promise<{ token: string; expiresAt: string }> { const expiresInSeconds = positiveSafeInteger.parse(params.expiresInSeconds); - const issuedAt = Math.floor(Date.now() / 1000); + const now = params.now ?? new Date(); + if (!Number.isFinite(now.getTime())) throw new Error('Invalid token issuance time'); + const issuedAt = Math.floor(now.getTime() / 1000); const expiresAt = issuedAt + expiresInSeconds; if (!Number.isSafeInteger(expiresAt)) { throw new Error('Token expiration exceeds the safe integer range'); diff --git a/packages/worker-utils/src/kilo-token.test.ts b/packages/worker-utils/src/kilo-token.test.ts index daf6525638..ee57f98bd4 100644 --- a/packages/worker-utils/src/kilo-token.test.ts +++ b/packages/worker-utils/src/kilo-token.test.ts @@ -144,6 +144,28 @@ describe('signKiloToken', () => { expect(kiloTokenPayload.parse(payload)).toEqual(payload); }); + + it('accepts admin organization roles and runtime claims', async () => { + const { token } = await signKiloToken({ + userId: 'user-runtime', + pepper: 'pepper-runtime', + secret: SECRET, + expiresInSeconds: 60, + extra: { + organizationRole: 'admin', + runtimeAdmission: { + source: 'user', + authorizationUserId: 'user-runtime', + authorizationPepper: 'pepper-runtime', + }, + }, + }); + + await expect(verifyKiloToken(token, SECRET)).resolves.toMatchObject({ + organizationRole: 'admin', + runtimeAdmission: { source: 'user' }, + }); + }); }); describe('verifyKiloToken', () => { diff --git a/packages/worker-utils/src/kilo-token.ts b/packages/worker-utils/src/kilo-token.ts index c4c030e21f..71e9b163d2 100644 --- a/packages/worker-utils/src/kilo-token.ts +++ b/packages/worker-utils/src/kilo-token.ts @@ -2,6 +2,17 @@ import { SignJWT, jwtVerify } from 'jose'; import { z } from 'zod'; export const KILO_TOKEN_VERSION = 3; +const organizationRole = z.enum(['owner', 'admin', 'member', 'billing_manager']); +const runtimeAdmission = z.object({ + source: z.enum(['user', 'automation']), + authorizationUserId: z.string().min(1), + authorizationPepper: z.string().nullable(), +}); +const runtimeAuthorization = z.object({ + id: z.string().uuid(), + resourceKind: z.enum(['cloud-agent-next', 'gastown']), + resourceId: z.string().min(1), +}); /** * All known fields that can appear in a Kilo user JWT, sourced from @@ -20,16 +31,16 @@ export const kiloTokenPayload = z.object({ gastownAccess: z.boolean().optional(), botId: z.string().optional(), organizationId: z.string().optional(), - organizationRole: z.enum(['owner', 'member', 'billing_manager']).optional(), + organizationRole: organizationRole.optional(), internalApiUse: z.boolean().optional(), createdOnPlatform: z.string().optional(), tokenSource: z.string().optional(), deviceAuthRequestCode: z.string().optional(), deviceSessionId: z.string().optional(), // Org memberships (baked into gastown tokens to avoid DB lookups) - orgMemberships: z - .array(z.object({ orgId: z.string(), role: z.enum(['owner', 'member', 'billing_manager']) })) - .optional(), + orgMemberships: z.array(z.object({ orgId: z.string(), role: organizationRole })).optional(), + runtimeAdmission: runtimeAdmission.optional(), + runtimeAuthorization: runtimeAuthorization.optional(), // Standard JWT claims iat: z.number().optional(), exp: z.number().optional(), @@ -55,6 +66,8 @@ export type SignKiloTokenExtra = Pick< | 'deviceAuthRequestCode' | 'deviceSessionId' | 'orgMemberships' + | 'runtimeAdmission' + | 'runtimeAuthorization' >; export async function signKiloToken(params: { diff --git a/packages/worker-utils/src/runtime-authorization-contract.ts b/packages/worker-utils/src/runtime-authorization-contract.ts new file mode 100644 index 0000000000..3648b70c47 --- /dev/null +++ b/packages/worker-utils/src/runtime-authorization-contract.ts @@ -0,0 +1,82 @@ +import { z } from 'zod'; + +export const RuntimeResourceKindSchema = z.enum(['cloud-agent-next', 'gastown']); +export type RuntimeResourceKind = z.infer; + +export function runtimeAuthorizationMaximumLifetimeMs(resourceKind: RuntimeResourceKind): number { + switch (resourceKind) { + case 'cloud-agent-next': + return 24 * 60 * 60_000; + case 'gastown': + return 30 * 24 * 60 * 60_000; + default: { + const exhaustive: never = resourceKind; + return exhaustive; + } + } +} + +export const RuntimeAdmissionSchema = z.object({ + source: z.enum(['user', 'automation']), + authorizationUserId: z.string().min(1), + authorizationPepper: z.string().nullable(), +}); +export type RuntimeAdmission = z.infer; + +const pepperDigest = z.string().regex(/^[a-f0-9]{64}$/); +const nullablePepperDigest = z.union([pepperDigest, z.literal('null')]); + +export const RuntimeAuthorizationSchema = z + .object({ + version: z.literal(1), + id: z.string().uuid(), + resourceKind: RuntimeResourceKindSchema, + resourceId: z.string().min(1), + userId: z.string().min(1), + authorizationUserId: z.string().min(1), + organizationId: z.string().min(1).optional(), + issuedAt: z.string().datetime(), + delegationExpiresAt: z.string().datetime(), + state: z.enum(['active', 'revoked']), + bindings: z + .object({ + userPepperDigest: nullablePepperDigest, + authorizationPepperDigest: nullablePepperDigest, + userMembershipId: z.string().min(1).optional(), + authorizationUserMembershipId: z.string().min(1).optional(), + }) + .strict(), + source: z + .object({ + tokenSource: z.string().optional(), + botId: z.string().optional(), + createdOnPlatform: z.string().optional(), + admissionSource: z.enum(['user', 'automation']), + }) + .strict(), + env: z.string().optional(), + }) + .strict() + .superRefine((authorization, ctx) => { + const issuedAt = Date.parse(authorization.issuedAt); + const delegationExpiresAt = Date.parse(authorization.delegationExpiresAt); + if (delegationExpiresAt <= issuedAt) { + ctx.addIssue({ + code: 'custom', + message: 'Delegation expiration must follow issuance', + path: ['delegationExpiresAt'], + }); + } + if ( + delegationExpiresAt > + issuedAt + runtimeAuthorizationMaximumLifetimeMs(authorization.resourceKind) + ) { + ctx.addIssue({ + code: 'custom', + message: 'Delegation expiration exceeds the resource maximum lifetime', + path: ['delegationExpiresAt'], + }); + } + }); + +export type RuntimeAuthorization = z.infer; diff --git a/packages/worker-utils/src/runtime-authorization.test.ts b/packages/worker-utils/src/runtime-authorization.test.ts new file mode 100644 index 0000000000..42b213c076 --- /dev/null +++ b/packages/worker-utils/src/runtime-authorization.test.ts @@ -0,0 +1,346 @@ +import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; +import { decodeJwt } from 'jose'; +import { signModernKiloToken } from './kilo-token-policy.js'; +import { + createRuntimeAuthorization, + renewRuntimeAuthorization, + sealRuntimeAuthorization, + unsealRuntimeAuthorization, + RuntimeAuthorizationExpiredError, + RuntimeAuthorizationRevokedError, +} from './runtime-authorization.js'; +import type { + RuntimeAuthorizationAdapters, + RuntimeAuthorizationPrincipal, +} from './runtime-authorization.js'; + +const secret = 'runtime-authorization-test-secret-at-least-32'; +const principals = new Map([ + [ + 'user', + { + id: 'user', + apiTokenPepper: 'user-pepper', + blockedAt: null, + blockedReason: null, + isBot: false, + }, + ], + [ + 'bot', + { id: 'bot', apiTokenPepper: 'bot-pepper', blockedAt: null, blockedReason: null, isBot: true }, + ], +]); +const memberships = new Map([ + ['org:user', { id: 'membership-user', role: 'admin', organizationDeletedAt: null }], + ['org:bot', { id: 'membership-bot', role: 'member', organizationDeletedAt: null }], +]); + +function adapters(): RuntimeAuthorizationAdapters { + return { + getPrincipal: vi.fn(async ({ userId }) => principals.get(userId) ?? null), + getMembership: vi.fn( + async ({ organizationId, userId }) => memberships.get(`${organizationId}:${userId}`) ?? null + ), + }; +} + +async function admission( + options: { + userId?: string; + source?: 'user' | 'automation'; + audience?: string | string[]; + expiresInSeconds?: number; + } = {} +) { + const source = options.source ?? 'user'; + const userId = options.userId ?? 'user'; + return signModernKiloToken({ + userId, + pepper: principals.get(userId)?.apiTokenPepper ?? null, + secret, + expiresInSeconds: options.expiresInSeconds ?? 300, + audience: options.audience ?? 'cloud-agent-next', + tokenPurpose: source === 'user' ? 'human-api' : 'internal-service', + credentialExchange: false, + extra: { + organizationId: 'org', + runtimeAdmission: { + source, + authorizationUserId: 'user', + authorizationPepper: 'user-pepper', + }, + }, + }); +} + +describe('runtime authorization', () => { + beforeEach(() => { + principals.set('user', { + id: 'user', + apiTokenPepper: 'user-pepper', + blockedAt: null, + blockedReason: null, + isBot: false, + }); + principals.set('bot', { + id: 'bot', + apiTokenPepper: 'bot-pepper', + blockedAt: null, + blockedReason: null, + isBot: true, + }); + memberships.set('org:user', { + id: 'membership-user', + role: 'admin', + organizationDeletedAt: null, + }); + memberships.set('org:bot', { + id: 'membership-bot', + role: 'member', + organizationDeletedAt: null, + }); + }); + afterEach(() => vi.useRealTimers()); + it('creates a bound authorization and issues a runtime-only audience token', async () => { + const control = await admission(); + const result = await createRuntimeAuthorization({ + token: control.token, + secret, + connectionString: 'postgres://unused', + resourceKind: 'cloud-agent-next', + resourceId: 'session', + organizationId: 'org', + adapters: adapters(), + }); + + expect(result.authorization.bindings.userPepperDigest).not.toContain('user-pepper'); + expect(decodeJwt(result.token)).toMatchObject({ + aud: ['kilo-api', 'kilo-gateway', 'session-ingest'], + tokenPurpose: 'delegated-workload', + credentialExchange: false, + runtimeAuthorization: { id: result.authorization.id, resourceId: 'session' }, + }); + }); + + it('sets a fixed resource-specific delegation deadline independently of control expiry', async () => { + vi.useFakeTimers(); + const now = new Date('2026-01-01T00:00:00.000Z'); + vi.setSystemTime(now); + const cloudAgent = await createRuntimeAuthorization({ + token: (await admission({ expiresInSeconds: 60 })).token, + secret, + connectionString: 'postgres://unused', + resourceKind: 'cloud-agent-next', + resourceId: 'session', + organizationId: 'org', + adapters: adapters(), + now, + }); + const gastown = await createRuntimeAuthorization({ + token: (await admission({ audience: 'gastown', expiresInSeconds: 60 })).token, + secret, + connectionString: 'postgres://unused', + resourceKind: 'gastown', + resourceId: 'town', + organizationId: 'org', + adapters: adapters(), + now, + }); + + expect(cloudAgent.authorization.delegationExpiresAt).toBe('2026-01-02T00:00:00.000Z'); + expect(gastown.authorization.delegationExpiresAt).toBe('2026-01-31T00:00:00.000Z'); + expect(decodeJwt(cloudAgent.token).exp).toBe(Date.UTC(2026, 0, 1, 1) / 1000); + }); + + it('caps runtime-token expiration to the positive duration remaining before the deadline', async () => { + vi.useFakeTimers(); + vi.setSystemTime(new Date('2026-01-01T00:00:00.000Z')); + const created = await createRuntimeAuthorization({ + token: (await admission()).token, + secret, + connectionString: 'postgres://unused', + resourceKind: 'cloud-agent-next', + resourceId: 'session', + organizationId: 'org', + adapters: adapters(), + }); + const renewalTime = new Date('2026-01-01T23:40:00.000Z'); + vi.setSystemTime(renewalTime); + + const renewed = await renewRuntimeAuthorization({ + authorization: created.authorization, + secret, + connectionString: 'postgres://unused', + adapters: adapters(), + now: renewalTime, + }); + + expect(decodeJwt(renewed.token).exp).toBe(Date.UTC(2026, 0, 2) / 1000); + expect(renewed.expiresAt).toBe('2026-01-02T00:00:00.000Z'); + }); + + it('rejects renewal at and after the fixed delegation deadline', async () => { + vi.useFakeTimers(); + const createdAt = new Date('2026-01-01T00:00:00.000Z'); + vi.setSystemTime(createdAt); + const created = await createRuntimeAuthorization({ + token: (await admission()).token, + secret, + connectionString: 'postgres://unused', + resourceKind: 'cloud-agent-next', + resourceId: 'session', + organizationId: 'org', + adapters: adapters(), + now: createdAt, + }); + const deadline = new Date(created.authorization.delegationExpiresAt); + + await expect( + renewRuntimeAuthorization({ + authorization: created.authorization, + secret, + connectionString: 'postgres://unused', + adapters: adapters(), + now: deadline, + }) + ).rejects.toBeInstanceOf(RuntimeAuthorizationExpiredError); + await expect( + renewRuntimeAuthorization({ + authorization: created.authorization, + secret, + connectionString: 'postgres://unused', + adapters: adapters(), + now: new Date(deadline.getTime() + 1), + }) + ).rejects.toBeInstanceOf(RuntimeAuthorizationExpiredError); + }); + + it('rejects normal, multi-audience, expired, and mismatched control tokens', async () => { + const valid = await admission(); + const common = { + secret, + connectionString: 'postgres://unused', + resourceKind: 'cloud-agent-next' as const, + resourceId: 'session', + organizationId: 'org', + adapters: adapters(), + }; + await expect( + createRuntimeAuthorization({ ...common, token: valid.token.replace(/[^.]+$/, 'tampered') }) + ).rejects.toThrow(); + await expect( + createRuntimeAuthorization({ + ...common, + token: (await admission({ audience: ['cloud-agent-next', 'kilo-api'] })).token, + }) + ).rejects.toThrow(); + const ordinary = await signModernKiloToken({ + userId: 'user', + pepper: 'user-pepper', + secret, + expiresInSeconds: 300, + audience: 'cloud-agent-next', + tokenPurpose: 'human-api', + credentialExchange: false, + }); + await expect( + createRuntimeAuthorization({ ...common, token: ordinary.token }) + ).rejects.toThrow(); + await expect( + createRuntimeAuthorization({ + ...common, + token: (await admission({ audience: 'gastown' })).token, + }) + ).rejects.toThrow(); + const expiring = await admission({ expiresInSeconds: 60 }); + vi.setSystemTime(new Date(Date.now() + 61_000)); + await expect( + createRuntimeAuthorization({ ...common, token: expiring.token }) + ).rejects.toThrow(); + }); + + it('binds renewals to current peppers, blocks, and membership identities', async () => { + const result = await createRuntimeAuthorization({ + token: (await admission({ userId: 'bot', source: 'automation' })).token, + secret, + connectionString: 'postgres://unused', + resourceKind: 'cloud-agent-next', + resourceId: 'session', + organizationId: 'org', + adapters: adapters(), + }); + principals.set('bot', { + id: 'bot', + apiTokenPepper: 'rotated', + blockedAt: null, + blockedReason: null, + isBot: true, + }); + await expect( + renewRuntimeAuthorization({ + authorization: result.authorization, + secret, + connectionString: 'postgres://unused', + adapters: adapters(), + }) + ).rejects.toBeInstanceOf(RuntimeAuthorizationRevokedError); + principals.set('bot', { + id: 'bot', + apiTokenPepper: 'bot-pepper', + blockedAt: null, + blockedReason: null, + isBot: true, + }); + memberships.set('org:bot', { + id: 're-added-membership', + role: 'member', + organizationDeletedAt: null, + }); + await expect( + renewRuntimeAuthorization({ + authorization: result.authorization, + secret, + connectionString: 'postgres://unused', + adapters: adapters(), + }) + ).rejects.toBeInstanceOf(RuntimeAuthorizationRevokedError); + }); + + it('seals only the exact intended runtime record', async () => { + const result = await createRuntimeAuthorization({ + token: (await admission({ audience: 'gastown' })).token, + secret, + connectionString: 'postgres://unused', + resourceKind: 'gastown', + resourceId: 'town', + organizationId: 'org', + adapters: adapters(), + }); + const sealed = await sealRuntimeAuthorization(result.authorization, secret); + await expect( + unsealRuntimeAuthorization(sealed, secret, { + resourceKind: 'gastown', + resourceId: 'town', + userId: 'user', + organizationId: 'org', + }) + ).resolves.toEqual(result.authorization); + await expect( + unsealRuntimeAuthorization(sealed, secret, { + resourceKind: 'gastown', + resourceId: 'other', + userId: 'user', + organizationId: 'org', + }) + ).rejects.toThrow(); + await expect( + unsealRuntimeAuthorization(`${sealed}x`, secret, { + resourceKind: 'gastown', + resourceId: 'town', + userId: 'user', + organizationId: 'org', + }) + ).rejects.toThrow(); + }); +}); diff --git a/packages/worker-utils/src/runtime-authorization.ts b/packages/worker-utils/src/runtime-authorization.ts new file mode 100644 index 0000000000..a1ed3aafb8 --- /dev/null +++ b/packages/worker-utils/src/runtime-authorization.ts @@ -0,0 +1,368 @@ +import { getWorkerDb, kilocode_users, organization_memberships, organizations } from '@kilocode/db'; +import { and, eq } from 'drizzle-orm'; +import { jwtVerify, SignJWT } from 'jose'; +import { z } from 'zod'; +import { signModernKiloToken, verifyKiloTokenForPolicy } from './kilo-token-policy.js'; +import type { RuntimeAuthorization } from './runtime-authorization-contract.js'; +import { + RuntimeAuthorizationSchema, + runtimeAuthorizationMaximumLifetimeMs, + RuntimeResourceKindSchema, +} from './runtime-authorization-contract.js'; +export { RuntimeAuthorizationSchema }; +export type { RuntimeAuthorization }; + +const membershipRole = z.enum(['owner', 'admin', 'member']); +const principalSchema = z.object({ + id: z.string().min(1), + apiTokenPepper: z.string().nullable(), + blockedAt: z.string().nullable(), + blockedReason: z.string().nullable(), + isBot: z.boolean(), +}); +const membershipSchema = z.object({ + id: z.string().min(1), + role: z.string(), + organizationDeletedAt: z.string().nullable(), +}); + +export type RuntimeAuthorizationPrincipal = z.infer; +export type RuntimeAuthorizationMembership = z.infer; +export type RuntimeAuthorizationAdapters = { + getPrincipal?: (input: { + connectionString: string; + userId: string; + }) => Promise; + getMembership?: (input: { + connectionString: string; + userId: string; + organizationId: string; + }) => Promise; +}; + +type CommonInput = { + secret: string; + connectionString: string; + adapters?: RuntimeAuthorizationAdapters; +}; + +export class RuntimeAuthorizationRevokedError extends Error { + constructor() { + super('Runtime authorization has been revoked'); + this.name = 'RuntimeAuthorizationRevokedError'; + } +} + +export class RuntimeAuthorizationExpiredError extends Error { + constructor() { + super('Runtime authorization delegation has expired'); + this.name = 'RuntimeAuthorizationExpiredError'; + } +} + +function currentDate(now?: Date): Date { + const value = now ?? new Date(); + if (!Number.isFinite(value.getTime())) throw new Error('Invalid runtime authorization time'); + return new Date(value.getTime()); +} + +async function digest(value: string | null): Promise { + if (value === null) return 'null'; + const bytes = new Uint8Array( + await crypto.subtle.digest('SHA-256', new TextEncoder().encode(value)) + ); + return Array.from(bytes, byte => byte.toString(16).padStart(2, '0')).join(''); +} + +function isBlocked(principal: RuntimeAuthorizationPrincipal): boolean { + return principal.blockedAt !== null || principal.blockedReason !== null; +} + +async function getPrincipal( + input: CommonInput, + userId: string +): Promise { + if (input.adapters?.getPrincipal) { + const principal = await input.adapters.getPrincipal({ + connectionString: input.connectionString, + userId, + }); + return principal === null ? null : principalSchema.parse(principal); + } + const db = getWorkerDb(input.connectionString); + const [row] = await db + .select({ + id: kilocode_users.id, + apiTokenPepper: kilocode_users.api_token_pepper, + blockedAt: kilocode_users.blocked_at, + blockedReason: kilocode_users.blocked_reason, + isBot: kilocode_users.is_bot, + }) + .from(kilocode_users) + .where(eq(kilocode_users.id, userId)); + return row ? principalSchema.parse(row) : null; +} + +async function getMembership( + input: CommonInput, + userId: string, + organizationId: string +): Promise { + if (input.adapters?.getMembership) { + const membership = await input.adapters.getMembership({ + connectionString: input.connectionString, + userId, + organizationId, + }); + return membership === null ? null : membershipSchema.parse(membership); + } + const db = getWorkerDb(input.connectionString); + const [row] = await db + .select({ + id: organization_memberships.id, + role: organization_memberships.role, + organizationDeletedAt: organizations.deleted_at, + }) + .from(organization_memberships) + .innerJoin(organizations, eq(organization_memberships.organization_id, organizations.id)) + .where( + and( + eq(organization_memberships.kilo_user_id, userId), + eq(organization_memberships.organization_id, organizationId) + ) + ); + return row ? membershipSchema.parse(row) : null; +} + +async function requireBindings( + input: CommonInput, + record: RuntimeAuthorization +): Promise<{ + user: RuntimeAuthorizationPrincipal; + authorizationUser: RuntimeAuthorizationPrincipal; +}> { + const user = await getPrincipal(input, record.userId); + const authorizationUser = await getPrincipal(input, record.authorizationUserId); + if ( + user === null || + authorizationUser === null || + isBlocked(user) || + isBlocked(authorizationUser) || + (await digest(user.apiTokenPepper)) !== record.bindings.userPepperDigest || + (await digest(authorizationUser.apiTokenPepper)) !== record.bindings.authorizationPepperDigest + ) + throw new RuntimeAuthorizationRevokedError(); + if (record.organizationId) { + const userMembership = await getMembership(input, user.id, record.organizationId); + const authorizationMembership = await getMembership( + input, + authorizationUser.id, + record.organizationId + ); + if ( + userMembership === null || + authorizationMembership === null || + userMembership.organizationDeletedAt !== null || + authorizationMembership.organizationDeletedAt !== null || + !membershipRole.safeParse(userMembership.role).success || + !membershipRole.safeParse(authorizationMembership.role).success || + userMembership.id !== record.bindings.userMembershipId || + authorizationMembership.id !== record.bindings.authorizationUserMembershipId + ) + throw new RuntimeAuthorizationRevokedError(); + } + return { user, authorizationUser }; +} + +export async function createRuntimeAuthorization( + input: CommonInput & { + token: string; + resourceKind: RuntimeAuthorization['resourceKind']; + resourceId: string; + organizationId?: string; + now?: Date; + } +): Promise<{ authorization: RuntimeAuthorization; token: string; expiresAt: string }> { + const resourceKind = RuntimeResourceKindSchema.parse(input.resourceKind); + const resourceId = z.string().min(1).parse(input.resourceId); + const auth = await verifyKiloTokenForPolicy(input.token, input.secret, { + audience: resourceKind, + mode: 'allow-legacy', + }); + const claims = auth.claims; + const admission = claims.runtimeAdmission; + const soleAudience = + typeof claims.aud === 'string' + ? claims.aud + : claims.aud?.length === 1 + ? claims.aud[0] + : undefined; + if ( + admission === undefined || + claims.runtimeAuthorization !== undefined || + soleAudience !== resourceKind || + claims.credentialExchange !== false || + claims.apiTokenPepper === undefined || + (admission.source === 'user' && + claims.tokenPurpose !== 'human-api' && + claims.tokenPurpose !== 'device-access') || + (admission.source === 'automation' && claims.tokenPurpose !== 'internal-service') || + claims.tokenPurpose === 'delegated-workload' || + claims.kiloUserId.length === 0 || + input.organizationId !== claims.organizationId + ) + throw new Error('Invalid runtime admission'); + const user = await getPrincipal(input, claims.kiloUserId); + const authorizationUser = await getPrincipal(input, admission.authorizationUserId); + if ( + user === null || + authorizationUser === null || + isBlocked(user) || + isBlocked(authorizationUser) || + user.apiTokenPepper !== claims.apiTokenPepper || + authorizationUser.apiTokenPepper !== admission.authorizationPepper || + (user.id !== authorizationUser.id && (admission.source !== 'automation' || !user.isBot)) + ) + throw new Error('Invalid runtime admission'); + let userMembershipId: string | undefined; + let authorizationUserMembershipId: string | undefined; + if (input.organizationId) { + const userMembership = await getMembership(input, user.id, input.organizationId); + const authorizationMembership = await getMembership( + input, + authorizationUser.id, + input.organizationId + ); + if ( + userMembership === null || + authorizationMembership === null || + userMembership.organizationDeletedAt !== null || + authorizationMembership.organizationDeletedAt !== null || + !membershipRole.safeParse(userMembership.role).success || + !membershipRole.safeParse(authorizationMembership.role).success + ) + throw new Error('Invalid runtime admission'); + userMembershipId = userMembership.id; + authorizationUserMembershipId = authorizationMembership.id; + } + const now = currentDate(input.now); + const authorization = RuntimeAuthorizationSchema.parse({ + version: 1, + id: crypto.randomUUID(), + resourceKind, + resourceId, + userId: user.id, + authorizationUserId: authorizationUser.id, + organizationId: input.organizationId, + issuedAt: now.toISOString(), + delegationExpiresAt: new Date( + now.getTime() + runtimeAuthorizationMaximumLifetimeMs(resourceKind) + ).toISOString(), + state: 'active', + bindings: { + userPepperDigest: await digest(user.apiTokenPepper), + authorizationPepperDigest: await digest(authorizationUser.apiTokenPepper), + userMembershipId, + authorizationUserMembershipId, + }, + source: { + tokenSource: claims.tokenSource, + botId: claims.botId, + createdOnPlatform: claims.createdOnPlatform, + admissionSource: admission.source, + }, + env: claims.env, + }); + const signed = await issueRuntimeToken(authorization, user.apiTokenPepper, input.secret, now); + return { authorization, ...signed }; +} + +async function issueRuntimeToken( + authorization: RuntimeAuthorization, + pepper: string | null, + secret: string, + now: Date +): Promise<{ token: string; expiresAt: string }> { + const maximumExpiration = Math.floor(Date.parse(authorization.delegationExpiresAt) / 1000); + const issuedAt = Math.floor(now.getTime() / 1000); + const expiresInSeconds = Math.min(60 * 60, maximumExpiration - issuedAt); + if (expiresInSeconds <= 0) throw new RuntimeAuthorizationExpiredError(); + const audiences = + authorization.resourceKind === 'cloud-agent-next' + ? ['kilo-api', 'kilo-gateway', 'session-ingest'] + : ['kilo-api', 'kilo-gateway']; + return signModernKiloToken({ + userId: authorization.userId, + pepper, + secret, + expiresInSeconds, + now, + audience: audiences, + tokenPurpose: 'delegated-workload', + credentialExchange: false, + env: authorization.env, + extra: { + organizationId: authorization.organizationId, + tokenSource: authorization.source.tokenSource, + botId: authorization.source.botId, + createdOnPlatform: authorization.source.createdOnPlatform, + runtimeAuthorization: { + id: authorization.id, + resourceKind: authorization.resourceKind, + resourceId: authorization.resourceId, + }, + }, + }); +} + +export async function renewRuntimeAuthorization( + input: CommonInput & { authorization: RuntimeAuthorization; now?: Date } +): Promise<{ token: string; expiresAt: string }> { + const authorization = RuntimeAuthorizationSchema.parse(input.authorization); + if (authorization.state !== 'active') throw new RuntimeAuthorizationRevokedError(); + const now = currentDate(input.now); + if (now.getTime() >= Date.parse(authorization.delegationExpiresAt)) { + throw new RuntimeAuthorizationExpiredError(); + } + const { user } = await requireBindings(input, authorization); + return issueRuntimeToken(authorization, user.apiTokenPepper, input.secret, now); +} + +const sealedRecord = z.object({ + type: z.literal('runtime-authorization-record'), + authorization: RuntimeAuthorizationSchema, +}); + +export async function sealRuntimeAuthorization( + record: RuntimeAuthorization, + secret: string +): Promise { + const authorization = RuntimeAuthorizationSchema.parse(record); + return new SignJWT({ type: 'runtime-authorization-record', authorization }) + .setProtectedHeader({ alg: 'HS256', typ: 'JWT' }) + .setAudience(`${authorization.resourceKind}:runtime-authorization-record`) + .setIssuedAt() + .setExpirationTime('5m') + .sign(new TextEncoder().encode(secret)); +} + +export async function unsealRuntimeAuthorization( + sealed: string, + secret: string, + expected: Pick +): Promise { + const { payload } = await jwtVerify(sealed, new TextEncoder().encode(secret), { + algorithms: ['HS256'], + }); + const parsed = sealedRecord.parse(payload); + const record = parsed.authorization; + if ( + record.resourceKind !== expected.resourceKind || + record.resourceId !== expected.resourceId || + record.userId !== expected.userId || + record.organizationId !== expected.organizationId || + payload.aud !== `${expected.resourceKind}:runtime-authorization-record` + ) + throw new Error('Runtime authorization seal binding mismatch'); + return record; +} diff --git a/packages/worker-utils/src/runtime-proxy-attestation.test.ts b/packages/worker-utils/src/runtime-proxy-attestation.test.ts new file mode 100644 index 0000000000..a3859c9e41 --- /dev/null +++ b/packages/worker-utils/src/runtime-proxy-attestation.test.ts @@ -0,0 +1,65 @@ +import { SignJWT } from 'jose'; +import { describe, expect, it } from 'vitest'; +import { + issueRuntimeProxyAttestation, + verifyRuntimeProxyAttestation, +} from './runtime-proxy-attestation.js'; + +const secret = 'runtime-proxy-attestation-secret'; +const now = new Date('2026-01-01T00:00:00.000Z'); +const input = { + secret, + audience: 'kilo-api' as const, + userId: 'usr_123', + authorizationId: '11111111-1111-4111-8111-111111111111', + resourceId: 'agent_123', + bearer: 'exact-runtime-bearer', + now, +}; + +describe('runtime proxy attestation', () => { + it('binds a short-lived proof to the exact bearer and runtime identity', async () => { + const proof = await issueRuntimeProxyAttestation(input); + await expect(verifyRuntimeProxyAttestation({ ...input, value: proof, now })).resolves.toBe( + true + ); + }); + + it.each([ + ['wrong bearer', { bearer: 'other-bearer' }], + ['wrong user', { userId: 'usr_other' }], + ['wrong authorization', { authorizationId: '22222222-2222-4222-8222-222222222222' }], + ['wrong resource', { resourceId: 'agent_other' }], + ['wrong audience', { audience: 'kilo-gateway' as const }], + ])('rejects a proof with a %s binding', async (_name, changed) => { + const proof = await issueRuntimeProxyAttestation(input); + await expect( + verifyRuntimeProxyAttestation({ ...input, ...changed, value: proof, now }) + ).resolves.toBe(false); + }); + + it('rejects forged and expired proofs', async () => { + const forged = await new SignJWT({ + ...input, + version: 1, + type: 'cloud-agent-next-runtime-proxy', + }) + .setProtectedHeader({ alg: 'HS256', typ: 'JWT' }) + .setIssuer('kilocode-runtime-proxy') + .setAudience('kilo-api') + .setIssuedAt(Math.floor(now.getTime() / 1000)) + .setExpirationTime(Math.floor(now.getTime() / 1000) + 30) + .sign(new TextEncoder().encode('wrong-secret')); + const proof = await issueRuntimeProxyAttestation(input); + await expect(verifyRuntimeProxyAttestation({ ...input, value: forged, now })).resolves.toBe( + false + ); + await expect( + verifyRuntimeProxyAttestation({ + ...input, + value: proof, + now: new Date(now.getTime() + 31_000), + }) + ).resolves.toBe(false); + }); +}); diff --git a/packages/worker-utils/src/runtime-proxy-attestation.ts b/packages/worker-utils/src/runtime-proxy-attestation.ts new file mode 100644 index 0000000000..f6c075666e --- /dev/null +++ b/packages/worker-utils/src/runtime-proxy-attestation.ts @@ -0,0 +1,125 @@ +import { jwtVerify, SignJWT } from 'jose'; +import { z } from 'zod'; + +export const RUNTIME_PROXY_ATTESTATION_HEADER = 'X-Kilo-Runtime-Proxy-Attestation'; +export const RUNTIME_PROXY_ATTESTATION_ISSUER = 'kilocode-runtime-proxy'; +export const RUNTIME_PROXY_ATTESTATION_TYPE = 'cloud-agent-next-runtime-proxy'; +export const RUNTIME_PROXY_ATTESTATION_VERSION = 1; +export const RUNTIME_PROXY_ATTESTATION_MAX_AGE_SECONDS = 30; + +export const RuntimeProxyAttestationAudienceSchema = z.enum([ + 'kilo-api', + 'kilo-gateway', + 'session-ingest', +]); +export type RuntimeProxyAttestationAudience = z.infer; + +export const CloudAgentNextRuntimeAuthorizationClaimSchema = z + .object({ + id: z.string().uuid(), + resourceKind: z.literal('cloud-agent-next'), + resourceId: z.string().min(1), + }) + .strict(); +export type CloudAgentNextRuntimeAuthorizationClaim = z.infer< + typeof CloudAgentNextRuntimeAuthorizationClaimSchema +>; + +const runtimeProxyAttestationClaims = z + .object({ + iss: z.literal(RUNTIME_PROXY_ATTESTATION_ISSUER), + type: z.literal(RUNTIME_PROXY_ATTESTATION_TYPE), + version: z.literal(RUNTIME_PROXY_ATTESTATION_VERSION), + aud: RuntimeProxyAttestationAudienceSchema, + userId: z.string().min(1), + authorizationId: z.string().uuid(), + resourceId: z.string().min(1), + bearerDigest: z.string().regex(/^[a-f0-9]{64}$/), + iat: z.number().int().nonnegative(), + exp: z.number().int().positive(), + }) + .strict() + .superRefine((claims, ctx) => { + if ( + claims.exp <= claims.iat || + claims.exp - claims.iat > RUNTIME_PROXY_ATTESTATION_MAX_AGE_SECONDS + ) { + ctx.addIssue({ code: 'custom', message: 'Invalid runtime proxy attestation lifetime' }); + } + }); + +export type RuntimeProxyAttestationClaims = z.infer; + +async function bearerDigest(bearer: string): Promise { + const digest = await crypto.subtle.digest('SHA-256', new TextEncoder().encode(bearer)); + return Array.from(new Uint8Array(digest), byte => byte.toString(16).padStart(2, '0')).join(''); +} + +function signingKey(secret: string): Uint8Array { + return new TextEncoder().encode(secret); +} + +export async function issueRuntimeProxyAttestation(input: { + secret: string; + audience: RuntimeProxyAttestationAudience; + userId: string; + authorizationId: string; + resourceId: string; + bearer: string; + now?: Date; +}): Promise { + const now = input.now ?? new Date(); + const issuedAt = Math.floor(now.getTime() / 1000); + if (!Number.isSafeInteger(issuedAt) || issuedAt < 0) + throw new Error('Invalid attestation issuance time'); + const claims = runtimeProxyAttestationClaims.parse({ + iss: RUNTIME_PROXY_ATTESTATION_ISSUER, + type: RUNTIME_PROXY_ATTESTATION_TYPE, + version: RUNTIME_PROXY_ATTESTATION_VERSION, + aud: RuntimeProxyAttestationAudienceSchema.parse(input.audience), + userId: input.userId, + authorizationId: input.authorizationId, + resourceId: input.resourceId, + bearerDigest: await bearerDigest(input.bearer), + iat: issuedAt, + exp: issuedAt + RUNTIME_PROXY_ATTESTATION_MAX_AGE_SECONDS, + }); + return new SignJWT(claims) + .setProtectedHeader({ alg: 'HS256', typ: 'JWT' }) + .setIssuedAt(claims.iat) + .setExpirationTime(claims.exp) + .setIssuer(claims.iss) + .setAudience(claims.aud) + .sign(signingKey(input.secret)); +} + +export async function verifyRuntimeProxyAttestation(input: { + value: string | null | undefined; + secret: string; + audience: RuntimeProxyAttestationAudience; + userId: string; + authorizationId: string; + resourceId: string; + bearer: string; + now?: Date; +}): Promise { + if (!input.value || input.value.length > 4096) return false; + try { + const { payload } = await jwtVerify(input.value, signingKey(input.secret), { + algorithms: ['HS256'], + issuer: RUNTIME_PROXY_ATTESTATION_ISSUER, + audience: input.audience, + currentDate: input.now, + }); + const claims = runtimeProxyAttestationClaims.parse(payload); + return ( + claims.aud === input.audience && + claims.userId === input.userId && + claims.authorizationId === input.authorizationId && + claims.resourceId === input.resourceId && + claims.bearerDigest === (await bearerDigest(input.bearer)) + ); + } catch { + return false; + } +} 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 3daa1a1244..9ba0d10b96 100644 --- a/services/session-ingest/src/middleware/kilo-jwt-auth.test.ts +++ b/services/session-ingest/src/middleware/kilo-jwt-auth.test.ts @@ -10,6 +10,10 @@ import { SESSION_INGEST_AUDIENCE, SESSION_INGEST_USER_DELETION_AUDIENCE, } from '@kilocode/worker-utils/internal-service-token-audiences'; +import { + issueRuntimeProxyAttestation, + RUNTIME_PROXY_ATTESTATION_HEADER, +} from '@kilocode/worker-utils/runtime-proxy-attestation'; import { kiloJwtAuthMiddleware, type KiloJwtAuthVariables } from './kilo-jwt-auth'; @@ -149,6 +153,26 @@ async function signModernInternalToken(audience: string): Promise { return token; } +async function signCloudAgentRuntimeToken(): Promise { + const { token } = await signModernKiloToken({ + userId: 'usr_123', + pepper: 'pepper-current', + secret: TEST_JWT_SECRET, + expiresInSeconds: 3600, + audience: SESSION_INGEST_AUDIENCE, + tokenPurpose: 'delegated-workload', + credentialExchange: false, + extra: { + runtimeAuthorization: { + id: '11111111-1111-4111-8111-111111111111', + resourceKind: 'cloud-agent-next', + resourceId: 'agent_123', + }, + }, + }); + return token; +} + async function signClaims( claims: Record, secret = TEST_JWT_SECRET @@ -183,17 +207,26 @@ function request( query?: string; ticketStore?: TicketStore; secret?: string | null; + runtimeProxyAttestation?: string; + env?: TestEnv; } = {} ) { const app = makeApp(); const url = `http://local${options.path ?? '/api/me'}${options.query ?? ''}`; const headers = new Headers(); if (token !== null) headers.set('Authorization', `Bearer ${token}`); + if (options.runtimeProxyAttestation) { + headers.set(RUNTIME_PROXY_ATTESTATION_HEADER, options.runtimeProxyAttestation); + } if (options.websocket) headers.set('Upgrade', 'websocket'); return { response: app.fetch( new Request(url, { method: options.method, headers }), - makeEnv('secret' in options ? (options.secret ?? null) : TEST_JWT_SECRET, options.ticketStore) + options.env ?? + makeEnv( + 'secret' in options ? (options.secret ?? null) : TEST_JWT_SECRET, + options.ticketStore + ) ), }; } @@ -255,6 +288,59 @@ describe('kiloJwtAuthMiddleware', () => { expect(await res.json()).toEqual({ user_id: 'usr_123' }); }); + it('requires a proxy attestation for a cloud-agent-next runtime bearer after pepper validation', async () => { + userRowByUserId.set('usr_123', { pepper: 'pepper-current', blockedReason: null }); + const token = await signCloudAgentRuntimeToken(); + + const missingProof = await request(token).response; + expect(missingProof.status).toBe(401); + await expect(missingProof.json()).resolves.toEqual({ + success: false, + error: 'Invalid runtime proxy attestation', + }); + const proof = await issueRuntimeProxyAttestation({ + secret: TEST_JWT_SECRET, + audience: SESSION_INGEST_AUDIENCE, + userId: 'usr_123', + authorizationId: '11111111-1111-4111-8111-111111111111', + resourceId: 'agent_123', + bearer: token, + }); + expect((await request(token, { runtimeProxyAttestation: proof }).response).status).toBe(200); + }); + + it.each([ + [ + 'rejects', + async (env: TestEnv) => { + env.NEXTAUTH_SECRET_PROD.get = async () => { + throw new Error('secrets store unavailable'); + }; + }, + ], + [ + 'returns no secret', + async (env: TestEnv) => { + env.NEXTAUTH_SECRET_PROD.get = async () => null; + }, + ], + ])( + 'returns sanitized retryable 503 when the runtime attestation secret lookup %s', + async (_name, alter) => { + userRowByUserId.set('usr_123', { pepper: 'pepper-current', blockedReason: null }); + const env = makeEnv(TEST_JWT_SECRET); + await alter(env); + + const response = await request(await signCloudAgentRuntimeToken(), { env }).response; + + expect(response.status).toBe(503); + await expect(response.json()).resolves.toEqual({ + success: false, + error: 'Service temporarily unavailable', + }); + } + ); + 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'); diff --git a/services/session-ingest/src/middleware/kilo-jwt-auth.ts b/services/session-ingest/src/middleware/kilo-jwt-auth.ts index 9399b1c23d..84292af233 100644 --- a/services/session-ingest/src/middleware/kilo-jwt-auth.ts +++ b/services/session-ingest/src/middleware/kilo-jwt-auth.ts @@ -1,10 +1,16 @@ import { createMiddleware } from 'hono/factory'; -import { extractBearerToken } from '@kilocode/worker-utils'; +import { extractBearerToken, getCachedSecret } from '@kilocode/worker-utils'; import { SESSION_INGEST_AUDIENCE, SESSION_INGEST_USER_DELETION_AUDIENCE, } from '@kilocode/worker-utils/internal-service-token-audiences'; import { verifyKiloBearerAgainstCurrentPepper } from '@kilocode/worker-utils/kilo-token-auth'; +import { + CloudAgentNextRuntimeAuthorizationClaimSchema, + RUNTIME_PROXY_ATTESTATION_HEADER, + verifyRuntimeProxyAttestation, +} from '@kilocode/worker-utils/runtime-proxy-attestation'; +import { decodeJwt } from 'jose'; import type { Env } from '../env'; @@ -94,6 +100,35 @@ export const kiloJwtAuthMiddleware = createMiddleware<{ return c.json({ success: false, error: 'Invalid or expired token' }, 401); } + const runtimeAuthorization = CloudAgentNextRuntimeAuthorizationClaimSchema.safeParse( + decodeJwt(token).runtimeAuthorization + ); + if (runtimeAuthorization.success) { + let attested: boolean; + try { + const secret = await getCachedSecret(c.env.NEXTAUTH_SECRET_PROD, 'NEXTAUTH_SECRET'); + attested = await verifyRuntimeProxyAttestation({ + value: c.req.header(RUNTIME_PROXY_ATTESTATION_HEADER), + secret, + audience: SESSION_INGEST_AUDIENCE, + userId: auth.userId, + authorizationId: runtimeAuthorization.data.id, + resourceId: runtimeAuthorization.data.resourceId, + bearer: token, + }); + } catch (error) { + console.error('Auth infrastructure failure', { + operation: 'runtime-proxy-attestation-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 (!attested) { + return c.json({ success: false, error: 'Invalid runtime proxy attestation' }, 401); + } + } + c.set('user_id', auth.userId); c.set('deletionAudience', deletionAudience); return next(); From f4a4857a1a6362acb66898856a1d2cb24dd73906 Mon Sep 17 00:00:00 2001 From: Florian Hines Date: Tue, 8 Sep 2026 12:03:39 -0500 Subject: [PATCH 2/9] feat(cloud-agent): extract scoped runtime credentials --- .../triggers/prepare-fix-payload.test.ts | 65 ++ .../auto-fix/triggers/prepare-fix-payload.ts | 9 +- .../triggers/prepare-triage-payload.test.ts | 64 ++ .../triggers/prepare-triage-payload.ts | 9 +- .../cloud-agent-next/worktree-chat.test.ts | 8 +- .../src/lib/cloud-agent-next/worktree-chat.ts | 9 +- .../triggers/prepare-review-payload.ts | 16 +- .../services/analysis-service.test.ts | 37 + .../analysis-service.token-source.test.ts | 79 +- .../services/analysis-service.ts | 42 +- apps/web/src/routers/app-builder-router.ts | 27 +- .../routers/app-builder-token-source.test.ts | 42 +- .../routers/cli-sessions-v2-router.test.ts | 4 + .../web/src/routers/cli-sessions-v2-router.ts | 36 +- .../routers/cli-sessions-v2-worktree.test.ts | 7 +- .../routers/cloud-agent-next-router.test.ts | 4 +- .../src/routers/cloud-agent-next-router.ts | 50 +- .../organization-app-builder-router.ts | 52 +- ...ganization-cloud-agent-next-router.test.ts | 4 +- .../organization-cloud-agent-next-router.ts | 111 ++- services/cloud-agent-next/.dev.vars.example | 3 + .../src/kilo-facade/user-kilo-facade.ts | 2 +- .../__fixtures__/runtime-url-normalization.ts | 27 + .../src/kilo/kilo-targets.test.ts | 20 + .../kilo/runtime-credential-proxy-routes.ts | 218 ++++++ .../src/kilo/wrapper-client.ts | 8 + .../src/persistence/CloudAgentSession.ts | 375 ++++++++++ .../src/persistence/SandboxControl.ts | 183 +++++ services/cloud-agent-next/src/router.test.ts | 41 +- .../src/router/handlers/session-execution.ts | 22 +- .../src/router/handlers/session-prepare.ts | 8 +- .../src/router/handlers/session-send.ts | 7 +- .../src/router/handlers/session-start.ts | 24 +- .../router/handlers/session-worktree.test.ts | 164 +++- .../src/router/handlers/session-worktree.ts | 115 ++- .../src/runtime-credential-proxy-rpc.test.ts | 357 +++++++++ .../src/runtime-credential-proxy-rpc.ts | 200 +++++ .../src/runtime-credential-proxy.test.ts | 390 ++++++++++ .../src/runtime-credential-proxy.ts | 328 ++++++++ .../src/sandbox-control/frames.ts | 2 + .../src/sandbox-control/lifecycle.test.ts | 76 ++ .../session-credentials.test.ts | 121 ++- .../sandbox-control/session-credentials.ts | 129 +++- .../src/sandbox-control/socket.test.ts | 20 +- .../src/sandbox-control/socket.ts | 9 + .../vercel-network-policy.test.ts | 143 ++++ .../sandbox-control/vercel-network-policy.ts | 124 +++- .../src/sandbox-session/SandboxSession.ts | 455 +++++++++++- .../src/sandbox-session/control-rpc.test.ts | 39 + .../src/sandbox-session/control-rpc.ts | 34 +- .../src/sandbox-session/session-client.ts | 6 + .../session-message-queue.test.ts | 162 ++++ .../src/sandbox-session/session-operation.ts | 12 +- .../src/sandbox-session/terminal-lifecycle.ts | 18 + .../src/server-stream-ticket.test.ts | 16 +- services/cloud-agent-next/src/server.test.ts | 702 +++++++++++++++++- services/cloud-agent-next/src/server.ts | 272 ++++++- .../src/session-service.test.ts | 102 +++ .../cloud-agent-next/src/session-service.ts | 100 ++- .../src/session/agent-runtime.ts | 8 + .../src/session/queue-message.test.ts | 62 ++ .../src/session/queue-message.ts | 83 ++- .../runtime-authorization-persistence.test.ts | 266 +++++++ .../runtime-authorization-persistence.ts | 146 ++++ .../src/session/session-message-state.ts | 6 + .../src/session/session-prepare.test.ts | 16 + .../src/session/session-registration.ts | 53 +- .../src/shared/sandbox-control-protocol.ts | 18 + .../src/shared/wrapper-bootstrap.ts | 55 ++ .../cloud-agent-next/src/terminal/access.ts | 1 + services/cloud-agent-next/src/types.ts | 1 + services/cloud-agent-next/test/e2e/README.md | 44 +- services/cloud-agent-next/test/e2e/client.ts | 4 + .../test/e2e/fake-llm-server.ts | 5 + .../cloud-agent-next/test/e2e/lifecycle.ts | 48 +- .../test/e2e/sandbox-control.ts | 29 +- .../test/e2e/smoke-cleanup.ts | 38 + services/cloud-agent-next/test/e2e/smoke.ts | 71 +- .../runtime-authorization-recovery.test.ts | 518 +++++++++++++ .../test/integration/sandbox-control.test.ts | 101 ++- .../session/idle-reconciliation.test.ts | 81 ++ .../test/unit/fake-llm-server.test.ts | 22 +- .../test/unit/smoke-cleanup.test.ts | 140 ++++ .../worktree-credential-refresh.test.ts | 6 +- .../worker-configuration.d.ts | 69 +- services/cloud-agent-next/wrangler.jsonc | 2 + services/cloud-agent-next/wrangler.test.jsonc | 3 + .../wrapper/src/control/apply-attach.test.ts | 75 +- .../wrapper/src/control/apply-attach.ts | 1 + .../control/control-event-transport.test.ts | 26 + .../src/control/control-event-transport.ts | 4 +- .../src/control/control-test-fixtures.ts | 16 +- .../src/control/delete-worktree.test.ts | 103 ++- .../wrapper/src/control/delete-worktree.ts | 89 ++- .../wrapper/src/control/main.ts | 22 +- .../src/control/native-observations.ts | 47 +- .../wrapper/src/control/operation-registry.ts | 12 +- .../control/sandbox-control-client.test.ts | 6 + .../src/control/sandbox-control-client.ts | 2 + .../control/sandbox-control-handlers.test.ts | 185 +++-- .../src/control/sandbox-control-handlers.ts | 79 +- .../control/session-operation-cleanup.test.ts | 10 +- .../src/control/terminal-runtime.test.ts | 56 +- .../wrapper/src/control/terminal-runtime.ts | 54 +- .../worktree-mutation-notifications.test.ts | 35 +- .../worktree-mutation-notifications.ts | 6 +- .../src/control/worktree-runtime-cleanup.ts | 4 +- .../src/control/worktree-runtime.test.ts | 435 ++++++++--- .../wrapper/src/control/worktree-runtime.ts | 219 ++++-- .../wrapper/src/session-bootstrap.test.ts | 51 ++ 110 files changed, 8521 insertions(+), 621 deletions(-) create mode 100644 apps/web/src/lib/auto-fix/triggers/prepare-fix-payload.test.ts create mode 100644 apps/web/src/lib/auto-triage/triggers/prepare-triage-payload.test.ts create mode 100644 services/cloud-agent-next/src/kilo/__fixtures__/runtime-url-normalization.ts create mode 100644 services/cloud-agent-next/src/kilo/runtime-credential-proxy-routes.ts create mode 100644 services/cloud-agent-next/src/runtime-credential-proxy-rpc.test.ts create mode 100644 services/cloud-agent-next/src/runtime-credential-proxy-rpc.ts create mode 100644 services/cloud-agent-next/src/runtime-credential-proxy.test.ts create mode 100644 services/cloud-agent-next/src/runtime-credential-proxy.ts create mode 100644 services/cloud-agent-next/src/sandbox-session/control-rpc.test.ts create mode 100644 services/cloud-agent-next/src/session/runtime-authorization-persistence.test.ts create mode 100644 services/cloud-agent-next/src/session/runtime-authorization-persistence.ts create mode 100644 services/cloud-agent-next/test/e2e/smoke-cleanup.ts create mode 100644 services/cloud-agent-next/test/integration/runtime-authorization-recovery.test.ts create mode 100644 services/cloud-agent-next/test/unit/smoke-cleanup.test.ts diff --git a/apps/web/src/lib/auto-fix/triggers/prepare-fix-payload.test.ts b/apps/web/src/lib/auto-fix/triggers/prepare-fix-payload.test.ts new file mode 100644 index 0000000000..03d7382171 --- /dev/null +++ b/apps/web/src/lib/auto-fix/triggers/prepare-fix-payload.test.ts @@ -0,0 +1,65 @@ +const mockGetFixTicketById = jest.fn(); + +jest.mock('@/lib/drizzle', () => ({ + db: { + select: () => ({ + from: () => ({ + where: () => ({ + limit: () => [{ id: 'user-1', api_token_pepper: 'pepper' }], + }), + }), + }), + }, +})); + +jest.mock('@/lib/tokens', () => ({ + generateCloudAgentWorkflowToken: jest.fn(() => 'workflow-token'), + TOKEN_EXPIRY: { default: 3600 }, +})); + +jest.mock('../db/fix-tickets', () => ({ + getFixTicketById: (...args: unknown[]) => mockGetFixTicketById(...args), +})); + +jest.mock('@sentry/nextjs', () => ({ captureException: jest.fn() })); + +import { prepareFixPayload } from './prepare-fix-payload'; +import { generateCloudAgentWorkflowToken } from '@/lib/tokens'; + +const mockGenerateCloudAgentWorkflowToken = jest.mocked(generateCloudAgentWorkflowToken); + +const organizationId = 'aaaaaaaa-aaaa-4aaa-8aaa-aaaaaaaaaaaa'; + +beforeEach(() => { + jest.clearAllMocks(); + mockGetFixTicketById.mockResolvedValue({ + repo_full_name: 'kilo/repo', + issue_number: 42, + issue_title: 'Fix it', + issue_body: 'Please fix', + trigger_source: 'label', + }); +}); + +describe('prepareFixPayload workflow token ownership', () => { + it.each([ + [{ type: 'org', id: organizationId, userId: 'user-1' }, organizationId], + [{ type: 'user', id: 'user-1', userId: 'user-1' }, undefined], + ] as const)( + 'passes the exact owner organization or undefined', + async (owner, expectedOrganizationId) => { + await prepareFixPayload({ + ticketId: 'ticket-1', + owner, + agentConfig: { + config: { enabled_for_issues: true, repository_selection_mode: 'all' }, + }, + }); + + expect(mockGenerateCloudAgentWorkflowToken).toHaveBeenCalledWith( + expect.objectContaining({ id: 'user-1' }), + expect.objectContaining({ organizationId: expectedOrganizationId }) + ); + } + ); +}); diff --git a/apps/web/src/lib/auto-fix/triggers/prepare-fix-payload.ts b/apps/web/src/lib/auto-fix/triggers/prepare-fix-payload.ts index af094b6274..4260a507e1 100644 --- a/apps/web/src/lib/auto-fix/triggers/prepare-fix-payload.ts +++ b/apps/web/src/lib/auto-fix/triggers/prepare-fix-payload.ts @@ -9,7 +9,7 @@ import { captureException } from '@sentry/nextjs'; import { db } from '@/lib/drizzle'; import { kilocode_users } from '@kilocode/db/schema'; import { eq } from 'drizzle-orm'; -import { generateApiToken } from '@/lib/tokens'; +import { generateCloudAgentWorkflowToken, TOKEN_EXPIRY } from '@/lib/tokens'; import { getFixTicketById } from '../db/fix-tickets'; import type { Owner } from '../core/schemas'; import type { DispatchFixRequest } from '../core/schemas'; @@ -51,7 +51,12 @@ export async function prepareFixPayload(params: PreparePayloadParams): Promise ({ + db: { + select: () => ({ + from: () => ({ + where: () => ({ + limit: () => [{ id: 'user-1', api_token_pepper: 'pepper' }], + }), + }), + }), + }, +})); + +jest.mock('@/lib/tokens', () => ({ + generateCloudAgentWorkflowToken: jest.fn(() => 'workflow-token'), + TOKEN_EXPIRY: { default: 3600 }, +})); + +jest.mock('../db/triage-tickets', () => ({ + getTriageTicketById: (...args: unknown[]) => mockGetTriageTicketById(...args), +})); + +jest.mock('@sentry/nextjs', () => ({ captureException: jest.fn() })); + +import { prepareTriagePayload } from './prepare-triage-payload'; +import { generateCloudAgentWorkflowToken } from '@/lib/tokens'; + +const mockGenerateCloudAgentWorkflowToken = jest.mocked(generateCloudAgentWorkflowToken); + +const organizationId = 'aaaaaaaa-aaaa-4aaa-8aaa-aaaaaaaaaaaa'; + +beforeEach(() => { + jest.clearAllMocks(); + mockGetTriageTicketById.mockResolvedValue({ + repo_full_name: 'kilo/repo', + issue_number: 42, + issue_title: 'Triage it', + issue_body: 'Please triage', + }); +}); + +describe('prepareTriagePayload workflow token ownership', () => { + it.each([ + [{ type: 'org', id: organizationId, userId: 'user-1' }, organizationId], + [{ type: 'user', id: 'user-1', userId: 'user-1' }, undefined], + ] as const)( + 'passes the exact owner organization or undefined', + async (owner, expectedOrganizationId) => { + await prepareTriagePayload({ + ticketId: 'ticket-1', + owner, + agentConfig: { + config: { enabled_for_issues: true, repository_selection_mode: 'all' }, + }, + }); + + expect(mockGenerateCloudAgentWorkflowToken).toHaveBeenCalledWith( + expect.objectContaining({ id: 'user-1' }), + expect.objectContaining({ organizationId: expectedOrganizationId }) + ); + } + ); +}); diff --git a/apps/web/src/lib/auto-triage/triggers/prepare-triage-payload.ts b/apps/web/src/lib/auto-triage/triggers/prepare-triage-payload.ts index 3bbc586f56..7eecb6903e 100644 --- a/apps/web/src/lib/auto-triage/triggers/prepare-triage-payload.ts +++ b/apps/web/src/lib/auto-triage/triggers/prepare-triage-payload.ts @@ -9,7 +9,7 @@ import { captureException } from '@sentry/nextjs'; import { db } from '@/lib/drizzle'; import { kilocode_users } from '@kilocode/db/schema'; import { eq } from 'drizzle-orm'; -import { generateApiToken } from '@/lib/tokens'; +import { generateCloudAgentWorkflowToken, TOKEN_EXPIRY } from '@/lib/tokens'; import { getTriageTicketById } from '../db/triage-tickets'; import type { Owner } from '../core'; import type { AutoTriageAgentConfig, DispatchTriageRequest } from '../core/schemas'; @@ -53,7 +53,12 @@ export async function prepareTriagePayload( } // 3. Generate auth token for cloud agent with bot identifier - const authToken = generateApiToken(user, { botId: 'auto-triage' }); + const authToken = generateCloudAgentWorkflowToken(user, { + organizationId: owner.type === 'org' ? owner.id : undefined, + tokenSource: 'auto-triage', + botId: 'auto-triage', + expiresIn: TOKEN_EXPIRY.default, + }); // 4. Get config values const config = agentConfig.config as AutoTriageAgentConfig; diff --git a/apps/web/src/lib/cloud-agent-next/worktree-chat.test.ts b/apps/web/src/lib/cloud-agent-next/worktree-chat.test.ts index fefa9edca5..43e388a6b9 100644 --- a/apps/web/src/lib/cloud-agent-next/worktree-chat.test.ts +++ b/apps/web/src/lib/cloud-agent-next/worktree-chat.test.ts @@ -30,14 +30,14 @@ const mockWorkerCreateWorktreeChat = const mockCreateCloudAgentNextClient = jest.fn(() => ({ createWorktreeChat: mockWorkerCreateWorktreeChat, })); -const mockGenerateCloudAgentToken = jest.fn(() => 'cloud-agent-token'); +const mockCreateControlTokenForRequest = jest.fn(async () => ({ token: 'cloud-agent-token' })); jest.mock('@/lib/drizzle', () => ({ db: { select: mockSelect }, })); -jest.mock('@/lib/tokens', () => ({ - generateCloudAgentToken: mockGenerateCloudAgentToken, +jest.mock('@/lib/auth/resource-delegation', () => ({ + createControlTokenForRequest: mockCreateControlTokenForRequest, })); jest.mock('./cloud-agent-client', () => ({ @@ -160,7 +160,7 @@ describe('createWorktreeChat', () => { expect(where.params).toContain(organizationId); expect(where.sql).toContain('"organization_memberships"."id" is not null'); } - expect(mockGenerateCloudAgentToken).not.toHaveBeenCalled(); + expect(mockCreateControlTokenForRequest).not.toHaveBeenCalled(); expect(mockWorkerCreateWorktreeChat).not.toHaveBeenCalled(); }); diff --git a/apps/web/src/lib/cloud-agent-next/worktree-chat.ts b/apps/web/src/lib/cloud-agent-next/worktree-chat.ts index f4e33862d2..2814cc254d 100644 --- a/apps/web/src/lib/cloud-agent-next/worktree-chat.ts +++ b/apps/web/src/lib/cloud-agent-next/worktree-chat.ts @@ -7,7 +7,7 @@ import { TRPCError } from '@trpc/server'; import { and, eq, isNotNull, isNull } from 'drizzle-orm'; import * as z from 'zod'; import { db } from '@/lib/drizzle'; -import { generateCloudAgentToken } from '@/lib/tokens'; +import { createControlTokenForRequest } from '@/lib/auth/resource-delegation'; import { isMobileClient } from '@/lib/trpc/min-version'; import { createCloudAgentNextClient, type CreateWorktreeChatOutput } from './cloud-agent-client'; @@ -151,8 +151,13 @@ export async function createWorktreeChat({ }); } + const { token } = await createControlTokenForRequest(user, 'cloud-agent-next', { + headers: headersList ?? new Headers(), + organizationId, + tokenSource: 'cloud-agent', + }); try { - return await createCloudAgentNextClient(generateCloudAgentToken(user)).createWorktreeChat({ + return await createCloudAgentNextClient(token).createWorktreeChat({ sourceKiloSessionId, sourceCloudAgentSessionId: source.cloudAgentSessionId, operationKey, diff --git a/apps/web/src/lib/code-reviews/triggers/prepare-review-payload.ts b/apps/web/src/lib/code-reviews/triggers/prepare-review-payload.ts index 65ee241b82..2ac2738090 100644 --- a/apps/web/src/lib/code-reviews/triggers/prepare-review-payload.ts +++ b/apps/web/src/lib/code-reviews/triggers/prepare-review-payload.ts @@ -12,7 +12,7 @@ import { z } from 'zod'; import { db } from '@/lib/drizzle'; import { kilocode_users } from '@kilocode/db/schema'; import { eq } from 'drizzle-orm'; -import { generateApiToken } from '@/lib/tokens'; +import { generateCloudAgentWorkflowToken, TOKEN_EXPIRY } from '@/lib/tokens'; import { generateGitHubInstallationToken, findKiloReviewComment, @@ -298,7 +298,12 @@ export async function prepareReviewPayload( expectedHeadSha: expectedHeadSha.data, } ); - const authToken = generateApiToken(user, { botId: 'reviewer' }); + const authToken = generateCloudAgentWorkflowToken(user, { + organizationId: owner.type === 'org' ? owner.id : undefined, + tokenSource: 'code-review', + botId: 'reviewer', + expiresIn: TOKEN_EXPIRY.default, + }); // Single source for the standard reviewer's model so the session input and the // forward-shaped `reviewAgents[0]` can never drift apart. const standardModel = config.model_slug || DEFAULT_CODE_REVIEW_MODEL; @@ -712,7 +717,12 @@ export async function prepareReviewPayload( ]); // 5. Generate auth token for cloud agent with bot identifier - const authToken = generateApiToken(user, { botId: 'reviewer' }); + const authToken = generateCloudAgentWorkflowToken(user, { + organizationId: owner.type === 'org' ? owner.id : undefined, + tokenSource: 'code-review', + botId: 'reviewer', + expiresIn: TOKEN_EXPIRY.default, + }); // A council run replaces the standard sub-agent sharding policy with a coordinator // contract (one sub-agent per specialist, no self-review), so the base prompt must OMIT diff --git a/apps/web/src/lib/security-agent/services/analysis-service.test.ts b/apps/web/src/lib/security-agent/services/analysis-service.test.ts index 49217642ad..1166278a1b 100644 --- a/apps/web/src/lib/security-agent/services/analysis-service.test.ts +++ b/apps/web/src/lib/security-agent/services/analysis-service.test.ts @@ -22,6 +22,12 @@ const mockTriageSecurityFinding = jest.fn() as jest.MockedFunction< typeof triageModule.triageSecurityFinding >; const mockGenerateApiToken = jest.fn() as jest.MockedFunction; +const mockGenerateCloudAgentWorkflowToken = jest.fn() as jest.MockedFunction< + typeof tokensModule.generateCloudAgentWorkflowToken +>; +const mockGenerateWorkflowGatewayToken = jest.fn() as jest.MockedFunction< + typeof tokensModule.generateWorkflowGatewayToken +>; // eslint-disable-next-line @typescript-eslint/no-explicit-any const mockPrepareSession = jest.fn(); // eslint-disable-next-line @typescript-eslint/no-explicit-any @@ -57,6 +63,9 @@ jest.mock('./triage-service', () => ({ jest.mock('@/lib/tokens', () => ({ generateApiToken: mockGenerateApiToken, + generateCloudAgentWorkflowToken: mockGenerateCloudAgentWorkflowToken, + generateWorkflowGatewayToken: mockGenerateWorkflowGatewayToken, + TOKEN_EXPIRY: { default: 157_680_000 }, })); jest.mock('@/lib/cloud-agent-next/cloud-agent-client', () => ({ @@ -99,6 +108,8 @@ describe('analysis-service', () => { beforeEach(() => { jest.clearAllMocks(); mockTryAcquireAnalysisStartLease.mockResolvedValue(true); + mockGenerateCloudAgentWorkflowToken.mockReturnValue('cloud-agent-token'); + mockGenerateWorkflowGatewayToken.mockReturnValue('gateway-token'); }); it('does not start when start lease cannot be acquired', async () => { @@ -160,6 +171,8 @@ describe('analysis-service', () => { const mockFinding = { id: findingId, + owned_by_organization_id: organizationId, + owned_by_user_id: null, source: 'dependabot', source_id: '42', status: 'open', @@ -256,6 +269,30 @@ describe('analysis-service', () => { }); }); + it('rejects a personal finding owned by another user before minting credentials', async () => { + const user = { id: 'user-1', google_user_email: 'test@example.com' } as User; + mockGetSecurityFindingById.mockResolvedValue({ + id: 'finding-other-user', + owned_by_organization_id: null, + owned_by_user_id: 'user-2', + status: 'open', + } as Awaited>); + + await expect( + startSecurityAnalysis({ + findingId: 'finding-other-user', + user, + githubRepo: 'acme/repo', + githubToken: 'gh-token', + }) + ).resolves.toEqual({ + started: false, + error: 'Analysis organization does not match the finding owner', + }); + expect(mockGenerateCloudAgentWorkflowToken).not.toHaveBeenCalled(); + expect(mockGenerateWorkflowGatewayToken).not.toHaveBeenCalled(); + }); + it('uses triageModel for triage and analysisModel for sandbox session', async () => { const findingId = 'finding-model-split'; const user = { id: 'user-1', google_user_email: 'test@example.com' } as User; diff --git a/apps/web/src/lib/security-agent/services/analysis-service.token-source.test.ts b/apps/web/src/lib/security-agent/services/analysis-service.token-source.test.ts index 7fe95e100d..4ac31a1071 100644 --- a/apps/web/src/lib/security-agent/services/analysis-service.token-source.test.ts +++ b/apps/web/src/lib/security-agent/services/analysis-service.token-source.test.ts @@ -1,4 +1,5 @@ import { beforeAll, beforeEach, describe, expect, it, jest } from '@jest/globals'; +import jwt from 'jsonwebtoken'; import type { SecurityFinding, User } from '@kilocode/db/schema'; import type * as securityAnalysisModule from '../db/security-analysis'; import type * as securityFindingsModule from '../db/security-findings'; @@ -6,7 +7,19 @@ import type * as triageModule from './triage-service'; import type { startSecurityAnalysis as startSecurityAnalysisType } from './analysis-service'; import type { CloudAgentNextClient } from '@/lib/cloud-agent-next/cloud-agent-client'; import { insertTestUser } from '@/tests/helpers/user.helper'; -import { expectNonExchangeableSystemToken } from '@/tests/helpers/system-token.helper'; +import { + isKiloCredentialExchangeEligible, + verifyKiloTokenForPolicy, +} from '@kilocode/worker-utils/kilo-token-policy'; + +const shared = { enabled: true }; +const tokenSecret = 'security-agent-token-source-test-secret'; + +jest.mock('@/lib/config.server', () => ({ + NEXTAUTH_SECRET: 'security-agent-token-source-test-secret', + CALLBACK_TOKEN_SECRET: 'test-callback-token-secret', + isResourceTokenIssuanceEnabled: () => shared.enabled, +})); const mockGetSecurityFindingById = jest.fn(); const mockUpdateAnalysisStatus = jest.fn(); @@ -124,6 +137,7 @@ function createFinding(user: User): SecurityFinding { describe('startSecurityAnalysis token source', () => { beforeEach(() => { jest.clearAllMocks(); + shared.enabled = true; mockTryAcquireAnalysisStartLease.mockResolvedValue(true); mockUpdateAnalysisStatus.mockResolvedValue(true); mockTriageSecurityFinding.mockResolvedValue({ @@ -147,9 +161,14 @@ describe('startSecurityAnalysis token source', () => { }); }); - it('uses one non-exchangeable security-agent token for triage and sandbox analysis', async () => { + it('uses separate modern gateway and sandbox security-agent credentials', async () => { const user = await insertTestUser({ api_token_pepper: crypto.randomUUID() }); - const finding = createFinding(user); + const organizationId = crypto.randomUUID(); + const finding = { + ...createFinding(user), + owned_by_organization_id: organizationId, + owned_by_user_id: null, + }; mockGetSecurityFindingById.mockResolvedValue(finding); const result = await startSecurityAnalysis({ @@ -157,6 +176,7 @@ describe('startSecurityAnalysis token source', () => { user, githubRepo: finding.repo_full_name, githubToken: 'github-token', + organizationId, }); expect(result).toEqual({ started: true, triageOnly: false }); @@ -164,11 +184,60 @@ describe('startSecurityAnalysis token source', () => { const cloudAgentToken = mockCreateCloudAgentNextClient.mock.calls[0]?.[0]; if (!triageInput) throw new Error('Expected triage to receive an input'); expect(triageInput.authToken).toEqual(expect.any(String)); - expect(cloudAgentToken).toBe(triageInput.authToken); - await expectNonExchangeableSystemToken(triageInput.authToken, user, 'security-agent'); + const gatewayClaims = jwt.verify(triageInput.authToken, tokenSecret) as jwt.JwtPayload; + const cloudAgentClaims = jwt.decode(cloudAgentToken); + if (!cloudAgentClaims || typeof cloudAgentClaims === 'string') { + throw new Error('Expected sandbox JWT claims'); + } + expect(gatewayClaims).toMatchObject({ + aud: 'kilo-gateway', + tokenPurpose: 'delegated-workload', + credentialExchange: false, + organizationId, + tokenSource: 'security-agent', + }); + expect(gatewayClaims.exp! - gatewayClaims.iat!).toBeLessThanOrEqual(60 * 60); + expect(cloudAgentClaims).toMatchObject({ + aud: 'cloud-agent-next', + tokenPurpose: 'internal-service', + credentialExchange: false, + organizationId, + }); + expect(cloudAgentToken).not.toBe(triageInput.authToken); + const tokenPolicy = await verifyKiloTokenForPolicy(triageInput.authToken, tokenSecret, { + audience: 'kilo-gateway', + mode: 'required', + }); + expect(isKiloCredentialExchangeEligible(tokenPolicy, { legacy: 'five-year-api' })).toBe(false); expect(mockPrepareSession).toHaveBeenCalledTimes(1); expect(mockInitiateFromPreparedSession).toHaveBeenCalledWith({ cloudAgentSessionId: 'agent-session-123', }); }); + + it('preserves the legacy gateway token shape when shared issuance is disabled', async () => { + shared.enabled = false; + const user = await insertTestUser({ api_token_pepper: crypto.randomUUID() }); + const finding = createFinding(user); + mockGetSecurityFindingById.mockResolvedValue(finding); + + await startSecurityAnalysis({ + findingId: finding.id, + user, + githubRepo: finding.repo_full_name, + githubToken: 'github-token', + }); + + const triageInput = mockTriageSecurityFinding.mock.calls[0]?.[0]; + if (!triageInput) throw new Error('Expected triage to receive an input'); + const claims = jwt.verify(triageInput.authToken, tokenSecret) as jwt.JwtPayload; + expect(claims).toMatchObject({ + kiloUserId: user.id, + apiTokenPepper: user.api_token_pepper, + tokenSource: 'security-agent', + }); + expect(claims).not.toHaveProperty('aud'); + expect(claims).not.toHaveProperty('tokenPurpose'); + expect(claims).not.toHaveProperty('credentialExchange'); + }); }); diff --git a/apps/web/src/lib/security-agent/services/analysis-service.ts b/apps/web/src/lib/security-agent/services/analysis-service.ts index d30fa653c9..123673e806 100644 --- a/apps/web/src/lib/security-agent/services/analysis-service.ts +++ b/apps/web/src/lib/security-agent/services/analysis-service.ts @@ -4,7 +4,11 @@ import { createCloudAgentNextClient, InsufficientCreditsError, } from '@/lib/cloud-agent-next/cloud-agent-client'; -import { generateApiToken } from '@/lib/tokens'; +import { + generateCloudAgentWorkflowToken, + generateWorkflowGatewayToken, + TOKEN_EXPIRY, +} from '@/lib/tokens'; import { getSecurityFindingById } from '../db/security-findings'; import { updateAnalysisStatus, @@ -254,6 +258,14 @@ export async function startSecurityAnalysis(params: { return { started: false, error: `Finding not found: ${findingId}` }; } const findingDataSnapshot = buildSecurityFindingAnalysisInput(finding); + const findingOrganizationId = finding.owned_by_organization_id ?? undefined; + const findingUserId = finding.owned_by_user_id ?? undefined; + if ( + organizationId !== findingOrganizationId || + (!findingOrganizationId && findingUserId !== undefined && findingUserId !== user.id) + ) { + return { started: false, error: 'Analysis organization does not match the finding owner' }; + } const leaseAcquired = await tryAcquireAnalysisStartLease(findingId); if (!leaseAcquired) { @@ -290,7 +302,15 @@ export async function startSecurityAnalysis(params: { const analysisStartTime = Date.now(); try { - const authToken = generateApiToken(user, { tokenSource: 'security-agent' }); + const cloudAgentToken = generateCloudAgentWorkflowToken(user, { + organizationId: findingOrganizationId, + tokenSource: 'security-agent', + expiresIn: TOKEN_EXPIRY.default, + }); + const gatewayToken = generateWorkflowGatewayToken(user, { + organizationId: findingOrganizationId, + tokenSource: 'security-agent', + }); let triage: SecurityFindingTriage; @@ -306,7 +326,7 @@ export async function startSecurityAnalysis(params: { trackSecurityAgentAnalysisStarted({ distinctId: user.id, userId: user.id, - organizationId, + organizationId: findingOrganizationId, findingId, model: analysisModel, triageModel, @@ -319,7 +339,7 @@ export async function startSecurityAnalysis(params: { trackSecurityAgentAnalysisStarted({ distinctId: user.id, userId: user.id, - organizationId, + organizationId: findingOrganizationId, findingId, model: analysisModel, triageModel, @@ -330,11 +350,11 @@ export async function startSecurityAnalysis(params: { const tier1Start = performance.now(); triage = await triageSecurityFinding({ finding, - authToken, + authToken: gatewayToken, model: triageModel, correlationId, userId: user.id, - organizationId, + organizationId: findingOrganizationId, }); const tier1DurationMs = Math.round(performance.now() - tier1Start); @@ -391,7 +411,7 @@ export async function startSecurityAnalysis(params: { trackSecurityAgentAnalysisCompleted({ distinctId: user.id, userId: user.id, - organizationId, + organizationId: findingOrganizationId, findingId, model: triageModel, triageModel, @@ -403,7 +423,9 @@ export async function startSecurityAnalysis(params: { durationMs: Date.now() - analysisStartTime, }); - const owner: SecurityReviewOwner = organizationId ? { organizationId } : { userId: user.id }; + const owner: SecurityReviewOwner = findingOrganizationId + ? { organizationId: findingOrganizationId } + : { userId: user.id }; void maybeAutoDismissAnalysis({ findingId, @@ -438,7 +460,7 @@ export async function startSecurityAnalysis(params: { await updateAnalysisStatus(findingId, 'pending', { analysis: partialAnalysis }); const prompt = buildAnalysisPrompt(finding); - const client = createCloudAgentNextClient(authToken); + const client = createCloudAgentNextClient(cloudAgentToken); const callbackUrl = `${APP_URL}/api/internal/security-analysis-callback/${findingId}`; const callbackToken = await deriveCallbackToken({ @@ -453,7 +475,7 @@ export async function startSecurityAnalysis(params: { model: analysisModel, githubRepo, githubToken, - kilocodeOrganizationId: organizationId, + kilocodeOrganizationId: findingOrganizationId, createdOnPlatform: 'security-agent', callbackTarget: { url: callbackUrl, diff --git a/apps/web/src/routers/app-builder-router.ts b/apps/web/src/routers/app-builder-router.ts index 4879c460b3..a5eadbbadd 100644 --- a/apps/web/src/routers/app-builder-router.ts +++ b/apps/web/src/routers/app-builder-router.ts @@ -1,6 +1,6 @@ import 'server-only'; import { baseProcedure, createTRPCRouter } from '@/lib/trpc/init'; -import { generateApiToken } from '@/lib/tokens'; +import { createControlTokenForRequest } from '@/lib/auth/resource-delegation'; import * as appBuilderService from '@/lib/app-builder/app-builder-service'; import { createProjectBaseSchema, @@ -22,7 +22,10 @@ export const appBuilderRouter = createTRPCRouter({ */ createProject: baseProcedure.input(createProjectBaseSchema).mutation(async ({ ctx, input }) => { const owner = { type: 'user' as const, id: ctx.user.id }; - const authToken = generateApiToken(ctx.user, { tokenSource: 'app-builder' }); + const { token: authToken } = await createControlTokenForRequest(ctx.user, 'cloud-agent-next', { + headers: ctx.headersList, + tokenSource: 'app-builder', + }); return appBuilderService.createProject({ owner, @@ -58,7 +61,10 @@ export const appBuilderRouter = createTRPCRouter({ * Get a single project with all messages and session state */ getProject: baseProcedure.input(projectIdBaseSchema).query(async ({ ctx, input }) => { - const authToken = generateApiToken(ctx.user, { tokenSource: 'app-builder' }); + const { token: authToken } = await createControlTokenForRequest(ctx.user, 'cloud-agent-next', { + headers: ctx.headersList, + tokenSource: 'app-builder', + }); return appBuilderService.getProject( input.projectId, { type: 'user', id: ctx.user.id }, @@ -136,7 +142,10 @@ export const appBuilderRouter = createTRPCRouter({ */ interruptSession: baseProcedure.input(projectIdBaseSchema).mutation(async ({ ctx, input }) => { const owner = { type: 'user' as const, id: ctx.user.id }; - const authToken = generateApiToken(ctx.user, { tokenSource: 'app-builder' }); + const { token: authToken } = await createControlTokenForRequest(ctx.user, 'cloud-agent-next', { + headers: ctx.headersList, + tokenSource: 'app-builder', + }); const result = await appBuilderService.interruptSession(input.projectId, owner, authToken); return { success: result.success }; }), @@ -172,7 +181,10 @@ export const appBuilderRouter = createTRPCRouter({ */ startSession: baseProcedure.input(projectIdBaseSchema).mutation(async ({ ctx, input }) => { const owner = { type: 'user' as const, id: ctx.user.id }; - const authToken = generateApiToken(ctx.user, { tokenSource: 'app-builder' }); + const { token: authToken } = await createControlTokenForRequest(ctx.user, 'cloud-agent-next', { + headers: ctx.headersList, + tokenSource: 'app-builder', + }); const result = await appBuilderService.startSessionForProject({ projectId: input.projectId, @@ -194,7 +206,10 @@ export const appBuilderRouter = createTRPCRouter({ */ sendMessage: baseProcedure.input(sendMessageBaseSchema).mutation(async ({ ctx, input }) => { const owner = { type: 'user' as const, id: ctx.user.id }; - const authToken = generateApiToken(ctx.user, { tokenSource: 'app-builder' }); + const { token: authToken } = await createControlTokenForRequest(ctx.user, 'cloud-agent-next', { + headers: ctx.headersList, + tokenSource: 'app-builder', + }); const result = await appBuilderService.sendMessage({ projectId: input.projectId, diff --git a/apps/web/src/routers/app-builder-token-source.test.ts b/apps/web/src/routers/app-builder-token-source.test.ts index af5019d8b0..764c2dc548 100644 --- a/apps/web/src/routers/app-builder-token-source.test.ts +++ b/apps/web/src/routers/app-builder-token-source.test.ts @@ -6,6 +6,12 @@ const appBuilderServiceMocks = { sendMessage: jest.fn(), }; +const mockCreateControlTokenForRequest = jest.fn(); + +jest.mock('@/lib/auth/resource-delegation', () => ({ + createControlTokenForRequest: (...args: unknown[]) => mockCreateControlTokenForRequest(...args), +})); + jest.mock('@/lib/redis', () => ({ redisClient: { get: jest.fn(async () => null) }, })); @@ -23,11 +29,7 @@ import { beforeEach, describe, expect, it } from '@jest/globals'; import { createCallerForUser } from '@/routers/test-utils'; import { db } from '@/lib/drizzle'; import { organizations, organization_memberships } from '@kilocode/db/schema'; -import { expectNonExchangeableSystemToken } from '@/tests/helpers/system-token.helper'; import { insertTestUser } from '@/tests/helpers/user.helper'; -import { CLOUD_AGENT_NEXT_AUDIENCE } from '@kilocode/worker-utils/internal-service-token-audiences'; -import { verifyKiloTokenForResource } from '@kilocode/worker-utils/kilo-token-policy'; -import { NEXTAUTH_SECRET } from '@/lib/config.server'; import type { Owner } from '@/lib/integrations/core/types'; import type { User } from '@kilocode/db/schema'; @@ -114,11 +116,17 @@ describe('App Builder system tokens', () => { beforeEach(async () => { jest.clearAllMocks(); + mockCreateControlTokenForRequest.mockImplementation( + async (currentUser: User, _resource: string, options?: { organizationId?: string }) => ({ + token: `app-builder-token-${options?.organizationId ?? 'personal'}`, + user: currentUser, + }) + ); mockAppBuilderService(); user = await insertTestUser({ api_token_pepper: 'app-builder-token-source-pepper' }); }); - it('forwards non-exchangeable app-builder tokens for personal and organization operations', async () => { + it('forwards verified request-derived app-builder tokens for personal and organization operations', async () => { const [organization] = await db .insert(organizations) .values({ @@ -156,14 +164,22 @@ describe('App Builder system tokens', () => { { type: 'org', id: organization.id }, ]); - for (const { token } of tokens) { - await expectNonExchangeableSystemToken(token, user, 'app-builder'); - await expect( - verifyKiloTokenForResource(token, NEXTAUTH_SECRET, { - audience: CLOUD_AGENT_NEXT_AUDIENCE, - mode: 'allow-legacy', - }) - ).resolves.toMatchObject({ kiloUserId: user.id, tokenSource: 'app-builder' }); + expect(tokens.map(({ token }) => token)).toEqual([ + 'app-builder-token-personal', + `app-builder-token-${organization.id}`, + 'app-builder-token-personal', + `app-builder-token-${organization.id}`, + 'app-builder-token-personal', + `app-builder-token-${organization.id}`, + 'app-builder-token-personal', + `app-builder-token-${organization.id}`, + 'app-builder-token-personal', + `app-builder-token-${organization.id}`, + ]); + expect(mockCreateControlTokenForRequest).toHaveBeenCalledTimes(10); + for (const [, resource, options] of mockCreateControlTokenForRequest.mock.calls) { + expect(resource).toBe('cloud-agent-next'); + expect(options).toMatchObject({ tokenSource: 'app-builder' }); } }); }); diff --git a/apps/web/src/routers/cli-sessions-v2-router.test.ts b/apps/web/src/routers/cli-sessions-v2-router.test.ts index ffd127f20c..6a8c16d009 100644 --- a/apps/web/src/routers/cli-sessions-v2-router.test.ts +++ b/apps/web/src/routers/cli-sessions-v2-router.test.ts @@ -39,6 +39,10 @@ jest.mock('@/lib/cloud-agent-next/cloud-agent-client', () => ({ }, })); +jest.mock('@/lib/auth/resource-delegation', () => ({ + createControlTokenForRequest: jest.fn(async () => ({ token: 'test-cloud-agent-control-token' })), +})); + jest.mock('@/lib/tokens', () => { const actual: Record = jest.requireActual('@/lib/tokens'); return { diff --git a/apps/web/src/routers/cli-sessions-v2-router.ts b/apps/web/src/routers/cli-sessions-v2-router.ts index 88d3d058b9..49451813da 100644 --- a/apps/web/src/routers/cli-sessions-v2-router.ts +++ b/apps/web/src/routers/cli-sessions-v2-router.ts @@ -29,7 +29,7 @@ import { type CliSessionV2, } from '@kilocode/db/schema'; import { createCloudAgentNextClient } from '@/lib/cloud-agent-next/cloud-agent-client'; -import { generateCloudAgentToken } from '@/lib/tokens'; +import { createControlTokenForRequest } from '@/lib/auth/resource-delegation'; import { fetchSessionSnapshot, fetchSessionMessagesPage, @@ -1044,7 +1044,15 @@ export const cliSessionsV2Router = createTRPCRouter({ .input(WorktreeInputSchema) .output(z.object({ success: z.literal(true), deletedSessionIds: z.array(z.string()) })) .mutation(async ({ ctx, input }) => { - const client = createCloudAgentNextClient(generateCloudAgentToken(ctx.user)); + const client = createCloudAgentNextClient( + ( + await createControlTokenForRequest(ctx.user, 'cloud-agent-next', { + headers: ctx.headersList, + organizationId: input.organizationId ?? undefined, + tokenSource: 'cloud-agent', + }) + ).token + ); try { return await client.deleteWorktree({ worktreeId: input.worktreeId, @@ -1423,7 +1431,13 @@ export const cliSessionsV2Router = createTRPCRouter({ let watermarkEventId: number | null = null; if (!input.cursor && session.cloud_agent_session_id) { try { - const authToken = generateCloudAgentToken(ctx.user); + const authToken = ( + await createControlTokenForRequest(ctx.user, 'cloud-agent-next', { + headers: ctx.headersList, + organizationId: session.organization_id ?? undefined, + tokenSource: 'cloud-agent', + }) + ).token; const client = createCloudAgentNextClient(authToken); const sessionState = await client.getSession(session.cloud_agent_session_id); watermarkEventId = sessionState.latestEventId ?? null; @@ -1589,7 +1603,13 @@ export const cliSessionsV2Router = createTRPCRouter({ if (session.cloud_agent_session_id) { try { - const authToken = generateCloudAgentToken(ctx.user); + const authToken = ( + await createControlTokenForRequest(ctx.user, 'cloud-agent-next', { + headers: ctx.headersList, + organizationId: session.organization_id ?? undefined, + tokenSource: 'cloud-agent', + }) + ).token; const client = createCloudAgentNextClient(authToken); runtimeState = await client.getSession(session.cloud_agent_session_id); } catch (error) { @@ -1975,7 +1995,13 @@ export const cliSessionsV2Router = createTRPCRouter({ const session = await getSessionWithAccessCheck(session_id, ctx); if (session.cloud_agent_session_id) { - const authToken = generateCloudAgentToken(ctx.user); + const authToken = ( + await createControlTokenForRequest(ctx.user, 'cloud-agent-next', { + headers: ctx.headersList, + organizationId: session.organization_id ?? undefined, + tokenSource: 'cloud-agent', + }) + ).token; const client = createCloudAgentNextClient(authToken); try { const result = await client.deleteSession(session.cloud_agent_session_id); diff --git a/apps/web/src/routers/cli-sessions-v2-worktree.test.ts b/apps/web/src/routers/cli-sessions-v2-worktree.test.ts index cdedcab6a8..6457ca2ff2 100644 --- a/apps/web/src/routers/cli-sessions-v2-worktree.test.ts +++ b/apps/web/src/routers/cli-sessions-v2-worktree.test.ts @@ -73,7 +73,12 @@ jest.mock('@/lib/cloud-agent-next/cloud-agent-client', () => ({ jest.mock('@/lib/tokens', () => ({ generateApiToken: jest.fn(() => 'cloud-agent-token'), - generateCloudAgentToken: mockGenerateCloudAgentToken, +})); + +jest.mock('@/lib/auth/resource-delegation', () => ({ + createControlTokenForRequest: jest.fn(async (user: User) => ({ + token: mockGenerateCloudAgentToken(user), + })), })); jest.mock('@/lib/session-ingest-client', () => ({ diff --git a/apps/web/src/routers/cloud-agent-next-router.test.ts b/apps/web/src/routers/cloud-agent-next-router.test.ts index 1af7167ce2..f2cf4dd0ec 100644 --- a/apps/web/src/routers/cloud-agent-next-router.test.ts +++ b/apps/web/src/routers/cloud-agent-next-router.test.ts @@ -138,8 +138,8 @@ const mockOrderRepositoriesByUsage = }) => Promise >(); -jest.mock('@/lib/tokens', () => ({ - generateCloudAgentToken: jest.fn(() => 'cloud-agent-token'), +jest.mock('@/lib/auth/resource-delegation', () => ({ + createControlTokenForRequest: jest.fn(async () => ({ token: 'cloud-agent-token' })), })); jest.mock('@/lib/cloud-agent-next/cloud-agent-client', () => ({ diff --git a/apps/web/src/routers/cloud-agent-next-router.ts b/apps/web/src/routers/cloud-agent-next-router.ts index 2d2039f209..64ec082a2c 100644 --- a/apps/web/src/routers/cloud-agent-next-router.ts +++ b/apps/web/src/routers/cloud-agent-next-router.ts @@ -8,7 +8,8 @@ import { import { computeCloudAgentNextBalanceCheckEligibility } from '@/lib/cloud-agent-next/balance-check-eligibility'; import { rethrowAsTerminalError } from '@/lib/cloud-agent-next/terminal-errors'; import { createWorktreeChat } from '@/lib/cloud-agent-next/worktree-chat'; -import { generateCloudAgentToken } from '@/lib/tokens'; +import { createControlTokenForRequest } from '@/lib/auth/resource-delegation'; +import type { User } from '@kilocode/db/schema'; import { isFeatureFlagEnabledOrDevelopment } from '@/lib/posthog-feature-flags'; import { fetchGitHubRepositoriesForUser } from '@/lib/cloud-agent/github-integration-helpers'; import { @@ -121,6 +122,15 @@ async function assertUserOwnsSession(userId: string, cloudAgentSessionId: string } } +async function createCloudAgentControlToken(user: User, headersList?: Headers): Promise { + return ( + await createControlTokenForRequest(user, 'cloud-agent-next', { + headers: headersList, + tokenSource: 'cloud-agent', + }) + ).token; +} + /** * Cloud Agent Next Router (Personal Context) * @@ -154,7 +164,7 @@ export const cloudAgentNextRouter = createTRPCRouter({ }); } - const authToken = generateCloudAgentToken(ctx.user); + const authToken = await createCloudAgentControlToken(ctx.user, ctx.headersList); const eligibility = await computeCloudAgentNextBalanceCheckEligibility({ fromDb: db, user: ctx.user, @@ -235,7 +245,7 @@ export const cloudAgentNextRouter = createTRPCRouter({ .output(baseInitiateSessionNextOutputSchema) .mutation(async ({ ctx, input }) => { await assertUserOwnsSession(ctx.user.id, input.cloudAgentSessionId); - const authToken = generateCloudAgentToken(ctx.user); + const authToken = await createCloudAgentControlToken(ctx.user, ctx.headersList); const client = createCloudAgentNextClient(authToken); // No token fetch needed: prepare and initiate happen back-to-back, @@ -262,7 +272,7 @@ export const cloudAgentNextRouter = createTRPCRouter({ .output(baseInitiateSessionNextOutputSchema) .mutation(async ({ ctx, input }) => { await assertUserOwnsSession(ctx.user.id, input.cloudAgentSessionId); - const authToken = generateCloudAgentToken(ctx.user); + const authToken = await createCloudAgentControlToken(ctx.user, ctx.headersList); // Prompt turns carry their own model; command turns run the session's // stored model, so resolve it to apply the same free/BYOK eligibility // to every follow-up that queues a model-using turn. If the worker @@ -328,7 +338,8 @@ export const cloudAgentNextRouter = createTRPCRouter({ .output(getWorktreeChangesOutputSchema) .query(async ({ ctx, input }) => { await assertUserOwnsSession(ctx.user.id, input.cloudAgentSessionId); - const client = createCloudAgentNextClient(generateCloudAgentToken(ctx.user)); + const authToken = await createCloudAgentControlToken(ctx.user, ctx.headersList); + const client = createCloudAgentNextClient(authToken); return await client.getWorktreeChanges(input.cloudAgentSessionId); }), @@ -337,7 +348,8 @@ export const cloudAgentNextRouter = createTRPCRouter({ .output(refreshWorktreeChangesOutputSchema) .mutation(async ({ ctx, input }) => { await assertUserOwnsSession(ctx.user.id, input.cloudAgentSessionId); - const client = createCloudAgentNextClient(generateCloudAgentToken(ctx.user)); + const authToken = await createCloudAgentControlToken(ctx.user, ctx.headersList); + const client = createCloudAgentNextClient(authToken); return await client.refreshWorktreeChanges(input.cloudAgentSessionId); }), @@ -348,7 +360,7 @@ export const cloudAgentNextRouter = createTRPCRouter({ await assertUserOwnsSession(ctx.user.id, input.cloudAgentSessionId); try { - const authToken = generateCloudAgentToken(ctx.user); + const authToken = await createCloudAgentControlToken(ctx.user, ctx.headersList); const client = createCloudAgentNextClient(authToken); const result = await client.createTerminal(input); const terminalTicket = createTerminalTicket({ @@ -387,7 +399,7 @@ export const cloudAgentNextRouter = createTRPCRouter({ await assertUserOwnsSession(ctx.user.id, input.cloudAgentSessionId); try { - const authToken = generateCloudAgentToken(ctx.user); + const authToken = await createCloudAgentControlToken(ctx.user, ctx.headersList); const client = createCloudAgentNextClient(authToken); return await client.resizeTerminal(input); } catch (error) { @@ -402,7 +414,7 @@ export const cloudAgentNextRouter = createTRPCRouter({ await assertUserOwnsSession(ctx.user.id, input.cloudAgentSessionId); try { - const authToken = generateCloudAgentToken(ctx.user); + const authToken = await createCloudAgentControlToken(ctx.user, ctx.headersList); const client = createCloudAgentNextClient(authToken); return await client.closeTerminal(input); } catch (error) { @@ -497,7 +509,7 @@ export const cloudAgentNextRouter = createTRPCRouter({ ) .mutation(async ({ ctx, input }) => { await assertUserOwnsSession(ctx.user.id, input.sessionId); - const authToken = generateCloudAgentToken(ctx.user); + const authToken = await createCloudAgentControlToken(ctx.user, ctx.headersList); const client = createCloudAgentNextClient(authToken); return await client.interruptSession(input.sessionId); @@ -513,7 +525,7 @@ export const cloudAgentNextRouter = createTRPCRouter({ .output(z.object({ dropped: z.boolean() })) .mutation(async ({ ctx, input }) => { await assertUserOwnsSession(ctx.user.id, input.sessionId); - const authToken = generateCloudAgentToken(ctx.user); + const authToken = await createCloudAgentControlToken(ctx.user, ctx.headersList); const client = createCloudAgentNextClient(authToken); return await client.cancelQueuedMessage(input.sessionId, input.messageId); @@ -524,7 +536,7 @@ export const cloudAgentNextRouter = createTRPCRouter({ .output(z.object({ success: z.boolean() })) .mutation(async ({ ctx, input }) => { await assertUserOwnsSession(ctx.user.id, input.sessionId); - const authToken = generateCloudAgentToken(ctx.user); + const authToken = await createCloudAgentControlToken(ctx.user, ctx.headersList); const client = createCloudAgentNextClient(authToken); return await client.answerQuestion(input); }), @@ -534,7 +546,7 @@ export const cloudAgentNextRouter = createTRPCRouter({ .output(z.object({ success: z.boolean() })) .mutation(async ({ ctx, input }) => { await assertUserOwnsSession(ctx.user.id, input.sessionId); - const authToken = generateCloudAgentToken(ctx.user); + const authToken = await createCloudAgentControlToken(ctx.user, ctx.headersList); const client = createCloudAgentNextClient(authToken); return await client.rejectQuestion(input); }), @@ -544,7 +556,7 @@ export const cloudAgentNextRouter = createTRPCRouter({ .output(z.object({ success: z.boolean() })) .mutation(async ({ ctx, input }) => { await assertUserOwnsSession(ctx.user.id, input.sessionId); - const authToken = generateCloudAgentToken(ctx.user); + const authToken = await createCloudAgentControlToken(ctx.user, ctx.headersList); const client = createCloudAgentNextClient(authToken); return await client.answerPermission(input); }), @@ -558,7 +570,7 @@ export const cloudAgentNextRouter = createTRPCRouter({ .output(baseGetSessionNextOutputSchema) .query(async ({ ctx, input }) => { await assertUserOwnsSession(ctx.user.id, input.cloudAgentSessionId); - const authToken = generateCloudAgentToken(ctx.user); + const authToken = await createCloudAgentControlToken(ctx.user, ctx.headersList); const client = createCloudAgentNextClient(authToken); return await client.getSession(input.cloudAgentSessionId); @@ -574,9 +586,9 @@ export const cloudAgentNextRouter = createTRPCRouter({ message: 'Session not found or access denied', }); }); - return await createCloudAgentNextClient(generateCloudAgentToken(ctx.user)).getSandboxStatus( - input.cloudAgentSessionId - ); + return await createCloudAgentNextClient( + await createCloudAgentControlToken(ctx.user, ctx.headersList) + ).getSandboxStatus(input.cloudAgentSessionId); }), getComputeBillingStatus: baseProcedure @@ -584,7 +596,7 @@ export const cloudAgentNextRouter = createTRPCRouter({ .query(async ({ ctx, input }) => { await assertUserOwnsSession(ctx.user.id, input.cloudAgentSessionId); return await createCloudAgentNextClient( - generateCloudAgentToken(ctx.user) + await createCloudAgentControlToken(ctx.user, ctx.headersList) ).getComputeBillingStatus(input.cloudAgentSessionId); }), diff --git a/apps/web/src/routers/organizations/organization-app-builder-router.ts b/apps/web/src/routers/organizations/organization-app-builder-router.ts index c381da2821..b2186d0fdd 100644 --- a/apps/web/src/routers/organizations/organization-app-builder-router.ts +++ b/apps/web/src/routers/organizations/organization-app-builder-router.ts @@ -1,6 +1,6 @@ import 'server-only'; import { createTRPCRouter } from '@/lib/trpc/init'; -import { generateApiToken } from '@/lib/tokens'; +import { createControlTokenForRequest } from '@/lib/auth/resource-delegation'; import { organizationMemberProcedure, organizationMemberMutationProcedure, @@ -39,7 +39,15 @@ export const organizationAppBuilderRouter = createTRPCRouter({ .input(createProjectSchema) .mutation(async ({ ctx, input }) => { const owner = { type: 'org' as const, id: input.organizationId }; - const authToken = generateApiToken(ctx.user, { tokenSource: 'app-builder' }); + const { token: authToken } = await createControlTokenForRequest( + ctx.user, + 'cloud-agent-next', + { + headers: ctx.headersList, + organizationId: input.organizationId, + tokenSource: 'app-builder', + } + ); return appBuilderService.createProject({ owner, @@ -81,7 +89,15 @@ export const organizationAppBuilderRouter = createTRPCRouter({ .input(projectWithOrgIdSchema) .query(async ({ ctx, input }) => { const owner = { type: 'org' as const, id: input.organizationId }; - const authToken = generateApiToken(ctx.user, { tokenSource: 'app-builder' }); + const { token: authToken } = await createControlTokenForRequest( + ctx.user, + 'cloud-agent-next', + { + headers: ctx.headersList, + organizationId: input.organizationId, + tokenSource: 'app-builder', + } + ); return appBuilderService.getProject(input.projectId, owner, authToken); }), @@ -173,7 +189,15 @@ export const organizationAppBuilderRouter = createTRPCRouter({ .input(projectWithOrgIdSchema) .mutation(async ({ ctx, input }) => { const owner = { type: 'org' as const, id: input.organizationId }; - const authToken = generateApiToken(ctx.user, { tokenSource: 'app-builder' }); + const { token: authToken } = await createControlTokenForRequest( + ctx.user, + 'cloud-agent-next', + { + headers: ctx.headersList, + organizationId: input.organizationId, + tokenSource: 'app-builder', + } + ); const result = await appBuilderService.interruptSession(input.projectId, owner, authToken); return { success: result.success }; }), @@ -212,7 +236,15 @@ export const organizationAppBuilderRouter = createTRPCRouter({ .input(projectWithOrgIdSchema) .mutation(async ({ ctx, input }) => { const owner = { type: 'org' as const, id: input.organizationId }; - const authToken = generateApiToken(ctx.user, { tokenSource: 'app-builder' }); + const { token: authToken } = await createControlTokenForRequest( + ctx.user, + 'cloud-agent-next', + { + headers: ctx.headersList, + organizationId: input.organizationId, + tokenSource: 'app-builder', + } + ); const result = await appBuilderService.startSessionForProject({ projectId: input.projectId, @@ -236,7 +268,15 @@ export const organizationAppBuilderRouter = createTRPCRouter({ .input(sendMessageSchema) .mutation(async ({ ctx, input }) => { const owner = { type: 'org' as const, id: input.organizationId }; - const authToken = generateApiToken(ctx.user, { tokenSource: 'app-builder' }); + const { token: authToken } = await createControlTokenForRequest( + ctx.user, + 'cloud-agent-next', + { + headers: ctx.headersList, + organizationId: input.organizationId, + tokenSource: 'app-builder', + } + ); const result = await appBuilderService.sendMessage({ projectId: input.projectId, diff --git a/apps/web/src/routers/organizations/organization-cloud-agent-next-router.test.ts b/apps/web/src/routers/organizations/organization-cloud-agent-next-router.test.ts index 91ed51437c..6aeeb63493 100644 --- a/apps/web/src/routers/organizations/organization-cloud-agent-next-router.test.ts +++ b/apps/web/src/routers/organizations/organization-cloud-agent-next-router.test.ts @@ -165,8 +165,8 @@ const mockEnsureOrganizationAccess = jest.fn(); const mockRequireActiveSubscription = jest.fn<() => void>(); -jest.mock('@/lib/tokens', () => ({ - generateCloudAgentToken: jest.fn(() => 'cloud-agent-token'), +jest.mock('@/lib/auth/resource-delegation', () => ({ + createControlTokenForRequest: jest.fn(async () => ({ token: 'cloud-agent-token' })), })); jest.mock('@/lib/cloud-agent-next/cloud-agent-client', () => ({ diff --git a/apps/web/src/routers/organizations/organization-cloud-agent-next-router.ts b/apps/web/src/routers/organizations/organization-cloud-agent-next-router.ts index 4592c9472e..96d6283bfa 100644 --- a/apps/web/src/routers/organizations/organization-cloud-agent-next-router.ts +++ b/apps/web/src/routers/organizations/organization-cloud-agent-next-router.ts @@ -8,7 +8,8 @@ import { import { computeCloudAgentNextBalanceCheckEligibility } from '@/lib/cloud-agent-next/balance-check-eligibility'; import { rethrowAsTerminalError } from '@/lib/cloud-agent-next/terminal-errors'; import { createWorktreeChat } from '@/lib/cloud-agent-next/worktree-chat'; -import { generateCloudAgentToken } from '@/lib/tokens'; +import { createControlTokenForRequest } from '@/lib/auth/resource-delegation'; +import type { User } from '@kilocode/db/schema'; import { isFeatureFlagEnabledOrDevelopment } from '@/lib/posthog-feature-flags'; import { ensureOrganizationAccess, @@ -64,6 +65,20 @@ import { import { linkPendingUploads, releasePendingUploads } from '@/lib/r2/cloud-agent-pending-uploads'; import * as z from 'zod'; import { PLATFORM } from '@/lib/integrations/core/constants'; + +async function createCloudAgentControlToken( + user: User, + headersList: Headers | undefined, + organizationId: string +): Promise { + return ( + await createControlTokenForRequest(user, 'cloud-agent-next', { + headers: headersList, + organizationId, + tokenSource: 'cloud-agent', + }) + ).token; +} import { signStreamTicket } from '@/lib/cloud-agent/stream-ticket'; import { db } from '@/lib/drizzle'; import { verifyOrgOwnsSessionV2ByCloudAgentId } from '@/lib/cloud-agent/session-ownership'; @@ -266,7 +281,11 @@ export const organizationCloudAgentNextRouter = createTRPCRouter({ }); } - const authToken = generateCloudAgentToken(ctx.user); + const authToken = await createCloudAgentControlToken( + ctx.user, + ctx.headersList, + input.organizationId + ); const eligibility = await computeCloudAgentNextBalanceCheckEligibility({ fromDb: db, user: ctx.user, @@ -371,7 +390,11 @@ export const organizationCloudAgentNextRouter = createTRPCRouter({ userId: ctx.user.id, cloudAgentSessionId: input.cloudAgentSessionId, }); - const authToken = generateCloudAgentToken(ctx.user); + const authToken = await createCloudAgentControlToken( + ctx.user, + ctx.headersList, + input.organizationId + ); const client = createCloudAgentNextClient(authToken); // No token fetch needed: prepare and initiate happen back-to-back, @@ -402,7 +425,11 @@ export const organizationCloudAgentNextRouter = createTRPCRouter({ userId: ctx.user.id, cloudAgentSessionId: input.cloudAgentSessionId, }); - const authToken = generateCloudAgentToken(ctx.user); + const authToken = await createCloudAgentControlToken( + ctx.user, + ctx.headersList, + input.organizationId + ); // Prompt turns carry their own model; command turns run the session's // stored model, so resolve it to apply the same free/BYOK eligibility // to every follow-up that queues a model-using turn. If the worker @@ -476,7 +503,12 @@ export const organizationCloudAgentNextRouter = createTRPCRouter({ userId: ctx.user.id, cloudAgentSessionId: input.cloudAgentSessionId, }); - const client = createCloudAgentNextClient(generateCloudAgentToken(ctx.user)); + const authToken = await createCloudAgentControlToken( + ctx.user, + ctx.headersList, + input.organizationId + ); + const client = createCloudAgentNextClient(authToken); return await client.getWorktreeChanges(input.cloudAgentSessionId); }), @@ -489,7 +521,12 @@ export const organizationCloudAgentNextRouter = createTRPCRouter({ userId: ctx.user.id, cloudAgentSessionId: input.cloudAgentSessionId, }); - const client = createCloudAgentNextClient(generateCloudAgentToken(ctx.user)); + const authToken = await createCloudAgentControlToken( + ctx.user, + ctx.headersList, + input.organizationId + ); + const client = createCloudAgentNextClient(authToken); return await client.refreshWorktreeChanges(input.cloudAgentSessionId); }), @@ -504,7 +541,11 @@ export const organizationCloudAgentNextRouter = createTRPCRouter({ }); try { - const authToken = generateCloudAgentToken(ctx.user); + const authToken = await createCloudAgentControlToken( + ctx.user, + ctx.headersList, + input.organizationId + ); const client = createCloudAgentNextClient(authToken); const result = await client.createTerminal({ cloudAgentSessionId: input.cloudAgentSessionId, @@ -557,7 +598,11 @@ export const organizationCloudAgentNextRouter = createTRPCRouter({ }); try { - const authToken = generateCloudAgentToken(ctx.user); + const authToken = await createCloudAgentControlToken( + ctx.user, + ctx.headersList, + input.organizationId + ); const client = createCloudAgentNextClient(authToken); return await client.resizeTerminal({ cloudAgentSessionId: input.cloudAgentSessionId, @@ -581,7 +626,11 @@ export const organizationCloudAgentNextRouter = createTRPCRouter({ }); try { - const authToken = generateCloudAgentToken(ctx.user); + const authToken = await createCloudAgentControlToken( + ctx.user, + ctx.headersList, + input.organizationId + ); const client = createCloudAgentNextClient(authToken); return await client.closeTerminal({ cloudAgentSessionId: input.cloudAgentSessionId, @@ -669,7 +718,11 @@ export const organizationCloudAgentNextRouter = createTRPCRouter({ userId: ctx.user.id, cloudAgentSessionId: input.sessionId, }); - const authToken = generateCloudAgentToken(ctx.user); + const authToken = await createCloudAgentControlToken( + ctx.user, + ctx.headersList, + input.organizationId + ); const client = createCloudAgentNextClient(authToken); return await client.interruptSession(input.sessionId); @@ -689,7 +742,11 @@ export const organizationCloudAgentNextRouter = createTRPCRouter({ userId: ctx.user.id, cloudAgentSessionId: input.sessionId, }); - const authToken = generateCloudAgentToken(ctx.user); + const authToken = await createCloudAgentControlToken( + ctx.user, + ctx.headersList, + input.organizationId + ); const client = createCloudAgentNextClient(authToken); return await client.cancelQueuedMessage(input.sessionId, input.messageId); @@ -704,7 +761,11 @@ export const organizationCloudAgentNextRouter = createTRPCRouter({ userId: ctx.user.id, cloudAgentSessionId: input.sessionId, }); - const authToken = generateCloudAgentToken(ctx.user); + const authToken = await createCloudAgentControlToken( + ctx.user, + ctx.headersList, + input.organizationId + ); const client = createCloudAgentNextClient(authToken); return await client.answerQuestion({ sessionId: input.sessionId, @@ -722,7 +783,11 @@ export const organizationCloudAgentNextRouter = createTRPCRouter({ userId: ctx.user.id, cloudAgentSessionId: input.sessionId, }); - const authToken = generateCloudAgentToken(ctx.user); + const authToken = await createCloudAgentControlToken( + ctx.user, + ctx.headersList, + input.organizationId + ); const client = createCloudAgentNextClient(authToken); return await client.rejectQuestion({ sessionId: input.sessionId, @@ -739,7 +804,11 @@ export const organizationCloudAgentNextRouter = createTRPCRouter({ userId: ctx.user.id, cloudAgentSessionId: input.sessionId, }); - const authToken = generateCloudAgentToken(ctx.user); + const authToken = await createCloudAgentControlToken( + ctx.user, + ctx.headersList, + input.organizationId + ); const client = createCloudAgentNextClient(authToken); return await client.answerPermission({ sessionId: input.sessionId, @@ -761,7 +830,11 @@ export const organizationCloudAgentNextRouter = createTRPCRouter({ userId: ctx.user.id, cloudAgentSessionId: input.cloudAgentSessionId, }); - const authToken = generateCloudAgentToken(ctx.user); + const authToken = await createCloudAgentControlToken( + ctx.user, + ctx.headersList, + input.organizationId + ); const client = createCloudAgentNextClient(authToken); return await client.getSession(input.cloudAgentSessionId); @@ -787,9 +860,9 @@ export const organizationCloudAgentNextRouter = createTRPCRouter({ message: 'Session not found or access denied', }); } - return await createCloudAgentNextClient(generateCloudAgentToken(ctx.user)).getSandboxStatus( - input.cloudAgentSessionId - ); + return await createCloudAgentNextClient( + await createCloudAgentControlToken(ctx.user, ctx.headersList, input.organizationId) + ).getSandboxStatus(input.cloudAgentSessionId); }), getComputeBillingStatus: organizationMemberProcedure @@ -801,7 +874,7 @@ export const organizationCloudAgentNextRouter = createTRPCRouter({ cloudAgentSessionId: input.cloudAgentSessionId, }); return await createCloudAgentNextClient( - generateCloudAgentToken(ctx.user) + await createCloudAgentControlToken(ctx.user, ctx.headersList, input.organizationId) ).getComputeBillingStatus(input.cloudAgentSessionId); }), diff --git a/services/cloud-agent-next/.dev.vars.example b/services/cloud-agent-next/.dev.vars.example index 996288d2a0..1d30afe89c 100644 --- a/services/cloud-agent-next/.dev.vars.example +++ b/services/cloud-agent-next/.dev.vars.example @@ -44,6 +44,9 @@ PER_SESSION_SANDBOX_ORG_IDS= CONTROL_PLANE_IDS=* WORKTREE_CREATION_ENABLED_IDS=* +# Non-secret rollout control for per-session Kilo runtime isolation. +RUNTIME_ISOLATION_ENABLED=false + CREDENTIAL_CONTAINMENT_ENABLED=false # Repo snapshot org IDs (optional) diff --git a/services/cloud-agent-next/src/kilo-facade/user-kilo-facade.ts b/services/cloud-agent-next/src/kilo-facade/user-kilo-facade.ts index 2f10846b47..8f23d02242 100644 --- a/services/cloud-agent-next/src/kilo-facade/user-kilo-facade.ts +++ b/services/cloud-agent-next/src/kilo-facade/user-kilo-facade.ts @@ -845,7 +845,7 @@ async function admitBasicPrompt(params: { try { return await preflightAndAdmitPromptMessage( { cloudAgentSessionId: params.cloudAgentSessionId, ...command }, - { env: params.env, userId: params.userId }, + { env: params.env, userId: params.userId, authToken: params.authToken }, 'kilo.prompt_async', () => admitPrompt({ diff --git a/services/cloud-agent-next/src/kilo/__fixtures__/runtime-url-normalization.ts b/services/cloud-agent-next/src/kilo/__fixtures__/runtime-url-normalization.ts new file mode 100644 index 0000000000..07aa1fc959 --- /dev/null +++ b/services/cloud-agent-next/src/kilo/__fixtures__/runtime-url-normalization.ts @@ -0,0 +1,27 @@ +/** + * Immutable compatibility fixtures copied from Kilo CLI `route()` behavior. + * Sources: Kilo v7.4.20 release artifact and current CLI main (2026-09-04). + * + * Both versions normalize a configured base then append absolute API paths. + * Provider-specific `/api/gateway/v1/...` compatibility aliases are tested by + * the Cloud facade resolver, not asserted as behavior of `route()` itself. + * Keep this independent of the mutable CLI checkout: older installed CLIs are + * part of the Cloud Agent compatibility contract. + */ +export const historicalRouteFixtures = [ + { version: '7.4.20', path: '/api/profile' }, + { version: '7.4.20', path: '/api/defaults' }, + { version: '7.4.20', path: '/api/openrouter/models' }, + { version: '7.4.20', path: '/api/session' }, + { version: 'current', path: '/api/profile' }, + { version: 'current', path: '/api/defaults' }, + { version: 'current', path: '/api/openrouter/models' }, + { version: 'current', path: '/api/session' }, +] as const; + +/** Exact relevant `route()` normalization contract from both sources. */ +export function historicalCliRoute(base: string, path: string): string { + const url = new URL(base); + url.pathname = `${url.pathname.replace(/\/+$/, '')}/${path.replace(/^\/+/, '')}`; + return url.toString(); +} diff --git a/services/cloud-agent-next/src/kilo/kilo-targets.test.ts b/services/cloud-agent-next/src/kilo/kilo-targets.test.ts index 6873bb3953..e22fcc235e 100644 --- a/services/cloud-agent-next/src/kilo/kilo-targets.test.ts +++ b/services/cloud-agent-next/src/kilo/kilo-targets.test.ts @@ -4,6 +4,26 @@ import { deriveKiloSandboxTargets, providerBaseUrlEncodedInToken, } from './kilo-targets.js'; +import { + historicalCliRoute, + historicalRouteFixtures, +} from './__fixtures__/runtime-url-normalization.js'; + +describe('immutable historical Kilo route normalization', () => { + it.each([ + ['origin facade', 'https://worker.example.test'], + ['safe prefixed facade', 'https://worker.example.test/runtime-proxy'], + ])('%s keeps v7.4.20 and current requests inside the Worker facade', (_name, facade) => { + for (const fixture of historicalRouteFixtures) { + const routed = new URL(historicalCliRoute(facade, fixture.path)); + expect(routed.origin).toBe('https://worker.example.test'); + expect(routed.pathname).toBe( + `${new URL(facade).pathname.replace(/\/+$/, '')}${fixture.path}` + ); + if (new URL(facade).pathname !== '/') expect(routed.pathname).not.toBe(fixture.path); + } + }); +}); describe('providerBaseUrlEncodedInToken', () => { it('extracts and normalizes a provider base while preserving the full token separately', () => { diff --git a/services/cloud-agent-next/src/kilo/runtime-credential-proxy-routes.ts b/services/cloud-agent-next/src/kilo/runtime-credential-proxy-routes.ts new file mode 100644 index 0000000000..c10972aed1 --- /dev/null +++ b/services/cloud-agent-next/src/kilo/runtime-credential-proxy-routes.ts @@ -0,0 +1,218 @@ +export type RuntimeCredentialProxyTargets = { + backendBaseUrl: string; + providerBaseUrl: string; + sessionIngestBaseUrl: string; +}; + +export type RuntimeCredentialProxyRoute = 'backend' | 'provider' | 'ingest'; + +type ResolveRuntimeCredentialProxyRouteInput = { + targets: RuntimeCredentialProxyTargets; + route: RuntimeCredentialProxyRoute; + method: string; + pathname: string; + search: string; + kiloSessionId: string; + organizationId?: string; + contentType?: string | null; + bodyText?: string; +}; + +const ID = /^[A-Za-z0-9_-]{1,256}$/; +const ORGANIZATION_ID = /^[A-Za-z0-9][A-Za-z0-9._-]*$/; + +function safePathname(value: string): string | null { + if (!value.startsWith('/') || value.includes('\\') || /%(?:2f|5c)/i.test(value)) return null; + let decoded: string; + try { + decoded = decodeURIComponent(value); + } catch { + return null; + } + if ( + decoded.includes('\\') || + decoded.includes('//') || + /%(?:2e|2f|5c)/i.test(decoded) || + decoded.split('/').some(segment => segment === '.' || segment === '..') + ) { + return null; + } + return decoded; +} + +function targetUrl(base: string, pathname: string, search: string): URL | null { + try { + const url = new URL(base); + if ( + (url.protocol !== 'https:' && url.protocol !== 'http:') || + url.username || + url.password || + url.search || + url.hash + ) { + return null; + } + url.pathname = `${url.pathname.replace(/\/+$/, '')}${pathname}`; + url.search = search; + return url; + } catch { + return null; + } +} + +function isAllowedBackendRoute( + method: string, + path: string, + organizationId: string | undefined +): boolean { + if (method === 'GET') { + if ( + [ + '/api/user', + '/api/profile', + '/api/profile/balance', + '/api/defaults', + '/api/users/notifications', + ].includes(path) + ) { + return true; + } + const match = /^\/api\/organizations\/([A-Za-z0-9._-]+)\/(models|defaults|modes)$/.exec(path); + return match !== null && match[1] === organizationId; + } + return ( + method === 'POST' && + organizationId !== undefined && + path === `/api/organizations/${organizationId}/models/validate` + ); +} + +function isAllowedProviderRoute(method: string, path: string): boolean { + if (path === '/models') return method === 'GET'; + if (path === '/models/validate') return method === 'POST'; + return ( + method === 'POST' && + [ + '/chat/completions', + '/messages', + '/responses', + '/embeddings', + '/v1/chat/completions', + '/v1/responses', + ].includes(path) + ); +} + +function logicalProviderPath(pathname: string): string | null { + const prefix = ['/api/openrouter', '/api/gateway'].find(value => + pathname.startsWith(`${value}/`) + ); + if (!prefix) return null; + const path = pathname.slice(prefix.length); + if (path.startsWith('/v1/') && prefix !== '/api/gateway') return null; + return isAllowedProviderRoute('GET', path) || isAllowedProviderRoute('POST', path) ? path : null; +} + +/** Identifies the facade plane only; the route resolver still enforces its allowlist. */ +export function inferRuntimeCredentialProxyRoute( + pathname: string +): RuntimeCredentialProxyRoute | null { + if (pathname.startsWith('/api/openrouter/') || pathname.startsWith('/api/gateway/')) { + return 'provider'; + } + if ( + pathname === '/api/session' || + /^\/api\/session\/[A-Za-z0-9_-]+\/(?:export|ingest|title)$/.test(pathname) + ) { + return 'ingest'; + } + return pathname.startsWith('/api/') ? 'backend' : null; +} + +function providerTargetUrl(base: string, path: string, search: string): URL | null { + let provider: URL; + try { + provider = new URL(base); + } catch { + return null; + } + if ( + (provider.protocol !== 'https:' && provider.protocol !== 'http:') || + provider.username || + provider.password || + provider.search || + provider.hash + ) { + return null; + } + + if (provider.origin === 'https://api.kilo.ai' && provider.pathname === '/') { + provider.pathname = `/api/gateway${path}`; + provider.search = search; + return provider; + } + return targetUrl(base, path, search); +} + +function hasAuthorizedSessionBody(input: ResolveRuntimeCredentialProxyRouteInput): boolean { + if (input.contentType?.split(';', 1)[0]?.trim().toLowerCase() !== 'application/json') + return false; + if (input.bodyText === undefined || input.bodyText.length > 8192) return false; + try { + const body: unknown = JSON.parse(input.bodyText); + if (typeof body !== 'object' || body === null || Array.isArray(body)) return false; + const sessionId = (body as Record).sessionId; + return ( + Object.keys(body).length === 1 && + typeof sessionId === 'string' && + ID.test(sessionId) && + sessionId === input.kiloSessionId + ); + } catch { + return false; + } +} + +function isAllowedIngestRoute( + input: ResolveRuntimeCredentialProxyRouteInput, + path: string +): boolean { + if (path === '/api/session') return input.method === 'POST' && hasAuthorizedSessionBody(input); + const match = /^\/api\/session\/([A-Za-z0-9_-]+)\/(export|ingest|title)$/.exec(path); + if (!match || match[1] !== input.kiloSessionId) return false; + return ( + (match[2] === 'export' && input.method === 'GET') || + (match[2] !== 'export' && input.method === 'POST') + ); +} + +/** Resolves only exact credential-bearing routes; unrecognized input fails closed. */ +export function resolveRuntimeCredentialProxyRoute( + input: ResolveRuntimeCredentialProxyRouteInput +): URL | null { + if ( + !ID.test(input.kiloSessionId) || + (input.organizationId !== undefined && !ORGANIZATION_ID.test(input.organizationId)) + ) { + return null; + } + const path = safePathname(input.pathname); + if (!path) return null; + + if ( + input.route === 'backend' && + isAllowedBackendRoute(input.method, path, input.organizationId) + ) { + return targetUrl(input.targets.backendBaseUrl, path, input.search); + } + if (input.route === 'provider') { + const logicalPath = logicalProviderPath(path); + if (logicalPath && isAllowedProviderRoute(input.method, logicalPath)) { + return providerTargetUrl(input.targets.providerBaseUrl, logicalPath, input.search); + } + } + if (input.route === 'ingest' && isAllowedIngestRoute(input, path)) { + return targetUrl(input.targets.sessionIngestBaseUrl, path, input.search); + } + return null; +} diff --git a/services/cloud-agent-next/src/kilo/wrapper-client.ts b/services/cloud-agent-next/src/kilo/wrapper-client.ts index 858584969a..6d09a7b328 100644 --- a/services/cloud-agent-next/src/kilo/wrapper-client.ts +++ b/services/cloud-agent-next/src/kilo/wrapper-client.ts @@ -1231,6 +1231,10 @@ export class WrapperClient { return this.request('GET', '/health'); } + async listTerminals(): Promise { + return this.request('GET', '/pty'); + } + /** * Get current job status. */ @@ -1291,6 +1295,10 @@ export class WrapperContainerClient { return this.request('GET', '/health'); } + async listTerminals(): Promise { + return this.request('GET', '/pty'); + } + async createTerminal(size?: { cols: number; rows: number }): Promise { return this.request('POST', '/pty', size); } diff --git a/services/cloud-agent-next/src/persistence/CloudAgentSession.ts b/services/cloud-agent-next/src/persistence/CloudAgentSession.ts index b567e6aae9..9369a0c180 100644 --- a/services/cloud-agent-next/src/persistence/CloudAgentSession.ts +++ b/services/cloud-agent-next/src/persistence/CloudAgentSession.ts @@ -8,6 +8,12 @@ import { DurableObject } from 'cloudflare:workers'; import type { CloudAgentQueueReport } from '@kilocode/worker-utils/cloud-agent-queue-report'; import { generateBranchSlug } from '@kilocode/worker-utils/deployment-slug'; import type { OperationResult } from './types.js'; +import { + renewRuntimeAuthorization, + unsealRuntimeAuthorization, +} from '@kilocode/worker-utils/runtime-authorization'; +import type { RuntimeAuthorization } from '@kilocode/worker-utils/runtime-authorization-contract'; +import { RuntimeAuthorizationSchema } from '@kilocode/worker-utils/runtime-authorization-contract'; import { getSandboxProvider, parseSessionMetadata, @@ -100,6 +106,21 @@ import { deriveSharedSandboxId, generateSandboxId } from '../sandbox-id.js'; import { recordSharedSandboxFailover } from '../shared-sandbox-route.js'; import { nextMetadataAfterAdmittedAgentModel } from './persist-admitted-agent-model.js'; import { dispatchedKilocodeModelId } from './model-utils.js'; +import { + getRuntimeAuthorizationStatus, + getRuntimeAuthorizationRecoveryState, + renewStoredRuntimeAuthorization, + RUNTIME_AUTHORIZATION_RECOVERY_KEY, + RUNTIME_AUTHORIZATION_KEY, + runtimeAuthorizationRecoveryLockSchema, +} from '../session/runtime-authorization-persistence.js'; +import { RUNTIME_PROXY_GRANT_KEY } from '../runtime-credential-proxy.js'; +import { getEffectiveCredentialContainment } from './session-metadata.js'; +import { + issuePersistedRuntimeProxyGrant, + resolvePersistedRuntimeProxyCredential, +} from '../runtime-credential-proxy-rpc.js'; +import type { RuntimeProxyFence } from '../runtime-credential-proxy.js'; import { resolveSecret, validateStreamTicket, STREAM_TICKET_AUDIENCE } from '../auth.js'; import { isAllowedStreamWebSocketOrigin } from './ws-origin.js'; @@ -134,6 +155,7 @@ import { import { createQueuedSessionMessageState, getSessionMessageState, + hasNonTerminalSessionMessage, listNonTerminalAcceptedMessages, markAgentActivityObserved, markMessageAccepted, @@ -251,6 +273,7 @@ function extractAssistantTextFromParts(parts: AssistantMessagePart[]): string { type GroupedRegisterSessionInput = { identity: SessionMetadata['identity']; auth: SessionMetadata['auth']; + runtimeAuthorizationSeal?: string; clone?: SessionMetadata['clone']; /** Omitted for a clone-only create: no synthetic initial turn is registered. */ message?: { @@ -1667,6 +1690,313 @@ export class CloudAgentSession extends DurableObject { return this.getStoredMetadata(); } + async getRuntimeToken(): Promise { + const metadata = await this.getMetadata(); + const secret = await resolveSecret(this.env.NEXTAUTH_SECRET); + if (!secret) throw new Error('NEXTAUTH_SECRET is not configured on the worker'); + return renewStoredRuntimeAuthorization({ + metadata, + getAuthorization: () => this.ctx.storage.get(RUNTIME_AUTHORIZATION_KEY), + putAuthorization: authorization => + this.ctx.storage.put(RUNTIME_AUTHORIZATION_KEY, authorization), + getMetadata: () => this.getMetadata(), + putMetadata: updated => this.ctx.storage.put('metadata', updated), + renew: authorization => + renewRuntimeAuthorization({ + authorization, + secret, + connectionString: this.env.HYPERDRIVE.connectionString, + }), + }); + } + + async getRuntimeAuthorizationStatus(): Promise<'legacy' | 'active' | 'revoked'> { + return getRuntimeAuthorizationStatus({ + metadata: await this.getMetadata(), + getAuthorization: () => this.ctx.storage.get(RUNTIME_AUTHORIZATION_KEY), + }); + } + + async getRuntimeAuthorizationRecoveryState(): Promise<{ + state: 'legacy' | 'revoked' | 'active' | 'expired'; + id?: string; + recoveryId?: string; + }> { + const state = await getRuntimeAuthorizationRecoveryState({ + metadata: await this.getMetadata(), + getAuthorization: () => this.ctx.storage.get(RUNTIME_AUTHORIZATION_KEY), + }); + const lock = runtimeAuthorizationRecoveryLockSchema.safeParse( + await this.ctx.storage.get(RUNTIME_AUTHORIZATION_RECOVERY_KEY) + ); + return state.state === 'expired' && lock.success && lock.data.expectedOldId === state.id + ? { ...state, recoveryId: lock.data.recoveryId } + : state; + } + + async isRuntimeAuthorizationRecoveryInProgress(): Promise { + return (await this.ctx.storage.get(RUNTIME_AUTHORIZATION_RECOVERY_KEY)) !== undefined; + } + + async recoverExpiredRuntimeAuthorization(input: { + ownerId: string; + expectedOldId: string; + recoveryId: string; + runtimeAuthorizationSeal: string; + runtimeToken: string; + }): Promise<{ status: 'recovered' | 'not-needed' | 'denied' | 'busy' | 'retry' }> { + const metadata = await this.getMetadata(); + if (!metadata || metadata.identity.userId !== input.ownerId) return { status: 'denied' }; + const secret = await resolveSecret(this.env.NEXTAUTH_SECRET); + if (!secret) return { status: 'denied' }; + let fresh: RuntimeAuthorization; + try { + fresh = await unsealRuntimeAuthorization(input.runtimeAuthorizationSeal, secret, { + resourceKind: 'cloud-agent-next', + resourceId: metadata.identity.sessionId, + userId: metadata.identity.userId, + organizationId: metadata.identity.orgId, + }); + } catch { + return { status: 'denied' }; + } + if (fresh.state !== 'active') return { status: 'denied' }; + const state = await this.getRuntimeAuthorizationRecoveryState(); + if (state.state === 'legacy' || state.state === 'active') return { status: 'not-needed' }; + if (state.state !== 'expired' || state.id !== input.expectedOldId) return { status: 'denied' }; + const [active, pending] = await Promise.all([ + hasNonTerminalSessionMessage(this.ctx.storage), + countPendingSessionMessages(this.ctx.storage), + ]); + if ( + active || + pending > 0 || + isWrapperRunFinalizing(await getWrapperRuntimeState(this.ctx.storage)) + ) { + return { status: 'busy' }; + } + const acquired = await this.ctx.storage.transaction(async transaction => { + if ( + (await countPendingSessionMessages(transaction)) > 0 || + (await hasNonTerminalSessionMessage(transaction)) + ) { + return false; + } + const existingLock = runtimeAuthorizationRecoveryLockSchema.safeParse( + await transaction.get(RUNTIME_AUTHORIZATION_RECOVERY_KEY) + ); + if ( + existingLock.success && + (existingLock.data.recoveryId !== input.recoveryId || + existingLock.data.expectedOldId !== input.expectedOldId) + ) { + return false; + } + await transaction.put(RUNTIME_AUTHORIZATION_RECOVERY_KEY, { + expectedOldId: input.expectedOldId, + recoveryId: input.recoveryId, + }); + return true; + }); + if (!acquired) { + return { status: 'retry' }; + } + try { + const observation = this.physicalWrapperObserver + ? await this.physicalWrapperObserver() + : this.orchestrator + ? { status: 'absent' as const } + : await createAgentSandbox( + this.env, + metadata, + this.getAgentSandboxRuntimeContext() + ).observeWrappersWithoutWaking(); + if (observation.status === 'inspection-failed') return { status: 'retry' }; + if (observation.status === 'present') { + const terminal = await this.getTerminalClient(); + if (!terminal.success || !terminal.data) return { status: 'retry' }; + if ((await terminal.data.client.listTerminals()).length > 0) return { status: 'busy' }; + const supervisor = this.getWrapperSupervisor(); + await supervisor.requestPhysicalWrapperStop('idle-timeout', { kind: 'session' }); + await supervisor.runMaintenance(Date.now()); + const stopped = this.physicalWrapperObserver + ? await this.physicalWrapperObserver() + : this.orchestrator + ? { status: 'absent' as const } + : await createAgentSandbox( + this.env, + metadata, + this.getAgentSandboxRuntimeContext() + ).observeWrappersWithoutWaking(); + if (stopped.status !== 'absent') return { status: 'retry' }; + } + const latest = await this.getRuntimeAuthorizationRecoveryState(); + if (latest.state !== 'expired' || latest.id !== input.expectedOldId) { + return latest.state === 'active' || latest.state === 'legacy' + ? { status: 'not-needed' } + : { status: 'denied' }; + } + await this.ctx.storage.transaction(async transaction => { + const current = RuntimeAuthorizationSchema.safeParse( + await transaction.get(RUNTIME_AUTHORIZATION_KEY) + ); + if ( + !current.success || + current.data.id !== input.expectedOldId || + current.data.state !== 'active' || + Date.parse(current.data.delegationExpiresAt) > Date.now() + ) { + throw new Error('runtime_authorization_recovery_cas_failed'); + } + const currentMetadata = await transaction.get('metadata'); + const latestMetadata = currentMetadata ? parseSessionMetadata(currentMetadata) : null; + if ( + !latestMetadata || + latestMetadata.identity.sessionId !== metadata.identity.sessionId || + latestMetadata.identity.userId !== metadata.identity.userId || + latestMetadata.identity.orgId !== metadata.identity.orgId + ) { + throw new Error('runtime_authorization_recovery_cas_failed'); + } + await transaction.put(RUNTIME_AUTHORIZATION_KEY, fresh); + await transaction.put( + 'metadata', + serializeSessionMetadata({ + ...latestMetadata, + auth: { ...latestMetadata.auth, kilocodeToken: input.runtimeToken }, + }) + ); + await transaction.delete(RUNTIME_PROXY_GRANT_KEY); + await transaction.delete(RUNTIME_AUTHORIZATION_RECOVERY_KEY); + }); + await clearWrapperRuntimeIdentity(this.ctx.storage, {}, { incrementGeneration: true }); + return { status: 'recovered' }; + } catch (error) { + if (error instanceof Error && error.message === 'runtime_authorization_recovery_cas_failed') { + return { status: 'retry' }; + } + return { status: 'retry' }; + } + } + + private async runtimeProxyFence(): Promise { + const [metadata, runtime, lease] = await Promise.all([ + this.getMetadata(), + getWrapperRuntimeState(this.ctx.storage), + getWrapperLease(this.ctx.storage), + ]); + if ( + !metadata || + !metadata.workspace?.sandboxId || + !runtime.wrapperRunId || + !runtime.wrapperConnectionId || + lease.state !== 'owns_wrapper' || + lease.instance.instanceGeneration !== runtime.wrapperGeneration + ) { + return null; + } + return { + plane: 'legacy', + generation: runtime.wrapperGeneration, + allocationId: lease.instance.instanceId, + wrapperRunId: runtime.wrapperRunId, + wrapperConnectionId: runtime.wrapperConnectionId, + }; + } + + async issueRuntimeCredentialProxyGrant(fence: { + wrapperRunId: string; + wrapperGeneration: number; + wrapperConnectionId: string; + }): Promise { + const currentFence = await this.runtimeProxyFence(); + if ( + !currentFence || + currentFence.plane !== 'legacy' || + currentFence.wrapperRunId !== fence.wrapperRunId || + currentFence.generation !== fence.wrapperGeneration || + currentFence.wrapperConnectionId !== fence.wrapperConnectionId + ) { + return null; + } + const token = await this.getRuntimeToken(); + const [metadata, storedAuthorization, latestFence] = await Promise.all([ + this.getMetadata(), + this.ctx.storage.get(RUNTIME_AUTHORIZATION_KEY), + this.runtimeProxyFence(), + ]); + if ( + !latestFence || + latestFence.plane !== 'legacy' || + latestFence.wrapperRunId !== fence.wrapperRunId || + latestFence.generation !== fence.wrapperGeneration || + latestFence.wrapperConnectionId !== fence.wrapperConnectionId + ) { + return null; + } + const authorization = RuntimeAuthorizationSchema.safeParse(storedAuthorization); + return issuePersistedRuntimeProxyGrant({ + env: this.env, + storage: this.ctx.storage, + metadata, + authorization: authorization.success ? authorization.data : null, + fence: latestFence, + token, + mode: + metadata && getEffectiveCredentialContainment(metadata).kilocode ? 'contained' : 'direct', + }); + } + + async resolveRuntimeCredentialProxyGrant(handle: string): Promise<{ + token: string; + organizationId?: string; + runtimeAuthorization: { userId: string; authorizationId: string; resourceId: string }; + } | null> { + return resolvePersistedRuntimeProxyCredential({ + env: this.env, + storage: this.ctx.storage, + handle, + metadata: () => this.getMetadata(), + authorization: async () => { + const parsed = RuntimeAuthorizationSchema.safeParse( + await this.ctx.storage.get(RUNTIME_AUTHORIZATION_KEY) + ); + return parsed.success ? parsed.data : null; + }, + fence: () => this.runtimeProxyFence(), + token: () => this.getRuntimeToken(), + }); + } + + async reauthorizeRuntimeAuthorization(input: { + ownerId: string; + expectedOldId: string; + runtimeAuthorizationSeal: string; + }): Promise { + const metadata = await this.getMetadata(); + if (!metadata || metadata.identity.userId !== input.ownerId) return false; + const secret = await resolveSecret(this.env.NEXTAUTH_SECRET); + if (!secret) return false; + let authorization: RuntimeAuthorization; + try { + authorization = await unsealRuntimeAuthorization(input.runtimeAuthorizationSeal, secret, { + resourceKind: 'cloud-agent-next', + resourceId: metadata.identity.sessionId, + userId: metadata.identity.userId, + organizationId: metadata.identity.orgId, + }); + } catch { + return false; + } + if (authorization.state !== 'active') return false; + const current = RuntimeAuthorizationSchema.safeParse( + await this.ctx.storage.get(RUNTIME_AUTHORIZATION_KEY) + ); + if (!current.success || current.data.id !== input.expectedOldId) return false; + await this.ctx.storage.put(RUNTIME_AUTHORIZATION_KEY, authorization); + return true; + } + async getRuntimeLocation(): Promise { const metadata = await this.getStoredMetadata(); return metadata ? sessionRuntimeLocator(metadata) : null; @@ -2157,6 +2487,9 @@ export class CloudAgentSession extends DurableObject { } async createTerminal(input: TerminalCreateInput): Promise> { + if (await this.ctx.storage.get(RUNTIME_AUTHORIZATION_RECOVERY_KEY)) { + return { success: false, error: 'Runtime authorization recovery is in progress' }; + } const terminal = await this.getTerminalClient(); if (!terminal.success || !terminal.data) { return { success: false, error: terminal.error }; @@ -2186,6 +2519,9 @@ export class CloudAgentSession extends DurableObject { cols: number; rows: number; }): Promise> { + if (await this.ctx.storage.get(RUNTIME_AUTHORIZATION_RECOVERY_KEY)) { + return { success: false, error: 'Runtime authorization recovery is in progress' }; + } const terminal = await this.getTerminalClient(); if (!terminal.success || !terminal.data) { return { success: false, error: terminal.error }; @@ -2497,6 +2833,28 @@ export class CloudAgentSession extends DurableObject { if (existing) { return { success: false, error: 'Session already registered' }; } + let runtimeAuthorization: RuntimeAuthorization | undefined; + if (input.runtimeAuthorizationSeal) { + const secret = await resolveSecret(this.env.NEXTAUTH_SECRET); + if (!secret) return { success: false, error: 'Authentication unavailable' }; + let authorization: RuntimeAuthorization; + try { + authorization = await unsealRuntimeAuthorization(input.runtimeAuthorizationSeal, secret, { + resourceKind: 'cloud-agent-next', + resourceId: input.identity.sessionId, + userId: input.identity.userId, + organizationId: input.identity.orgId, + }); + } catch { + return { success: false, error: 'Invalid runtime authorization' }; + } + if (authorization.state !== 'active') + return { success: false, error: 'Runtime authorization revoked' }; + if (await this.ctx.storage.get(RUNTIME_AUTHORIZATION_KEY)) { + return { success: false, error: 'Runtime authorization already installed' }; + } + runtimeAuthorization = authorization; + } const routeAssignmentError = await validateSharedSandboxRouteAssignment(input.workspace ?? {}); if (routeAssignmentError) { return { success: false, error: `Invalid metadata: ${routeAssignmentError}` }; @@ -2583,6 +2941,9 @@ export class CloudAgentSession extends DurableObject { return { success: false, error: modeError }; } + if (runtimeAuthorization) { + await this.ctx.storage.put(RUNTIME_AUTHORIZATION_KEY, runtimeAuthorization); + } await this.ctx.storage.put('metadata', serialized); await this.updateLastActivity(); await this.ensureAlarmScheduled(); @@ -3657,6 +4018,13 @@ export class CloudAgentSession extends DurableObject { async admitSubmittedMessage( request: SubmittedSessionMessageRequest ): Promise { + if (await this.ctx.storage.get(RUNTIME_AUTHORIZATION_RECOVERY_KEY)) { + return { + success: false, + code: 'COMPUTE_STOPPING', + error: 'Runtime authorization recovery is in progress', + }; + } const deletionPending = await this.deletionPendingAdmissionFailure(); if (deletionPending) return deletionPending; const result = await this.getSessionMessageQueue().admitSubmittedMessage(request); @@ -3720,6 +4088,13 @@ export class CloudAgentSession extends DurableObject { async admitPreparedInitialMessage( request: LegacyRegisteredInitialAdmissionRequest ): Promise { + if (await this.ctx.storage.get(RUNTIME_AUTHORIZATION_RECOVERY_KEY)) { + return { + success: false, + code: 'COMPUTE_STOPPING', + error: 'Runtime authorization recovery is in progress', + }; + } const deletionPending = await this.deletionPendingAdmissionFailure(); if (deletionPending) return deletionPending; const metadata = await this.getMetadata(); diff --git a/services/cloud-agent-next/src/persistence/SandboxControl.ts b/services/cloud-agent-next/src/persistence/SandboxControl.ts index 59e8326d9f..dee95672b8 100644 --- a/services/cloud-agent-next/src/persistence/SandboxControl.ts +++ b/services/cloud-agent-next/src/persistence/SandboxControl.ts @@ -178,6 +178,7 @@ import { } from '../sandbox-control/session-credentials.js'; import { adaptSessionAttachPayloadForWrapper } from '../sandbox-session/attach-payload.js'; import { parseControlPlaneCredential } from '../sandbox-control/managed-credential.js'; +import { verifyRuntimeCredentialProxyHandle } from '../runtime-credential-proxy.js'; import { diagnosticCause, diagnosticConnection, @@ -307,6 +308,15 @@ export type SandboxControlStatus = { work: WorkState; wrapperInstanceId?: string; operationResults?: true; + runtimeRecovery?: true; +}; + +export type ControlRuntimeCredentialProxyFence = { + plane: 'control'; + allocationId: string; + providerInstanceId: string; + connectionId: string; + wrapperInstanceId: string; }; export type RuntimeQuarantineResult = @@ -628,6 +638,77 @@ export class SandboxControl extends DurableObject { return this.readOwner(); } + async getRuntimeCredentialProxyFence(input: { + ownerId: string; + sessionId: string; + kiloSessionId: string; + directory: string; + }): Promise { + await this.ensureOperationalInitialized(); + if ( + typeof input.ownerId !== 'string' || + typeof input.sessionId !== 'string' || + typeof input.kiloSessionId !== 'string' || + typeof input.directory !== 'string' + ) { + return null; + } + const [ownerId, routes, physical, grants] = await Promise.all([ + this.readOwner(), + loadRouteTable(this.ctx.storage), + loadPhysicalRecord(this.ctx.storage), + loadSessionCredentialGrants(this.ctx.storage), + ]); + if (ownerId !== input.ownerId) return null; + const route = routes.get(input.sessionId); + const provisioned = grants.some( + grant => + grant.userId === input.ownerId && + grant.directory === input.directory && + grant.expiresAt > Date.now() && + grant.members.some( + member => + member.sessionId === input.sessionId && member.kiloSessionId === input.kiloSessionId + ) + ); + if ( + (!route && !provisioned) || + (route && + (route.ownerId !== input.ownerId || + route.kiloSessionId !== input.kiloSessionId || + route.directory !== input.directory)) + ) { + return null; + } + const worktreeId = route?.worktreeId ?? this.worktreeIdFromDirectory(input.directory); + if ( + this.runtimeDeleted || + this.exclusiveDeletionWorktreeId || + (worktreeId && this.deletingWorktrees.has(worktreeId)) || + physical.state !== 'running' || + physical.stopTombstone !== null || + physical.createIntent === null || + physical.providerRef === null + ) { + return null; + } + const runtime = this.readyWrapperRuntime(); + if ( + !runtime || + !runtime.wrapperInstanceId || + runtime.providerInstanceId !== physical.providerRef + ) { + return null; + } + return { + plane: 'control', + allocationId: physical.createIntent.intentId, + providerInstanceId: runtime.providerInstanceId, + connectionId: runtime.connectionId, + wrapperInstanceId: runtime.wrapperInstanceId, + }; + } + async request(input: SandboxControlOutboundRequest): Promise { await this.ensureOperationalInitialized(); if (input.operation === 'session.git.summary') return this.requestWorktreeChanges(input); @@ -703,6 +784,14 @@ export class SandboxControl extends DurableObject { ) { throw new Error('Sandbox wrapper runtime changed'); } + if ( + input.expectedConnection && + (runtime.connectionId !== input.expectedConnection.connectionId || + runtime.providerInstanceId !== input.expectedConnection.providerInstanceId || + runtime.wrapperInstanceId !== input.expectedConnection.wrapperInstanceId) + ) { + throw new Error('Sandbox control connection changed'); + } if ( authorization?.success && authorization.data.wrapperInstanceId !== runtime.wrapperInstanceId @@ -759,6 +848,15 @@ export class SandboxControl extends DurableObject { if (input.operation === 'session.attach' || input.operation === 'session.prompt') { const payload = parseOperationPayload(input.operation, input.payload); if (!payload.ok) throw new Error(payload.error.message); + const attach = + input.operation === 'session.attach' + ? sessionAttachPayloadSchema.parse(payload.payload) + : undefined; + if (attach?.runtimeIsolation === 'per-session') { + if (runtime.runtimeIsolation !== true) { + throw new Error('Sandbox wrapper does not support per-session runtime isolation'); + } + } const identity = sessionRequestIdentitySchema.safeParse(input.session); if (!identity.success) throw new Error('session identity is required'); await this.ctx.storage.transaction(async () => { @@ -1850,6 +1948,90 @@ export class SandboxControl extends DurableObject { }); } + async bindRuntimeCredentialProxyHandle(input: { + ownerId: string; + sessionId: string; + kiloSessionId: string; + directory: string; + handle: string; + }): Promise<{ bound: true }> { + return this.withCredentialUpdate(async () => { + if ( + typeof input.handle !== 'string' || + input.handle.length === 0 || + input.handle.length > 4096 + ) { + throw new Error('Invalid runtime credential proxy handle'); + } + const ownerId = await this.requireOwner(); + if (ownerId !== input.ownerId) throw new Error('Sandbox owner mismatch'); + if (this.providerKind !== 'vercel' || this.runtimeDeleted) { + throw new Error('Sandbox credential containment mismatch'); + } + const grants = await loadSessionCredentialGrants(this.ctx.storage); + const now = Date.now(); + const index = grants.findIndex( + grant => + grant.userId === ownerId && + grant.directory === input.directory && + grant.expiresAt > now && + grant.members.some( + member => + member.sessionId === input.sessionId && member.kiloSessionId === input.kiloSessionId + ) && + grant.kilo.runtimeProxy !== undefined + ); + if (index < 0) throw new Error('Session has no matching runtime proxy credential grant'); + const grant = grants[index]; + if (!grant) throw new Error('Session has no matching runtime proxy credential grant'); + const claims = await verifyRuntimeCredentialProxyHandle(this.env, input.handle); + if ( + !claims || + !('sessionId' in claims) || + claims.userId !== ownerId || + claims.sessionId !== input.sessionId || + claims.kiloSessionId !== input.kiloSessionId + ) { + throw new Error('Invalid runtime credential proxy member handle'); + } + const existingProxy = grant.kilo.runtimeProxy; + if (!existingProxy) throw new Error('Session has no matching runtime proxy credential grant'); + const updated = grants.map((value, current) => + current === index + ? { + ...value, + kilo: { + ...value.kilo, + runtimeProxy: value.kilo.runtimeProxy + ? { + ...value.kilo.runtimeProxy, + members: [ + ...value.kilo.runtimeProxy.members.filter( + member => member.sessionId !== input.sessionId + ), + { + sessionId: input.sessionId, + kiloSessionId: input.kiloSessionId, + handle: input.handle, + }, + ], + } + : undefined, + }, + } + : value + ); + await saveSessionCredentialGrants(this.ctx.storage, updated); + await this.updateNetworkPolicy({ + ownerId, + networkPolicy: buildControlNetworkPolicy(updated.filter(value => value.expiresAt > now)), + requiredContainment: WORKTREE_CREDENTIAL_CONTAINMENT, + }); + await this.ctx.storage.delete(CREDENTIAL_POLICY_DIRTY_KEY); + return { bound: true }; + }); + } + async detachSession(sessionId: string): Promise<{ existed: boolean }> { await this.ensureOperationalInitialized(); const route = (await loadRouteTable(this.ctx.storage)).get(sessionId); @@ -2338,6 +2520,7 @@ export class SandboxControl extends DurableObject { this.socketHandler.supportsOperationResults() ? { operationResults: true as const } : {}), + ...(runtime?.runtimeRecovery ? { runtimeRecovery: true as const } : {}), }; } diff --git a/services/cloud-agent-next/src/router.test.ts b/services/cloud-agent-next/src/router.test.ts index 624012bd6b..a1c3a5ab92 100644 --- a/services/cloud-agent-next/src/router.test.ts +++ b/services/cloud-agent-next/src/router.test.ts @@ -2701,7 +2701,7 @@ describe('legacy V2 execution response compatibility', () => { preflightPreparedInitialPromptModelMock.mockResolvedValue(undefined); }); - function createLegacyExecutionCaller() { + function createLegacyExecutionCaller(options?: { authToken?: string }) { const admitPreparedInitialMessage = vi.fn().mockResolvedValue({ success: true, outcome: 'queued', @@ -2716,10 +2716,11 @@ describe('legacy V2 execution response compatibility', () => { }); const hasMessageAdmission = vi.fn().mockResolvedValue(false); const replayPreparedInitialMessage = vi.fn().mockResolvedValue(undefined); + const getRuntimeAuthorizationRecoveryState = vi.fn().mockResolvedValue({ state: 'active' }); recordCloudAgentSessionFailureMock.mockReset().mockResolvedValue({}); const context = { userId: 'test-user-123', - authToken: 'test-token', + authToken: options?.authToken ?? 'test-token', botId: undefined, request: new Request('https://cloud-agent-next.test/trpc'), env: { @@ -2730,6 +2731,7 @@ describe('legacy V2 execution response compatibility', () => { admitSubmittedMessage, hasMessageAdmission, replayPreparedInitialMessage, + getRuntimeAuthorizationRecoveryState, })), }, SESSION_INGEST: {}, @@ -2742,6 +2744,7 @@ describe('legacy V2 execution response compatibility', () => { admitSubmittedMessage, hasMessageAdmission, replayPreparedInitialMessage, + getRuntimeAuthorizationRecoveryState, recordCloudAgentSessionFailure: recordCloudAgentSessionFailureMock, }; } @@ -2845,7 +2848,10 @@ describe('legacy V2 execution response compatibility', () => { }); it('returns an admitted prepared initial retry without repeating model preflight', async () => { - const { caller, replayPreparedInitialMessage } = createLegacyExecutionCaller(); + const { caller, replayPreparedInitialMessage, getRuntimeAuthorizationRecoveryState } = + createLegacyExecutionCaller({ + authToken: 'eyJhbGciOiJIUzI1NiJ9.eyJhdWQiOiJraWxvIn0.signature', + }); replayPreparedInitialMessage.mockResolvedValue({ success: true, outcome: 'queued', @@ -2862,9 +2868,38 @@ describe('legacy V2 execution response compatibility', () => { expect(result).toMatchObject({ executionId: acceptedMessageId, delivery: 'queued' }); expect(replayPreparedInitialMessage).toHaveBeenCalledTimes(1); + expect(getRuntimeAuthorizationRecoveryState).not.toHaveBeenCalled(); expect(preflightPreparedInitialPromptModelMock).not.toHaveBeenCalled(); }); + it('runs prepared admission recovery before model preflight and admission', async () => { + const { + caller, + replayPreparedInitialMessage, + getRuntimeAuthorizationRecoveryState, + admitPreparedInitialMessage, + } = createLegacyExecutionCaller({ + authToken: 'eyJhbGciOiJIUzI1NiJ9.eyJhdWQiOiJraWxvIn0.signature', + }); + + await caller.initiateFromKilocodeSessionV2({ cloudAgentSessionId: validSessionId }); + + expect(replayPreparedInitialMessage).toHaveBeenCalledOnce(); + expect(getRuntimeAuthorizationRecoveryState).toHaveBeenCalledOnce(); + expect(preflightPreparedInitialPromptModelMock).toHaveBeenCalledOnce(); + expect(admitPreparedInitialMessage).toHaveBeenCalledOnce(); + expect(replayPreparedInitialMessage.mock.invocationCallOrder[0]).toBeLessThan( + getRuntimeAuthorizationRecoveryState.mock.invocationCallOrder[0] ?? Number.POSITIVE_INFINITY + ); + expect(getRuntimeAuthorizationRecoveryState.mock.invocationCallOrder[0]).toBeLessThan( + preflightPreparedInitialPromptModelMock.mock.invocationCallOrder[0] ?? + Number.POSITIVE_INFINITY + ); + expect(preflightPreparedInitialPromptModelMock.mock.invocationCallOrder[0]).toBeLessThan( + admitPreparedInitialMessage.mock.invocationCallOrder[0] ?? Number.POSITIVE_INFINITY + ); + }); + it('sendMessageV2 preserves sent delivery when a runtime-accepted admission is replayed', async () => { const { caller, admitSubmittedMessage } = createLegacyExecutionCaller(); admitSubmittedMessage.mockResolvedValue({ diff --git a/services/cloud-agent-next/src/router/handlers/session-execution.ts b/services/cloud-agent-next/src/router/handlers/session-execution.ts index 0a1b371624..6973e2cdc7 100644 --- a/services/cloud-agent-next/src/router/handlers/session-execution.ts +++ b/services/cloud-agent-next/src/router/handlers/session-execution.ts @@ -20,7 +20,11 @@ import { LegacyExecutionResponse, } from '../schemas.js'; import type { SessionId } from '../../types/ids.js'; -import { preflightAndQueuePromptMessage, queueMessage } from '../../session/queue-message.js'; +import { + preflightAndQueuePromptMessage, + preflightRuntimeAuthorizationRecovery, + queueMessage, +} from '../../session/queue-message.js'; import { admitLegacyPreparedInitialMessage, replayLegacyPreparedInitialMessageIfAlreadyAdmitted, @@ -58,13 +62,20 @@ export function createSessionExecutionV2Handlers() { validatedSessionAccess: ctx.validatedSessionAccess, }); const admissionInput = { cloudAgentSessionId: input.cloudAgentSessionId }; - const admissionContext = { env: ctx.env, userId: ctx.userId, botId: ctx.botId }; + const admissionContext = { + env: ctx.env, + userId: ctx.userId, + botId: ctx.botId, + authToken: ctx.authToken, + }; const replay = await replayLegacyPreparedInitialMessageIfAlreadyAdmitted( admissionInput, admissionContext ); if (replay) return withLegacyExecutionId(replay); + await preflightRuntimeAuthorizationRecovery(input.cloudAgentSessionId, admissionContext); + await preflightPreparedInitialPromptModel({ env: ctx.env, userId: ctx.userId, @@ -140,7 +151,12 @@ export function createSessionExecutionV2Handlers() { condenseOnComplete: input.condenseOnComplete, } satisfies TurnFinalization, }; - const admissionContext = { env: ctx.env, userId: ctx.userId, botId: ctx.botId }; + const admissionContext = { + env: ctx.env, + userId: ctx.userId, + botId: ctx.botId, + authToken: ctx.authToken, + }; const ack = turn.type === 'prompt' ? await preflightAndQueuePromptMessage( diff --git a/services/cloud-agent-next/src/router/handlers/session-prepare.ts b/services/cloud-agent-next/src/router/handlers/session-prepare.ts index 007ab40850..ce4fc2472f 100644 --- a/services/cloud-agent-next/src/router/handlers/session-prepare.ts +++ b/services/cloud-agent-next/src/router/handlers/session-prepare.ts @@ -49,6 +49,7 @@ import type { SessionCreateRequest } from '../../session/session-requests.js'; import { assertKiloModelAvailable } from '../../model-validation.js'; import { assertRepositoryAccessBeforeSessionCreation } from '../../session/validate-repository-access.js'; import { assertOrganizationMembership } from './organization-membership.js'; +import jwt from 'jsonwebtoken'; type SessionPrepareHandlers = { prepareSession: typeof prepareSessionHandler; @@ -371,7 +372,12 @@ const prepareSessionHandler = internalApiProtectedProcedure }); } - if (requestWithProfile.initialTurn?.type === 'prompt') { + const claims = jwt.decode(ctx.authToken); + const isPolicyBearing = + claims !== null && + typeof claims === 'object' && + ('aud' in claims || 'tokenPurpose' in claims || 'credentialExchange' in claims); + if (requestWithProfile.initialTurn?.type === 'prompt' && !isPolicyBearing) { await assertKiloModelAvailable({ env: ctx.env, submittedModel: requestWithProfile.agent.model, diff --git a/services/cloud-agent-next/src/router/handlers/session-send.ts b/services/cloud-agent-next/src/router/handlers/session-send.ts index 1cce1e9c74..5b5c6ab9b4 100644 --- a/services/cloud-agent-next/src/router/handlers/session-send.ts +++ b/services/cloud-agent-next/src/router/handlers/session-send.ts @@ -47,7 +47,12 @@ const sendMessageHandler = protectedProcedure agent: input.agent, finalization: input.finalization, }; - const admissionContext = { env: ctx.env, userId: ctx.userId, botId: ctx.botId }; + const admissionContext = { + env: ctx.env, + userId: ctx.userId, + botId: ctx.botId, + authToken: ctx.authToken, + }; return preflightAndQueuePromptMessage(queuedMessage, admissionContext, 'send'); }); }); diff --git a/services/cloud-agent-next/src/router/handlers/session-start.ts b/services/cloud-agent-next/src/router/handlers/session-start.ts index 109f2b6d87..7c2ffd1200 100644 --- a/services/cloud-agent-next/src/router/handlers/session-start.ts +++ b/services/cloud-agent-next/src/router/handlers/session-start.ts @@ -26,6 +26,7 @@ import type { SessionCreateRequest } from '../../session/session-requests.js'; import { assertKiloModelAvailable } from '../../model-validation.js'; import { assertRepositoryAccessBeforeSessionCreation } from '../../session/validate-repository-access.js'; import { assertOrganizationMembership } from './organization-membership.js'; +import jwt from 'jsonwebtoken'; type SessionStartHandlers = { start: typeof startSessionHandler; @@ -126,14 +127,21 @@ const startSessionHandler = protectedProcedure requestWithProfile.agent.mode, requestWithProfile.profile?.resolved ?? {} ); - await assertKiloModelAvailable({ - env: ctx.env, - submittedModel: requestWithProfile.agent.model, - originalToken: ctx.authToken, - originalOrganizationId: requestWithProfile.options?.kilocodeOrganizationId, - createdOnPlatform: requestWithProfile.options?.createdOnPlatform, - procedure: 'start', - }); + const claims = jwt.decode(ctx.authToken); + const isPolicyBearing = + claims !== null && + typeof claims === 'object' && + ('aud' in claims || 'tokenPurpose' in claims || 'credentialExchange' in claims); + if (!isPolicyBearing) { + await assertKiloModelAvailable({ + env: ctx.env, + submittedModel: requestWithProfile.agent.model, + originalToken: ctx.authToken, + originalOrganizationId: requestWithProfile.options?.kilocodeOrganizationId, + createdOnPlatform: requestWithProfile.options?.createdOnPlatform, + procedure: 'start', + }); + } const registration = await startNewSession( requestWithProfile, diff --git a/services/cloud-agent-next/src/router/handlers/session-worktree.test.ts b/services/cloud-agent-next/src/router/handlers/session-worktree.test.ts index 3e987e076d..235c939912 100644 --- a/services/cloud-agent-next/src/router/handlers/session-worktree.test.ts +++ b/services/cloud-agent-next/src/router/handlers/session-worktree.test.ts @@ -23,6 +23,9 @@ const { withDORetryMock, createSessionForCloudAgentMock, deleteSessionForCloudAgentMock, + createRuntimeAuthorizationMock, + sealRuntimeAuthorizationMock, + verifyKiloTokenForPolicyMock, } = vi.hoisted(() => ({ admitOperationMock: vi.fn(), markReconcilePendingMock: vi.fn(), @@ -35,6 +38,9 @@ const { withDORetryMock: vi.fn(), createSessionForCloudAgentMock: vi.fn(), deleteSessionForCloudAgentMock: vi.fn(), + createRuntimeAuthorizationMock: vi.fn(), + sealRuntimeAuthorizationMock: vi.fn(), + verifyKiloTokenForPolicyMock: vi.fn(), })); vi.mock('@kilocode/db/operation-ledger', () => ({ @@ -70,6 +76,15 @@ vi.mock('../../utils/do-retry.js', () => ({ ) => withDORetryMock(getStub, operation, operationName), })); +vi.mock('@kilocode/worker-utils/runtime-authorization', () => ({ + createRuntimeAuthorization: createRuntimeAuthorizationMock, + sealRuntimeAuthorization: sealRuntimeAuthorizationMock, +})); + +vi.mock('@kilocode/worker-utils/kilo-token-policy', () => ({ + verifyKiloTokenForPolicy: verifyKiloTokenForPolicyMock, +})); + const USER_ID = 'oauth/google:1234'; const ORGANIZATION_ID = '11111111-1111-4111-8111-111111111111'; const OTHER_ORGANIZATION_ID = '22222222-2222-4222-8222-222222222222'; @@ -225,7 +240,9 @@ function fixture(options?: { ownershipResults?: OwnershipFixture[][]; internalSecret?: string; controlPlaneIds?: string; + runtimeIsolationEnabled?: string; botId?: string; + authToken?: string; }) { const userId = options?.userId ?? USER_ID; const metadata = @@ -243,6 +260,7 @@ function fixture(options?: { const sourceStub = { getMetadata: vi.fn().mockResolvedValue(metadata) }; const destinationStub = { getMetadata: vi.fn().mockResolvedValue(null), + getRuntimeAuthorizationStatus: vi.fn().mockResolvedValue('legacy'), registerSession: vi.fn().mockResolvedValue({ success: true }), createSessionWithInitialAdmission: vi.fn(), }; @@ -259,13 +277,15 @@ function fixture(options?: { }); const context = { userId, - authToken: CURRENT_AUTH_TOKEN, + authToken: options?.authToken ?? CURRENT_AUTH_TOKEN, ...(options?.botId ? { botId: options.botId } : {}), request: { headers } as Request, env: { INTERNAL_API_SECRET: INTERNAL_SECRET, + NEXTAUTH_SECRET: 'runtime-authorization-test-secret', CONTROL_PLANE_IDS: options?.controlPlaneIds ?? '*', WORKTREE_CREATION_ENABLED_IDS: '', + RUNTIME_ISOLATION_ENABLED: options?.runtimeIsolationEnabled ?? 'true', HYPERDRIVE: { connectionString: 'postgres://worktree-handler-test' }, SANDBOX_SESSION: sandboxSessionNamespace, CLOUD_AGENT_SESSION: legacySessionNamespace, @@ -333,6 +353,12 @@ beforeEach(() => { ); settleOperationMock.mockResolvedValue({ settled: true }); deleteSessionForCloudAgentMock.mockResolvedValue(undefined); + verifyKiloTokenForPolicyMock.mockResolvedValue({ claims: {} }); + createRuntimeAuthorizationMock.mockResolvedValue({ + authorization: { id: 'destination-authorization', state: 'active' }, + token: 'destination-delegated-token', + }); + sealRuntimeAuthorizationMock.mockResolvedValue('destination-seal'); createSessionForCloudAgentMock.mockResolvedValue({ status: 'ready', clone: { sessionId: DESTINATION_KILO_SESSION_ID, copiedItemCount: 0 }, @@ -521,6 +547,69 @@ describe('createWorktreeChat request validation and authorization', () => { }); describe('createWorktreeChat ownership, metadata, and control-plane routing', () => { + it('rejects a new destination before recording ownership when runtime isolation is disabled', async () => { + const { caller, input, destinationStub } = fixture({ runtimeIsolationEnabled: 'false' }); + verifyKiloTokenForPolicyMock.mockResolvedValue({ claims: { runtimeAdmission: {} } }); + + await expect(caller.createWorktreeChat(input)).rejects.toMatchObject({ + code: 'PRECONDITION_FAILED', + message: 'runtime_isolation_unavailable', + }); + expect(recordOperationProgressMock).not.toHaveBeenCalled(); + expect(createSessionForCloudAgentMock).not.toHaveBeenCalled(); + expect(destinationStub.registerSession).not.toHaveBeenCalled(); + }); + + it('preserves legacy registration for audience-bound tokens without modern policy markers', async () => { + const controlToken = 'legacy.header.signature'; + verifyKiloTokenForPolicyMock.mockResolvedValue({ + claims: { aud: 'cloud-agent-next' }, + }); + const { caller, input, destinationStub } = fixture({ authToken: controlToken }); + + await caller.createWorktreeChat(input); + + expect(createRuntimeAuthorizationMock).not.toHaveBeenCalled(); + expect(destinationStub.registerSession).toHaveBeenCalledWith( + expect.objectContaining({ + auth: { kiloSessionId: DESTINATION_KILO_SESSION_ID, kilocodeToken: controlToken }, + }) + ); + }); + + it('derives and seals authority for the exact destination, never persisting control authority', async () => { + const controlToken = 'header.payload.signature'; + verifyKiloTokenForPolicyMock.mockResolvedValue({ + claims: { aud: 'cloud-agent-next', runtimeAdmission: {} }, + }); + const { caller, input, destinationStub } = fixture({ authToken: controlToken }); + + await caller.createWorktreeChat(input); + + expect(createRuntimeAuthorizationMock).toHaveBeenCalledWith( + expect.objectContaining({ + token: controlToken, + resourceKind: 'cloud-agent-next', + resourceId: DESTINATION_WORKSPACE_ID, + }) + ); + expect(destinationStub.registerSession).toHaveBeenCalledWith( + expect.objectContaining({ + auth: { + kiloSessionId: DESTINATION_KILO_SESSION_ID, + kilocodeToken: 'destination-delegated-token', + }, + runtimeAuthorizationSeal: 'destination-seal', + }) + ); + const progress = recordOperationProgressMock.mock.calls[0]?.[2] as Record; + const completed = settleOperationMock.mock.calls.at(-1)?.[1] as Record; + expect(JSON.stringify(progress)).not.toContain(controlToken); + expect(JSON.stringify(progress)).not.toContain('destination-seal'); + expect(JSON.stringify(completed)).not.toContain(controlToken); + expect(JSON.stringify(completed)).not.toContain('destination-seal'); + }); + it.each([ { autoCommit: true, condenseOnComplete: true }, { autoCommit: false, condenseOnComplete: true }, @@ -795,6 +884,79 @@ describe('createWorktreeChat operation-ledger replay and conflict handling', () }); describe('createWorktreeChat registration rollback and unknown-outcome reconciliation', () => { + it('recreates a fresh destination seal after a lost response', async () => { + const controlToken = 'header.payload.signature'; + verifyKiloTokenForPolicyMock.mockResolvedValue({ + claims: { aud: 'cloud-agent-next', runtimeAdmission: {} }, + }); + const { caller, input, metadata, destinationStub } = fixture({ + authToken: controlToken, + ownershipResults: [[ownershipRow()], [ownershipRow()], [destinationOwnershipRow()]], + }); + destinationStub.registerSession.mockRejectedValueOnce(new Error('lost response')); + + await expect(caller.createWorktreeChat(input)).rejects.toThrow('lost response'); + const progress = recordOperationProgressMock.mock.calls[0]?.[2] as Record; + admitOperationMock.mockResolvedValueOnce({ + admission: 'duplicate_reconcile_pending', + row: ledgerRow({ status: 'reconcile_pending', canonical_result: progress }), + }); + destinationStub.getMetadata.mockResolvedValueOnce(null); + sealRuntimeAuthorizationMock.mockResolvedValueOnce('fresh-recovery-seal'); + + await caller.createWorktreeChat(input); + + expect(createRuntimeAuthorizationMock).toHaveBeenCalledTimes(2); + expect(destinationStub.registerSession.mock.calls[1]?.[0]).toMatchObject({ + auth: { kilocodeToken: 'destination-delegated-token' }, + runtimeAuthorizationSeal: 'fresh-recovery-seal', + }); + expect(JSON.stringify(metadata)).not.toContain('fresh-recovery-seal'); + }); + + it('fails closed when current control authority is revoked during recovery', async () => { + const controlToken = 'header.payload.signature'; + verifyKiloTokenForPolicyMock.mockResolvedValueOnce({ + claims: { aud: 'cloud-agent-next', runtimeAdmission: {} }, + }); + const { caller, input, destinationStub } = fixture({ + authToken: controlToken, + ownershipResults: [[ownershipRow()], [ownershipRow()], [destinationOwnershipRow()]], + }); + destinationStub.registerSession.mockRejectedValueOnce(new Error('lost response')); + await expect(caller.createWorktreeChat(input)).rejects.toThrow('lost response'); + + const progress = recordOperationProgressMock.mock.calls[0]?.[2] as Record; + admitOperationMock.mockResolvedValueOnce({ + admission: 'duplicate_reconcile_pending', + row: ledgerRow({ status: 'reconcile_pending', canonical_result: progress }), + }); + verifyKiloTokenForPolicyMock.mockRejectedValueOnce(new Error('revoked')); + + await expect(caller.createWorktreeChat(input)).rejects.toMatchObject({ code: 'FORBIDDEN' }); + expect(destinationStub.registerSession).toHaveBeenCalledTimes(1); + }); + + it('does not accept a modern committed destination without its private authorization record', async () => { + const controlToken = 'header.payload.signature'; + verifyKiloTokenForPolicyMock.mockResolvedValue({ + claims: { aud: 'cloud-agent-next', runtimeAdmission: {} }, + }); + const { caller, input, metadata, destinationStub } = fixture({ authToken: controlToken }); + destinationStub.getMetadata.mockResolvedValueOnce(destinationMetadata(metadata)); + destinationStub.getRuntimeAuthorizationStatus.mockResolvedValueOnce('revoked'); + admitOperationMock.mockResolvedValueOnce({ + admission: 'duplicate_reconcile_pending', + row: ledgerRow({ + status: 'reconcile_pending', + canonical_result: await progressFor(input), + }), + }); + + await expect(caller.createWorktreeChat(input)).rejects.toMatchObject({ code: 'FORBIDDEN' }); + expect(destinationStub.registerSession).not.toHaveBeenCalled(); + }); + it('rolls back only the empty ownership row after an explicit registration rejection', async () => { const { caller, input, destinationStub } = fixture(); destinationStub.registerSession.mockResolvedValueOnce({ diff --git a/services/cloud-agent-next/src/router/handlers/session-worktree.ts b/services/cloud-agent-next/src/router/handlers/session-worktree.ts index f4fbeed710..3ecfeb9b52 100644 --- a/services/cloud-agent-next/src/router/handlers/session-worktree.ts +++ b/services/cloud-agent-next/src/router/handlers/session-worktree.ts @@ -13,6 +13,11 @@ import { type CloudAgentWorktreeId, } from '@kilocode/session-ingest-contracts'; import { normalizeGitUrl } from '@kilocode/worker-utils'; +import { + createRuntimeAuthorization, + sealRuntimeAuthorization, +} from '@kilocode/worker-utils/runtime-authorization'; +import { verifyKiloTokenForPolicy } from '@kilocode/worker-utils/kilo-token-policy'; import { and, eq } from 'drizzle-orm'; import { z } from 'zod'; @@ -26,6 +31,7 @@ import { getSandboxSessionStub } from '../../sandbox-session/session-stub.js'; import { generateSessionId, isControlPlaneOwner } from '../../session-plane.js'; import { assertSessionOperationIdentity, + assertRuntimeIsolationAdmission, SESSION_CREATE_INTENT_FINGERPRINT_KEY, } from '../../session/session-registration.js'; import type { TRPCContext } from '../../types.js'; @@ -34,6 +40,7 @@ import { generateKiloSessionId } from '../../utils/kilo-session-id.js'; import { sha256Hex } from '../../utils/sha256.js'; import { getWorktreeWorkspacePath } from '../../workspace.js'; import { internalApiProtectedProcedure } from '../auth.js'; +import { resolveSecret } from '../../auth.js'; import { assertOrganizationMembership } from './organization-membership.js'; const workspaceSessionIdSchema = z.templateLiteral(['workspace_', z.uuid()]); @@ -297,6 +304,74 @@ function resultFromProgress(progress: OperationProgress, replayed = false): Work return replayed ? { ...result, replayed: true } : result; } +type WorktreeRuntimeAuthorization = { token: string; seal: string } | undefined; + +/** + * Establish whether the current caller presented modern control authority. + * This deliberately verifies the token instead of inferring its type from an + * untrusted decoded JWT payload. The actual authorization is re-created for + * each registration RPC below, which also re-checks principal and membership + * bindings. + */ +async function requiresRuntimeAuthorization(ctx: TRPCContext): Promise { + const secret = await resolveSecret(ctx.env.NEXTAUTH_SECRET); + if (!secret) { + throw new TRPCError({ code: 'INTERNAL_SERVER_ERROR', message: 'Authentication unavailable' }); + } + try { + const verified = await verifyKiloTokenForPolicy(ctx.authToken, secret, { + audience: 'cloud-agent-next', + mode: 'allow-legacy', + }); + const modern = + verified.claims.tokenPurpose !== undefined || + verified.claims.credentialExchange !== undefined || + verified.claims.runtimeAdmission !== undefined; + if (modern && verified.claims.runtimeAdmission === undefined) { + throw new Error('Missing runtime admission'); + } + return modern; + } catch { + throw new TRPCError({ code: 'FORBIDDEN', message: 'Runtime authorization denied' }); + } +} + +async function createDestinationRuntimeAuthorization( + ctx: TRPCContext, + progress: OperationProgress, + organizationId: string | undefined +): Promise { + if (!(await requiresRuntimeAuthorization(ctx))) return undefined; + + const secret = await resolveSecret(ctx.env.NEXTAUTH_SECRET); + if (!secret) { + throw new TRPCError({ code: 'INTERNAL_SERVER_ERROR', message: 'Authentication unavailable' }); + } + try { + const created = await createRuntimeAuthorization({ + token: ctx.authToken, + secret, + connectionString: ctx.env.HYPERDRIVE.connectionString, + resourceKind: 'cloud-agent-next', + resourceId: progress.cloudAgentSessionId, + ...(organizationId ? { organizationId } : {}), + }); + return { + token: created.token, + seal: await sealRuntimeAuthorization(created.authorization, secret), + }; + } catch { + // Membership, principal, and token admission can all change between a + // lost response and a replay. Never revive a destination on stale control + // authority. + throw new TRPCError({ code: 'FORBIDDEN', message: 'Runtime authorization denied' }); + } +} + +async function assertNewDestinationRuntimeIsolation(ctx: TRPCContext): Promise { + if (await requiresRuntimeAuthorization(ctx)) assertRuntimeIsolationAdmission(ctx.env); +} + function sourceWorktreeBranchName(source: WorktreeSource): string { return ( source.workspace.branchName ?? @@ -308,7 +383,8 @@ function sourceWorktreeBranchName(source: WorktreeSource): string { function buildRegistrationInput( source: WorktreeSource, ctx: TRPCContext, - progress: OperationProgress + progress: OperationProgress, + runtimeAuthorization: WorktreeRuntimeAuthorization ): Parameters['registerSession']>[0] { const repository = { ...source.repository }; if ('token' in repository) delete repository.token; @@ -319,15 +395,34 @@ function buildRegistrationInput( return { identity: { ...source.metadata.identity, sessionId: progress.cloudAgentSessionId }, - auth: { kiloSessionId: progress.kiloSessionId, kilocodeToken: ctx.authToken }, + auth: { + kiloSessionId: progress.kiloSessionId, + kilocodeToken: runtimeAuthorization?.token ?? ctx.authToken, + }, agent: source.metadata.agent, repository, workspace, ...(source.metadata.profile ? { profile: source.metadata.profile } : {}), ...(source.metadata.finalization ? { finalization: source.metadata.finalization } : {}), + ...(runtimeAuthorization ? { runtimeAuthorizationSeal: runtimeAuthorization.seal } : {}), }; } +async function assertDestinationRuntimeAuthorizationActive( + ctx: TRPCContext, + sessionId: string +): Promise { + if (!(await requiresRuntimeAuthorization(ctx))) return; + const status = await withDORetry( + () => getSandboxSessionStub(ctx.env, ctx.userId, sessionId), + stub => stub.getRuntimeAuthorizationStatus(), + 'getRuntimeAuthorizationStatus' + ); + if (status !== 'active') { + throw new TRPCError({ code: 'FORBIDDEN', message: 'Runtime authorization denied' }); + } +} + function assertRegisteredMetadata( rawMetadata: unknown, source: WorktreeSource, @@ -571,7 +666,6 @@ async function registerWorktreeSession( cloudAgentSessionId: progress.cloudAgentSessionId, kiloSessionId: progress.kiloSessionId, }; - const registrationInput = buildRegistrationInput(source, ctx, progress); let registrationAttempted = false; let response: unknown; @@ -583,6 +677,7 @@ async function registerWorktreeSession( const existing = await stub.getMetadata(); if (existing) { assertRegisteredMetadata(existing, source, progress); + await assertDestinationRuntimeAuthorizationActive(ctx, progress.cloudAgentSessionId); logControlDiagnostic('worktree_chat_reconciliation', { ...diagnostic, result: 'registration_recovered', @@ -592,7 +687,15 @@ async function registerWorktreeSession( } } registrationAttempted = true; - return stub.registerSession(registrationInput); + await assertNewDestinationRuntimeIsolation(ctx); + const runtimeAuthorization = await createDestinationRuntimeAuthorization( + ctx, + progress, + source.ownership.organizationId ?? undefined + ); + return stub.registerSession( + buildRegistrationInput(source, ctx, progress, runtimeAuthorization) + ); }, 'registerSession' ); @@ -703,6 +806,7 @@ async function executeWorktreeCreate( fingerprint: string ): Promise { const startedAt = Date.now(); + await assertNewDestinationRuntimeIsolation(ctx); const progress = operationProgressSchema.parse({ cloudAgentSessionId: generateSessionId('control'), kiloSessionId: generateKiloSessionId(), @@ -804,6 +908,9 @@ async function reconcileWorktreeCreate( ); throw error; } + await assertDestinationRuntimeAuthorizationActive(ctx, progress.cloudAgentSessionId); + } else { + await assertNewDestinationRuntimeIsolation(ctx); } const ownership = await findOwnershipRow( diff --git a/services/cloud-agent-next/src/runtime-credential-proxy-rpc.test.ts b/services/cloud-agent-next/src/runtime-credential-proxy-rpc.test.ts new file mode 100644 index 0000000000..7ec66e57d2 --- /dev/null +++ b/services/cloud-agent-next/src/runtime-credential-proxy-rpc.test.ts @@ -0,0 +1,357 @@ +import { describe, expect, it, vi } from 'vitest'; +import jwt from 'jsonwebtoken'; +import type { RuntimeAuthorization } from '@kilocode/worker-utils/runtime-authorization-contract'; +import type { SessionMetadata } from './persistence/session-metadata.js'; +import { + issuePersistedRuntimeProxyGrant, + resolvePersistedRuntimeProxyCredential, +} from './runtime-credential-proxy-rpc.js'; +import { RUNTIME_PROXY_GRANT_KEY, type RuntimeProxyGrant } from './runtime-credential-proxy.js'; + +const secret = 'test-secret'; +const env = { NEXTAUTH_SECRET: secret } as never; +const authorizationId = '00000000-0000-4000-8000-000000000001'; + +type Fence = { + plane: 'legacy'; + generation: number; + allocationId: string; + wrapperRunId: string; + wrapperConnectionId: string; +}; + +function authorization(state: 'active' | 'revoked' = 'active'): RuntimeAuthorization { + const issuedAt = new Date(Date.now()); + return { + version: 1, + id: authorizationId, + resourceKind: 'cloud-agent-next', + resourceId: 'agent_1', + userId: 'user_1', + authorizationUserId: 'user_1', + organizationId: 'org_1', + issuedAt: issuedAt.toISOString(), + delegationExpiresAt: new Date(issuedAt.getTime() + 24 * 60 * 60_000).toISOString(), + state, + bindings: { userPepperDigest: 'a'.repeat(64), authorizationPepperDigest: 'b'.repeat(64) }, + source: { admissionSource: 'user' }, + }; +} + +function metadata(kiloSessionId = 'kilo_1'): SessionMetadata { + return { + metadataSchemaVersion: 2, + identity: { sessionId: 'agent_1', userId: 'user_1', orgId: 'org_1' }, + auth: { kiloSessionId }, + lifecycle: { version: 1, timestamp: 0 }, + }; +} + +function fence(generation = 1): Fence { + return { + plane: 'legacy', + generation, + allocationId: 'allocation_1', + wrapperRunId: 'run_1', + wrapperConnectionId: 'connection_1', + }; +} + +function signedToken(expiresAt: number, nonce?: string): string { + return jwt.sign({ exp: Math.floor(expiresAt / 1000), ...(nonce ? { nonce } : {}) }, secret, { + noTimestamp: true, + }); +} + +function storage() { + const values = new Map(); + return { + get: async (key: string) => values.get(key) as T | undefined, + put: async (key: string, value: unknown) => { + values.set(key, value); + }, + }; +} + +async function issue(input: { + store: ReturnType; + currentMetadata?: SessionMetadata | null; + currentAuthorization?: RuntimeAuthorization | null; + currentFence?: Fence | null; + token?: string | null; +}) { + return issuePersistedRuntimeProxyGrant({ + env, + storage: input.store, + metadata: input.currentMetadata === undefined ? metadata() : input.currentMetadata, + authorization: + input.currentAuthorization === undefined ? authorization() : input.currentAuthorization, + fence: input.currentFence === undefined ? fence() : input.currentFence, + token: input.token === undefined ? signedToken(Date.now() + 2 * 60 * 60_000) : input.token, + mode: 'contained', + }); +} + +describe('persisted runtime credential proxy RPC', () => { + it('issues only for active authorization, current session metadata and fence, and a live token', async () => { + vi.useFakeTimers(); + vi.setSystemTime(new Date('2026-01-01T00:00:00.000Z')); + const valid = await issue({ store: storage() }); + expect(valid).toEqual(expect.any(String)); + + const expired = signedToken(Date.now() - 1_000); + for (const input of [ + { currentAuthorization: authorization('revoked') }, + { currentMetadata: null }, + { currentMetadata: metadata('') }, + { currentFence: null }, + { token: expired }, + ]) { + await expect(issue({ store: storage(), ...input })).resolves.toBeNull(); + } + vi.useRealTimers(); + }); + + it('returns the same stable handle across backing-token renewal', async () => { + vi.useFakeTimers(); + vi.setSystemTime(new Date('2026-01-01T00:00:00.000Z')); + const store = storage(); + const first = await issue({ + store, + token: signedToken(Date.now() + 2 * 60 * 60_000, 'backing-token-before-renewal'), + }); + vi.advanceTimersByTime(60_000); + const second = await issue({ + store, + token: signedToken(Date.now() + 4 * 60 * 60_000, 'backing-token-after-renewal'), + }); + expect(second).toBe(first); + vi.useRealTimers(); + }); + + it('bounds its independent transport lease at one day even with a backing token beyond one hour', async () => { + vi.useFakeTimers(); + vi.setSystemTime(new Date('2026-01-01T00:00:00.000Z')); + const store = storage(); + await issue({ store, token: signedToken(Date.now() + 2 * 60 * 60_000) }); + + const grant = await store.get(RUNTIME_PROXY_GRANT_KEY); + expect(grant?.leaseExpiresAt).toBe(Date.now() + 24 * 60 * 60_000); + vi.useRealTimers(); + }); + + it('caps a proxy lease to the remaining delegation duration and issues no handle at its deadline', async () => { + vi.useFakeTimers(); + vi.setSystemTime(new Date('2026-01-01T23:40:00.000Z')); + const store = storage(); + const expiringAuthorization = authorization(); + expiringAuthorization.delegationExpiresAt = '2026-01-02T00:00:00.000Z'; + + await expect(issue({ store, currentAuthorization: expiringAuthorization })).resolves.toEqual( + expect.any(String) + ); + const grant = await store.get(RUNTIME_PROXY_GRANT_KEY); + expect(grant?.leaseExpiresAt).toBe(Date.UTC(2026, 0, 2)); + + vi.setSystemTime(new Date('2026-01-02T00:00:00.000Z')); + await expect( + issue({ store: storage(), currentAuthorization: expiringAuthorization }) + ).resolves.toBeNull(); + vi.useRealTimers(); + }); + + it('resolves a valid handle with its existing backing token', async () => { + const store = storage(); + const token = signedToken(Date.now() + 10 * 60_000); + const handle = await issue({ store, token }); + const getToken = vi.fn(async () => token); + + await expect( + resolvePersistedRuntimeProxyCredential({ + env, + storage: store, + handle: handle!, + metadata: async () => metadata(), + authorization: async () => authorization(), + fence: async () => fence(), + token: getToken, + }) + ).resolves.toEqual({ + token, + organizationId: 'org_1', + runtimeAuthorization: { + userId: 'user_1', + authorizationId, + resourceId: 'agent_1', + }, + }); + expect(getToken).toHaveBeenCalledTimes(1); + }); + + it('renews a near-expiry backing token without changing the persisted handle grant lease', async () => { + vi.useFakeTimers(); + vi.setSystemTime(new Date('2026-01-01T00:00:00.000Z')); + const store = storage(); + const nearExpiry = signedToken(Date.now() + 5 * 60_000); + const renewed = signedToken(Date.now() + 2 * 60 * 60_000); + const handle = await issue({ store, token: nearExpiry }); + const before = await store.get(RUNTIME_PROXY_GRANT_KEY); + const getToken = vi.fn(async () => (getToken.mock.calls.length === 1 ? nearExpiry : renewed)); + + await expect( + resolvePersistedRuntimeProxyCredential({ + env, + storage: store, + handle: handle!, + metadata: async () => metadata(), + authorization: async () => authorization(), + fence: async () => fence(), + token: getToken, + }) + ).resolves.toMatchObject({ token: renewed }); + expect(getToken).toHaveBeenCalledTimes(2); + expect(await store.get(RUNTIME_PROXY_GRANT_KEY)).toEqual(before); + vi.useRealTimers(); + }); + + it('cannot extend a transport lease through repeated resolve calls', async () => { + vi.useFakeTimers(); + vi.setSystemTime(new Date('2026-01-01T00:00:00.000Z')); + const store = storage(); + const token = signedToken(Date.now() + 2 * 60 * 60_000); + const handle = await issue({ store, token }); + const original = await store.get(RUNTIME_PROXY_GRANT_KEY); + const input = { + env, + storage: store, + handle: handle!, + metadata: async () => metadata(), + authorization: async () => authorization(), + fence: async () => fence(), + token: async () => token, + }; + + await resolvePersistedRuntimeProxyCredential(input); + vi.advanceTimersByTime(60 * 60_000); + await resolvePersistedRuntimeProxyCredential(input); + expect(await store.get(RUNTIME_PROXY_GRANT_KEY)).toEqual(original); + vi.useRealTimers(); + }); + + it('denies a revoked authorization', async () => { + const store = storage(); + const token = signedToken(Date.now() + 10 * 60_000); + const handle = await issue({ store, token }); + + await expect( + resolvePersistedRuntimeProxyCredential({ + env, + storage: store, + handle: handle!, + metadata: async () => metadata(), + authorization: async () => authorization('revoked'), + fence: async () => fence(), + token: async () => token, + }) + ).resolves.toBeNull(); + }); + + it('rejects a proxy credential when authorization is revoked while token I/O is pending', async () => { + const store = storage(); + const token = signedToken(Date.now() + 10 * 60_000); + const handle = await issue({ store, token }); + let currentAuthorization = authorization(); + const started = Promise.withResolvers(); + const release = Promise.withResolvers(); + + const result = resolvePersistedRuntimeProxyCredential({ + env, + storage: store, + handle: handle!, + metadata: async () => metadata(), + authorization: async () => currentAuthorization, + fence: async () => fence(), + token: async () => { + started.resolve(); + await release.promise; + return token; + }, + }); + await started.promise; + currentAuthorization = authorization('revoked'); + release.resolve(); + + await expect(result).resolves.toBeNull(); + }); + + it('rejects a proxy credential when authorization expires while token I/O is pending', async () => { + vi.useFakeTimers(); + vi.setSystemTime(new Date('2026-01-01T00:00:00.000Z')); + const store = storage(); + const expiringAuthorization = authorization(); + expiringAuthorization.delegationExpiresAt = new Date(Date.now() + 1_000).toISOString(); + const token = signedToken(Date.now() + 10 * 60_000); + const handle = await issue({ store, currentAuthorization: expiringAuthorization, token }); + const started = Promise.withResolvers(); + const release = Promise.withResolvers(); + + const result = resolvePersistedRuntimeProxyCredential({ + env, + storage: store, + handle: handle!, + metadata: async () => metadata(), + authorization: async () => expiringAuthorization, + fence: async () => fence(), + token: async () => { + started.resolve(); + await release.promise; + return token; + }, + }); + await started.promise; + vi.advanceTimersByTime(1_000); + release.resolve(); + + await expect(result).resolves.toBeNull(); + vi.useRealTimers(); + }); + + it.each(['fence', 'grant'] as const)( + '%s changes are handled according to the stable grant fence', + async replacement => { + vi.useFakeTimers(); + vi.setSystemTime(new Date('2026-01-01T00:00:00.000Z')); + const store = storage(); + const nearExpiry = signedToken(Date.now() + 5 * 60_000); + const renewed = signedToken(Date.now() + 2 * 60 * 60_000); + const handle = await issue({ store, token: nearExpiry }); + let currentFence = fence(); + let calls = 0; + + const resolved = await resolvePersistedRuntimeProxyCredential({ + env, + storage: store, + handle: handle!, + metadata: async () => metadata(), + authorization: async () => authorization(), + fence: async () => currentFence, + token: async () => { + calls += 1; + if (calls === 2 && replacement === 'fence') currentFence = fence(2); + if (calls === 2 && replacement === 'grant') { + await issue({ store, token: renewed }); + } + return calls === 1 ? nearExpiry : renewed; + }, + }); + if (replacement === 'fence') { + expect(resolved).toBeNull(); + } else { + // Reissuing against the same runtime fence returns the original handle + // and does not invalidate a request that is concurrently renewing. + expect(resolved).toMatchObject({ token: renewed }); + } + vi.useRealTimers(); + } + ); +}); diff --git a/services/cloud-agent-next/src/runtime-credential-proxy-rpc.ts b/services/cloud-agent-next/src/runtime-credential-proxy-rpc.ts new file mode 100644 index 0000000000..b68f67273f --- /dev/null +++ b/services/cloud-agent-next/src/runtime-credential-proxy-rpc.ts @@ -0,0 +1,200 @@ +import jwt from 'jsonwebtoken'; +import type { RuntimeAuthorization } from '@kilocode/worker-utils/runtime-authorization-contract'; +import { + createRuntimeProxyGrant, + issueRuntimeCredentialProxyHandle, + matchesRuntimeProxyGrant, + resolveRuntimeProxyCredential, + RUNTIME_PROXY_GRANT_KEY, + runtimeProxyGrantSchema, + verifyRuntimeCredentialProxyHandle, + type RuntimeProxyFence, + type RuntimeProxyGrant, +} from './runtime-credential-proxy.js'; +import type { SessionMetadata } from './persistence/session-metadata.js'; +import type { Env } from './types.js'; + +type Storage = { + get(key: string): Promise; + put(key: string, value: unknown): Promise; +}; + +const RUNTIME_PROXY_LEASE_MS = 24 * 60 * 60_000; + +function tokenExpiry(token: string): number | null { + const decoded = jwt.decode(token); + return typeof decoded === 'object' && decoded !== null && typeof decoded.exp === 'number' + ? decoded.exp * 1000 + : null; +} + +function context(metadata: SessionMetadata, fence: RuntimeProxyFence) { + const kiloSessionId = metadata.auth.kiloSessionId; + if (!kiloSessionId) return null; + return { + sessionId: metadata.identity.sessionId, + kiloSessionId, + userId: metadata.identity.userId, + ...(metadata.identity.orgId ? { orgId: metadata.identity.orgId } : {}), + fence, + }; +} + +function sameFence(left: RuntimeProxyFence, right: RuntimeProxyFence): boolean { + if (left.plane !== right.plane || left.allocationId !== right.allocationId) return false; + return left.plane === 'legacy' && right.plane === 'legacy' + ? left.generation === right.generation && + left.wrapperRunId === right.wrapperRunId && + left.wrapperConnectionId === right.wrapperConnectionId + : left.plane === 'control' && + right.plane === 'control' && + left.providerInstanceId === right.providerInstanceId && + left.connectionId === right.connectionId && + left.wrapperInstanceId === right.wrapperInstanceId; +} + +/** + * Shared private-DO grant lifecycle. The transport lease outlives individual + * backing tokens so renewal remains transparent to Kilo, but proxy requests + * can never lengthen the persisted lease. + */ +export async function issuePersistedRuntimeProxyGrant(input: { + env: Pick; + storage: Storage; + metadata: SessionMetadata | null; + authorization: RuntimeAuthorization | null; + fence: RuntimeProxyFence | null; + token: string | null; + mode: RuntimeProxyGrant['mode']; + now?: number; +}): Promise { + const now = input.now ?? Date.now(); + const current = input.metadata && input.fence ? context(input.metadata, input.fence) : null; + const tokenExpiresAt = input.token ? tokenExpiry(input.token) : null; + const delegationExpiresAt = + input.authorization === null ? null : Date.parse(input.authorization.delegationExpiresAt); + if ( + !current || + !tokenExpiresAt || + tokenExpiresAt <= now || + input.authorization?.state !== 'active' || + !delegationExpiresAt || + delegationExpiresAt <= now + ) + return null; + const existing = await input.storage.get(RUNTIME_PROXY_GRANT_KEY); + const parsedExisting = runtimeProxyGrantSchema.safeParse(existing); + if ( + parsedExisting.success && + parsedExisting.data.issuedAt !== undefined && + parsedExisting.data.authorizationId === input.authorization.id && + parsedExisting.data.sessionId === current.sessionId && + parsedExisting.data.kiloSessionId === current.kiloSessionId && + parsedExisting.data.userId === current.userId && + parsedExisting.data.orgId === current.orgId && + parsedExisting.data.allocationId === current.fence.allocationId && + sameFence(parsedExisting.data, current.fence) && + parsedExisting.data.mode === input.mode && + parsedExisting.data.leaseExpiresAt > now && + parsedExisting.data.leaseExpiresAt <= delegationExpiresAt + ) { + return issueRuntimeCredentialProxyHandle( + input.env, + parsedExisting.data, + parsedExisting.data.issuedAt + ); + } + const issuedAt = now; + const { fence, ...identity } = current; + const grant = createRuntimeProxyGrant({ + authorizationId: input.authorization.id, + ...identity, + ...fence, + mode: input.mode, + leaseExpiresAt: Math.min(issuedAt + RUNTIME_PROXY_LEASE_MS, delegationExpiresAt), + state: 'active', + issuedAt, + }); + await input.storage.put(RUNTIME_PROXY_GRANT_KEY, grant); + return issueRuntimeCredentialProxyHandle(input.env, grant, issuedAt); +} + +export async function resolvePersistedRuntimeProxyCredential(input: { + env: Pick; + storage: Storage; + handle: string; + metadata: () => Promise; + authorization: () => Promise; + fence: () => Promise; + token: () => Promise; + now?: number; +}): Promise<{ + token: string; + organizationId?: string; + runtimeAuthorization: { userId: string; authorizationId: string; resourceId: string }; +} | null> { + const now = input.now ?? Date.now(); + const claims = await verifyRuntimeCredentialProxyHandle(input.env, input.handle); + if (!claims || !('sessionId' in claims)) return null; + const [metadata, authorization, fence, grant] = await Promise.all([ + input.metadata(), + input.authorization(), + input.fence(), + input.storage.get(RUNTIME_PROXY_GRANT_KEY), + ]); + const current = metadata && fence ? context(metadata, fence) : null; + if (!current || !authorization) return null; + if ( + !matchesRuntimeProxyGrant(grant, claims, { + ...current, + authorizationId: authorization.id, + now, + }) + ) + return null; + const backingToken = await input.token(); + if (!backingToken) return null; + const resolved = await resolveRuntimeProxyCredential({ + env: input.env, + handle: input.handle, + grant, + authorization, + context: current, + token: backingToken, + now, + renew: async () => (await input.token()) ?? '', + }); + if (!resolved?.token) return null; + + // Renewal awaits external I/O. Re-read all durable fences before exposing it. + const [latestMetadata, latestAuthorization, latestFence, latestGrant] = await Promise.all([ + input.metadata(), + input.authorization(), + input.fence(), + input.storage.get(RUNTIME_PROXY_GRANT_KEY), + ]); + const latest = latestMetadata && latestFence ? context(latestMetadata, latestFence) : null; + const latestNow = Date.now(); + if ( + !latest || + !latestAuthorization || + latestAuthorization.state !== 'active' || + Date.parse(latestAuthorization.delegationExpiresAt) <= latestNow || + !matchesRuntimeProxyGrant(latestGrant, claims, { + ...latest, + authorizationId: latestAuthorization.id, + now: latestNow, + }) + ) { + return null; + } + return { + token: resolved.token, + ...(latest.orgId ? { organizationId: latest.orgId } : {}), + runtimeAuthorization: { + userId: latestAuthorization.userId, + authorizationId: latestAuthorization.id, + resourceId: latestAuthorization.resourceId, + }, + }; +} diff --git a/services/cloud-agent-next/src/runtime-credential-proxy.test.ts b/services/cloud-agent-next/src/runtime-credential-proxy.test.ts new file mode 100644 index 0000000000..e6047562c8 --- /dev/null +++ b/services/cloud-agent-next/src/runtime-credential-proxy.test.ts @@ -0,0 +1,390 @@ +import { describe, expect, it } from 'vitest'; +import jwt from 'jsonwebtoken'; +import { + createRuntimeProxyGrant, + issueRuntimeCredentialProxyHandle, + matchesRuntimeProxyGrant, + runtimeProxyGrantSchema, + runtimeCredentialProxyBaseUrl, + runtimeCredentialProxyFacadeBaseUrl, + runtimeCredentialProxyUpstream, + verifyRuntimeCredentialProxyHandle, +} from './runtime-credential-proxy.js'; +import { resolveRuntimeCredentialProxyRoute } from './kilo/runtime-credential-proxy-routes.js'; + +const targets = { + backendBaseUrl: 'https://backend.example.test/base', + providerBaseUrl: 'https://provider.example.test/api/openrouter', + sessionIngestBaseUrl: 'https://ingest.example.test', +}; + +describe('runtime credential proxy', () => { + it('issues a bounded strict handle for an active runtime proxy grant', async () => { + const env = { NEXTAUTH_SECRET: 'test-secret' } as never; + const now = Date.now(); + const grant = createRuntimeProxyGrant({ + plane: 'legacy', + authorizationId: '11111111-1111-4111-8111-111111111111', + sessionId: 'agent_1', + kiloSessionId: 'kilo_1', + userId: 'user_1', + mode: 'contained', + generation: 1, + allocationId: 'allocation_1', + wrapperRunId: 'run_1', + wrapperConnectionId: 'connection_1', + leaseExpiresAt: now + 60_000, + state: 'active', + }); + const handle = await issueRuntimeCredentialProxyHandle(env, grant); + + const claims = await verifyRuntimeCredentialProxyHandle(env, handle); + expect(claims).toMatchObject({ + aud: 'cloud-agent-next:runtime-credential-proxy', + grantId: grant.grantId, + authorizationId: grant.authorizationId, + sessionId: grant.sessionId, + kiloSessionId: grant.kiloSessionId, + userId: grant.userId, + nonce: grant.nonce, + exp: Math.floor(grant.leaseExpiresAt / 1000), + }); + expect(claims?.iat).toBeGreaterThan(0); + if (grant.plane !== 'legacy') throw new Error('Expected legacy grant'); + expect( + claims && + matchesRuntimeProxyGrant(grant, claims, { + authorizationId: grant.authorizationId, + sessionId: grant.sessionId, + kiloSessionId: grant.kiloSessionId, + userId: grant.userId, + fence: { + plane: 'legacy', + generation: grant.generation, + allocationId: grant.allocationId, + wrapperRunId: grant.wrapperRunId, + wrapperConnectionId: grant.wrapperConnectionId, + }, + now, + }) + ).toBe(true); + await expect(verifyRuntimeCredentialProxyHandle(env, `${handle}x`)).resolves.toBeNull(); + }); + + it('rejects expired, unsigned, non-HS256, and unknown claims', async () => { + const env = { NEXTAUTH_SECRET: 'test-secret' } as never; + const expired = createRuntimeProxyGrant({ + plane: 'legacy', + authorizationId: '11111111-1111-4111-8111-111111111111', + sessionId: 'agent_1', + kiloSessionId: 'kilo_1', + userId: 'user_1', + mode: 'direct', + generation: 1, + allocationId: 'allocation_1', + wrapperRunId: 'run_1', + wrapperConnectionId: 'connection_1', + leaseExpiresAt: Date.now() - 1, + state: 'active', + }); + await expect(issueRuntimeCredentialProxyHandle(env, expired)).rejects.toThrow('has expired'); + await expect( + verifyRuntimeCredentialProxyHandle(env, 'eyJhbGciOiJub25lIn0.e30.') + ).resolves.toBeNull(); + const claims = { + aud: 'cloud-agent-next:runtime-credential-proxy', + grantId: '11111111-1111-4111-8111-111111111111', + authorizationId: '22222222-2222-4222-8222-222222222222', + sessionId: 'agent_1', + kiloSessionId: 'kilo_1', + userId: 'user_1', + nonce: 'a'.repeat(43), + iat: Math.floor(Date.now() / 1000) - 60, + exp: Math.floor(Date.now() / 1000) - 1, + }; + await expect( + verifyRuntimeCredentialProxyHandle( + env, + jwt.sign(claims, 'test-secret', { algorithm: 'HS256' }) + ) + ).resolves.toBeNull(); + await expect( + verifyRuntimeCredentialProxyHandle( + env, + jwt.sign({ ...claims, exp: claims.iat + 120, extra: true }, 'test-secret', { + algorithm: 'HS256', + }) + ) + ).resolves.toBeNull(); + await expect( + verifyRuntimeCredentialProxyHandle( + env, + jwt.sign({ ...claims, exp: claims.iat + 120 }, 'test-secret', { algorithm: 'HS384' }) + ) + ).resolves.toBeNull(); + }); + + it('keeps control and legacy runtime grant fences disjoint and decodes v1 as legacy only', async () => { + const env = { NEXTAUTH_SECRET: 'test-secret' } as never; + const now = Date.now(); + const control = createRuntimeProxyGrant({ + plane: 'control', + authorizationId: '11111111-1111-4111-8111-111111111111', + sessionId: 'agent_1', + kiloSessionId: 'kilo_1', + userId: 'user_1', + mode: 'contained', + allocationId: 'allocation_1', + providerInstanceId: 'provider_1', + connectionId: 'connection_1', + wrapperInstanceId: 'wrapper_1', + leaseExpiresAt: now + 60_000, + state: 'active', + }); + if (control.plane !== 'control') throw new Error('Expected control grant'); + const claims = await verifyRuntimeCredentialProxyHandle( + env, + await issueRuntimeCredentialProxyHandle(env, control) + ); + expect(claims).not.toBeNull(); + expect( + claims && + matchesRuntimeProxyGrant(control, claims, { + authorizationId: control.authorizationId, + sessionId: control.sessionId, + kiloSessionId: control.kiloSessionId, + userId: control.userId, + fence: { + plane: 'legacy', + generation: 1, + allocationId: control.allocationId, + wrapperRunId: 'run_1', + wrapperConnectionId: 'connection_1', + }, + now, + }) + ).toBe(false); + const { + providerInstanceId: _providerInstanceId, + connectionId: _connectionId, + wrapperInstanceId: _wrapperInstanceId, + plane: _plane, + ...common + } = control; + const legacyV1 = runtimeProxyGrantSchema.parse({ + ...common, + version: 1, + generation: 1, + wrapperRunId: 'run_1', + wrapperConnectionId: 'connection_1', + }); + expect(legacyV1.plane).toBe('legacy'); + }); + + it('maps only exact method-scoped provider, backend, and ingest routes', () => { + expect(runtimeCredentialProxyBaseUrl('https://worker.example.test/root/')).toBe( + 'https://worker.example.test/root/api/runtime-credential-proxy' + ); + expect( + runtimeCredentialProxyUpstream( + targets, + 'provider', + 'POST', + 'api/openrouter/chat/completions', + '?stream=true', + 'agent_1' + )?.toString() + ).toBe('https://provider.example.test/api/openrouter/chat/completions?stream=true'); + expect( + runtimeCredentialProxyUpstream( + targets, + 'backend', + 'GET', + 'api/user', + '', + 'agent_1' + )?.toString() + ).toBe('https://backend.example.test/base/api/user'); + expect( + runtimeCredentialProxyUpstream( + targets, + 'ingest', + 'POST', + 'api/session/agent_1/ingest', + '', + 'agent_1' + )?.toString() + ).toBe('https://ingest.example.test/api/session/agent_1/ingest'); + expect( + runtimeCredentialProxyUpstream(targets, 'backend', 'GET', 'trpc/admin', '', 'agent_1') + ).toBeNull(); + expect( + runtimeCredentialProxyUpstream( + targets, + 'ingest', + 'POST', + 'api/session/other/ingest', + '', + 'agent_1' + ) + ).toBeNull(); + }); + + it.each([ + ['origin', 'https://worker.example.test', 'https://worker.example.test'], + ['safe prefix', 'https://worker.example.test/runtime', 'https://worker.example.test/runtime'], + [ + 'trailing slash', + 'https://worker.example.test/runtime/', + 'https://worker.example.test/runtime', + ], + ['api prefix', 'https://worker.example.test/runtime/api', null], + ['non-default port', 'https://worker.example.test:8443/runtime', null], + ['HTTP', 'http://worker.example.test/runtime', null], + ['credentials', 'https://user@worker.example.test/runtime', null], + ['query', 'https://worker.example.test/runtime?next=x', null], + ['fragment', 'https://worker.example.test/runtime#next', null], + ['traversal', 'https://worker.example.test/runtime/../other', null], + ['encoded traversal', 'https://worker.example.test/runtime/%252e%252e/other', null], + ])('accepts only a safe facade %s', (_name, workerUrl, expected) => { + expect(runtimeCredentialProxyFacadeBaseUrl(workerUrl)).toBe(expected); + }); + + it('requires exact identity and strict JSON for session creation', () => { + expect( + resolveRuntimeCredentialProxyRoute({ + targets, + route: 'ingest', + method: 'POST', + pathname: '/api/session', + search: '', + kiloSessionId: 'kilo_1', + contentType: 'application/json; charset=utf-8', + bodyText: '{"sessionId":"kilo_1"}', + })?.toString() + ).toBe('https://ingest.example.test/api/session'); + for (const body of [ + '{}', + '{"sessionId":"other"}', + '{"sessionId":"kilo_1","extra":true}', + '[]', + ]) { + expect( + resolveRuntimeCredentialProxyRoute({ + targets, + route: 'ingest', + method: 'POST', + pathname: '/api/session', + search: '', + kiloSessionId: 'kilo_1', + contentType: 'application/json', + bodyText: body, + }) + ).toBeNull(); + } + }); + + it('fails closed for ambiguous paths and mismatched organization routes', () => { + for (const pathname of [ + '/api/openrouter', + '/chat/completions', + '/api/organizations/allowed/models', + '/api/openrouter/unknown', + '/api/openrouter/chat%2fcompletions', + '/api/openrouter/chat%252fcompletions', + '/api/openrouter/../chat/completions', + ]) { + expect( + resolveRuntimeCredentialProxyRoute({ + targets, + route: 'provider', + method: 'POST', + pathname, + search: '', + kiloSessionId: 'kilo_1', + }) + ).toBeNull(); + } + expect( + resolveRuntimeCredentialProxyRoute({ + targets, + route: 'provider', + method: 'POST', + pathname: '/api/gateway/chat/completions', + search: '', + kiloSessionId: 'kilo_1', + })?.toString() + ).toBe('https://provider.example.test/api/openrouter/chat/completions'); + expect( + resolveRuntimeCredentialProxyRoute({ + targets: { ...targets, providerBaseUrl: 'https://api.kilo.ai' }, + route: 'provider', + method: 'POST', + pathname: '/api/gateway/v1/chat/completions', + search: '', + kiloSessionId: 'kilo_1', + })?.toString() + ).toBe('https://api.kilo.ai/api/gateway/v1/chat/completions'); + expect( + resolveRuntimeCredentialProxyRoute({ + targets, + route: 'provider', + method: 'POST', + pathname: '/api/gateway/v1/responses', + search: '', + kiloSessionId: 'kilo_1', + })?.toString() + ).toBe('https://provider.example.test/api/openrouter/v1/responses'); + expect( + resolveRuntimeCredentialProxyRoute({ + targets, + route: 'provider', + method: 'GET', + pathname: '/api/gateway/v1/chat/completions', + search: '', + kiloSessionId: 'kilo_1', + }) + ).toBeNull(); + expect( + resolveRuntimeCredentialProxyRoute({ + targets: { ...targets, providerBaseUrl: 'https://api.kilo.ai' }, + route: 'provider', + method: 'POST', + pathname: '/api/openrouter/chat/completions', + search: '?stream=true', + kiloSessionId: 'kilo_1', + })?.toString() + ).toBe('https://api.kilo.ai/api/gateway/chat/completions?stream=true'); + expect( + resolveRuntimeCredentialProxyRoute({ + targets, + route: 'backend', + method: 'GET', + pathname: '/api/organizations/other/models', + search: '', + kiloSessionId: 'kilo_1', + organizationId: 'allowed', + }) + ).toBeNull(); + expect( + resolveRuntimeCredentialProxyRoute({ + targets, + route: 'provider', + method: 'POST', + pathname: '/api/openrouter/embeddings', + search: '', + kiloSessionId: 'kilo_1', + })?.toString() + ).toBe('https://provider.example.test/api/openrouter/embeddings'); + expect( + resolveRuntimeCredentialProxyRoute({ + targets: { ...targets, providerBaseUrl: 'ftp://provider.example.test/api/openrouter' }, + route: 'provider', + method: 'POST', + pathname: '/api/openrouter/embeddings', + search: '', + kiloSessionId: 'kilo_1', + }) + ).toBeNull(); + }); +}); diff --git a/services/cloud-agent-next/src/runtime-credential-proxy.ts b/services/cloud-agent-next/src/runtime-credential-proxy.ts new file mode 100644 index 0000000000..5567492ae8 --- /dev/null +++ b/services/cloud-agent-next/src/runtime-credential-proxy.ts @@ -0,0 +1,328 @@ +import jwt from 'jsonwebtoken'; +import { Buffer } from 'node:buffer'; +import { z } from 'zod'; +import { resolveSecret } from './auth.js'; +import { + resolveRuntimeCredentialProxyRoute, + type RuntimeCredentialProxyRoute, +} from './kilo/runtime-credential-proxy-routes.js'; +import type { Env } from './types.js'; + +const AUDIENCE = 'cloud-agent-next:runtime-credential-proxy'; +const sessionClaimsSchema = z + .object({ + aud: z.literal(AUDIENCE), + grantId: z.string().uuid(), + authorizationId: z.string().uuid(), + sessionId: z.string().min(1), + kiloSessionId: z.string().min(1), + userId: z.string().min(1), + nonce: z.string().regex(/^[A-Za-z0-9_-]{43}$/), + }) + .extend({ + iat: z.number().int().positive(), + exp: z.number().int().positive(), + }) + .strict() + .refine(claims => claims.exp > claims.iat); + +export type RuntimeProxyHandleClaims = z.infer; + +export const RUNTIME_PROXY_GRANT_KEY = 'runtime_proxy_grant'; +const runtimeProxyGrantBaseSchema = z + .object({ + grantId: z.string().uuid(), + authorizationId: z.string().uuid(), + sessionId: z.string().min(1), + kiloSessionId: z.string().min(1), + userId: z.string().min(1), + orgId: z.string().min(1).optional(), + nonce: z.string().regex(/^[A-Za-z0-9_-]{43}$/), + mode: z.enum(['direct', 'contained']), + allocationId: z.string().min(1), + issuedAt: z.number().int().positive().optional(), + leaseExpiresAt: z.number().int().positive(), + state: z.literal('active'), + }) + .strict(); + +const legacyRuntimeProxyGrantV2Schema = runtimeProxyGrantBaseSchema.extend({ + version: z.literal(2), + plane: z.literal('legacy'), + generation: z.number().int().nonnegative(), + wrapperRunId: z.string().min(1), + wrapperConnectionId: z.string().min(1), +}); +const controlRuntimeProxyGrantV2Schema = runtimeProxyGrantBaseSchema.extend({ + version: z.literal(2), + plane: z.literal('control'), + providerInstanceId: z.string().min(1), + connectionId: z.string().min(1), + wrapperInstanceId: z.string().min(1), +}); +const legacyRuntimeProxyGrantV1Schema = runtimeProxyGrantBaseSchema + .extend({ + version: z.literal(1), + generation: z.number().int().nonnegative(), + wrapperRunId: z.string().min(1), + wrapperConnectionId: z.string().min(1), + }) + .transform(grant => ({ ...grant, plane: 'legacy' as const })); + +/** v1 persisted grants are accepted temporarily as legacy only; never control. */ +export const runtimeProxyGrantSchema = z.union([ + legacyRuntimeProxyGrantV2Schema, + controlRuntimeProxyGrantV2Schema, + legacyRuntimeProxyGrantV1Schema, +]); +export type RuntimeProxyGrant = z.infer; +export type RuntimeProxyFence = + | { + plane: 'legacy'; + generation: number; + allocationId: string; + wrapperRunId: string; + wrapperConnectionId: string; + } + | { + plane: 'control'; + allocationId: string; + providerInstanceId: string; + connectionId: string; + wrapperInstanceId: string; + }; +type RuntimeProxyGrantInput = + | Omit, 'version' | 'grantId' | 'nonce'> + | Omit, 'version' | 'grantId' | 'nonce'>; + +export function createRuntimeProxyGrant(input: RuntimeProxyGrantInput): RuntimeProxyGrant { + return runtimeProxyGrantSchema.parse({ + version: 2, + grantId: crypto.randomUUID(), + nonce: Buffer.from(crypto.getRandomValues(new Uint8Array(32))).toString('base64url'), + ...input, + }); +} + +export async function issueRuntimeCredentialProxyHandle( + env: Pick, + input: RuntimeProxyGrant, + issuedAt = Date.now() +): Promise { + const grant = runtimeProxyGrantSchema.parse(input); + const secret = await resolveSecret(env.NEXTAUTH_SECRET); + if (!secret) throw new Error('Authentication unavailable'); + const iat = Math.floor(issuedAt / 1000); + const exp = Math.floor(grant.leaseExpiresAt / 1000); + if (exp <= iat) throw new Error('Runtime proxy grant has expired'); + return jwt.sign( + { + aud: AUDIENCE, + grantId: grant.grantId, + authorizationId: grant.authorizationId, + sessionId: grant.sessionId, + kiloSessionId: grant.kiloSessionId, + userId: grant.userId, + nonce: grant.nonce, + iat, + exp, + }, + secret, + { algorithm: 'HS256' } + ); +} + +export async function verifyRuntimeCredentialProxyHandle( + env: Pick, + value: string +): Promise { + const secret = await resolveSecret(env.NEXTAUTH_SECRET); + if (!secret || value.length > 4096) return null; + try { + const verified = jwt.verify(value, secret, { algorithms: ['HS256'] }); + const claims = sessionClaimsSchema.safeParse(verified); + return claims.success ? claims.data : null; + } catch { + return null; + } +} + +export function matchesRuntimeProxyGrant( + value: unknown, + handle: RuntimeProxyHandleClaims, + context: { + authorizationId: string; + sessionId: string; + kiloSessionId: string; + userId: string; + orgId?: string; + fence: RuntimeProxyFence; + now: number; + } +): value is RuntimeProxyGrant { + const grant = runtimeProxyGrantSchema.safeParse(value); + if (!grant.success) return false; + const current = grant.data; + return ( + current.state === 'active' && + current.leaseExpiresAt > context.now && + Math.floor(current.leaseExpiresAt / 1000) === handle.exp && + current.grantId === handle.grantId && + current.authorizationId === handle.authorizationId && + current.sessionId === handle.sessionId && + current.kiloSessionId === handle.kiloSessionId && + current.userId === handle.userId && + current.nonce === handle.nonce && + current.authorizationId === context.authorizationId && + current.sessionId === context.sessionId && + current.kiloSessionId === context.kiloSessionId && + current.userId === context.userId && + current.orgId === context.orgId && + current.plane === context.fence.plane && + current.allocationId === context.fence.allocationId && + (current.plane === 'legacy' && context.fence.plane === 'legacy' + ? current.generation === context.fence.generation && + current.wrapperRunId === context.fence.wrapperRunId && + current.wrapperConnectionId === context.fence.wrapperConnectionId + : current.plane === 'control' && + context.fence.plane === 'control' && + current.providerInstanceId === context.fence.providerInstanceId && + current.connectionId === context.fence.connectionId && + current.wrapperInstanceId === context.fence.wrapperInstanceId) + ); +} + +export async function resolveRuntimeProxyCredential(input: { + handle: string; + env: Pick; + grant: unknown; + authorization: { id: string; state: string; delegationExpiresAt: string } | null; + context: { + sessionId: string; + kiloSessionId: string; + userId: string; + orgId?: string; + fence: RuntimeProxyFence; + }; + now?: number; + token: string; + renew: () => Promise; +}): Promise<{ token: string; transportProofRequired: boolean } | null> { + const claims = await verifyRuntimeCredentialProxyHandle(input.env, input.handle); + const now = input.now ?? Date.now(); + const parsedGrant = runtimeProxyGrantSchema.safeParse(input.grant); + if ( + !claims || + !('sessionId' in claims) || + !input.authorization || + input.authorization.state !== 'active' || + Date.parse(input.authorization.delegationExpiresAt) <= now || + !parsedGrant.success || + parsedGrant.data.leaseExpiresAt > Date.parse(input.authorization.delegationExpiresAt) || + !matchesRuntimeProxyGrant(input.grant, claims, { + ...input.context, + authorizationId: input.authorization.id, + now, + }) + ) { + return null; + } + const decoded = jwt.decode(input.token); + const exp = + typeof decoded === 'object' && decoded !== null && typeof decoded.exp === 'number' + ? decoded.exp * 1000 + : 0; + const token = exp > now + 5 * 60_000 ? input.token : await input.renew(); + return { + token, + transportProofRequired: parsedGrant.data.mode === 'contained', + }; +} + +export function runtimeCredentialProxyBaseUrl(workerUrl: string): string | null { + try { + const url = new URL(workerUrl); + if (url.username || url.password || url.search || url.hash) return null; + return `${url.origin}${url.pathname.replace(/\/+$/, '')}/api/runtime-credential-proxy`; + } catch { + return null; + } +} + +/** + * The backwards-compatible facade is deliberately an origin (with an optional + * deployment prefix), never an API route. Older Kilo releases append their own + * `/api/...` paths while newer releases normalize configured bases first. + */ +export function runtimeCredentialProxyFacadeBaseUrl(workerUrl: string): string | null { + if (hasUnsafeFacadePathEncoding(workerUrl)) return null; + try { + const url = new URL(workerUrl); + const prefix = url.pathname.replace(/\/+$/, ''); + const segments = prefix.split('/').filter(Boolean); + if ( + url.protocol !== 'https:' || + url.port || + url.username || + url.password || + url.search || + url.hash || + prefix.includes('\\') || + prefix.includes('//') || + segments.some(segment => segment === 'api' || segment === '.' || segment === '..') + ) { + return null; + } + return `${url.origin}${prefix}`; + } catch { + return null; + } +} + +function hasUnsafeFacadePathEncoding(value: string): boolean { + let decoded = value; + for (let depth = 0; depth < 8; depth++) { + if ( + decoded.includes('\\') || + /%(?:2f|5c)/i.test(decoded) || + /(?:^|\/)(?:\.|%2e){1,2}(?=\/|[?#]|$)/i.test(decoded) || + [...decoded].some( + character => character.charCodeAt(0) <= 0x20 || character.charCodeAt(0) === 0x7f + ) + ) { + return true; + } + try { + const next = decodeURIComponent(decoded); + if (next === decoded) return false; + decoded = next; + } catch { + return true; + } + } + return true; +} + +export function runtimeCredentialProxyUpstream( + targets: { backendBaseUrl: string; providerBaseUrl: string; sessionIngestBaseUrl: string }, + route: RuntimeCredentialProxyRoute, + method: string, + pathname: string, + search: string, + kiloSessionId: string, + organizationId?: string, + contentType?: string | null, + bodyText?: string +): URL | null { + return resolveRuntimeCredentialProxyRoute({ + targets, + route, + method, + pathname: `/${pathname.replace(/^\/+/, '')}`, + search, + kiloSessionId, + organizationId, + contentType, + bodyText, + }); +} diff --git a/services/cloud-agent-next/src/sandbox-control/frames.ts b/services/cloud-agent-next/src/sandbox-control/frames.ts index 3118eb669e..6e83161f3c 100644 --- a/services/cloud-agent-next/src/sandbox-control/frames.ts +++ b/services/cloud-agent-next/src/sandbox-control/frames.ts @@ -15,6 +15,7 @@ import { sessionAbortPayloadSchema, sessionAttachPayloadSchema, sessionDetachPayloadSchema, + sessionRuntimeRetirePayloadSchema, sessionEventPayloadSchema, sessionGitSummaryPayloadSchema, sessionPreparingPayloadSchema, @@ -59,6 +60,7 @@ const REQUEST_PAYLOAD_SCHEMAS: Record = { 'session.sync': sessionSyncPayloadSchema, 'session.git.summary': sessionGitSummaryPayloadSchema, 'session.detach': sessionDetachPayloadSchema, + 'session.runtime.retire': sessionRuntimeRetirePayloadSchema, 'session.terminal.create': sessionTerminalCreatePayloadSchema, 'session.terminal.resize': sessionTerminalResizePayloadSchema, 'session.terminal.close': sessionTerminalClosePayloadSchema, diff --git a/services/cloud-agent-next/src/sandbox-control/lifecycle.test.ts b/services/cloud-agent-next/src/sandbox-control/lifecycle.test.ts index ef0e75b28b..e790b2b6df 100644 --- a/services/cloud-agent-next/src/sandbox-control/lifecycle.test.ts +++ b/services/cloud-agent-next/src/sandbox-control/lifecycle.test.ts @@ -1674,6 +1674,82 @@ describe('SandboxControl lifecycle boundaries', () => { } ); + it.each([ + [ + 'connection', + (identity: SandboxControlConnectionIdentity) => ({ + ...identity, + connectionId: crypto.randomUUID(), + }), + ], + [ + 'provider', + (identity: SandboxControlConnectionIdentity) => ({ + ...identity, + providerInstanceId: 'different-provider-instance', + }), + ], + ] as const)( + 'rejects an expected %s connection mismatch even when the wrapper instance is unchanged', + async (_field, changed) => { + const h = await harness(); + await h.create(); + const identity = await h.ready(); + + await expect( + h.control.request({ + operation: 'session.attach', + session: ROUTE, + payload: {}, + expectedWrapperInstanceId: identity.wrapperInstanceId, + expectedConnection: changed(identity), + }) + ).rejects.toThrow('Sandbox control connection changed'); + expect(h.sendRequest).not.toHaveBeenCalled(); + } + ); + + it('returns a runtime credential proxy fence only for the current ready routed allocation', async () => { + const h = await harness(); + const input = { + ownerId: OWNER, + sessionId: ROUTE.sessionId, + kiloSessionId: ROUTE.kiloSessionId, + directory: ROUTE.directory, + }; + await h.create(); + const identity = await h.ready(); + const physical = await h.control.getPhysicalRecord(); + expect(physical.createIntent).not.toBeNull(); + expect(await h.control.getRuntimeCredentialProxyFence(input)).toEqual({ + plane: 'control', + allocationId: physical.createIntent?.intentId, + providerInstanceId: identity.providerInstanceId, + connectionId: identity.connectionId, + wrapperInstanceId: identity.wrapperInstanceId, + }); + await expect( + h.control.getRuntimeCredentialProxyFence({ ...input, directory: '/workspace/other' }) + ).resolves.toBeNull(); + + await h.control.beginStop('idle'); + await expect(h.control.getRuntimeCredentialProxyFence(input)).resolves.toBeNull(); + await h.control.recordStopAttempt(); + await h.flush(); + await h.create(); + const replacement = await h.ready(); + expect(await h.control.getRuntimeCredentialProxyFence(input)).toMatchObject({ + providerInstanceId: replacement.providerInstanceId, + connectionId: replacement.connectionId, + wrapperInstanceId: replacement.wrapperInstanceId, + }); + vi.spyOn(h.socket, 'getConnectionIdentity').mockReturnValue({ + ...replacement, + connectionId: crypto.randomUUID(), + }); + await expect(h.control.getRuntimeCredentialProxyFence(input)).resolves.toBeNull(); + }); + it.each(['session.attach', 'session.prompt'] as const)( 'protects validated %s demand at the idle boundary without renewing heartbeat supervision', async operation => { diff --git a/services/cloud-agent-next/src/sandbox-control/session-credentials.test.ts b/services/cloud-agent-next/src/sandbox-control/session-credentials.test.ts index 111008469f..2f9e1aceff 100644 --- a/services/cloud-agent-next/src/sandbox-control/session-credentials.test.ts +++ b/services/cloud-agent-next/src/sandbox-control/session-credentials.test.ts @@ -1451,7 +1451,6 @@ describe('direct worktree Kilo token stability', () => { ['organization', { organizationId: INTEGRATION_ID }], ['role', { organizationRole: 'member' }], ['bot', { botId: 'bot-a' }], - ['audience', { aud: 'internal-service' }], ['unknown authorization', { futurePermission: 'restricted' }], ])('does not substitute the retained token after a change to %s', async (_name, claims) => { const first = await prepareDirect(env, data(token())); @@ -1461,6 +1460,14 @@ describe('direct worktree Kilo token stability', () => { expect(refreshed.payload.env?.KILOCODE_TOKEN).toBe(changed); }); + it('fails closed for unsupported audience-bearing direct credentials at this checkpoint', async () => { + const first = await prepareDirect(env, data(token())); + const changed = token(NOW + 1000, { aud: 'internal-service' }); + await expect(prepareDirect(env, data(changed), first.grant, NOW + 2000)).rejects.toThrow( + 'Invalid contained worktree credentials' + ); + }); + it.each([ ['expired', () => token(NOW - HOUR, { exp: NOW / 1000 })], ['future-issued', () => token(NOW + HOUR)], @@ -1816,6 +1823,63 @@ describe('credential resolution boundaries', () => { }); describe('native Vercel worktree policies', () => { + it('uses only Worker proxy targets for a modern runtime and never serializes its backing authority', async () => { + const { broker } = createBroker(); + const backingToken = jwt.sign( + { runtimeAuthorization: { id: '11111111-1111-4111-8111-111111111111' } }, + 'backing-secret' + ); + const env = { ...environment(broker), WORKER_URL: 'https://worker.example.test' }; + const { grant } = await prepare( + env, + vercelMetadata({ auth: { kiloSessionId: ROOT_ID, kilocodeToken: backingToken } }) + ); + + const staticPolicy = buildControlNetworkPolicy([grant]); + expect(JSON.stringify(staticPolicy)).not.toContain(backingToken); + expect(staticPolicy.allowedDomains).toContain('worker.example.test'); + expect(staticPolicy.injectionRules).toEqual( + expect.not.arrayContaining([ + expect.objectContaining({ + headers: expect.objectContaining({ + authorization: expect.stringContaining('backing-secret'), + }), + }), + ]) + ); + + const policy = buildControlNetworkPolicy([ + { + ...grant, + kilo: { + ...grant.kilo, + runtimeProxy: { + ...grant.kilo.runtimeProxy!, + members: [{ sessionId: SESSION_ID, kiloSessionId: ROOT_ID, handle: 'member-handle' }], + }, + }, + }, + ]); + const serialized = JSON.stringify(policy); + expect(serialized).toContain('member-handle'); + expect(serialized).not.toContain(backingToken); + expect(serialized).not.toContain('runtimeAuthorization'); + expect( + policy.injectionRules + .filter(rule => rule.headers.authorization === 'Bearer member-handle') + .every(rule => rule.domain === 'worker.example.test') + ).toBe(true); + expect( + findMatchingCredentialInjectionRule(policy.injectionRules, { + url: new URL( + 'https://worker.example.test/api/runtime-credential-proxy/provider/anything-else' + ), + method: 'POST', + headers: new Headers({ authorization: 'Bearer member-handle' }), + }) + ).toBeUndefined(); + }); + it('composes sibling worktree rules without losing registered roots or broadening aliases', async () => { const { broker } = createBroker(); const env = environment(broker); @@ -1864,6 +1928,61 @@ describe('native Vercel worktree policies', () => { }); }); + it('keeps distinct modern sibling runtime credentials isolated', async () => { + const { broker } = createBroker(); + const env = { ...environment(broker), WORKER_URL: 'https://worker.example.test' }; + const firstToken = jwt.sign( + { runtimeAuthorization: { id: '11111111-1111-4111-8111-111111111111' } }, + 'first-backing-secret' + ); + const secondToken = jwt.sign( + { runtimeAuthorization: { id: '22222222-2222-4222-8222-222222222222' } }, + 'second-backing-secret' + ); + const first = await prepare( + env, + vercelMetadata({ auth: { kiloSessionId: ROOT_ID, kilocodeToken: firstToken } }) + ); + const second = await prepare( + env, + vercelMetadata({ + identity: secondRoot().identity, + auth: { kiloSessionId: SECOND_ROOT_ID, kilocodeToken: secondToken }, + }), + first.grant + ); + expect(second.grant.kilo.token).not.toBe(firstToken); + expect(second.grant.members).toEqual([ + { sessionId: SESSION_ID, kiloSessionId: ROOT_ID }, + { sessionId: SECOND_SESSION_ID, kiloSessionId: SECOND_ROOT_ID }, + ]); + expect(JSON.stringify(buildControlNetworkPolicy([second.grant]))).not.toContain(secondToken); + }); + + it('rejects a persisted singleton runtime proxy handle', async () => { + const { broker } = createBroker(); + const backingToken = jwt.sign( + { runtimeAuthorization: { id: '11111111-1111-4111-8111-111111111111' } }, + 'backing-secret' + ); + const { grant } = await prepare( + { ...environment(broker), WORKER_URL: 'https://worker.example.test' }, + vercelMetadata({ auth: { kiloSessionId: ROOT_ID, kilocodeToken: backingToken } }) + ); + const runtimeProxy = grant.kilo.runtimeProxy; + if (!runtimeProxy) throw new Error('Expected runtime proxy'); + + expect( + sessionCredentialGrantSchema.safeParse({ + ...grant, + kilo: { + ...grant.kilo, + runtimeProxy: { ...runtimeProxy, handle: 'old-singleton-handle' }, + }, + }).success + ).toBe(false); + }); + it('refreshes native GitHub injection on trusted prepare while preserving the integration pin and alias', async () => { const { broker } = createBroker(); const env = environment(broker); diff --git a/services/cloud-agent-next/src/sandbox-control/session-credentials.ts b/services/cloud-agent-next/src/sandbox-control/session-credentials.ts index 90e547bbcd..2daea97e3d 100644 --- a/services/cloud-agent-next/src/sandbox-control/session-credentials.ts +++ b/services/cloud-agent-next/src/sandbox-control/session-credentials.ts @@ -23,6 +23,8 @@ import { resolveManagedBitbucketToken, } from '../services/git-token-service-client.js'; import { readProfileBundle } from '../session-profile.js'; +import { hasModernRuntimeAuthorization } from '../session/runtime-authorization-persistence.js'; +import { runtimeCredentialProxyFacadeBaseUrl } from '../runtime-credential-proxy.js'; import type { SessionAttachPayload } from '../shared/sandbox-control-protocol.js'; import { parseCanonicalBitbucketCloneUrl, sessionIdSchema, type Env } from '../types.js'; import { createControlPlaneCredential, parseControlPlaneCredential } from './managed-credential.js'; @@ -71,6 +73,19 @@ const targetsSchema = z sessionIngestBaseUrl: z.string().url(), }) .strict(); +const runtimeProxySchema = z + .object({ + targets: targetsSchema, + /** Exact root capabilities substituted only by the Vercel policy. */ + members: z + .array( + memberSchema.extend({ + handle: tokenSchema.max(4096), + }) + ) + .default([]), + }) + .strict(); const capabilitySchema = z .object({ credential: tokenSchema.regex(/^(?:kka1|kgh2|kgl2|kbb1)\./), @@ -175,6 +190,7 @@ export const sessionCredentialGrantSchema = z token: realTokenSchema, tokenSelectedAt: timestampSchema.optional(), targets: targetsSchema, + runtimeProxy: runtimeProxySchema.optional(), capabilities: z.record(z.string(), capabilitySchema), }) .strict(), @@ -198,7 +214,11 @@ export const sessionCredentialGrantSchema = z reject(); } if (grant.containmentEnabled === false) { - if (grant.kilo.alias !== undefined || Object.keys(grant.kilo.capabilities).length > 0) + if ( + grant.kilo.alias !== undefined || + grant.kilo.runtimeProxy !== undefined || + Object.keys(grant.kilo.capabilities).length > 0 + ) reject(); } else { const kiloAlias = grant.kilo.alias && parseControlPlaneCredential(grant.kilo.alias); @@ -260,6 +280,25 @@ export const sessionCredentialGrantSchema = z ) { reject(); } + if (grant.kilo.runtimeProxy) { + const targets = deriveRuntimeProxyTargets(grant.kilo.runtimeProxy.targets); + const members = grant.kilo.runtimeProxy.members; + if ( + !targets || + new Set(members.map(member => member.sessionId)).size !== members.length || + new Set(members.map(member => member.kiloSessionId)).size !== members.length || + members.some( + member => + !grant.members.some( + expected => + expected.sessionId === member.sessionId && + expected.kiloSessionId === member.kiloSessionId + ) + ) + ) { + reject(); + } + } } else { const prefixes = { github: 'kgh2.', gitlab: 'kgl2.', bitbucket: 'kbb1.' }; if ( @@ -294,7 +333,7 @@ export function isContainedSessionCredentialGrant( } type CredentialEnv = Parameters[0] & - Partial> & + Partial> & KiloTargetEnv; type PreparedSessionAttachPayload = SessionAttachPayload & { @@ -329,6 +368,48 @@ function safeUrl(value: string): URL | null { } } +function deriveRuntimeProxyTargets(targets: z.infer): { facade: URL } | null { + const parsed = [ + targets.backendBaseUrl, + targets.providerBaseUrl, + targets.sessionIngestBaseUrl, + ].map(safeUrl); + const [backend, provider, ingest] = parsed; + if ( + !backend || + !provider || + !ingest || + backend.protocol !== 'https:' || + backend.port !== '' || + provider.protocol !== 'https:' || + provider.port !== '' || + ingest.protocol !== 'https:' || + ingest.port !== '' || + backend.toString().replace(/\/+$/, '') !== provider.toString().replace(/\/+$/, '') || + backend.toString().replace(/\/+$/, '') !== ingest.toString().replace(/\/+$/, '') || + backend.search || + provider.search || + ingest.search || + runtimeCredentialProxyFacadeBaseUrl(backend.toString()) !== + backend.toString().replace(/\/+$/, '') + ) { + return null; + } + return { facade: backend }; +} + +function runtimeProxyTargets(env: CredentialEnv): z.infer | null { + if (!env.WORKER_URL) return null; + const base = runtimeCredentialProxyFacadeBaseUrl(env.WORKER_URL); + if (!base) return null; + const targets = { + backendBaseUrl: base, + providerBaseUrl: base, + sessionIngestBaseUrl: base, + }; + return deriveRuntimeProxyTargets(targets) ? targets : null; +} + function canonicalRepositoryUrl(value: string, managed = false): string { const url = safeUrl(value); if (!url || url.protocol !== 'https:' || url.search) invalidCredentials(); @@ -575,6 +656,14 @@ async function selectDirectKiloToken( existing: SessionCredentialGrant | undefined, now: number ): Promise { + const decoded = jwt.decode(token); + if (decoded !== null && typeof decoded === 'object') { + if ('runtimeAdmission' in decoded) invalidCredentials(); + if ('runtimeAuthorization' in decoded) invalidCredentials(); + if ('aud' in decoded || 'tokenPurpose' in decoded || 'credentialExchange' in decoded) { + invalidCredentials(); + } + } if ( !existing || existing.kilo.token === token || @@ -845,6 +934,18 @@ export async function prepareSessionCredentials(input: { ) { invalidCredentials(); } + const modernRuntimeProxy = + provider === 'vercel' && containmentEnabled && hasModernRuntimeAuthorization(metadata) + ? runtimeProxyTargets(env) + : null; + if ( + provider === 'vercel' && + containmentEnabled && + hasModernRuntimeAuthorization(metadata) && + !modernRuntimeProxy + ) { + invalidCredentials(); + } const kiloToken = containmentEnabled ? token.data : await selectDirectKiloToken(env, token.data, existing, now); @@ -876,6 +977,14 @@ export async function prepareSessionCredentials(input: { } : {}), targets: targets.targets, + ...(modernRuntimeProxy + ? { + runtimeProxy: { + targets: modernRuntimeProxy, + members: existing?.kilo.runtimeProxy?.members ?? [], + }, + } + : {}), capabilities: existing?.kilo.token === kiloToken ? existing.kilo.capabilities : {}, }, ...(existing?.scm ? { scm: existing.scm } : {}), @@ -951,6 +1060,16 @@ export function removeSessionCredentialMembership( capabilities: Object.fromEntries( Object.entries(grant.kilo.capabilities).filter(([id]) => id !== sessionId) ), + ...(grant.kilo.runtimeProxy + ? { + runtimeProxy: { + ...grant.kilo.runtimeProxy, + members: grant.kilo.runtimeProxy.members.filter( + member => member.sessionId !== sessionId + ), + }, + } + : {}), }, }, ]; @@ -976,6 +1095,7 @@ export function buildControlNetworkPolicy( targets: grant.kilo.targets, rootSessionIds: grant.members.map(member => member.kiloSessionId), organizationId: grant.orgId, + ...(grant.kilo.runtimeProxy ? { runtimeProxy: grant.kilo.runtimeProxy } : {}), }, ...(grant.repository?.type === 'github' && grant.scm?.nativeToken ? { @@ -996,7 +1116,10 @@ export function buildControlNetworkPolicy( const injectionRules = policies.flatMap(policy => policy.injectionRules); return { mode: 'custom', - allowedDomains: [...new Set(injectionRules.map(rule => rule.domain)), '*'], + allowedDomains: [ + ...new Set(policies.flatMap(policy => policy.allowedDomains)), + ...(policies.length === 0 ? ['*'] : []), + ], injectionRules, }; } diff --git a/services/cloud-agent-next/src/sandbox-control/socket.test.ts b/services/cloud-agent-next/src/sandbox-control/socket.test.ts index 0ca4ea5d4d..145ef1c35a 100644 --- a/services/cloud-agent-next/src/sandbox-control/socket.test.ts +++ b/services/cloud-agent-next/src/sandbox-control/socket.test.ts @@ -63,7 +63,11 @@ function helloFrame( providerInstanceId: string, wrapperInstanceId?: string, requestId = 'req_hello', - capabilities?: { nativeRuntimeRetirement?: boolean; workingBranches?: boolean } + capabilities?: { + nativeRuntimeRetirement?: boolean; + runtimeIsolation?: true; + workingBranches?: boolean; + } ): string { return JSON.stringify({ type: 'request', @@ -100,6 +104,20 @@ describe('sandbox control socket handler', () => { expect(JSON.stringify(parsed)).not.toContain('private'); } }); + it('retains the optional runtime isolation capability without requiring it from old wrappers', async () => { + const incoming = createFakeWebSocket(); + const handler = createSandboxControlSocketHandler(createFakeState([incoming]), 'sbx_test'); + + await handler.handleMessage( + asWs(incoming), + helloFrame('inst_1', WRAPPER_INSTANCE_ID, 'req_isolation', { runtimeIsolation: true }) + ); + + expect(handler.getConnectionIdentity()).toMatchObject({ + providerInstanceId: 'inst_1', + runtimeIsolation: true, + }); + }); it.each([ ['2.4.0', '2.4.0'], [undefined, null], diff --git a/services/cloud-agent-next/src/sandbox-control/socket.ts b/services/cloud-agent-next/src/sandbox-control/socket.ts index 1ac426e9b5..fded550897 100644 --- a/services/cloud-agent-next/src/sandbox-control/socket.ts +++ b/services/cloud-agent-next/src/sandbox-control/socket.ts @@ -62,6 +62,7 @@ export type SandboxControlOutboundRequest = { authorization?: SessionOperationAuthorization; timeoutMs?: number; expectedWrapperInstanceId?: string; + expectedConnection?: SandboxControlConnectionIdentity; deadlineAt?: number; }; @@ -70,6 +71,8 @@ export type SandboxControlConnectionIdentity = { providerInstanceId: string; wrapperInstanceId?: string; recoveryCapable?: boolean; + runtimeIsolation?: true; + runtimeRecovery?: true; }; export type SandboxControlEventResult = { applied: boolean; retryable?: boolean }; @@ -226,6 +229,8 @@ function readConnectionIdentity( providerInstanceId: attachment.providerInstanceId, ...(attachment.recoveryCapable ? { recoveryCapable: true } : {}), ...(attachment.wrapperInstanceId ? { wrapperInstanceId: attachment.wrapperInstanceId } : {}), + ...(attachment.runtimeIsolation ? { runtimeIsolation: true } : {}), + ...(attachment.runtimeRecovery ? { runtimeRecovery: true } : {}), }; } @@ -541,6 +546,8 @@ export function createSandboxControlSocketHandler( providerInstanceId: payload.providerInstanceId, ...(payload.wrapperInstanceId ? { wrapperInstanceId: payload.wrapperInstanceId } : {}), ...(payload.capabilities?.connectionRecovery === true ? { recoveryCapable: true } : {}), + ...(payload.capabilities?.runtimeIsolation === true ? { runtimeIsolation: true } : {}), + ...(payload.capabilities?.runtimeRecovery === true ? { runtimeRecovery: true } : {}), }; const completed: SandboxControlSocketAttachment = { handshakeComplete: true, @@ -552,6 +559,8 @@ export function createSandboxControlSocketHandler( providerInstanceId: identity.providerInstanceId, ...(payload.capabilities ? { capabilities: payload.capabilities } : {}), ...(identity.wrapperInstanceId ? { wrapperInstanceId: identity.wrapperInstanceId } : {}), + ...(identity.runtimeIsolation ? { runtimeIsolation: true } : {}), + ...(identity.runtimeRecovery ? { runtimeRecovery: true } : {}), }; const superseded: WebSocket[] = []; let replaced = false; diff --git a/services/cloud-agent-next/src/sandbox-control/vercel-network-policy.test.ts b/services/cloud-agent-next/src/sandbox-control/vercel-network-policy.test.ts index 7013e35164..6f21bde4bd 100644 --- a/services/cloud-agent-next/src/sandbox-control/vercel-network-policy.test.ts +++ b/services/cloud-agent-next/src/sandbox-control/vercel-network-policy.test.ts @@ -101,6 +101,149 @@ describe('buildVercelCredentialNetworkPolicy', () => { }); }); + it('authenticates only packaged CLI provider paths through the runtime proxy', () => { + const policy = buildVercelCredentialNetworkPolicy({ + kilo: kiloInput({ + runtimeProxy: { + members: [ + { + sessionId: 'workspace_a', + kiloSessionId: ROOT_SESSION_ID, + handle: 'runtime-proxy-handle', + }, + ], + targets: { + backendBaseUrl: 'https://worker.example.com', + providerBaseUrl: 'https://worker.example.com', + sessionIngestBaseUrl: 'https://worker.example.com', + }, + }, + }), + }); + const authorization = 'Bearer runtime-proxy-handle'; + const base = 'https://worker.example.com/api/openrouter'; + + for (const [path, method] of [ + ['/models', 'GET'], + ['/models/validate', 'POST'], + ['/chat/completions', 'POST'], + ['/messages', 'POST'], + ['/responses', 'POST'], + ['/embeddings', 'POST'], + ]) { + expect(effectiveAuthorization(policy, { url: `${base}${path}`, authorization, method })).toBe( + authorization + ); + } + for (const path of ['/v1/chat/completions', '/v1/responses']) { + expect( + effectiveAuthorization(policy, { + url: `https://worker.example.com/api/gateway${path}`, + authorization, + method: 'POST', + }) + ).toBe(authorization); + expect( + effectiveAuthorization(policy, { + url: `${base}${path}`, + authorization, + method: 'POST', + }) + ).toBeUndefined(); + expect( + effectiveAuthorization(policy, { + url: `https://worker.example.com/api/gateway${path}`, + authorization, + method: 'GET', + }) + ).toBeUndefined(); + } + + for (const [path, method] of [ + ['/models', 'POST'], + ['/chat/completions', 'GET'], + ['/api/organizations/trusted-org/models', 'GET'], + ['/unknown', 'POST'], + ['/../chat/completions', 'POST'], + ['/models', 'DELETE'], + ]) { + expect( + effectiveAuthorization(policy, { url: `${base}${path}`, authorization, method }) + ).toBeUndefined(); + } + expect( + effectiveAuthorization(policy, { + url: 'https://worker.example.com/chat/completions', + authorization, + method: 'POST', + }) + ).toBeUndefined(); + }); + + it('uses the owning member handle for every proxy route', () => { + const policy = buildVercelCredentialNetworkPolicy({ + kilo: kiloInput({ + rootSessionIds: [ROOT_SESSION_ID, OTHER_SESSION_ID], + runtimeProxy: { + members: [ + { sessionId: 'workspace_a', kiloSessionId: ROOT_SESSION_ID, handle: 'member-a' }, + { sessionId: 'workspace_b', kiloSessionId: OTHER_SESSION_ID, handle: 'member-b' }, + ], + targets: { + backendBaseUrl: 'https://worker.example.com', + providerBaseUrl: 'https://worker.example.com', + sessionIngestBaseUrl: 'https://worker.example.com', + }, + }, + }), + }); + const authorization = 'Bearer member-a'; + expect( + effectiveAuthorization(policy, { + url: 'https://worker.example.com/api/openrouter/chat/completions', + method: 'POST', + authorization, + }) + ).toBe('Bearer member-a'); + expect( + effectiveAuthorization(policy, { + url: `https://worker.example.com/api/session/${ROOT_SESSION_ID}/ingest`, + method: 'POST', + authorization, + }) + ).toBe('Bearer member-a'); + expect( + effectiveAuthorization(policy, { + url: `https://worker.example.com/api/session/${OTHER_SESSION_ID}/title`, + method: 'POST', + authorization: 'Bearer member-b', + }) + ).toBe('Bearer member-b'); + expect( + effectiveAuthorization(policy, { + url: 'https://worker.example.com/api/session/unrelated/ingest', + method: 'POST', + authorization, + }) + ).toBeUndefined(); + // The policy can only key `/api/session` on the injected handle; the + // facade validates its JSON body against that handle's root identity. + expect( + effectiveAuthorization(policy, { + url: 'https://worker.example.com/api/session', + method: 'POST', + authorization: 'Bearer member-b', + }) + ).toBe('Bearer member-b'); + expect( + effectiveAuthorization(policy, { + url: `https://worker.example.com/api/session/${ROOT_SESSION_ID}/export`, + method: 'GET', + authorization: 'Bearer member-b', + }) + ).toBeUndefined(); + }); + it.each([ ['trusted organization', 'trusted-org', 'trusted-org'], ['personal account', undefined, ''], diff --git a/services/cloud-agent-next/src/sandbox-control/vercel-network-policy.ts b/services/cloud-agent-next/src/sandbox-control/vercel-network-policy.ts index 60a7551324..5b829b7f32 100644 --- a/services/cloud-agent-next/src/sandbox-control/vercel-network-policy.ts +++ b/services/cloud-agent-next/src/sandbox-control/vercel-network-policy.ts @@ -4,6 +4,7 @@ import type { VercelSandboxNetworkPolicy, } from '../agent-sandbox/vercel/vercel-sandbox-rest-client.js'; import { deriveKiloSandboxTargets, type KiloSandboxTargets } from '../kilo/kilo-targets.js'; +import { runtimeCredentialProxyFacadeBaseUrl } from '../runtime-credential-proxy.js'; export type VercelCredentialPolicyInput = { kilo?: { @@ -12,6 +13,10 @@ export type VercelCredentialPolicyInput = { targets: KiloSandboxTargets; rootSessionIds: string[]; organizationId?: string; + runtimeProxy?: { + members?: Array<{ sessionId: string; kiloSessionId: string; handle: string }>; + targets: { backendBaseUrl: string; providerBaseUrl: string; sessionIngestBaseUrl: string }; + }; }; github?: { token: string; @@ -200,6 +205,109 @@ export function buildKiloCredentialInjectionRules( return rules; } +function runtimeProxyInjectionRules( + kilo: NonNullable, + organizationId: string | undefined +): VercelSandboxInjectionRule[] { + const input = kilo.runtimeProxy; + if (!input) return []; + const targets = [ + new URL(input.targets.backendBaseUrl), + new URL(input.targets.providerBaseUrl), + new URL(input.targets.sessionIngestBaseUrl), + ]; + if ( + targets.some( + target => + target.protocol !== 'https:' || + target.port !== '' || + target.username || + target.password || + target.search || + target.hash + ) + ) { + invalidPolicy(); + } + const [backend, provider, ingest] = targets; + if ( + backend.toString().replace(/\/+$/, '') !== provider.toString().replace(/\/+$/, '') || + backend.toString().replace(/\/+$/, '') !== ingest.toString().replace(/\/+$/, '') || + runtimeCredentialProxyFacadeBaseUrl(backend.toString()) !== + backend.toString().replace(/\/+$/, '') + ) { + invalidPolicy(); + } + const rules: VercelSandboxInjectionRule[] = []; + const add = (target: URL, path: string, methods: string[], authorization: string) => + rules.push( + createCredentialRule({ + target, + path: { exact: path }, + methods, + expectedAuthorization: authorization, + injectedAuthorization: authorization, + }) + ); + for (const member of input.members ?? []) { + const authorization = `Bearer ${member.handle}`; + for (const path of [ + '/api/user', + '/api/profile', + '/api/profile/balance', + '/api/defaults', + '/api/users/notifications', + ]) + add(backend, `${basePath(backend)}${path}`, ['GET'], authorization); + if (organizationId !== undefined) { + for (const path of ['models', 'defaults', 'modes']) { + add( + backend, + `${basePath(backend)}/api/organizations/${organizationId}/${path}`, + ['GET'], + authorization + ); + } + add( + backend, + `${basePath(backend)}/api/organizations/${organizationId}/models/validate`, + ['POST'], + authorization + ); + } + for (const [path, methods] of [ + ['/models', ['GET']], + ['/models/validate', ['POST']], + ['/chat/completions', ['POST']], + ['/messages', ['POST']], + ['/responses', ['POST']], + ['/embeddings', ['POST']], + ] as const) { + for (const prefix of ['/api/openrouter', '/api/gateway']) { + add(provider, `${basePath(provider)}${prefix}${path}`, [...methods], authorization); + } + } + for (const path of ['/v1/chat/completions', '/v1/responses']) { + add(provider, `${basePath(provider)}/api/gateway${path}`, ['POST'], authorization); + } + add(ingest, `${basePath(ingest)}/api/session`, ['POST'], authorization); + const sessionId = member.kiloSessionId; + for (const [suffix, methods] of [ + ['export', ['GET']], + ['ingest', ['POST']], + ['title', ['POST']], + ] as const) { + add( + ingest, + `${basePath(ingest)}/api/session/${sessionId}/${suffix}`, + [...methods], + authorization + ); + } + } + return rules; +} + function githubInjectionRules( input: NonNullable ): VercelSandboxInjectionRule[] { @@ -299,13 +407,25 @@ export function buildVercelCredentialNetworkPolicy( } const injectionRules = [ - ...(input.kilo === undefined ? [] : buildKiloCredentialInjectionRules(input.kilo)), + ...(input.kilo === undefined + ? [] + : input.kilo.runtimeProxy + ? runtimeProxyInjectionRules(input.kilo, input.kilo.organizationId) + : buildKiloCredentialInjectionRules(input.kilo)), ...(input.github === undefined ? [] : githubInjectionRules(input.github)), ]; return { mode: 'custom', - allowedDomains: [...new Set(injectionRules.map(rule => rule.domain)), '*'], + allowedDomains: [ + ...new Set([ + ...injectionRules.map(rule => rule.domain), + ...(input.kilo?.runtimeProxy + ? [new URL(input.kilo.runtimeProxy.targets.backendBaseUrl).hostname] + : []), + ]), + '*', + ], injectionRules, }; } diff --git a/services/cloud-agent-next/src/sandbox-session/SandboxSession.ts b/services/cloud-agent-next/src/sandbox-session/SandboxSession.ts index 8d46dc1c5a..5de62d0afa 100644 --- a/services/cloud-agent-next/src/sandbox-session/SandboxSession.ts +++ b/services/cloud-agent-next/src/sandbox-session/SandboxSession.ts @@ -6,6 +6,23 @@ import type { import { generateBranchSlug } from '@kilocode/worker-utils/deployment-slug'; import { TRPCError } from '@trpc/server'; import { withTimeout } from '@kilocode/worker-utils'; +import { + renewRuntimeAuthorization, + unsealRuntimeAuthorization, +} from '@kilocode/worker-utils/runtime-authorization'; +import type { RuntimeAuthorization } from '@kilocode/worker-utils/runtime-authorization-contract'; +import { RuntimeAuthorizationSchema } from '@kilocode/worker-utils/runtime-authorization-contract'; +import { resolveSecret } from '../auth.js'; +import { + issuePersistedRuntimeProxyGrant, + resolvePersistedRuntimeProxyCredential, +} from '../runtime-credential-proxy-rpc.js'; +import { + runtimeCredentialProxyFacadeBaseUrl, + runtimeProxyGrantSchema, + RUNTIME_PROXY_GRANT_KEY, + verifyRuntimeCredentialProxyHandle, +} from '../runtime-credential-proxy.js'; import { z } from 'zod'; import { diagnosticSyncStatus } from '../shared/control-diagnostics.js'; import { @@ -85,6 +102,15 @@ import { sandboxControlRpc } from './control-rpc.js'; import { getSandboxControlStub } from '../sandbox-control/stub.js'; import { DEADLINE_MS } from '../sandbox-control/deadlines.js'; import { createMessageId } from '../session/message-id.js'; +import { + getRuntimeAuthorizationStatus, + getRuntimeAuthorizationRecoveryState, + hasModernRuntimeAuthorization, + renewStoredRuntimeAuthorization, + RUNTIME_AUTHORIZATION_RECOVERY_KEY, + RUNTIME_AUTHORIZATION_KEY, + runtimeAuthorizationRecoveryLockSchema, +} from '../session/runtime-authorization-persistence.js'; import { validateControlSessionOptions } from './attach-payload.js'; import { pendingInputProjection } from './session-input-projection.js'; import { @@ -120,12 +146,14 @@ import { sessionOperationExpiresAt, sessionOperationResultHash, sessionPromptResultSchema, + sessionRuntimeRetireResultSchema, sessionSyncResultSchema, sessionPermissionResolveResultSchema, sessionQuestionResolveResultSchema, sessionAbortResultSchema, sameSessionOperation, wrapperInstanceIdSchema, + type SessionAttachPayload, type SessionOperationAck, type SessionOperationAuthorization, type SessionRequestIdentity, @@ -262,6 +290,7 @@ type ControlEventEvaluationRequest = type SandboxSessionRegistrationInput = { identity: SessionMetadata['identity']; auth: SessionMetadata['auth']; + runtimeAuthorizationSeal?: string; agent: SessionMetadata['agent']; repository?: SessionMetadata['repository']; workspace?: SessionMetadata['workspace']; @@ -1013,11 +1042,299 @@ export class SandboxSession extends DurableObject { } async getMetadata(): Promise { - return this.deletedWorktreeId || this.terminalLifecycle.isDeleted() + return this.deletedWorktreeId || this.terminalLifecycle.isBlocked() ? null : this.terminalLifecycle.getStoredMetadata(); } + async getRuntimeToken(): Promise { + const metadata = await this.getMetadata(); + const secret = await resolveSecret(this.env.NEXTAUTH_SECRET); + if (!secret) throw new Error('NEXTAUTH_SECRET is not configured on the worker'); + return renewStoredRuntimeAuthorization({ + metadata, + getAuthorization: async () => this.ctx.storage.kv.get(RUNTIME_AUTHORIZATION_KEY), + putAuthorization: async authorization => { + this.ctx.storage.kv.put(RUNTIME_AUTHORIZATION_KEY, authorization); + }, + getMetadata: () => this.getMetadata(), + putMetadata: async updated => { + this.ctx.storage.kv.put(METADATA_KEY, updated); + }, + renew: authorization => + renewRuntimeAuthorization({ + authorization, + secret, + connectionString: this.env.HYPERDRIVE.connectionString, + }), + }); + } + + async getRuntimeAuthorizationStatus(): Promise<'legacy' | 'active' | 'revoked'> { + return getRuntimeAuthorizationStatus({ + metadata: await this.getMetadata(), + getAuthorization: async () => this.ctx.storage.kv.get(RUNTIME_AUTHORIZATION_KEY), + }); + } + + async getRuntimeAuthorizationRecoveryState(): Promise<{ + state: 'legacy' | 'revoked' | 'active' | 'expired'; + id?: string; + recoveryId?: string; + }> { + const state = await getRuntimeAuthorizationRecoveryState({ + metadata: await this.getMetadata(), + getAuthorization: async () => this.ctx.storage.kv.get(RUNTIME_AUTHORIZATION_KEY), + }); + const lock = runtimeAuthorizationRecoveryLockSchema.safeParse( + this.ctx.storage.kv.get(RUNTIME_AUTHORIZATION_RECOVERY_KEY) + ); + return state.state === 'expired' && lock.success && lock.data.expectedOldId === state.id + ? { ...state, recoveryId: lock.data.recoveryId } + : state; + } + + async isRuntimeAuthorizationRecoveryInProgress(): Promise { + return this.ctx.storage.kv.get(RUNTIME_AUTHORIZATION_RECOVERY_KEY) !== undefined; + } + + async recoverExpiredRuntimeAuthorization(input: { + ownerId: string; + expectedOldId: string; + recoveryId: string; + runtimeAuthorizationSeal: string; + runtimeToken: string; + }): Promise<{ status: 'recovered' | 'not-needed' | 'denied' | 'busy' | 'retry' }> { + const metadata = await this.getMetadata(); + if (!metadata || metadata.identity.userId !== input.ownerId) return { status: 'denied' }; + const secret = await resolveSecret(this.env.NEXTAUTH_SECRET); + if (!secret) return { status: 'denied' }; + let fresh: RuntimeAuthorization; + try { + fresh = await unsealRuntimeAuthorization(input.runtimeAuthorizationSeal, secret, { + resourceKind: 'cloud-agent-next', + resourceId: metadata.identity.sessionId, + userId: metadata.identity.userId, + organizationId: metadata.identity.orgId, + }); + } catch { + return { status: 'denied' }; + } + if (fresh.state !== 'active') return { status: 'denied' }; + if (!metadata.auth.kiloSessionId) return { status: 'denied' }; + const current = await this.getRuntimeAuthorizationRecoveryState(); + if (current.state === 'legacy' || current.state === 'active') return { status: 'not-needed' }; + if (current.state !== 'expired' || current.id !== input.expectedOldId) + return { status: 'denied' }; + if ( + this.loadMessages().some( + message => message.state === 'accepted' || message.state === 'queued' + ) + ) { + return { status: 'busy' }; + } + const held = runtimeAuthorizationRecoveryLockSchema.safeParse( + this.ctx.storage.kv.get(RUNTIME_AUTHORIZATION_RECOVERY_KEY) + ); + if ( + held.success && + (held.data.recoveryId !== input.recoveryId || held.data.expectedOldId !== input.expectedOldId) + ) { + return { status: 'retry' }; + } + this.ctx.storage.kv.put(RUNTIME_AUTHORIZATION_RECOVERY_KEY, { + expectedOldId: input.expectedOldId, + recoveryId: input.recoveryId, + }); + try { + const sandboxId = metadata.workspace?.sandboxId; + if (sandboxId) { + const control = sandboxControlRpc(this.env, sandboxId); + const status = await control.getStatus(); + if (status.physical === 'running') { + if ( + status.connection !== 'ready' || + !status.wrapperInstanceId || + status.runtimeRecovery !== true + ) { + return { status: 'busy' }; + } + const attached = this.terminalLifecycle.getAttachedWrapperInstanceId(); + if (attached && attached !== status.wrapperInstanceId) return { status: 'retry' }; + const retired = sessionRuntimeRetireResultSchema.parse( + controlRequestResult( + await control.request({ + operation: 'session.runtime.retire', + session: { + sessionId: metadata.identity.sessionId, + kiloSessionId: metadata.auth.kiloSessionId ?? '', + directory: this.directory(metadata), + }, + expectedWrapperInstanceId: status.wrapperInstanceId, + payload: { recoveryId: input.recoveryId }, + }) + ) + ); + if (retired.recoveryId !== input.recoveryId || !retired.retired) + return { status: 'retry' }; + if ( + attached && + !this.terminalLifecycle.clearAttachedWrapperAfterRecovery(status.wrapperInstanceId) + ) { + return { status: 'busy' }; + } + } else if (status.physical !== 'stopped') { + return { status: 'retry' }; + } + } + const latest = await this.getRuntimeAuthorizationRecoveryState(); + if (latest.state !== 'expired' || latest.id !== input.expectedOldId) + return { status: 'retry' }; + this.ctx.storage.transactionSync(() => { + const stored = RuntimeAuthorizationSchema.safeParse( + this.ctx.storage.kv.get(RUNTIME_AUTHORIZATION_KEY) + ); + if ( + !stored.success || + stored.data.id !== input.expectedOldId || + stored.data.state !== 'active' || + Date.parse(stored.data.delegationExpiresAt) > Date.now() + ) { + throw new Error('runtime_authorization_recovery_cas_failed'); + } + const currentMetadata = this.terminalLifecycle.getStoredMetadata(); + if ( + !currentMetadata || + currentMetadata.identity.sessionId !== metadata.identity.sessionId || + currentMetadata.identity.userId !== metadata.identity.userId || + currentMetadata.identity.orgId !== metadata.identity.orgId + ) { + throw new Error('runtime_authorization_recovery_cas_failed'); + } + this.ctx.storage.kv.put(RUNTIME_AUTHORIZATION_KEY, fresh); + this.ctx.storage.kv.put( + METADATA_KEY, + serializeSessionMetadata({ + ...currentMetadata, + auth: { ...currentMetadata.auth, kilocodeToken: input.runtimeToken }, + }) + ); + this.ctx.storage.kv.delete(RUNTIME_PROXY_GRANT_KEY); + this.ctx.storage.kv.delete(RUNTIME_AUTHORIZATION_RECOVERY_KEY); + }); + return { status: 'recovered' }; + } catch { + return { status: 'retry' }; + } + } + + async issueRuntimeCredentialProxyGrant(_fence: { + wrapperRunId: string; + wrapperGeneration: number; + wrapperConnectionId: string; + }): Promise { + const metadata = await this.getMetadata(); + const kiloSessionId = metadata?.auth.kiloSessionId; + const sandboxId = metadata?.workspace?.sandboxId; + if (!metadata || !kiloSessionId || !sandboxId) return null; + const control = sandboxControlRpc(this.env, sandboxId); + const readFence = () => + control.getRuntimeCredentialProxyFence({ + ownerId: metadata.identity.userId, + sessionId: metadata.identity.sessionId, + kiloSessionId, + directory: this.directory(metadata), + }); + const fence = await readFence(); + if (!fence) return null; + const token = await this.getRuntimeToken(); + const [latestMetadata, storedAuthorization, latestFence] = await Promise.all([ + this.getMetadata(), + this.ctx.storage.kv.get(RUNTIME_AUTHORIZATION_KEY), + readFence(), + ]); + const authorization = RuntimeAuthorizationSchema.safeParse(storedAuthorization); + if ( + !latestFence || + latestFence.allocationId !== fence.allocationId || + latestFence.providerInstanceId !== fence.providerInstanceId || + latestFence.connectionId !== fence.connectionId || + latestFence.wrapperInstanceId !== fence.wrapperInstanceId + ) { + return null; + } + return issuePersistedRuntimeProxyGrant({ + env: this.env, + storage: this.ctx.storage, + metadata: latestMetadata, + authorization: authorization.success ? authorization.data : null, + fence: latestFence, + token, + mode: 'contained', + }); + } + + async resolveRuntimeCredentialProxyGrant(_handle: string): Promise<{ + token: string; + organizationId?: string; + runtimeAuthorization: { userId: string; authorizationId: string; resourceId: string }; + } | null> { + return resolvePersistedRuntimeProxyCredential({ + env: this.env, + storage: this.ctx.storage, + handle: _handle, + metadata: () => this.getMetadata(), + authorization: async () => { + const parsed = RuntimeAuthorizationSchema.safeParse( + await this.ctx.storage.kv.get(RUNTIME_AUTHORIZATION_KEY) + ); + return parsed.success ? parsed.data : null; + }, + fence: async () => { + const metadata = await this.getMetadata(); + const kiloSessionId = metadata?.auth.kiloSessionId; + const sandboxId = metadata?.workspace?.sandboxId; + if (!metadata || !kiloSessionId || !sandboxId) return null; + return sandboxControlRpc(this.env, sandboxId).getRuntimeCredentialProxyFence({ + ownerId: metadata.identity.userId, + sessionId: metadata.identity.sessionId, + kiloSessionId, + directory: this.directory(metadata), + }); + }, + token: () => this.getRuntimeToken(), + }); + } + + async reauthorizeRuntimeAuthorization(input: { + ownerId: string; + expectedOldId: string; + runtimeAuthorizationSeal: string; + }): Promise { + const metadata = await this.getMetadata(); + if (!metadata || metadata.identity.userId !== input.ownerId) return false; + const secret = await resolveSecret(this.env.NEXTAUTH_SECRET); + if (!secret) return false; + let authorization: RuntimeAuthorization; + try { + authorization = await unsealRuntimeAuthorization(input.runtimeAuthorizationSeal, secret, { + resourceKind: 'cloud-agent-next', + resourceId: metadata.identity.sessionId, + userId: metadata.identity.userId, + organizationId: metadata.identity.orgId, + }); + } catch { + return false; + } + if (authorization.state !== 'active') return false; + const current = RuntimeAuthorizationSchema.safeParse( + this.ctx.storage.kv.get(RUNTIME_AUTHORIZATION_KEY) + ); + if (!current.success || current.data.id !== input.expectedOldId) return false; + this.ctx.storage.kv.put(RUNTIME_AUTHORIZATION_KEY, authorization); + return true; + } + async getCredentialMetadata(): Promise { if (this.deletedWorktreeId) return null; return this.terminalLifecycle.isBlocked() ? null : this.terminalLifecycle.getStoredMetadata(); @@ -1470,6 +1787,9 @@ export class SandboxSession extends DurableObject { rows?: number; operationId?: string; }): Promise> { + if (this.ctx.storage.kv.get(RUNTIME_AUTHORIZATION_RECOVERY_KEY)) { + return { success: false, error: 'Runtime authorization recovery is in progress' }; + } if (this.pendingRuntimeCleanup()) return { success: false, error: 'Runtime cleanup is pending' }; return this.trackOperation(this.terminalLifecycle.createTerminal(input)); @@ -1480,6 +1800,9 @@ export class SandboxSession extends DurableObject { cols?: number; rows?: number; }): Promise> { + if (this.ctx.storage.kv.get(RUNTIME_AUTHORIZATION_RECOVERY_KEY)) { + return { success: false, error: 'Runtime authorization recovery is in progress' }; + } return this.trackOperation(this.terminalLifecycle.resizeTerminal(input)); } @@ -1776,6 +2099,28 @@ export class SandboxSession extends DurableObject { } return { success: true }; } + let runtimeAuthorization: RuntimeAuthorization | undefined; + if (input.runtimeAuthorizationSeal) { + const secret = await resolveSecret(this.env.NEXTAUTH_SECRET); + if (!secret) return { success: false, error: 'Authentication unavailable' }; + let authorization: RuntimeAuthorization; + try { + authorization = await unsealRuntimeAuthorization(input.runtimeAuthorizationSeal, secret, { + resourceKind: 'cloud-agent-next', + resourceId: input.identity.sessionId, + userId: input.identity.userId, + organizationId: input.identity.orgId, + }); + } catch { + return { success: false, error: 'Invalid runtime authorization' }; + } + if (authorization.state !== 'active') + return { success: false, error: 'Runtime authorization revoked' }; + if (this.ctx.storage.kv.get(RUNTIME_AUTHORIZATION_KEY)) { + return { success: false, error: 'Runtime authorization already installed' }; + } + runtimeAuthorization = authorization; + } try { validateControlSessionOptions(input); } catch (error) { @@ -1810,6 +2155,9 @@ export class SandboxSession extends DurableObject { }); if (this.deletedWorktreeId) return { success: false, error: 'worktree_deleting' }; if (this.terminalLifecycle.isBlocked()) return { success: false, error: 'Session not found' }; + if (runtimeAuthorization) { + this.ctx.storage.kv.put(RUNTIME_AUTHORIZATION_KEY, runtimeAuthorization); + } this.ctx.storage.kv.put(METADATA_KEY, serializeSessionMetadata(metadata)); this.ctx.storage.kv.put(PENDING_INTERACTIONS_KEY, { revision: 0, @@ -2169,6 +2517,13 @@ export class SandboxSession extends DurableObject { input: ControlSessionMessageInput, origin: 'initial' | 'followup' ): Promise { + if (this.ctx.storage.kv.get(RUNTIME_AUTHORIZATION_RECOVERY_KEY)) { + return { + success: false, + code: 'COMPUTE_STOPPING', + error: 'Runtime authorization recovery is in progress', + }; + } const epoch = this.terminalLifecycle.captureEpoch(); const metadata = this.terminalLifecycle.getStoredMetadata(); if (epoch === null || !metadata || this.deletedWorktreeId) { @@ -2334,6 +2689,7 @@ export class SandboxSession extends DurableObject { messageId: string, options?: { allowCreate?: boolean } ): Promise { + if (this.ctx.storage.kv.get(RUNTIME_AUTHORIZATION_RECOVERY_KEY)) return; if (this.deletedWorktreeId) return; const metadata = this.terminalLifecycle.getStoredMetadata(); const epoch = this.terminalLifecycle.captureEpoch(); @@ -2452,7 +2808,12 @@ export class SandboxSession extends DurableObject { }; const dispatchAuthorized = async ( operation: SessionOperationAuthorization['operation'], - payload: unknown + payload: unknown, + expectedConnection?: { + providerInstanceId: string; + connectionId: string; + wrapperInstanceId: string; + } ) => { if (!wrapperInstanceId) throw new Error('Wrapper identity is missing'); const authorization: SessionOperationAuthorization = { @@ -2466,7 +2827,7 @@ export class SandboxSession extends DurableObject { let dispatched: Awaited>; try { dispatched = await dispatchSessionOperation( - { authorization, payload }, + { authorization, payload, expectedConnection }, { read: () => this.loadMessages(), commit: messages => this.saveMessages(messages, epoch, 'wrapper_outcome'), @@ -2704,10 +3065,73 @@ export class SandboxSession extends DurableObject { throw new Error('Contained session attachment is unavailable'); const attachPayload = { ...status.attachment, + ...(hasModernRuntimeAuthorization(metadata) + ? { runtimeIsolation: 'per-session' as const } + : {}), ...(needsPreparation ? { preparation: { attemptId: recorder.attemptId, triggerMessageId: messageId } } : {}), }; + const authorization = RuntimeAuthorizationSchema.safeParse( + this.ctx.storage.kv.get(RUNTIME_AUTHORIZATION_KEY) + ); + let proxyKilo: SessionAttachPayload['kilo'] | undefined; + let proxyFence: + | { + providerInstanceId: string; + connectionId: string; + wrapperInstanceId: string; + } + | undefined; + if (authorization.success && authorization.data.state === 'active') { + const proxyBaseUrl = this.env.WORKER_URL + ? runtimeCredentialProxyFacadeBaseUrl(this.env.WORKER_URL) + : null; + if (!proxyBaseUrl) throw new Error('Runtime credential proxy is unavailable'); + const handle = await wait( + () => + this.issueRuntimeCredentialProxyGrant({ + wrapperRunId: '', + wrapperGeneration: 0, + wrapperConnectionId: '', + }), + SANDBOX_CONTROL_ATTACH_TIMEOUT_MS + ); + if (!handle) throw new Error('Runtime credential proxy grant is unavailable'); + const claims = await verifyRuntimeCredentialProxyHandle(this.env, handle); + if (!claims) throw new Error('Runtime credential proxy grant is invalid'); + const grant = runtimeProxyGrantSchema.safeParse( + this.ctx.storage.kv.get(RUNTIME_PROXY_GRANT_KEY) + ); + if (!grant.success || grant.data.plane !== 'control') { + throw new Error('Runtime credential proxy grant is unavailable'); + } + proxyFence = { + providerInstanceId: grant.data.providerInstanceId, + connectionId: grant.data.connectionId, + wrapperInstanceId: grant.data.wrapperInstanceId, + }; + proxyKilo = { + ...status.attachment.kilo, + token: handle, + targets: { + backendBaseUrl: proxyBaseUrl, + providerBaseUrl: proxyBaseUrl, + sessionIngestBaseUrl: proxyBaseUrl, + }, + }; + if (getSandboxProvider(metadata) === 'vercel') { + await wait(() => + control.bindRuntimeCredentialProxyHandle({ + ownerId: metadata.identity.userId, + sessionId, + kiloSessionId, + directory: session.directory, + handle, + }) + ); + } + } phase = needsPreparation ? 'preparing' : 'attach'; await wait(() => control.attachSession({ @@ -2726,7 +3150,21 @@ export class SandboxSession extends DurableObject { return; } if (operationResults) { - if ((await dispatchAuthorized('session.attach', attachPayload)).state === 'running') { + if ( + ( + await dispatchAuthorized( + 'session.attach', + proxyKilo + ? { + ...attachPayload, + env: { ...attachPayload.env, KILOCODE_TOKEN: proxyKilo.token }, + kilo: proxyKilo, + } + : attachPayload, + proxyFence + ) + ).state === 'running' + ) { await this.armQueueRetry(Math.min(deadlineAt, Date.now() + QUEUE_RETRY_MS)); return; } @@ -2740,7 +3178,14 @@ export class SandboxSession extends DurableObject { operation: 'session.attach', session, expectedWrapperInstanceId: wrapperInstanceId, - payload: attachPayload, + ...(proxyFence ? { expectedConnection: proxyFence } : {}), + payload: proxyKilo + ? { + ...attachPayload, + env: { ...attachPayload.env, KILOCODE_TOKEN: proxyKilo.token }, + kilo: proxyKilo, + } + : attachPayload, timeoutMs: SANDBOX_CONTROL_ATTACH_TIMEOUT_MS, }) ) diff --git a/services/cloud-agent-next/src/sandbox-session/control-rpc.test.ts b/services/cloud-agent-next/src/sandbox-session/control-rpc.test.ts new file mode 100644 index 0000000000..0f133f87b1 --- /dev/null +++ b/services/cloud-agent-next/src/sandbox-session/control-rpc.test.ts @@ -0,0 +1,39 @@ +import { describe, expect, it, vi } from 'vitest'; +import { sandboxControlRpc } from './control-rpc.js'; + +describe('sandboxControlRpc', () => { + it('retries an ambiguous bind response without changing the request identity', async () => { + const input = { + ownerId: 'user_1', + sessionId: 'workspace_11111111-1111-4111-8111-111111111111', + kiloSessionId: 'ses_abcdefghijklmnopqrstuvwxyz', + directory: '/workspace/worktree', + handle: 'session-runtime-proxy-handle', + }; + const responseLost = Object.assign(new Error('response lost after bind'), { retryable: true }); + const bind = vi + .fn() + .mockRejectedValueOnce(responseLost) + .mockResolvedValueOnce('worktree-runtime-proxy-handle'); + const getByName = vi + .fn() + .mockReturnValueOnce({ bindRuntimeCredentialProxyHandle: bind }) + .mockReturnValueOnce({ bindRuntimeCredentialProxyHandle: bind }); + vi.stubGlobal('scheduler', { wait: vi.fn().mockResolvedValue(undefined) }); + + try { + await expect( + sandboxControlRpc( + { SANDBOX_CONTROL: { getByName } } as never, + 'sandbox_credential_proxy' + ).bindRuntimeCredentialProxyHandle(input) + ).resolves.toBe('worktree-runtime-proxy-handle'); + expect(getByName).toHaveBeenCalledTimes(2); + expect(bind).toHaveBeenCalledTimes(2); + expect(bind).toHaveBeenNthCalledWith(1, input); + expect(bind).toHaveBeenNthCalledWith(2, input); + } finally { + vi.unstubAllGlobals(); + } + }); +}); diff --git a/services/cloud-agent-next/src/sandbox-session/control-rpc.ts b/services/cloud-agent-next/src/sandbox-session/control-rpc.ts index 52fc7c653b..6541e525f5 100644 --- a/services/cloud-agent-next/src/sandbox-session/control-rpc.ts +++ b/services/cloud-agent-next/src/sandbox-session/control-rpc.ts @@ -16,7 +16,11 @@ import type { } from '../sandbox-control/terminal-billing.js'; import type { SessionOperationAuthorization } from '../shared/sandbox-control-protocol.js'; import type { Env } from '../types.js'; -import type { RuntimeQuarantineResult, SandboxAcquisition } from '../persistence/SandboxControl.js'; +import type { + ControlRuntimeCredentialProxyFence, + RuntimeQuarantineResult, + SandboxAcquisition, +} from '../persistence/SandboxControl.js'; import type { SandboxBillingInput } from '../container-usage-context.js'; import { getSandboxControlStub } from '../sandbox-control/stub.js'; import { withDORetry } from '../utils/do-retry.js'; @@ -39,6 +43,7 @@ type SandboxControlRpc = { physical: PhysicalState; wrapperInstanceId?: string; operationResults?: true; + runtimeRecovery?: true; attachment?: SessionAttachPayload; }>; getStatus(): Promise<{ @@ -46,7 +51,14 @@ type SandboxControlRpc = { physical: PhysicalState; wrapperInstanceId?: string; operationResults?: true; + runtimeRecovery?: true; }>; + getRuntimeCredentialProxyFence(input: { + ownerId: string; + sessionId: string; + kiloSessionId: string; + directory: string; + }): Promise; quarantineRuntime(input: { ownerId: string; sessionId: string; @@ -56,6 +68,13 @@ type SandboxControlRpc = { authorization?: SessionOperationAuthorization; }): Promise; attachSession(input: AttachRouteInput): Promise; + bindRuntimeCredentialProxyHandle(input: { + ownerId: string; + sessionId: string; + kiloSessionId: string; + directory: string; + handle: string; + }): Promise<{ bound: true }>; detachSession(sessionId: string): Promise<{ existed: boolean }>; validateTerminalAccess(input: SandboxTerminalAccessInput): Promise; recordTerminalActivity(input: SandboxTerminalAccessInput): Promise; @@ -92,8 +111,21 @@ export function sandboxControlRpc( ), ensureReady: input => stub().ensureReady(input), getStatus: () => withDORetry(stub, control => control.getStatus(), 'getStatus', config()), + getRuntimeCredentialProxyFence: input => + withDORetry( + stub, + control => control.getRuntimeCredentialProxyFence(input), + 'getRuntimeCredentialProxyFence', + config() + ), quarantineRuntime: input => stub().quarantineRuntime(input), attachSession: input => stub().attachSession(input), + bindRuntimeCredentialProxyHandle: input => + withDORetry( + stub, + control => control.bindRuntimeCredentialProxyHandle(input), + 'bindRuntimeCredentialProxyHandle' + ), detachSession: sessionId => withDORetry(stub, control => control.detachSession(sessionId), 'detachSession'), validateTerminalAccess: input => diff --git a/services/cloud-agent-next/src/sandbox-session/session-client.ts b/services/cloud-agent-next/src/sandbox-session/session-client.ts index d8779a83ca..54a944d40a 100644 --- a/services/cloud-agent-next/src/sandbox-session/session-client.ts +++ b/services/cloud-agent-next/src/sandbox-session/session-client.ts @@ -6,6 +6,12 @@ export type SessionClient = Pick< | 'fetch' | 'closeOrgStreams' | 'getMetadata' + | 'getRuntimeAuthorizationStatus' + | 'getRuntimeAuthorizationRecoveryState' + | 'isRuntimeAuthorizationRecoveryInProgress' + | 'recoverExpiredRuntimeAuthorization' + | 'issueRuntimeCredentialProxyGrant' + | 'resolveRuntimeCredentialProxyGrant' | 'validateKiloGlobalFeedProducer' | 'getLatestAssistantMessage' | 'getLatestEventId' diff --git a/services/cloud-agent-next/src/sandbox-session/session-message-queue.test.ts b/services/cloud-agent-next/src/sandbox-session/session-message-queue.test.ts index 775157f538..ce472a67f7 100644 --- a/services/cloud-agent-next/src/sandbox-session/session-message-queue.test.ts +++ b/services/cloud-agent-next/src/sandbox-session/session-message-queue.test.ts @@ -1,3 +1,4 @@ +import jwt from 'jsonwebtoken'; import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; import { normalizeCliEvent } from '../../../../packages/cloud-agent-sdk/src/normalizer'; import { createServiceState } from '../../../../packages/cloud-agent-sdk/src/service-state'; @@ -31,6 +32,7 @@ import { DEADLINE_MS } from '../sandbox-control/deadlines.js'; import { createControlPlaneCredential } from '../sandbox-control/managed-credential.js'; import { logger } from '../logger.js'; import { SESSION_DELIVERY_TIMEOUT_MS } from './control-dispatch.js'; +import { RUNTIME_AUTHORIZATION_KEY } from '../session/runtime-authorization-persistence.js'; import { createControlStopRequest } from '../shared/control-plane-session.js'; import type { AcceptedCommandTurn, @@ -1157,6 +1159,13 @@ function sessionFixture(overrides: Partial = {}, sharedControl? }); const control = { getStatus: vi.fn(async (): Promise => ({ ...status })), + getRuntimeCredentialProxyFence: vi.fn(async () => ({ + plane: 'control' as const, + allocationId: 'allocation_1', + providerInstanceId: 'provider_1', + connectionId: 'connection_1', + wrapperInstanceId: RUNTIME_ID, + })), ensureReady: vi.fn( async (_input: Parameters[0]): Promise => ({ ...status, @@ -1164,6 +1173,7 @@ function sessionFixture(overrides: Partial = {}, sharedControl? }) ), attachSession: vi.fn(async () => ({})), + bindRuntimeCredentialProxyHandle: vi.fn(async () => ({ bound: true as const })), detachSession: vi.fn(async () => ({ existed: true })), quarantineRuntime: vi.fn( async ( @@ -1181,6 +1191,8 @@ function sessionFixture(overrides: Partial = {}, sharedControl? } satisfies Control; const env = { SANDBOX_CONTROL: { getByName: () => sharedControl ?? control }, + WORKER_URL: 'https://worker.example.test', + NEXTAUTH_SECRET: 'test-secret', CLOUD_AGENT_CONTAINER_BILLING_ENABLED: 'true', CLOUD_AGENT_CONTAINER_BILLING_ORG_IDS: 'org_1', CLOUD_AGENT_CONTAINER_BILLING_USER_IDS: 'user_1', @@ -1191,6 +1203,7 @@ function sessionFixture(overrides: Partial = {}, sharedControl? return session; }, control, + env, metadata, storage, values, @@ -1270,6 +1283,38 @@ function sessionFixture(overrides: Partial = {}, sharedControl? }; } +function installModernRuntimeAuthorization(fixture: ReturnType) { + const authorizationId = '44444444-4444-4444-8444-444444444444'; + const token = jwt.sign( + { + runtimeAuthorization: { id: authorizationId }, + exp: Math.floor(Date.now() / 1000) + 60 * 60, + }, + 'test-secret' + ); + fixture.storage.kv.put( + 'session_metadata', + serializeSessionMetadata({ + ...fixture.metadata, + auth: { ...fixture.metadata.auth, kilocodeToken: token }, + }) + ); + fixture.storage.kv.put(RUNTIME_AUTHORIZATION_KEY, { + version: 1, + id: authorizationId, + resourceKind: 'cloud-agent-next', + resourceId: SESSION_ID, + userId: 'user_1', + authorizationUserId: 'user_1', + issuedAt: '2026-01-01T00:00:00.000Z', + delegationExpiresAt: '2026-01-02T00:00:00.000Z', + state: 'active', + bindings: { userPepperDigest: 'a'.repeat(64), authorizationPepperDigest: 'b'.repeat(64) }, + source: { admissionSource: 'user' }, + }); + return token; +} + function delegateRequest( fixture: ReturnType, operation: SandboxControlOutboundRequest['operation'], @@ -4429,6 +4474,103 @@ describe('SandboxSession orchestration', () => { } ); + it.each(['revoked', 'deleted'] as const)( + 'immediately denies runtime proxy issue and resolution after terminal lifecycle is %s despite pending or failed detach', + async lifecycle => { + const fixture = sessionFixture({ + identity: { + sessionId: SESSION_ID, + userId: 'user_1', + orgId: 'org_1', + billingOrigin: 'cloud-agent-web', + }, + }); + installModernRuntimeAuthorization(fixture); + const handle = await fixture.session.issueRuntimeCredentialProxyGrant({ + wrapperRunId: 'ignored', + wrapperGeneration: 0, + wrapperConnectionId: 'ignored', + }); + expect(handle).toEqual(expect.any(String)); + + const detach = deferred<{ existed: boolean }>(); + fixture.control.detachSession.mockImplementationOnce(() => detach.promise); + const blocked = + lifecycle === 'revoked' + ? fixture.session.closeOrgStreams('org_1') + : fixture.session.deleteSession(); + + await expect( + fixture.session.issueRuntimeCredentialProxyGrant({ + wrapperRunId: 'ignored', + wrapperGeneration: 0, + wrapperConnectionId: 'ignored', + }) + ).resolves.toBeNull(); + await expect(fixture.session.resolveRuntimeCredentialProxyGrant(handle!)).resolves.toBeNull(); + + detach.reject(new Error('detach failed')); + await expect(blocked).rejects.toThrow('detach failed'); + await expect(fixture.session.resolveRuntimeCredentialProxyGrant(handle!)).resolves.toBeNull(); + } + ); + + it.each([false, true])( + 'sends modern attach credentials through a fenced proxy grant with operation results %s', + async operationResults => { + const fixture = sessionFixture(); + const backingToken = installModernRuntimeAuthorization(fixture); + + const nativeRuntimeId = '44444444-4444-4444-8444-444444444444'; + fixture.setStatus({ + physical: 'running', + connection: 'ready', + wrapperInstanceId: RUNTIME_ID, + ...(operationResults ? { operationResults: true as const } : {}), + }); + delegateRequest(fixture, 'session.attach', async () => + controlResponse({ attached: true, nativeRuntimeId }) + ); + + await fixture.admit('modern-proxy'); + await fixture.flush(); + + const attach = fixture.control.request.mock.calls.find( + ([input]) => input.operation === 'session.attach' + )?.[0]; + expect(attach).toMatchObject({ + expectedConnection: { + providerInstanceId: 'provider_1', + connectionId: 'connection_1', + wrapperInstanceId: RUNTIME_ID, + }, + payload: { + kilo: { + scopeId: SESSION_ID, + token: expect.stringMatching(/^[A-Za-z0-9_-]+\.[A-Za-z0-9_-]+\.[A-Za-z0-9_-]+$/), + targets: { + backendBaseUrl: 'https://worker.example.test', + providerBaseUrl: 'https://worker.example.test', + sessionIngestBaseUrl: 'https://worker.example.test', + }, + }, + }, + }); + if (operationResults) { + expect(fixture.values.get('native_runtime_fence')).toMatchObject({ + sandboxId: SANDBOX_ID, + wrapperInstanceId: RUNTIME_ID, + nativeRuntimeId, + authorization: attach?.authorization, + }); + } + const serialized = JSON.stringify(attach?.payload); + expect(serialized).not.toContain(KILO_CREDENTIAL); + expect(serialized).not.toContain(backingToken); + expect(serialized).not.toContain('test-secret'); + } + ); + it.each(['completed', 'failed', 'cancelled'] as const)( 'settles an early %s outcome once without resurrecting work on acknowledgement', async status => { @@ -5661,6 +5803,26 @@ describe('SandboxSession orchestration', () => { expect(await fixture.snapshot()).toMatchObject({ preparationSnapshots: coldPreparation }); }); + it('keeps a persisted modern attachment isolated after the rollout flag is disabled', async () => { + const fixture = sessionFixture({ + auth: { + kiloSessionId: 'kilo_root', + kilocodeToken: 'eyJhbGciOiJub25lIn0.eyJydW50aW1lQXV0aG9yaXphdGlvbiI6e319.', + }, + }); + fixture.env.RUNTIME_ISOLATION_ENABLED = 'false'; + + await fixture.admit('modern'); + await fixture.flush(); + + expect(fixture.control.request).toHaveBeenCalledWith( + expect.objectContaining({ + operation: 'session.attach', + payload: expect.objectContaining({ runtimeIsolation: 'per-session' }), + }) + ); + }); + it.each(['cloudflare', 'vercel'] as const)( 'refreshes direct credentials without warm preparation across eviction on %s', async sandboxProvider => { diff --git a/services/cloud-agent-next/src/sandbox-session/session-operation.ts b/services/cloud-agent-next/src/sandbox-session/session-operation.ts index 83a1e76ab7..6f2ec25024 100644 --- a/services/cloud-agent-next/src/sandbox-session/session-operation.ts +++ b/services/cloud-agent-next/src/sandbox-session/session-operation.ts @@ -1,5 +1,8 @@ import type { DORetryScope } from '@kilocode/worker-utils'; -import type { SandboxControlOutboundRequest } from '../sandbox-control/socket.js'; +import type { + SandboxControlConnectionIdentity, + SandboxControlOutboundRequest, +} from '../sandbox-control/socket.js'; import type { EventQueries } from '../session/queries/index.js'; import type { StoredEvent } from '../websocket/types.js'; import { @@ -148,7 +151,11 @@ export async function reconcileSessionOperation( } export async function dispatchSessionOperation( - input: { authorization: SessionOperationAuthorization; payload: unknown }, + input: { + authorization: SessionOperationAuthorization; + payload: unknown; + expectedConnection?: SandboxControlConnectionIdentity; + }, messages: OperationMessages, effects: SessionOperationEffects & { isCurrent: () => boolean } ): Promise { @@ -220,6 +227,7 @@ export async function dispatchSessionOperation( authorization, session: authorization.session, expectedWrapperInstanceId: authorization.wrapperInstanceId, + ...(input.expectedConnection ? { expectedConnection: input.expectedConnection } : {}), payload, timeoutMs, deadlineAt, diff --git a/services/cloud-agent-next/src/sandbox-session/terminal-lifecycle.ts b/services/cloud-agent-next/src/sandbox-session/terminal-lifecycle.ts index 93149fae6b..aa752577ea 100644 --- a/services/cloud-agent-next/src/sandbox-session/terminal-lifecycle.ts +++ b/services/cloud-agent-next/src/sandbox-session/terminal-lifecycle.ts @@ -815,6 +815,23 @@ export function createSandboxTerminalLifecycle(deps: TerminalLifecycleDeps) { } } + function clearAttachedWrapperAfterRecovery(wrapperInstanceId: string): boolean { + const attached = readAttachedSession(); + if (!attached || attached.wrapperInstanceId !== wrapperInstanceId) return false; + for (const [, raw] of storage.kv.list({ prefix: TERMINAL_PREFIX })) { + const terminal = terminalRecordSchema.safeParse(raw); + if ( + terminal.success && + terminal.data.state === 'running' && + terminal.data.wrapperInstanceId === wrapperInstanceId + ) { + return false; + } + } + storage.kv.delete(ATTACHED_SESSION_KEY); + return true; + } + function purgeDeletedState(): void { if (readFence()?.state !== 'deleted') return; const keys = Array.from(storage.kv.list(), ([key]) => key); @@ -826,6 +843,7 @@ export function createSandboxTerminalLifecycle(deps: TerminalLifecycleDeps) { return { beginDeletion, beginRevocation, + clearAttachedWrapperAfterRecovery, captureEpoch: () => snapshot()?.epoch ?? null, cleanupSession, closeTerminal, 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 1feea831ba..282521947c 100644 --- a/services/cloud-agent-next/src/server-stream-ticket.test.ts +++ b/services/cloud-agent-next/src/server-stream-ticket.test.ts @@ -308,7 +308,10 @@ describe('server /stream ticket nonce consume', () => { installNonceStore(env); const doFetch = vi.fn().mockResolvedValue(new Response('ok', { status: 200 })); env.SANDBOX_SESSION.idFromName.mockReturnValue('sandbox-session-do-id'); - env.SANDBOX_SESSION.get.mockReturnValue({ fetch: doFetch }); + env.SANDBOX_SESSION.get.mockReturnValue({ + fetch: doFetch, + isRuntimeAuthorizationRecoveryInProgress: vi.fn().mockResolvedValue(false), + }); const ticket = signStreamTicket({ cloudAgentSessionId: sessionId }); const response = await fetchWorker( @@ -395,6 +398,7 @@ describe('server /terminal ticket nonce consume', () => { env.CLOUD_AGENT_SESSION.idFromName.mockReturnValue('session-do-id'); env.CLOUD_AGENT_SESSION.get.mockReturnValue({ + isRuntimeAuthorizationRecoveryInProgress: vi.fn().mockResolvedValue(false), getMetadata: vi.fn().mockResolvedValue({ metadataSchemaVersion: 2, identity: { @@ -426,7 +430,10 @@ describe('server /terminal ticket nonce consume', () => { const consumedNonces = installNonceStore(env); const sessionFetch = vi.fn().mockResolvedValue(new Response('bridged', { status: 200 })); env.SANDBOX_SESSION.idFromName.mockReturnValue('sandbox-session-do-id'); - env.SANDBOX_SESSION.get.mockReturnValue({ fetch: sessionFetch }); + env.SANDBOX_SESSION.get.mockReturnValue({ + fetch: sessionFetch, + isRuntimeAuthorizationRecoveryInProgress: vi.fn().mockResolvedValue(false), + }); const ticket = signTerminalTicket({ cloudAgentSessionId: sessionId }); const first = await fetchWorker(terminalRequest(ticket, sessionId), env); @@ -452,7 +459,10 @@ describe('server /terminal ticket nonce consume', () => { const consumedNonces = installNonceStore(env); const sessionFetch = vi.fn().mockResolvedValue(new Response('bridged', { status: 200 })); env.SANDBOX_SESSION.idFromName.mockReturnValue('sandbox-session-do-id'); - env.SANDBOX_SESSION.get.mockReturnValue({ fetch: sessionFetch }); + env.SANDBOX_SESSION.get.mockReturnValue({ + fetch: sessionFetch, + isRuntimeAuthorizationRecoveryInProgress: vi.fn().mockResolvedValue(false), + }); requireCurrentSessionAccessMock.mockRejectedValueOnce( Object.assign(new Error('Session access denied'), { code: 'FORBIDDEN' }) ); diff --git a/services/cloud-agent-next/src/server.test.ts b/services/cloud-agent-next/src/server.test.ts index 9cb42466b9..86bda7b1f7 100644 --- a/services/cloud-agent-next/src/server.test.ts +++ b/services/cloud-agent-next/src/server.test.ts @@ -4,6 +4,14 @@ import { VERCEL_SANDBOX_UNAVAILABLE_MESSAGE } from './agent-sandbox/vercel/verce import type { Env } from './types.js'; import { mintWrapperDispatchTicket, type WrapperDispatchTicketClaims } from './auth.js'; import { mintControlLogUploadGrant } from './sandbox-control/log-upload-grant.js'; +import { + createRuntimeProxyGrant, + issueRuntimeCredentialProxyHandle, +} from './runtime-credential-proxy.js'; +import { + RUNTIME_PROXY_ATTESTATION_HEADER, + verifyRuntimeProxyAttestation, +} from '@kilocode/worker-utils/runtime-proxy-attestation'; const { getRunningTerminalClientMock, @@ -162,7 +170,7 @@ function createEnv(): MockEnv { INTERNAL_API_SECRET: 'test-internal-secret', CLOUD_AGENT_SESSION: { idFromName: vi.fn(), - get: vi.fn(), + get: vi.fn(() => ({ getRuntimeAuthorizationStatus: vi.fn().mockResolvedValue('legacy') })), }, USER_KILO_FACADE: { idFromName: vi.fn(), @@ -177,7 +185,9 @@ function createEnv(): MockEnv { }, SANDBOX_SESSION: { idFromName: vi.fn(), - get: vi.fn(), + get: vi.fn(() => ({ + isRuntimeAuthorizationRecoveryInProgress: vi.fn().mockResolvedValue(false), + })), }, }; } @@ -395,7 +405,10 @@ describe('server /terminal', () => { const sessionResponse = new Response('bridged', { status: 200 }); const sessionFetch = vi.fn().mockResolvedValue(sessionResponse); env.SANDBOX_SESSION.idFromName.mockReturnValue('sandbox-session-do-id'); - env.SANDBOX_SESSION.get.mockReturnValue({ fetch: sessionFetch }); + env.SANDBOX_SESSION.get.mockReturnValue({ + fetch: sessionFetch, + isRuntimeAuthorizationRecoveryInProgress: vi.fn().mockResolvedValue(false), + }); const request = new Request( `http://worker.test/terminal?cloudAgentSessionId=${sessionId}&ptyId=pty_123&ticket=${encodeURIComponent(ticket)}&role=wrapper&ownerId=attacker`, { @@ -451,6 +464,31 @@ describe('server /terminal', () => { expect(forwarded.headers.get('x-forwarded-user')).toBeNull(); }); + it('rejects a control-plane browser upgrade during runtime authorization recovery', async () => { + const sessionId = 'workspace_aaaaaaaa-aaaa-aaaa-aaaa-aaaaaaaaaaaa'; + const env = createEnv(); + const consume = installTerminalNonceConsumer(env); + const sessionFetch = vi.fn(); + env.SANDBOX_SESSION.idFromName.mockReturnValue('sandbox-session-do-id'); + env.SANDBOX_SESSION.get.mockReturnValue({ + fetch: sessionFetch, + isRuntimeAuthorizationRecoveryInProgress: vi.fn().mockResolvedValue(true), + }); + + const response = await fetchWorker( + new Request( + `http://worker.test/terminal?cloudAgentSessionId=${sessionId}&ptyId=pty_123&ticket=${encodeURIComponent(signTerminalTicket(sessionId))}`, + { headers: { Upgrade: 'websocket' } } + ), + env + ); + + expect(response.status).toBe(503); + await expect(response.text()).resolves.toBe('Runtime authorization recovery is in progress'); + expect(consume).toHaveBeenCalledOnce(); + expect(sessionFetch).not.toHaveBeenCalled(); + }); + it('rejects revoked control-plane access before consuming the browser ticket nonce', async () => { const sessionId = 'workspace_aaaaaaaa-aaaa-aaaa-aaaa-aaaaaaaaaaaa'; const ticket = signTerminalTicket(sessionId); @@ -540,7 +578,11 @@ describe('server /terminal', () => { const getMetadata = vi.fn().mockResolvedValue(metadata); const fetch = vi.fn(); env.CLOUD_AGENT_SESSION.idFromName.mockReturnValue('do-id'); - env.CLOUD_AGENT_SESSION.get.mockReturnValue({ fetch, getMetadata }); + env.CLOUD_AGENT_SESSION.get.mockReturnValue({ + fetch, + getMetadata, + isRuntimeAuthorizationRecoveryInProgress: vi.fn().mockResolvedValue(false), + }); const request = new Request( `http://worker.test/terminal?cloudAgentSessionId=session-1&ptyId=pty_123&ticket=${encodeURIComponent(ticket)}`, @@ -612,6 +654,7 @@ describe('server /terminal', () => { }); env.CLOUD_AGENT_SESSION.idFromName.mockReturnValue('do-id'); env.CLOUD_AGENT_SESSION.get.mockReturnValue({ + isRuntimeAuthorizationRecoveryInProgress: vi.fn().mockResolvedValue(false), getMetadata: vi.fn().mockResolvedValue({ metadataSchemaVersion: 2, identity: { @@ -724,6 +767,533 @@ describe('server /terminal', () => { }); }); +describe('server runtime credential proxy', () => { + async function handle(): Promise { + return issueRuntimeCredentialProxyHandle( + { NEXTAUTH_SECRET: secret } as never, + createRuntimeProxyGrant({ + plane: 'legacy', + authorizationId: '11111111-1111-4111-8111-111111111111', + sessionId: 'agent_proxy', + kiloSessionId: 'kilo_proxy', + userId: 'usr_proxy', + orgId: 'org_proxy', + mode: 'contained', + generation: 1, + allocationId: 'allocation_proxy', + wrapperRunId: 'run_proxy', + wrapperConnectionId: 'connection_proxy', + leaseExpiresAt: Date.now() + 60_000, + state: 'active', + }) + ); + } + + it('denies invalid handles before resolving a session or fetching', async () => { + const env = createEnv(); + const upstream = vi.fn(); + vi.stubGlobal('fetch', upstream); + const response = await fetchWorker( + new Request('https://worker.test/api/runtime-credential-proxy/provider/models', { + headers: { Authorization: 'Bearer invalid' }, + }), + env + ); + expect(response.status).toBe(401); + expect(env.CLOUD_AGENT_SESSION.get).not.toHaveBeenCalled(); + expect(upstream).not.toHaveBeenCalled(); + vi.unstubAllGlobals(); + }); + + it('routes verified facade requests and leaves non-handles to ordinary routing', async () => { + const env = Object.assign(createEnv(), { WORKER_URL: 'https://worker.test' }); + const resolve = vi.fn().mockResolvedValue({ + token: 'https://api.kilo.ai:backing-token', + organizationId: 'org_proxy', + runtimeAuthorization: { + userId: 'usr_proxy', + authorizationId: '11111111-1111-4111-8111-111111111111', + resourceId: 'agent_proxy', + }, + }); + env.CLOUD_AGENT_SESSION.get.mockReturnValue({ resolveRuntimeCredentialProxyGrant: resolve }); + const upstream = vi.fn().mockResolvedValue(new Response('ok')); + vi.stubGlobal('fetch', upstream); + try { + const authorization = `Bearer ${await handle()}`; + const requests = [ + ['GET', '/api/profile', undefined], + ['GET', '/api/defaults', undefined], + ['GET', '/api/openrouter/models', undefined], + ['POST', '/api/openrouter/chat/completions', '{}'], + ['POST', '/api/gateway/chat/completions', '{}'], + ['POST', '/api/gateway/v1/chat/completions', '{}'], + ['POST', '/api/gateway/v1/responses', '{}'], + ['POST', '/api/session', '{"sessionId":"kilo_proxy"}'], + ['GET', '/api/session/kilo_proxy/export', undefined], + ['POST', '/api/session/kilo_proxy/ingest', '{}'], + ['POST', '/api/session/kilo_proxy/title', '{}'], + ] as const; + for (const [method, path, body] of requests) { + const response = await fetchWorker( + new Request(`https://worker.test${path}`, { + method, + headers: { + Authorization: authorization, + ...(body === undefined ? {} : { 'Content-Type': 'application/json' }), + }, + ...(body === undefined ? {} : { body }), + }), + env + ); + expect(response.status).toBe(200); + } + expect( + upstream.mock.calls.map(([request]) => new URL((request as Request).url).pathname) + ).toEqual([ + '/api/profile', + '/api/defaults', + '/api/gateway/models', + '/api/gateway/chat/completions', + '/api/gateway/chat/completions', + '/api/gateway/v1/chat/completions', + '/api/gateway/v1/responses', + '/api/session', + '/api/session/kilo_proxy/export', + '/api/session/kilo_proxy/ingest', + '/api/session/kilo_proxy/title', + ]); + const createRequest = upstream.mock.calls[7]?.[0] as Request; + expect(await createRequest.text()).toBe('{"sessionId":"kilo_proxy"}'); + + const invalid = await fetchWorker( + new Request('https://worker.test/api/profile', { + headers: { Authorization: 'Bearer invalid' }, + }), + env + ); + const unknown = await fetchWorker( + new Request('https://worker.test/api/not-allowed', { + headers: { Authorization: authorization }, + }), + env + ); + expect(invalid.status).toBe(404); + expect(unknown.status).toBe(404); + expect(upstream).toHaveBeenCalledTimes(requests.length); + } finally { + vi.unstubAllGlobals(); + } + }); + + it('does not let one root handle access a sibling path or body-bound ingest request', async () => { + const env = Object.assign(createEnv(), { WORKER_URL: 'https://worker.test' }); + const firstHandle = await handle(); + const secondHandle = await issueRuntimeCredentialProxyHandle( + { NEXTAUTH_SECRET: secret } as never, + createRuntimeProxyGrant({ + plane: 'control', + authorizationId: '22222222-2222-4222-8222-222222222222', + sessionId: 'agent_sibling', + kiloSessionId: 'kilo_sibling', + userId: 'usr_proxy', + mode: 'contained', + allocationId: 'allocation_proxy', + providerInstanceId: 'provider_proxy', + connectionId: 'connection_proxy', + wrapperInstanceId: 'wrapper_proxy', + leaseExpiresAt: Date.now() + 60_000, + state: 'active', + }) + ); + const firstResolve = vi.fn().mockResolvedValue({ + token: 'https://api.kilo.ai:backing-token', + runtimeAuthorization: { + userId: 'usr_proxy', + authorizationId: '11111111-1111-4111-8111-111111111111', + resourceId: 'agent_proxy', + }, + }); + const secondResolve = vi.fn().mockResolvedValue({ + token: 'https://api.kilo.ai:backing-token', + runtimeAuthorization: { + userId: 'usr_proxy', + authorizationId: '22222222-2222-4222-8222-222222222222', + resourceId: 'agent_sibling', + }, + }); + env.CLOUD_AGENT_SESSION.idFromName.mockImplementation((name: string) => name); + env.CLOUD_AGENT_SESSION.get.mockImplementation((id: string) => + id === 'usr_proxy:agent_proxy' + ? { resolveRuntimeCredentialProxyGrant: firstResolve } + : { resolveRuntimeCredentialProxyGrant: secondResolve } + ); + const upstream = vi.fn().mockResolvedValue(new Response('ok')); + vi.stubGlobal('fetch', upstream); + try { + for (const [path, body] of [ + ['/api/session/kilo_sibling/export', undefined], + ['/api/session', '{"sessionId":"kilo_sibling"}'], + ] as const) { + const response = await fetchWorker( + new Request(`https://worker.test${path}`, { + method: body === undefined ? 'GET' : 'POST', + headers: { + Authorization: `Bearer ${firstHandle}`, + ...(body === undefined ? {} : { 'Content-Type': 'application/json' }), + }, + ...(body === undefined ? {} : { body }), + }), + env + ); + expect(response.status).toBe(404); + } + expect(upstream).not.toHaveBeenCalled(); + + firstResolve.mockResolvedValue(null); + + const surviving = await fetchWorker( + new Request('https://worker.test/api/session/kilo_sibling/export', { + headers: { Authorization: `Bearer ${secondHandle}` }, + }), + env + ); + expect(surviving.status).toBe(200); + expect(secondResolve).toHaveBeenCalledOnce(); + expect(upstream).toHaveBeenCalledOnce(); + } finally { + vi.unstubAllGlobals(); + } + }); + + it('recognizes only paths under a safe configured facade prefix', async () => { + const env = Object.assign(createEnv(), { WORKER_URL: 'https://worker.test/runtime' }); + env.CLOUD_AGENT_SESSION.get.mockReturnValue({ + resolveRuntimeCredentialProxyGrant: vi.fn().mockResolvedValue({ + token: 'https://api.kilo.ai:backing-token', + runtimeAuthorization: { + userId: 'usr_proxy', + authorizationId: '11111111-1111-4111-8111-111111111111', + resourceId: 'agent_proxy', + }, + }), + }); + const upstream = vi.fn().mockResolvedValue(new Response('ok')); + vi.stubGlobal('fetch', upstream); + try { + const authorization = `Bearer ${await handle()}`; + expect( + ( + await fetchWorker( + new Request('https://worker.test/runtime/api/profile', { + headers: { Authorization: authorization }, + }), + env + ) + ).status + ).toBe(200); + expect( + ( + await fetchWorker( + new Request('https://worker.test/api/profile', { + headers: { Authorization: authorization }, + }), + env + ) + ).status + ).toBe(404); + expect(upstream).toHaveBeenCalledOnce(); + } finally { + vi.unstubAllGlobals(); + } + }); + + it('replaces caller credentials, enforces organization identity, and preserves request shape', async () => { + const env = createEnv(); + const resolve = vi.fn().mockResolvedValue({ + token: 'https://provider.example.test/api/openrouter:backing-token', + organizationId: 'org_proxy', + runtimeAuthorization: { + userId: 'usr_proxy', + authorizationId: '11111111-1111-4111-8111-111111111111', + resourceId: 'agent_proxy', + }, + }); + env.CLOUD_AGENT_SESSION.get.mockReturnValue({ resolveRuntimeCredentialProxyGrant: resolve }); + const upstream = vi.fn(async (request: Request) => { + expect(request.method).toBe('POST'); + expect(new URL(request.url).pathname).toBe('/api/openrouter/chat/completions'); + expect(await request.text()).toBe('{"stream":true}'); + expect(request.headers.get('authorization')).toMatch(/^Bearer /); + expect(request.headers.get('authorization')).not.toBe('Bearer caller-token'); + expect(request.headers.get(RUNTIME_PROXY_ATTESTATION_HEADER)).not.toBe( + 'caller-supplied-proof' + ); + expect(request.headers.get('cookie')).toBeNull(); + expect(request.headers.get('x-kilocode-organizationid')).toBe('org_proxy'); + await expect( + verifyRuntimeProxyAttestation({ + value: request.headers.get(RUNTIME_PROXY_ATTESTATION_HEADER), + secret, + audience: 'kilo-gateway', + userId: 'usr_proxy', + authorizationId: '11111111-1111-4111-8111-111111111111', + resourceId: 'agent_proxy', + bearer: 'https://provider.example.test/api/openrouter:backing-token', + }) + ).resolves.toBe(true); + return new Response('stream-body', { + status: 307, + headers: { Location: 'https://other.test' }, + }); + }); + vi.stubGlobal('fetch', upstream); + const response = await fetchWorker( + new Request( + 'https://worker.test/api/runtime-credential-proxy/provider/api/openrouter/chat/completions?stream=true', + { + method: 'POST', + headers: { + Authorization: `Bearer ${await handle()}`, + Cookie: 'session=caller', + 'X-Kilocode-OrganizationId': 'attacker-org', + [RUNTIME_PROXY_ATTESTATION_HEADER]: 'caller-supplied-proof', + 'Content-Type': 'application/json', + }, + body: '{"stream":true}', + } + ), + env + ); + expect(response.status).toBe(307); + expect(response.headers.get('location')).toBe('https://other.test'); + await expect(response.text()).resolves.toBe('stream-body'); + expect(resolve).toHaveBeenCalledOnce(); + vi.unstubAllGlobals(); + }); + + it('removes every adjacent prohibited header before injecting runtime credentials', async () => { + const env = createEnv(); + env.CLOUD_AGENT_SESSION.get.mockReturnValue({ + resolveRuntimeCredentialProxyGrant: vi.fn().mockResolvedValue({ + token: 'https://provider.example.test/api/openrouter:backing-token', + organizationId: 'org_proxy', + runtimeAuthorization: { + userId: 'usr_proxy', + authorizationId: '11111111-1111-4111-8111-111111111111', + resourceId: 'agent_proxy', + }, + }), + }); + const prohibited = [ + 'forwarded', + 'proxy-connection', + 'proxy-a', + 'proxy-b', + 'x-forwarded-a', + 'x-forwarded-b', + 'x-internal-a', + 'x-internal-b', + 'x-kilo-a', + 'x-kilo-b', + 'x-kilocode-a', + 'x-kilocode-b', + 'x-real-ip', + ]; + const upstream = vi.fn().mockResolvedValue(new Response('ok')); + vi.stubGlobal('fetch', upstream); + try { + const response = await fetchWorker( + new Request( + 'https://worker.test/api/runtime-credential-proxy/provider/api/openrouter/models', + { + headers: { + ...Object.fromEntries(prohibited.map(name => [name, 'untrusted'])), + Authorization: `Bearer ${await handle()}`, + 'X-Kilocode-OrganizationId': 'attacker-org', + 'X-Client-Request-Id': 'request_proxy', + }, + } + ), + env + ); + expect(response.status).toBe(200); + expect(upstream).toHaveBeenCalledOnce(); + const forwarded = upstream.mock.calls[0][0] as Request; + for (const name of prohibited) expect(forwarded.headers.get(name)).toBeNull(); + expect(forwarded.headers.get('authorization')).toBe( + 'Bearer https://provider.example.test/api/openrouter:backing-token' + ); + expect(forwarded.headers.get('x-kilocode-organizationid')).toBe('org_proxy'); + expect(forwarded.headers.get('x-client-request-id')).toBe('request_proxy'); + } finally { + vi.unstubAllGlobals(); + } + }); + + it('removes generic caller credential headers before injecting runtime credentials', async () => { + const env = createEnv(); + env.CLOUD_AGENT_SESSION.get.mockReturnValue({ + resolveRuntimeCredentialProxyGrant: vi.fn().mockResolvedValue({ + token: 'https://provider.example.test/api/openrouter:backing-token', + runtimeAuthorization: { + userId: 'usr_proxy', + authorizationId: '11111111-1111-4111-8111-111111111111', + resourceId: 'agent_proxy', + }, + }), + }); + const callerCredentials = ['x-api-key', 'api-key', 'x-auth-token']; + const upstream = vi.fn().mockResolvedValue(new Response('ok')); + vi.stubGlobal('fetch', upstream); + try { + const response = await fetchWorker( + new Request( + 'https://worker.test/api/runtime-credential-proxy/provider/api/openrouter/models', + { + headers: { + ...Object.fromEntries(callerCredentials.map(name => [name, 'caller-credential'])), + Authorization: `Bearer ${await handle()}`, + }, + } + ), + env + ); + expect(response.status).toBe(200); + const forwarded = upstream.mock.calls[0][0] as Request; + for (const name of callerCredentials) expect(forwarded.headers.get(name)).toBeNull(); + expect(forwarded.headers.get('authorization')).toBe( + 'Bearer https://provider.example.test/api/openrouter:backing-token' + ); + } finally { + vi.unstubAllGlobals(); + } + }); + + it('removes unsafe upstream response headers while preserving redirects and streaming', async () => { + const env = createEnv(); + env.CLOUD_AGENT_SESSION.get.mockReturnValue({ + resolveRuntimeCredentialProxyGrant: vi.fn().mockResolvedValue({ + token: 'https://provider.example.test/api/openrouter:backing-token', + runtimeAuthorization: { + userId: 'usr_proxy', + authorizationId: '11111111-1111-4111-8111-111111111111', + resourceId: 'agent_proxy', + }, + }), + }); + const unsafeHeaders = [ + 'connection', + 'proxy-connection', + 'keep-alive', + 'proxy-authenticate', + 'proxy-authorization', + 'te', + 'trailer', + 'transfer-encoding', + 'upgrade', + 'set-cookie', + ]; + const stream = new ReadableStream({ + start(controller) { + controller.enqueue(new TextEncoder().encode('stream-body')); + controller.close(); + }, + }); + const upstream = vi.fn().mockResolvedValue( + new Response(stream, { + status: 307, + statusText: 'Temporary Redirect', + headers: { + ...Object.fromEntries(unsafeHeaders.map(name => [name, 'unsafe'])), + 'content-type': 'text/event-stream', + 'cache-control': 'no-cache', + location: 'https://other.test/continue', + [RUNTIME_PROXY_ATTESTATION_HEADER]: 'upstream-proof', + }, + }) + ); + vi.stubGlobal('fetch', upstream); + try { + const response = await fetchWorker( + new Request( + 'https://worker.test/api/runtime-credential-proxy/provider/api/openrouter/models', + { headers: { Authorization: `Bearer ${await handle()}` } } + ), + env + ); + expect(response.status).toBe(307); + expect(response.statusText).toBe('Temporary Redirect'); + expect(response.headers.get('location')).toBe('https://other.test/continue'); + expect(response.headers.get('content-type')).toBe('text/event-stream'); + expect(response.headers.get('cache-control')).toBe('no-cache'); + for (const name of unsafeHeaders) expect(response.headers.get(name)).toBeNull(); + expect(response.headers.get(RUNTIME_PROXY_ATTESTATION_HEADER)).toBeNull(); + await expect(response.text()).resolves.toBe('stream-body'); + expect(upstream.mock.calls[0][1]).toEqual({ redirect: 'manual' }); + } finally { + vi.unstubAllGlobals(); + } + }); + + it('forwards packaged CLI inference requests to the production gateway', async () => { + const env = createEnv(); + const resolve = vi.fn().mockResolvedValue({ + token: jwt.sign({ exp: 4_000_000_000 }, secret), + runtimeAuthorization: { + userId: 'usr_proxy', + authorizationId: '11111111-1111-4111-8111-111111111111', + resourceId: 'agent_proxy', + }, + }); + env.CLOUD_AGENT_SESSION.get.mockReturnValue({ resolveRuntimeCredentialProxyGrant: resolve }); + const upstream = vi.fn(async (request: Request) => { + expect(request.method).toBe('POST'); + expect(request.url).toBe('https://api.kilo.ai/api/gateway/chat/completions?stream=true'); + return new Response('ok'); + }); + vi.stubGlobal('fetch', upstream); + + const response = await fetchWorker( + new Request( + 'https://worker.test/api/runtime-credential-proxy/provider/api/openrouter/chat/completions?stream=true', + { method: 'POST', headers: { Authorization: `Bearer ${await handle()}` } } + ), + env + ); + + expect(response.status).toBe(200); + expect(upstream).toHaveBeenCalledOnce(); + vi.unstubAllGlobals(); + }); + + it('validates the body-bound ingest identity before fetching upstream', async () => { + const env = createEnv(); + const resolve = vi.fn().mockResolvedValue({ + token: jwt.sign({ exp: 4_000_000_000 }, secret), + runtimeAuthorization: { + userId: 'usr_proxy', + authorizationId: '11111111-1111-4111-8111-111111111111', + resourceId: 'agent_proxy', + }, + }); + env.CLOUD_AGENT_SESSION.get.mockReturnValue({ resolveRuntimeCredentialProxyGrant: resolve }); + const upstream = vi.fn(); + vi.stubGlobal('fetch', upstream); + const response = await fetchWorker( + new Request('https://worker.test/api/runtime-credential-proxy/ingest/api/session', { + method: 'POST', + headers: { Authorization: `Bearer ${await handle()}`, 'Content-Type': 'application/json' }, + body: '{"sessionId":"other"}', + }), + env + ); + expect(response.status).toBe(404); + expect(upstream).not.toHaveBeenCalled(); + vi.unstubAllGlobals(); + }); +}); + describe('server /kilo facade route', () => { for (const path of ['/kilo', '/kilo/event']) { it(`returns 401 before facade dispatch when auth is missing for ${path}`, async () => { @@ -887,7 +1457,9 @@ describe('server raw global feed route', () => { const validateKiloGlobalFeedProducer = vi.fn(async () => ({ success: true as const })); const facadeFetch = vi.fn().mockResolvedValue(new Response('accepted', { status: 200 })); env.CLOUD_AGENT_SESSION.idFromName.mockReturnValue('session-do-id'); - env.CLOUD_AGENT_SESSION.get.mockReturnValue({ validateKiloGlobalFeedProducer }); + env.CLOUD_AGENT_SESSION.get.mockReturnValue({ + validateKiloGlobalFeedProducer, + }); env.USER_KILO_FACADE.idFromName.mockReturnValue('facade-id'); env.USER_KILO_FACADE.get.mockReturnValue({ fetch: facadeFetch }); const token = signKiloToken('usr_feed'); @@ -910,7 +1482,9 @@ describe('server raw global feed route', () => { const env = createEnv(); const validateKiloGlobalFeedProducer = vi.fn(async () => ({ success: true as const })); env.CLOUD_AGENT_SESSION.idFromName.mockReturnValue('session-do-id'); - env.CLOUD_AGENT_SESSION.get.mockReturnValue({ validateKiloGlobalFeedProducer }); + env.CLOUD_AGENT_SESSION.get.mockReturnValue({ + validateKiloGlobalFeedProducer, + }); const facadeFetch = vi.fn<(request: Request) => Promise>( async () => new Response('accepted', { status: 200 }) ); @@ -1131,6 +1705,9 @@ describe('server wrapper ingest route', () => { it('keeps current session ownership enforcement for an audience-less legacy raw Kilo JWT', async () => { const env = createEnv(); + const doFetch = vi.fn(); + const getRuntimeAuthorizationStatus = vi.fn().mockResolvedValue('legacy'); + env.CLOUD_AGENT_SESSION.get.mockReturnValue({ getRuntimeAuthorizationStatus, fetch: doFetch }); const token = signKiloToken('usr_feed'); requireCurrentSessionAccessMock.mockRejectedValue( Object.assign(new Error('Session access denied'), { code: 'FORBIDDEN' }) @@ -1149,14 +1726,18 @@ describe('server wrapper ingest route', () => { kiloUserId: 'usr_feed', cloudAgentSessionId: 'agent_live', }); - expect(env.CLOUD_AGENT_SESSION.idFromName).not.toHaveBeenCalled(); + expect(getRuntimeAuthorizationStatus).toHaveBeenCalledOnce(); + expect(doFetch).not.toHaveBeenCalled(); }); it('accepts a valid wrapper dispatch ticket and forwards to the session Durable Object', async () => { const env = createEnv(); 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 }); + env.CLOUD_AGENT_SESSION.get.mockReturnValue({ + getRuntimeAuthorizationStatus: vi.fn().mockResolvedValue('legacy'), + fetch: doFetch, + }); const ticket = signWrapperDispatchTicket(); const response = await fetchWorker( @@ -1179,7 +1760,10 @@ describe('server wrapper ingest route', () => { const env = createEnv(); 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 }); + env.CLOUD_AGENT_SESSION.get.mockReturnValue({ + getRuntimeAuthorizationStatus: vi.fn().mockResolvedValue('legacy'), + fetch: doFetch, + }); const token = signKiloToken('usr_feed'); const response = await fetchWorker( @@ -1231,6 +1815,80 @@ describe('server wrapper ingest route', () => { }); describe('server wrapper log upload route', () => { + it('rejects a legacy wrapper token before ingest reaches a runtime-authorized session', async () => { + const env = createEnv(); + const fetch = vi.fn(); + env.CLOUD_AGENT_SESSION.idFromName.mockReturnValue('session-do-id'); + env.CLOUD_AGENT_SESSION.get.mockReturnValue({ + getRuntimeAuthorizationStatus: vi.fn().mockResolvedValue('active'), + fetch, + }); + + const response = await fetchWorker( + new Request('http://worker.test/sessions/usr_feed/agent_live/ingest', { + headers: { + Upgrade: 'websocket', + Authorization: `Bearer ${signKiloToken('usr_feed')}`, + }, + }), + env + ); + + expect(response.status).toBe(401); + expect(fetch).not.toHaveBeenCalled(); + expect(requireCurrentSessionAccessMock).not.toHaveBeenCalled(); + }); + + it('keeps legacy wrapper tokens working for fenced global feed dispatch', async () => { + const env = createEnv(); + const validateKiloGlobalFeedProducer = vi.fn().mockResolvedValue({ success: true }); + const facadeFetch = vi.fn().mockResolvedValue(new Response(null, { status: 200 })); + env.CLOUD_AGENT_SESSION.idFromName.mockReturnValue('session-do-id'); + env.CLOUD_AGENT_SESSION.get.mockReturnValue({ + validateKiloGlobalFeedProducer, + }); + env.USER_KILO_FACADE.idFromName.mockReturnValue('facade-id'); + env.USER_KILO_FACADE.get.mockReturnValue({ fetch: facadeFetch }); + + const response = await fetchWorker( + new Request( + 'http://worker.test/sessions/usr_feed/agent_live/kilo-global-ingest?kiloSessionId=ses_12345678901234567890123456&wrapperRunId=wr_1&wrapperGeneration=2&wrapperConnectionId=conn_1', + { + headers: { + Upgrade: 'websocket', + Authorization: `Bearer ${signKiloToken('usr_feed')}`, + }, + } + ), + env + ); + + expect(response.status).toBe(200); + expect(validateKiloGlobalFeedProducer).toHaveBeenCalledOnce(); + expect(facadeFetch).toHaveBeenCalledOnce(); + }); + + it('rejects a legacy wrapper token before writing a runtime-authorized log archive', async () => { + const env = Object.assign(createEnv(), { R2_BUCKET: { put: vi.fn() } }); + env.CLOUD_AGENT_SESSION.idFromName.mockReturnValue('session-do-id'); + env.CLOUD_AGENT_SESSION.get.mockReturnValue({ + getRuntimeAuthorizationStatus: vi.fn().mockResolvedValue('revoked'), + }); + + const response = await fetchWorker( + new Request('http://worker.test/sessions/usr_feed/agent_live/logs/session/logs.tar.gz', { + method: 'PUT', + headers: { Authorization: `Bearer ${signKiloToken('usr_feed')}` }, + body: 'archive', + }), + env + ); + + expect(response.status).toBe(401); + expect(env.R2_BUCKET.put).not.toHaveBeenCalled(); + expect(requireCurrentSessionAccessMock).not.toHaveBeenCalled(); + }); + it('does not accept legacy raw archives for control-plane sessions', async () => { const env = Object.assign(createEnv(), { R2_BUCKET: { put: vi.fn() } }); const response = await fetchWorker( @@ -1697,7 +2355,10 @@ describe('server /sandbox-terminal', () => { const sessionResponse = new Response('wrapper bridged', { status: 200 }); const sessionFetch = vi.fn().mockResolvedValue(sessionResponse); env.SANDBOX_SESSION.idFromName.mockReturnValue('sandbox-session-do-id'); - env.SANDBOX_SESSION.get.mockReturnValue({ fetch: sessionFetch }); + env.SANDBOX_SESSION.get.mockReturnValue({ + fetch: sessionFetch, + isRuntimeAuthorizationRecoveryInProgress: vi.fn().mockResolvedValue(false), + }); const request = new Request( `http://worker.test/sandbox-terminal/${encodeURIComponent(ownerId)}/${sessionId}/pty_123?ticket=browser-secret&ptyId=attacker&role=browser`, { @@ -1744,6 +2405,27 @@ describe('server /sandbox-terminal', () => { expect(forwarded.headers.get('x-internal-role')).toBeNull(); expect(forwarded.headers.get('x-forwarded-user')).toBeNull(); }); + + it('rejects a valid producer WebSocket before forwarding during runtime authorization recovery', async () => { + const env = createEnv(); + const sessionFetch = vi.fn(); + env.SANDBOX_SESSION.idFromName.mockReturnValue('sandbox-session-do-id'); + env.SANDBOX_SESSION.get.mockReturnValue({ + fetch: sessionFetch, + isRuntimeAuthorizationRecoveryInProgress: vi.fn().mockResolvedValue(true), + }); + + const response = await fetchWorker( + new Request(`http://worker.test/sandbox-terminal/user-1/${sessionId}/pty_123`, { + headers: { Upgrade: 'websocket', Authorization: 'Bearer producer-capability' }, + }), + env + ); + + expect(response.status).toBe(503); + await expect(response.text()).resolves.toBe('Runtime authorization recovery in progress'); + expect(sessionFetch).not.toHaveBeenCalled(); + }); }); describe('server control log routes', () => { diff --git a/services/cloud-agent-next/src/server.ts b/services/cloud-agent-next/src/server.ts index fb5620859d..77bdc8756e 100644 --- a/services/cloud-agent-next/src/server.ts +++ b/services/cloud-agent-next/src/server.ts @@ -41,6 +41,7 @@ import { import { getSandboxControlStub, isSandboxControlId } from './sandbox-control/stub.js'; import { getSandboxSessionStub, resolveSessionStub } from './sandbox-session/session-stub.js'; import { sessionPlaneFromId } from './session-plane.js'; +import { withDORetry } from './utils/do-retry.js'; import { generateSandboxCredential, hashSandboxCredential, @@ -48,6 +49,19 @@ import { } from './sandbox-control/credential.js'; import { PtyIdSchema, sessionIdSchema } from './router/schemas.js'; import { registerControlLogRoutes } from './sandbox-control/log-routes.js'; +import { + runtimeCredentialProxyFacadeBaseUrl, + runtimeCredentialProxyUpstream, + type RuntimeProxyHandleClaims, + verifyRuntimeCredentialProxyHandle, +} from './runtime-credential-proxy.js'; +import { deriveKiloSandboxTargets } from './kilo/kilo-targets.js'; +import { inferRuntimeCredentialProxyRoute } from './kilo/runtime-credential-proxy-routes.js'; +import { + issueRuntimeProxyAttestation, + RUNTIME_PROXY_ATTESTATION_HEADER, + type RuntimeProxyAttestationAudience, +} from '@kilocode/worker-utils/runtime-proxy-attestation'; const app = new Hono(); @@ -181,10 +195,16 @@ async function handleTerminalWebSocket(request: Request, env: Env): Promise, next: Next) => { }); }); +// Kilo 7.4.20 and current releases both append their own `/api/...` route to +// configured targets. Only a verified opaque handle may turn that otherwise +// ordinary Worker path into a facade request. +app.use('*', async (c: Context, next: Next) => { + const facadeBase = c.env.WORKER_URL + ? runtimeCredentialProxyFacadeBaseUrl(c.env.WORKER_URL) + : null; + if (!facadeBase) return next(); + const facadePath = new URL(facadeBase).pathname.replace(/\/+$/, ''); + const requestPath = new URL(c.req.url).pathname; + const path = facadePath + ? requestPath.startsWith(`${facadePath}/`) + ? requestPath.slice(facadePath.length) + : null + : requestPath; + if (!path) return next(); + // This route remains available for already-issued configurations. A prefixed + // deployment never supported it because Hono registers the legacy route at + // the Worker root; root deployments retain the transition behavior. + if (path.startsWith('/api/runtime-credential-proxy/')) return next(); + const handle = runtimeProxyAuthorization(c.req.raw); + if (!handle) return next(); + const claims = await verifyRuntimeCredentialProxyHandle(c.env, handle); + if (!claims) return next(); + const route = inferRuntimeCredentialProxyRoute(path); + if (!route) return c.text('Not found', 404); + return forwardRuntimeCredentialProxy(c, handle, claims, route, path); +}); + app.get('/health', (c: Context) => { return c.json({ status: 'ok', @@ -214,6 +263,10 @@ app.get('/health', (c: Context) => { }); }); +// Handle authentication is bearer-only; the URL contains neither a credential +// nor an upstream authority. Register before the broad facade routes. +app.all('/api/runtime-credential-proxy/:route/*', routeRuntimeCredentialProxy); + function requireInternalApi(c: Context): Response | null { if (!c.env.INTERNAL_API_SECRET) { return c.text('Internal API secret not configured', 500); @@ -286,6 +339,9 @@ app.get('/sandbox-terminal/:ownerId/:sessionId/:ptyId', async (c: Context { + if (!request.body) return null; + const reader = request.body.getReader(); + const decoder = new TextDecoder(); + let size = 0; + let value = ''; + try { + while (true) { + const chunk = await reader.read(); + if (chunk.done) return value + decoder.decode(); + const bytes = new Uint8Array(chunk.value); + size += bytes.byteLength; + if (size > maximumBytes) { + await reader.cancel(); + return null; + } + value += decoder.decode(bytes, { stream: true }); + } + } catch { + return null; + } +} + +async function routeRuntimeCredentialProxy(c: Context): Promise { + const handle = runtimeProxyAuthorization(c.req.raw); + if (!handle) return c.text('Unauthorized', 401); + const claims = await verifyRuntimeCredentialProxyHandle(c.env, handle); + if (!claims) return c.text('Unauthorized', 401); + const route = c.req.param('route'); + const prefix = `/api/runtime-credential-proxy/${route}/`; + const requestPath = new URL(c.req.url).pathname; + if (!requestPath.startsWith(prefix)) return c.text('Not found', 404); + const path = `/${requestPath.slice(prefix.length).replace(/^\/+/, '')}`; + if (route !== 'backend' && route !== 'provider' && route !== 'ingest') + return c.text('Not found', 404); + return forwardRuntimeCredentialProxy(c, handle, claims, route, path); +} + +async function forwardRuntimeCredentialProxy( + c: Context, + handle: string, + claims: RuntimeProxyHandleClaims, + route: 'backend' | 'provider' | 'ingest', + path: string +): Promise { + if (!claims) return c.text('Unauthorized', 401); + let bodyText: string | undefined; + if (route === 'ingest' && path === '/api/session' && c.req.method === 'POST') { + bodyText = (await readBoundedBody(c.req.raw, 8192)) ?? undefined; + if (bodyText === undefined) return c.text('Not found', 404); + } + let credential: { + token: string; + organizationId?: string; + runtimeAuthorization: { userId: string; authorizationId: string; resourceId: string }; + } | null; + try { + credential = await withDORetry( + () => resolveSessionStub(c.env, claims.userId, claims.sessionId), + session => session.resolveRuntimeCredentialProxyGrant(handle), + 'resolveRuntimeCredentialProxyGrant' + ); + } catch { + return c.text('Credential unavailable', 503); + } + if (!credential) return c.text('Unauthorized', 401); + const targets = deriveKiloSandboxTargets(c.env, credential.token, { requireHttps: true }); + if (!targets.success) return c.text('Not found', 404); + const upstream = runtimeCredentialProxyUpstream( + targets.targets, + route, + c.req.method, + path, + new URL(c.req.url).search, + claims.kiloSessionId, + credential.organizationId, + c.req.header('content-type'), + bodyText + ); + if (!upstream) return c.text('Not found', 404); + try { + // The route allowlist is resolved above before a proof is issued. + const audience: RuntimeProxyAttestationAudience = + route === 'backend' ? 'kilo-api' : route === 'provider' ? 'kilo-gateway' : 'session-ingest'; + const proof = await issueRuntimeProxyAttestation({ + secret: await resolveSecret(c.env.NEXTAUTH_SECRET).then(value => { + if (!value) throw new Error('Authentication unavailable'); + return value; + }), + audience, + bearer: credential.token, + ...credential.runtimeAuthorization, + }); + const headers = runtimeProxyHeaders(c.req.raw, credential.token, credential.organizationId); + headers.set(RUNTIME_PROXY_ATTESTATION_HEADER, proof); + const response = await fetch( + createSanitizedForwardRequest(c.req.raw, upstream, headers, bodyText), + { redirect: 'manual' } + ); + return sanitizeRuntimeProxyResponse(response); + } catch { + return c.text('Upstream unavailable', 502); + } +} + function parseOptionalWrapperGeneration(raw: string | null): number | undefined { if (raw === null) return undefined; const parsed = Number(raw); @@ -362,6 +595,23 @@ function stripPublicCredentialHeaders(headers: Headers): Headers { return sanitized; } +async function rejectLegacyWrapperTokenForRuntimeGrant( + env: Env, + claims: WrapperAuthClaims, + userId: string, + sessionId: string +): Promise { + if (claims.type !== 'legacy_kilo_token') return null; + const status = await withDORetry( + () => resolveSessionStub(env, userId, sessionId), + stub => stub.getRuntimeAuthorizationStatus(), + 'getRuntimeAuthorizationStatus' + ); + return status === 'legacy' + ? null + : new Response('Legacy wrapper token is not authorized', { status: 401 }); +} + async function routeToUserKiloFacade( c: Context, userId: string, @@ -611,6 +861,14 @@ app.all('/sessions/:userId/:sessionId/ingest', async (c: Context) = return c.text('Token does not match session user', 403); } + const legacyRejection = await rejectLegacyWrapperTokenForRuntimeGrant( + c.env, + authResult.claims, + userId, + sessionId + ); + if (legacyRejection) return legacyRejection; + const url = new URL(c.req.url); const wrapperGenerationParam = url.searchParams.get('wrapperGeneration'); const wrapperGeneration = parseOptionalWrapperGeneration(wrapperGenerationParam); @@ -688,6 +946,14 @@ app.put( return c.text('Token does not match session user', 403); } + const legacyRejection = await rejectLegacyWrapperTokenForRuntimeGrant( + c.env, + authResult.claims, + userId, + sessionId + ); + if (legacyRejection) return legacyRejection; + const kiloSessionId = new URL(c.req.url).searchParams.get('kiloSessionId'); if (!kiloSessionId && authResult.claims.type === 'wrapper_dispatch_ticket') { return c.text('Missing kiloSessionId parameter', 400); diff --git a/services/cloud-agent-next/src/session-service.test.ts b/services/cloud-agent-next/src/session-service.test.ts index 004483e548..7035501fe5 100644 --- a/services/cloud-agent-next/src/session-service.test.ts +++ b/services/cloud-agent-next/src/session-service.test.ts @@ -4,6 +4,7 @@ import type * as DevContainerModule from './kilo/devcontainer.js'; import type * as GitTokenServiceClientModule from './services/git-token-service-client.js'; import { validateWrapperDispatchTicket } from './auth.js'; import { deriveKiloSandboxTargets } from './kilo/kilo-targets.js'; +import jwt from 'jsonwebtoken'; import { ExecutionError } from './execution/errors.js'; import { createPendingSessionMessage, @@ -2459,6 +2460,107 @@ describe('SessionService.buildWrapperSessionReadyAndPromptRequests', () => { expect(result.readyRequest.workspace.workspacePath).toBe(workspacePath); }); + it('uses a fenced Worker proxy handle for modern authorization without exposing its backing token', async () => { + const service = new SessionService(); + const env = createEnv(); + env.WORKER_URL = 'https://cloud-agent.example.com'; + const backingToken = jwt.sign( + { + runtimeAuthorization: { id: '11111111-1111-4111-8111-111111111111' }, + exp: Math.floor(Date.now() / 1000) + 3600, + }, + 'backing-secret' + ); + const issueRuntimeCredentialProxyGrant = vi.fn().mockResolvedValue('stable-proxy-handle'); + env.CLOUD_AGENT_SESSION.get = vi.fn(() => ({ issueRuntimeCredentialProxyGrant })) as never; + const metadata = createMetadata({ kilocodeToken: backingToken }); + + const result = await service.buildWrapperSessionReadyAndPromptRequests({ + env, + plan: { + scope: { sessionId: 'agent_test', userId: 'user_test' }, + turn: { + type: 'prompt', + messageId: 'msg_018f1e2d3c4bModernProxyAAAA', + prompt: 'Do the work', + }, + agent: { mode: 'code', model: 'test-model' }, + workspace: { sandboxId: 'ses-abcdef', metadata }, + wrapper: { + fence: { + wrapperRunId: 'wr_modern', + wrapperGeneration: 3, + wrapperConnectionId: 'conn_modern', + }, + }, + } satisfies FencedWrapperDispatchRequest, + }); + + expect(issueRuntimeCredentialProxyGrant).toHaveBeenCalledWith({ + wrapperRunId: 'wr_modern', + wrapperGeneration: 3, + wrapperConnectionId: 'conn_modern', + }); + expect(result.readyRequest.runtimeCredentialProxy).toEqual({ + handle: 'stable-proxy-handle', + targets: { + backendBaseUrl: 'https://cloud-agent.example.com', + providerBaseUrl: 'https://cloud-agent.example.com', + sessionIngestBaseUrl: 'https://cloud-agent.example.com', + }, + }); + expect(result.readyRequest.materialized.env.KILOCODE_TOKEN).toBe('stable-proxy-handle'); + expect(JSON.parse(result.readyRequest.materialized.env.KILO_AUTH_CONTENT)).toEqual({ + kilo: { type: 'api', key: 'stable-proxy-handle' }, + }); + expect(JSON.stringify(result.readyRequest)).not.toContain(backingToken); + expect(JSON.stringify(result.readyRequest)).not.toContain('backing-secret'); + expect(tokenMocks.issueCloudAgentGitLabSessionCapability).not.toHaveBeenCalled(); + }); + + it('fails closed when the session cannot issue a modern runtime proxy handle', async () => { + const service = new SessionService(); + const env = createEnv(); + env.WORKER_URL = 'https://cloud-agent.example.com'; + const backingToken = jwt.sign( + { + runtimeAuthorization: { id: '11111111-1111-4111-8111-111111111111' }, + exp: Math.floor(Date.now() / 1000) + 3600, + }, + 'backing-secret' + ); + const issueRuntimeCredentialProxyGrant = vi.fn().mockResolvedValue(null); + env.CLOUD_AGENT_SESSION.get = vi.fn(() => ({ issueRuntimeCredentialProxyGrant })) as never; + const metadata = createMetadata({ kilocodeToken: backingToken }); + + await expect( + service.buildWrapperSessionReadyAndPromptRequests({ + env, + plan: { + scope: { sessionId: 'agent_test', userId: 'user_test' }, + turn: { + type: 'prompt', + messageId: 'msg_018f1e2d3c4bNoProxyHandleAAAA', + prompt: 'Do the work', + }, + agent: { mode: 'code', model: 'test-model' }, + workspace: { sandboxId: 'ses-abcdef', metadata }, + wrapper: { + fence: { + wrapperRunId: 'wr_modern', + wrapperGeneration: 3, + wrapperConnectionId: 'conn_modern', + }, + }, + } satisfies FencedWrapperDispatchRequest, + }) + ).rejects.toMatchObject({ + code: 'INVALID_REQUEST', + message: 'Runtime credential proxy grant is unavailable', + }); + expect(tokenMocks.issueCloudAgentGitLabSessionCapability).not.toHaveBeenCalled(); + }); + it.each([undefined, 1])( 'uses the persisted readable branch in wrapper readiness with preparedAt=%s', async preparedAt => { diff --git a/services/cloud-agent-next/src/session-service.ts b/services/cloud-agent-next/src/session-service.ts index 07d1e1e3f1..6f2be44f4e 100644 --- a/services/cloud-agent-next/src/session-service.ts +++ b/services/cloud-agent-next/src/session-service.ts @@ -26,6 +26,7 @@ import { resolveManagedGitLabToken, } from './services/git-token-service-client.js'; import { deriveKiloSandboxTargets } from './kilo/kilo-targets.js'; +import { runtimeCredentialProxyFacadeBaseUrl } from './runtime-credential-proxy.js'; import { ExecutionError } from './execution/errors.js'; import { checkDiskAndCleanBeforeSetup, @@ -61,6 +62,7 @@ import { } from './persistence/session-metadata.js'; import { withDORetry } from './utils/do-retry.js'; import { resolveSessionStub } from './sandbox-session/session-stub.js'; +import { hasModernRuntimeAuthorization } from './session/runtime-authorization-persistence.js'; import { decryptWithPrivateKey, mergeEnvVarsWithSecrets } from './utils/encryption.js'; import { codeReviewIdFromCallbackTarget, type MCPSecretValue } from './router/schemas.js'; import type { SessionProfileBundle } from './session-profile.js'; @@ -85,6 +87,7 @@ import { type WrapperBootstrapRepoSource, type WrapperCommandRequest, type WrapperPromptRequest, + type WrapperRuntimeCredentialProxyConfig, type WrapperSessionReadyRequest, type WrapperWorkspaceReady, } from './shared/wrapper-bootstrap.js'; @@ -1279,6 +1282,7 @@ export class SessionService { workspacePath, env: opts.env, kiloCapability: opts.kiloCapability, + kiloBackendBaseUrl: opts.kiloBackendBaseUrl, kiloProviderBaseUrl: opts.kiloProviderBaseUrl, kiloSessionIngestBaseUrl: opts.kiloSessionIngestBaseUrl, kilocodeModel: opts.kilocodeModel, @@ -1307,6 +1311,7 @@ export class SessionService { workspacePath, env, kiloCapability, + kiloBackendBaseUrl, kiloProviderBaseUrl, kiloSessionIngestBaseUrl, kilocodeModel, @@ -1639,8 +1644,9 @@ export class SessionService { envVars.KILOCODE_ORGANIZATION_ID = kilocodeOrganizationId; } - if (env.KILOCODE_BACKEND_BASE_URL) { - const sandboxUrl = backendUrlForSandbox(env.KILOCODE_BACKEND_BASE_URL); + if (kiloBackendBaseUrl || env.KILOCODE_BACKEND_BASE_URL) { + const sandboxUrl = + kiloBackendBaseUrl ?? backendUrlForSandbox(env.KILOCODE_BACKEND_BASE_URL ?? ''); envVars.KILOCODE_BACKEND_BASE_URL = sandboxUrl; // Used by kilo server to check user auth to send to ingest envVars.KILO_API_URL = sandboxUrl; @@ -1967,7 +1973,12 @@ export class SessionService { kilocodeContainment: boolean; userToken: string; } - ): Promise<{ capability: string; providerBaseUrl?: string; sessionIngestBaseUrl?: string }> { + ): Promise<{ + capability: string; + backendBaseUrl?: string; + providerBaseUrl?: string; + sessionIngestBaseUrl?: string; + }> { if (params.sandboxId.startsWith('dind-')) { return { capability: params.userToken }; } @@ -2005,6 +2016,7 @@ export class SessionService { // `upstream_not_allowed`, even though the capability itself is valid. return { capability: issued.value.capability, + backendBaseUrl: derivedTargets.targets.backendBaseUrl, providerBaseUrl: derivedTargets.targets.providerBaseUrl, sessionIngestBaseUrl: derivedTargets.targets.sessionIngestBaseUrl, }; @@ -2020,7 +2032,12 @@ export class SessionService { sandboxId: string; userToken: string; } - ): Promise<{ capability: string; providerBaseUrl?: string; sessionIngestBaseUrl?: string }> { + ): Promise<{ + capability: string; + backendBaseUrl?: string; + providerBaseUrl?: string; + sessionIngestBaseUrl?: string; + }> { return this.issueKiloSessionCapability(env, { userId: params.userId, cloudAgentSessionId: params.cloudAgentSessionId, @@ -2047,7 +2064,6 @@ export class SessionService { const { scope, turn, agent, finalization, workspace, wrapper } = plan; const { sessionId, userId, orgId } = scope; const { sandboxId, metadata } = workspace; - if (!metadata.auth.kilocodeToken) { throw ExecutionError.invalidRequest('Missing kilocodeToken in session metadata'); } @@ -2058,17 +2074,54 @@ export class SessionService { if (!nextAuthSecret) { throw ExecutionError.invalidRequest('NEXTAUTH_SECRET is not configured on the worker'); } - const { - capability: kiloCapability, - providerBaseUrl: kiloProviderBaseUrl, - sessionIngestBaseUrl: kiloSessionIngestBaseUrl, - } = await this.resolveKiloCapability(env, metadata, { - userId, - cloudAgentSessionId: sessionId, - kiloSessionId: metadata.auth.kiloSessionId, - sandboxId, - userToken: metadata.auth.kilocodeToken, - }); + const modernRuntimeAuthorization = hasModernRuntimeAuthorization(metadata); + let runtimeCredentialProxy: WrapperRuntimeCredentialProxyConfig | undefined; + let kiloCapability: string; + let kiloBackendBaseUrl: string | undefined; + let kiloProviderBaseUrl: string | undefined; + let kiloSessionIngestBaseUrl: string | undefined; + if (modernRuntimeAuthorization) { + const workerUrl = env.WORKER_URL; + const proxyBaseUrl = workerUrl ? runtimeCredentialProxyFacadeBaseUrl(workerUrl) : null; + const targets = deriveKiloSandboxTargets(env, metadata.auth.kilocodeToken); + if (!proxyBaseUrl || !targets.success) { + throw ExecutionError.invalidRequest( + 'Runtime credential proxy configuration is unavailable' + ); + } + // The session-owned RPC checks the current persisted runtime fence. It is + // deliberately called only after this delivery plan carries every fence field. + const handle = await withDORetry( + () => resolveSessionStub(env, userId, sessionId), + stub => stub.issueRuntimeCredentialProxyGrant(plan.wrapper.fence), + 'issueRuntimeCredentialProxyGrant' + ); + if (!handle) { + throw ExecutionError.invalidRequest('Runtime credential proxy grant is unavailable'); + } + const proxyTargets = { + backendBaseUrl: proxyBaseUrl, + providerBaseUrl: proxyBaseUrl, + sessionIngestBaseUrl: proxyBaseUrl, + }; + runtimeCredentialProxy = { handle, targets: proxyTargets }; + kiloCapability = handle; + kiloBackendBaseUrl = proxyTargets.backendBaseUrl; + kiloProviderBaseUrl = proxyTargets.providerBaseUrl; + kiloSessionIngestBaseUrl = proxyTargets.sessionIngestBaseUrl; + } else { + const kiloCredential = await this.resolveKiloCapability(env, metadata, { + userId, + cloudAgentSessionId: sessionId, + kiloSessionId: metadata.auth.kiloSessionId, + sandboxId, + userToken: metadata.auth.kilocodeToken, + }); + kiloCapability = kiloCredential.capability; + kiloBackendBaseUrl = kiloCredential.backendBaseUrl; + kiloProviderBaseUrl = kiloCredential.providerBaseUrl; + kiloSessionIngestBaseUrl = kiloCredential.sessionIngestBaseUrl; + } const devcontainerRequested = metadata.workspace?.devcontainerRequested === true || metadata.devcontainer !== undefined; @@ -2118,6 +2171,7 @@ export class SessionService { workspacePath, env, kiloCapability, + kiloBackendBaseUrl, kiloProviderBaseUrl, kiloSessionIngestBaseUrl, kilocodeModel: agent.model, @@ -2200,6 +2254,7 @@ export class SessionService { requireSnapshot: metadata.clone !== undefined, }, ...(repo ? { repo } : {}), + ...(runtimeCredentialProxy ? { runtimeCredentialProxy } : {}), ...(devcontainerRequested ? { devcontainer: { @@ -2331,6 +2386,7 @@ export class SessionService { } const { capability: kiloCapability, + backendBaseUrl: kiloBackendBaseUrl, providerBaseUrl: kiloProviderBaseUrl, sessionIngestBaseUrl: kiloSessionIngestBaseUrl, } = await this.resolveKiloCapability(env, metadata, { @@ -2401,6 +2457,7 @@ export class SessionService { context, env, kiloCapability, + kiloBackendBaseUrl, kiloProviderBaseUrl, kiloSessionIngestBaseUrl, kilocodeModel: options.kilocodeModel, @@ -2810,9 +2867,11 @@ export class SessionService { options: RestoreRuntimeOptions, restoreTokenFilePath: string | undefined ): Record { - const backendUrl = options.env.KILOCODE_BACKEND_BASE_URL - ? backendUrlForSandbox(options.env.KILOCODE_BACKEND_BASE_URL) - : undefined; + const backendUrl = + options.kiloBackendBaseUrl ?? + (options.env.KILOCODE_BACKEND_BASE_URL + ? backendUrlForSandbox(options.env.KILOCODE_BACKEND_BASE_URL) + : undefined); return { KILOCODE_TOKEN_FILE: restoreTokenFilePath, KILO_SESSION_INGEST_URL: @@ -3033,6 +3092,7 @@ export type GetOrCreateSessionOptions = { context: SessionContext; env: PersistenceEnv; kiloCapability: string; + kiloBackendBaseUrl?: string; kiloProviderBaseUrl?: string; kiloSessionIngestBaseUrl?: string; kilocodeModel?: string; @@ -3050,6 +3110,7 @@ type RestoreRuntimeOptions = { dockerEnv?: Record; env: PersistenceEnv; kiloCapability: string; + kiloBackendBaseUrl?: string; kiloSessionIngestBaseUrl?: string; runtimeEnv: Record; sessionHome: string; @@ -3061,6 +3122,7 @@ type GetSaferEnvVarsOptions = { workspacePath: string; env: PersistenceEnv; kiloCapability: string; + kiloBackendBaseUrl?: string; kiloProviderBaseUrl?: string; kiloSessionIngestBaseUrl?: string; kilocodeModel?: string; diff --git a/services/cloud-agent-next/src/session/agent-runtime.ts b/services/cloud-agent-next/src/session/agent-runtime.ts index a3ca72d9b4..9ab6afe235 100644 --- a/services/cloud-agent-next/src/session/agent-runtime.ts +++ b/services/cloud-agent-next/src/session/agent-runtime.ts @@ -23,6 +23,7 @@ import type { WrapperCommand } from '../shared/protocol.js'; import type { Env as WorkerEnv } from '../types.js'; import { resolveSessionStub } from '../sandbox-session/session-stub.js'; import { WrapperCleanupBlockedError } from './wrapper-cleanup-blocked-error.js'; +import { RUNTIME_AUTHORIZATION_RECOVERY_KEY } from './runtime-authorization-persistence.js'; import { allocateWrapperRuntimeState, clearAllocatedWrapperRuntimeState, @@ -318,6 +319,13 @@ export function createAgentRuntime(dependencies: AgentRuntimeDependencies): Agen plan: MessageDeliveryRequest, hooks: AgentRuntimeSendHooks = {} ): Promise { + if (await storage.get(RUNTIME_AUTHORIZATION_RECOVERY_KEY)) { + return { + success: false, + code: 'WRAPPER_FINALIZING', + error: 'Runtime authorization recovery is in progress', + }; + } if (canUseSandboxRuntime && !(await canUseSandboxRuntime())) { return { success: false, code: 'INTERNAL', error: 'Session deletion is in progress' }; } diff --git a/services/cloud-agent-next/src/session/queue-message.test.ts b/services/cloud-agent-next/src/session/queue-message.test.ts index 9a348c73c3..9d98836840 100644 --- a/services/cloud-agent-next/src/session/queue-message.test.ts +++ b/services/cloud-agent-next/src/session/queue-message.test.ts @@ -3,6 +3,7 @@ import { TRPCError } from '@trpc/server'; import { preflightAndQueuePromptMessage, + preflightRuntimeAuthorizationRecovery, queueMessage, type QueueMessageInput, } from './queue-message.js'; @@ -81,6 +82,33 @@ describe('preflightAndQueuePromptMessage', () => { expect(idFromName).toHaveBeenCalledWith('user_abc:agent_existing'); expect(getSandboxSession).not.toHaveBeenCalled(); }); + + it('runs foreground runtime authorization recovery once before normal admission', async () => { + const { stub, admitSubmittedMessage } = makeDoStub({ + success: true, + outcome: 'queued', + messageId: 'msg_018f1e2d3c4bAbCdEfGhIjKlMn', + compatibilityDelivery: 'queued', + }); + const getRuntimeAuthorizationRecoveryState = vi.fn().mockResolvedValue({ state: 'active' }); + Object.assign(stub, { getRuntimeAuthorizationRecoveryState }); + + await preflightAndQueuePromptMessage( + { + cloudAgentSessionId: 'agent_existing', + turn: { type: 'prompt', id: 'msg_018f1e2d3c4bAbCdEfGhIjKlMn', prompt: 'follow up' }, + }, + { + env: makeEnv(stub) as Env, + userId: 'user_abc', + authToken: 'eyJhbGciOiJIUzI1NiJ9.eyJhdWQiOiJraWxvIn0.signature', + }, + 'send' + ); + + expect(getRuntimeAuthorizationRecoveryState).toHaveBeenCalledOnce(); + expect(admitSubmittedMessage).toHaveBeenCalledOnce(); + }); }); describe('queueMessage', () => { @@ -459,3 +487,37 @@ describe('queueMessage', () => { } }); }); + +describe('preflightRuntimeAuthorizationRecovery', () => { + const currentAuthToken = 'eyJhbGciOiJIUzI1NiJ9.eyJhdWQiOiJraWxvIn0.signature'; + + it('denies an explicitly revoked modern runtime authorization before admission', async () => { + const admitSubmittedMessage = vi.fn(); + const getRuntimeAuthorizationRecoveryState = vi.fn().mockResolvedValue({ state: 'revoked' }); + + await expect( + preflightRuntimeAuthorizationRecovery('agent_revoked', { + env: makeEnv({ admitSubmittedMessage, getRuntimeAuthorizationRecoveryState }) as Env, + userId: 'user_abc', + authToken: currentAuthToken, + }) + ).rejects.toMatchObject({ code: 'FORBIDDEN', message: 'Runtime authorization denied' }); + + expect(getRuntimeAuthorizationRecoveryState).toHaveBeenCalledOnce(); + expect(admitSubmittedMessage).not.toHaveBeenCalled(); + }); + + it('does nothing for a legacy runtime authorization', async () => { + const getRuntimeAuthorizationRecoveryState = vi.fn().mockResolvedValue({ state: 'legacy' }); + + await expect( + preflightRuntimeAuthorizationRecovery('agent_legacy', { + env: makeEnv({ getRuntimeAuthorizationRecoveryState }) as Env, + userId: 'user_abc', + authToken: currentAuthToken, + }) + ).resolves.toBeUndefined(); + + expect(getRuntimeAuthorizationRecoveryState).toHaveBeenCalledOnce(); + }); +}); diff --git a/services/cloud-agent-next/src/session/queue-message.ts b/services/cloud-agent-next/src/session/queue-message.ts index 770b1d6939..6f4753490d 100644 --- a/services/cloud-agent-next/src/session/queue-message.ts +++ b/services/cloud-agent-next/src/session/queue-message.ts @@ -24,6 +24,13 @@ import { sessionPlaneFromId } from '../session-plane.js'; import { logger } from '../logger.js'; import { preflightExistingPromptModel } from './model-preflight.js'; import { createMessageId } from './message-id.js'; +import { + createRuntimeAuthorization, + sealRuntimeAuthorization, +} from '@kilocode/worker-utils/runtime-authorization'; +import { resolveSecret } from '../auth.js'; +import { fetchSessionMetadata } from '../session-service.js'; +import jwt from 'jsonwebtoken'; /** Retryable error codes that should map to 503 Service Unavailable. */ const RETRYABLE_CODES: readonly RetryableResultCode[] = [ @@ -93,8 +100,73 @@ export type QueueMessageContext = { env: Env; userId: string; botId?: string; + authToken?: string; }; +export async function preflightRuntimeAuthorizationRecovery( + cloudAgentSessionId: string, + ctx: QueueMessageContext +): Promise { + if (!ctx.authToken) return; + const claims = jwt.decode(ctx.authToken); + if ( + !claims || + typeof claims !== 'object' || + !('aud' in claims || 'tokenPurpose' in claims || 'credentialExchange' in claims) + ) { + return; + } + const sessionId = cloudAgentSessionId as SessionId; + const stub = resolveSessionStub(ctx.env, ctx.userId, sessionId); + const state = await withDORetry( + () => stub, + target => target.getRuntimeAuthorizationRecoveryState(), + 'getRuntimeAuthorizationRecoveryState' + ); + if (state.state === 'legacy' || state.state === 'active') return; + if (state.state === 'revoked') { + throw new TRPCError({ code: 'FORBIDDEN', message: 'Runtime authorization denied' }); + } + if (!state.id) return; + const expectedOldId = state.id; + const recoveryId = state.recoveryId ?? crypto.randomUUID(); + const metadata = await fetchSessionMetadata(ctx.env, ctx.userId, cloudAgentSessionId); + if (!metadata || metadata.identity.userId !== ctx.userId) return; + const secret = await resolveSecret(ctx.env.NEXTAUTH_SECRET); + if (!secret) + throw new TRPCError({ code: 'INTERNAL_SERVER_ERROR', message: 'Authentication unavailable' }); + const created = await createRuntimeAuthorization({ + token: ctx.authToken, + secret, + connectionString: ctx.env.HYPERDRIVE.connectionString, + resourceKind: 'cloud-agent-next', + resourceId: metadata.identity.sessionId, + ...(metadata.identity.orgId ? { organizationId: metadata.identity.orgId } : {}), + }); + const runtimeAuthorizationSeal = await sealRuntimeAuthorization(created.authorization, secret); + const result = await withDORetry( + () => stub, + target => + target.recoverExpiredRuntimeAuthorization({ + ownerId: ctx.userId, + expectedOldId, + recoveryId, + runtimeAuthorizationSeal, + runtimeToken: created.token, + }), + 'recoverExpiredRuntimeAuthorization' + ); + if (result.status === 'denied') { + throw new TRPCError({ code: 'FORBIDDEN', message: 'Runtime authorization denied' }); + } + if (result.status === 'busy' || result.status === 'retry') { + throw new TRPCError({ + code: 'CONFLICT', + message: 'Runtime authorization recovery is waiting for the session runtime to become idle', + }); + } +} + /** * Admit a user message via `CloudAgentSession.admitSubmittedMessage`. * @@ -132,6 +204,7 @@ export async function preflightAndAdmitPromptMessage( procedure: string, admit: (input: QueueMessageInput, ctx: QueueMessageContext) => Promise ): Promise { + await preflightRuntimeAuthorizationRecovery(input.cloudAgentSessionId, ctx); if (sessionPlaneFromId(input.cloudAgentSessionId) === 'control') return admit(input, ctx); if (await hasMessageAdmission(input, ctx)) return admit(input, ctx); @@ -151,12 +224,20 @@ export function preflightAndQueuePromptMessage( ctx: QueueMessageContext, procedure: string ): Promise { - return preflightAndAdmitPromptMessage(input, ctx, procedure, queueMessage); + return preflightAndAdmitPromptMessage(input, ctx, procedure, queueMessageAfterRecoveryPreflight); } export async function queueMessage( input: QueueMessageInput, ctx: QueueMessageContext +): Promise { + await preflightRuntimeAuthorizationRecovery(input.cloudAgentSessionId, ctx); + return queueMessageAfterRecoveryPreflight(input, ctx); +} + +async function queueMessageAfterRecoveryPreflight( + input: QueueMessageInput, + ctx: QueueMessageContext ): Promise { const sessionId = input.cloudAgentSessionId as SessionId; const request: SubmittedSessionMessageRequest = { diff --git a/services/cloud-agent-next/src/session/runtime-authorization-persistence.test.ts b/services/cloud-agent-next/src/session/runtime-authorization-persistence.test.ts new file mode 100644 index 0000000000..ed55649f34 --- /dev/null +++ b/services/cloud-agent-next/src/session/runtime-authorization-persistence.test.ts @@ -0,0 +1,266 @@ +import { describe, expect, it, vi } from 'vitest'; +import jwt from 'jsonwebtoken'; +import type { RuntimeAuthorization } from '@kilocode/worker-utils/runtime-authorization-contract'; +import { + RuntimeAuthorizationExpiredError, + RuntimeAuthorizationRevokedError, +} from '@kilocode/worker-utils/runtime-authorization'; +import { + getRuntimeAuthorizationStatus, + getRuntimeAuthorizationRecoveryState, + renewStoredRuntimeAuthorization, +} from './runtime-authorization-persistence.js'; +import type { SessionMetadata } from '../persistence/session-metadata.js'; + +const authorization = ( + id: string, + state: 'active' | 'revoked' = 'active' +): RuntimeAuthorization => ({ + version: 1, + id, + resourceKind: 'cloud-agent-next', + resourceId: 'agent_1', + userId: 'user_1', + authorizationUserId: 'user_1', + organizationId: 'org_1', + issuedAt: '2026-01-01T00:00:00.000Z', + delegationExpiresAt: '2026-01-02T00:00:00.000Z', + state, + bindings: { + userPepperDigest: 'a'.repeat(64), + authorizationPepperDigest: 'b'.repeat(64), + userMembershipId: 'membership_1', + authorizationUserMembershipId: 'membership_1', + }, + source: { admissionSource: 'user' }, +}); + +const metadata = (token?: string): SessionMetadata => + ({ + metadataSchemaVersion: 2, + identity: { sessionId: 'agent_1', userId: 'user_1', orgId: 'org_1' }, + auth: token ? { kilocodeToken: token } : {}, + agent: {}, + workspace: {}, + lifecycle: { version: 1, timestamp: 0 }, + }) as SessionMetadata; + +describe('runtime authorization persistence', () => { + it('leaves legacy tokens unchanged and reports legacy without a private record', async () => { + const stored = metadata('legacy-token'); + const token = await renewStoredRuntimeAuthorization({ + metadata: stored, + getAuthorization: async () => undefined, + putAuthorization: async () => {}, + getMetadata: async () => stored, + putMetadata: async () => {}, + renew: async () => ({ token: 'renewed-token' }), + }); + + expect(token).toBe('legacy-token'); + await expect( + getRuntimeAuthorizationStatus({ metadata: stored, getAuthorization: async () => undefined }) + ).resolves.toBe('legacy'); + }); + + it('does not publish a renewal after a newer authorization replaces the record', async () => { + let record: RuntimeAuthorization = authorization('00000000-0000-4000-8000-000000000001'); + const writes: SessionMetadata[] = []; + await expect( + renewStoredRuntimeAuthorization({ + metadata: metadata(), + getAuthorization: async () => record, + putAuthorization: async value => { + record = value; + }, + getMetadata: async () => metadata(), + putMetadata: async value => { + writes.push(value); + }, + renew: async () => { + record = authorization('00000000-0000-4000-8000-000000000002'); + return { token: 'new-token' }; + }, + now: Date.UTC(2026, 0, 1), + }) + ).rejects.toBeInstanceOf(RuntimeAuthorizationRevokedError); + expect(record.id).toBe('00000000-0000-4000-8000-000000000002'); + expect(writes).toEqual([]); + }); + + it('reuses a matching modern token outside the renewal window', async () => { + const record = authorization('00000000-0000-4000-8000-000000000005'); + const now = Date.UTC(2026, 0, 1); + const token = jwt.sign( + { + runtimeAuthorization: { id: record.id }, + exp: Math.floor((now + 10 * 60_000) / 1000), + }, + 'test-secret', + { noTimestamp: true } + ); + const renew = vi.fn(async () => ({ token: 'unexpected' })); + + await expect( + renewStoredRuntimeAuthorization({ + metadata: metadata(token), + getAuthorization: async () => record, + putAuthorization: async () => {}, + getMetadata: async () => metadata(token), + putMetadata: async () => {}, + renew, + now, + }) + ).resolves.toBe(token); + expect(renew).not.toHaveBeenCalled(); + }); + + it.each([ + { label: 'near expiry', expiresIn: 5 * 60_000 }, + { label: 'wrong authorization', expiresIn: 10 * 60_000, authorizationId: 'other' }, + ])('renews a modern token at $label', async ({ expiresIn, authorizationId }) => { + const record = authorization('00000000-0000-4000-8000-000000000006'); + const now = Date.UTC(2026, 0, 1); + const token = jwt.sign( + { + runtimeAuthorization: { id: authorizationId ?? record.id }, + exp: Math.floor((now + expiresIn) / 1000), + }, + 'test-secret', + { noTimestamp: true } + ); + let stored = metadata(token); + + await expect( + renewStoredRuntimeAuthorization({ + metadata: stored, + getAuthorization: async () => record, + putAuthorization: async () => {}, + getMetadata: async () => stored, + putMetadata: async value => { + stored = value; + }, + renew: async () => ({ token: 'renewed-token' }), + now, + }) + ).resolves.toBe('renewed-token'); + expect(stored.auth.kilocodeToken).toBe('renewed-token'); + }); + + it('rejects a revoked private authorization before renewal', async () => { + const record = authorization('00000000-0000-4000-8000-000000000007', 'revoked'); + const renew = vi.fn(async () => ({ token: 'unexpected' })); + + await expect( + renewStoredRuntimeAuthorization({ + metadata: metadata(), + getAuthorization: async () => record, + putAuthorization: async () => {}, + getMetadata: async () => metadata(), + putMetadata: async () => {}, + renew, + }) + ).rejects.toBeInstanceOf(RuntimeAuthorizationRevokedError); + expect(renew).not.toHaveBeenCalled(); + }); + + it('marks only the matching revoked record and never renews it', async () => { + let record = authorization('00000000-0000-4000-8000-000000000003'); + await expect( + renewStoredRuntimeAuthorization({ + metadata: metadata(), + getAuthorization: async () => record, + putAuthorization: async value => { + record = value; + }, + getMetadata: async () => metadata(), + putMetadata: async () => {}, + renew: async () => { + throw new RuntimeAuthorizationRevokedError(); + }, + now: Date.UTC(2026, 0, 1), + }) + ).rejects.toBeInstanceOf(RuntimeAuthorizationRevokedError); + expect(record.state).toBe('revoked'); + }); + + it('fails closed without revoking the current record at the delegation deadline', async () => { + let record = { + ...authorization('00000000-0000-4000-8000-000000000008'), + delegationExpiresAt: '2026-01-02T00:00:00.000Z', + }; + const renew = vi.fn(async () => ({ token: 'unexpected' })); + + await expect( + renewStoredRuntimeAuthorization({ + metadata: metadata(), + getAuthorization: async () => record, + putAuthorization: async value => { + record = value; + }, + getMetadata: async () => metadata(), + putMetadata: async () => {}, + renew, + now: Date.UTC(2026, 0, 2), + }) + ).rejects.toBeInstanceOf(RuntimeAuthorizationExpiredError); + expect(record.state).toBe('active'); + expect(renew).not.toHaveBeenCalled(); + }); + + it('keeps an active record recoverable when a background renewal observes expiration', async () => { + let record = authorization('00000000-0000-4000-8000-000000000009'); + + await expect( + renewStoredRuntimeAuthorization({ + metadata: metadata(), + getAuthorization: async () => record, + putAuthorization: async value => { + record = value; + }, + getMetadata: async () => metadata(), + putMetadata: async () => {}, + renew: async () => { + throw new RuntimeAuthorizationExpiredError(); + }, + now: Date.UTC(2026, 0, 1), + }) + ).rejects.toBeInstanceOf(RuntimeAuthorizationExpiredError); + expect(record.state).toBe('active'); + }); + + it('distinguishes natural expiry from explicit revocation for foreground recovery', async () => { + const expired = authorization('00000000-0000-4000-8000-000000000010'); + await expect( + getRuntimeAuthorizationRecoveryState({ + metadata: metadata(), + getAuthorization: async () => expired, + now: Date.UTC(2026, 0, 2), + }) + ).resolves.toEqual({ state: 'expired', id: expired.id }); + await expect( + getRuntimeAuthorizationRecoveryState({ + metadata: metadata(), + getAuthorization: async () => ({ ...expired, state: 'revoked' }), + now: Date.UTC(2026, 0, 2), + }) + ).resolves.toEqual({ state: 'revoked' }); + }); + + it('treats modern-token metadata without a valid private record as revoked', async () => { + const token = jwt.sign( + { runtimeAuthorization: { id: '00000000-0000-4000-8000-000000000004' } }, + 'test-secret' + ); + await expect( + renewStoredRuntimeAuthorization({ + metadata: metadata(token), + getAuthorization: async () => undefined, + putAuthorization: async () => {}, + getMetadata: async () => metadata(token), + putMetadata: async () => {}, + renew: async () => ({ token: 'unexpected' }), + }) + ).rejects.toBeInstanceOf(RuntimeAuthorizationRevokedError); + }); +}); diff --git a/services/cloud-agent-next/src/session/runtime-authorization-persistence.ts b/services/cloud-agent-next/src/session/runtime-authorization-persistence.ts new file mode 100644 index 0000000000..39e6ef3ddc --- /dev/null +++ b/services/cloud-agent-next/src/session/runtime-authorization-persistence.ts @@ -0,0 +1,146 @@ +import jwt from 'jsonwebtoken'; +import { + RuntimeAuthorizationExpiredError, + RuntimeAuthorizationRevokedError, +} from '@kilocode/worker-utils/runtime-authorization'; +import { + RuntimeAuthorizationSchema, + type RuntimeAuthorization, +} from '@kilocode/worker-utils/runtime-authorization-contract'; +import { serializeSessionMetadata, type SessionMetadata } from '../persistence/session-metadata.js'; +import { z } from 'zod'; + +export const RUNTIME_AUTHORIZATION_KEY = 'runtime_authorization'; +export const RUNTIME_AUTHORIZATION_RECOVERY_KEY = 'runtime_authorization_recovery'; +export const runtimeAuthorizationRecoveryLockSchema = z + .object({ expectedOldId: z.string().uuid(), recoveryId: z.string().uuid() }) + .strict(); +const RUNTIME_TOKEN_RENEWAL_WINDOW_MS = 5 * 60_000; + +function runtimeAuthorizationId(value: unknown): string | null { + if (typeof value !== 'object' || value === null || !('id' in value)) return null; + return typeof value.id === 'string' ? value.id : null; +} + +export function hasModernRuntimeAuthorization(metadata: SessionMetadata): boolean { + const token = metadata.auth.kilocodeToken; + if (!token) return false; + const decoded = jwt.decode(token); + return ( + typeof decoded === 'object' && + decoded !== null && + 'runtimeAuthorization' in decoded && + typeof decoded.runtimeAuthorization === 'object' && + decoded.runtimeAuthorization !== null + ); +} + +export async function getRuntimeAuthorizationStatus(input: { + metadata: SessionMetadata | null; + getAuthorization: () => Promise; + now?: number; +}): Promise<'legacy' | 'active' | 'revoked'> { + const authorization = RuntimeAuthorizationSchema.safeParse(await input.getAuthorization()); + if (authorization.success) { + return authorization.data.state === 'active' && + Date.parse(authorization.data.delegationExpiresAt) <= (input.now ?? Date.now()) + ? 'revoked' + : authorization.data.state; + } + return input.metadata && hasModernRuntimeAuthorization(input.metadata) ? 'revoked' : 'legacy'; +} + +export type RuntimeAuthorizationRecoveryState = { + state: 'legacy' | 'revoked' | 'active' | 'expired'; + id?: string; +}; + +export async function getRuntimeAuthorizationRecoveryState(input: { + metadata: SessionMetadata | null; + getAuthorization: () => Promise; + now?: number; +}): Promise { + const authorization = RuntimeAuthorizationSchema.safeParse(await input.getAuthorization()); + if (!authorization.success) { + return input.metadata && hasModernRuntimeAuthorization(input.metadata) + ? { state: 'revoked' } + : { state: 'legacy' }; + } + if (authorization.data.state !== 'active') return { state: 'revoked' }; + return Date.parse(authorization.data.delegationExpiresAt) <= (input.now ?? Date.now()) + ? { state: 'expired', id: authorization.data.id } + : { state: 'active', id: authorization.data.id }; +} + +export async function renewStoredRuntimeAuthorization(input: { + metadata: SessionMetadata | null; + getAuthorization: () => Promise; + putAuthorization: (authorization: RuntimeAuthorization) => Promise; + getMetadata: () => Promise; + putMetadata: (metadata: SessionMetadata) => Promise; + renew: (authorization: RuntimeAuthorization) => Promise<{ token: string }>; + now?: number; +}): Promise { + const metadata = input.metadata; + if (!metadata) return null; + const authorization = RuntimeAuthorizationSchema.safeParse(await input.getAuthorization()); + if (!authorization.success) { + if (hasModernRuntimeAuthorization(metadata)) throw new RuntimeAuthorizationRevokedError(); + return metadata.auth.kilocodeToken ?? null; + } + if (authorization.data.state !== 'active') throw new RuntimeAuthorizationRevokedError(); + const now = input.now ?? Date.now(); + const revokeIfCurrent = async () => { + const current = RuntimeAuthorizationSchema.safeParse(await input.getAuthorization()); + if ( + current.success && + current.data.id === authorization.data.id && + current.data.state === 'active' + ) { + await input.putAuthorization({ ...current.data, state: 'revoked' }); + } + }; + if (Date.parse(authorization.data.delegationExpiresAt) <= now) { + throw new RuntimeAuthorizationExpiredError(); + } + const token = metadata.auth.kilocodeToken; + const decoded = token ? jwt.decode(token) : null; + if ( + typeof decoded === 'object' && + decoded !== null && + typeof decoded.exp === 'number' && + decoded.exp * 1000 > now + RUNTIME_TOKEN_RENEWAL_WINDOW_MS && + decoded.exp * 1000 <= Date.parse(authorization.data.delegationExpiresAt) && + runtimeAuthorizationId(decoded.runtimeAuthorization) === authorization.data.id + ) { + return token ?? null; + } + try { + const renewed = await input.renew(authorization.data); + const currentAuthorization = RuntimeAuthorizationSchema.safeParse( + await input.getAuthorization() + ); + const currentMetadata = await input.getMetadata(); + if ( + !currentAuthorization.success || + currentAuthorization.data.id !== authorization.data.id || + currentAuthorization.data.state !== 'active' || + !currentMetadata || + currentMetadata.identity.sessionId !== metadata.identity.sessionId || + currentMetadata.identity.userId !== metadata.identity.userId || + currentMetadata.identity.orgId !== metadata.identity.orgId + ) { + throw new RuntimeAuthorizationRevokedError(); + } + await input.putMetadata( + serializeSessionMetadata({ + ...currentMetadata, + auth: { ...currentMetadata.auth, kilocodeToken: renewed.token }, + }) + ); + return renewed.token; + } catch (error) { + if (error instanceof RuntimeAuthorizationRevokedError) await revokeIfCurrent(); + throw error; + } +} diff --git a/services/cloud-agent-next/src/session/session-message-state.ts b/services/cloud-agent-next/src/session/session-message-state.ts index 3507df9ae1..6e2587f70a 100644 --- a/services/cloud-agent-next/src/session/session-message-state.ts +++ b/services/cloud-agent-next/src/session/session-message-state.ts @@ -618,6 +618,12 @@ export async function listNonTerminalAcceptedMessages( return entries.filter(state => state.status === 'accepted'); } +export async function hasNonTerminalSessionMessage( + storage: SessionMessageStorage +): Promise { + return (await listSessionMessageStates(storage)).some(state => !isTerminalStatus(state.status)); +} + async function listIndexedMessagesForWrapperRun( storage: SessionMessageStorage, wrapperRunId: string diff --git a/services/cloud-agent-next/src/session/session-prepare.test.ts b/services/cloud-agent-next/src/session/session-prepare.test.ts index d6e7f8adb2..349268d961 100644 --- a/services/cloud-agent-next/src/session/session-prepare.test.ts +++ b/services/cloud-agent-next/src/session/session-prepare.test.ts @@ -384,6 +384,22 @@ describe('createSessionWithLedger admission ladder', () => { }); }); + it('rejects a new modern control-plane session before durable session effects when isolation is disabled', async () => { + generateSessionIdMock.mockReturnValue(WORKSPACE_SESSION_ID); + const doStub = makeDoStub(); + const ctx = makeContext(doStub); + ctx.env.RUNTIME_ISOLATION_ENABLED = 'false'; + ctx.authToken = 'eyJhbGciOiJub25lIn0.eyJhdWQiOiJhcGkifQ.'; + + await expect(runCreate(ctx)).rejects.toMatchObject({ + code: 'PRECONDITION_FAILED', + message: 'runtime_isolation_unavailable', + }); + expect(createSessionReportMock).not.toHaveBeenCalled(); + expect(recordSandboxIdentityMock).not.toHaveBeenCalled(); + expect(doStub.createSessionWithInitialAdmission).not.toHaveBeenCalled(); + }); + describe.each([undefined, 'true', 'false', '', 'False', '0'] as const)( 'workspace containment samples CREDENTIAL_CONTAINMENT_ENABLED=%s', flag => { diff --git a/services/cloud-agent-next/src/session/session-registration.ts b/services/cloud-agent-next/src/session/session-registration.ts index ec9a6671b3..6fdc5b5f80 100644 --- a/services/cloud-agent-next/src/session/session-registration.ts +++ b/services/cloud-agent-next/src/session/session-registration.ts @@ -33,6 +33,11 @@ import { type CloudAgentWorktreeId, } from '@kilocode/session-ingest-contracts'; import { normalizeGitUrl } from '@kilocode/worker-utils'; +import { + createRuntimeAuthorization, + sealRuntimeAuthorization, +} from '@kilocode/worker-utils/runtime-authorization'; +import jwt from 'jsonwebtoken'; import type { Env, SandboxId } from '../types.js'; import type { CloudAgentSession } from '../persistence/CloudAgentSession.js'; @@ -62,6 +67,7 @@ import { resolveSharedSandboxAssignment } from '../shared-sandbox-route.js'; import { generateKiloSessionId } from '../utils/kilo-session-id.js'; import { sha256Hex } from '../utils/sha256.js'; import { createMessageId } from './message-id.js'; +import { assertKiloModelAvailable } from '../model-validation.js'; import type { MessageResultRPCResponse } from './message-result.js'; import type { AcceptedExecutionTurn, @@ -179,6 +185,7 @@ type NewSessionAllocation = SessionRegistrationResult & { credentialContainment: CredentialContainment; sessionService: SessionService; rollbackCliSession: () => Promise; + runtimeAuthorization?: { token: string; seal: string }; }; // ----- operation-ledger boundary (P1-A-08b) ----------------------------------- @@ -502,6 +509,11 @@ function worktreeEnabledForCreate( return sessionPlaneForNewOwner(ctx.env, owner) === 'control' && isWorktreeOwner(ctx.env, owner); } +export function assertRuntimeIsolationAdmission(env: Pick): void { + if (env.RUNTIME_ISOLATION_ENABLED === 'true') return; + throw new TRPCError({ code: 'PRECONDITION_FAILED', message: 'runtime_isolation_unavailable' }); +} + function finalizationVersionForCreate(row: OperationLedgerRow): 1 | 2 { const recorded = row.canonical_result?.[SESSION_CREATE_FINALIZATION_VERSION_KEY]; if (recorded !== undefined) { @@ -547,6 +559,41 @@ async function allocateNewSession( ) : undefined; const createdOnPlatform = input.options?.createdOnPlatform ?? 'cloud-agent'; + let runtimeAuthorization: NewSessionAllocation['runtimeAuthorization']; + const claims = jwt.decode(ctx.authToken); + const isPolicyBearing = + claims !== null && + typeof claims === 'object' && + ('aud' in claims || 'tokenPurpose' in claims || 'credentialExchange' in claims); + if (isPolicyBearing) { + if (cloudAgentSessionId.startsWith('workspace_')) assertRuntimeIsolationAdmission(ctx.env); + const secret = ctx.env.NEXTAUTH_SECRET; + const nextAuthSecret = typeof secret === 'string' ? secret : await secret.get(); + if (!nextAuthSecret) + throw new TRPCError({ code: 'INTERNAL_SERVER_ERROR', message: 'Authentication unavailable' }); + const created = await createRuntimeAuthorization({ + token: ctx.authToken, + secret: nextAuthSecret, + connectionString: ctx.env.HYPERDRIVE.connectionString, + resourceKind: 'cloud-agent-next', + resourceId: cloudAgentSessionId, + ...(orgId ? { organizationId: orgId } : {}), + }); + runtimeAuthorization = { + token: created.token, + seal: await sealRuntimeAuthorization(created.authorization, nextAuthSecret), + }; + if (initialTurn?.type === 'prompt') { + await assertKiloModelAvailable({ + env: ctx.env, + submittedModel: input.agent.model, + originalToken: runtimeAuthorization.token, + originalOrganizationId: orgId, + createdOnPlatform: input.options?.createdOnPlatform, + procedure: 'runtime_authorized_session_create', + }); + } + } try { if (ledger) { @@ -773,6 +820,7 @@ async function allocateNewSession( .error('Failed to rollback cli_sessions_v2 record'); } }, + ...(runtimeAuthorization ? { runtimeAuthorization } : {}), }; } @@ -908,8 +956,11 @@ function buildSessionRegistrationCommand( }, auth: { kiloSessionId: allocation.kiloSessionId, - kilocodeToken: ctx.authToken, + kilocodeToken: allocation.runtimeAuthorization?.token ?? ctx.authToken, }, + ...(allocation.runtimeAuthorization + ? { runtimeAuthorizationSeal: allocation.runtimeAuthorization.seal } + : {}), clone: input.clone ? { cloneFromKiloSessionId: input.clone.cloneFromKiloSessionId, diff --git a/services/cloud-agent-next/src/shared/sandbox-control-protocol.ts b/services/cloud-agent-next/src/shared/sandbox-control-protocol.ts index 984f6db703..20808f9d93 100644 --- a/services/cloud-agent-next/src/shared/sandbox-control-protocol.ts +++ b/services/cloud-agent-next/src/shared/sandbox-control-protocol.ts @@ -55,6 +55,7 @@ export const SESSION_OPERATIONS = [ 'session.sync', 'session.git.summary', 'session.detach', + 'session.runtime.retire', 'session.terminal.create', 'session.terminal.resize', 'session.terminal.close', @@ -172,6 +173,8 @@ export const sandboxHelloPayloadSchema = z.object({ nativeRuntimeRetirement: z.boolean().optional(), connectionRecovery: z.boolean().optional(), eventReceipts: z.boolean().optional(), + runtimeIsolation: z.literal(true).optional(), + runtimeRecovery: z.literal(true).optional(), workingBranches: z.boolean().optional(), }) .optional(), @@ -188,6 +191,8 @@ export const sandboxHelloResultSchema = z.object({ nativeRuntimeRetirement: z.boolean().optional(), connectionRecovery: z.boolean().optional(), eventReceipts: z.boolean().optional(), + runtimeIsolation: z.literal(true).optional(), + runtimeRecovery: z.literal(true).optional(), }) .optional(), }); @@ -363,6 +368,7 @@ export const sessionAttachPayloadSchema = z .optional(), env: z.record(z.string().max(256), z.string().max(8192)).optional(), setupCommands: z.array(z.string().max(500)).max(20).optional(), + runtimeIsolation: z.enum(['per-session']).optional(), preparation: z .object({ attemptId: z.string().min(1).max(128), @@ -561,6 +567,14 @@ export const sessionDetachResultSchema = z }) .strict(); +export const sessionRuntimeRetirePayloadSchema = z + .object({ recoveryId: z.string().uuid() }) + .strict(); + +export const sessionRuntimeRetireResultSchema = z + .object({ recoveryId: z.string().uuid(), retired: z.literal(true) }) + .strict(); + const terminalSizeSchema = z.object({ cols: z.number().int().min(2).max(500), rows: z.number().int().min(2).max(200), @@ -711,6 +725,8 @@ export type SessionSyncPayload = z.infer; export type SessionSyncResult = z.infer; export type SessionDetachPayload = z.infer; export type SessionDetachResult = z.infer; +export type SessionRuntimeRetirePayload = z.infer; +export type SessionRuntimeRetireResult = z.infer; export type SessionTerminalCreatePayload = z.infer; export type SessionTerminalCreateResult = z.infer; export type SessionTerminalResizePayload = z.infer; @@ -905,6 +921,8 @@ export const sandboxControlSocketAttachmentSchema = z.object({ .optional(), providerInstanceId: z.string().min(1).max(256).optional(), wrapperInstanceId: wrapperInstanceIdSchema.optional(), + runtimeIsolation: z.literal(true).optional(), + runtimeRecovery: z.literal(true).optional(), observation: sandboxControlObservationSchema.optional(), }); diff --git a/services/cloud-agent-next/src/shared/wrapper-bootstrap.ts b/services/cloud-agent-next/src/shared/wrapper-bootstrap.ts index 1a2b7eea55..71cf94574f 100644 --- a/services/cloud-agent-next/src/shared/wrapper-bootstrap.ts +++ b/services/cloud-agent-next/src/shared/wrapper-bootstrap.ts @@ -58,6 +58,16 @@ export type WrapperBootstrapMaterializedConfig = { runtimeSkills?: WrapperBootstrapRuntimeSkill[]; }; +export type WrapperRuntimeCredentialProxyConfig = { + /** Opaque Worker-issued handle; never the renewable backing credential. */ + handle: string; + targets: { + backendBaseUrl: string; + providerBaseUrl: string; + sessionIngestBaseUrl: string; + }; +}; + export type WrapperDevContainerMetadata = { workspacePath: string; innerWorkspaceFolder: string; @@ -129,6 +139,7 @@ export type WrapperSessionReadyRequest = { repo?: WrapperBootstrapRepoSource; devcontainer?: WrapperBootstrapDevContainer; materialized: WrapperBootstrapMaterializedConfig; + runtimeCredentialProxy?: WrapperRuntimeCredentialProxyConfig; session: WrapperSessionBinding; preparation?: { attemptId: string; @@ -247,6 +258,43 @@ function hasString(value: Record, key: string): boolean { return typeof value[key] === 'string' && value[key].length > 0; } +function isRuntimeCredentialProxyConfig( + value: unknown +): value is WrapperRuntimeCredentialProxyConfig { + if (!isRecord(value) || !hasString(value, 'handle') || !isRecord(value.targets)) return false; + const targets = value.targets; + if ( + Object.keys(value).length !== 2 || + !Object.hasOwn(value, 'handle') || + !Object.hasOwn(value, 'targets') || + Object.keys(targets).length !== 3 + ) { + return false; + } + const targetKeys = ['backendBaseUrl', 'providerBaseUrl', 'sessionIngestBaseUrl'] as const; + if (!targetKeys.every(key => Object.hasOwn(targets, key) && hasString(targets, key))) { + return false; + } + const values = Object.values(targets); + if (new Set(values).size !== 1) return false; + return targetKeys.every(key => { + try { + const target = new URL(targets[key] as string); + return ( + target.protocol === 'https:' && + !target.username && + !target.password && + !target.search && + !target.hash && + !target.port && + !target.pathname.split('/').filter(Boolean).includes('api') + ); + } catch { + return false; + } + }); +} + function isWrapperDevContainerMetadata(value: unknown): value is WrapperDevContainerMetadata { if (!isRecord(value)) return false; if (!hasString(value, 'workspacePath')) return false; @@ -297,6 +345,13 @@ export function isWrapperSessionReadyRequest(value: unknown): value is WrapperSe const materialized = value.materialized; if (!isRecord(materialized) || !isRecord(materialized.env)) return false; + if ( + value.runtimeCredentialProxy !== undefined && + !isRuntimeCredentialProxyConfig(value.runtimeCredentialProxy) + ) { + return false; + } + const session = value.session; if (!isRecord(session)) return false; if (!hasString(session, 'ingestUrl')) return false; diff --git a/services/cloud-agent-next/src/terminal/access.ts b/services/cloud-agent-next/src/terminal/access.ts index 97f95e72e9..ef9fb83eb6 100644 --- a/services/cloud-agent-next/src/terminal/access.ts +++ b/services/cloud-agent-next/src/terminal/access.ts @@ -41,6 +41,7 @@ export function validateTerminalMetadata( export type TerminalWrapperClient = { health(): Promise; + listTerminals(): Promise; createTerminal(size?: { cols: number; rows: number }): Promise; resizeTerminal(ptyId: string, size: { cols: number; rows: number }): Promise; closeTerminal(ptyId: string): Promise<{ success: boolean }>; diff --git a/services/cloud-agent-next/src/types.ts b/services/cloud-agent-next/src/types.ts index 83783b1acb..1044076680 100644 --- a/services/cloud-agent-next/src/types.ts +++ b/services/cloud-agent-next/src/types.ts @@ -598,6 +598,7 @@ export type Env = { /** Comma-separated user or org IDs admitted to the call-home control plane. `*` includes personal. */ CONTROL_PLANE_IDS?: string; WORKTREE_CREATION_ENABLED_IDS?: string; + RUNTIME_ISOLATION_ENABLED?: string; CREDENTIAL_CONTAINMENT_ENABLED?: string; /** Comma-separated org IDs that receive workspace repo snapshots, or '*' for all */ REPO_SNAPSHOT_ORG_IDS?: string; diff --git a/services/cloud-agent-next/test/e2e/README.md b/services/cloud-agent-next/test/e2e/README.md index 390d918228..177ca0c61a 100644 --- a/services/cloud-agent-next/test/e2e/README.md +++ b/services/cloud-agent-next/test/e2e/README.md @@ -166,10 +166,18 @@ tsx services/cloud-agent-next/test/e2e/smoke.ts ``` The matrix starts with `cold-hot`, which pays one cold sandbox boot and then -runs several hot same-session turns. Fresh sessions use per-session sandboxes -in local dev, so the harness identifies each newly-created sandbox instead of -killing every sandbox between cases. Kill scenarios only terminate the sandbox -family created for that scenario. +runs several hot same-session turns. The matrix tracks the session IDs returned by its own start/prepare calls. +After each scenario, including failures, it interrupts those sessions before +stopping sandboxes with proven exclusive ownership. It does not kill unrelated +or previous-run sandboxes at startup. Cleanup failures stop the matrix instead +of allowing pending work to contaminate later scenarios. Kill scenarios inject +their intentional fault before interruption, then cancel remaining work during +cleanup. + +Tracking requires a returned session ID. If unified `start` allocates ownership +but fails before returning that ID, the driver cannot automatically cancel it. +Use the failed run's user ID and ownership logs to identify and interrupt only +those sessions; do not infer ownership from container creation time. Per-run overrides via env vars. Defaults assume a zero-offset session; for any other offset, compute the real ports from `pnpm dev:status --json` @@ -225,7 +233,8 @@ source of directive truth is `test/e2e/fake-llm-server.ts`. | `slow::` | `n` content chunks `` apart, then stop + `[DONE]`. Used for pacing/timing probes. | | `idle` | One empty-delta chunk, then stop + `[DONE]`. | | `hang` | Opens the SSE stream but emits nothing and never closes. Drives abort/timeout paths. | -| `error:` | HTTP 402 with OpenAI-shaped error body carrying ``. Exercises kilo's error propagation. | +| `error-terminal:` | HTTP 400 with OpenAI-shaped error body carrying ``. Exercises nonretryable provider-error propagation through the gateway. | +| `error:` | HTTP 402 with OpenAI-shaped error body carrying ``. The non-BYOK gateway converts this to retryable HTTP 503. | | `gate:` | Opens the SSE stream, emits no chunks, blocks until the driver calls `POST /test/release?tag=`. On release, emits `"done"` + stop + `[DONE]`. | Unknown `__fake__:` directives produce HTTP 402 with @@ -265,7 +274,7 @@ These are wrapped by `releaseGate()`, `waitForGateEngaged()`, | `queue-rapid-fire-no-gate` | Send immediate follow-ups behind `echo:first` and assert they reach their terminal FIFO state without gate coordination. | | `queue-overflow` | Block on `gate:overflow`, fill the pending queue until enqueue fails with HTTP 429, release gate, drain. | | `queue-interrupt-clears` | Block on `gate:`, enqueue two, `interruptSession`, assert `cloud.message.failed` with `reason: 'interrupted'` for each. | -| `llm-error` | Return fake provider HTTP 402 and assert the turn reaches a failed terminal event instead of hanging. | +| `llm-error` | Return nonretryable fake provider HTTP 400 and assert the turn reaches a failed terminal event instead of hanging. | | `chunked-streaming` | Stream delayed fake chunks and assert multiple downstream `message.part.delta` events survive. | | `empty-response` | Run `idle`, assert completion, and assert no downstream `message.part.delta` is emitted. | | `interrupt-mid-stream` | Interrupt an actively gated fake request and assert the active message is interrupted, not a queued message. | @@ -323,19 +332,16 @@ the newer `start` / `send` procedures. `prepareSession` requires `pnpm dev:restart cloudflare-git-token-service`, then confirm the entry reappears. The fake LLM is irrelevant here — kilo never gets far enough to dial it. -- **Matrix fails intermittently with `preparing×N` and no terminal** — - environmental, not a regression. The `@cloudflare/containers` library's - container control connection sometimes returns 503 under Docker Desktop - load, triggering exponential-backoff retries that consume the scenario - timeout. The `smoke.ts` matrix now kills stale containers before starting - and uses a 120s per-scenario timeout, but this is not always enough. If - the matrix is flaky: (1) stop any competing dev session from another - worktree that also runs Cloud Agent sandboxes; (2) prune stopped containers - (`docker ps -a --filter status=exited --format '{{.Names}}' | rg - workerd-cloud-agent | xargs -r docker rm -f`); (3) restart - `cloud-agent-next` to clear stale DO alarm timers; (4) re-run the failing - scenario standalone — if it passes alone, the matrix failure was Docker - contention, not code. +- **Matrix fails with `preparing×N` and no terminal** — Correlate the failed + message with Worker and wrapper logs before classifying the cause. Container + startup failures happen before wrapper bootstrap; a `post-bootstrap kilo + session lookup begin` without an end identifies a later native lookup stall. + Matrix cleanup interrupts its tracked sessions before stopping exclusively + owned sandboxes. For older runs or an interrupted driver, cancel only the + recorded run-owned sessions before any owned-family teardown: killing a + container alone leaves queued work able to recreate it after a Worker restart. + Preserve the failed result and rerun the scenario in isolation; a successful + retry does not erase the original failure. - **`releaseGate` returned 404** — the gate already went away, usually because the wrapper's request was aborted (e.g. by an `interruptSession`). Queue-interrupt-clears tolerates this; other scenarios treat it as an diff --git a/services/cloud-agent-next/test/e2e/client.ts b/services/cloud-agent-next/test/e2e/client.ts index f85a831304..08c697844d 100644 --- a/services/cloud-agent-next/test/e2e/client.ts +++ b/services/cloud-agent-next/test/e2e/client.ts @@ -32,6 +32,8 @@ export type CallbackTarget = { }; export type DriverConfig = { + /** Track returned session IDs, including prepare success followed by initiation failure. */ + onSessionCreated?: (sessionId: string) => void; workerUrl: string; expectControlPlane?: boolean; user: TestUser; @@ -193,6 +195,7 @@ export async function startSession( api === 'legacy' ? await startSessionLegacy(config, args) : await startSessionUnified(config, args); + if (api === 'unified') config.onSessionCreated?.(result.cloudAgentSessionId); if (config.expectControlPlane && !result.cloudAgentSessionId.startsWith('workspace_')) { throw new Error( `Started ${result.cloudAgentSessionId}, but expected an enrolled workspace_* session; do not retry start` @@ -279,6 +282,7 @@ async function startSessionLegacy( streamUrl?: string; status?: string; }; + config.onSessionCreated?.(prepared.cloudAgentSessionId); const initiated = await trpcCall(config, 'initiateFromKilocodeSessionV2', { cloudAgentSessionId: prepared.cloudAgentSessionId, }); diff --git a/services/cloud-agent-next/test/e2e/fake-llm-server.ts b/services/cloud-agent-next/test/e2e/fake-llm-server.ts index 6a25f3c558..fb0d4d5ab5 100644 --- a/services/cloud-agent-next/test/e2e/fake-llm-server.ts +++ b/services/cloud-agent-next/test/e2e/fake-llm-server.ts @@ -752,6 +752,11 @@ export const scenarioRegistry: Record = { }); }, + 'error-terminal'(args, ctx) { + const message = args[0] ?? 'simulated error'; + writeJsonError(ctx.res, 400, message, 'invalid_request'); + }, + error(args, ctx) { const message = args[0] ?? 'simulated error'; writeJsonError(ctx.res, 402, message, 'insufficient_quota'); diff --git a/services/cloud-agent-next/test/e2e/lifecycle.ts b/services/cloud-agent-next/test/e2e/lifecycle.ts index 5143df32e5..290dae7dd9 100644 --- a/services/cloud-agent-next/test/e2e/lifecycle.ts +++ b/services/cloud-agent-next/test/e2e/lifecycle.ts @@ -205,6 +205,13 @@ async function stopOwnedSandboxFamily(sandbox: SandboxContainer, sessionId: stri return killed; } +/** Only tear down sandboxes whose exclusive session ownership can be proven. */ +export async function stopOwnedSessionSandboxes(sessionId: string): Promise { + for (const sandbox of await findOwnedSandboxes(sessionId, new Set())) { + await stopOwnedSandboxFamily(sandbox, sessionId); + } +} + async function sendRecoveryTurn( config: DriverConfig, sessionId: string, @@ -818,20 +825,36 @@ export async function lifecycleWorktreeShared(args: LifecycleArgs): Promise` so the fake returns HTTP 402 with - * an OpenAI-shape error body. Assert the worker terminalizes with a failure + * llm-error: drives `__fake__:error-terminal:` so the fake returns HTTP 400 + * with an OpenAI-shape error body. The gateway converts upstream 402 to retryable + * 503 for non-BYOK requests. Assert the worker terminalizes with a failure * (not `complete`), and the sandbox doesn't hang indefinitely. * * Conversation arg is the error message (e.g. `llm-error boom`). @@ -2204,7 +2228,11 @@ export async function lifecycleLlmError(args: LifecycleArgs): Promise value.startsWith('HOME='))?.slice(5); + if (!home || !path.isAbsolute(home) || fs.realpathSync(home) !== home) continue; + listeners.push({ serverUrl, processId, directory, home }); } catch { continue; } @@ -179,7 +185,7 @@ function rootResult(root, serverUrl) { if (matches.length !== 1 || serverUrl !== request.serverUrl || root.directory !== request.directory) { return { ok: false, reason: 'Kilo listener identity changed' }; } - return { ok: true, id: root.id, directory: root.directory, processId: request.processId }; + return { ok: true, id: root.id, directory: root.directory, home: request.home, processId: request.processId }; } async function run() { @@ -219,7 +225,11 @@ async function run() { } const listeners = kiloListeners().filter(listener => sameListener(listener, request)); - if (listeners.length !== 1 || !(await getRoot(request.serverUrl, request.ownerKiloSessionId))) { + if ( + listeners.length !== 1 || + listeners[0].home !== request.home || + !(await getRoot(request.serverUrl, request.ownerKiloSessionId)) + ) { return { ok: false, reason: 'Owned Kilo listener identity did not match' }; } @@ -447,6 +457,7 @@ function requireControlPlaneKiloRoot( if ( result.id !== kiloSessionId || typeof result.directory !== 'string' || + typeof result.home !== 'string' || typeof result.processId !== 'number' || !Number.isSafeInteger(result.processId) || result.processId <= 0 @@ -456,6 +467,7 @@ function requireControlPlaneKiloRoot( return { id: result.id, directory: result.directory, + home: result.home, processId: result.processId, }; } @@ -488,6 +500,8 @@ export async function findControlPlaneKiloRuntime( !/^http:\/\/(?:127\.0\.0\.1|\[::1\]):\d+$/.test(result.serverUrl) || typeof result.directory !== 'string' || !result.directory.startsWith('/') || + typeof result.home !== 'string' || + !result.home.startsWith('/') || typeof result.processId !== 'number' || !Number.isSafeInteger(result.processId) || result.processId <= 0 || @@ -500,6 +514,7 @@ export async function findControlPlaneKiloRuntime( kiloSessionId, serverUrl: result.serverUrl, directory: result.directory, + home: result.home, processId: result.processId, ...(typeof result.logPath === 'string' ? { logPath: result.logPath } : {}), }); @@ -528,6 +543,7 @@ export async function stopOwnedControlPlaneSandbox( ownerKiloSessionId: kiloSessionId, serverUrl: runtime.serverUrl, directory: runtime.directory, + home: runtime.home, processId: runtime.processId, }, executeDocker @@ -564,6 +580,7 @@ export async function inspectControlPlaneKiloRoot( kiloSessionId, serverUrl: runtime.serverUrl, directory: runtime.directory, + home: runtime.home, processId: runtime.processId, ownerKiloSessionId: runtime.kiloSessionId, }); @@ -579,6 +596,7 @@ export async function controlPlaneKiloRootExists( kiloSessionId, serverUrl: runtime.serverUrl, directory: runtime.directory, + home: runtime.home, processId: runtime.processId, ownerKiloSessionId: runtime.kiloSessionId, }); @@ -597,6 +615,7 @@ export async function inspectControlPlaneWorkspaceFile( kiloSessionId: input.kiloSessionId, serverUrl: runtime.serverUrl, directory: runtime.directory, + home: runtime.home, processId: runtime.processId, ownerKiloSessionId: runtime.kiloSessionId, filePath: input.filePath, @@ -627,6 +646,7 @@ export async function inspectControlPlaneQuestions( questionId: input.questionId, serverUrl: runtime.serverUrl, directory: runtime.directory, + home: runtime.home, processId: runtime.processId, ownerKiloSessionId: runtime.kiloSessionId, }); @@ -667,6 +687,7 @@ export async function importControlPlaneKiloRoot( sourceKiloSessionId: runtime.kiloSessionId, serverUrl: runtime.serverUrl, directory: runtime.directory, + home: runtime.home, processId: runtime.processId, ownerKiloSessionId: runtime.kiloSessionId, }); @@ -683,6 +704,7 @@ export async function promptControlPlaneKiloRoot( sourceKiloSessionId: runtime.kiloSessionId, serverUrl: runtime.serverUrl, directory: runtime.directory, + home: runtime.home, processId: runtime.processId, ownerKiloSessionId: runtime.kiloSessionId, messageId: input.messageId, @@ -705,6 +727,7 @@ export async function waitForControlPlaneKiloCompletion( kiloSessionId: input.kiloSessionId, serverUrl: runtime.serverUrl, directory: runtime.directory, + home: runtime.home, processId: runtime.processId, ownerKiloSessionId: runtime.kiloSessionId, messageId: input.messageId, diff --git a/services/cloud-agent-next/test/e2e/smoke-cleanup.ts b/services/cloud-agent-next/test/e2e/smoke-cleanup.ts new file mode 100644 index 0000000000..bc0169a9ae --- /dev/null +++ b/services/cloud-agent-next/test/e2e/smoke-cleanup.ts @@ -0,0 +1,38 @@ +import type { InterruptResult } from './client.js'; + +type CleanupDependencies = { + interrupt: (sessionId: string) => Promise; + stopOwnedSandboxes: (sessionId: string) => Promise; +}; + +/** Settle durable demand before killing containers that alarms could recreate. */ +export async function cleanupOwnedSessions( + sessionIds: ReadonlySet, + deps: CleanupDependencies +): Promise { + const interrupted: string[] = []; + const errors: Error[] = []; + for (const sessionId of sessionIds) { + try { + const result = await deps.interrupt(sessionId); + if ( + !result.success && + result.message !== 'No accepted wrapper messages or pending queued messages' && + result.message !== 'No session work to interrupt' + ) { + throw new Error(`Interruption was not confirmed for ${sessionId}`); + } + interrupted.push(sessionId); + } catch { + errors.push(new Error(`Failed to interrupt owned session ${sessionId}; skipped teardown`)); + } + } + for (const sessionId of interrupted) { + try { + await deps.stopOwnedSandboxes(sessionId); + } catch { + errors.push(new Error(`Failed to stop owned sandboxes for ${sessionId}`)); + } + } + if (errors.length > 0) throw new AggregateError(errors, 'Owned session cleanup failed'); +} diff --git a/services/cloud-agent-next/test/e2e/smoke.ts b/services/cloud-agent-next/test/e2e/smoke.ts index cc0b718842..9733d17938 100644 --- a/services/cloud-agent-next/test/e2e/smoke.ts +++ b/services/cloud-agent-next/test/e2e/smoke.ts @@ -14,15 +14,14 @@ import path from 'node:path'; import { fileURLToPath } from 'node:url'; import { ensureTestUser, loadDevVars, loadRepoEnvFiles, DRIVER_USER_EMAIL_SUFFIX } from './auth.js'; -import { DEFAULT_CONFIG, type ApiVersion, type DriverConfig } from './client.js'; -import { LIFECYCLE_SCENARIOS, type LifecycleResult } from './lifecycle.js'; -import { printResult } from './run.js'; +import { DEFAULT_CONFIG, interruptSession, type ApiVersion, type DriverConfig } from './client.js'; import { - killSandboxFamily, - listSandboxContainers, - waitForSandboxFamilyGone, - type SandboxContainer, -} from './sandbox-control.js'; + LIFECYCLE_SCENARIOS, + stopOwnedSessionSandboxes, + type LifecycleResult, +} from './lifecycle.js'; +import { cleanupOwnedSessions } from './smoke-cleanup.js'; +import { printResult } from './run.js'; const SERVICE_PACKAGE_DIR = path.resolve(path.dirname(fileURLToPath(import.meta.url)), '../..'); @@ -67,32 +66,6 @@ const DEFAULT_MATRIX: Case[] = [ { lifecycle: 'kill-mid-flight', conversation: 'hang' }, ]; -function sandboxFamilyKey(container: SandboxContainer): string { - return container.isProxy ? container.name.replace(/-proxy$/, '') : container.name; -} - -async function cleanupMatrixSandboxes(baselineSandboxIds: Set): Promise { - const createdSandboxes = (await listSandboxContainers()).filter( - container => !baselineSandboxIds.has(container.id) - ); - const sandboxFamilies = new Map(); - for (const container of createdSandboxes) { - const key = sandboxFamilyKey(container); - const existing = sandboxFamilies.get(key); - if (!existing || (existing.isProxy && !container.isProxy)) { - sandboxFamilies.set(key, container); - } - } - - for (const sandbox of sandboxFamilies.values()) { - await killSandboxFamily(sandbox); - const gone = await waitForSandboxFamilyGone(sandbox, 30_000); - if (!gone) { - console.warn(`smoke: sandbox family ${sandboxFamilyKey(sandbox)} remained after cleanup`); - } - } -} - async function main(): Promise { loadRepoEnvFiles(SERVICE_PACKAGE_DIR); const devVars = loadDevVars(SERVICE_PACKAGE_DIR); @@ -101,7 +74,11 @@ async function main(): Promise { const user = await ensureTestUser(process.env.DATABASE_URL, email); console.log(`driver user: ${user.id} (${user.email})`); + const ownedSessionIds = new Set(); const config: DriverConfig = { + onSessionCreated: sessionId => { + ownedSessionIds.add(sessionId); + }, ...DEFAULT_CONFIG, user, nextAuthSecret: devVars.NEXTAUTH_SECRET ?? '', @@ -111,26 +88,6 @@ async function main(): Promise { model: process.env.E2E_MODEL ?? DEFAULT_CONFIG.model, }; - // Kill stale sandbox containers from previous runs before starting. - // Accumulated stopped/running containers degrade Docker Desktop performance - // and cause the preparing×7 wrapper-startup stall pattern. The baseline - // snapshot below only identifies *new* containers, so leftovers from prior - // runs would be skipped by cleanupMatrixSandboxes. - const staleContainers = await listSandboxContainers(); - if (staleContainers.length > 0) { - console.log(`smoke: cleaning ${staleContainers.length} stale sandbox container(s)`); - for (const container of staleContainers) { - await killSandboxFamily(container); - const gone = await waitForSandboxFamilyGone(container, 30_000); - if (!gone) { - console.warn(`smoke: sandbox family ${sandboxFamilyKey(container)} remained after cleanup`); - } - } - } - - const baselineSandboxIds = new Set( - (await listSandboxContainers()).map(container => container.id) - ); const results: LifecycleResult[] = []; for (const { lifecycle, conversation, api = 'unified' } of DEFAULT_MATRIX) { const scenarioFn = LIFECYCLE_SCENARIOS[lifecycle]; @@ -144,7 +101,11 @@ async function main(): Promise { printResult(result); results.push(result); } finally { - await cleanupMatrixSandboxes(baselineSandboxIds); + await cleanupOwnedSessions(ownedSessionIds, { + interrupt: sessionId => interruptSession(config, sessionId), + stopOwnedSandboxes: stopOwnedSessionSandboxes, + }); + ownedSessionIds.clear(); } } diff --git a/services/cloud-agent-next/test/integration/runtime-authorization-recovery.test.ts b/services/cloud-agent-next/test/integration/runtime-authorization-recovery.test.ts new file mode 100644 index 0000000000..c39bb1e83d --- /dev/null +++ b/services/cloud-agent-next/test/integration/runtime-authorization-recovery.test.ts @@ -0,0 +1,518 @@ +import { env, listDurableObjectIds, runInDurableObject } from 'cloudflare:test'; +import { beforeEach, describe, expect, it } from 'vitest'; +import jwt from 'jsonwebtoken'; +import { sealRuntimeAuthorization } from '@kilocode/worker-utils/runtime-authorization'; +import type { RuntimeAuthorization } from '@kilocode/worker-utils/runtime-authorization-contract'; +import { RUNTIME_PROXY_GRANT_KEY } from '../../src/runtime-credential-proxy.js'; +import { + RUNTIME_AUTHORIZATION_KEY, + RUNTIME_AUTHORIZATION_RECOVERY_KEY, +} from '../../src/session/runtime-authorization-persistence.js'; +import { + allocateWrapperRuntimeState, + getWrapperRuntimeState, +} from '../../src/session/wrapper-runtime-state.js'; +import { registerReadySession } from '../helpers/session-setup.js'; + +const organizationId = '11111111-1111-4111-8111-111111111111'; + +function authorization(input: { + id: string; + sessionId: string; + userId: string; + expiresAt: string; + issuedAt?: string; + state?: 'active' | 'revoked'; +}): RuntimeAuthorization { + return { + version: 1, + id: input.id, + resourceKind: 'cloud-agent-next', + resourceId: input.sessionId, + userId: input.userId, + authorizationUserId: input.userId, + organizationId, + issuedAt: input.issuedAt ?? new Date(Date.now() - 60_000).toISOString(), + delegationExpiresAt: input.expiresAt, + state: input.state ?? 'active', + bindings: { + userPepperDigest: 'a'.repeat(64), + authorizationPepperDigest: 'b'.repeat(64), + userMembershipId: 'membership_1', + authorizationUserMembershipId: 'membership_1', + }, + source: { admissionSource: 'user' }, + }; +} + +async function secret() { + return typeof env.NEXTAUTH_SECRET === 'string' ? env.NEXTAUTH_SECRET : env.NEXTAUTH_SECRET.get(); +} + +async function seal(value: RuntimeAuthorization) { + return sealRuntimeAuthorization(value, await secret()); +} + +async function runtimeToken(value: RuntimeAuthorization) { + return jwt.sign({ runtimeAuthorization: { id: value.id } }, await secret(), { + algorithm: 'HS256', + expiresIn: '30 minutes', + }); +} + +beforeEach(async () => { + const namespaces = [env.CLOUD_AGENT_SESSION, env.SANDBOX_SESSION]; + await Promise.all( + namespaces.flatMap(async namespace => { + const ids = await listDurableObjectIds(namespace); + return Promise.all( + ids.map(id => + runInDurableObject(namespace.get(id), instance => instance.ctx.storage.deleteAll()) + ) + ); + }) + ); +}); + +describe('runtime authorization recovery', () => { + it('rejects prepared and ordinary admission while a recovery lock is held', async () => { + const userId = 'user_cloud_recovery_lock'; + const sessionId = 'agent_cloud_recovery_lock'; + const stub = env.CLOUD_AGENT_SESSION.get( + env.CLOUD_AGENT_SESSION.idFromName(`${userId}:${sessionId}`) + ); + + const result = await runInDurableObject(stub, async instance => { + await registerReadySession(instance, { + sessionId, + userId, + kiloSessionId: 'aaaaaaaa-aaaa-4aaa-aaaa-aaaaaaaaaaab', + prompt: 'initial', + mode: 'code', + model: 'test-model', + initialMessageId: 'msg_018f1e2d3c4bRecoveryLockAb', + }); + await instance.ctx.storage.put(RUNTIME_AUTHORIZATION_RECOVERY_KEY, { + expectedOldId: '00000000-0000-4000-8000-000000000001', + recoveryId: '00000000-0000-4000-8000-000000000002', + }); + return { + ordinary: await instance.admitSubmittedMessage({ + userId, + turn: { type: 'prompt', id: 'msg_018f1e2d3c4bRecoveryBusyAb', prompt: 'follow up' }, + }), + prepared: await instance.admitPreparedInitialMessage({ userId }), + }; + }); + + expect(result.ordinary).toEqual({ + success: false, + code: 'COMPUTE_STOPPING', + error: 'Runtime authorization recovery is in progress', + }); + expect(result.prepared).toEqual(result.ordinary); + }); + + it('recovers an expired CloudAgentSession authorization only after confirmed idle retirement', async () => { + const userId = 'user_cloud_recovery'; + const sessionId = 'agent_cloud_recovery'; + const old = authorization({ + id: '00000000-0000-4000-8000-000000000101', + sessionId, + userId, + expiresAt: new Date(Date.now() - 24 * 60 * 60_000).toISOString(), + issuedAt: new Date(Date.now() - 26 * 60 * 60_000).toISOString(), + }); + const fresh = authorization({ + id: '00000000-0000-4000-8000-000000000102', + sessionId, + userId, + expiresAt: new Date(Date.now() + 60 * 60_000).toISOString(), + }); + const stub = env.CLOUD_AGENT_SESSION.get( + env.CLOUD_AGENT_SESSION.idFromName(`${userId}:${sessionId}`) + ); + + const result = await runInDurableObject(stub, async instance => { + await registerReadySession(instance, { + sessionId, + userId, + orgId: organizationId, + kiloSessionId: 'aaaaaaaa-aaaa-4aaa-aaaa-aaaaaaaaaaaa', + prompt: 'initial', + mode: 'code', + model: 'test-model', + kilocodeToken: 'expired-token', + initialMessageId: 'msg_018f1e2d3c4bAbCdEfGhIjKlMn', + }); + const previousRuntime = await allocateWrapperRuntimeState(instance.ctx.storage); + await instance.ctx.storage.put(RUNTIME_AUTHORIZATION_KEY, old); + await instance.ctx.storage.put(RUNTIME_PROXY_GRANT_KEY, { cached: 'old-grant' }); + const stops: string[] = []; + instance['getTerminalClient'] = async () => + ({ success: true, data: { client: { listTerminals: async () => [] } } }) as never; + instance['physicalWrapperStopper'] = async request => { + stops.push(request.reason); + return { status: 'absent' }; + }; + let observations = 0; + instance['physicalWrapperObserver'] = async () => + ++observations === 1 ? { status: 'present' } : { status: 'absent' }; + + const outcome = await instance.recoverExpiredRuntimeAuthorization({ + ownerId: userId, + expectedOldId: old.id, + recoveryId: '00000000-0000-4000-8000-000000000103', + runtimeAuthorizationSeal: await seal(fresh), + runtimeToken: 'fresh-token', + }); + const preparedAdmission = await instance.admitPreparedInitialMessage({ userId }); + return { + outcome, + preparedAdmission, + stops, + metadata: await instance.getMetadata(), + authorization: await instance.ctx.storage.get(RUNTIME_AUTHORIZATION_KEY), + grant: await instance.ctx.storage.get(RUNTIME_PROXY_GRANT_KEY), + recovery: await instance.ctx.storage.get(RUNTIME_AUTHORIZATION_RECOVERY_KEY), + runtime: await getWrapperRuntimeState(instance.ctx.storage), + previousRuntime: previousRuntime.state, + }; + }); + + expect(result.outcome).toEqual({ status: 'recovered' }); + expect(result.preparedAdmission).toMatchObject({ + success: true, + messageId: 'msg_018f1e2d3c4bAbCdEfGhIjKlMn', + }); + expect(result.stops).toEqual(['idle-timeout']); + expect(result.metadata).toMatchObject({ + identity: { userId, sessionId, orgId: organizationId }, + auth: { kilocodeToken: 'fresh-token' }, + }); + expect(result.authorization).toMatchObject({ id: fresh.id, state: 'active' }); + expect(result.grant).toBeUndefined(); + expect(result.recovery).toBeUndefined(); + expect(result.runtime).toEqual({ + wrapperGeneration: result.previousRuntime.wrapperGeneration + 1, + }); + }); + + it('denies invalid recovery, preserves active PTYs, and fences terminal mutations while locked', async () => { + const userId = 'user_cloud_recovery_guards'; + const sessionId = 'agent_cloud_recovery_guards'; + const old = authorization({ + id: '00000000-0000-4000-8000-000000000201', + sessionId, + userId, + expiresAt: new Date(Date.now() - 60_000).toISOString(), + issuedAt: new Date(Date.now() - 2 * 60 * 60_000).toISOString(), + }); + const fresh = authorization({ + id: '00000000-0000-4000-8000-000000000202', + sessionId, + userId, + expiresAt: new Date(Date.now() + 60 * 60_000).toISOString(), + }); + const stub = env.CLOUD_AGENT_SESSION.get( + env.CLOUD_AGENT_SESSION.idFromName(`${userId}:${sessionId}`) + ); + + const outcomes = await runInDurableObject(stub, async instance => { + await registerReadySession(instance, { + sessionId, + userId, + orgId: organizationId, + prompt: 'initial', + mode: 'code', + model: 'test-model', + }); + await instance.ctx.storage.put(RUNTIME_AUTHORIZATION_KEY, old); + const input = { + ownerId: userId, + expectedOldId: old.id, + recoveryId: '00000000-0000-4000-8000-000000000203', + runtimeAuthorizationSeal: await seal(fresh), + runtimeToken: 'fresh-token', + }; + const foreign = await instance.recoverExpiredRuntimeAuthorization({ + ...input, + ownerId: 'other', + }); + const revoked = await instance.recoverExpiredRuntimeAuthorization({ + ...input, + runtimeAuthorizationSeal: await seal({ ...fresh, state: 'revoked' }), + }); + await instance.admitSubmittedMessage({ + userId, + turn: { type: 'prompt', id: 'msg_018f1e2d3c4bBusyRecoveryAb', prompt: 'queued' }, + }); + const busy = await instance.recoverExpiredRuntimeAuthorization(input); + return { foreign, revoked, busy }; + }); + + expect(outcomes).toEqual({ + foreign: { status: 'denied' }, + revoked: { status: 'denied' }, + busy: { status: 'busy' }, + }); + + { + const userId = 'user_cloud_recovery_pty'; + const sessionId = 'agent_cloud_recovery_pty'; + const old = authorization({ + id: '00000000-0000-4000-8000-000000000251', + sessionId, + userId, + expiresAt: new Date(Date.now() - 60_000).toISOString(), + issuedAt: new Date(Date.now() - 2 * 60 * 60_000).toISOString(), + }); + const fresh = authorization({ + id: '00000000-0000-4000-8000-000000000252', + sessionId, + userId, + expiresAt: new Date(Date.now() + 60 * 60_000).toISOString(), + }); + const stub = env.CLOUD_AGENT_SESSION.get( + env.CLOUD_AGENT_SESSION.idFromName(`${userId}:${sessionId}`) + ); + + const result = await runInDurableObject(stub, async instance => { + await registerReadySession(instance, { + sessionId, + userId, + orgId: organizationId, + prompt: 'initial', + mode: 'code', + model: 'test-model', + kilocodeToken: 'expired-token', + }); + await instance.ctx.storage.put(RUNTIME_AUTHORIZATION_KEY, old); + await instance.ctx.storage.put(RUNTIME_PROXY_GRANT_KEY, { cached: 'old-grant' }); + const createTerminal = () => { + throw new Error('creation must be blocked by the recovery lock'); + }; + const resizeTerminal = () => { + throw new Error('reconnection must be blocked by the recovery lock'); + }; + const closeTerminal = async () => ({ success: true }); + instance['getTerminalClient'] = async () => + ({ + success: true, + data: { + client: { + listTerminals: async () => [{ id: 'pty_active' }], + createTerminal, + resizeTerminal, + closeTerminal, + }, + }, + }) as never; + instance['physicalWrapperObserver'] = async () => ({ status: 'present' }); + let stops = 0; + instance['physicalWrapperStopper'] = async () => { + stops += 1; + return { status: 'absent' }; + }; + const recovery = await instance.recoverExpiredRuntimeAuthorization({ + ownerId: userId, + expectedOldId: old.id, + recoveryId: '00000000-0000-4000-8000-000000000253', + runtimeAuthorizationSeal: await seal(fresh), + runtimeToken: 'fresh-token', + }); + return { + recovery, + stops, + authorization: await instance.ctx.storage.get(RUNTIME_AUTHORIZATION_KEY), + grant: await instance.ctx.storage.get(RUNTIME_PROXY_GRANT_KEY), + locked: await instance.isRuntimeAuthorizationRecoveryInProgress(), + create: await instance.createTerminal({}), + resize: await instance.resizeTerminal({ ptyId: 'pty_active', cols: 80, rows: 24 }), + close: await instance.closeTerminal({ ptyId: 'pty_active' }), + }; + }); + + expect(result.recovery).toEqual({ status: 'busy' }); + expect(result.stops).toBe(0); + expect(result.authorization).toMatchObject({ id: old.id }); + expect(result.grant).toEqual({ cached: 'old-grant' }); + expect(result.locked).toBe(true); + expect(result.create).toEqual({ + success: false, + error: 'Runtime authorization recovery is in progress', + }); + expect(result.resize).toEqual({ + success: false, + error: 'Runtime authorization recovery is in progress', + }); + expect(result.close).toEqual({ success: true, data: { success: true } }); + } + }); + + it('retires an idle SandboxSession runtime, clears its attachment fence, and reattaches on dispatch', async () => { + const userId = 'user_sandbox_recovery'; + const sessionId = 'workspace_sandbox_recovery'; + const sandboxId = 'ses-11111111111141118111111111111111'; + const wrapperInstanceId = '22222222-2222-4222-8222-222222222222'; + const old = authorization({ + id: '00000000-0000-4000-8000-000000000301', + sessionId, + userId, + expiresAt: new Date(Date.now() - 60_000).toISOString(), + issuedAt: new Date(Date.now() - 2 * 60 * 60_000).toISOString(), + }); + const fresh = authorization({ + id: '00000000-0000-4000-8000-000000000302', + sessionId, + userId, + expiresAt: new Date(Date.now() + 60 * 60_000).toISOString(), + }); + const stub = env.SANDBOX_SESSION.getByName(`${userId}:${sessionId}`); + + const result = await runInDurableObject(stub, async instance => { + const requests: string[] = []; + let retirementAttempts = 0; + const control = { + getStatus: async () => ({ + physical: 'running' as const, + connection: 'ready' as const, + work: 'idle' as const, + wrapperInstanceId, + runtimeRecovery: true as const, + }), + getRuntimeCredentialProxyFence: async () => ({ + allocationId: 'allocation_1', + plane: 'control' as const, + providerInstanceId: 'provider_1', + connectionId: 'connection_1', + wrapperInstanceId, + }), + request: async (request: { operation: string; payload?: unknown }) => { + requests.push(request.operation); + if (request.operation === 'session.runtime.retire') { + retirementAttempts += 1; + if (retirementAttempts === 1) throw new Error('retirement acknowledgement lost'); + return { + type: 'response' as const, + requestId: 'retire', + ok: true as const, + result: { + retired: true, + recoveryId: (request.payload as { recoveryId: string }).recoveryId, + }, + }; + } + if (request.operation === 'session.attach') { + return { + type: 'response' as const, + requestId: 'attach', + ok: true as const, + result: { attached: true }, + }; + } + return { + type: 'response' as const, + requestId: 'prompt', + ok: true as const, + result: { messageId: 'msg_018f1e2d3c4bFreshDispatchAb', status: 'accepted' }, + }; + }, + ensureReady: async () => ({ + physical: 'running' as const, + connection: 'ready' as const, + wrapperInstanceId, + attachment: { + directory: '/workspace/recovery', + env: { KILOCODE_TOKEN: 'control-token' }, + kilo: { + scopeId: sessionId, + token: 'control-token', + targets: { + backendBaseUrl: 'https://backend.example.test', + providerBaseUrl: 'https://provider.example.test', + sessionIngestBaseUrl: 'https://ingest.example.test', + }, + }, + }, + }), + attachSession: async () => ({}), + }; + instance['env'].SANDBOX_CONTROL = { getByName: () => control } as never; + instance['env'].WORKER_URL = 'https://worker.example.test'; + expect( + await instance.registerSession({ + identity: { sessionId, userId, orgId: organizationId }, + auth: { + kiloSessionId: 'bbbbbbbb-bbbb-4bbb-bbbb-bbbbbbbbbbbb', + kilocodeToken: 'expired-token', + }, + runtimeAuthorizationSeal: await seal(old), + agent: { mode: 'code', model: 'test-model' }, + workspace: { sandboxId, workspacePath: '/workspace/recovery' }, + }) + ).toEqual({ success: true }); + instance['terminalLifecycle'].recordAttachment({ + metadata: (await instance.getMetadata())!, + sandboxId, + wrapperInstanceId, + epoch: 0, + }); + instance.ctx.storage.kv.put(RUNTIME_PROXY_GRANT_KEY, { cached: 'old-grant' }); + const recoveryInput = { + ownerId: userId, + expectedOldId: old.id, + recoveryId: '00000000-0000-4000-8000-000000000303', + runtimeAuthorizationSeal: await seal(fresh), + runtimeToken: await runtimeToken(fresh), + }; + const lostAcknowledgement = await instance.recoverExpiredRuntimeAuthorization(recoveryInput); + const retainedRecovery = instance.ctx.storage.kv.get(RUNTIME_AUTHORIZATION_RECOVERY_KEY); + const recovered = await instance.recoverExpiredRuntimeAuthorization(recoveryInput); + const afterRecovery = { + metadata: await instance.getMetadata(), + authorization: instance.ctx.storage.kv.get(RUNTIME_AUTHORIZATION_KEY), + grant: instance.ctx.storage.kv.get(RUNTIME_PROXY_GRANT_KEY), + attachment: instance['terminalLifecycle'].getAttachedWrapperInstanceId(), + }; + const admitted = await instance.admitSubmittedMessage({ + userId, + turn: { + type: 'command', + id: 'msg_018f1e2d3c4bFreshDispatchAb', + command: 'status', + arguments: '', + }, + }); + await instance['dispatchQueued']('msg_018f1e2d3c4bFreshDispatchAb', { allowCreate: true }); + return { + lostAcknowledgement, + retainedRecovery, + recovered, + afterRecovery, + admitted, + requests, + }; + }); + + expect(result.lostAcknowledgement).toEqual({ status: 'retry' }); + expect(result.retainedRecovery).toEqual({ + expectedOldId: old.id, + recoveryId: '00000000-0000-4000-8000-000000000303', + }); + expect(result.recovered).toEqual({ status: 'recovered' }); + expect(result.afterRecovery).toMatchObject({ + metadata: { identity: { userId, sessionId, orgId: organizationId } }, + authorization: { id: fresh.id }, + }); + expect(result.afterRecovery.grant).toBeUndefined(); + expect(result.afterRecovery.attachment).toBeUndefined(); + expect(result.admitted).toMatchObject({ success: true, outcome: 'queued' }); + expect(result.requests).toEqual([ + 'session.runtime.retire', + 'session.runtime.retire', + 'session.attach', + 'session.prompt', + ]); + }); +}); diff --git a/services/cloud-agent-next/test/integration/sandbox-control.test.ts b/services/cloud-agent-next/test/integration/sandbox-control.test.ts index 08ef322418..b56c45616c 100644 --- a/services/cloud-agent-next/test/integration/sandbox-control.test.ts +++ b/services/cloud-agent-next/test/integration/sandbox-control.test.ts @@ -50,6 +50,10 @@ import { type SessionCredentialGrant, } from '../../src/sandbox-control/session-credentials.js'; import { findMatchingCredentialInjectionRule } from '../../src/sandbox-control/vercel-network-policy.js'; +import { + createRuntimeProxyGrant, + issueRuntimeCredentialProxyHandle, +} from '../../src/runtime-credential-proxy.js'; import { MANAGED_SCM_OUTBOUND_HANDLER } from '../../src/sandbox-id.js'; import { SandboxSession } from '../../src/sandbox-session/SandboxSession.js'; import { @@ -1086,6 +1090,7 @@ async function credentialFixture( const environment = { ...env, ...VERCEL_ENV, + NEXTAUTH_SECRET: 'integration-runtime-proxy-secret', GIT_TOKEN_SERVICE: broker.binding, WORKER_URL: 'https://worker.test', KILOCODE_BACKEND_BASE_URL: CONTAINMENT_TARGETS.backendBaseUrl, @@ -3699,6 +3704,79 @@ describe('SandboxControl mandatory worktree credentials', () => { }); describe('SandboxControl native worktree containment', () => { + it('keeps one member mapping and policy after an ambiguous bind retry', async () => { + const fixture = await credentialTerminalFixture('vercel'); + const { control, registration, socket, vercel } = fixture; + try { + const [prepared] = await storedGrants(control); + const proxyTargets = { + backendBaseUrl: 'https://worker.test', + providerBaseUrl: 'https://worker.test', + sessionIngestBaseUrl: 'https://worker.test', + }; + await runInDurableObject(control, async (_instance, state) => { + await saveSessionCredentialGrants(state.storage, [ + { + ...prepared, + kilo: { + ...prepared.kilo, + runtimeProxy: { targets: proxyTargets, members: [] }, + }, + }, + ]); + }); + const fence = await control.getRuntimeCredentialProxyFence({ + ownerId: registration.identity.userId, + sessionId: registration.identity.sessionId, + kiloSessionId: registration.auth.kiloSessionId, + directory: prepared.directory, + }); + if (!fence) throw new Error('Expected active runtime proxy fence'); + const memberHandle = await issueRuntimeCredentialProxyHandle( + fixture.environment, + createRuntimeProxyGrant({ + plane: 'control', + authorizationId: '11111111-1111-4111-8111-111111111111', + sessionId: registration.identity.sessionId, + kiloSessionId: registration.auth.kiloSessionId, + userId: registration.identity.userId, + ...(registration.identity.orgId ? { orgId: registration.identity.orgId } : {}), + mode: 'contained', + leaseExpiresAt: Date.now() + HOUR, + state: 'active', + ...fence, + }) + ); + const input = { + ownerId: registration.identity.userId, + sessionId: registration.identity.sessionId, + kiloSessionId: registration.auth.kiloSessionId, + directory: prepared.directory, + handle: memberHandle, + }; + + const first = await control.bindRuntimeCredentialProxyHandle(input); + const firstPolicy = vercel.runtime.policy; + const second = await control.bindRuntimeCredentialProxyHandle(input); + const [stored] = await storedGrants(control); + + expect(first).toEqual({ bound: true }); + expect(second).toEqual({ bound: true }); + expect(stored.kilo.runtimeProxy).toMatchObject({ + members: [ + { + sessionId: registration.identity.sessionId, + kiloSessionId: registration.auth.kiloSessionId, + handle: memberHandle, + }, + ], + }); + expect(vercel.runtime.policy).toEqual(firstPolicy); + } finally { + socket.close(); + } + }); + it('installs, refreshes, and removes the combined Vercel policy for exact worktree roots', async () => { const fixture = await credentialFixture('vercel'); const { control, registration, session, broker, vercel } = fixture; @@ -6437,6 +6515,8 @@ describe('SandboxControl passive status', () => { } satisfies SessionMessageRecord, ]); }); + const routing = Promise.withResolvers(); + const forwardingTasks: Promise[] = []; try { await runInDurableObject(stub, async (instance, state) => { await receiveHeartbeat(instance, state); @@ -6455,6 +6535,19 @@ describe('SandboxControl passive status', () => { const records = await state.storage.list(); const alarm = await state.storage.getAlarm(); const send = vi.spyOn(socket, 'send'); + const waitUntil = state.waitUntil.bind(state); + vi.spyOn(state, 'waitUntil').mockImplementation(promise => { + forwardingTasks.push(promise); + waitUntil(promise); + }); + // Route/physical reads can still be pending when the legacy event handler returns. + // Hold forwarding before enqueue so the persistence check cannot win that race. + const forward = fresh['forwardRoutedSessionFrame'].bind(fresh); + fresh['forwardRoutedSessionFrame'] = async (...args) => { + await routing.promise; + return forward(...args); + }; + await fresh.webSocketMessage( socket, JSON.stringify({ @@ -6467,7 +6560,8 @@ describe('SandboxControl passive status', () => { }, }) ); - await Promise.all(fresh['sessionForwarding'].values()); + expect([...fresh['sessionForwarding'].values()]).toEqual([]); + expect(forwardingTasks.length).toBeGreaterThan(0); expect(await fresh.getSandboxStatus(statusInput)).toMatchObject({ status: 'active', estimatedSleepAt: null, @@ -6496,6 +6590,9 @@ describe('SandboxControl passive status', () => { }); expect(renew).toHaveBeenCalledTimes(1); }); + routing.resolve(); + // The waitUntil task includes routing, enqueue, and the session persistence RPC. + await Promise.all(forwardingTasks); await runInDurableObject(session, async (_instance, state) => { const events = createEventQueries( drizzle(state.storage, { logger: false }), @@ -6508,6 +6605,8 @@ describe('SandboxControl passive status', () => { }); }); } finally { + routing.resolve(); + await Promise.all(forwardingTasks); ws.close(); } }); diff --git a/services/cloud-agent-next/test/integration/session/idle-reconciliation.test.ts b/services/cloud-agent-next/test/integration/session/idle-reconciliation.test.ts index a8a7e7f246..34cddba022 100644 --- a/services/cloud-agent-next/test/integration/session/idle-reconciliation.test.ts +++ b/services/cloud-agent-next/test/integration/session/idle-reconciliation.test.ts @@ -9,6 +9,10 @@ import { putSessionMessageState, } from '../../../src/session/session-message-state.js'; import { registerReadySession } from '../../helpers/session-setup.js'; +import { sealRuntimeAuthorization } from '@kilocode/worker-utils/runtime-authorization'; +import type { RuntimeAuthorization } from '@kilocode/worker-utils/runtime-authorization-contract'; +import { RUNTIME_PROXY_GRANT_KEY } from '../../../src/runtime-credential-proxy.js'; +import { RUNTIME_AUTHORIZATION_KEY } from '../../../src/session/runtime-authorization-persistence.js'; // Registered sessions leave dispatch, alarm, and fire-and-forget publication // work in the session DO. Interrupt every session a test touched, clear its @@ -56,6 +60,83 @@ describe('idle lifecycle integration', () => { ); }); + it('keeps queued work fenced and never retires or replaces expired authority', async () => { + const userId = 'user_recovery_queued'; + const sessionId = 'agent_recovery_queued'; + const oldAuthorization: RuntimeAuthorization = { + version: 1, + id: '00000000-0000-4000-8000-000000000101', + resourceKind: 'cloud-agent-next', + resourceId: sessionId, + userId, + authorizationUserId: userId, + issuedAt: '2026-01-01T00:00:00.000Z', + delegationExpiresAt: '2026-01-02T00:00:00.000Z', + state: 'active', + bindings: { + userPepperDigest: 'a'.repeat(64), + authorizationPepperDigest: 'b'.repeat(64), + userMembershipId: 'membership_1', + authorizationUserMembershipId: 'membership_1', + }, + source: { admissionSource: 'user' }, + }; + const fresh = { + ...oldAuthorization, + id: '00000000-0000-4000-8000-000000000102', + issuedAt: '2026-01-03T00:00:00.000Z', + delegationExpiresAt: '2026-01-04T00:00:00.000Z', + }; + const stub = env.CLOUD_AGENT_SESSION.get( + env.CLOUD_AGENT_SESSION.idFromName(`${userId}:${sessionId}`) + ); + const result = await runInDurableObject(stub, async instance => { + await registerReadySession(instance, { + sessionId, + userId, + kiloSessionId: 'aaaaaaaa-aaaa-4aaa-aaaa-aaaaaaaaaaaa', + prompt: 'initial prompt', + mode: 'code', + model: 'test-model', + kilocodeToken: 'expired-cached-token', + }); + await instance.ctx.storage.put(RUNTIME_AUTHORIZATION_KEY, oldAuthorization); + await instance.ctx.storage.put(RUNTIME_PROXY_GRANT_KEY, { cached: 'old-grant' }); + await putSessionMessageState(instance.ctx.storage, { + messageId: 'msg_018f1e2d3c4babcdefghijklmN', + status: 'queued', + prompt: 'queued', + createdAt: Date.now(), + }); + let stops = 0; + instance['physicalWrapperStopper'] = async () => { + stops += 1; + return { status: 'absent' }; + }; + const secret = + typeof env.NEXTAUTH_SECRET === 'string' + ? env.NEXTAUTH_SECRET + : await env.NEXTAUTH_SECRET.get(); + const outcome = await instance.recoverExpiredRuntimeAuthorization({ + ownerId: userId, + expectedOldId: oldAuthorization.id, + recoveryId: '00000000-0000-4000-8000-000000000103', + runtimeAuthorizationSeal: await sealRuntimeAuthorization(fresh, secret), + runtimeToken: 'fresh-cached-token', + }); + return { + outcome, + stops, + authorization: await instance.ctx.storage.get(RUNTIME_AUTHORIZATION_KEY), + grant: await instance.ctx.storage.get(RUNTIME_PROXY_GRANT_KEY), + }; + }); + expect(result.outcome).toEqual({ status: 'busy' }); + expect(result.stops).toBe(0); + expect(result.authorization).toMatchObject({ id: oldAuthorization.id }); + expect(result.grant).toEqual({ cached: 'old-grant' }); + }); + it('persists raw root idle without using it as a success boundary', async () => { const userId = 'user_idle_no_fallback'; const sessionId = 'agent_idle_no_fallback'; diff --git a/services/cloud-agent-next/test/unit/fake-llm-server.test.ts b/services/cloud-agent-next/test/unit/fake-llm-server.test.ts index d83aab6a80..6deb56d95b 100644 --- a/services/cloud-agent-next/test/unit/fake-llm-server.test.ts +++ b/services/cloud-agent-next/test/unit/fake-llm-server.test.ts @@ -406,6 +406,15 @@ describe('fake-llm-server HTTP', () => { expect(chunks[chunks.length - 1].data).toBe('[DONE]'); }); + it('error-terminal returns HTTP 400 with an OpenAI-shaped error', async () => { + const h = await start(); + const res = await postChat(h.url, '__fake__:error-terminal:boom'); + expect(res.status).toBe(400); + expect(await res.json()).toEqual({ + error: { message: 'boom', code: 400, type: 'invalid_request' }, + }); + }); + it('error scenario returns HTTP 402 with OpenAI-shaped error', async () => { const h = await start(); const res = await postChat(h.url, '__fake__:error:too broke'); @@ -1083,6 +1092,7 @@ type ProcFixture = { port: number; inode: string; directory: string; + home: string; roots: string[]; address?: string; state?: string; @@ -1139,6 +1149,8 @@ function discoveryFixture(processes: ProcFixture[], log?: string) { entry.command ?? ['/usr/local/bin/kilo', 'serve', '--hostname=127.0.0.1', '--port=0'] ).join('\0') + '\0' ); + const environmentEntry = entries.find(item => filename === `/proc/${item.pid}/environ`); + if (environmentEntry) return `HOME=${environmentEntry.home}\0PATH=/usr/local/bin\0`; if (filename === '/tmp/kilocode-control-wrapper.log' && log !== undefined) return log; forbiddenReads.push(filename); throw new Error('Unexpected filesystem read'); @@ -1230,6 +1242,7 @@ function directoryProcesses(): ProcFixture[] { port: 41001, inode: '501', directory: '/workspace/worktrees/worktree-a', + home: '/tmp/kilo-worktrees/a1b2c3d4', roots: ['ses_a', 'ses_sibling'], }, { @@ -1237,6 +1250,7 @@ function directoryProcesses(): ProcFixture[] { port: 41002, inode: '502', directory: '/workspace/worktrees/worktree-b', + home: '/tmp/kilo-worktrees/e5f6a7b8', roots: ['ses_b'], address: '00000000000000000000000001000000', }, @@ -1254,13 +1268,19 @@ describe('per-directory Kilo discovery', () => { container: { id: 'owned' }, processId: 101, directory: '/workspace/worktrees/worktree-a', + home: '/tmp/kilo-worktrees/a1b2c3d4', serverUrl: 'http://127.0.0.1:41001', }); - expect(sibling).toMatchObject({ processId: first?.processId, directory: first?.directory }); + expect(sibling).toMatchObject({ + processId: first?.processId, + directory: first?.directory, + home: first?.home, + }); expect(second).toMatchObject({ container: { id: 'owned' }, processId: 202, directory: '/workspace/worktrees/worktree-b', + home: '/tmp/kilo-worktrees/e5f6a7b8', serverUrl: 'http://[::1]:41002', }); expect(first?.logPath).toBeUndefined(); diff --git a/services/cloud-agent-next/test/unit/smoke-cleanup.test.ts b/services/cloud-agent-next/test/unit/smoke-cleanup.test.ts new file mode 100644 index 0000000000..0dd0352025 --- /dev/null +++ b/services/cloud-agent-next/test/unit/smoke-cleanup.test.ts @@ -0,0 +1,140 @@ +import { afterEach, describe, expect, it, vi } from 'vitest'; +import { DEFAULT_CONFIG, startSession, type DriverConfig } from '../e2e/client.js'; +import { cleanupOwnedSessions } from '../e2e/smoke-cleanup.js'; + +vi.mock('../e2e/auth.js', () => ({ mintApiToken: () => 'test-token' })); + +afterEach(() => vi.unstubAllGlobals()); + +describe('owned session cleanup', () => { + it.each([ + 'No accepted wrapper messages or pending queued messages', + 'No session work to interrupt', + ])('allows the documented idle response: %s', async message => { + const stop = vi.fn().mockResolvedValue(undefined); + await cleanupOwnedSessions(new Set(['idle']), { + interrupt: async () => ({ success: false, message }), + stopOwnedSandboxes: stop, + }); + expect(stop.mock.calls).toEqual([['idle']]); + }); + + it('cancels pending demand before teardown so a failed scenario cannot respawn', async () => { + const pending = new Set(['owned-a', 'owned-b', 'unrelated']); + const actions: string[] = []; + const recreated: string[] = []; + await cleanupOwnedSessions(new Set(['owned-a', 'owned-b']), { + interrupt: async id => { + actions.push(`interrupt:${id}`); + pending.delete(id); + return { success: true }; + }, + stopOwnedSandboxes: async id => { + actions.push(`stop:${id}`); + if (pending.has(id)) recreated.push(id); + }, + }); + expect(actions).toEqual([ + 'interrupt:owned-a', + 'interrupt:owned-b', + 'stop:owned-a', + 'stop:owned-b', + ]); + expect(recreated).toEqual([]); + expect([...pending]).toEqual(['unrelated']); + }); + + it('reports failed cancellation, skips its teardown, and cleans other owned sessions', async () => { + const stop = vi.fn().mockResolvedValue(undefined); + await expect( + cleanupOwnedSessions(new Set(['failed', 'idle']), { + interrupt: async id => + id === 'failed' + ? { success: false, message: 'interruption failed' } + : { + success: false, + message: 'No accepted wrapper messages or pending queued messages', + }, + stopOwnedSandboxes: stop, + }) + ).rejects.toThrow('Owned session cleanup failed'); + expect(stop.mock.calls).toEqual([['idle']]); + }); + + it('does not tear down a session after an interruption transport failure', async () => { + const stop = vi.fn(); + await expect( + cleanupOwnedSessions(new Set(['owned']), { + interrupt: async () => { + throw new Error('transport failed'); + }, + stopOwnedSandboxes: stop, + }) + ).rejects.toThrow('Owned session cleanup failed'); + expect(stop).not.toHaveBeenCalled(); + }); +}); + +describe('smoke session ownership tracking', () => { + function config(onSessionCreated: (sessionId: string) => void): DriverConfig { + return { + ...DEFAULT_CONFIG, + user: { id: 'matrix-user', email: 'matrix@example.test', api_token_pepper: 'test' }, + nextAuthSecret: 'test', + internalApiSecret: 'test', + onSessionCreated, + }; + } + + it('retains the owned unified session when a post-start assertion fails', async () => { + const owned = new Set(); + vi.stubGlobal( + 'fetch', + vi.fn().mockResolvedValue( + Response.json({ + result: { data: { cloudAgentSessionId: 'agent_owned' } }, + }) + ) + ); + await expect( + startSession( + { + ...config(id => owned.add(id)), + expectControlPlane: true, + }, + { prompt: 'test' } + ) + ).rejects.toThrow('expected an enrolled'); + expect([...owned]).toEqual(['agent_owned']); + }); + + it('retains a legacy prepared session when initiation fails', async () => { + const owned = new Set(); + vi.stubGlobal( + 'fetch', + vi + .fn() + .mockResolvedValueOnce( + Response.json({ + result: { + data: { + cloudAgentSessionId: 'agent_prepared', + kiloSessionId: 'ses_prepared', + }, + }, + }) + ) + .mockResolvedValueOnce( + Response.json({ error: { message: 'initiation failed' } }, { status: 503 }) + ) + ); + await expect( + startSession( + config(id => owned.add(id)), + { prompt: 'test' }, + 'legacy' + ) + ).rejects.toThrow('initiation failed'); + expect([...owned]).toEqual(['agent_prepared']); + }); +}); diff --git a/services/cloud-agent-next/test/unit/wrapper/worktree-credential-refresh.test.ts b/services/cloud-agent-next/test/unit/wrapper/worktree-credential-refresh.test.ts index f8ff8777e6..5e4603709f 100644 --- a/services/cloud-agent-next/test/unit/wrapper/worktree-credential-refresh.test.ts +++ b/services/cloud-agent-next/test/unit/wrapper/worktree-credential-refresh.test.ts @@ -834,8 +834,10 @@ describe('direct worktree credential refresh', () => { signal: runtimeLifetime.signal, }; const getRuntime = f.registry.get.bind(f.registry); - f.registry.get = directory => - directory === sibling.directory ? (fakeRuntime as typeof runtime) : getRuntime(directory); + f.registry.get = request => + (typeof request === 'string' ? request : request.directory) === sibling.directory + ? (fakeRuntime as typeof runtime) + : getRuntime(request); const siblingOp = deps.operations.start( sibling, undefined, diff --git a/services/cloud-agent-next/worker-configuration.d.ts b/services/cloud-agent-next/worker-configuration.d.ts index f79b1368dc..036c697413 100644 --- a/services/cloud-agent-next/worker-configuration.d.ts +++ b/services/cloud-agent-next/worker-configuration.d.ts @@ -1,5 +1,5 @@ /* eslint-disable */ -// Generated by Wrangler by running `wrangler types` (hash: 8d45d333d1f69f2f8355e9778e2fd5e3) +// Generated by Wrangler by running `wrangler types` (hash: 2af2de4738120214879df56be9e60d57) // Runtime types generated with workerd@1.20260714.1 2026-06-03 nodejs_compat interface __BaseEnv_Env { SHARED_SANDBOX_OVERRIDES: KVNamespace; @@ -9,35 +9,48 @@ interface __BaseEnv_Env { CALLBACK_QUEUE: Queue; CLOUD_AGENT_REPORT_QUEUE: Queue; INTERNAL_API_SECRET_PROD: SecretsStoreSecret; - KILOCODE_BACKEND_BASE_URL: "http://localhost:3000" | "https://api.kilo.ai"; - KILO_OPENROUTER_BASE?: "http://localhost:3000/api"; GITHUB_APP_SLUG: "kiloconnect-development" | "kiloconnect"; GITHUB_APP_BOT_USER_ID: "242397087" | "240665456"; GITHUB_LITE_APP_SLUG: "" | "kiloconnect-lite"; GITHUB_LITE_APP_BOT_USER_ID: "" | "257753004"; - WORKER_URL: "http://localhost:8794" | "https://cloud-agent-next.kilosessions.ai"; SANDBOX_TRANSPORT?: "rpc"; CLI_TIMEOUT_SECONDS: "900"; REAPER_INTERVAL_MS: "300000"; R2_ATTACHMENTS_BUCKET: "cloud-agent-attachments-dev" | "cloud-agent-attachments"; BACKUP_BUCKET_NAME: "kilocode-sessions-dev" | "kilocode-sessions"; CLOUDFLARE_R2_ACCOUNT_ID: "e115e769bcdd4c3d66af59d3332cb394"; - WS_ALLOWED_ORIGINS: "http://localhost:3000,http://host.docker.internal:3000" | "https://app.kilo.ai,https://api.kilo.ai"; - KILO_SESSION_INGEST_URL: "http://localhost:8800" | "https://ingest.kilosessions.ai"; PER_SESSION_SANDBOX_ORG_IDS?: "*"; CREDENTIAL_CONTAINMENT_ENABLED: "false" | "true"; REPO_SNAPSHOT_ORG_IDS?: ""; CONTAINER_BILLING_HEARTBEAT_SECONDS: "60" | "300"; - CONTROL_PLANE_IDS: ""; - WORKTREE_CREATION_ENABLED_IDS: ""; + CONTROL_PLANE_IDS: "*" | ""; + WORKTREE_CREATION_ENABLED_IDS: "*" | ""; + RUNTIME_ISOLATION_ENABLED: "false"; VERCEL_SANDBOX_ORG_IDS: ""; VERCEL_PROJECT_ID: ""; VERCEL_TEAM_ID: ""; VERCEL_SANDBOX_SNAPSHOT_ID: ""; VERCEL_SANDBOX_RUNTIME_BUILD_ID: ""; VERCEL_SANDBOX_RUNTIME: "node24"; - VERCEL_SANDBOX_INITIAL_TIMEOUT_MS: ""; - VERCEL_SANDBOX_EXTEND_DURATION_MS: ""; + NEXTAUTH_SECRET: string; + INTERNAL_API_SECRET: string; + KILOCODE_BACKEND_BASE_URL: string; + KILO_OPENROUTER_BASE: string; + CLOUD_AGENT_CONTAINER_BILLING_ENABLED: string; + CLOUD_AGENT_CONTAINER_BILLING_USER_IDS: string; + CLOUD_AGENT_CONTAINER_BILLING_ORG_IDS: string; + WORKER_URL: string; + AGENT_ENV_VARS_PRIVATE_KEY: string; + R2_ENDPOINT: string; + R2_ATTACHMENTS_READONLY_ACCESS_KEY_ID: string; + R2_ATTACHMENTS_READONLY_SECRET_ACCESS_KEY: string; + R2_ACCESS_KEY_ID: string; + R2_SECRET_ACCESS_KEY: string; + KILO_SESSION_INGEST_URL: string; + LOG_REJECTED_KILO_URLS: string; + WS_ALLOWED_ORIGINS: string; + VERCEL_SANDBOX_INITIAL_TIMEOUT_MS: string; + VERCEL_SANDBOX_EXTEND_DURATION_MS: string; Sandbox: DurableObjectNamespace; SandboxSmall: DurableObjectNamespace; SandboxDIND: DurableObjectNamespace; @@ -56,9 +69,6 @@ interface __BaseEnv_Env { CONTAINER_USAGE_METER: Service /* entrypoint ContainerUsageMeter from container-usage-meter */; TOOL_CGROUP_RESERVE_MB?: "2048"; TOOL_CGROUP_CPU_WEIGHT?: "50"; - CLOUD_AGENT_CONTAINER_BILLING_ENABLED?: "true"; - CLOUD_AGENT_CONTAINER_BILLING_USER_IDS?: "f1a848ca-bade-48d8-a5ad-1042d08651e6"; - CLOUD_AGENT_CONTAINER_BILLING_ORG_IDS?: "9d278969-5453-4ae3-a51f-a8d2274a7b56"; } declare namespace Cloudflare { interface GlobalProps { @@ -73,35 +83,48 @@ declare namespace Cloudflare { CALLBACK_QUEUE: Queue; CLOUD_AGENT_REPORT_QUEUE: Queue; INTERNAL_API_SECRET_PROD: SecretsStoreSecret; - KILOCODE_BACKEND_BASE_URL: "http://localhost:3000"; - KILO_OPENROUTER_BASE: "http://localhost:3000/api"; GITHUB_APP_SLUG: "kiloconnect-development"; GITHUB_APP_BOT_USER_ID: "242397087"; GITHUB_LITE_APP_SLUG: ""; GITHUB_LITE_APP_BOT_USER_ID: ""; - WORKER_URL: "http://localhost:8794"; SANDBOX_TRANSPORT: "rpc"; CLI_TIMEOUT_SECONDS: "900"; REAPER_INTERVAL_MS: "300000"; R2_ATTACHMENTS_BUCKET: "cloud-agent-attachments-dev"; BACKUP_BUCKET_NAME: "kilocode-sessions-dev"; CLOUDFLARE_R2_ACCOUNT_ID: "e115e769bcdd4c3d66af59d3332cb394"; - WS_ALLOWED_ORIGINS: "http://localhost:3000,http://host.docker.internal:3000"; - KILO_SESSION_INGEST_URL: "http://localhost:8800"; PER_SESSION_SANDBOX_ORG_IDS: "*"; CREDENTIAL_CONTAINMENT_ENABLED: "false"; REPO_SNAPSHOT_ORG_IDS: ""; CONTAINER_BILLING_HEARTBEAT_SECONDS: "60"; - CONTROL_PLANE_IDS: ""; - WORKTREE_CREATION_ENABLED_IDS: ""; + CONTROL_PLANE_IDS: "*"; + WORKTREE_CREATION_ENABLED_IDS: "*"; + RUNTIME_ISOLATION_ENABLED: "false"; VERCEL_SANDBOX_ORG_IDS: ""; VERCEL_PROJECT_ID: ""; VERCEL_TEAM_ID: ""; VERCEL_SANDBOX_SNAPSHOT_ID: ""; VERCEL_SANDBOX_RUNTIME_BUILD_ID: ""; VERCEL_SANDBOX_RUNTIME: "node24"; - VERCEL_SANDBOX_INITIAL_TIMEOUT_MS: ""; - VERCEL_SANDBOX_EXTEND_DURATION_MS: ""; + NEXTAUTH_SECRET: string; + INTERNAL_API_SECRET: string; + KILOCODE_BACKEND_BASE_URL: string; + KILO_OPENROUTER_BASE: string; + CLOUD_AGENT_CONTAINER_BILLING_ENABLED: string; + CLOUD_AGENT_CONTAINER_BILLING_USER_IDS: string; + CLOUD_AGENT_CONTAINER_BILLING_ORG_IDS: string; + WORKER_URL: string; + AGENT_ENV_VARS_PRIVATE_KEY: string; + R2_ENDPOINT: string; + R2_ATTACHMENTS_READONLY_ACCESS_KEY_ID: string; + R2_ATTACHMENTS_READONLY_SECRET_ACCESS_KEY: string; + R2_ACCESS_KEY_ID: string; + R2_SECRET_ACCESS_KEY: string; + KILO_SESSION_INGEST_URL: string; + LOG_REJECTED_KILO_URLS: string; + WS_ALLOWED_ORIGINS: string; + VERCEL_SANDBOX_INITIAL_TIMEOUT_MS: string; + VERCEL_SANDBOX_EXTEND_DURATION_MS: string; Sandbox: DurableObjectNamespace; SandboxSmall: DurableObjectNamespace; SandboxDIND: DurableObjectNamespace; @@ -126,7 +149,7 @@ type StringifyValues> = { [Binding in keyof EnvType]: EnvType[Binding] extends string ? EnvType[Binding] : string; }; declare namespace NodeJS { - interface ProcessEnv extends StringifyValues> {} + interface ProcessEnv extends StringifyValues> {} } declare module "*.sql" { const value: string; diff --git a/services/cloud-agent-next/wrangler.jsonc b/services/cloud-agent-next/wrangler.jsonc index bb5b41885b..a064065fc8 100644 --- a/services/cloud-agent-next/wrangler.jsonc +++ b/services/cloud-agent-next/wrangler.jsonc @@ -66,6 +66,7 @@ "CONTAINER_BILLING_HEARTBEAT_SECONDS": "300", "CONTROL_PLANE_IDS": "", "WORKTREE_CREATION_ENABLED_IDS": "", + "RUNTIME_ISOLATION_ENABLED": "false", "VERCEL_SANDBOX_ORG_IDS": "", "VERCEL_PROJECT_ID": "", "VERCEL_TEAM_ID": "", @@ -457,6 +458,7 @@ "CONTAINER_BILLING_HEARTBEAT_SECONDS": "60", "CONTROL_PLANE_IDS": "*", "WORKTREE_CREATION_ENABLED_IDS": "*", + "RUNTIME_ISOLATION_ENABLED": "false", "VERCEL_SANDBOX_ORG_IDS": "", "VERCEL_PROJECT_ID": "", "VERCEL_TEAM_ID": "", diff --git a/services/cloud-agent-next/wrangler.test.jsonc b/services/cloud-agent-next/wrangler.test.jsonc index b3dfdb01b1..ced0e63aa1 100644 --- a/services/cloud-agent-next/wrangler.test.jsonc +++ b/services/cloud-agent-next/wrangler.test.jsonc @@ -4,6 +4,9 @@ "main": "test/test-worker.ts", "compatibility_date": "2026-06-03", "compatibility_flags": ["nodejs_compat"], + "vars": { + "NEXTAUTH_SECRET": "cloud-agent-integration-test-secret", + }, "rules": [{ "type": "Text", "globs": ["**/*.sql"], "fallthrough": true }], "durable_objects": { "bindings": [ diff --git a/services/cloud-agent-next/wrapper/src/control/apply-attach.test.ts b/services/cloud-agent-next/wrapper/src/control/apply-attach.test.ts index c55c46291c..99c631b209 100644 --- a/services/cloud-agent-next/wrapper/src/control/apply-attach.test.ts +++ b/services/cloud-agent-next/wrapper/src/control/apply-attach.test.ts @@ -95,12 +95,15 @@ function fakeKiloRuntimes(overrides: Partial = {}): WorktreeK ...overrides, } as WrapperKiloClient; const runtimes = new Map(); + const key = (identity: typeof session) => + `${identity.sessionId}\0${identity.kiloSessionId}\0${identity.directory}`; return { - attach(identity, auth, environment) { + attach(identity, auth, environment, _canRefreshCredentials) { const { directory } = identity; - let runtime = runtimes.get(directory); + let runtime = runtimes.get(key(identity)); if (!runtime) { runtime = { + identity: { ...identity }, runtimeId: 'native_1', directory, scopeId: auth.scopeId, @@ -114,7 +117,7 @@ function fakeKiloRuntimes(overrides: Partial = {}): WorktreeK kiloClient, signal: new AbortController().signal, }; - runtimes.set(directory, runtime); + runtimes.set(key(identity), runtime); } return { ready: Promise.resolve(runtime), @@ -125,11 +128,14 @@ function fakeKiloRuntimes(overrides: Partial = {}): WorktreeK }; }, detach: () => true, + retireForRecovery: async () => 'retired', deleteDirectory: async directory => { - runtimes.delete(directory); + for (const [key, runtime] of runtimes) { + if (runtime.directory === directory) runtimes.delete(key); + } }, retireRuntime: async (directory, _deadlineAt, target) => { - const runtime = runtimes.get(directory); + const runtime = [...runtimes.values()].find(runtime => runtime.directory === directory); if ( !runtime || !target || @@ -137,13 +143,21 @@ function fakeKiloRuntimes(overrides: Partial = {}): WorktreeK target.client !== runtime.kiloClient ) return 'stale'; - runtimes.delete(directory); + if (runtime?.identity) runtimes.delete(key(runtime.identity)); return 'retired'; }, verifyQuiescence: async (directory, target, deadlineAt) => - runtimes.get(directory)?.kiloClient === target.client && Date.now() < deadlineAt, - getRetained: directory => runtimes.get(directory), - get: directory => runtimes.get(directory), + [...runtimes.values()].some( + runtime => runtime.directory === directory && runtime.kiloClient === target.client + ) && Date.now() < deadlineAt, + getRetained: directory => + [...runtimes.values()].find(runtime => runtime.directory === directory), + get: identity => + typeof identity === 'string' + ? [...runtimes.values()].find(runtime => runtime.directory === identity) + : runtimes.get(key(identity)), + getAll: directory => [...runtimes.values()].filter(runtime => runtime.directory === directory), + isCurrent: runtime => [...runtimes.values()].includes(runtime), isHealthy: () => true, shutdown: () => {}, }; @@ -1428,7 +1442,7 @@ describe('applySessionAttach', () => { const snapshotIdentity = 'snapshot_other'; const steps: string[] = []; const assertRegistration = () => { - const runtime = runtimes.get(session.directory); + const runtime = runtimes.get(session); if (!runtime) throw new Error('Expected worktree runtime'); const storage = path.join(runtime.env.XDG_DATA_HOME, 'kilo', 'storage', 'session_share'); expect( @@ -1511,7 +1525,7 @@ describe('applySessionAttach', () => { const runtimes = isolatedKiloRuntimes(); const deps = { ...noFs, kiloRuntimes: runtimes }; expect(await applySessionAttach(first, { kilo }, deps)).toMatchObject({ ok: true }); - const runtime = runtimes.get(directory); + const runtime = runtimes.get(first); const failed = await applySessionAttach( sibling, { kilo }, @@ -1529,7 +1543,7 @@ describe('applySessionAttach', () => { expect(failed).toMatchObject({ ok: false }); expect(rootForSession(sibling.kiloSessionId)).toBeUndefined(); expect(rootForSession(first.kiloSessionId)).toBe(first.kiloSessionId); - expect(runtimes.get(directory)).toBe(runtime); + expect(runtimes.get(first)).toBe(runtime); expect(runtime?.signal.aborted).toBe(false); expect(runtimes.detach(sibling)).toBe(false); expect(await applySessionAttach(sibling, { kilo }, deps)).toMatchObject({ ok: true }); @@ -1541,7 +1555,7 @@ describe('applySessionAttach', () => { const deps = { ...noFs, kiloRuntimes: runtimes }; expect(await applySessionAttach(identity, { kilo }, deps)).toMatchObject({ ok: true }); rememberChildSession({ childId: 'child_existing', parentId: identity.kiloSessionId }); - const runtime = runtimes.get(identity.directory); + const runtime = runtimes.get(identity); expect( await applySessionAttach( identity, @@ -1560,11 +1574,11 @@ describe('applySessionAttach', () => { ).toMatchObject({ ok: false }); expect(rootForSession(identity.kiloSessionId)).toBe(identity.kiloSessionId); expect(rootForSession('child_existing')).toBe(identity.kiloSessionId); - expect(runtimes.get(identity.directory)).toBe(runtime); + expect(runtimes.get(identity)).toBe(runtime); expect(runtime?.signal.aborted).toBe(false); }); - it('retires a cancelled pending root and keeps the original immutable grant for a sibling', async () => { + it('retires a cancelled pending root without coupling a sibling identity to its grant', async () => { const directory = path.join(homeRoot, 'shared'); const identity = { ...session, directory }; const sibling = { ...siblingSession, directory }; @@ -1575,7 +1589,7 @@ describe('applySessionAttach', () => { const grant = { ...kilo, targets: { ...kilo.targets } }; const attaching = applySessionAttach( identity, - { kilo: grant }, + { kilo: grant, runtimeIsolation: 'per-session' }, { ...noFs, kiloRuntimes: runtimes, @@ -1595,22 +1609,23 @@ describe('applySessionAttach', () => { ); try { await restoring.promise; - const runtime = runtimes.get(directory); + const runtime = runtimes.get(identity); grant.token = 'mutated-guest'; grant.targets.sessionIngestBaseUrl = 'https://other.example.test'; const deps = { ...noFs, kiloRuntimes: runtimes }; - expect(await applySessionAttach(sibling, { kilo: grant }, deps)).toMatchObject({ - ok: false, - error: { code: 'unauthorized' }, - }); - expect(await applySessionAttach(sibling, { kilo }, deps)).toMatchObject({ ok: true }); + expect( + await applySessionAttach(sibling, { kilo: grant, runtimeIsolation: 'per-session' }, deps) + ).toMatchObject({ ok: true }); + const siblingRuntime = runtimes.get(sibling); + expect(siblingRuntime).toBeDefined(); + expect(siblingRuntime).not.toBe(runtime); const storage = path.join( runtime?.env.XDG_DATA_HOME ?? '', 'kilo', 'storage', 'session_share' ); - for (const id of [identity.kiloSessionId, sibling.kiloSessionId]) { + for (const id of [identity.kiloSessionId]) { expect(JSON.parse(fs.readFileSync(path.join(storage, `${id}.json`), 'utf8'))).toEqual({ id, ingestPath: `/api/session/${id}/ingest`, @@ -1622,12 +1637,12 @@ describe('applySessionAttach', () => { expect(runtimes.detach(identity)).toBe(false); expect(rootForSession(identity.kiloSessionId)).toBeUndefined(); expect(rootForSession(sibling.kiloSessionId)).toBe(sibling.kiloSessionId); - expect(runtimes.get(directory)).toBe(runtime); - expect(runtime?.env.KILOCODE_TOKEN).toBe(kilo.token); - expect(runtime?.env.KILO_SESSION_INGEST_URL).toBe(kilo.targets.sessionIngestBaseUrl); - expect(runtime?.signal.aborted).toBe(false); + expect(runtimes.get(sibling)).toBe(siblingRuntime); + expect(siblingRuntime?.env.KILOCODE_TOKEN).toBe('mutated-guest'); + expect(siblingRuntime?.env.KILO_SESSION_INGEST_URL).toBe('https://other.example.test'); + expect(siblingRuntime?.signal.aborted).toBe(false); expect(runtimes.detach(sibling)).toBe(true); - expect(runtime?.signal.aborted).toBe(true); + expect(siblingRuntime?.signal.aborted).toBe(true); } finally { release.resolve(); await attaching; @@ -1710,7 +1725,7 @@ describe('applySessionAttach', () => { const runtimes = fakeKiloRuntimes(); // Pre-attach to create the runtime so defaultSessionExists can find kiloClient.serverUrl runtimes.attach(session, kilo, {}); - const runtime = runtimes.get(session.directory); + const runtime = runtimes.get(session); if (!runtime) throw new Error('Expected runtime'); const server = Bun.serve({ port: 0, @@ -1979,7 +1994,7 @@ describe('applySessionAttach', () => { { ...noFs, kiloRuntimes: runtimes, sessionExists: async () => true } ); expect(result).toEqual({ ok: true, result: { attached: true } }); - const home = runtimes.get(directory)?.env.HOME; + const home = runtimes.get({ ...session, directory })?.env.HOME; expect(home).toStartWith(homeRoot); expect(fs.readFileSync(path.join(directory, 'setup-env.txt'), 'utf8')).toBe( `${home}\n${kilo.token}\nabsent\nabsent\nopaque-bitbucket-token\nacme-workspace\nwidgets\n{33333333-3333-4333-8333-333333333333}\n{11111111-1111-4111-8111-111111111111}\n` diff --git a/services/cloud-agent-next/wrapper/src/control/apply-attach.ts b/services/cloud-agent-next/wrapper/src/control/apply-attach.ts index 4209c4db6a..4649a54c75 100644 --- a/services/cloud-agent-next/wrapper/src/control/apply-attach.ts +++ b/services/cloud-agent-next/wrapper/src/control/apply-attach.ts @@ -358,6 +358,7 @@ async function executeSessionAttach( attach.kilo, attach.env, deps.canRefreshCredentials, + attach.runtimeIsolation, deps.onMutation, deps.onCleanupTarget ); diff --git a/services/cloud-agent-next/wrapper/src/control/control-event-transport.test.ts b/services/cloud-agent-next/wrapper/src/control/control-event-transport.test.ts index a93d75e23d..e63bf1ddc9 100644 --- a/services/cloud-agent-next/wrapper/src/control/control-event-transport.test.ts +++ b/services/cloud-agent-next/wrapper/src/control/control-event-transport.test.ts @@ -279,6 +279,32 @@ describe('native-scoped control event failures', () => { } }); + it('selects the failed native runtime among isolated roots in the same directory', () => { + const first = { runtimeId: 'native_first' }; + const second = { runtimeId: 'native_second' }; + const retired = mock(); + const getRuntime = mock((directory: string, nativeRuntimeId: string) => + directory === session.directory + ? [first, second].find(runtime => runtime.runtimeId === nativeRuntimeId) + : undefined + ); + const handleFailure = createControlEventFailureHandler({ getRuntime, onFailure: retired }); + const failure: ControlEventOutboxFailure = { + reason: 'expired', + publication: { + event: 'session.event', + receiptId: 'receipt_second', + sequence: 1, + session: { ...session, nativeRuntimeId: second.runtimeId }, + payload, + }, + }; + handleFailure(failure); + expect(getRuntime).toHaveBeenCalledWith(session.directory, second.runtimeId); + expect(retired).toHaveBeenCalledWith(failure, second); + expect(retired).toHaveBeenCalledTimes(1); + }); + it('reports failures without native identity without guessing the current runtime', async () => { const retired = mock(); const getRuntime = mock(() => ({ runtimeId: crypto.randomUUID() })); diff --git a/services/cloud-agent-next/wrapper/src/control/control-event-transport.ts b/services/cloud-agent-next/wrapper/src/control/control-event-transport.ts index 7ebbe3f199..e9c55953f0 100644 --- a/services/cloud-agent-next/wrapper/src/control/control-event-transport.ts +++ b/services/cloud-agent-next/wrapper/src/control/control-event-transport.ts @@ -8,7 +8,7 @@ import { type EventKind = 'session.event' | 'session.preparing'; export function createControlEventFailureHandler(options: { - getRuntime: (directory: string) => Runtime | undefined; + getRuntime: (directory: string, nativeRuntimeId: string) => Runtime | undefined; onFailure: (failure: ControlEventOutboxFailure, runtime: Runtime) => void; }) { const failedRuntimes = new WeakSet(); @@ -17,7 +17,7 @@ export function createControlEventFailureHandler true, + retireForRecovery: async () => 'retired', deleteDirectory: async () => {}, - getRetained: directory => (directory === runtime.directory ? runtime : undefined), + getRetained: directory => + directory === runtime.directory && !nativeLifetime.signal.aborted ? runtime : undefined, retireRuntime: async (directory, _deadlineAt, target) => { if ( directory !== runtime.directory || @@ -137,8 +140,15 @@ export function createHandlerFixture( target.client === runtime.kiloClient && !nativeLifetime.signal.aborted && Date.now() < deadlineAt, - get: directory => - directory === runtime.directory && !nativeLifetime.signal.aborted ? runtime : undefined, + get: request => + (typeof request === 'string' + ? request === runtime.directory + : request.directory === runtime.directory) && !nativeLifetime.signal.aborted + ? runtime + : undefined, + getAll: directory => + directory === runtime.directory && !nativeLifetime.signal.aborted ? [runtime] : [], + isCurrent: candidate => candidate === runtime && !nativeLifetime.signal.aborted, isHealthy: () => true, shutdown: () => {}, } diff --git a/services/cloud-agent-next/wrapper/src/control/delete-worktree.test.ts b/services/cloud-agent-next/wrapper/src/control/delete-worktree.test.ts index 764b637e5b..69a63114dc 100644 --- a/services/cloud-agent-next/wrapper/src/control/delete-worktree.test.ts +++ b/services/cloud-agent-next/wrapper/src/control/delete-worktree.test.ts @@ -78,7 +78,7 @@ function fixture() { aborted, client, deps: { - client, + clients: [client], assertDirectory: async () => undefined, retireDirectory: async (dir: string) => { retired.push(dir); @@ -165,7 +165,7 @@ describe('Kilo 7.4.20 cleanup HTTP compatibility', () => { try { const deletion = deleteWorktree(input, { ...f.deps, - client: createWorktreeKiloCleanupClient(server.url.toString()), + clients: [createWorktreeKiloCleanupClient(server.url.toString())], }); await abortStarted.promise; expect(f.sessions.get(sessionId(0))?.active).toBe(true); @@ -240,6 +240,71 @@ describe('Kilo session DELETE confirmation', () => { }); describe('scoped worktree runtime deletion', () => { + test('discovers and cleans every owning runtime client before retiring the checkout', async () => { + const f = fixture(); + const events: string[] = []; + const createClient = (name: string, id: string): WorktreeKiloCleanupClient => { + const sessions = new Map([[id, { id, directory }]]); + return { + listSessionIds: async dir => (dir === directory ? [...sessions.keys()] : []), + getSession: async (_dir, sessionId) => sessions.get(sessionId) ?? null, + children: async () => [], + abortSession: async (_dir, sessionId) => { + events.push(`${name}:abort:${sessionId}`); + }, + stopSessionProcesses: async (_dir, sessionId) => { + events.push(`${name}:stop:${sessionId}`); + }, + deleteSession: async (_dir, sessionId) => { + events.push(`${name}:delete:${sessionId}`); + sessions.delete(sessionId); + }, + closeTerminals: async () => { + events.push(`${name}:terminals`); + }, + disposeDirectory: async () => { + events.push(`${name}:dispose`); + }, + }; + }; + const first = createClient('first', sessionId(0)); + const second = createClient('second', sessionId(1)); + + expect( + await deleteWorktree( + { worktreeId, directory, sessionIds: [sessionId(0), sessionId(1)] }, + { + ...f.deps, + clients: [first, second], + detachTerminals: async () => { + events.push('wrapper:terminals'); + }, + retireDirectory: async () => { + events.push('runtime:retire'); + }, + removeDirectory: async () => { + events.push('checkout:remove'); + }, + } + ) + ).toEqual({ deleted: true, sessionIds: [sessionId(0), sessionId(1)] }); + expect(events).toEqual([ + `first:abort:${sessionId(0)}`, + `second:abort:${sessionId(1)}`, + `first:stop:${sessionId(0)}`, + `second:stop:${sessionId(1)}`, + 'wrapper:terminals', + 'first:terminals', + `first:delete:${sessionId(0)}`, + 'second:terminals', + `second:delete:${sessionId(1)}`, + 'first:dispose', + 'second:dispose', + 'runtime:retire', + 'checkout:remove', + ]); + }); + test.each([true, false])( 'awaits terminal detachment and runtime retirement before removing the checkout: live=%s', async live => { @@ -259,7 +324,7 @@ describe('scoped worktree runtime deletion', () => { const input = { worktreeId, directory, sessionIds: [sessionId(0), sessionId(1)] }; const deletion = deleteWorktree(input, { ...f.deps, - client: live ? f.client : undefined, + clients: live ? [f.client] : [], detachTerminals: async dir => { expect(dir).toBe(directory); detaching.resolve(); @@ -316,7 +381,7 @@ describe('scoped worktree runtime deletion', () => { rememberAttachedRoot(sessionId(0), directory); rememberAttachedRoot(sessionId(2), otherDirectory); const input = { worktreeId, directory, sessionIds: [sessionId(0)] }; - const deps = { ...f.deps, client: undefined }; + const deps = { ...f.deps, clients: [] }; expect(await prepareWorktreeDeletion(input, deps)).toEqual(input.sessionIds); for (let attempt = 0; attempt < 2; attempt++) { expect(await deleteWorktree(input, deps)).toEqual({ @@ -344,7 +409,7 @@ describe('scoped worktree runtime deletion', () => { await rejects( deleteWorktree(input, { ...f.deps, - client: live ? f.client : undefined, + clients: live ? [f.client] : [], retireDirectory: async () => { throw new Error('Runtime retirement failed'); }, @@ -353,7 +418,7 @@ describe('scoped worktree runtime deletion', () => { ); expect(f.directories.has(directory)).toBe(true); expect(rootForSession(sessionId(0))).toBe(sessionId(0)); - expect(await deleteWorktree(input, { ...f.deps, client: undefined })).toEqual({ + expect(await deleteWorktree(input, { ...f.deps, clients: [] })).toEqual({ deleted: true, sessionIds: input.sessionIds, }); @@ -370,7 +435,7 @@ describe('scoped worktree runtime deletion', () => { { worktreeId, directory, sessionIds: [sessionId(0)] }, { ...f.deps, - client: undefined, + clients: [], detachTerminals: async () => { throw new Error('Terminal detachment failed'); }, @@ -393,7 +458,7 @@ describe('scoped worktree runtime deletion', () => { await rejects( deleteWorktree( { worktreeId, directory, sessionIds: [] }, - { ...f.deps, client: undefined, assertDirectory: undefined } + { ...f.deps, clients: [], assertDirectory: undefined } ), /Invalid worktree directory/ ); @@ -471,6 +536,7 @@ describe('scoped worktree runtime deletion', () => { deletePty: async (id, dir) => ptys.get(id)?.cwd === dir && ptys.delete(id), } as WrapperKiloClient; const kiloRuntime: WorktreeKiloRuntime = { + identity: roots[0], scopeId: 'test-scope', runtimeId: 'native_1', directory, @@ -479,6 +545,7 @@ describe('scoped worktree runtime deletion', () => { signal: new AbortController().signal, }; const otherKiloRuntime: WorktreeKiloRuntime = { + identity: roots[2], scopeId: 'test-scope-other', runtimeId: 'native_2', directory: otherDirectory, @@ -486,11 +553,21 @@ describe('scoped worktree runtime deletion', () => { kiloClient, signal: new AbortController().signal, }; + const siblingKiloRuntime: WorktreeKiloRuntime = { + ...kiloRuntime, + identity: roots[1], + }; const runtime = createControlTerminalRuntime({ controlUrl: 'ws://127.0.0.1:1/sandbox-control/sandbox', wrapperInstanceId: crypto.randomUUID(), - getKiloRuntime: dir => - dir === directory ? kiloRuntime : dir === otherDirectory ? otherKiloRuntime : undefined, + getKiloRuntime: identity => + identity.sessionId === roots[0].sessionId + ? kiloRuntime + : identity.sessionId === roots[1].sessionId + ? siblingKiloRuntime + : identity.sessionId === roots[2].sessionId + ? otherKiloRuntime + : undefined, }); try { @@ -612,7 +689,7 @@ describe('scoped worktree runtime deletion', () => { await rejects( deleteWorktree( { worktreeId, directory, sessionIds: [sessionId(0)] }, - { ...f.deps, client: live ? f.client : undefined } + { ...f.deps, clients: live ? [f.client] : [] } ), /directory conflict/ ); @@ -648,14 +725,14 @@ describe('scoped worktree runtime deletion', () => { expect(() => validateWorktreeDirectory({ worktreeId, directory: unsafe, sessionIds: [] }) ).toThrow('Invalid worktree directory'); - await rejects(prepareWorktreeDeletion(input, {}), /Invalid worktree directory/); + await rejects(prepareWorktreeDeletion(input, { clients: [] }), /Invalid worktree directory/); }); test.each([true, false])( 'waits for an in-flight directory operation and permanently fences later operations for that directory only: live=%s', async live => { const f = fixture(); - const deps = { ...f.deps, client: live ? f.client : undefined }; + const deps = { ...f.deps, clients: live ? [f.client] : [] }; const started = Promise.withResolvers(); const release = Promise.withResolvers(); const inflight = runDirectoryOperation(directory, async () => { diff --git a/services/cloud-agent-next/wrapper/src/control/delete-worktree.ts b/services/cloud-agent-next/wrapper/src/control/delete-worktree.ts index abfd975efe..3211ac1a35 100644 --- a/services/cloud-agent-next/wrapper/src/control/delete-worktree.ts +++ b/services/cloud-agent-next/wrapper/src/control/delete-worktree.ts @@ -155,7 +155,7 @@ async function assertNoSymlinks(directory: string): Promise { export type WorktreeCleanupDeps = { onDiagnostic?: ControlDiagnosticReporter; - client?: WorktreeKiloCleanupClient; + clients: readonly WorktreeKiloCleanupClient[]; assertDirectory?: (directory: string) => Promise; retireDirectory?: (directory: string) => Promise; removeDirectory?: (directory: string) => Promise; @@ -163,6 +163,24 @@ export type WorktreeCleanupDeps = { detachTerminals?: (directory: string) => Promise; }; +function uniqueClients(clients: readonly WorktreeKiloCleanupClient[]): WorktreeKiloCleanupClient[] { + return [...new Set(clients)]; +} + +async function ownersForSession( + clients: readonly WorktreeKiloCleanupClient[], + directory: string, + sessionId: string +): Promise> { + const sessions = await Promise.all( + clients.map(async client => ({ + client, + session: await client.getSession(directory, sessionId), + })) + ); + return sessions.flatMap(({ client, session }) => (session ? [{ client, session }] : [])); +} + export async function prepareWorktreeDeletion( raw: unknown, deps: WorktreeCleanupDeps @@ -187,11 +205,11 @@ export async function prepareWorktreeDeletion( await fenceDirectoryOperations(input.directory); stage = 'directory_validation'; await (deps.assertDirectory ?? assertNoSymlinks)(input.directory); - const { client } = deps; + const clients = uniqueClients(deps.clients); stage = 'manifest_discovery'; const sessionIds = new Set([ ...input.sessionIds, - ...(client ? await client.listSessionIds(input.directory) : []), + ...(await Promise.all(clients.map(client => client.listSessionIds(input.directory)))).flat(), ]); sessionCount = sessionIds.size; for (const sessionId of sessionIds) { @@ -199,19 +217,22 @@ export async function prepareWorktreeDeletion( const rememberedDirectory = directoryForSession(sessionId); if (rememberedDirectory && rememberedDirectory !== input.directory) throw new Error('Worktree session directory conflict'); - if (!client) continue; - const session = await client.getSession(input.directory, sessionId); - if (session && session.directory !== input.directory) + const owners = await ownersForSession(clients, input.directory, sessionId); + if (owners.some(({ session }) => session.directory !== input.directory)) throw new Error('Worktree session directory conflict'); + if (owners.length === 0) continue; stage = 'session_abort'; - await client.abortSession(input.directory, sessionId); - if (!session) continue; + await Promise.all( + owners.map(({ client }) => client.abortSession(input.directory, sessionId)) + ); stage = 'manifest_discovery'; - for (const child of await client.children(input.directory, sessionId)) { - if (child.directory !== input.directory) - throw new Error('Worktree child directory conflict'); - sessionIds.add(child.id); - sessionCount = sessionIds.size; + for (const { client } of owners) { + for (const child of await client.children(input.directory, sessionId)) { + if (child.directory !== input.directory) + throw new Error('Worktree child directory conflict'); + sessionIds.add(child.id); + sessionCount = sessionIds.size; + } } } stage = 'manifest_discovery'; @@ -249,29 +270,43 @@ export async function deleteWorktree( stage = 'manifest_growth'; throw new Error('Worktree cleanup manifest changed'); } - const { client } = deps; + const clients = uniqueClients(deps.clients); + const ownersBySession = new Map< + string, + Array<{ client: WorktreeKiloCleanupClient; session: CleanupSession }> + >(); + for (const sessionId of sessionIds) { + const owners = await ownersForSession(clients, input.directory, sessionId); + if (owners.some(({ session }) => session.directory !== input.directory)) + throw new Error('Worktree session directory conflict'); + ownersBySession.set(sessionId, owners); + } stage = 'process_cleanup'; - if (client) { - for (const sessionId of sessionIds) { - await client.stopSessionProcesses(input.directory, sessionId); - } + for (const [sessionId, owners] of ownersBySession) { + await Promise.all( + owners.map(({ client }) => client.stopSessionProcesses(input.directory, sessionId)) + ); } stage = 'terminal_cleanup'; await deps.detachTerminals?.(input.directory); - if (client) { + for (const client of clients) { await client.closeTerminals(input.directory); stage = 'session_delete'; - for (const sessionId of [...sessionIds].reverse()) { + for (const [sessionId, owners] of [...ownersBySession].reverse()) { + if (!owners.some(owner => owner.client === client)) continue; await client.deleteSession(input.directory, sessionId); } - stage = 'session_delete_confirmation'; - for (const sessionId of sessionIds) { - if (await client.getSession(input.directory, sessionId)) { - stage = 'session_delete_unconfirmed'; - throw new Error('Kilo session deletion was not confirmed'); - } + } + stage = 'session_delete_confirmation'; + for (const [sessionId, owners] of ownersBySession) { + for (const { client } of owners) { + if (!(await client.getSession(input.directory, sessionId))) continue; + stage = 'session_delete_unconfirmed'; + throw new Error('Kilo session deletion was not confirmed'); } - stage = 'directory_dispose'; + } + stage = 'directory_dispose'; + for (const client of clients) { await client.disposeDirectory(input.directory); } stage = 'runtime_retirement'; diff --git a/services/cloud-agent-next/wrapper/src/control/main.ts b/services/cloud-agent-next/wrapper/src/control/main.ts index 7be75bdd85..c927bd0f37 100644 --- a/services/cloud-agent-next/wrapper/src/control/main.ts +++ b/services/cloud-agent-next/wrapper/src/control/main.ts @@ -58,7 +58,12 @@ function main( sessionId: eventKiloSessionId(event.properties), runtimeDirectory: runtime.directory, }); - if (!identity?.rootKiloSessionId) return; + if ( + !identity?.rootKiloSessionId || + (runtime.isolation === 'per-session' && + identity.rootKiloSessionId !== runtime.identity?.kiloSessionId) + ) + return; updateSessionSnapshots(event, deps.sessions); deps.activity?.observeEvent( event.type, @@ -89,7 +94,7 @@ function main( onUnexpectedClose: failure => { logToFile(`Kilo worktree retired reason=${failure.reason} directory=${failure.directory}`); const stillCurrent = () => { - const current = kiloRuntimes.get(failure.directory); + const current = kiloRuntimes.get(failure.identity); return current === undefined || current.runtimeId === failure.runtimeId; }; if (failure.cleanup === 'unconfirmed' || !control?.reportNativeRuntimeRetirement) { @@ -119,7 +124,7 @@ function main( ? createControlTerminalRuntime({ controlUrl: controlConfig.SANDBOX_CONTROL_URL, wrapperInstanceId: controlConfig.wrapperInstanceId, - getKiloRuntime: directory => kiloRuntimes.get(directory), + getKiloRuntime: identity => kiloRuntimes.get(identity), }) : undefined; const deps = createControlHandlerDeps({ @@ -156,7 +161,11 @@ function main( const mutationNotifications = createWorktreeMutationNotifications({ sessions: deps.sessions, - kiloRuntimes, + kiloRuntimes: { + get: identity => kiloRuntimes.get(identity), + isCurrent: runtime => + kiloRuntimes.isCurrent?.(runtime) ?? kiloRuntimes.get(runtime.directory) === runtime, + }, signal: abort.signal, sendEvent: (event, payload, identity) => control?.sendEvent?.(event, payload, identity), }); @@ -335,7 +344,10 @@ function main( isReady: () => deps.kiloReady, onConnected: () => diagnostics.onDiagnostic('wrapper.lifecycle', { phase: 'ready', ok: true }), onEventReceiptFailure: createControlEventFailureHandler({ - getRuntime: directory => kiloRuntimes.get(directory), + getRuntime: (directory, nativeRuntimeId) => { + const runtime = kiloRuntimes.getRetained?.(directory, nativeRuntimeId); + return runtime && !runtime.signal.aborted ? runtime : undefined; + }, onFailure: (failure, runtime) => { reportOutboxRetirement(failure, runtime.runtimeId, 'started'); void deps.operations diff --git a/services/cloud-agent-next/wrapper/src/control/native-observations.ts b/services/cloud-agent-next/wrapper/src/control/native-observations.ts index 351b288c37..ae58b378de 100644 --- a/services/cloud-agent-next/wrapper/src/control/native-observations.ts +++ b/services/cloud-agent-next/wrapper/src/control/native-observations.ts @@ -18,7 +18,7 @@ type DirectoryObservation = { type NativeObservationDeps = { signal?: AbortSignal; roots: () => readonly RootSnapshot[]; - getRuntime: (directory: string) => WorktreeKiloRuntime | undefined; + getRuntime: (directory: string, kiloSessionId: string) => WorktreeKiloRuntime | undefined; reconcileActivity: ( statuses: Awaited>, roots: readonly string[] @@ -27,23 +27,35 @@ type NativeObservationDeps = { export function createNativeObservations(deps: NativeObservationDeps) { const observations = new Map(); + const runtimeIds = new WeakMap(); + let nextRuntimeId = 0; - function forget(directory: string): void { - observations.delete(directory); + function runtimeKey(runtime: WorktreeKiloRuntime): number { + const known = runtimeIds.get(runtime); + if (known !== undefined) return known; + const id = nextRuntimeId; + nextRuntimeId += 1; + runtimeIds.set(runtime, id); + return id; + } + + function forget(observationKey: string): void { + observations.delete(observationKey); } async function sampleDirectory( + observationKey: string, directory: string, roots: readonly RootSnapshot[], signal?: AbortSignal ): Promise { - const runtime = deps.getRuntime(directory); + const runtime = deps.getRuntime(directory, roots[0]?.kiloSessionId ?? ''); if (!runtime) { - forget(directory); + forget(observationKey); return; } const client = runtime.kiloClient; - let entry = observations.get(directory); + let entry = observations.get(observationKey); if ( entry && (entry.runtime !== runtime || @@ -59,7 +71,7 @@ export function createNativeObservations(deps: NativeObservationDeps) { } if (!entry) { entry = { runtime, client, roots: roots.map(root => ({ ...root })), pending: undefined }; - observations.set(directory, entry); + observations.set(observationKey, entry); } if (entry.pending) return entry.pending; const target = entry; @@ -70,8 +82,8 @@ export function createNativeObservations(deps: NativeObservationDeps) { const readSignal = AbortSignal.any(signals); const isCurrent = () => !readSignal.aborted && - observations.get(directory) === target && - deps.getRuntime(directory) === runtime && + observations.get(observationKey) === target && + deps.getRuntime(directory, capturedRoots[0]?.kiloSessionId ?? '') === runtime && runtime.kiloClient === client; const pending = (async () => { @@ -101,18 +113,23 @@ export function createNativeObservations(deps: NativeObservationDeps) { async function refresh(signal?: AbortSignal): Promise { if (deps.signal?.aborted || signal?.aborted) return; - const rootsByDirectory = new Map(); + const rootsByRuntime = new Map(); for (const root of deps.roots()) { if (!root.directory) continue; - const roots = rootsByDirectory.get(root.directory) ?? []; + const runtime = deps.getRuntime(root.directory, root.kiloSessionId); + if (!runtime) continue; + const key = `${root.directory}\0${runtimeKey(runtime)}`; + const roots = rootsByRuntime.get(key) ?? []; roots.push(root); - rootsByDirectory.set(root.directory, roots); + rootsByRuntime.set(key, roots); } - for (const directory of observations.keys()) { - if (!rootsByDirectory.has(directory)) forget(directory); + for (const key of observations.keys()) { + if (!rootsByRuntime.has(key)) forget(key); } await Promise.all( - [...rootsByDirectory].map(([directory, roots]) => sampleDirectory(directory, roots, signal)) + [...rootsByRuntime.entries()].map(([key, roots]) => + sampleDirectory(key, roots[0]?.directory ?? '', roots, signal) + ) ); } diff --git a/services/cloud-agent-next/wrapper/src/control/operation-registry.ts b/services/cloud-agent-next/wrapper/src/control/operation-registry.ts index f0df4fe7ad..3921125828 100644 --- a/services/cloud-agent-next/wrapper/src/control/operation-registry.ts +++ b/services/cloud-agent-next/wrapper/src/control/operation-registry.ts @@ -8,7 +8,7 @@ import { type SessionRequestIdentity, } from '../../../src/shared/sandbox-control-protocol.js'; import { rejectBeforeAdmission } from './control-handler-result.js'; -import type { WorktreeKiloRuntimes } from './worktree-runtime.js'; +import type { WorktreeKiloRuntime, WorktreeKiloRuntimes } from './worktree-runtime.js'; import type { NativeOperationTarget, NativeRetirement } from './session-operation-cleanup.js'; import { SessionOperation, @@ -19,12 +19,8 @@ import { type OperationRegistryDependencies = { native: { - get( - directory: string - ): WorktreeKiloRuntimes['get'] extends (directory: string) => infer Runtime ? Runtime : never; - getRetained( - directory: string - ): WorktreeKiloRuntimes['get'] extends (directory: string) => infer Runtime ? Runtime : never; + get(identity: SessionRequestIdentity): ReturnType; + getRetained(directory: string, runtimeId?: string): WorktreeKiloRuntime | undefined; prepareForNewWork?(directory: string): boolean; retireRuntime( directory: string, @@ -219,7 +215,7 @@ export function createOperationRegistry(deps: OperationRegistryDependencies) { const operation = new SessionOperation(identity, authorization, work, { ...effects, isCurrent: () => active.get(identity.kiloSessionId) === operation, - getRuntime: () => deps.native.get(identity.directory), + getRuntime: () => deps.native.get(identity), prepareForNewWork: () => deps.native.prepareForNewWork?.(identity.directory) ?? true, verifyQuiescence: (target, deadlineAt) => deps.native.verifyQuiescence(identity.directory, target, deadlineAt), diff --git a/services/cloud-agent-next/wrapper/src/control/sandbox-control-client.test.ts b/services/cloud-agent-next/wrapper/src/control/sandbox-control-client.test.ts index 732768fd7d..67284d19a7 100644 --- a/services/cloud-agent-next/wrapper/src/control/sandbox-control-client.test.ts +++ b/services/cloud-agent-next/wrapper/src/control/sandbox-control-client.test.ts @@ -132,8 +132,10 @@ function versionHeartbeat(version: string | null) { throw new Error('Unexpected attach'); }, detach: () => false, + retireForRecovery: async () => 'absent', deleteDirectory: async () => {}, get: () => undefined, + isCurrent: () => false, isHealthy: () => true, shutdown() {}, }, @@ -363,6 +365,8 @@ describe('createSandboxControlClient', () => { nativeRuntimeRetirement?: boolean; connectionRecovery?: boolean; eventReceipts?: boolean; + runtimeIsolation?: boolean; + runtimeRecovery?: boolean; workingBranches?: boolean; }; }; @@ -378,6 +382,8 @@ describe('createSandboxControlClient', () => { nativeRuntimeRetirement: true, connectionRecovery: true, eventReceipts: true, + runtimeIsolation: true, + runtimeRecovery: true, workingBranches: true, }, }); diff --git a/services/cloud-agent-next/wrapper/src/control/sandbox-control-client.ts b/services/cloud-agent-next/wrapper/src/control/sandbox-control-client.ts index d7313504df..462e281660 100644 --- a/services/cloud-agent-next/wrapper/src/control/sandbox-control-client.ts +++ b/services/cloud-agent-next/wrapper/src/control/sandbox-control-client.ts @@ -521,6 +521,8 @@ export function createSandboxControlClient( nativeRuntimeRetirement: true, connectionRecovery: true, eventReceipts: true, + runtimeIsolation: true, + runtimeRecovery: true, workingBranches: true, }, ...(wrapperInstanceId ? { wrapperInstanceId } : {}), diff --git a/services/cloud-agent-next/wrapper/src/control/sandbox-control-handlers.test.ts b/services/cloud-agent-next/wrapper/src/control/sandbox-control-handlers.test.ts index 6c98d37062..350097fc7a 100644 --- a/services/cloud-agent-next/wrapper/src/control/sandbox-control-handlers.test.ts +++ b/services/cloud-agent-next/wrapper/src/control/sandbox-control-handlers.test.ts @@ -10,6 +10,7 @@ import { SANDBOX_CONTROL_REQUEST_TIMEOUT_MS, sessionSyncResultSchema, type SessionEventPayload, + type SessionRequestIdentity, type SessionGitSummaryResult, } from '../../../src/shared/sandbox-control-protocol'; import { createWrapperKiloClient, type WrapperKiloClient, type WrapperPty } from '../kilo-api'; @@ -122,6 +123,7 @@ function deps( const client = Object.hasOwn(overrides, 'kiloClient') ? kiloClient : fakeKilo(); const runtime: WorktreeKiloRuntime | undefined = client ? { + identity: { ...identity }, scopeId: kilo.scopeId, runtimeId: 'native_1', directory: identity.directory, @@ -147,6 +149,7 @@ function deps( release: () => {}, }), detach: () => true, + retireForRecovery: async () => 'retired', deleteDirectory: async () => {}, getRetained: directory => (directory === runtime.directory ? runtime : undefined), retireRuntime: async (directory, _deadlineAt, target) => @@ -159,7 +162,16 @@ function deps( directory === runtime.directory && target.runtimeId === runtime.runtimeId && target.client === runtime.kiloClient, - get: directory => (directory === runtime.directory ? runtime : undefined), + get: request => + ( + typeof request === 'string' + ? request === runtime.directory + : request.directory === runtime.directory + ) + ? runtime + : undefined, + getAll: directory => (directory === runtime.directory ? [runtime] : []), + isCurrent: candidate => candidate === runtime, isHealthy: () => true, shutdown: () => {}, } as NonNullable) @@ -255,6 +267,9 @@ function fakeTerminalRuntime( rememberAttachedSession: () => {}, detachSession: async () => {}, detachDirectory: async () => {}, + hasActivePty: () => false, + beginRecoveryRetirement: () => {}, + endRecoveryRetirement: () => {}, create: async () => ({ pty }), resize: async () => ({ pty }), close: async () => ({ success: true }), @@ -817,6 +832,7 @@ describe('handleControlRequest', () => { }, }); runtimes.set(directory, { + identity: { ...session }, directory, scopeId: directory, runtimeId: crypto.randomUUID(), @@ -837,6 +853,7 @@ describe('handleControlRequest', () => { throw new Error('Unexpected startup'); }, detach: () => true, + retireForRecovery: async () => 'retired', deleteDirectory: async () => {}, getRetained: directory => runtimes.get(directory), retireRuntime: async (directory, _deadlineAt, target) => { @@ -856,7 +873,12 @@ describe('handleControlRequest', () => { target.client === runtime.kiloClient ); }, - get: directory => runtimes.get(directory), + get: identity => runtimes.get(typeof identity === 'string' ? identity : identity.directory), + getAll: directory => { + const runtime = runtimes.get(directory); + return runtime ? [runtime] : []; + }, + isCurrent: runtime => [...runtimes.values()].includes(runtime), isHealthy: () => true, shutdown: () => {}, } as NonNullable, @@ -1112,7 +1134,7 @@ describe('handleControlRequest', () => { }); expect(rootForSession(session.kiloSessionId)).toBe(session.kiloSessionId); expect(rootForSession(sibling.kiloSessionId)).toBe(sibling.kiloSessionId); - expect(handlerDeps.kiloRuntimes?.get(session.directory)).toBe(runtime); + expect(handlerDeps.kiloRuntimes?.get(session)).toBe(runtime); confirmed = true; expect(await handleControlRequest('session.detach', session, {}, handlerDeps)).toEqual({ ok: true, @@ -1219,7 +1241,11 @@ describe('handleControlRequest', () => { runtimes.shutdown = () => { shutdowns += 1; }; - const runtime = runtimes.get(session.directory); + const runtime = runtimes.get(session); + expect(await handleControlRequest('session.attach', sibling, { kilo }, handlerDeps)).toEqual({ + ok: true, + result: { attached: true }, + }); expect( await handleControlRequest('session.prompt', sibling, promptPayload, handlerDeps) ).toEqual({ @@ -1232,7 +1258,7 @@ describe('handleControlRequest', () => { result: { detached: true }, }); expect(shutdowns).toBe(0); - expect(runtimes.get(session.directory)).toBe(runtime); + expect(runtimes.getRetained?.(session.directory)).toBe(runtime); expect(handlerDeps.operations.active(sibling.kiloSessionId)?.signal.aborted).toBe(false); expect(events).toEqual([]); expect(aborted).toEqual([session.kiloSessionId]); @@ -1643,10 +1669,12 @@ describe('production worktree deletion routes', () => { const calls: string[] = []; const handlerDeps = deps({ kiloRuntimes: { - get: dir => { - calls.push(`get:${dir}`); + get: identity => { + calls.push(`get:${typeof identity === 'string' ? identity : identity.directory}`); return undefined; }, + getAll: () => [], + isCurrent: () => false, attach: () => { calls.push('attach'); throw new Error('Unexpected runtime startup'); @@ -1655,6 +1683,7 @@ describe('production worktree deletion routes', () => { calls.push('detach'); return true; }, + retireForRecovery: async () => 'retired', deleteDirectory: async dir => { calls.push(`delete:${dir}`); }, @@ -1756,7 +1785,7 @@ describe('production worktree deletion routes', () => { ); const aborted: string[] = []; const outcomes: Array<{ id: string; event: SessionEventPayload }> = []; - const lookups: string[] = []; + const lookups: Array = []; const handlerDeps = deps( { kiloClient: fakeKilo({ @@ -1778,7 +1807,7 @@ describe('production worktree deletion routes', () => { first ); const runtimes = handlerDeps.kiloRuntimes; - const selected = runtimes?.get(directory); + const selected = runtimes?.get(first); if (!runtimes || !selected) throw new Error('Expected selected worktree runtime'); const otherRuntime = { ...selected, @@ -1796,9 +1825,14 @@ describe('production worktree deletion routes', () => { }, }), }; - runtimes.get = dir => { - lookups.push(dir); - return dir === directory ? selected : dir === siblingDirectory ? otherRuntime : undefined; + runtimes.get = identity => { + const worktree = typeof identity === 'string' ? identity : identity.directory; + lookups.push(identity); + return worktree === directory + ? selected + : worktree === siblingDirectory + ? otherRuntime + : undefined; }; const attach = spyOn(runtimes, 'attach'); let preparation: ReturnType | undefined; @@ -1886,7 +1920,9 @@ describe('production worktree deletion routes', () => { ) ).toBe(true); expect(lookups.length).toBeGreaterThan(0); - expect(lookups.every(value => value === directory)).toBe(true); + expect( + lookups.every(value => typeof value !== 'string' && value.directory === directory) + ).toBe(true); expect(http.requests.every(request => request.directory === directory)).toBe(true); expect(http.requests).toContainEqual({ method: 'POST', @@ -1945,11 +1981,12 @@ describe('production worktree deletion routes', () => { first ); const runtimes = handlerDeps.kiloRuntimes; - const selected = runtimes?.get(directory); + const selected = runtimes?.get(first); if (!runtimes || !selected) throw new Error('Expected selected worktree runtime'); - runtimes.get = dir => { - lookups.push(dir); - return dir === directory ? selected : undefined; + runtimes.get = identity => { + const worktree = typeof identity === 'string' ? identity : identity.directory; + lookups.push(worktree); + return worktree === directory ? selected : undefined; }; runtimes.deleteDirectory = async dir => { retirements.push(dir); @@ -1981,7 +2018,7 @@ describe('production worktree deletion routes', () => { ok: true, result: { deleted: true, sessionIds: [sessionId(1), sessionId(2), sessionId(3)] }, }); - expect(lookups).toEqual([directory]); + expect(lookups).toEqual([]); expect(retirements).toEqual([directory]); expect(filesystem.remove.mock.calls).toEqual([[directory, { recursive: true, force: true }]]); expect(handlerDeps.sessions).toEqual([ @@ -2040,12 +2077,15 @@ describe('production worktree deletion routes', () => { lastActivityAt: 100, })), kiloRuntimes: { - get: dir => { - lookups.push(dir); - if (dir === directory) return undefined; + get: identity => { + const worktree = typeof identity === 'string' ? identity : identity.directory; + lookups.push(worktree); + if (worktree === directory) return undefined; forbidden.push('unscoped lookup'); return siblingRuntime; }, + getAll: directory => (directory === siblingDirectory ? [siblingRuntime] : []), + isCurrent: runtime => runtime === siblingRuntime, attach: () => { forbidden.push('startup'); throw new Error('Cleanup must not start a runtime'); @@ -2054,6 +2094,7 @@ describe('production worktree deletion routes', () => { forbidden.push('detach'); return false; }, + retireForRecovery: async () => 'retired', deleteDirectory: async dir => { retirements.push(dir); }, @@ -2080,7 +2121,7 @@ describe('production worktree deletion routes', () => { await handleControlRequest('worktree.delete', undefined, input, handlerDeps) ).toEqual({ ok: true, result: { deleted: true, sessionIds: input.sessionIds } }); } - expect(lookups).toEqual([directory, directory, directory, directory]); + expect(lookups).toEqual([]); expect(retirements).toEqual([directory, directory]); expect(detached).toEqual([directory, directory]); expect(filesystem.remove.mock.calls).toEqual([ @@ -3342,7 +3383,7 @@ describe('control cancellation and attachments', () => { }), }); const runtimes = handlerDeps.kiloRuntimes; - const firstRuntime = runtimes?.get(session.directory); + const firstRuntime = runtimes?.get(session); if (!runtimes || !firstRuntime) throw new Error('Expected preparation runtime'); const secondDeps = deps( { @@ -3367,12 +3408,20 @@ describe('control cancellation and attachments', () => { ok: true, result: { attached: true }, }); - const secondRuntime = secondDeps.kiloRuntimes?.get(secondSession.directory); + const secondRuntime = secondDeps.kiloRuntimes?.get(secondSession); if (!secondRuntime) throw new Error('Expected execution runtime'); - runtimes.get = directory => - directory === session.directory + runtimes.get = identity => + ( + typeof identity === 'string' + ? identity === session.directory + : identity.kiloSessionId === session.kiloSessionId + ) ? firstRuntime - : directory === secondSession.directory + : ( + typeof identity === 'string' + ? identity === secondSession.directory + : identity.kiloSessionId === secondSession.kiloSessionId + ) ? secondRuntime : undefined; const attaching = handleControlRequest('session.attach', session, { kilo }, handlerDeps); @@ -3561,6 +3610,12 @@ describe('control cancellation and attachments', () => { }), }); try { + expect(await handleControlRequest('session.attach', identity, { kilo }, handlerDeps)).toEqual( + { + ok: true, + result: { attached: true }, + } + ); const accepted = await handleControlRequest( 'session.prompt', identity, @@ -3637,6 +3692,12 @@ describe('control cancellation and attachments', () => { emitSessionEvent: (_session, event) => events.push(event), }); try { + expect(await handleControlRequest('session.attach', identity, { kilo }, handlerDeps)).toEqual( + { + ok: true, + result: { attached: true }, + } + ); await handleControlRequest( 'session.prompt', identity, @@ -4199,7 +4260,7 @@ describe('control interactions and sync', () => { ok: true, result: { status: { type: 'busy' }, questions: [], permissions: [] }, }); - const runtime = handlerDeps.kiloRuntimes?.get(session.directory); + const runtime = handlerDeps.kiloRuntimes?.get(session); if (!runtime) throw new Error('Expected worktree runtime'); (runtime as { kiloClient: WrapperKiloClient }).kiloClient = fakeKilo({ getQuestions: async () => { @@ -4673,7 +4734,7 @@ describe('refreshHeartbeatPayload', () => { } ); } - const runtime = handlerDeps.kiloRuntimes?.get(session.directory); + const runtime = handlerDeps.kiloRuntimes?.get(session); expect(runtime).toBeDefined(); const refresh = refreshHeartbeatPayload(handlerDeps); try { @@ -4687,7 +4748,7 @@ describe('refreshHeartbeatPayload', () => { result: { attached: true }, }); activity.attach(sibling.kiloSessionId); - expect(handlerDeps.kiloRuntimes?.get(session.directory)).toBe(runtime); + expect(handlerDeps.kiloRuntimes?.get(session)).toBe(runtime); statuses.resolve({ kilo_1: { type: 'busy' }, kilo_2: { type: 'busy', waitingOn: 'tool' } }); expect(await refresh).toEqual({ state: 'active', @@ -4794,16 +4855,30 @@ describe('refreshHeartbeatPayload', () => { true ); const runtimes = handlerDeps.kiloRuntimes; - const firstRuntime = runtimes?.get(first.directory); - const secondRuntime = otherDeps.kiloRuntimes?.get(second.directory); + const firstRuntime = runtimes?.get(first); + const secondRuntime = otherDeps.kiloRuntimes?.get(second); if (!runtimes || !firstRuntime || !secondRuntime) throw new Error('Expected directory runtimes'); - runtimes.get = directory => - directory === first.directory + runtimes.get = identity => + ( + typeof identity === 'string' + ? identity === first.directory + : identity.kiloSessionId === first.kiloSessionId + ) ? firstRuntime - : directory === second.directory + : ( + typeof identity === 'string' + ? identity === second.directory + : identity.kiloSessionId === second.kiloSessionId + ) ? secondRuntime : undefined; + runtimes.getAll = directory => + directory === first.directory + ? [firstRuntime] + : directory === second.directory + ? [secondRuntime] + : []; now = 150; const payload = await refreshHeartbeatPayload(handlerDeps); @@ -4903,16 +4978,22 @@ describe('refreshHeartbeatPayload', () => { true ); const runtimes = handlerDeps.kiloRuntimes; - const firstRuntime = runtimes?.get(first.directory); - const secondRuntime = otherDeps.kiloRuntimes?.get(second.directory); + const firstRuntime = runtimes?.get(first); + const secondRuntime = otherDeps.kiloRuntimes?.get(second); if (!runtimes || !firstRuntime || !secondRuntime) throw new Error('Expected directory runtimes'); - runtimes.get = directory => - directory === first.directory + runtimes.get = identity => + (typeof identity === 'string' ? identity : identity.directory) === first.directory ? firstRuntime - : directory === second.directory + : (typeof identity === 'string' ? identity : identity.directory) === second.directory ? secondRuntime : undefined; + runtimes.getAll = directory => + directory === first.directory + ? [firstRuntime] + : directory === second.directory + ? [secondRuntime] + : []; expect((await refreshHeartbeatPayload(handlerDeps)).sessions).toEqual([ { kiloSessionId: 'kilo_1', state: 'active', idleForMs: 0, waitingOn: 'tool' }, { kiloSessionId: 'kilo_2', state: 'finalizing', idleForMs: 0, waitingOn: 'finalizing' }, @@ -5006,16 +5087,30 @@ describe('refreshHeartbeatPayload', () => { other ); const runtimes = handlerDeps.kiloRuntimes; - const firstRuntime = runtimes?.get(session.directory); - const otherRuntime = otherDeps.kiloRuntimes?.get(other.directory); + const firstRuntime = runtimes?.get(session); + const otherRuntime = otherDeps.kiloRuntimes?.get(other); if (!runtimes || !firstRuntime || !otherRuntime) throw new Error('Expected directory runtimes'); - runtimes.get = directory => - directory === session.directory + runtimes.get = identity => + ( + typeof identity === 'string' + ? identity === session.directory + : identity.kiloSessionId === session.kiloSessionId + ) ? firstRuntime - : directory === other.directory + : ( + typeof identity === 'string' + ? identity === other.directory + : identity.kiloSessionId === other.kiloSessionId + ) ? otherRuntime : undefined; + runtimes.getAll = directory => + directory === session.directory + ? [firstRuntime] + : directory === other.directory + ? [otherRuntime] + : []; const payload = await refreshHeartbeatPayload(handlerDeps); expect(reads).toHaveLength(2); diff --git a/services/cloud-agent-next/wrapper/src/control/sandbox-control-handlers.ts b/services/cloud-agent-next/wrapper/src/control/sandbox-control-handlers.ts index eaea32d624..5a503fd8b4 100644 --- a/services/cloud-agent-next/wrapper/src/control/sandbox-control-handlers.ts +++ b/services/cloud-agent-next/wrapper/src/control/sandbox-control-handlers.ts @@ -16,6 +16,7 @@ import { sessionPermissionResolvePayloadSchema, sessionPromptPayloadSchema, sessionQuestionResolvePayloadSchema, + sessionRuntimeRetirePayloadSchema, sessionSyncPayloadSchema, sessionTerminalClosePayloadSchema, sessionTerminalCloseResultSchema, @@ -328,8 +329,9 @@ export function createControlHandlerDeps(input: Omit) const deps: HandlerDeps = Object.assign(input, { operations: createOperationRegistry({ native: { - get: directory => input.kiloRuntimes?.get(directory), - getRetained: directory => input.kiloRuntimes?.getRetained?.(directory), + get: identity => input.kiloRuntimes?.get(identity), + getRetained: (directory, runtimeId) => + input.kiloRuntimes?.getRetained?.(directory, runtimeId), prepareForNewWork: directory => input.kiloRuntimes?.prepareForNewWork?.(directory) ?? true, retireRuntime: (directory, deadlineAt, target) => input.kiloRuntimes?.retireRuntime?.(directory, deadlineAt, target) ?? @@ -367,7 +369,11 @@ export function createControlHandlerDeps(input: Omit) directory: directoryForSession(kiloSessionId), revision: deps.activity?.revision(kiloSessionId), })), - getRuntime: directory => deps.kiloRuntimes?.get(directory), + getRuntime: (directory, kiloSessionId) => + deps.kiloRuntimes + ?.getAll?.(directory) + .find(runtime => runtime.identity?.kiloSessionId === kiloSessionId) ?? + deps.kiloRuntimes?.get(directory), reconcileActivity: (statuses, roots) => deps.activity?.reconcile(statuses, roots), }); } @@ -454,13 +460,18 @@ export async function handleControlRequest( return fail('not_ready', 'Worktree cancellation is incomplete', true); } failureStage = 'runtime_lookup'; - const runtime = kiloRuntimes.get(input.directory); - const client = - deps.worktreeCleanupClient ?? - (runtime ? createWorktreeKiloCleanupClient(runtime.kiloClient.serverUrl) : undefined); + const runtimes = kiloRuntimes.getAll?.(input.directory) ?? []; + // An injected cleanup client is a fallback for a checkout whose runtime has + // already gone away. Live runtimes each retain their own Kilo state. + const clients = + runtimes.length > 0 + ? runtimes.map(runtime => createWorktreeKiloCleanupClient(runtime.kiloClient.serverUrl)) + : deps.worktreeCleanupClient + ? [deps.worktreeCleanupClient] + : []; const cleanupDeps = { onDiagnostic: deps.onDiagnostic, - client, + clients, detachRoot: (id: string) => { deps.activity?.detach(id); const index = deps.sessions.findIndex(snapshot => snapshot.kiloSessionId === id); @@ -565,6 +576,8 @@ async function handleSessionControlRequest( return handleAttach(session, payload, deps, authorization); case 'session.detach': return handleDetach(session, payload, deps); + case 'session.runtime.retire': + return handleRuntimeRetire(session, payload, deps); case 'session.prompt': return handlePrompt(session, payload, deps, authorization); case 'session.abort': @@ -631,7 +644,7 @@ function sessionKiloRuntime( rootForSession(session.kiloSessionId) !== session.kiloSessionId ) return undefined; - return deps.kiloRuntimes?.get(session.directory); + return deps.kiloRuntimes?.get(session); } function terminalFailure(error: unknown): ControlHandlerResult { @@ -730,6 +743,49 @@ async function handleDetach( } } +async function handleRuntimeRetire( + session: SessionRequestIdentity, + payload: unknown, + deps: HandlerDeps +): Promise { + const parsed = sessionRuntimeRetirePayloadSchema.safeParse(payload); + if (!parsed.success) return fail('protocol_error', 'Invalid payload', false); + const runtimes = deps.kiloRuntimes; + if (!runtimes) return missingKilo(); + let terminalRetirementStarted = false; + try { + deps.terminalRuntime?.beginRecoveryRetirement(session); + terminalRetirementStarted = true; + if (!runtimes.retireForRecovery) return missingKilo(); + const retirement = await runtimes.retireForRecovery(session, parsed.data.recoveryId, () => { + const task = deps.operations.active(session.kiloSessionId); + const active = deps.activity + ?.snapshots() + .some( + snapshot => snapshot.kiloSessionId === session.kiloSessionId && snapshot.state !== 'idle' + ); + if (task || active) { + throw new WorktreeKiloRuntimeError('session_busy', 'Session has work in progress', true); + } + if (deps.terminalRuntime?.hasActivePty(session)) { + throw new WorktreeKiloRuntimeError('session_busy', 'Session has an active PTY', true); + } + }); + if (retirement === 'retired') { + await deps.terminalRuntime?.detachSession(session); + forgetAttachedRoot(session.kiloSessionId, session.directory); + deps.activity?.detach(session.kiloSessionId); + const index = deps.sessions.findIndex(item => item.kiloSessionId === session.kiloSessionId); + if (index !== -1) deps.sessions.splice(index, 1); + } + return ok({ recoveryId: parsed.data.recoveryId, retired: true }); + } catch (error) { + return terminalFailure(error); + } finally { + if (terminalRetirementStarted) deps.terminalRuntime?.endRecoveryRetirement(session); + } +} + type RuntimeSchema = { safeParse(value: unknown): { success: true; data: Value } | { success: false }; }; @@ -916,7 +972,10 @@ async function handleAbort( const parsed = sessionAbortPayloadSchema.safeParse(payload ?? {}); if (!parsed.success) return fail('protocol_error', 'Invalid payload', false); if (parsed.data.nativeRuntimeId) { - const runtime = deps.kiloRuntimes?.getRetained?.(session.directory); + const runtime = deps.kiloRuntimes?.getRetained?.( + session.directory, + parsed.data.nativeRuntimeId + ); if (!runtime || runtime.runtimeId !== parsed.data.nativeRuntimeId) { return ok({ status: 'aborted', diff --git a/services/cloud-agent-next/wrapper/src/control/session-operation-cleanup.test.ts b/services/cloud-agent-next/wrapper/src/control/session-operation-cleanup.test.ts index b3b9578f83..0fe247b9ce 100644 --- a/services/cloud-agent-next/wrapper/src/control/session-operation-cleanup.test.ts +++ b/services/cloud-agent-next/wrapper/src/control/session-operation-cleanup.test.ts @@ -154,11 +154,13 @@ describe('SessionOperation cleanup', () => { release: () => {}, }), detach: () => true, + retireForRecovery: async () => 'retired', deleteDirectory: async () => {}, - getRetained: () => runtime, + getRetained: () => (runtime.signal.aborted ? undefined : runtime), retireRuntime: async () => 'unconfirmed', verifyQuiescence: async () => false, - get: () => runtime, + get: () => (runtime.signal.aborted ? undefined : runtime), + isCurrent: candidate => candidate === runtime && !runtime.signal.aborted, isHealthy: () => true, shutdown: () => {}, }; @@ -265,7 +267,7 @@ describe('SessionOperation cleanup', () => { ok: true, result: { status: 'aborted', quiescent: true }, }); - expect(handlerDeps.kiloRuntimes?.get(session.directory)).toBeDefined(); + expect(handlerDeps.kiloRuntimes?.getRetained?.(session.directory)).toBeDefined(); const authorizationB = operationAuthorization('session.prompt', 'message_b'); await handleControlRequest( @@ -304,7 +306,7 @@ describe('SessionOperation cleanup', () => { quiescent: true, }, }); - expect(handlerDeps.kiloRuntimes?.get(session.directory)?.runtimeId).toBe('native_1'); + expect(handlerDeps.kiloRuntimes?.get(session)?.runtimeId).toBe('native_1'); expect(operationB?.snapshot().local?.result).toEqual({ ok: true, result: {} }); }); diff --git a/services/cloud-agent-next/wrapper/src/control/terminal-runtime.test.ts b/services/cloud-agent-next/wrapper/src/control/terminal-runtime.test.ts index 70d46882be..33a7d3f891 100644 --- a/services/cloud-agent-next/wrapper/src/control/terminal-runtime.test.ts +++ b/services/cloud-agent-next/wrapper/src/control/terminal-runtime.test.ts @@ -66,18 +66,20 @@ function createRuntime( const runtime = createControlTerminalRuntime({ controlUrl, wrapperInstanceId, - getKiloRuntime: directory => { - let worktree = worktrees.get(directory); + getKiloRuntime: identity => { + const key = `${identity.sessionId}\0${identity.kiloSessionId}\0${identity.directory}`; + let worktree = worktrees.get(key); if (!worktree) { worktree = { + identity: { ...identity }, runtimeId: `native_${worktrees.size + 1}`, - scopeId: directory, - directory, - env: { WORKTREE_VALUE: directory }, + scopeId: identity.directory, + directory: identity.directory, + env: { WORKTREE_VALUE: identity.directory }, kiloClient, signal: new AbortController().signal, }; - worktrees.set(directory, worktree); + worktrees.set(key, worktree); } return worktree; }, @@ -546,12 +548,41 @@ describe('control terminal PTY ownership', () => { }); }); + it('rejects a same-directory sibling when lookup returns another root runtime', async () => { + const sibling = { ...secondSession, directory: firstSession.directory }; + const firstRuntime: WorktreeKiloRuntime = { + identity: { ...firstSession }, + runtimeId: 'native_first', + directory: firstSession.directory, + scopeId: firstSession.directory, + env: { HOME: '/home/first', KILOCODE_TOKEN: 'first-token' }, + kiloClient: fakeKilo(), + signal: new AbortController().signal, + }; + const runtime = createControlTerminalRuntime({ + controlUrl: 'ws://127.0.0.1:1/sandbox-control/sandbox', + wrapperInstanceId, + getKiloRuntime: identity => + identity.directory === firstSession.directory ? firstRuntime : undefined, + }); + activeRuntimes.add(runtime); + + attach(runtime, firstSession); + expect(() => runtime.rememberAttachedSession(sibling)).toThrow( + /Terminal session ownership mismatch/ + ); + expect(await runtime.create(firstSession, creationPayload())).toMatchObject({ + pty: { cwd: firstSession.directory }, + }); + }); + it('uses the owning worktree client and credentials for every PTY operation', async () => { const calls: Array<{ operation: string; directory: string; env?: Record }> = []; const worktrees = new Map(); for (const identity of [firstSession, secondSession]) { const directory = identity.directory; - worktrees.set(directory, { + worktrees.set(identity.kiloSessionId, { + identity: { ...identity }, runtimeId: `native_${identity.sessionId}`, directory, scopeId: directory, @@ -576,7 +607,7 @@ describe('control terminal PTY ownership', () => { const runtime = createControlTerminalRuntime({ controlUrl: 'ws://127.0.0.1:1/sandbox-control/sandbox', wrapperInstanceId, - getKiloRuntime: directory => worktrees.get(directory), + getKiloRuntime: identity => worktrees.get(identity.kiloSessionId), }); activeRuntimes.add(runtime); attach(runtime, firstSession); @@ -589,7 +620,7 @@ describe('control terminal PTY ownership', () => { expect(calls.at(-3)).toMatchObject({ operation: 'create', directory: identity.directory, - env: worktrees.get(identity.directory)?.env, + env: worktrees.get(identity.kiloSessionId)?.env, }); expect(calls.at(-3)?.env).not.toHaveProperty('SANDBOX_CONTROL_CREDENTIAL'); expect(calls.slice(-2)).toEqual([ @@ -598,7 +629,7 @@ describe('control terminal PTY ownership', () => { ]); } - worktrees.delete(firstSession.directory); + worktrees.delete(firstSession.kiloSessionId); expect( await terminalFailure(runtime.create(firstSession, creationPayload(crypto.randomUUID()))) ).toMatchObject({ code: 'not_ready', message: 'Kilo worktree is not available' }); @@ -764,7 +795,8 @@ describe('control terminal reverse WebSocket bridge', () => { [firstSession, firstServers, 'pty_first'], [secondSession, secondServers, 'pty_second'], ] as const) { - worktrees.set(identity.directory, { + worktrees.set(identity.kiloSessionId, { + identity: { ...identity }, runtimeId: `native_${identity.sessionId}`, scopeId: identity.directory, directory: identity.directory, @@ -779,7 +811,7 @@ describe('control terminal reverse WebSocket bridge', () => { const runtime = createControlTerminalRuntime({ controlUrl: firstServers.controlUrl, wrapperInstanceId, - getKiloRuntime: directory => worktrees.get(directory), + getKiloRuntime: identity => worktrees.get(identity.kiloSessionId), }); activeRuntimes.add(runtime); attach(runtime, firstSession); diff --git a/services/cloud-agent-next/wrapper/src/control/terminal-runtime.ts b/services/cloud-agent-next/wrapper/src/control/terminal-runtime.ts index bdfedfa578..97adbbd3c2 100644 --- a/services/cloud-agent-next/wrapper/src/control/terminal-runtime.ts +++ b/services/cloud-agent-next/wrapper/src/control/terminal-runtime.ts @@ -72,6 +72,9 @@ export type ControlTerminalRuntime = { rememberAttachedSession(identity: SessionRequestIdentity): void; detachSession(identity: SessionRequestIdentity): Promise; detachDirectory(directory: string): Promise; + hasActivePty(identity: SessionRequestIdentity): boolean; + beginRecoveryRetirement(identity: SessionRequestIdentity): void; + endRecoveryRetirement(identity: SessionRequestIdentity): void; create( identity: SessionRequestIdentity, payload: SessionTerminalCreatePayload @@ -150,7 +153,7 @@ function waitForSocketOpen(socket: WebSocket): Promise { export function createControlTerminalRuntime(options: { controlUrl: string; wrapperInstanceId: string; - getKiloRuntime: (directory: string) => WorktreeKiloRuntime | undefined; + getKiloRuntime: (identity: SessionRequestIdentity) => WorktreeKiloRuntime | undefined; }): ControlTerminalRuntime { const { wrapperInstanceId } = options; const controlOrigin = new URL(options.controlUrl).origin; @@ -158,11 +161,12 @@ export function createControlTerminalRuntime(options: { const terminals = new Map(); const operations = new Map(); const bridges = new Map(); + const recoveringSessions = new Set(); let shutDown = false; function requireAttached(identity: SessionRequestIdentity): AttachedTerminalSession { const attached = attachedSessions.get(identity.sessionId); - if (!attached || shutDown) { + if (!attached || shutDown || recoveringSessions.has(identity.sessionId)) { throw new ControlTerminalRuntimeError('not_ready', 'Terminal session is not attached', true); } if ( @@ -177,7 +181,12 @@ export function createControlTerminalRuntime(options: { false ); } - if (options.getKiloRuntime(identity.directory) !== attached.kiloRuntime) { + const runtime = options.getKiloRuntime(identity); + if ( + runtime !== attached.kiloRuntime || + (runtime?.isolation === 'per-session' && + (!runtime.identity || !sameSession(runtime.identity, identity))) + ) { throw new ControlTerminalRuntimeError('not_ready', 'Kilo worktree is not available', true); } return attached; @@ -402,12 +411,45 @@ export function createControlTerminalRuntime(options: { await Promise.allSettled(pending); } + function hasActivePty(identity: SessionRequestIdentity): boolean { + const attached = attachedSessions.get(identity.sessionId); + if (!attached || !sameSession(attached, identity)) return false; + for (const operation of operations.values()) { + if (sameSession(operation, attached)) return true; + } + for (const terminal of terminals.values()) { + if (sameSession(terminal, attached) && terminal.state === 'running') return true; + } + return false; + } + + function beginRecoveryRetirement(identity: SessionRequestIdentity): void { + const attached = attachedSessions.get(identity.sessionId); + if (attached && !sameSession(attached, identity)) { + throw new ControlTerminalRuntimeError( + 'unauthorized', + 'Terminal session ownership mismatch', + false + ); + } + if (hasActivePty(identity)) { + throw new ControlTerminalRuntimeError('session_busy', 'Session has an active PTY', true); + } + recoveringSessions.add(identity.sessionId); + } + + function endRecoveryRetirement(identity: SessionRequestIdentity): void { + recoveringSessions.delete(identity.sessionId); + } + return { rememberAttachedSession(identity) { - const kiloRuntime = options.getKiloRuntime(identity.directory); + const kiloRuntime = options.getKiloRuntime(identity); if ( shutDown || !kiloRuntime || + (kiloRuntime?.isolation === 'per-session' && + (!kiloRuntime.identity || !sameSession(kiloRuntime.identity, identity))) || directoryForSession(identity.kiloSessionId) !== identity.directory || rootForSession(identity.kiloSessionId) !== identity.kiloSessionId ) { @@ -603,6 +645,9 @@ export function createControlTerminalRuntime(options: { return connection; }, + hasActivePty, + beginRecoveryRetirement, + endRecoveryRetirement, shutdown() { if (shutDown) return; shutDown = true; @@ -616,6 +661,7 @@ export function createControlTerminalRuntime(options: { terminals.clear(); operations.clear(); attachedSessions.clear(); + recoveringSessions.clear(); }, }; } diff --git a/services/cloud-agent-next/wrapper/src/control/worktree-mutation-notifications.test.ts b/services/cloud-agent-next/wrapper/src/control/worktree-mutation-notifications.test.ts index 47568e5687..492d4a2a4b 100644 --- a/services/cloud-agent-next/wrapper/src/control/worktree-mutation-notifications.test.ts +++ b/services/cloud-agent-next/wrapper/src/control/worktree-mutation-notifications.test.ts @@ -244,19 +244,32 @@ const disposers: Array<() => void> = []; function setup(deliver?: SendEvent, signal?: AbortSignal) { const sessions: HandlerSessionSnapshot[] = []; const runtimes = new Map(); + const key = (identity: { sessionId: string; kiloSessionId: string; directory: string }) => + `${identity.sessionId}\0${identity.kiloSessionId}\0${identity.directory}`; const abort = new AbortController(); const sendEvent = mock(deliver ?? (() => true)); const notifications = createWorktreeMutationNotifications({ sessions, - kiloRuntimes: { get: directory => runtimes.get(directory) }, + kiloRuntimes: { + get: identity => + typeof identity === 'string' + ? [...runtimes.values()].find(runtime => runtime.directory === identity) + : runtimes.get(key(identity)), + isCurrent: runtime => + runtime.identity !== undefined && runtimes.get(key(runtime.identity)) === runtime, + }, signal: signal ?? abort.signal, sendEvent, }); disposers.push(notifications.dispose); - function addRuntime(directory: string) { + function addRuntime( + directory: string, + identity = { sessionId: 'workspace_root', kiloSessionId: 'root', directory } + ) { const controller = new AbortController(); let client = {} as WrapperKiloClient; const runtime: WorktreeKiloRuntime = { + identity, runtimeId: crypto.randomUUID(), directory, scopeId: directory, @@ -266,7 +279,7 @@ function setup(deliver?: SendEvent, signal?: AbortSignal) { return client; }, }; - runtimes.set(directory, runtime); + runtimes.set(key(identity), runtime); return { runtime, controller, @@ -596,6 +609,22 @@ describe('worktree mutation notifications', () => { expect(h.sendEvent.mock.calls).toEqual([expectedHint(), expectedHint('sibling')]); }); + it('fans a same-directory mutation from an isolated sibling runtime to both roots', () => { + const h = setup(); + const siblingIdentity = { + sessionId: 'workspace_sibling', + kiloSessionId: 'sibling', + directory, + }; + const siblingRuntime = h.addRuntime(directory, siblingIdentity).runtime; + h.attach('sibling'); + + h.notifications.observe(siblingRuntime, { ...fileEdited, directory }); + jest.advanceTimersByTime(5_000); + + expect(h.sendEvent.mock.calls).toEqual([expectedHint(), expectedHint('sibling')]); + }); + it('observes ambiguous sessionless feed events before routing without modifying original events or activity', async () => { const h = setup(); h.attach('sibling'); diff --git a/services/cloud-agent-next/wrapper/src/control/worktree-mutation-notifications.ts b/services/cloud-agent-next/wrapper/src/control/worktree-mutation-notifications.ts index ff2e1edf0f..36ec9eaf58 100644 --- a/services/cloud-agent-next/wrapper/src/control/worktree-mutation-notifications.ts +++ b/services/cloud-agent-next/wrapper/src/control/worktree-mutation-notifications.ts @@ -158,7 +158,9 @@ function mutationSessionId({ type, properties }: KiloEvent): string | undefined export function createWorktreeMutationNotifications(options: { sessions: readonly HandlerSessionSnapshot[]; - kiloRuntimes: Pick; + kiloRuntimes: Pick & { + isCurrent(runtime: WorktreeKiloRuntime): boolean; + }; signal: AbortSignal; sendEvent: ( event: 'session.event', @@ -175,7 +177,7 @@ export function createWorktreeMutationNotifications(options: { !disposed && !options.signal.aborted && !runtime.signal.aborted && - options.kiloRuntimes.get(runtime.directory) === runtime + options.kiloRuntimes.isCurrent(runtime) ); } diff --git a/services/cloud-agent-next/wrapper/src/control/worktree-runtime-cleanup.ts b/services/cloud-agent-next/wrapper/src/control/worktree-runtime-cleanup.ts index b59dd68cee..8973f8baf8 100644 --- a/services/cloud-agent-next/wrapper/src/control/worktree-runtime-cleanup.ts +++ b/services/cloud-agent-next/wrapper/src/control/worktree-runtime-cleanup.ts @@ -40,7 +40,9 @@ export function retireWorktreeRuntime, R const completion = Promise.withResolvers(); entry.retiring = completion.promise; const processes = - entry.processes?.stop(deadlineAt) ?? Promise.resolve(entry.processIssued !== true); + entry.processes?.stop(deadlineAt) ?? + entry.stopped?.then(() => true) ?? + Promise.resolve(entry.processIssued !== true); entry.abort.abort(); for (const root of [...entry.roots]) deps.unregisterRoot(root); const cleanup = async (): Promise => { diff --git a/services/cloud-agent-next/wrapper/src/control/worktree-runtime.test.ts b/services/cloud-agent-next/wrapper/src/control/worktree-runtime.test.ts index a75feda600..59109252d7 100644 --- a/services/cloud-agent-next/wrapper/src/control/worktree-runtime.test.ts +++ b/services/cloud-agent-next/wrapper/src/control/worktree-runtime.test.ts @@ -102,10 +102,13 @@ function asFetch( return Object.assign(fn, { preconnect: fetch.preconnect }); } -function createKiloStub(health: unknown = { healthy: true, version: '7.4.20' }) { +function createKiloStub( + health: unknown = { healthy: true, version: '7.4.20' }, + statuses: Record = {} +) { const requests: Array<{ pathname: string; directory: string | null; body?: unknown }> = []; const permissions: Awaited> = []; - const sessionStatuses: Record = {}; + const sessionStatuses: Record = { ...statuses }; const feeds = new Set>(); const encoder = new TextEncoder(); let feedConnections = 0; @@ -273,11 +276,25 @@ function createRegistry(overrides: Partial, + canRefreshCredentials?: () => boolean + ) { + return registry.attach(identity, kilo, env, canRefreshCredentials, 'per-session'); + }, get kiloCliVersion() { return registry.kiloCliVersion; }, async ensure(directory: string, kilo: WorktreeKiloAuth, env?: Record) { - const attachment = registry.attach(rootIdentity(directory), kilo, env); + const attachment = registry.attach( + rootIdentity(directory), + kilo, + env, + undefined, + 'per-session' + ); try { const runtime = await attachment.ready; attachment.commit(); @@ -287,6 +304,7 @@ function createRegistry(overrides: Partial registry.get(directory), + getKiloRuntime: identity => registry.get(identity), }); terminalRuntimes.push(terminalRuntime); return createControlHandlerDeps({ @@ -541,11 +559,122 @@ describe('worktree Kilo environments', () => { }); describe('worktree Kilo runtime registry', () => { + it('waits for exit, gates replacement, and preserves a recovery acknowledgement', async () => { + const stopped = Promise.withResolvers(); + const directory = path.join(tmpDir, 'recovery'); + const identity = rootIdentity(directory); + const stub = createKiloStub(undefined, { [identity.kiloSessionId]: { type: 'idle' } }); + const replacementStub = createKiloStub(undefined, { + [identity.kiloSessionId]: { type: 'idle' }, + }); + servers.push(stub); + servers.push(replacementStub); + let closes = 0; + let starts = 0; + const { rawRegistry } = createRegistry({ + startServer: async () => { + if (starts++ > 0) return { url: replacementStub.url, close: () => {} }; + return { + url: stub.url, + stopped: stopped.promise, + close: () => { + closes += 1; + }, + }; + }, + }); + const attachment = rawRegistry.attach(identity, auth, undefined, undefined, 'per-session'); + const runtime = await attachment.ready; + attachment.commit(); + const recoveryId = '11111111-1111-4111-8111-111111111111'; + const retirement = rawRegistry.retireForRecovery(identity, recoveryId, () => {}); + await waitUntil(() => closes === 1); + expect(rawRegistry.get(identity)).toBeUndefined(); + expect(() => rawRegistry.attach(identity, auth, undefined, undefined, 'per-session')).toThrow( + 'Kilo runtime is retiring' + ); + stopped.resolve(); + await retirement; + + const fresh = rawRegistry.attach(identity, auth, undefined, undefined, 'per-session'); + const freshRuntime = await fresh.ready; + fresh.commit(); + await rawRegistry.retireForRecovery(identity, recoveryId, () => { + throw new Error('Recovery acknowledgement must win'); + }); + expect(freshRuntime.signal.aborted).toBe(false); + expect(runtime.signal.aborted).toBe(true); + fresh.release(); + attachment.release(); + }); + + it('acknowledges an absent root without affecting a sibling or a later attachment', async () => { + const directory = path.join(tmpDir, 'cold-recovery'); + const absent = rootIdentity(directory, 'absent'); + const sibling = rootIdentity(directory, 'sibling'); + const stub = createKiloStub(undefined, { + [absent.kiloSessionId]: { type: 'idle' }, + [sibling.kiloSessionId]: { type: 'idle' }, + }); + servers.push(stub); + const { rawRegistry } = createRegistry({ + startServer: async () => ({ url: stub.url, close: () => {} }), + }); + const siblingAttachment = rawRegistry.attach( + sibling, + auth, + undefined, + undefined, + 'per-session' + ); + const siblingRuntime = await siblingAttachment.ready; + siblingAttachment.commit(); + const recoveryId = '22222222-2222-4222-8222-222222222222'; + expect(await rawRegistry.retireForRecovery(absent, recoveryId, () => {})).toBe('absent'); + expect(rawRegistry.get(sibling)).toBe(siblingRuntime); + + const pending = rawRegistry.attach(absent, auth, undefined, undefined, 'per-session'); + expect( + await rejected( + rawRegistry.retireForRecovery(absent, '33333333-3333-4333-8333-333333333333', () => {}) + ) + ).toMatchObject({ code: 'session_busy' }); + const freshRuntime = await pending.ready; + pending.commit(); + expect(await rawRegistry.retireForRecovery(absent, recoveryId, () => {})).toBe('acknowledged'); + expect(freshRuntime.signal.aborted).toBe(false); + pending.release(); + siblingAttachment.release(); + }); + + it('defaults missing runtime isolation to a directory-shared runtime', async () => { + const { rawRegistry, launches } = createRegistry(); + const directory = path.join(tmpDir, 'legacy-shared'); + const firstIdentity = rootIdentity(directory, 'first'); + const secondIdentity = rootIdentity(directory, 'second'); + const first = rawRegistry.attach(firstIdentity, auth); + const second = rawRegistry.attach(secondIdentity, auth); + try { + const [firstRuntime, secondRuntime] = await Promise.all([first.ready, second.ready]); + first.commit(); + second.commit(); + expect(secondRuntime).toBe(firstRuntime); + expect(firstRuntime.isolation).toBe('directory-shared'); + expect(launches).toHaveLength(1); + expect(rawRegistry.detach(firstIdentity)).toBe(true); + expect(secondRuntime.signal.aborted).toBe(false); + expect(rawRegistry.get(secondIdentity)).toBe(secondRuntime); + } finally { + first.release(); + second.release(); + } + }); + it('starts lazily and reuses one server and feed for concurrent same-worktree roots', async () => { const { registry, launches } = createRegistry(); const directory = path.join(tmpDir, 'worktree-a'); expect(launches).toEqual([]); - expect(registry.get(directory)).toBeUndefined(); + expect(registry.get(rootIdentity(directory))).toBeUndefined(); const [first, second, third] = await Promise.all([ registry.ensure(directory, auth), @@ -555,7 +684,7 @@ describe('worktree Kilo runtime registry', () => { expect(second).toBe(first); expect(third).toBe(first); expect(await registry.ensure(directory, auth)).toBe(first); - expect(registry.get(directory)).toBe(first); + expect(registry.get(rootIdentity(directory))).toBe(first); expect(launches).toHaveLength(1); expect(servers[0]?.feedConnections).toBe(1); @@ -657,7 +786,7 @@ describe('worktree Kilo runtime registry', () => { ]); expect(runtime.signal.aborted).toBe(true); expect(harness.registry.isHealthy()).toBe(true); - expect(harness.registry.get(runtime.directory)).toBeUndefined(); + expect(harness.registry.getRetained?.(runtime.directory)).toBeUndefined(); expect( fetchSpy.mock.calls.filter( ([request]) => @@ -671,7 +800,7 @@ describe('worktree Kilo runtime registry', () => { } ); - it('keeps the refreshed SDK event feed healthy after intentional old-process shutdown', async () => { + it('refreshes direct credentials only for their identity after intentional old-process shutdown', async () => { const received: string[] = []; const harness = createRegistry({ startServer: async options => { @@ -692,6 +821,7 @@ describe('worktree Kilo runtime registry', () => { }); const directAuth = { ...auth, containmentEnabled: false }; const identity = rootIdentity(path.join(tmpDir, 'shared')); + const siblingIdentity = rootIdentity(identity.directory, 'sibling'); const first = harness.registry.attach(identity, directAuth); const runtime = await first.ready; const cleanupOriginal = first.cleanup @@ -699,6 +829,11 @@ describe('worktree Kilo runtime registry', () => { : undefined; first.commit(); first.release(); + const siblingAttachment = harness.registry.attach(siblingIdentity, directAuth); + const sibling = await siblingAttachment.ready; + siblingAttachment.commit(); + siblingAttachment.release(); + const siblingClient = sibling.kiloClient; const originalClient = runtime.kiloClient; const originalRuntimeId = runtime.runtimeId; const refresh = harness.registry.attach( @@ -712,23 +847,26 @@ describe('worktree Kilo runtime registry', () => { expect(refreshed.runtimeId).not.toBe(originalRuntimeId); expect(runtime.signal.aborted).toBe(false); expect(await cleanupOriginal?.(Date.now() + 1_000)).toBe('stale'); - expect(harness.registry.get(identity.directory)).toBe(refreshed); + expect(harness.registry.get(identity)).toBe(refreshed); refresh.commit(); refresh.release(); expect(runtime.kiloClient).not.toBe(originalClient); - expect(servers).toHaveLength(2); - servers[1]?.emit({ payload: { type: 'server.heartbeat', properties: {} } }); - servers[1]?.emit({ payload: { type: 'session.updated', properties: {} } }); + expect(sibling.kiloClient).toBe(siblingClient); + expect(sibling.signal.aborted).toBe(false); + expect(harness.registry.get(siblingIdentity)).toBe(sibling); + expect(servers).toHaveLength(3); + servers[2]?.emit({ payload: { type: 'server.heartbeat', properties: {} } }); + servers[2]?.emit({ payload: { type: 'session.updated', properties: {} } }); await waitUntil(() => received.includes('session.updated')); expect(runtime.signal.aborted).toBe(false); expect(harness.registry.isHealthy()).toBe(true); - expect(harness.registry.get(identity.directory)).toBe(refreshed); + expect(harness.registry.get(identity)).toBe(refreshed); expect(harness.unexpectedCloses).toBe(0); servers[1]?.endFeeds(); await waitUntil(() => (servers[1]?.feedConnections ?? 0) === 2); expect(runtime.signal.aborted).toBe(false); expect(harness.registry.isHealthy()).toBe(true); - expect(harness.registry.get(identity.directory)).toBe(refreshed); + expect(harness.registry.get(identity)).toBe(refreshed); expect(harness.unexpectedCloses).toBe(0); }); @@ -776,7 +914,7 @@ describe('worktree Kilo runtime registry', () => { ); }); - it('routes separate SandboxSession roots sharing a worktree through one SDK runtime', async () => { + it('routes separate SandboxSession roots sharing a worktree through isolated SDK runtimes', async () => { const identities = [ { sessionId: 'workspace_first', @@ -812,7 +950,7 @@ describe('worktree Kilo runtime registry', () => { const terminals = createControlTerminalRuntime({ controlUrl: 'ws://127.0.0.1:1/sandbox-control/test', wrapperInstanceId: crypto.randomUUID(), - getKiloRuntime: directory => harness.registry.get(directory), + getKiloRuntime: identity => harness.registry.get(identity), }); const deps: HandlerDeps = createControlHandlerDeps({ kiloRuntimes: harness.registry, @@ -849,11 +987,24 @@ describe('worktree Kilo runtime registry', () => { ) ); expect(attached).toEqual(identities.map(() => ({ ok: true, result: { attached: true } }))); - expect(harness.launches).toHaveLength(2); - expect(servers.map(server => server.feedConnections)).toEqual([1, 1]); + expect(harness.launches).toHaveLength(3); + expect(servers.map(server => server.feedConnections)).toEqual([1, 1, 1]); + const sameDirectoryRuntimes = identities.slice(0, 2).map(identity => { + const runtime = harness.registry.get(identity); + if (!runtime) throw new Error('Expected same-directory runtime'); + return runtime; + }); + expect(sameDirectoryRuntimes[0].kiloClient).not.toBe(sameDirectoryRuntimes[1].kiloClient); + expect(sameDirectoryRuntimes[0].kiloClient.serverUrl).not.toBe( + sameDirectoryRuntimes[1].kiloClient.serverUrl + ); + expect(sameDirectoryRuntimes[0].env.HOME).not.toBe(sameDirectoryRuntimes[1].env.HOME); + expect(path.join(sameDirectoryRuntimes[0].env.XDG_DATA_HOME, 'kilo', 'auth.json')).not.toBe( + path.join(sameDirectoryRuntimes[1].env.XDG_DATA_HOME, 'kilo', 'auth.json') + ); for (const identity of identities) { - const runtime = harness.registry.get(identity.directory); + const runtime = harness.registry.get(identity); const server = servers.find(server => server.url === runtime?.kiloClient.serverUrl); if (!runtime || !server) throw new Error('Expected attached worktree runtime'); const before = server.requests.length; @@ -997,13 +1148,14 @@ describe('worktree Kilo runtime registry', () => { } const [first, second] = identities; - const runtime = harness.registry.get(first.directory); + const runtime = harness.registry.get(first); expect(await handleControlRequest('session.detach', second, {}, deps)).toEqual({ ok: true, result: { detached: true }, }); - expect(harness.registry.get(first.directory)).toBe(runtime); - expect(harness.closes).toBe(0); + expect(harness.registry.get(first)).toBe(runtime); + expect(harness.registry.get(second)).toBeUndefined(); + expect(harness.closes).toBe(1); expect(rootForSession(undefined, first.directory)).toBe(first.kiloSessionId); expect(rootForSession(`child_${second.kiloSessionId}`, first.directory)).toBeUndefined(); const survivingPrompt = { @@ -1028,15 +1180,15 @@ describe('worktree Kilo runtime registry', () => { expect((await handleControlRequest('session.prompt', second, survivingPrompt, deps)).ok).toBe( false ); - expect(harness.launches).toHaveLength(2); + expect(harness.launches).toHaveLength(3); expect(await handleControlRequest('session.detach', first, {}, deps)).toEqual({ ok: true, result: { detached: true }, }); expect(runtime?.signal.aborted).toBe(true); - expect(harness.registry.get(first.directory)).toBeUndefined(); - expect(harness.closes).toBe(1); + expect(harness.registry.get(first)).toBeUndefined(); + expect(harness.closes).toBe(2); expect(rootForSession(first.kiloSessionId)).toBeUndefined(); expect(rootForSession(`child_${first.kiloSessionId}`)).toBeUndefined(); expect( @@ -1079,7 +1231,7 @@ describe('worktree Kilo runtime registry', () => { ) ).ok ).toBe(true); - expect(harness.registry.get(first.directory)).not.toBe(runtime); + expect(harness.registry.get(first)).not.toBe(runtime); expect( ( await handleControlRequest( @@ -1092,14 +1244,14 @@ describe('worktree Kilo runtime registry', () => { ) ).ok ).toBe(true); - expect(harness.launches).toHaveLength(3); + expect(harness.launches).toHaveLength(4); expect(harness.unexpectedCloses).toBe(0); } finally { terminals.shutdown(); } }); - it('retires only the final root runtime and permits reuse without affecting another worktree', async () => { + it('retires each same-directory identity independently and permits reuse without affecting another worktree', async () => { const harness = createRegistry(); const directory = path.join(tmpDir, 'shared'); const first = await harness.registry.ensure(directory, auth); @@ -1113,18 +1265,20 @@ describe('worktree Kilo runtime registry', () => { fs.writeFileSync(marker, 'ready'); expect(harness.registry.detach(rootIdentity(directory))).toBe(true); - expect(harness.registry.get(directory)).toBe(first); + expect(first.signal.aborted).toBe(true); + expect(harness.registry.get(rootIdentity(directory))).toBeUndefined(); expect(sibling.signal.aborted).toBe(false); - expect(await sibling.ready).toBe(first); + const siblingRuntime = await sibling.ready; + expect(siblingRuntime).not.toBe(first); sibling.commit(); sibling.release(); - expect(harness.closes).toBe(0); + expect(harness.closes).toBe(1); expect(harness.registry.detach(siblingIdentity)).toBe(true); - expect(first.signal.aborted).toBe(true); - expect(harness.registry.get(directory)).toBeUndefined(); - expect(harness.closes).toBe(1); - expect(harness.registry.get(other.directory)).toBe(other); + expect(siblingRuntime.signal.aborted).toBe(true); + expect(harness.registry.get(siblingIdentity)).toBeUndefined(); + expect(harness.closes).toBe(2); + expect(harness.registry.get(rootIdentity(other.directory))).toBe(other); await other.kiloClient.abortSession({ sessionId: 'root_other' }); const otherServer = servers.find(server => server.url === other.kiloClient.serverUrl); expect(otherServer?.requests.at(-1)?.pathname).toBe('/session/root_other/abort'); @@ -1137,9 +1291,9 @@ describe('worktree Kilo runtime registry', () => { expect(replacement.env.HOME).toBe(first.env.HOME); expect(replacement.env.KILOCODE_TOKEN).toBe('replacement-token'); expect(fs.readFileSync(marker, 'utf8')).toBe('ready'); - expect(harness.registry.get(other.directory)).toBe(other); + expect(harness.registry.get(rootIdentity(other.directory))).toBe(other); expect(harness.unexpectedCloses).toBe(0); - expect(harness.launches).toHaveLength(3); + expect(harness.launches).toHaveLength(4); }); it('does not accumulate duplicate roots or release a committed root on a failed retry', async () => { @@ -1157,7 +1311,7 @@ describe('worktree Kilo runtime registry', () => { expect(await retry.ready).toBe(runtime); retry.release(); - expect(harness.registry.get(identity.directory)).toBe(runtime); + expect(harness.registry.get(identity)).toBe(runtime); expect(harness.closes).toBe(0); expect(harness.launches).toHaveLength(1); expect(harness.registry.detach(identity)).toBe(true); @@ -1188,36 +1342,42 @@ describe('worktree Kilo runtime registry', () => { expect(harness.closes).toBe(1); }); - it('rejects mismatched root identities while startup is pending without disturbing ownership', async () => { + it('creates independent runtimes for distinct identities without disturbing pending ownership', async () => { const harness = createRegistry(); const identity = rootIdentity(path.join(tmpDir, 'shared')); const attachment = harness.registry.attach(identity, auth); - for (const mismatch of [ + const distinct = [ { ...identity, directory: path.join(tmpDir, 'other') }, { ...identity, sessionId: 'workspace_foreign' }, { ...identity, kiloSessionId: 'root_foreign' }, - ]) { - expect(() => harness.registry.attach(mismatch, auth)).toThrow('Session identity mismatch'); - expect(() => harness.registry.detach(mismatch)).toThrow('Session identity mismatch'); - } + ]; + const attachments = distinct.map(candidate => harness.registry.attach(candidate, auth)); const runtime = await attachment.ready; attachment.commit(); - expect(harness.registry.get(identity.directory)).toBe(runtime); + expect(harness.registry.get(identity)).toBe(runtime); + for (const candidate of attachments) { + const candidateRuntime = await candidate.ready; + expect(candidateRuntime).not.toBe(runtime); + candidate.commit(); + } expect(harness.registry.detach(rootIdentity(identity.directory, 'unknown'))).toBe(false); expect(harness.closes).toBe(0); + for (const candidate of distinct) expect(harness.registry.detach(candidate)).toBe(true); + expect(harness.registry.detach(identity)).toBe(true); }); - it('keeps a pending sibling startup alive while cancelling only the detached root', async () => { + it('keeps a pending sibling startup alive while cancelling only its detached identity', async () => { const launched = Promise.withResolvers(); const release = Promise.withResolvers(); const server = createKiloStub(); servers.push(server); let closes = 0; + let starts = 0; const { registry } = createRegistry({ startServer: async options => { proveOwnedProcesses(options); launched.resolve(); - await release.promise; + if (starts++ === 0) await release.promise; return { url: server.url, close: () => { @@ -1229,23 +1389,24 @@ describe('worktree Kilo runtime registry', () => { const directory = path.join(tmpDir, 'shared'); const firstIdentity = rootIdentity(directory, 'first'); const first = registry.attach(firstIdentity, auth); + await launched.promise; const siblingIdentity = rootIdentity(directory, 'sibling'); const sibling = registry.attach(siblingIdentity, auth); try { - await launched.promise; expect(registry.detach(firstIdentity)).toBe(true); expect(first.signal.aborted).toBe(true); expect(sibling.signal.aborted).toBe(false); - release.resolve(); const runtime = await sibling.ready; - await first.ready; + expect(runtime.identity).toEqual(siblingIdentity); + release.resolve(); + expect(await rejected(first.ready)).toMatchObject({ code: 'not_ready' }); expect(() => first.commit()).toThrow(); first.release(); sibling.commit(); - expect(registry.get(directory)).toBe(runtime); - expect(closes).toBe(0); - expect(registry.detach(siblingIdentity)).toBe(true); + expect(registry.get(siblingIdentity)).toBe(runtime); expect(closes).toBe(1); + expect(registry.detach(siblingIdentity)).toBe(true); + expect(closes).toBe(2); } finally { release.resolve(); await Promise.allSettled([first.ready, sibling.ready]); @@ -1308,7 +1469,7 @@ describe('worktree Kilo runtime registry', () => { expect(steps).toEqual(['start-old', 'close-old', 'start-new']); expect(oldServer.feedConnections).toBe(0); expect(newServer.feedConnections).toBe(1); - expect(harness.registry.get(identity.directory)).toBe(runtime); + expect(harness.registry.get(identity)).toBe(runtime); expect( JSON.parse( fs.readFileSync(path.join(runtime.env.XDG_DATA_HOME, 'kilo', 'auth.json'), 'utf8') @@ -1316,12 +1477,13 @@ describe('worktree Kilo runtime registry', () => { ).toEqual({ kilo: { type: 'api', key: 'replacement' }, }); - expect(() => - harness.registry.attach(rootIdentity(path.join(tmpDir, 'other')), { - ...auth, - token: 'replacement', - }) - ).toThrow('Kilo worktree auth context mismatch'); + const independent = harness.registry.attach(rootIdentity(path.join(tmpDir, 'other')), { + ...auth, + token: 'replacement', + }); + expect((await independent.ready).env.KILOCODE_TOKEN).toBe('replacement'); + independent.commit(); + independent.release(); expect(harness.unexpectedCloses).toBe(0); } finally { release.resolve(); @@ -1330,7 +1492,7 @@ describe('worktree Kilo runtime registry', () => { } }); - it('rejects scope, directory, token, and target changes without launching another server', async () => { + it('rejects contained credential changes for one identity while allowing a distinct identity', async () => { const { registry, launches } = createRegistry(); const directory = path.join(tmpDir, 'worktree-a'); await registry.ensure(directory, auth); @@ -1350,10 +1512,9 @@ describe('worktree Kilo runtime registry', () => { retryable: false, }); } - expect(await rejected(registry.ensure(path.join(tmpDir, 'worktree-b'), auth))).toMatchObject({ - code: 'unauthorized', - }); - expect(launches).toHaveLength(1); + const distinct = await registry.ensure(path.join(tmpDir, 'worktree-b'), auth); + expect(distinct).not.toBe(registry.get(rootIdentity(directory))); + expect(launches).toHaveLength(2); }); it('releases a failed sole attachment and permits retry with fresh auth', async () => { @@ -1372,7 +1533,7 @@ describe('worktree Kilo runtime registry', () => { expect(await rejected(registry.ensure(directory, auth))).toMatchObject({ message: 'Kilo worktree failed to start', }); - expect(registry.get(directory)).toBeUndefined(); + expect(registry.get(rootIdentity(directory))).toBeUndefined(); const runtime = await registry.ensure(directory, { ...auth, token: 'changed-token' }); expect(runtime.scopeId).toBe(auth.scopeId); expect(runtime.env.KILOCODE_TOKEN).toBe('changed-token'); @@ -1398,7 +1559,8 @@ describe('worktree Kilo runtime registry', () => { expect(first.signal.aborted).toBe(true); expect(second.signal.aborted).toBe(true); expect(harness.closes).toBe(2); - expect(harness.registry.get(first.directory)).toBeUndefined(); + if (!first.identity) throw new Error('Expected runtime identity'); + expect(harness.registry.get(first.identity)).toBeUndefined(); expect(await rejected(harness.registry.ensure(first.directory, auth))).toMatchObject({ message: 'Kilo worktrees are closed', }); @@ -1444,7 +1606,8 @@ describe('worktree Kilo runtime registry', () => { servers[0]?.endFeeds(); await waitUntil(() => (servers[0]?.feedConnections ?? 0) === 2); expect(runtime.signal.aborted).toBe(false); - expect(harness.registry.get(runtime.directory)).toBe(runtime); + if (!runtime.identity) throw new Error('Expected runtime identity'); + expect(harness.registry.get(runtime.identity)).toBe(runtime); expect(harness.closes).toBe(0); expect(harness.unexpectedCloses).toBe(0); }); @@ -1467,13 +1630,14 @@ describe('worktree Kilo runtime registry', () => { expect(failures).toEqual([ expect.objectContaining({ directory: runtime.directory, + identity: runtime.identity, reason: 'process_exited', cleanup: 'confirmed', runtimeId: runtime.runtimeId, }), ]); expect(runtime.signal.aborted).toBe(true); - expect(harness.registry.get(runtime.directory)).toBeUndefined(); + expect(harness.registry.getRetained?.(runtime.directory)).toBeUndefined(); }); it('preserves an admitted operation and its runtime while a real feed recovers', async () => { @@ -1528,7 +1692,7 @@ describe('worktree Kilo runtime registry', () => { server.endFeeds(); await waitUntil(() => server.feedConnections === 2); expect(harness.registry.prepareForNewWork?.(identity.directory)).toBe(false); - expect(harness.registry.get(identity.directory)).toBe(runtime); + expect(harness.registry.get(identity)).toBe(runtime); expect(runtime.kiloClient).toBe(client); expect(runtime.kiloClient.serverUrl).toBe(client.serverUrl); expect( @@ -1578,7 +1742,7 @@ describe('worktree Kilo runtime registry', () => { server.emit({ payload: { type: 'server.heartbeat', properties: {} } }); await waitUntil(() => harness.registry.prepareForNewWork?.(identity.directory) === true); - expect(harness.registry.get(identity.directory)).toBe(runtime); + expect(harness.registry.get(identity)).toBe(runtime); expect(runtime.kiloClient).toBe(client); expect(harness.unexpectedCloses).toBe(0); } finally { @@ -1623,7 +1787,7 @@ describe('worktree directory deletion', () => { expect(closed).toEqual([directory]); expect(deleted).toBe(false); expect(fs.existsSync(authFile)).toBe(true); - expect(harness.registry.get(other.directory)).toBe(other); + expect(harness.registry.get(rootIdentity(other.directory))).toBe(other); expect(fs.existsSync(other.env.HOME)).toBe(true); expect(other.signal.aborted).toBe(false); fs.writeFileSync(path.join(runtime.env.HOME, 'last-write-before-exit'), 'stopping'); @@ -1708,7 +1872,7 @@ describe('worktree directory deletion', () => { const liveAuth = { ...auth, scopeId: 'new_scope', token: 'new_guest' }; const live = await harness.registry.ensure(directory, liveAuth); const liveHome = live.env.HOME; - expect(liveHome).not.toBe(retiredHome); + expect(liveHome).toBe(retiredHome); const siblingIdentity = rootIdentity(directory, 'sibling'); const sibling = harness.registry.attach(siblingIdentity, liveAuth); await sibling.ready; @@ -1726,7 +1890,7 @@ describe('worktree directory deletion', () => { fs.writeFileSync(checkoutFile, 'owned by checkout cleanup'); live.env.HOME = other.env.HOME; - expect(harness.registry.get(directory)).toBe(live); + expect(harness.registry.get(identity)).toBe(live); await harness.registry.deleteDirectory(directory); expect(fs.existsSync(retiredHome)).toBe(false); expect(fs.existsSync(liveHome)).toBe(false); @@ -1734,25 +1898,26 @@ describe('worktree directory deletion', () => { expect(fs.readFileSync(otherAuthFile, 'utf8')).toBe(otherAuthBefore); expect(live.signal.aborted).toBe(true); expect(sibling.signal.aborted).toBe(true); - expect(harness.registry.get(directory)).toBeUndefined(); + expect(harness.registry.get(identity)).toBeUndefined(); for (const id of [identity.kiloSessionId, siblingIdentity.kiloSessionId, 'deleted_child']) { expect(rootForSession(id)).toBeUndefined(); expect(directoryForSession(id)).toBeUndefined(); } - expect(harness.registry.get(other.directory)).toBe(other); + if (!other.identity) throw new Error('Expected runtime identity'); + expect(harness.registry.get(other.identity)).toBe(other); expect(other.signal.aborted).toBe(false); expect(rootForSession('surviving_child')).toBe(otherIdentity.kiloSessionId); expect(await other.kiloClient.abortSession({ sessionId: otherIdentity.kiloSessionId })).toBe( true ); - expect(harness.closes).toBe(3); + expect(harness.closes).toBe(4); expect(harness.registry.isHealthy()).toBe(true); expect(() => harness.registry.attach(rootIdentity(directory, 'new_root'), liveAuth)).toThrow( 'Kilo worktree is deleted' ); await harness.registry.deleteDirectory(directory); - expect(harness.launches).toHaveLength(4); - expect(harness.closes).toBe(3); + expect(harness.launches).toHaveLength(5); + expect(harness.closes).toBe(4); expect(harness.unexpectedCloses).toBe(0); }); @@ -1765,7 +1930,7 @@ describe('worktree directory deletion', () => { homes.push(runtime.env.HOME); expect(harness.registry.detach(rootIdentity(directory))).toBe(true); await new Promise(resolve => setImmediate(resolve)); - expect(harness.registry.get(directory)).toBeUndefined(); + expect(harness.registry.get(rootIdentity(directory, scopeId))).toBeUndefined(); expect(fs.existsSync(runtime.env.HOME)).toBe(true); } await Promise.all([ @@ -1792,7 +1957,7 @@ describe('worktree directory deletion', () => { ); await deletion; await harness.registry.deleteDirectory(directory); - expect(harness.registry.get(directory)).toBeUndefined(); + expect(harness.registry.get(rootIdentity(directory))).toBeUndefined(); expect(harness.registry.detach(rootIdentity(directory))).toBe(false); expect(fs.readFileSync(path.join(directory, 'keep'), 'utf8')).toBe('checkout'); expect(fs.existsSync(path.join(tmpDir, 'homes'))).toBe(false); @@ -1865,7 +2030,7 @@ describe('worktree directory deletion', () => { expect(steps).toEqual(['late-write', 'close', 'deleted']); expect(fs.existsSync(options.env.HOME)).toBe(false); expect(fs.readdirSync(path.join(tmpDir, 'homes'))).toEqual([]); - expect(harness.registry.get(directory)).toBeUndefined(); + expect(harness.registry.get(rootIdentity(directory))).toBeUndefined(); attachment.release(); for (const replacement of replacements) replacement.release(); expect(launches).toBe(1); @@ -1915,7 +2080,7 @@ describe('worktree attachment lifecycle', () => { } ); expect(result.ok).toBe(false); - expect(harness.registry.get(identity.directory)).toBeUndefined(); + expect(harness.registry.get(identity)).toBeUndefined(); expect(rootForSession(identity.kiloSessionId)).toBeUndefined(); expect(directoryForSession(identity.kiloSessionId)).toBeUndefined(); expect(harness.closes).toBe(1); @@ -1932,7 +2097,7 @@ describe('worktree attachment lifecycle', () => { ) ).ok ).toBe(true); - expect(harness.registry.get(identity.directory)?.env.KILOCODE_TOKEN).toBe('replacement'); + expect(harness.registry.get(identity)?.env.KILOCODE_TOKEN).toBe('replacement'); expect( ( await handleControlRequest( @@ -1949,7 +2114,7 @@ describe('worktree attachment lifecycle', () => { } ); - it('keeps a restoring sibling alive after the last committed root detaches', async () => { + it('keeps a restoring same-directory sibling alive after another identity detaches', async () => { const harness = createRegistry(); const deps = createHandlerDeps(harness.registry); const first = rootIdentity(path.join(tmpDir, 'shared'), 'first'); @@ -1957,7 +2122,7 @@ describe('worktree attachment lifecycle', () => { expect((await handleControlRequest('session.attach', first, { kilo: auth }, deps)).ok).toBe( true ); - const runtime = harness.registry.get(first.directory); + const runtime = harness.registry.get(first); const restoring = Promise.withResolvers(); const release = Promise.withResolvers(); const pending = applySessionAttach( @@ -1981,16 +2146,18 @@ describe('worktree attachment lifecycle', () => { try { await restoring.promise; expect((await handleControlRequest('session.detach', first, {}, deps)).ok).toBe(true); - expect(runtime?.signal.aborted).toBe(false); - expect(harness.closes).toBe(0); + expect(runtime?.signal.aborted).toBe(true); + expect(harness.closes).toBe(1); release.resolve(); expect(await pending).toEqual({ ok: true, result: { attached: true } }); - expect(harness.registry.get(first.directory)).toBe(runtime); + const siblingRuntime = harness.registry.get(sibling); + expect(siblingRuntime).toBeDefined(); + expect(siblingRuntime).not.toBe(runtime); expect(rootForSession(first.kiloSessionId)).toBeUndefined(); expect(rootForSession(sibling.kiloSessionId)).toBe(sibling.kiloSessionId); expect((await handleControlRequest('session.detach', sibling, {}, deps)).ok).toBe(true); - expect(runtime?.signal.aborted).toBe(true); - expect(harness.closes).toBe(1); + expect(siblingRuntime?.signal.aborted).toBe(true); + expect(harness.closes).toBe(2); } finally { release.resolve(); await pending; @@ -2021,10 +2188,10 @@ describe('worktree attachment lifecycle', () => { expect( (await handleControlRequest('session.attach', identity, { kilo: auth }, deps)).ok ).toBe(true); - const runtime = harness.registry.get(identity.directory); + const runtime = harness.registry.get(identity); release.resolve(); expect((await pending).ok).toBe(false); - expect(harness.registry.get(identity.directory)).toBe(runtime); + expect(harness.registry.get(identity)).toBe(runtime); expect(rootForSession(identity.kiloSessionId)).toBe(identity.kiloSessionId); expect(harness.closes).toBe(0); expect((await handleControlRequest('session.detach', identity, {}, deps)).ok).toBe(true); @@ -2044,7 +2211,7 @@ describe('worktree attachment lifecycle', () => { expect((await handleControlRequest('session.attach', sibling, { kilo: auth }, deps)).ok).toBe( true ); - const runtime = harness.registry.get(identity.directory); + const siblingRuntime = harness.registry.get(sibling); const restoring = Promise.withResolvers(); const release = Promise.withResolvers(); const pending = applySessionAttach( @@ -2075,14 +2242,16 @@ describe('worktree attachment lifecycle', () => { expect(rootForSession(identity.kiloSessionId)).toBeUndefined(); expect(rootForSession('old_child')).toBeUndefined(); expect(signal.aborted).toBe(true); - expect(runtime?.signal.aborted).toBe(false); + expect(siblingRuntime?.signal.aborted).toBe(false); expect( (await handleControlRequest('session.attach', identity, { kilo: auth }, deps)).ok ).toBe(true); rememberChildSession({ childId: 'replacement_child', parentId: identity.kiloSessionId }); release.resolve(); expect((await pending).ok).toBe(false); - expect(harness.registry.get(identity.directory)).toBe(runtime); + const replacementRuntime = harness.registry.get(identity); + expect(replacementRuntime).toBeDefined(); + expect(replacementRuntime).not.toBe(siblingRuntime); expect(rootForSession(identity.kiloSessionId)).toBe(identity.kiloSessionId); expect(rootForSession('replacement_child')).toBe(identity.kiloSessionId); expect(rootForSession('old_child')).toBeUndefined(); @@ -2099,8 +2268,8 @@ describe('worktree attachment lifecycle', () => { ) ).ok ).toBe(true); - expect(harness.closes).toBe(0); - expect(harness.launches).toHaveLength(1); + expect(harness.closes).toBe(1); + expect(harness.launches).toHaveLength(3); } finally { release.resolve(); await pending; @@ -2149,7 +2318,7 @@ describe('worktree attachment lifecycle', () => { expect(await replacement).toEqual({ ok: true, result: { attached: true } }); expect(markers).toBe(0); expect(rootForSession(identity.kiloSessionId)).toBe(identity.kiloSessionId); - expect(harness.registry.get(identity.directory)?.env.KILOCODE_TOKEN).toBe('replacement'); + expect(harness.registry.get(identity)?.env.KILOCODE_TOKEN).toBe('replacement'); } finally { release.resolve(); await pending; @@ -2328,8 +2497,70 @@ setInterval(() => {}, 1000); ).toBe('retired'); expect(first.signal.aborted).toBe(true); expect(sibling.signal.aborted).toBe(true); - expect(registry.get(sharedDirectory)).toBeUndefined(); - expect(registry.get(isolatedDirectory)).toBe(isolatedRuntime); + expect(registry.getRetained?.(sharedDirectory)).toBeUndefined(); + expect(registry.getRetained?.(isolatedDirectory)).toBe(isolatedRuntime); + }); + + it('aborts the second isolated runtime through handler dependencies without stopping its sibling', async () => { + const registry = createWorktreeKiloRuntimes({ + homeRoot: path.join(tmpDir, 'homes'), + inheritedEnv: inherited, + startServer: async options => { + const server = createKiloStub(); + servers.push(server); + options.onProcessScope?.({ + spawn: () => { + throw new Error('Unexpected process spawn'); + }, + run: operation => operation(), + seal: () => {}, + dispose: () => true, + observesOccupancy: () => true, + captureBaseline: async () => {}, + stop: async () => true, + verify: async () => true, + }); + return { url: server.url, close: () => {} }; + }, + onUnexpectedClose: () => {}, + }); + registries.push(registry); + const directory = path.join(tmpDir, 'shared'); + const firstIdentity = rootIdentity(directory, 'first'); + const secondIdentity = rootIdentity(directory, 'second'); + const first = registry.attach(firstIdentity, auth, {}, undefined, 'per-session'); + const firstRuntime = await first.ready; + first.commit(); + const second = registry.attach(secondIdentity, auth, {}, undefined, 'per-session'); + const secondRuntime = await second.ready; + second.commit(); + const target = { runtimeId: secondRuntime.runtimeId, client: secondRuntime.kiloClient }; + + expect(registry.getRetained?.(directory, secondRuntime.runtimeId)).toBe(secondRuntime); + expect(await registry.verifyQuiescence?.(directory, target, Date.now() + 1_000)).toBe(true); + const handlerDeps = createHandlerDeps(registry); + expect( + await handleControlRequest( + 'session.abort', + secondIdentity, + { nativeRuntimeId: secondRuntime.runtimeId, cleanupDeadlineAt: Date.now() + 1_000 }, + handlerDeps + ) + ).toMatchObject({ + ok: true, + result: { + status: 'aborted', + quiescent: true, + runtimeRetired: true, + nativeRuntimeId: secondRuntime.runtimeId, + }, + }); + expect(second.signal.aborted).toBe(true); + expect(first.signal.aborted).toBe(false); + expect(registry.get(firstIdentity)).toBe(firstRuntime); + expect(registry.getRetained?.(directory, secondRuntime.runtimeId)).toBeUndefined(); + expect(await registry.retireRuntime?.(directory, Date.now() + 1_000, target)).toBe('stale'); + expect(first.signal.aborted).toBe(false); }); it('retains failed native cleanup ownership without affecting another runtime', async () => { @@ -2370,7 +2601,7 @@ setInterval(() => {}, 1000); expect(first.signal.aborted).toBe(true); expect(sibling.signal.aborted).toBe(true); expect(registry.getRetained?.(failedDirectory)).toBe(runtime); - expect(registry.get(isolatedDirectory)).toBe(isolatedRuntime); + expect(registry.getRetained?.(isolatedDirectory)).toBe(isolatedRuntime); expect(() => registry.attach(rootIdentity(failedDirectory, 'retry'), auth)).toThrow( 'Native runtime retirement is unconfirmed' ); diff --git a/services/cloud-agent-next/wrapper/src/control/worktree-runtime.ts b/services/cloud-agent-next/wrapper/src/control/worktree-runtime.ts index 8b1a821fb9..4b40ab9a65 100644 --- a/services/cloud-agent-next/wrapper/src/control/worktree-runtime.ts +++ b/services/cloud-agent-next/wrapper/src/control/worktree-runtime.ts @@ -29,8 +29,10 @@ import { withKiloRequestDeadline, type KiloEventFeedError } from './sandbox-cont import { createWorktreeFeed, type KiloFeedEvent, type WorktreeFeed } from './worktree-feed.js'; export type WorktreeKiloAuth = NonNullable; +type RuntimeIsolation = 'directory-shared' | 'per-session'; type WorktreeKiloFailure = { + identity: SessionRequestIdentity; retirementId: string; directory: string; reason: KiloEventFeedError['reason'] | 'process_exited' | 'credential_refresh_failed'; @@ -40,6 +42,8 @@ type WorktreeKiloFailure = { }; export type WorktreeKiloRuntime = { + readonly identity?: SessionRequestIdentity; + readonly isolation?: RuntimeIsolation; readonly scopeId: string; readonly runtimeId: string; readonly directory: string; @@ -56,6 +60,8 @@ export type WorktreeKiloAttachment = { release(): void; }; +type RecoveryRetirement = 'retired' | 'absent' | 'acknowledged'; + export type WorktreeKiloRuntimes = { readonly kiloCliVersion?: string | null; attach( @@ -63,10 +69,16 @@ export type WorktreeKiloRuntimes = { kilo: WorktreeKiloAuth, env?: Record, canRefreshCredentials?: () => boolean, + runtimeIsolation?: RuntimeIsolation, beforeMutation?: () => void, onCleanupTarget?: (cleanup: (deadlineAt: number) => Promise) => void ): WorktreeKiloAttachment; detach(identity: SessionRequestIdentity): boolean; + retireForRecovery( + identity: SessionRequestIdentity, + recoveryId: string, + assertIdle: () => void + ): Promise; deleteDirectory(directory: string): Promise; retireRuntime?( directory: string, @@ -78,8 +90,10 @@ export type WorktreeKiloRuntimes = { target: NativeOperationTarget, deadlineAt: number ): Promise; - getRetained?(directory: string): WorktreeKiloRuntime | undefined; - get(directory: string): WorktreeKiloRuntime | undefined; + getRetained?(directory: string, runtimeId?: string): WorktreeKiloRuntime | undefined; + get(identity: SessionRequestIdentity | string): WorktreeKiloRuntime | undefined; + getAll?(directory: string): WorktreeKiloRuntime[]; + isCurrent?(runtime: WorktreeKiloRuntime): boolean; prepareForNewWork?(directory: string): boolean; isHealthy(): boolean; shutdown(): void; @@ -104,6 +118,8 @@ type WorktreeKiloServerHandle = Omit & { }; type RuntimeEntry = { + identity: SessionRequestIdentity; + isolation: RuntimeIsolation; kilo: WorktreeKiloAuth; directory: string; env: Record; @@ -338,6 +354,8 @@ export function createWorktreeKiloRuntimes(options: { const directoriesByScope = new Map(); const failedDirectories = new Set(); const roots = new Map(); + const recoveryGates = new Map>(); + const recoveryAcknowledgements = new Map>>(); const homesByDirectory = new Map>(); const deletedDirectories = new Set(); let observedVersion: string | null | undefined; @@ -348,24 +366,14 @@ export function createWorktreeKiloRuntimes(options: { observedVersion = observedVersion === undefined || observedVersion === version ? version : null; } + const identityKey = (identity: SessionRequestIdentity): string => + `${identity.sessionId}\0${identity.kiloSessionId}\0${identity.directory}`; + + const entryKey = (identity: SessionRequestIdentity, isolation: RuntimeIsolation): string => + isolation === 'per-session' ? identityKey(identity) : identity.directory; + function findRoot(identity: SessionRequestIdentity): RootAttachment | undefined { - for (const root of roots.values()) { - if ( - root.identity.kiloSessionId !== identity.kiloSessionId && - root.identity.sessionId !== identity.sessionId - ) { - continue; - } - if ( - root.identity.sessionId !== identity.sessionId || - root.identity.kiloSessionId !== identity.kiloSessionId || - root.identity.directory !== identity.directory - ) { - throw new WorktreeKiloRuntimeError('unauthorized', 'Session identity mismatch', false); - } - return root; - } - return undefined; + return roots.get(identityKey(identity)); } function cleanupDeadline(entry: RuntimeEntry, requested?: number): number { @@ -381,8 +389,8 @@ export function createWorktreeKiloRuntimes(options: { root.entry.roots.delete(root); root.attached = false; root.pending.clear(); - if (roots.get(root.identity.kiloSessionId) === root) { - roots.delete(root.identity.kiloSessionId); + if (roots.get(identityKey(root.identity)) === root) { + roots.delete(identityKey(root.identity)); forgetAttachedRoot(root.identity.kiloSessionId, root.identity.directory); } } @@ -396,8 +404,8 @@ export function createWorktreeKiloRuntimes(options: { cleanupDeadline, unregisterRoot, removeEntry: retiring => { - if (entries.get(retiring.directory) === retiring) { - entries.delete(retiring.directory); + if (entries.get(entryKey(retiring.identity, retiring.isolation)) === retiring) { + entries.delete(entryKey(retiring.identity, retiring.isolation)); if (directoriesByScope.get(retiring.kilo.scopeId) === retiring.directory) directoriesByScope.delete(retiring.kilo.scopeId); } @@ -406,7 +414,7 @@ export function createWorktreeKiloRuntimes(options: { } function removeRoot(root: RootAttachment): void { - if (roots.get(root.identity.kiloSessionId) !== root) return; + if (roots.get(identityKey(root.identity)) !== root) return; unregisterRoot(root); root.abort.abort(); if (root.entry.roots.size === 0) void retire(root.entry); @@ -419,6 +427,7 @@ export function createWorktreeKiloRuntimes(options: { void retire(entry, deadlineAt).then(result => { if (result === 'unconfirmed') failedDirectories.add(entry.directory); options.onUnexpectedClose({ + identity: entry.identity, retirementId: crypto.randomUUID(), directory: entry.directory, reason, @@ -516,6 +525,8 @@ export function createWorktreeKiloRuntimes(options: { }; entry.kiloClient = runtimeKiloClient; const runtime: WorktreeKiloRuntime = entry.runtime ?? { + identity: { ...entry.identity }, + isolation: entry.isolation, scopeId: entry.kilo.scopeId, get runtimeId() { return entry.runtimeId; @@ -540,7 +551,7 @@ export function createWorktreeKiloRuntimes(options: { signal: abort.signal, }, isCurrent: (runtimeId, client) => - entries.get(entry.directory) === entry && + entries.get(entryKey(entry.identity, entry.isolation)) === entry && entry.runtimeId === runtimeId && entry.kiloClient === client && entry.processAbort === abort, @@ -658,11 +669,24 @@ export function createWorktreeKiloRuntimes(options: { get kiloCliVersion() { return observedVersion ?? null; }, - attach(identity, kilo, environment, canRefreshCredentials, beforeMutation, onCleanupTarget) { + attach( + identity, + kilo, + environment, + canRefreshCredentials, + runtimeIsolation, + beforeMutation, + onCleanupTarget + ) { if (closed) { throw new WorktreeKiloRuntimeError('not_ready', 'Kilo worktrees are closed', false); } const { directory } = identity; + const isolation = runtimeIsolation ?? 'directory-shared'; + const key = entryKey(identity, isolation); + if (recoveryGates.has(identityKey(identity))) { + throw new WorktreeKiloRuntimeError('session_busy', 'Kilo runtime is retiring', true); + } if (!path.isAbsolute(directory) || path.resolve(directory) !== directory) { throw new WorktreeKiloRuntimeError('protocol_error', 'Invalid worktree directory', false); } @@ -670,8 +694,9 @@ export function createWorktreeKiloRuntimes(options: { throw new WorktreeKiloRuntimeError('not_ready', 'Kilo worktree is deleted', false); } let root = findRoot(identity); - const scopeDirectory = directoriesByScope.get(kilo.scopeId); - const previous = entries.get(directory); + const scopeDirectory = + isolation === 'directory-shared' ? directoriesByScope.get(kilo.scopeId) : undefined; + const previous = entries.get(key); if (previous?.retirementResult === 'unconfirmed') { throw new WorktreeKiloRuntimeError( 'not_ready', @@ -748,7 +773,7 @@ export function createWorktreeKiloRuntimes(options: { } if (!entry) { const homeId = createHash('sha256') - .update(kilo.scopeId) + .update(isolation === 'per-session' ? key : kilo.scopeId) .update('\0') .update(directory) .digest('hex'); @@ -757,6 +782,8 @@ export function createWorktreeKiloRuntimes(options: { homeId ); entry = { + identity: { ...identity }, + isolation, kilo: { ...kilo, targets: { ...kilo.targets } }, directory, env: buildWorktreeKiloEnvironment( @@ -774,10 +801,10 @@ export function createWorktreeKiloRuntimes(options: { const homes = homesByDirectory.get(directory) ?? new Set(); homes.add(home); homesByDirectory.set(directory, homes); - entries.set(directory, entry); + entries.set(key, entry); cleanupRequired = true; failedDirectories.delete(directory); - directoriesByScope.set(kilo.scopeId, directory); + if (isolation === 'directory-shared') directoriesByScope.set(kilo.scopeId, directory); } if (!root) { root = { @@ -787,7 +814,7 @@ export function createWorktreeKiloRuntimes(options: { attached: false, pending: new Set(), }; - roots.set(identity.kiloSessionId, root); + roots.set(identityKey(identity), root); entry.roots.add(root); rememberAttachedRoot(identity.kiloSessionId, directory); } @@ -829,13 +856,80 @@ export function createWorktreeKiloRuntimes(options: { removeRoot(root); return true; }, + async retireForRecovery(identity, recoveryId, assertIdle) { + const key = identityKey(identity); + const acknowledgements = recoveryAcknowledgements.get(key); + const acknowledged = acknowledgements?.get(recoveryId); + if (acknowledged) { + await acknowledged; + return 'acknowledged'; + } + if (recoveryGates.has(key)) { + throw new WorktreeKiloRuntimeError('session_busy', 'Kilo runtime is retiring', true); + } + const root = findRoot(identity); + if (!root) { + assertIdle(); + if (entries.has(entryKey(identity, 'per-session'))) { + throw new WorktreeKiloRuntimeError( + 'not_ready', + 'Session runtime is not recoverable', + false + ); + } + const absent = Promise.resolve('absent'); + const byRecovery = acknowledgements ?? new Map>(); + if (!acknowledgements) recoveryAcknowledgements.set(key, byRecovery); + byRecovery.set(recoveryId, absent); + return absent; + } + if (!root.attached) { + throw new WorktreeKiloRuntimeError('session_busy', 'Session runtime is attaching', true); + } + if (root.entry.isolation !== 'per-session') { + throw new WorktreeKiloRuntimeError( + 'not_ready', + 'Session runtime is not recoverable', + false + ); + } + const entry = root.entry; + if (!entry.runtime || !entry.kiloClient || !entry.stopped || entry.starting) { + throw new WorktreeKiloRuntimeError('not_ready', 'Kilo worktree is not ready', true); + } + let releaseGate: () => void = () => {}; + const gate = new Promise(resolve => { + releaseGate = resolve; + }); + recoveryGates.set(key, gate); + const retirement = (async (): Promise => { + try { + assertIdle(); + removeRoot(root); + await retire(entry); + return 'retired'; + } finally { + recoveryGates.delete(key); + releaseGate(); + } + })(); + const byRecovery = acknowledgements ?? new Map>(); + if (!acknowledgements) recoveryAcknowledgements.set(key, byRecovery); + byRecovery.set(recoveryId, retirement); + try { + return await retirement; + } catch (error) { + if (byRecovery.get(recoveryId) === retirement) byRecovery.delete(recoveryId); + throw error; + } + }, async deleteDirectory(directory) { deletedDirectories.add(directory); - const entry = entries.get(directory); for (const root of roots.values()) { if (root.identity.directory === directory) removeRoot(root); } - if (entry) { + for (const entry of [...entries.values()]) { + if (entry.directory !== directory) continue; const quiescent = await withTimeoutAndAbort(retire(entry), { timeoutMs: KILO_STARTUP_TIMEOUT_MS, timeoutMessage: 'Kilo worktree retirement timed out', @@ -849,12 +943,19 @@ export function createWorktreeKiloRuntimes(options: { homesByDirectory.delete(directory); }, async retireRuntime(directory, deadlineAt, target) { - const entry = entries.get(directory); + const entry = target + ? [...entries.values()].find( + entry => entry.directory === directory && entry.runtimeId === target.runtimeId + ) + : (entries.get(directory) ?? + [...entries.values()].find(entry => entry.directory === directory)); if (!entry) return 'stale'; return retire(entry, deadlineAt, target); }, async verifyQuiescence(directory, target, deadlineAt) { - const entry = entries.get(directory); + const entry = [...entries.values()].find( + entry => entry.directory === directory && entry.runtimeId === target.runtimeId + ); if ( !entry || entry.runtimeId !== target.runtimeId || @@ -866,26 +967,53 @@ export function createWorktreeKiloRuntimes(options: { return ( verified && Date.now() < deadlineAt && - entries.get(directory) === entry && + entries.get(entryKey(entry.identity, entry.isolation)) === entry && entry.runtimeId === target.runtimeId && entry.kiloClient === target.client && !entry.abort.signal.aborted ); }, - getRetained(directory) { - return entries.get(directory)?.runtime; + getRetained(directory, runtimeId) { + if (runtimeId === undefined) return entries.get(directory)?.runtime; + return [...entries.values()].find( + entry => entry.directory === directory && entry.runtimeId === runtimeId + )?.runtime; }, - get(directory) { - const entry = entries.get(directory); + get(identity) { + const entry = + typeof identity === 'string' + ? (entries.get(identity) ?? + [...entries.values()].find(entry => entry.directory === identity)) + : (entries.get(entryKey(identity, 'per-session')) ?? entries.get(identity.directory)); const runtime = entry?.starting ? undefined : entry?.runtime; return !closed && runtime && !runtime.signal.aborted ? runtime : undefined; }, prepareForNewWork(directory) { - const entry = entries.get(directory); + return [...entries.values()] + .filter(entry => entry.directory === directory) + .every( + entry => + !entry.abort.signal.aborted && + (entry.runtime === undefined || entry.feed?.prepareForNewWork() === true) + ); + }, + getAll(directory) { + return [...entries.values()] + .flatMap(entry => { + const runtime = entry.starting ? undefined : entry.runtime; + return !closed && runtime && !runtime.signal.aborted && entry.feed?.isFresh() + ? [runtime] + : []; + }) + .filter(runtime => runtime.directory === directory); + }, + isCurrent(runtime) { return ( - !entry || - (!entry.abort.signal.aborted && - (entry.runtime === undefined || entry.feed?.prepareForNewWork() === true)) + !closed && + runtime.identity !== undefined && + entries.get(entryKey(runtime.identity, runtime.isolation ?? 'directory-shared')) + ?.runtime === runtime && + !runtime.signal.aborted ); }, isHealthy() { @@ -902,7 +1030,6 @@ export function createWorktreeKiloRuntimes(options: { closed = true; for (const root of roots.values()) removeRoot(root); for (const entry of entries.values()) void retire(entry); - directoriesByScope.clear(); }, }; } diff --git a/services/cloud-agent-next/wrapper/src/session-bootstrap.test.ts b/services/cloud-agent-next/wrapper/src/session-bootstrap.test.ts index 11801d4119..6d3fb79abc 100644 --- a/services/cloud-agent-next/wrapper/src/session-bootstrap.test.ts +++ b/services/cloud-agent-next/wrapper/src/session-bootstrap.test.ts @@ -21,6 +21,7 @@ import type { } from '../../src/shared/wrapper-bootstrap'; import { buildCloudAgentRules } from '../../src/shared/cloud-agent-rules.js'; import { PNPM_STORE_DIR, PNPM_STORE_ENV_VAR } from '../../src/shared/runtime-environment.js'; +import { isWrapperSessionReadyRequest } from '../../src/shared/wrapper-bootstrap.js'; function makeRequest(tmpDir: string, overrides: Partial = {}) { const request: WrapperSessionReadyRequest = { @@ -3227,3 +3228,53 @@ describe('prepareWrapperBootstrapWorkspace', () => { expect(await fsp.readFile(localPath, 'utf8')).toBe('zip-payload'); }); }); + +describe('runtime credential proxy ready request validation', () => { + it('accepts only a stable handle and Worker proxy targets', () => { + const request = makeRequest(fs.mkdtempSync(path.join(os.tmpdir(), 'wrapper-proxy-request-'))); + request.runtimeCredentialProxy = { + handle: 'stable-proxy-handle', + targets: { + backendBaseUrl: 'https://worker.example.com', + providerBaseUrl: 'https://worker.example.com', + sessionIngestBaseUrl: 'https://worker.example.com', + }, + }; + expect(isWrapperSessionReadyRequest(request)).toBe(true); + expect( + isWrapperSessionReadyRequest({ + ...request, + runtimeCredentialProxy: { ...request.runtimeCredentialProxy, credential: 'backing-token' }, + }) + ).toBe(false); + for (const [requiredKey, wrongKey] of [ + ['backendBaseUrl', 'backendUrl'], + ['providerBaseUrl', 'providerUrl'], + ['sessionIngestBaseUrl', 'ingestUrl'], + ]) { + expect( + isWrapperSessionReadyRequest({ + ...request, + runtimeCredentialProxy: { + handle: 'stable-proxy-handle', + targets: Object.fromEntries( + Object.entries(request.runtimeCredentialProxy.targets).map(([key, value]) => [ + key === requiredKey ? wrongKey : key, + value, + ]) + ), + }, + }) + ).toBe(false); + } + expect( + isWrapperSessionReadyRequest({ + ...request, + runtimeCredentialProxy: { + handle: 'stable-proxy-handle', + targets: { backendBaseUrl: 'not-a-url' }, + }, + }) + ).toBe(false); + }); +}); From 1bca2576df24eb756f3ed3dae058ed3b9e58288c Mon Sep 17 00:00:00 2001 From: Florian Hines Date: Tue, 8 Sep 2026 12:09:23 -0500 Subject: [PATCH 3/9] fix(auth): include existing consumer type adapters --- services/gastown/src/middleware/auth.middleware.ts | 8 ++++++-- services/wasteland/src/middleware/auth.middleware.ts | 7 +++++-- 2 files changed, 11 insertions(+), 4 deletions(-) diff --git a/services/gastown/src/middleware/auth.middleware.ts b/services/gastown/src/middleware/auth.middleware.ts index fa8b694074..8d0e2accea 100644 --- a/services/gastown/src/middleware/auth.middleware.ts +++ b/services/gastown/src/middleware/auth.middleware.ts @@ -1,13 +1,15 @@ +import type { KiloAuthVariables } from '@kilocode/worker-utils/kilo-auth-middleware'; import type { Context } from 'hono'; import { createMiddleware } from 'hono/factory'; import { extractBearerToken } from '@kilocode/worker-utils'; +import type { KiloTokenPayload } from '@kilocode/worker-utils/kilo-token'; import { verifyAgentJWT, verifyContainerJWT, type AgentJWTPayload } from '../util/jwt.util'; import { resError } from '../util/res.util'; import type { GastownEnv } from '../gastown.worker'; -export type JwtOrgMembership = { orgId: string; role: 'owner' | 'member' | 'billing_manager' }; +export type JwtOrgMembership = NonNullable[number]; -export type AuthVariables = { +export type AuthVariables = KiloAuthVariables & { agentJWT: AgentJWTPayload; townId: string; kiloUserId: string; @@ -15,6 +17,8 @@ export type AuthVariables = { kiloApiTokenPepper: string | null; kiloGastownAccess: boolean; kiloOrgMemberships: JwtOrgMembership[]; + kiloControlToken: string; + kiloUsesModernToken: boolean; requestStartTime: number; orgId?: string; orgRole?: string; diff --git a/services/wasteland/src/middleware/auth.middleware.ts b/services/wasteland/src/middleware/auth.middleware.ts index 330900066b..5b27127a71 100644 --- a/services/wasteland/src/middleware/auth.middleware.ts +++ b/services/wasteland/src/middleware/auth.middleware.ts @@ -1,6 +1,9 @@ -export type JwtOrgMembership = { orgId: string; role: 'owner' | 'member' | 'billing_manager' }; +import type { KiloAuthVariables } from '@kilocode/worker-utils/kilo-auth-middleware'; +import type { KiloTokenPayload } from '@kilocode/worker-utils/kilo-token'; -export type AuthVariables = { +export type JwtOrgMembership = NonNullable[number]; + +export type AuthVariables = KiloAuthVariables & { kiloUserId: string; kiloIsAdmin: boolean; kiloApiTokenPepper: string | null; From c2f50a54b7ce5f05424f202cc2b2a5feaaa80a29 Mon Sep 17 00:00:00 2001 From: Florian Hines Date: Tue, 8 Sep 2026 12:23:21 -0500 Subject: [PATCH 4/9] feat(cloud-agent): enable runtime isolation readiness --- ENVIRONMENT.md | 2 +- services/cloud-agent-next/.dev.vars.example | 2 +- services/cloud-agent-next/worker-configuration.d.ts | 6 +++--- services/cloud-agent-next/wrangler.jsonc | 4 ++-- 4 files changed, 7 insertions(+), 7 deletions(-) diff --git a/ENVIRONMENT.md b/ENVIRONMENT.md index c7b7aa3e73..601f7829d7 100644 --- a/ENVIRONMENT.md +++ b/ENVIRONMENT.md @@ -90,7 +90,7 @@ Manage shared web env var additions and rotations with `pnpm web:env set Date: Tue, 8 Sep 2026 16:59:16 -0500 Subject: [PATCH 5/9] test(cloud-agent): replace unused CLI route emulator --- .../__fixtures__/runtime-url-normalization.ts | 27 --------- .../src/kilo/kilo-targets.test.ts | 56 ++++++++++++++----- 2 files changed, 41 insertions(+), 42 deletions(-) delete mode 100644 services/cloud-agent-next/src/kilo/__fixtures__/runtime-url-normalization.ts diff --git a/services/cloud-agent-next/src/kilo/__fixtures__/runtime-url-normalization.ts b/services/cloud-agent-next/src/kilo/__fixtures__/runtime-url-normalization.ts deleted file mode 100644 index 07aa1fc959..0000000000 --- a/services/cloud-agent-next/src/kilo/__fixtures__/runtime-url-normalization.ts +++ /dev/null @@ -1,27 +0,0 @@ -/** - * Immutable compatibility fixtures copied from Kilo CLI `route()` behavior. - * Sources: Kilo v7.4.20 release artifact and current CLI main (2026-09-04). - * - * Both versions normalize a configured base then append absolute API paths. - * Provider-specific `/api/gateway/v1/...` compatibility aliases are tested by - * the Cloud facade resolver, not asserted as behavior of `route()` itself. - * Keep this independent of the mutable CLI checkout: older installed CLIs are - * part of the Cloud Agent compatibility contract. - */ -export const historicalRouteFixtures = [ - { version: '7.4.20', path: '/api/profile' }, - { version: '7.4.20', path: '/api/defaults' }, - { version: '7.4.20', path: '/api/openrouter/models' }, - { version: '7.4.20', path: '/api/session' }, - { version: 'current', path: '/api/profile' }, - { version: 'current', path: '/api/defaults' }, - { version: 'current', path: '/api/openrouter/models' }, - { version: 'current', path: '/api/session' }, -] as const; - -/** Exact relevant `route()` normalization contract from both sources. */ -export function historicalCliRoute(base: string, path: string): string { - const url = new URL(base); - url.pathname = `${url.pathname.replace(/\/+$/, '')}/${path.replace(/^\/+/, '')}`; - return url.toString(); -} diff --git a/services/cloud-agent-next/src/kilo/kilo-targets.test.ts b/services/cloud-agent-next/src/kilo/kilo-targets.test.ts index e22fcc235e..4cdc00643f 100644 --- a/services/cloud-agent-next/src/kilo/kilo-targets.test.ts +++ b/services/cloud-agent-next/src/kilo/kilo-targets.test.ts @@ -5,23 +5,49 @@ import { providerBaseUrlEncodedInToken, } from './kilo-targets.js'; import { - historicalCliRoute, - historicalRouteFixtures, -} from './__fixtures__/runtime-url-normalization.js'; + inferRuntimeCredentialProxyRoute, + resolveRuntimeCredentialProxyRoute, +} from './runtime-credential-proxy-routes.js'; -describe('immutable historical Kilo route normalization', () => { +// Literal requests emitted by supported CLI versions exercise the production +// facade resolver, without emulating CLI URL construction in this test. +describe('CLI route compatibility', () => { it.each([ - ['origin facade', 'https://worker.example.test'], - ['safe prefixed facade', 'https://worker.example.test/runtime-proxy'], - ])('%s keeps v7.4.20 and current requests inside the Worker facade', (_name, facade) => { - for (const fixture of historicalRouteFixtures) { - const routed = new URL(historicalCliRoute(facade, fixture.path)); - expect(routed.origin).toBe('https://worker.example.test'); - expect(routed.pathname).toBe( - `${new URL(facade).pathname.replace(/\/+$/, '')}${fixture.path}` - ); - if (new URL(facade).pathname !== '/') expect(routed.pathname).not.toBe(fixture.path); - } + ['https://worker.example.test/api/profile', 'GET', 'https://api.kilo.ai/api/profile'], + ['https://worker.example.test/api/defaults', 'GET', 'https://api.kilo.ai/api/defaults'], + [ + 'https://worker.example.test/api/openrouter/models', + 'GET', + 'https://api.kilo.ai/api/gateway/models', + ], + [ + 'https://worker.example.test/api/gateway/v1/chat/completions', + 'POST', + 'https://api.kilo.ai/api/gateway/v1/chat/completions', + ], + [ + 'https://worker.example.test/api/session', + 'POST', + 'https://ingest.kilosessions.ai/api/session', + ], + ])('routes CLI request %s (%s) to %s', (cliUrl, method, expectedUrl) => { + const derived = deriveKiloSandboxTargets({}, 'user-token'); + if (!derived.success) throw new Error('Expected default targets'); + const url = new URL(cliUrl); + const route = inferRuntimeCredentialProxyRoute(url.pathname); + if (!route) throw new Error('Expected a supported CLI route'); + expect( + resolveRuntimeCredentialProxyRoute({ + targets: derived.targets, + route, + method, + pathname: url.pathname, + search: url.search, + kiloSessionId: 'ses_cli_compatibility', + contentType: 'application/json', + bodyText: '{"sessionId":"ses_cli_compatibility"}', + })?.href + ).toBe(expectedUrl); }); }); From 8599d647b6117890bdb6a17f128464a4ef819779 Mon Sep 17 00:00:00 2001 From: Florian Hines Date: Tue, 8 Sep 2026 16:59:17 -0500 Subject: [PATCH 6/9] fix(cloud-agent): diagnose stalled authorization recovery --- .../src/persistence/CloudAgentSession.ts | 111 ++++++++++++++--- .../runtime-authorization-persistence.test.ts | 68 +++++++++++ .../runtime-authorization-persistence.ts | 39 ++++++ .../runtime-authorization-recovery.test.ts | 115 +++++++++++++++++- 4 files changed, 316 insertions(+), 17 deletions(-) diff --git a/services/cloud-agent-next/src/persistence/CloudAgentSession.ts b/services/cloud-agent-next/src/persistence/CloudAgentSession.ts index 9369a0c180..e548ebdefb 100644 --- a/services/cloud-agent-next/src/persistence/CloudAgentSession.ts +++ b/services/cloud-agent-next/src/persistence/CloudAgentSession.ts @@ -113,6 +113,8 @@ import { RUNTIME_AUTHORIZATION_RECOVERY_KEY, RUNTIME_AUTHORIZATION_KEY, runtimeAuthorizationRecoveryLockSchema, + inspectRuntimeAuthorizationRecoveryLock, + RUNTIME_AUTHORIZATION_RECOVERY_DIAGNOSTICS_KEY, } from '../session/runtime-authorization-persistence.js'; import { RUNTIME_PROXY_GRANT_KEY } from '../runtime-credential-proxy.js'; import { getEffectiveCredentialContainment } from './session-metadata.js'; @@ -1726,16 +1728,48 @@ export class CloudAgentSession extends DurableObject { metadata: await this.getMetadata(), getAuthorization: () => this.ctx.storage.get(RUNTIME_AUTHORIZATION_KEY), }); - const lock = runtimeAuthorizationRecoveryLockSchema.safeParse( - await this.ctx.storage.get(RUNTIME_AUTHORIZATION_RECOVERY_KEY) - ); + const lock = this.inspectRuntimeAuthorizationRecovery(); return state.state === 'expired' && lock.success && lock.data.expectedOldId === state.id ? { ...state, recoveryId: lock.data.recoveryId } : state; } + private inspectRuntimeAuthorizationRecovery() { + // Synchronous read/update prevents concurrent inspections from duplicating warnings + // or overwriting a replacement lock after an await. + const lock = runtimeAuthorizationRecoveryLockSchema.safeParse( + this.ctx.storage.kv.get(RUNTIME_AUTHORIZATION_RECOVERY_KEY) + ); + if (lock.success) { + const inspection = inspectRuntimeAuthorizationRecoveryLock( + lock.data, + this.ctx.storage.kv.get(RUNTIME_AUTHORIZATION_RECOVERY_DIAGNOSTICS_KEY), + Date.now() + ); + if (inspection.changed) { + this.ctx.storage.kv.put( + RUNTIME_AUTHORIZATION_RECOVERY_DIAGNOSTICS_KEY, + inspection.diagnostics + ); + } + if (inspection.warn) { + logger + .withFields({ + sessionId: this.sessionId, + expectedOldId: lock.data.expectedOldId, + recoveryId: lock.data.recoveryId, + reason: 'prolonged_recovery_lock', + lockAgeMs: Date.now() - inspection.diagnostics.startedAt, + }) + .warn('Runtime authorization recovery requires attention'); + } + } + return lock; + } + async isRuntimeAuthorizationRecoveryInProgress(): Promise { - return (await this.ctx.storage.get(RUNTIME_AUTHORIZATION_RECOVERY_KEY)) !== undefined; + this.inspectRuntimeAuthorizationRecovery(); + return this.ctx.storage.kv.get(RUNTIME_AUTHORIZATION_RECOVERY_KEY) !== undefined; } async recoverExpiredRuntimeAuthorization(input: { @@ -1748,7 +1782,12 @@ export class CloudAgentSession extends DurableObject { const metadata = await this.getMetadata(); if (!metadata || metadata.identity.userId !== input.ownerId) return { status: 'denied' }; const secret = await resolveSecret(this.env.NEXTAUTH_SECRET); - if (!secret) return { status: 'denied' }; + if (!secret) { + logger + .withFields({ sessionId: metadata.identity.sessionId, reason: 'missing_secret' }) + .error('Runtime authorization recovery denied'); + return { status: 'denied' }; + } let fresh: RuntimeAuthorization; try { fresh = await unsealRuntimeAuthorization(input.runtimeAuthorizationSeal, secret, { @@ -1792,15 +1831,41 @@ export class CloudAgentSession extends DurableObject { ) { return false; } - await transaction.put(RUNTIME_AUTHORIZATION_RECOVERY_KEY, { - expectedOldId: input.expectedOldId, - recoveryId: input.recoveryId, - }); + if (!existingLock.success) { + await transaction.put(RUNTIME_AUTHORIZATION_RECOVERY_KEY, { + expectedOldId: input.expectedOldId, + recoveryId: input.recoveryId, + }); + await transaction.put(RUNTIME_AUTHORIZATION_RECOVERY_DIAGNOSTICS_KEY, { + expectedOldId: input.expectedOldId, + recoveryId: input.recoveryId, + startedAt: Date.now(), + }); + } return true; }); if (!acquired) { return { status: 'retry' }; } + type RecoveryFailureReason = + | 'physical_inspection_failed' + | 'terminal_inspection_failed' + | 'physical_retirement_failed' + | 'physical_absence_unconfirmed' + | 'authorization_state_changed' + | 'authorization_commit_failed' + | 'wrapper_identity_clear_failed'; + const diagnostic = (reason: RecoveryFailureReason) => { + logger + .withFields({ + sessionId: metadata.identity.sessionId, + expectedOldId: input.expectedOldId, + recoveryId: input.recoveryId, + reason, + }) + .warn('Runtime authorization recovery incomplete'); + }; + let failureReason: RecoveryFailureReason = 'physical_inspection_failed'; try { const observation = this.physicalWrapperObserver ? await this.physicalWrapperObserver() @@ -1811,11 +1876,19 @@ export class CloudAgentSession extends DurableObject { metadata, this.getAgentSandboxRuntimeContext() ).observeWrappersWithoutWaking(); - if (observation.status === 'inspection-failed') return { status: 'retry' }; + if (observation.status === 'inspection-failed') { + diagnostic('physical_inspection_failed'); + return { status: 'retry' }; + } if (observation.status === 'present') { + failureReason = 'terminal_inspection_failed'; const terminal = await this.getTerminalClient(); - if (!terminal.success || !terminal.data) return { status: 'retry' }; + if (!terminal.success || !terminal.data) { + diagnostic('terminal_inspection_failed'); + return { status: 'retry' }; + } if ((await terminal.data.client.listTerminals()).length > 0) return { status: 'busy' }; + failureReason = 'physical_retirement_failed'; const supervisor = this.getWrapperSupervisor(); await supervisor.requestPhysicalWrapperStop('idle-timeout', { kind: 'session' }); await supervisor.runMaintenance(Date.now()); @@ -1828,14 +1901,20 @@ export class CloudAgentSession extends DurableObject { metadata, this.getAgentSandboxRuntimeContext() ).observeWrappersWithoutWaking(); - if (stopped.status !== 'absent') return { status: 'retry' }; + if (stopped.status !== 'absent') { + diagnostic('physical_absence_unconfirmed'); + return { status: 'retry' }; + } } + failureReason = 'authorization_state_changed'; const latest = await this.getRuntimeAuthorizationRecoveryState(); if (latest.state !== 'expired' || latest.id !== input.expectedOldId) { + diagnostic('authorization_state_changed'); return latest.state === 'active' || latest.state === 'legacy' ? { status: 'not-needed' } : { status: 'denied' }; } + failureReason = 'authorization_commit_failed'; await this.ctx.storage.transaction(async transaction => { const current = RuntimeAuthorizationSchema.safeParse( await transaction.get(RUNTIME_AUTHORIZATION_KEY) @@ -1868,13 +1947,13 @@ export class CloudAgentSession extends DurableObject { ); await transaction.delete(RUNTIME_PROXY_GRANT_KEY); await transaction.delete(RUNTIME_AUTHORIZATION_RECOVERY_KEY); + await transaction.delete(RUNTIME_AUTHORIZATION_RECOVERY_DIAGNOSTICS_KEY); }); + failureReason = 'wrapper_identity_clear_failed'; await clearWrapperRuntimeIdentity(this.ctx.storage, {}, { incrementGeneration: true }); return { status: 'recovered' }; - } catch (error) { - if (error instanceof Error && error.message === 'runtime_authorization_recovery_cas_failed') { - return { status: 'retry' }; - } + } catch { + diagnostic(failureReason); return { status: 'retry' }; } } diff --git a/services/cloud-agent-next/src/session/runtime-authorization-persistence.test.ts b/services/cloud-agent-next/src/session/runtime-authorization-persistence.test.ts index ed55649f34..b8bcc317cf 100644 --- a/services/cloud-agent-next/src/session/runtime-authorization-persistence.test.ts +++ b/services/cloud-agent-next/src/session/runtime-authorization-persistence.test.ts @@ -7,6 +7,9 @@ import { } from '@kilocode/worker-utils/runtime-authorization'; import { getRuntimeAuthorizationStatus, + inspectRuntimeAuthorizationRecoveryLock, + runtimeAuthorizationRecoveryLockSchema, + RUNTIME_AUTHORIZATION_RECOVERY_WARNING_MS, getRuntimeAuthorizationRecoveryState, renewStoredRuntimeAuthorization, } from './runtime-authorization-persistence.js'; @@ -264,3 +267,68 @@ describe('runtime authorization persistence', () => { ).rejects.toBeInstanceOf(RuntimeAuthorizationRevokedError); }); }); + +describe('recovery lock diagnostics', () => { + const lock = { + expectedOldId: '00000000-0000-4000-8000-000000000001', + recoveryId: '00000000-0000-4000-8000-000000000002', + }; + const start = 1_000; + const threshold = RUNTIME_AUTHORIZATION_RECOVERY_WARNING_MS; + + it('starts aging legacy locks without modifying their strict two-field contract', () => { + const original = { ...lock }; + const first = inspectRuntimeAuthorizationRecoveryLock(lock, undefined, start); + expect(first).toEqual({ + diagnostics: { ...lock, startedAt: start }, + changed: true, + warn: false, + }); + expect(lock).toEqual(original); + expect(runtimeAuthorizationRecoveryLockSchema.parse(lock)).toEqual(original); + }); + + it('preserves the durable start and warning cadence across repeated inspections', () => { + const first = inspectRuntimeAuthorizationRecoveryLock(lock, undefined, start); + const early = inspectRuntimeAuthorizationRecoveryLock( + lock, + first.diagnostics, + start + threshold - 1 + ); + expect(early.changed).toBe(false); + expect(early.warn).toBe(false); + const warning = inspectRuntimeAuthorizationRecoveryLock( + lock, + early.diagnostics, + start + threshold + ); + expect(warning.warn).toBe(true); + // Round-trip storage as on a new DO instance; warnings are not an in-memory timer. + const stored = JSON.parse(JSON.stringify(warning.diagnostics)); + expect( + inspectRuntimeAuthorizationRecoveryLock(lock, stored, start + 2 * threshold - 1).warn + ).toBe(false); + const next = inspectRuntimeAuthorizationRecoveryLock(lock, stored, start + 2 * threshold); + expect(next.warn).toBe(true); + expect(next.diagnostics.startedAt).toBe(start); + expect(lock).toEqual({ expectedOldId: lock.expectedOldId, recoveryId: lock.recoveryId }); + }); + + it.each([ + undefined, + { ...lock, startedAt: 'invalid' }, + { ...lock, recoveryId: '00000000-0000-4000-8000-000000000003', startedAt: 0 }, + { ...lock, expectedOldId: '00000000-0000-4000-8000-000000000003', startedAt: 0 }, + ])( + 'does not inherit age from missing, malformed or differently bound diagnostics: %j', + diagnostics => { + expect(inspectRuntimeAuthorizationRecoveryLock(lock, diagnostics, start + threshold)).toEqual( + { + diagnostics: { ...lock, startedAt: start + threshold }, + changed: true, + warn: false, + } + ); + } + ); +}); diff --git a/services/cloud-agent-next/src/session/runtime-authorization-persistence.ts b/services/cloud-agent-next/src/session/runtime-authorization-persistence.ts index 39e6ef3ddc..e43fbad6ad 100644 --- a/services/cloud-agent-next/src/session/runtime-authorization-persistence.ts +++ b/services/cloud-agent-next/src/session/runtime-authorization-persistence.ts @@ -15,6 +15,45 @@ export const RUNTIME_AUTHORIZATION_RECOVERY_KEY = 'runtime_authorization_recover export const runtimeAuthorizationRecoveryLockSchema = z .object({ expectedOldId: z.string().uuid(), recoveryId: z.string().uuid() }) .strict(); +export const RUNTIME_AUTHORIZATION_RECOVERY_DIAGNOSTICS_KEY = + 'runtime_authorization_recovery_diagnostics'; +const recoveryDiagnosticsSchema = runtimeAuthorizationRecoveryLockSchema.extend({ + startedAt: z.number().int().nonnegative(), + lastWarningAt: z.number().int().nonnegative().optional(), +}); +// Observation-driven diagnostics only; this is never a lock expiry deadline. +export const RUNTIME_AUTHORIZATION_RECOVERY_WARNING_MS = 5 * 60_000; + +export function inspectRuntimeAuthorizationRecoveryLock( + lock: z.infer, + diagnostics: unknown, + now: number +) { + const parsed = recoveryDiagnosticsSchema.safeParse(diagnostics); + const current = + parsed.success && + parsed.data.expectedOldId === lock.expectedOldId && + parsed.data.recoveryId === lock.recoveryId + ? parsed.data + : undefined; + // For legacy locks the true start is unknown; persist the first observation. + const startedAt = current?.startedAt ?? now; + const warn = + now - startedAt >= RUNTIME_AUTHORIZATION_RECOVERY_WARNING_MS && + (current?.lastWarningAt === undefined || + now - current.lastWarningAt >= RUNTIME_AUTHORIZATION_RECOVERY_WARNING_MS); + return { + diagnostics: { + ...lock, + startedAt, + ...(current?.lastWarningAt !== undefined ? { lastWarningAt: current.lastWarningAt } : {}), + ...(warn ? { lastWarningAt: now } : {}), + }, + changed: !current || warn, + warn, + }; +} + const RUNTIME_TOKEN_RENEWAL_WINDOW_MS = 5 * 60_000; function runtimeAuthorizationId(value: unknown): string | null { diff --git a/services/cloud-agent-next/test/integration/runtime-authorization-recovery.test.ts b/services/cloud-agent-next/test/integration/runtime-authorization-recovery.test.ts index c39bb1e83d..f9c0e56f79 100644 --- a/services/cloud-agent-next/test/integration/runtime-authorization-recovery.test.ts +++ b/services/cloud-agent-next/test/integration/runtime-authorization-recovery.test.ts @@ -1,5 +1,5 @@ import { env, listDurableObjectIds, runInDurableObject } from 'cloudflare:test'; -import { beforeEach, describe, expect, it } from 'vitest'; +import { beforeEach, describe, expect, it, vi } from 'vitest'; import jwt from 'jsonwebtoken'; import { sealRuntimeAuthorization } from '@kilocode/worker-utils/runtime-authorization'; import type { RuntimeAuthorization } from '@kilocode/worker-utils/runtime-authorization-contract'; @@ -7,11 +7,14 @@ import { RUNTIME_PROXY_GRANT_KEY } from '../../src/runtime-credential-proxy.js'; import { RUNTIME_AUTHORIZATION_KEY, RUNTIME_AUTHORIZATION_RECOVERY_KEY, + RUNTIME_AUTHORIZATION_RECOVERY_DIAGNOSTICS_KEY, + RUNTIME_AUTHORIZATION_RECOVERY_WARNING_MS, } from '../../src/session/runtime-authorization-persistence.js'; import { allocateWrapperRuntimeState, getWrapperRuntimeState, } from '../../src/session/wrapper-runtime-state.js'; +import { logger } from '../../src/logger.js'; import { registerReadySession } from '../helpers/session-setup.js'; const organizationId = '11111111-1111-4111-8111-111111111111'; @@ -113,6 +116,81 @@ describe('runtime authorization recovery', () => { expect(result.prepared).toEqual(result.ordinary); }); + it('logs missing configuration and thrown recovery failures without credential or error data', async () => { + const userId = 'user_recovery_diagnostics'; + const sessionId = 'agent_recovery_diagnostics'; + const stub = env.CLOUD_AGENT_SESSION.getByName(`${userId}:${sessionId}`); + await runInDurableObject(stub, async instance => { + await registerReadySession(instance, { + sessionId, + userId, + orgId: organizationId, + prompt: 'initial', + mode: 'code', + model: 'test-model', + }); + const old = authorization({ + id: '00000000-0000-4000-8000-000000000401', + sessionId, + userId, + expiresAt: new Date(Date.now() - 60_000).toISOString(), + issuedAt: new Date(Date.now() - 120_000).toISOString(), + }); + const fresh = authorization({ + id: '00000000-0000-4000-8000-000000000402', + sessionId, + userId, + expiresAt: new Date(Date.now() + 60_000).toISOString(), + }); + await instance.ctx.storage.put(RUNTIME_AUTHORIZATION_KEY, old); + const input = { + ownerId: userId, + expectedOldId: old.id, + recoveryId: '00000000-0000-4000-8000-000000000403', + runtimeAuthorizationSeal: await seal(fresh), + runtimeToken: 'private-bearer', + }; + const originalSecret = instance['env'].NEXTAUTH_SECRET; + const fields = vi.spyOn(logger, 'withFields').mockReturnValue(logger); + const error = vi.spyOn(logger, 'error').mockImplementation(() => {}); + const warn = vi.spyOn(logger, 'warn').mockImplementation(() => {}); + try { + instance['env'].NEXTAUTH_SECRET = ''; + expect(await instance.recoverExpiredRuntimeAuthorization(input)).toEqual({ + status: 'denied', + }); + expect(fields).toHaveBeenLastCalledWith({ sessionId, reason: 'missing_secret' }); + expect(error).toHaveBeenCalledWith('Runtime authorization recovery denied'); + expect(await instance.ctx.storage.get(RUNTIME_AUTHORIZATION_RECOVERY_KEY)).toBeUndefined(); + instance['env'].NEXTAUTH_SECRET = originalSecret; + instance['physicalWrapperObserver'] = async () => { + throw new Error('private-error-with-bearer'); + }; + expect(await instance.recoverExpiredRuntimeAuthorization(input)).toEqual({ + status: 'retry', + }); + expect(fields).toHaveBeenLastCalledWith({ + sessionId, + expectedOldId: old.id, + recoveryId: input.recoveryId, + reason: 'physical_inspection_failed', + }); + expect(warn).toHaveBeenCalledWith('Runtime authorization recovery incomplete'); + expect(await instance.ctx.storage.get(RUNTIME_AUTHORIZATION_RECOVERY_KEY)).toEqual({ + expectedOldId: old.id, + recoveryId: input.recoveryId, + }); + expect(await instance.ctx.storage.get(RUNTIME_AUTHORIZATION_KEY)).toEqual(old); + expect( + JSON.stringify([fields.mock.calls, error.mock.calls, warn.mock.calls]) + ).not.toContain('private-'); + } finally { + instance['env'].NEXTAUTH_SECRET = originalSecret; + vi.restoreAllMocks(); + } + }); + }); + it('recovers an expired CloudAgentSession authorization only after confirmed idle retirement', async () => { const userId = 'user_cloud_recovery'; const sessionId = 'agent_cloud_recovery'; @@ -314,6 +392,21 @@ describe('runtime authorization recovery', () => { stops += 1; return { status: 'absent' }; }; + const lock = { + expectedOldId: old.id, + recoveryId: '00000000-0000-4000-8000-000000000253', + }; + // Legacy lock first observation must keep the exact old-reader contract. + await instance.ctx.storage.put(RUNTIME_AUTHORIZATION_RECOVERY_KEY, lock); + await instance.getRuntimeAuthorizationRecoveryState(); + expect(await instance.ctx.storage.get(RUNTIME_AUTHORIZATION_RECOVERY_KEY)).toEqual(lock); + const startedAt = Date.now() - RUNTIME_AUTHORIZATION_RECOVERY_WARNING_MS - 1; + await instance.ctx.storage.put(RUNTIME_AUTHORIZATION_RECOVERY_DIAGNOSTICS_KEY, { + ...lock, + startedAt, + }); + const fields = vi.spyOn(logger, 'withFields').mockReturnValue(logger); + const warning = vi.spyOn(logger, 'warn').mockImplementation(() => {}); const recovery = await instance.recoverExpiredRuntimeAuthorization({ ownerId: userId, expectedOldId: old.id, @@ -321,6 +414,26 @@ describe('runtime authorization recovery', () => { runtimeAuthorizationSeal: await seal(fresh), runtimeToken: 'fresh-token', }); + await instance.getRuntimeAuthorizationRecoveryState(); + await instance.isRuntimeAuthorizationRecoveryInProgress(); + expect( + warning.mock.calls.filter( + ([message]) => message === 'Runtime authorization recovery requires attention' + ) + ).toHaveLength(1); + expect(fields).toHaveBeenCalledWith({ + sessionId, + expectedOldId: old.id, + recoveryId: lock.recoveryId, + reason: 'prolonged_recovery_lock', + lockAgeMs: expect.any(Number), + }); + warning.mockRestore(); + fields.mockRestore(); + expect(await instance.ctx.storage.get(RUNTIME_AUTHORIZATION_RECOVERY_KEY)).toEqual(lock); + expect( + await instance.ctx.storage.get(RUNTIME_AUTHORIZATION_RECOVERY_DIAGNOSTICS_KEY) + ).toMatchObject({ ...lock, startedAt, lastWarningAt: expect.any(Number) }); return { recovery, stops, From 69ba94d691e5a7bcabc1817900604ae6850d26f3 Mon Sep 17 00:00:00 2001 From: Florian Hines Date: Tue, 8 Sep 2026 16:59:17 -0500 Subject: [PATCH 7/9] fix(cloud-agent): restore runtime authorization on create retries --- .../src/session/session-prepare.test.ts | 165 ++++++++++++++++++ .../src/session/session-registration.ts | 74 +++++--- 2 files changed, 218 insertions(+), 21 deletions(-) diff --git a/services/cloud-agent-next/src/session/session-prepare.test.ts b/services/cloud-agent-next/src/session/session-prepare.test.ts index 349268d961..fe1f7723fa 100644 --- a/services/cloud-agent-next/src/session/session-prepare.test.ts +++ b/services/cloud-agent-next/src/session/session-prepare.test.ts @@ -8,6 +8,12 @@ * deterministically; `startNewSession` and the reconcile ladder run real code. */ import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; +import jwt from 'jsonwebtoken'; +import { + createRuntimeAuthorization, + sealRuntimeAuthorization, +} from '@kilocode/worker-utils/runtime-authorization'; +import { assertKiloModelAvailable } from '../model-validation.js'; import type { WorkerDb } from '@kilocode/db/client'; import type { OperationLedgerRow } from '@kilocode/db/schema'; @@ -64,6 +70,38 @@ const { generateSandboxRoutingTargetMock: vi.fn(), })); +vi.mock('@kilocode/worker-utils/runtime-authorization', () => ({ + createRuntimeAuthorization: vi.fn(), + sealRuntimeAuthorization: vi.fn(), +})); +vi.mock('../model-validation.js', () => ({ assertKiloModelAvailable: vi.fn() })); + +function modernContext( + ctx: SessionRegistrationContext, + organizationId?: string +): SessionRegistrationContext { + ctx.env.NEXTAUTH_SECRET = 'registration-test-secret'; + ctx.env.RUNTIME_ISOLATION_ENABLED = 'true'; + ctx.authToken = jwt.sign( + { + kiloUserId: ctx.userId, + ...(organizationId ? { organizationId } : {}), + aud: 'cloud-agent-next', + tokenPurpose: 'human-api', + credentialExchange: false, + }, + 'registration-test-secret' + ); + vi.mocked(createRuntimeAuthorization).mockResolvedValue({ + token: 'runtime-token', + authorization: { id: 'authorization-test' }, + expiresAt: '2099-01-01T00:00:00Z', + } as Awaited>); + vi.mocked(sealRuntimeAuthorization).mockResolvedValue('runtime-seal'); + vi.mocked(assertKiloModelAvailable).mockResolvedValue(undefined); + return ctx; +} + vi.mock('@kilocode/db/operation-ledger', () => ({ admitOperation: admitOperationMock, settleOperation: settleOperationMock, @@ -2149,6 +2187,95 @@ describe('createSessionWithLedger worktree rollout and ownership reconciliation' }; } + it.each([false, true])( + 'restores modern authorization after a lost ownership response (registered=%s)', + async registered => { + const input = request(); + const stub = makeDoStub({ + getMetadata: vi + .fn() + .mockResolvedValue(registered ? { identity: { sessionId: WORKSPACE_SESSION_ID } } : null), + }); + const ctx = modernContext(context(stub)); + createCliSessionMock.mockRejectedValueOnce(new Error('ownership response lost after commit')); + await expect(runCreate(ctx, input)).rejects.toThrow('ownership response lost after commit'); + const progress = Object.assign( + {}, + ...recordOperationProgressMock.mock.calls.map(call => call[2]) + ); + admitOperationMock.mockResolvedValueOnce({ + admission: 'duplicate_reconcile_pending', + row: makeLedgerRow({ status: 'reconcile_pending', canonical_result: progress }), + }); + getPgDbMock.mockReturnValue(makeDb([[ownershipRow()], [{ email: 'test@example.com' }]])); + vi.mocked(createRuntimeAuthorization).mockClear(); + vi.mocked(assertKiloModelAvailable).mockClear(); + await expect(runCreate(ctx, input)).resolves.toMatchObject({ + cloudAgentSessionId: WORKSPACE_SESSION_ID, + kiloSessionId: KILO_SESSION_ID, + replayed: true, + }); + if (registered) { + expect(createRuntimeAuthorization).not.toHaveBeenCalled(); + expect(stub.createSessionWithInitialAdmission).not.toHaveBeenCalled(); + } else { + expect(createRuntimeAuthorization).toHaveBeenCalledWith( + expect.objectContaining({ + token: ctx.authToken, + resourceId: WORKSPACE_SESSION_ID, + resourceKind: 'cloud-agent-next', + }) + ); + expect(stub.createSessionWithInitialAdmission).toHaveBeenCalledWith( + expect.objectContaining({ + auth: expect.objectContaining({ + kilocodeToken: 'runtime-token', + kiloSessionId: KILO_SESSION_ID, + }), + runtimeAuthorizationSeal: 'runtime-seal', + message: { initialTurn: expect.objectContaining({ messageId: INITIAL_MESSAGE_ID }) }, + }) + ); + expect(assertKiloModelAvailable).toHaveBeenCalledWith( + expect.objectContaining({ originalToken: 'runtime-token' }) + ); + } + expect(createCliSessionMock).toHaveBeenCalledTimes(1); + expect(generateSessionIdMock).toHaveBeenCalledTimes(1); + expect(deleteCliSessionMock).not.toHaveBeenCalled(); + } + ); + + it.each(['admission', 'isolation', 'model'] as const)( + 'fails closed on resumed modern %s rejection', + async failure => { + const input = request(); + const stub = makeDoStub({ getMetadata: vi.fn().mockResolvedValue(null) }); + const ctx = modernContext(context(stub)); + const progress = await canonicalProgress(input); + admitOperationMock.mockResolvedValueOnce({ + admission: 'duplicate_reconcile_pending', + row: makeLedgerRow({ status: 'reconcile_pending', canonical_result: progress }), + }); + getPgDbMock.mockReturnValue(makeDb([[ownershipRow()], [{ email: 'test@example.com' }]])); + if (failure === 'admission') { + vi.mocked(createRuntimeAuthorization).mockRejectedValueOnce( + new Error('Invalid runtime admission') + ); + } else if (failure === 'isolation') { + ctx.env.RUNTIME_ISOLATION_ENABLED = 'false'; + } else { + vi.mocked(assertKiloModelAvailable).mockRejectedValueOnce(new Error('Model unavailable')); + } + await expect(runCreate(ctx, input)).rejects.toThrow(); + expect(stub.createSessionWithInitialAdmission).not.toHaveBeenCalled(); + expect(stub.registerSession).not.toHaveBeenCalled(); + expect(createCliSessionMock).not.toHaveBeenCalled(); + expect(deleteCliSessionMock).not.toHaveBeenCalled(); + expect(settleOperationMock).not.toHaveBeenCalled(); + } + ); + it.each([true, false, undefined])( 'recovers committed ownership with autoCommit=%s after rollout changes without changing the initial turn', async autoCommit => { @@ -3606,6 +3733,44 @@ describe('createSessionWithLedger clone reconciliation', () => { }); } + it.each([undefined, 'org-registration-test'])( + 'restores modern authorization for the canonical resumed clone (organization=%s)', + async organizationId => { + const input = cloneRequest({ + options: { operationKey: OPERATION_KEY, kilocodeOrganizationId: organizationId }, + }); + admitOperationMock.mockResolvedValueOnce({ + admission: 'takeover', + row: { ...(await cloneRow({}, input)), organization_id: organizationId ?? null }, + }); + createCliSessionMock.mockResolvedValueOnce({ + status: 'ready', + clone: { sessionId: KILO_SESSION_ID, copiedItemCount: 3 }, + }); + getPgDbMock.mockReturnValue(makeDb([[], [{ email: 'test@example.com' }]])); + const stub = makeDoStub(); + const ctx = modernContext(makeContext(stub), organizationId); + await runCreate(ctx, input); + expect(createRuntimeAuthorization).toHaveBeenCalledWith({ + token: ctx.authToken, + secret: 'registration-test-secret', + connectionString: ctx.env.HYPERDRIVE.connectionString, + resourceKind: 'cloud-agent-next', + resourceId: CLOUD_AGENT_SESSION_ID, + ...(organizationId ? { organizationId } : {}), + }); + expect(stub.registerSession).toHaveBeenCalledWith( + expect.objectContaining({ + identity: expect.objectContaining({ userId: USER_ID, orgId: organizationId }), + auth: expect.objectContaining({ kilocodeToken: 'runtime-token' }), + runtimeAuthorizationSeal: 'runtime-seal', + }) + ); + expect(assertKiloModelAvailable).not.toHaveBeenCalled(); + expect(generateSessionIdMock).not.toHaveBeenCalled(); + } + ); + it.each([undefined, '2026-08-01T10:00:00.000Z'])( 'resumes stored clone IDs without changing or inventing reporting age (%s)', async reportingCreatedAt => { diff --git a/services/cloud-agent-next/src/session/session-registration.ts b/services/cloud-agent-next/src/session/session-registration.ts index 6fdc5b5f80..b258476dad 100644 --- a/services/cloud-agent-next/src/session/session-registration.ts +++ b/services/cloud-agent-next/src/session/session-registration.ts @@ -535,31 +535,17 @@ function effectiveSessionRegistrationInput( }; } -async function allocateNewSession( +async function issueSessionRuntimeAuthorization( input: SessionRegistrationInput, ctx: SessionRegistrationContext, - options?: { billingOrigin?: string }, - ledger?: SessionCreationLedgerHooks -): Promise { - const sessionService = new SessionService(); - const initialTurn = input.initialTurn ? acceptInitialTurn(input.initialTurn) : undefined; + cloudAgentSessionId: string, + initialTurn: AcceptedExecutionTurn | undefined +): Promise { const orgId = input.options?.kilocodeOrganizationId; - const cloudAgentSessionId = generateSessionId( - sessionPlaneForNewOwner(ctx.env, { userId: ctx.userId, orgId }) - ); - const kiloSessionId = generateKiloSessionId(); - const reportingCreatedAt = - input.clone && !initialTurn && cloudAgentSessionId.startsWith('agent_') - ? new Date().toISOString() - : undefined; - const worktreeId = - ledger?.worktreeEnabled && cloudAgentSessionId.startsWith('workspace_') - ? cloudAgentWorktreeIdSchema.parse( - `worktree_${cloudAgentSessionId.slice('workspace_'.length)}` - ) - : undefined; - const createdOnPlatform = input.options?.createdOnPlatform ?? 'cloud-agent'; let runtimeAuthorization: NewSessionAllocation['runtimeAuthorization']; + // authMiddleware has verified this bearer (including legacy tokens) against + // its audience and current pepper. Decode only selects the compatibility path; + // createRuntimeAuthorization re-verifies modern claims and runtime admission. const claims = jwt.decode(ctx.authToken); const isPolicyBearing = claims !== null && @@ -595,6 +581,40 @@ async function allocateNewSession( } } + return runtimeAuthorization; +} + +async function allocateNewSession( + input: SessionRegistrationInput, + ctx: SessionRegistrationContext, + options?: { billingOrigin?: string }, + ledger?: SessionCreationLedgerHooks +): Promise { + const sessionService = new SessionService(); + const initialTurn = input.initialTurn ? acceptInitialTurn(input.initialTurn) : undefined; + const orgId = input.options?.kilocodeOrganizationId; + const cloudAgentSessionId = generateSessionId( + sessionPlaneForNewOwner(ctx.env, { userId: ctx.userId, orgId }) + ); + const kiloSessionId = generateKiloSessionId(); + const reportingCreatedAt = + input.clone && !initialTurn && cloudAgentSessionId.startsWith('agent_') + ? new Date().toISOString() + : undefined; + const worktreeId = + ledger?.worktreeEnabled && cloudAgentSessionId.startsWith('workspace_') + ? cloudAgentWorktreeIdSchema.parse( + `worktree_${cloudAgentSessionId.slice('workspace_'.length)}` + ) + : undefined; + const createdOnPlatform = input.options?.createdOnPlatform ?? 'cloud-agent'; + const runtimeAuthorization = await issueSessionRuntimeAuthorization( + input, + ctx, + cloudAgentSessionId, + initialTurn + ); + try { if (ledger) { // Record progress immediately after ID generation (plan P1-A-08b step 3): @@ -1799,6 +1819,12 @@ async function resumeCloneCreate( // `ready` continues. const allocation = rebuildRecordedSessionAllocation(input, ctx, row); + allocation.runtimeAuthorization = await issueSessionRuntimeAuthorization( + input, + ctx, + allocation.cloudAgentSessionId, + allocation.initialTurn + ); const billingOrigin = { billingOrigin: options.billingOrigin }; const result = input.initialTurn === undefined @@ -1866,6 +1892,12 @@ async function resumeFirstWorktreeCreate( } } + allocation.runtimeAuthorization = await issueSessionRuntimeAuthorization( + input, + ctx, + allocation.cloudAgentSessionId, + allocation.initialTurn + ); const result = await registerAndAdmitInitialTurn( input, ctx, From 9c326ff57b1204d38b811bdd5662e03d4e0b62a8 Mon Sep 17 00:00:00 2001 From: Florian Hines Date: Wed, 9 Sep 2026 12:32:27 -0500 Subject: [PATCH 8/9] fix(cloud-agent): attest runtime model validation requests --- .../model-validation.integration.test.ts | 418 ++++++++++++++++++ .../src/model-validation.test.ts | 146 ++++++ .../cloud-agent-next/src/model-validation.ts | 70 ++- .../src/session/model-preflight.test.ts | 66 +++ .../src/session/model-preflight.ts | 21 +- 5 files changed, 718 insertions(+), 3 deletions(-) create mode 100644 apps/web/src/app/api/gateway/model-validation.integration.test.ts diff --git a/apps/web/src/app/api/gateway/model-validation.integration.test.ts b/apps/web/src/app/api/gateway/model-validation.integration.test.ts new file mode 100644 index 0000000000..511f7b3f6b --- /dev/null +++ b/apps/web/src/app/api/gateway/model-validation.integration.test.ts @@ -0,0 +1,418 @@ +import type { RuntimeAuthorization } from '@kilocode/worker-utils/runtime-authorization-contract'; +const mockFetchSessionMetadata = jest.fn(); +const mockGetRuntimeToken = jest.fn(); +jest.mock('../../../../../../services/cloud-agent-next/src/session-service', () => ({ + fetchSessionMetadata: (...args: unknown[]) => mockFetchSessionMetadata(...args), +})); +jest.mock('../../../../../../services/cloud-agent-next/src/sandbox-session/session-stub', () => ({ + resolveSessionStub: () => ({ getRuntimeToken: () => mockGetRuntimeToken() }), +})); +let currentHeaders = new Headers(); +const mockHeaders = jest.fn(() => currentHeaders); +const mockGetServerSession = jest.fn(); + +jest.mock('next/headers', () => ({ + headers: () => mockHeaders(), + cookies: jest.fn(), +})); + +jest.mock('next-auth', () => ({ + __esModule: true, + ...jest.requireActual('next-auth'), + getServerSession: (...args: unknown[]) => mockGetServerSession(...args), +})); + +jest.mock('@sentry/nextjs', () => ({ + ...jest.requireActual('@sentry/nextjs'), + captureException: jest.fn(), +})); +jest.mock('@/lib/redis', () => ({ redisClient: { get: jest.fn(async () => null) } })); +jest.mock('@/lib/ai-gateway/providers/openrouter', () => ({ + getEnhancedOpenRouterModels: jest.fn(), +})); +jest.mock('@/lib/ai-gateway/providers/direct-byok', () => ({ + getDirectByokModelsForUser: jest.fn(), +})); +jest.mock('@/lib/ai-gateway/experiments/list-available-experiment-models', () => ({ + listAvailableExperimentModels: jest.fn(), +})); +jest.mock('@/lib/ai-gateway/byok', () => ({ + addUserByokAvailability: jest.fn(), + getUserByokProviderIds: jest.fn(), +})); +jest.mock('@/lib/organizations/organization-models', () => ({ + getAvailableModelsForOrganization: jest.fn(), +})); + +import { beforeEach, describe, expect, test } from '@jest/globals'; +import { NextRequest } from 'next/server'; +import { POST as personalValidator } from '@/app/api/openrouter/models/validate/route'; +import { POST as organizationValidator } from '@/app/api/organizations/[id]/models/validate/route'; +import { createTestOrganization } from '@/tests/helpers/organization.helper'; +import { signModernKiloToken } from '@kilocode/worker-utils/kilo-token-policy'; +import { RUNTIME_PROXY_ATTESTATION_HEADER } from '@kilocode/worker-utils/runtime-proxy-attestation'; + +jest.mock('../../../../../../services/cloud-agent-next/src/logger', () => ({ + logger: { withFields: () => ({ info: jest.fn(), warn: jest.fn(), error: jest.fn() }) }, +})); +import jwt from 'jsonwebtoken'; +import { JWT_TOKEN_VERSION } from '@/lib/tokens'; +import { NEXTAUTH_SECRET } from '@/lib/config.server'; +import { insertTestUser } from '@/tests/helpers/user.helper'; +import { + KILO_API_AUDIENCE, + KILO_GATEWAY_AUDIENCE, +} from '@kilocode/worker-utils/internal-service-token-audiences'; + +const { getEnhancedOpenRouterModels } = jest.requireMock('@/lib/ai-gateway/providers/openrouter'); +const { getDirectByokModelsForUser } = jest.requireMock('@/lib/ai-gateway/providers/direct-byok'); +const { listAvailableExperimentModels } = jest.requireMock( + '@/lib/ai-gateway/experiments/list-available-experiment-models' +); +const { addUserByokAvailability, getUserByokProviderIds } = + jest.requireMock('@/lib/ai-gateway/byok'); +const { getAvailableModelsForOrganization } = jest.requireMock( + '@/lib/organizations/organization-models' +); + +const publicCatalog = { data: [{ id: 'public/model' }] }; + +// Load the real Worker implementations at runtime without adding Worker ambient +// types and their transitive service graph to the web TypeScript project. +type ValidationEnv = { KILOCODE_BACKEND_BASE_URL: string; NEXTAUTH_SECRET: string }; +type TestSessionMetadata = { + metadataSchemaVersion: 2; + identity: { sessionId: string; userId: string }; + auth: { kilocodeToken: string }; + agent: { model: string }; + workspace: Record; + lifecycle: { version: number; timestamp: number }; +}; +const { assertKiloModelAvailable } = jest.requireActual<{ + assertKiloModelAvailable(input: { + env: ValidationEnv; + submittedModel: string; + originalToken: string; + originalOrganizationId?: string; + procedure: string; + }): Promise; +}>('../../../../../../services/cloud-agent-next/src/model-validation'); +const { preflightExistingPromptModel } = jest.requireActual<{ + preflightExistingPromptModel(input: { + env: ValidationEnv; + userId: string; + cloudAgentSessionId: string; + procedure: string; + }): Promise; +}>('../../../../../../services/cloud-agent-next/src/session/model-preflight'); +const { renewStoredRuntimeAuthorization } = jest.requireActual<{ + renewStoredRuntimeAuthorization(input: { + metadata: TestSessionMetadata; + getAuthorization(): Promise; + putAuthorization(authorization: RuntimeAuthorization): Promise; + getMetadata(): Promise; + putMetadata(metadata: TestSessionMetadata): Promise; + renew(authorization: RuntimeAuthorization): Promise<{ token: string }>; + }): Promise; +}>('../../../../../../services/cloud-agent-next/src/session/runtime-authorization-persistence'); + +const backend = 'https://backend.kilo.test'; +const env = { KILOCODE_BACKEND_BASE_URL: backend, NEXTAUTH_SECRET }; +const privateModel = 'private/byok'; +let proofMode: 'valid' | 'absent' | 'bad' = 'valid'; + +async function runtimeToken( + user: { id: string; api_token_pepper: string | null }, + organizationId?: string +) { + return ( + await signModernKiloToken({ + userId: user.id, + pepper: user.api_token_pepper, + secret: NEXTAUTH_SECRET, + expiresInSeconds: 3600, + env: process.env.NODE_ENV, + audience: [KILO_API_AUDIENCE, KILO_GATEWAY_AUDIENCE, 'session-ingest'], + tokenPurpose: 'delegated-workload', + credentialExchange: false, + extra: { + ...(organizationId ? { organizationId } : {}), + runtimeAuthorization: { + id: crypto.randomUUID(), + resourceKind: 'cloud-agent-next', + resourceId: 'agent-session', + }, + }, + }) + ).token; +} + +describe('cloud-agent model validation through real web authentication', () => { + beforeEach(() => { + jest.clearAllMocks(); + proofMode = 'valid'; + currentHeaders = new Headers(); + mockGetServerSession.mockResolvedValue(null); + getEnhancedOpenRouterModels.mockResolvedValue(publicCatalog); + listAvailableExperimentModels.mockResolvedValue([]); + getDirectByokModelsForUser.mockResolvedValue([]); + getUserByokProviderIds.mockResolvedValue([]); + addUserByokAvailability.mockImplementation(async (models: unknown[]) => models); + getAvailableModelsForOrganization.mockResolvedValue({ data: [] }); + jest.spyOn(global, 'fetch').mockImplementation(async (url, init) => { + currentHeaders = new Headers(init?.headers); + if (proofMode === 'absent') currentHeaders.delete(RUNTIME_PROXY_ATTESTATION_HEADER); + if (proofMode === 'bad') + currentHeaders.set(RUNTIME_PROXY_ATTESTATION_HEADER, 'invalid-proof'); + const request = new NextRequest(String(url), { + ...init, + signal: init?.signal ?? undefined, + headers: currentHeaders, + }); + const match = new URL(String(url)).pathname.match( + /^\/api\/organizations\/([^/]+)\/models\/validate$/ + ); + return match + ? organizationValidator(request, { + params: Promise.resolve({ id: decodeURIComponent(match[1]) }), + }) + : personalValidator(request); + }); + }); + afterEach(() => jest.restoreAllMocks()); + + test.each(['runtime_authorized_session_create', 'send'])( + 'personal private model passes %s with a signed runtime credential', + async procedure => { + const user = await insertTestUser({ api_token_pepper: crypto.randomUUID() }); + getDirectByokModelsForUser.mockImplementation(async (id: string) => + id === user.id ? [{ id: privateModel }] : [] + ); + await expect( + assertKiloModelAvailable({ + env, + submittedModel: privateModel, + originalToken: await runtimeToken(user), + procedure, + }) + ).resolves.toBeUndefined(); + expect(getDirectByokModelsForUser).toHaveBeenCalledWith(user.id); + } + ); + + test.each(['runtime_authorized_session_create', 'send'])( + 'organization private model passes %s with a signed runtime credential', + async procedure => { + const user = await insertTestUser({ api_token_pepper: crypto.randomUUID() }); + const org = await createTestOrganization('model-validation', user.id, 0); + getAvailableModelsForOrganization.mockImplementation(async (id: string) => ({ + data: id === org.id ? [{ id: privateModel }] : [], + })); + await expect( + assertKiloModelAvailable({ + env, + submittedModel: privateModel, + originalToken: await runtimeToken(user, org.id), + originalOrganizationId: org.id, + procedure, + }) + ).resolves.toBeUndefined(); + expect(getAvailableModelsForOrganization).toHaveBeenCalledWith(org.id, { + type: 'member', + kiloUserId: user.id, + allowNonMember: true, + }); + } + ); + + test.each(['active', 'revoked', 'expired'] as const)( + 'expired backing JWT preflight uses the %s delegation', + async state => { + const user = await insertTestUser({ api_token_pepper: crypto.randomUUID() }); + const now = Date.now(); + const id = crypto.randomUUID(); + const extra = { + runtimeAuthorization: { + id, + resourceKind: 'cloud-agent-next' as const, + resourceId: 'agent-session', + }, + }; + const sign = (issuedAt: Date) => + signModernKiloToken({ + userId: user.id, + pepper: user.api_token_pepper, + secret: NEXTAUTH_SECRET, + expiresInSeconds: 3600, + env: process.env.NODE_ENV, + audience: [KILO_API_AUDIENCE, KILO_GATEWAY_AUDIENCE], + tokenPurpose: 'delegated-workload', + credentialExchange: false, + extra, + now: issuedAt, + }); + let metadata: TestSessionMetadata = { + metadataSchemaVersion: 2, + identity: { sessionId: 'agent-session', userId: user.id }, + auth: { kilocodeToken: (await sign(new Date(now - 7200000))).token }, + agent: { model: privateModel }, + workspace: {}, + lifecycle: { version: 1, timestamp: now }, + }; + const record: RuntimeAuthorization = { + version: 1, + id, + resourceKind: 'cloud-agent-next', + resourceId: 'agent-session', + userId: user.id, + authorizationUserId: user.id, + issuedAt: new Date(now - 10800000).toISOString(), + delegationExpiresAt: new Date( + state === 'expired' ? now - 1000 : now + 10800000 + ).toISOString(), + state: state === 'revoked' ? 'revoked' : 'active', + bindings: { userPepperDigest: 'a'.repeat(64), authorizationPepperDigest: 'a'.repeat(64) }, + source: { admissionSource: 'user' }, + }; + const renew = jest.fn(async () => sign(new Date())); + mockFetchSessionMetadata.mockImplementation(async () => metadata); + // Exercise the same persistence/renewal function used by the owning DO's RPC. + mockGetRuntimeToken.mockImplementation(() => + renewStoredRuntimeAuthorization({ + metadata, + getAuthorization: async () => record, + putAuthorization: async () => {}, + getMetadata: async () => metadata, + putMetadata: async value => { + metadata = value; + }, + renew, + }) + ); + getDirectByokModelsForUser.mockResolvedValue([{ id: privateModel }]); + const validation = preflightExistingPromptModel({ + env, + userId: user.id, + cloudAgentSessionId: 'agent-session', + procedure: 'send', + }); + if (state === 'active') { + await expect(validation).resolves.toBeUndefined(); + expect(renew).toHaveBeenCalledTimes(1); + expect(getDirectByokModelsForUser).toHaveBeenCalledWith(user.id); + } else { + await expect(validation).rejects.toThrow(); + expect(renew).not.toHaveBeenCalled(); + expect(global.fetch).not.toHaveBeenCalled(); + } + expect(mockGetRuntimeToken).toHaveBeenCalledTimes(1); + } + ); + + test('legacy personal and organization private catalogs remain available', async () => { + const user = await insertTestUser({ api_token_pepper: crypto.randomUUID() }); + const org = await createTestOrganization('legacy-model-validation', user.id, 0); + getDirectByokModelsForUser.mockResolvedValue([{ id: privateModel }]); + getAvailableModelsForOrganization.mockResolvedValue({ data: [{ id: privateModel }] }); + const bearer = jwt.sign( + { + version: JWT_TOKEN_VERSION, + kiloUserId: user.id, + apiTokenPepper: user.api_token_pepper, + env: process.env.NODE_ENV, + }, + NEXTAUTH_SECRET + ); + for (const organizationId of [undefined, org.id]) { + await expect( + assertKiloModelAvailable({ + env, + submittedModel: privateModel, + originalToken: bearer, + originalOrganizationId: organizationId, + procedure: 'start', + }) + ).resolves.toBeUndefined(); + } + }); + + test('organization members cannot select an unavailable private model', async () => { + const user = await insertTestUser({ api_token_pepper: crypto.randomUUID() }); + const org = await createTestOrganization('unavailable-model-validation', user.id, 0); + getAvailableModelsForOrganization.mockResolvedValue({ + data: [{ id: 'another/private-model' }], + }); + await expect( + assertKiloModelAvailable({ + env, + submittedModel: privateModel, + originalToken: await runtimeToken(user, org.id), + originalOrganizationId: org.id, + procedure: 'send', + }) + ).rejects.toMatchObject({ code: 'BAD_REQUEST' }); + }); + + test('a foreign user cannot select another personal private model', async () => { + const owner = await insertTestUser({ api_token_pepper: crypto.randomUUID() }); + const foreign = await insertTestUser({ api_token_pepper: crypto.randomUUID() }); + getDirectByokModelsForUser.mockImplementation(async (id: string) => + id === owner.id ? [{ id: privateModel }] : [] + ); + await expect( + assertKiloModelAvailable({ + env, + submittedModel: privateModel, + originalToken: await runtimeToken(foreign), + procedure: 'send', + }) + ).rejects.toMatchObject({ code: 'BAD_REQUEST' }); + }); + + test('a non-member cannot validate an organization private model', async () => { + const owner = await insertTestUser({ api_token_pepper: crypto.randomUUID() }); + const foreign = await insertTestUser({ api_token_pepper: crypto.randomUUID() }); + const org = await createTestOrganization('foreign-model-validation', owner.id, 0); + getAvailableModelsForOrganization.mockResolvedValue({ data: [{ id: privateModel }] }); + await expect( + assertKiloModelAvailable({ + env, + submittedModel: privateModel, + originalToken: await runtimeToken(foreign, org.id), + originalOrganizationId: org.id, + procedure: 'send', + }) + ).rejects.toMatchObject({ code: 'FORBIDDEN' }); + expect(getAvailableModelsForOrganization).not.toHaveBeenCalled(); + }); + + test.each(['absent', 'bad'] as const)( + '%s proof cannot authenticate private catalogs', + async mode => { + const user = await insertTestUser({ api_token_pepper: crypto.randomUUID() }); + const org = await createTestOrganization('proof-model-validation', user.id, 0); + getDirectByokModelsForUser.mockResolvedValue([{ id: privateModel }]); + getAvailableModelsForOrganization.mockResolvedValue({ data: [{ id: privateModel }] }); + proofMode = mode; + await expect( + assertKiloModelAvailable({ + env, + submittedModel: privateModel, + originalToken: await runtimeToken(user), + procedure: 'start', + }) + ).rejects.toMatchObject({ code: 'BAD_REQUEST' }); + await expect( + assertKiloModelAvailable({ + env, + submittedModel: privateModel, + originalToken: await runtimeToken(user, org.id), + originalOrganizationId: org.id, + procedure: 'send', + }) + ).rejects.toMatchObject({ code: 'FORBIDDEN' }); + expect(getDirectByokModelsForUser).not.toHaveBeenCalled(); + expect(getAvailableModelsForOrganization).not.toHaveBeenCalled(); + } + ); +}); diff --git a/services/cloud-agent-next/src/model-validation.test.ts b/services/cloud-agent-next/src/model-validation.test.ts index c9016e11fa..da40a4d450 100644 --- a/services/cloud-agent-next/src/model-validation.test.ts +++ b/services/cloud-agent-next/src/model-validation.test.ts @@ -1,3 +1,8 @@ +import { signModernKiloToken } from '@kilocode/worker-utils/kilo-token-policy'; +import { + verifyRuntimeProxyAttestation, + RUNTIME_PROXY_ATTESTATION_HEADER, +} from '@kilocode/worker-utils/runtime-proxy-attestation'; import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; import { TRPCError } from '@trpc/server'; import type { Env } from './types.js'; @@ -315,3 +320,144 @@ describe('buildKiloOverrideValidationUrl', () => { ).toBe('http://localhost:8811/api/organizations/org-1/models/validate'); }); }); + +describe('modern model validation trust boundary', () => { + const secret = 'model-validation-test-secret'; + const env = { KILOCODE_BACKEND_BASE_URL: 'https://backend.test', NEXTAUTH_SECRET: secret }; + const authorizationId = '11111111-1111-4111-8111-111111111111'; + async function token( + options: { secret?: string; audience?: string; organizationId?: string } = {} + ) { + return ( + await signModernKiloToken({ + userId: 'oauth/test', + secret: options.secret ?? secret, + expiresInSeconds: 3600, + audience: options.audience ?? ['kilo-api', 'kilo-gateway'], + tokenPurpose: 'delegated-workload', + credentialExchange: false, + extra: { + organizationId: options.organizationId, + runtimeAuthorization: { + id: authorizationId, + resourceKind: 'cloud-agent-next', + resourceId: 'session-1', + }, + }, + }) + ).token; + } + afterEach(() => vi.restoreAllMocks()); + + it.each([undefined, 'org-1'])( + 'issues a bearer-bound proof for the intended audience (%s)', + async organizationId => { + const bearer = await token({ organizationId }); + const fetchMock = vi.spyOn(global, 'fetch').mockResolvedValue(Response.json({ valid: true })); + await assertKiloModelAvailable({ + env, + submittedModel: 'private/model', + originalToken: bearer, + originalOrganizationId: organizationId, + procedure: 'send', + }); + const init = fetchMock.mock.calls[0][1]; + const headers = new Headers(init?.headers); + expect(init?.redirect).toBe('manual'); + expect( + await verifyRuntimeProxyAttestation({ + secret, + audience: organizationId ? 'kilo-api' : 'kilo-gateway', + userId: 'oauth/test', + authorizationId, + resourceId: 'session-1', + bearer, + value: headers.get(RUNTIME_PROXY_ATTESTATION_HEADER), + }) + ).toBe(true); + } + ); + + it.each([ + 'missing-secret', + 'bad-signature', + 'wrong-audience', + 'foreign-organization', + 'override', + 'encoded-url', + ])('rejects %s before sending credentials', async scenario => { + let bearer = await token({ + secret: scenario === 'bad-signature' ? 'foreign-secret' : undefined, + audience: scenario === 'wrong-audience' ? 'session-ingest' : undefined, + }); + if (scenario === 'encoded-url') bearer = `https://evil.test/api/openrouter:${bearer}`; + const fetchMock = vi.spyOn(global, 'fetch').mockResolvedValue(Response.json({ valid: true })); + await expect( + assertKiloModelAvailable({ + env: { + ...env, + ...(scenario === 'missing-secret' ? { NEXTAUTH_SECRET: undefined } : {}), + ...(scenario === 'override' ? { KILO_OPENROUTER_BASE: 'https://evil.test/api' } : {}), + }, + submittedModel: 'public/model', + originalToken: bearer, + originalOrganizationId: scenario === 'foreign-organization' ? 'other-org' : undefined, + procedure: 'start', + }) + ).rejects.toMatchObject({ code: 'FORBIDDEN' }); + expect(fetchMock).not.toHaveBeenCalled(); + }); + + it.each([401, 404, 302])('does not anonymously retry or skip modern HTTP %s', async status => { + const fetchMock = vi.spyOn(global, 'fetch').mockResolvedValue(new Response(null, { status })); + await expect( + assertKiloModelAvailable({ + env, + submittedModel: 'public/model', + originalToken: await token(), + procedure: 'send', + }) + ).rejects.toMatchObject({ code: status === 401 ? 'FORBIDDEN' : 'SERVICE_UNAVAILABLE' }); + expect(fetchMock).toHaveBeenCalledTimes(1); + }); + + it('preserves non-runtime typed credential fallback without issuing proof', async () => { + const bearer = ( + await signModernKiloToken({ + userId: 'oauth/test', + secret, + expiresInSeconds: 3600, + audience: 'cloud-agent-next', + tokenPurpose: 'human-api', + credentialExchange: false, + }) + ).token; + const fetchMock = vi + .spyOn(global, 'fetch') + .mockResolvedValueOnce(new Response(null, { status: 401 })) + .mockResolvedValueOnce(Response.json({ valid: true })); + await assertKiloModelAvailable({ + env, + submittedModel: 'public/model', + originalToken: bearer, + procedure: 'start', + }); + expect(fetchMock).toHaveBeenCalledTimes(2); + expect( + new Headers(fetchMock.mock.calls[0][1]?.headers).has(RUNTIME_PROXY_ATTESTATION_HEADER) + ).toBe(false); + }); + + it('allows an override that resolves to the exact trusted backend route', async () => { + const fetchMock = vi.spyOn(global, 'fetch').mockResolvedValue(Response.json({ valid: true })); + await assertKiloModelAvailable({ + env: { ...env, KILO_OPENROUTER_BASE: 'https://backend.test/api' }, + submittedModel: 'private/model', + originalToken: await token(), + procedure: 'start', + }); + expect( + new Headers(fetchMock.mock.calls[0][1]?.headers).has(RUNTIME_PROXY_ATTESTATION_HEADER) + ).toBe(true); + }); +}); diff --git a/services/cloud-agent-next/src/model-validation.ts b/services/cloud-agent-next/src/model-validation.ts index 5b4e9d8dcc..4c97cbc0f8 100644 --- a/services/cloud-agent-next/src/model-validation.ts +++ b/services/cloud-agent-next/src/model-validation.ts @@ -1,5 +1,12 @@ import { TRPCError } from '@trpc/server'; import { z } from 'zod'; +import jwt from 'jsonwebtoken'; +import { verifyKiloTokenForPolicy } from '@kilocode/worker-utils/kilo-token-policy'; +import { + issueRuntimeProxyAttestation, + RUNTIME_PROXY_ATTESTATION_HEADER, +} from '@kilocode/worker-utils/runtime-proxy-attestation'; +import { resolveSecret } from './auth.js'; import { DEFAULT_BACKEND_URL } from './constants.js'; import { logger } from './logger.js'; import { dispatchedKilocodeModelId } from './persistence/model-utils.js'; @@ -14,12 +21,15 @@ const MODEL_VALIDATION_UNAVAILABLE_MESSAGE = 'Model availability could not be ve type ModelValidationEnv = Pick< PersistenceEnv, 'KILOCODE_BACKEND_BASE_URL' | 'KILO_OPENROUTER_BASE' | 'KILOCODE_ORG_ID_OVERRIDE' ->; +> & + Partial>; type EffectiveCatalogContext = { token?: string; organizationId?: string; feature: string; + runtimeCredential?: boolean; + proof?: string; }; type ModelValidationResult = @@ -60,6 +70,7 @@ function requestHeaders(context: EffectiveCatalogContext): Headers { const headers = new Headers(); headers.set('Content-Type', 'application/json'); headers.set('X-KiloCode-Feature', context.feature); + if (context.proof) headers.set(RUNTIME_PROXY_ATTESTATION_HEADER, context.proof); if (context.token) headers.set('Authorization', `Bearer ${context.token}`); if (context.organizationId) { headers.set('X-KiloCode-OrganizationId', context.organizationId); @@ -151,13 +162,18 @@ async function validateFromOfficialSource( const result = await validateEndpoint(officialValidationUrl(env, context.organizationId), { method: 'POST', headers: requestHeaders(context), + ...(context.runtimeCredential ? { redirect: 'manual' as const } : {}), body: JSON.stringify({ modelId }), }); if (result.type === 'unavailable') { return { type: 'validation-unavailable', source: 'official' }; } if (result.type === 'http-error') { - if (result.status === 404) return { type: 'skipped', source: 'official' }; + if (result.status === 404 && !context.runtimeCredential) + return { type: 'skipped', source: 'official' }; + if (result.status === 401 && context.runtimeCredential) { + return { type: 'access-denied', source: 'official' }; + } if (result.status === 401 && (context.token || context.organizationId)) { return validateFromOfficialSource(env, modelId, anonymousCatalogContext(context.feature)); } @@ -216,12 +232,16 @@ async function validateFromOverrideSource( const result = await validateEndpoint(validationUrl, { method: 'POST', headers: requestHeaders(context), + ...(context.runtimeCredential ? { redirect: 'manual' as const } : {}), body: JSON.stringify({ modelId }), }); if (result.type === 'unavailable') { return { type: 'validation-unavailable', source: 'override' }; } if (result.type === 'http-error') { + if (result.status === 401 && context.runtimeCredential) { + return { type: 'access-denied', source: 'override' }; + } if (result.status === 401 && (context.token || context.organizationId)) { return validateFromOfficialSource(env, modelId, anonymousCatalogContext(context.feature)); } @@ -233,6 +253,46 @@ async function validateFromOverrideSource( : { type: 'unavailable-model', source: 'override' }; } +// Decoding only selects the stricter path; every claim used to issue proof is verified below. +async function attestCatalogContext( + env: ModelValidationEnv, + context: EffectiveCatalogContext, + selectedUrl: string +): Promise { + if (!context.token) return; + const encodedBase = catalogBaseUrlEncodedInToken(context.token); + const token = encodedBase + ? context.token.slice(context.token.lastIndexOf(':') + 1) + : context.token; + const decoded = jwt.decode(token); + if (!decoded || typeof decoded === 'string' || !('runtimeAuthorization' in decoded)) return; + context.runtimeCredential = true; + const deny = () => + new TRPCError({ code: 'FORBIDDEN', message: 'Model catalog authentication unavailable' }); + // Only the exact configured backend route may receive a runtime proof or bearer. + if (encodedBase || selectedUrl !== officialValidationUrl(env, context.organizationId)) + throw deny(); + const secret = await resolveSecret(env.NEXTAUTH_SECRET); + if (!secret) throw deny(); + try { + const audience = context.organizationId ? 'kilo-api' : 'kilo-gateway'; + const verified = await verifyKiloTokenForPolicy(token, secret, { audience, mode: 'required' }); + if (verified.claims.organizationId !== context.organizationId) throw deny(); + const authorization = verified.claims.runtimeAuthorization; + if (!authorization || authorization.resourceKind !== 'cloud-agent-next') throw deny(); + context.proof = await issueRuntimeProxyAttestation({ + secret, + audience, + userId: verified.userId, + authorizationId: authorization.id, + resourceId: authorization.resourceId, + bearer: token, + }); + } catch { + throw deny(); + } +} + export async function assertKiloModelAvailable( input: AssertKiloModelAvailableInput ): Promise { @@ -247,6 +307,12 @@ export async function assertKiloModelAvailable( const context = effectiveCatalogContext(input); const startTime = Date.now(); const tokenSelectedBaseUrl = catalogBaseUrlEncodedInToken(context.token); + const selectedUrl = tokenSelectedBaseUrl + ? `${tokenSelectedBaseUrl}/models/validate` + : input.env.KILO_OPENROUTER_BASE + ? buildKiloOverrideValidationUrl(input.env.KILO_OPENROUTER_BASE, context.organizationId) + : officialValidationUrl(input.env, context.organizationId); + await attestCatalogContext(input.env, context, selectedUrl); const result = tokenSelectedBaseUrl ? await validateFromOverrideSource(input.env, tokenSelectedBaseUrl, modelId, context, true) : input.env.KILO_OPENROUTER_BASE diff --git a/services/cloud-agent-next/src/session/model-preflight.test.ts b/services/cloud-agent-next/src/session/model-preflight.test.ts index dc8503f9dd..530965ea7d 100644 --- a/services/cloud-agent-next/src/session/model-preflight.test.ts +++ b/services/cloud-agent-next/src/session/model-preflight.test.ts @@ -1,3 +1,5 @@ +import jwt from 'jsonwebtoken'; +import { resolveSessionStub } from '../sandbox-session/session-stub.js'; import { beforeEach, describe, expect, it, vi } from 'vitest'; import { fetchSessionMetadata } from '../session-service.js'; import { assertKiloModelAvailable } from '../model-validation.js'; @@ -7,6 +9,7 @@ import { preflightPreparedInitialPromptModel, } from './model-preflight.js'; +vi.mock('../sandbox-session/session-stub.js', () => ({ resolveSessionStub: vi.fn() })); vi.mock('../session-service.js', () => ({ fetchSessionMetadata: vi.fn() })); vi.mock('../model-validation.js', () => ({ assertKiloModelAvailable: vi.fn() })); @@ -44,6 +47,7 @@ describe('model preflight for stored sessions', () => { procedure: 'send', }); + expect(resolveSessionStub).not.toHaveBeenCalled(); expect(assertKiloModelAvailable).toHaveBeenCalledWith({ env, submittedModel: 'override/model', @@ -128,3 +132,65 @@ describe('model preflight for stored sessions', () => { expect(assertKiloModelAvailable).not.toHaveBeenCalled(); }); }); + +describe('modern preflight token renewal', () => { + const getRuntimeToken = vi.fn<() => Promise>(); + const expiredToken = jwt.sign( + { runtimeAuthorization: { id: '11111111-1111-4111-8111-111111111111' }, exp: 1 }, + 'test-secret' + ); + beforeEach(() => { + vi.clearAllMocks(); + vi.mocked(fetchSessionMetadata).mockResolvedValue({ + ...metadata, + auth: { kilocodeToken: expiredToken }, + }); + vi.mocked(resolveSessionStub).mockReturnValue({ getRuntimeToken } as unknown as ReturnType< + typeof resolveSessionStub + >); + }); + + it.each([preflightExistingPromptModel, preflightPreparedInitialPromptModel])( + 'renews the expired backing token before prompt validation', + async preflight => { + getRuntimeToken.mockResolvedValue('renewed-token'); + await preflight({ + env, + userId: 'user-1', + cloudAgentSessionId: metadata.identity.sessionId, + procedure: 'send', + }); + expect(resolveSessionStub).toHaveBeenCalledWith(env, 'user-1', metadata.identity.sessionId); + expect(getRuntimeToken).toHaveBeenCalledTimes(1); + expect(assertKiloModelAvailable).toHaveBeenCalledWith( + expect.objectContaining({ originalToken: 'renewed-token' }) + ); + } + ); + + it.each(['revoked', 'expired'])('fails closed when delegation is %s', async reason => { + getRuntimeToken.mockRejectedValue(new Error(reason)); + await expect( + preflightExistingPromptModel({ + env, + userId: 'user-1', + cloudAgentSessionId: metadata.identity.sessionId, + procedure: 'send', + }) + ).rejects.toThrow(reason); + expect(assertKiloModelAvailable).not.toHaveBeenCalled(); + }); + + it('fails closed when the owning session returns no credential', async () => { + getRuntimeToken.mockResolvedValue(null); + await expect( + preflightExistingPromptModel({ + env, + userId: 'user-1', + cloudAgentSessionId: metadata.identity.sessionId, + procedure: 'send', + }) + ).rejects.toMatchObject({ code: 'FORBIDDEN' }); + expect(assertKiloModelAvailable).not.toHaveBeenCalled(); + }); +}); diff --git a/services/cloud-agent-next/src/session/model-preflight.ts b/services/cloud-agent-next/src/session/model-preflight.ts index f4074e7443..04b8c9561f 100644 --- a/services/cloud-agent-next/src/session/model-preflight.ts +++ b/services/cloud-agent-next/src/session/model-preflight.ts @@ -2,6 +2,9 @@ import { TRPCError } from '@trpc/server'; import { assertKiloModelAvailable } from '../model-validation.js'; import type { CloudAgentSessionState, PersistenceEnv } from '../persistence/types.js'; import { fetchSessionMetadata } from '../session-service.js'; +import { resolveSessionStub } from '../sandbox-session/session-stub.js'; +import { withDORetry } from '../utils/do-retry.js'; +import { hasModernRuntimeAuthorization } from './runtime-authorization-persistence.js'; type StoredSessionPreflightInput = { env: PersistenceEnv; @@ -29,10 +32,26 @@ async function assertModelFromStoredContext( metadata: CloudAgentSessionState, submittedModel: string | undefined ): Promise { + let token = metadata.auth.kilocodeToken; + if (hasModernRuntimeAuthorization(metadata)) { + // The owning DO renews short-lived backing tokens within the existing delegation. + const runtimeToken = await withDORetry( + () => resolveSessionStub(input.env, input.userId, input.cloudAgentSessionId), + stub => stub.getRuntimeToken(), + 'getRuntimeTokenForModelPreflight' + ); + if (!runtimeToken) { + throw new TRPCError({ + code: 'FORBIDDEN', + message: 'Model catalog authentication unavailable', + }); + } + token = runtimeToken; + } await assertKiloModelAvailable({ env: input.env, submittedModel, - originalToken: metadata.auth.kilocodeToken, + originalToken: token, originalOrganizationId: metadata.identity.orgId, createdOnPlatform: metadata.identity.createdOnPlatform, procedure: input.procedure, From 505b627e1a472f95f6955d95cd4a642a8bb388ba Mon Sep 17 00:00:00 2001 From: Florian Hines Date: Wed, 9 Sep 2026 12:32:38 -0500 Subject: [PATCH 9/9] fix(cloud-agent): attest control-plane session ingest --- .github/workflows/ci.yml | 6 + .../src/sandbox-session/SandboxSession.ts | 87 +++++- .../sandbox-session/control-plane-ingest.ts | 62 +++++ .../session-message-queue.test.ts | 130 ++++++++- .../postgres/control-plane-ingest.test.ts | 247 ++++++++++++++++++ .../session-ingest/vitest.postgres.config.ts | 18 ++ 6 files changed, 539 insertions(+), 11 deletions(-) create mode 100644 services/session-ingest/test/postgres/control-plane-ingest.test.ts create mode 100644 services/session-ingest/vitest.postgres.config.ts diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 660c63b32c..c9aef707c6 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -311,6 +311,12 @@ jobs: # suffix that computeDatabaseUrl adds whenever NODE_ENV=test. run: NODE_ENV=development pnpm drizzle migrate && pnpm --filter container-usage-meter test:postgres + - name: Run control-plane ingest Postgres and Durable Object regression + env: + NODE_ENV: development + CLOUDFLARE_HYPERDRIVE_LOCAL_CONNECTION_STRING_HYPERDRIVE: ${{ env.POSTGRES_URL }} + run: pnpm --filter cloudflare-session-ingest exec vitest run --config vitest.postgres.config.ts + build: needs: [changes, typecheck, lint, format-check, drizzle-check] if: needs.changes.outputs.kilocode_backend == 'true' diff --git a/services/cloud-agent-next/src/sandbox-session/SandboxSession.ts b/services/cloud-agent-next/src/sandbox-session/SandboxSession.ts index 5de62d0afa..e17d45ce40 100644 --- a/services/cloud-agent-next/src/sandbox-session/SandboxSession.ts +++ b/services/cloud-agent-next/src/sandbox-session/SandboxSession.ts @@ -1,3 +1,4 @@ +import jwt from 'jsonwebtoken'; import { DurableObject } from 'cloudflare:workers'; import type { GetWorktreeChangesOutput, @@ -8,6 +9,8 @@ import { TRPCError } from '@trpc/server'; import { withTimeout } from '@kilocode/worker-utils'; import { renewRuntimeAuthorization, + RuntimeAuthorizationExpiredError, + RuntimeAuthorizationRevokedError, unsealRuntimeAuthorization, } from '@kilocode/worker-utils/runtime-authorization'; import type { RuntimeAuthorization } from '@kilocode/worker-utils/runtime-authorization-contract'; @@ -771,11 +774,46 @@ export class SandboxSession extends DurableObject { if (!isChild || internalSecret) { const publication = this.ingestPublicationChain .catch(() => undefined) - .then(() => { + .then(async () => { if (!this.terminalLifecycle.isCurrent(epoch)) return; + const storedAuthorization = this.ctx.storage.kv.get(RUNTIME_AUTHORIZATION_KEY); + const decoded = jwt.decode(token); + // Classification only: the bridge cryptographically verifies every + // modern claim before granting any attestation. + const modern = + storedAuthorization !== undefined || + (typeof decoded === 'object' && + decoded !== null && + 'runtimeAuthorization' in decoded); + const currentToken = modern ? await this.getRuntimeToken() : token; + const secret = modern ? await resolveSecret(this.env.NEXTAUTH_SECRET) : undefined; + const currentMetadata = await this.getMetadata(); + if ( + !currentToken || + (modern && !secret) || + !currentMetadata || + currentMetadata.auth.kiloSessionId !== rootKiloSessionId + ) + return; + const authorization = this.ctx.storage.kv.get(RUNTIME_AUTHORIZATION_KEY); return publishControlPlaneSessionIngest({ fetchIngest: request => this.env.SESSION_INGEST.fetch(request), - token, + token: currentToken, + runtimeContext: + modern && secret + ? { + secret, + userId: currentMetadata.identity.userId, + organizationId: currentMetadata.identity.orgId, + authorization, + isCurrent: () => + this.terminalLifecycle.isCurrent(epoch) && + !this.deletedWorktreeId && + !this.ctx.storage.kv.get(RUNTIME_AUTHORIZATION_RECOVERY_KEY) && + JSON.stringify(this.ctx.storage.kv.get(RUNTIME_AUTHORIZATION_KEY)) === + JSON.stringify(authorization), + } + : undefined, rootKiloSessionId, eventKiloSessionId, cloudAgentSessionId: metadata.identity.sessionId, @@ -784,8 +822,13 @@ export class SandboxSession extends DurableObject { items: ingestItems, }); }); - this.ingestPublicationChain = publication; - this.ctx.waitUntil(publication); + const settledPublication = publication.catch(() => { + logger + .withFields({ sessionId: this.sessionId }) + .warn('Control-plane session ingest authorization unavailable'); + }); + this.ingestPublicationChain = settledPublication; + this.ctx.waitUntil(settledPublication); } } const applied = this.terminalLifecycle.isCurrent(epoch); @@ -2565,10 +2608,37 @@ export class SandboxSession extends DurableObject { let validationFailure: Extract | undefined; if (intent?.turn.type === 'prompt' && origin === 'followup') { try { + let validationToken = metadata.auth.kilocodeToken; + if ( + this.ctx.storage.kv.get(RUNTIME_AUTHORIZATION_KEY) !== undefined || + hasModernRuntimeAuthorization(metadata) + ) { + try { + validationToken = (await this.getRuntimeToken()) ?? undefined; + } catch (error) { + throw new TRPCError({ + code: + error instanceof RuntimeAuthorizationExpiredError || + error instanceof RuntimeAuthorizationRevokedError + ? 'FORBIDDEN' + : 'SERVICE_UNAVAILABLE', + message: 'Runtime credential unavailable for model validation', + }); + } + if (!validationToken) + throw new TRPCError({ code: 'FORBIDDEN', message: 'Runtime credential unavailable' }); + if (this.ctx.storage.kv.get(RUNTIME_AUTHORIZATION_RECOVERY_KEY)) { + return { + success: false, + code: 'COMPUTE_STOPPING', + error: 'Runtime authorization recovery is in progress', + }; + } + } await assertKiloModelAvailable({ env: this.env, submittedModel: intent.agent.model, - originalToken: metadata.auth.kilocodeToken, + originalToken: validationToken, originalOrganizationId: metadata.identity.orgId, createdOnPlatform: metadata.identity.createdOnPlatform, procedure: 'admitSubmittedMessage', @@ -2590,6 +2660,13 @@ export class SandboxSession extends DurableObject { } let admitted = false; const result = this.ctx.storage.transactionSync((): SessionMessageAdmissionResult => { + if (this.ctx.storage.kv.get(RUNTIME_AUTHORIZATION_RECOVERY_KEY)) { + return { + success: false, + code: 'COMPUTE_STOPPING', + error: 'Runtime authorization recovery is in progress', + }; + } const latestMetadata = this.terminalLifecycle.getStoredMetadata(); if (!this.terminalLifecycle.isCurrent(epoch) || !latestMetadata) { return { success: false, code: 'NOT_FOUND', error: 'Session not found' }; diff --git a/services/cloud-agent-next/src/sandbox-session/control-plane-ingest.ts b/services/cloud-agent-next/src/sandbox-session/control-plane-ingest.ts index dbc2313c43..05d156b4f8 100644 --- a/services/cloud-agent-next/src/sandbox-session/control-plane-ingest.ts +++ b/services/cloud-agent-next/src/sandbox-session/control-plane-ingest.ts @@ -1,3 +1,9 @@ +import { verifyKiloTokenForPolicy } from '@kilocode/worker-utils/kilo-token-policy'; +import { RuntimeAuthorizationSchema } from '@kilocode/worker-utils/runtime-authorization-contract'; +import { + issueRuntimeProxyAttestation, + RUNTIME_PROXY_ATTESTATION_HEADER, +} from '@kilocode/worker-utils/runtime-proxy-attestation'; import { z } from 'zod'; import { logger } from '../logger.js'; import { @@ -90,6 +96,14 @@ export async function publishControlPlaneSessionIngest(params: { directory?: string; internalSecret?: string; items: IngestItem[]; + // Supplied only by the owning DO, never from event payloads or JWT claims. + runtimeContext?: { + secret: string; + userId: string; + organizationId?: string; + authorization: unknown; + isCurrent: () => boolean; + }; }): Promise { if (params.items.length === 0) return; const eventKiloSessionId = params.eventKiloSessionId ?? params.rootKiloSessionId; @@ -135,6 +149,53 @@ export async function publishControlPlaneSessionIngest(params: { if (lineage) headers.set(cloudAgentSessionScopeHeaders.trustedLineage, '1'); } + if (params.runtimeContext) { + try { + const context = params.runtimeContext; + const { claims } = await verifyKiloTokenForPolicy(params.token, context.secret, { + audience: 'session-ingest', + mode: 'allow-legacy', + }); + if (claims.kiloUserId !== context.userId || claims.organizationId !== context.organizationId) + return; + if (claims.runtimeAuthorization) { + const authorization = RuntimeAuthorizationSchema.parse(context.authorization); + const reference = claims.runtimeAuthorization; + if ( + authorization.state !== 'active' || + authorization.resourceKind !== 'cloud-agent-next' || + authorization.resourceId !== params.cloudAgentSessionId || + authorization.userId !== context.userId || + authorization.organizationId !== context.organizationId || + reference.id !== authorization.id || + reference.resourceKind !== authorization.resourceKind || + reference.resourceId !== authorization.resourceId || + Date.parse(authorization.issuedAt) > Date.now() || + Date.parse(authorization.delegationExpiresAt) <= Date.now() || + claims.exp * 1000 > Date.parse(authorization.delegationExpiresAt) + ) + return; + headers.set( + RUNTIME_PROXY_ATTESTATION_HEADER, + await issueRuntimeProxyAttestation({ + secret: context.secret, + audience: 'session-ingest', + userId: authorization.userId, + authorizationId: authorization.id, + resourceId: authorization.resourceId, + bearer: params.token, + }) + ); + } else if (context.authorization !== undefined && context.authorization !== null) { + return; + } + if (!context.isCurrent()) return; + } catch { + logger.warn('Control-plane session ingest authorization failed'); + return; + } + } + const body = JSON.stringify({ data: params.items }); const ingestHeaders = new Headers(headers); ingestHeaders.set('Content-Length', String(new TextEncoder().encode(body).byteLength)); @@ -165,6 +226,7 @@ export async function publishControlPlaneSessionIngest(params: { } } + if (params.runtimeContext && !params.runtimeContext.isCurrent()) return; const response = await params.fetchIngest( new Request(ingestUrl, { method: 'POST', diff --git a/services/cloud-agent-next/src/sandbox-session/session-message-queue.test.ts b/services/cloud-agent-next/src/sandbox-session/session-message-queue.test.ts index ce472a67f7..5992299531 100644 --- a/services/cloud-agent-next/src/sandbox-session/session-message-queue.test.ts +++ b/services/cloud-agent-next/src/sandbox-session/session-message-queue.test.ts @@ -1,3 +1,9 @@ +import { signModernKiloToken } from '@kilocode/worker-utils/kilo-token-policy'; +import { + verifyRuntimeProxyAttestation, + RUNTIME_PROXY_ATTESTATION_HEADER, +} from '@kilocode/worker-utils/runtime-proxy-attestation'; +import { assertKiloModelAvailable } from '../model-validation.js'; import jwt from 'jsonwebtoken'; import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest'; import { normalizeCliEvent } from '../../../../packages/cloud-agent-sdk/src/normalizer'; @@ -1084,6 +1090,7 @@ function sessionFixture(overrides: Partial = {}, sharedControl? const values = new Map(); let alarmAt: number | null = null; const errors: unknown[] = []; + const background: Promise[] = []; const kv: SyncKvStorage = { get: (key: string): T | undefined => structuredClone(values.get(key)) as T | undefined, put: (key: string, value: T) => { @@ -1136,6 +1143,7 @@ function sessionFixture(overrides: Partial = {}, sharedControl? blockConcurrencyWhile: async (callback: () => Promise) => callback(), getWebSockets: () => [], waitUntil: (promise: Promise) => { + background.push(promise); void promise.catch(error => { errors.push(error); }); @@ -1204,6 +1212,7 @@ function sessionFixture(overrides: Partial = {}, sharedControl? }, control, env, + settleBackground: () => Promise.all(background), metadata, storage, values, @@ -1755,6 +1764,119 @@ describe('SandboxSession orchestration', () => { expect(fixture.record('b')?.deliveryDeadlineAt).toBe(Date.now() + SESSION_DELIVERY_TIMEOUT_MS); }); + it.each(['modern', 'legacy'] as const)( + 'wires %s control events through the production bridge', + async mode => { + vi.setSystemTime(new Date('2026-01-01T12:00:00.000Z')); + const fixture = sessionFixture(); + await fixture.admit('bridge'); + await fixture.flush(); + const requests: Request[] = []; + fixture.env.SESSION_INGEST = { + fetch: async (request: Request) => { + requests.push(request); + return Response.json({ success: true }); + }, + } as Env['SESSION_INGEST']; + let token = fixture.metadata.auth.kilocodeToken; + if (mode === 'modern') { + installModernRuntimeAuthorization(fixture); + const signed = await signModernKiloToken({ + secret: 'test-secret', + userId: 'user_1', + pepper: null, + audience: ['kilo-api', 'kilo-gateway', 'session-ingest'], + tokenPurpose: 'delegated-workload', + credentialExchange: false, + expiresInSeconds: 3600, + extra: { + runtimeAuthorization: { + id: '44444444-4444-4444-8444-444444444444', + resourceKind: 'cloud-agent-next', + resourceId: SESSION_ID, + }, + }, + }); + token = signed.token; + fixture.storage.kv.put( + 'session_metadata', + serializeSessionMetadata({ + ...fixture.metadata, + auth: { ...fixture.metadata.auth, kilocodeToken: token }, + }) + ); + } else { + fixture.env.NEXTAUTH_SECRET = ''; + } + await expect( + fixture.session.receiveSandboxControlEvent( + receiptedEvent(1, { + type: 'message.updated', + properties: { info: { id: 'msg_bridge', sessionID: 'kilo_root', role: 'user' } }, + }) + ) + ).resolves.toEqual({ applied: true }); + await fixture.flush(); + await fixture.settleBackground(); + expect(requests).toHaveLength(1); + expect(requests[0].headers.get('Authorization')).toBe(`Bearer ${token}`); + if (mode === 'modern') { + expect( + await verifyRuntimeProxyAttestation({ + value: requests[0].headers.get(RUNTIME_PROXY_ATTESTATION_HEADER), + secret: 'test-secret', + audience: 'session-ingest', + userId: 'user_1', + authorizationId: '44444444-4444-4444-8444-444444444444', + resourceId: SESSION_ID, + bearer: token ?? '', + }) + ).toBe(true); + } else expect(requests[0].headers.has(RUNTIME_PROXY_ATTESTATION_HEADER)).toBe(false); + } + ); + + it('refreshes the modern followup model-validation credential and preserves the legacy path', async () => { + const fixture = sessionFixture(); + const runtimeToken = vi + .spyOn(fixture.session, 'getRuntimeToken') + .mockResolvedValue('refreshed-token'); + await fixture.admit('legacy-preflight'); + expect(runtimeToken).not.toHaveBeenCalled(); + installModernRuntimeAuthorization(fixture); + await fixture.admit('modern-preflight'); + expect(runtimeToken).toHaveBeenCalled(); + expect(assertKiloModelAvailable).toHaveBeenLastCalledWith( + expect.objectContaining({ originalToken: 'refreshed-token' }) + ); + }); + + it('maps modern credential infrastructure failure to retryable model-validation failure', async () => { + const fixture = sessionFixture(); + installModernRuntimeAuthorization(fixture); + vi.spyOn(fixture.session, 'getRuntimeToken').mockRejectedValue(new Error('unavailable')); + await expect(fixture.admit('unavailable')).resolves.toMatchObject({ + success: false, + code: 'MODEL_VALIDATION_UNAVAILABLE', + }); + expect(fixture.record('unavailable')).toBeUndefined(); + }); + + it('does not admit when recovery starts during model validation', async () => { + const fixture = sessionFixture(); + vi.mocked(assertKiloModelAvailable).mockImplementationOnce(async () => { + fixture.storage.kv.put('runtime_authorization_recovery', { + expectedOldId: 'old', + recoveryId: 'new', + }); + }); + await expect(fixture.admit('recovering')).resolves.toMatchObject({ + success: false, + code: 'COMPUTE_STOPPING', + }); + expect(fixture.record('recovering')).toBeUndefined(); + }); + it('commits one canonical outcome with its receipt and replays it after a lost response', async () => { const fixture = sessionFixture(); await fixture.admit('receipt'); @@ -5804,12 +5926,8 @@ describe('SandboxSession orchestration', () => { }); it('keeps a persisted modern attachment isolated after the rollout flag is disabled', async () => { - const fixture = sessionFixture({ - auth: { - kiloSessionId: 'kilo_root', - kilocodeToken: 'eyJhbGciOiJub25lIn0.eyJydW50aW1lQXV0aG9yaXphdGlvbiI6e319.', - }, - }); + const fixture = sessionFixture(); + installModernRuntimeAuthorization(fixture); fixture.env.RUNTIME_ISOLATION_ENABLED = 'false'; await fixture.admit('modern'); diff --git a/services/session-ingest/test/postgres/control-plane-ingest.test.ts b/services/session-ingest/test/postgres/control-plane-ingest.test.ts new file mode 100644 index 0000000000..fedc46df2b --- /dev/null +++ b/services/session-ingest/test/postgres/control-plane-ingest.test.ts @@ -0,0 +1,247 @@ +import { cloudAgentSessionScopeHeaders } from '@kilocode/session-ingest-contracts'; +import { Socket } from 'node:net'; +import { env } from 'cloudflare:test'; +import { afterEach, describe, expect, it, vi } from 'vitest'; +import { eq } from 'drizzle-orm'; +import { getWorkerDb } from '@kilocode/db/client'; +import type * as DbClient from '@kilocode/db/client'; +import { cli_sessions_v2, kilocode_users } from '@kilocode/db/schema'; +import { createRuntimeAuthorization } from '@kilocode/worker-utils/runtime-authorization'; +import { signModernKiloToken } from '@kilocode/worker-utils/kilo-token-policy'; +import { signKiloToken } from '@kilocode/worker-utils/kilo-token'; +import { RUNTIME_PROXY_ATTESTATION_HEADER } from '@kilocode/worker-utils/runtime-proxy-attestation'; +import { publishControlPlaneSessionIngest } from '../../../cloud-agent-next/src/sandbox-session/control-plane-ingest'; +import { app } from '../../src/app'; +import { getSessionExport } from '../../src/services/session-export'; +import type { Env } from '../../src/env'; + +// The Workers test module loader cannot require pg-cloudflare's optional CJS +// export. Supply the runtime's real node:net socket; SQL and storage stay real. +vi.mock('@kilocode/db/client', async importOriginal => { + const original = await importOriginal(); + return { + ...original, + getWorkerDb: (connectionString: string) => + original.getWorkerDb(connectionString, { + stream: () => new Socket(), + } as Parameters[1]), + }; +}); + +// Requires a migrated test PostgreSQL via the test Hyperdrive binding. All HTTP +// middleware, SQL authorization/lineage, R2 and Durable Objects are production code. +const secret = 'control-plane-ingest-integration-secret'; +const db = () => getWorkerDb(env.HYPERDRIVE.connectionString); +const users: string[] = []; +afterEach(async () => { + for (const userId of users.splice(0)) { + const rows = await db() + .select() + .from(cli_sessions_v2) + .where(eq(cli_sessions_v2.kilo_user_id, userId)); + for (const row of rows.filter(row => row.parent_session_id !== null)) { + await db().delete(cli_sessions_v2).where(eq(cli_sessions_v2.session_id, row.session_id)); + } + await db().delete(cli_sessions_v2).where(eq(cli_sessions_v2.kilo_user_id, userId)); + await db().delete(kilocode_users).where(eq(kilocode_users.id, userId)); + } +}); + +async function fixture() { + const userId = `oauth/github:bridge-${crypto.randomUUID()}`; + users.push(userId); + const rootKiloSessionId = `ses_${crypto.randomUUID().replaceAll('-', '').slice(0, 26)}`; + const eventKiloSessionId = `ses_${crypto.randomUUID().replaceAll('-', '').slice(0, 26)}`; + const cloudAgentSessionId = crypto.randomUUID(); + await db() + .insert(kilocode_users) + .values({ + id: userId, + google_user_email: `${crypto.randomUUID()}@example.test`, + google_user_name: 'Bridge test', + google_user_image_url: '', + stripe_customer_id: '', + api_token_pepper: null, + }); + await db().insert(cli_sessions_v2).values({ + session_id: rootKiloSessionId, + kilo_user_id: userId, + cloud_agent_session_id: cloudAgentSessionId, + }); + const admission = await signModernKiloToken({ + userId, + pepper: null, + secret, + audience: 'cloud-agent-next', + tokenPurpose: 'human-api', + credentialExchange: false, + expiresInSeconds: 3600, + extra: { + runtimeAdmission: { source: 'user', authorizationUserId: userId, authorizationPepper: null }, + }, + }); + const runtime = await createRuntimeAuthorization({ + token: admission.token, + secret, + connectionString: env.HYPERDRIVE.connectionString, + resourceKind: 'cloud-agent-next', + resourceId: cloudAgentSessionId, + }); + const bindings = { + ...env, + NEXTAUTH_SECRET_PROD: { get: async () => secret }, + INTERNAL_API_SECRET_PROD: { get: async () => 'bridge-internal-secret' }, + DIRECT_INGEST_PERCENT: '100', + DIRECT_INGEST_USER_IDS: '', + DIRECT_INGEST_MAX_BYTES: '1048576', + } as Env; + const requests: Request[] = []; + const responses: Response[] = []; + const params = { + fetchIngest: async (request: Request) => { + requests.push(request.clone()); + const response = await app.fetch(request, bindings); + responses.push(response.clone()); + return response; + }, + token: runtime.token, + rootKiloSessionId, + eventKiloSessionId, + cloudAgentSessionId, + directory: '/workspace/bridge', + internalSecret: 'bridge-internal-secret', + runtimeContext: { secret, userId, authorization: runtime.authorization, isCurrent: () => true }, + items: [ + { + type: 'session', + data: { + id: eventKiloSessionId, + parentID: rootKiloSessionId, + directory: '/workspace/bridge', + }, + }, + { type: 'message', data: { id: 'msg_child', sessionID: eventKiloSessionId, role: 'user' } }, + { + type: 'part', + data: { + id: 'prt_child', + messageID: 'msg_child', + type: 'text', + text: 'persisted child message', + }, + }, + ], + }; + return { params, runtime, bindings, requests, responses, userId }; +} + +describe('authorized control-plane ingest bridge', () => { + it.each(['modern', 'legacy'] as const)( + 'persists and exports %s child lineage/messages through the actual app and DO', + async mode => { + const f = await fixture(); + if (mode === 'legacy') { + const legacy = await signKiloToken({ + userId: f.userId, + pepper: null, + secret, + expiresInSeconds: 3600, + }); + f.params.token = legacy.token; + await publishControlPlaneSessionIngest({ ...f.params, runtimeContext: undefined }); + } else { + await publishControlPlaneSessionIngest(f.params); + } + expect(f.responses.map(response => response.status)).toEqual([200, 200]); + expect( + f.requests.every(request => request.headers.has(RUNTIME_PROXY_ATTESTATION_HEADER)) + ).toBe(mode === 'modern'); + const [child] = await db() + .select() + .from(cli_sessions_v2) + .where(eq(cli_sessions_v2.session_id, f.params.eventKiloSessionId)); + expect(child).toMatchObject({ + parent_session_id: f.params.rootKiloSessionId, + cloud_agent_session_scope_id: f.params.cloudAgentSessionId, + }); + const stream = await getSessionExport(f.bindings, f.params.eventKiloSessionId, f.userId); + expect(stream).not.toBeNull(); + const exported = await new Response(stream).json(); + expect(exported).toMatchObject({ + messages: [{ info: { id: 'msg_child' }, parts: [{ text: 'persisted child message' }] }], + }); + } + ); + + it('rejects an unattested runtime bearer at the actual middleware', async () => { + const f = await fixture(); + await publishControlPlaneSessionIngest({ ...f.params, runtimeContext: undefined }); + expect(f.responses.map(response => response.status)).toEqual([401]); + expect( + await db() + .select() + .from(cli_sessions_v2) + .where(eq(cli_sessions_v2.session_id, f.params.eventKiloSessionId)) + ).toEqual([]); + }); + + it('rejects foreign root scope at the actual child-create route even with valid proof', async () => { + const f = await fixture(); + await publishControlPlaneSessionIngest({ + ...f.params, + fetchIngest: request => { + const headers = new Headers(request.headers); + headers.set(cloudAgentSessionScopeHeaders.cloudAgentSessionId, crypto.randomUUID()); + return f.params.fetchIngest(new Request(request, { headers })); + }, + }); + expect(f.responses[0].status).toBe(404); + expect( + await db() + .select() + .from(cli_sessions_v2) + .where(eq(cli_sessions_v2.session_id, f.params.eventKiloSessionId)) + ).toEqual([]); + }); + + it.each([ + 'foreign-resource', + 'foreign-owner', + 'revoked', + 'expired', + 'replaced', + 'wrong-audience', + 'forged-signature', + ] as const)('does not issue proof or publish with %s authority', async failure => { + const f = await fixture(); + if (failure === 'foreign-resource') f.params.cloudAgentSessionId = crypto.randomUUID(); + if (failure === 'foreign-owner') f.params.runtimeContext.userId = 'foreign-owner'; + if (failure === 'revoked') f.params.runtimeContext.authorization.state = 'revoked'; + if (failure === 'expired') + f.params.runtimeContext.authorization.delegationExpiresAt = new Date( + Date.now() - 1000 + ).toISOString(); + if (failure === 'replaced') f.params.runtimeContext.isCurrent = () => false; + if (failure === 'wrong-audience' || failure === 'forged-signature') { + const token = await signModernKiloToken({ + userId: f.userId, + pepper: null, + secret: failure === 'forged-signature' ? 'wrong-secret' : secret, + audience: failure === 'wrong-audience' ? 'kilo-api' : 'session-ingest', + tokenPurpose: 'delegated-workload', + credentialExchange: false, + expiresInSeconds: 3600, + extra: { + runtimeAuthorization: { + id: f.runtime.authorization.id, + resourceKind: 'cloud-agent-next', + resourceId: f.params.cloudAgentSessionId, + }, + }, + }); + f.params.token = token.token; + } + await publishControlPlaneSessionIngest(f.params); + expect(f.requests).toEqual([]); + }); +}); diff --git a/services/session-ingest/vitest.postgres.config.ts b/services/session-ingest/vitest.postgres.config.ts new file mode 100644 index 0000000000..0a71cb41f7 --- /dev/null +++ b/services/session-ingest/vitest.postgres.config.ts @@ -0,0 +1,18 @@ +import { defineConfig } from 'vitest/config'; +import workersConfig from './vitest.workers.config'; + +// Kept separate from the SQLite-only suite: this requires migrated PostgreSQL. +if (!process.env.CLOUDFLARE_HYPERDRIVE_LOCAL_CONNECTION_STRING_HYPERDRIVE) { + throw new Error( + 'Set CLOUDFLARE_HYPERDRIVE_LOCAL_CONNECTION_STRING_HYPERDRIVE to a migrated test database' + ); +} + +export default defineConfig({ + ...workersConfig, + test: { + ...workersConfig.test, + name: 'postgres', + include: ['test/postgres/**/*.test.ts'], + }, +});