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/ENVIRONMENT.md b/ENVIRONMENT.md index e1514d364a..33c5de1641 100644 --- a/ENVIRONMENT.md +++ b/ENVIRONMENT.md @@ -90,7 +90,7 @@ Manage shared web env var additions and rotations with `pnpm web:env set ({ + 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/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 65c07a4087..83fdcf1c6d 100644 --- a/apps/web/src/routers/cloud-agent-next-router.test.ts +++ b/apps/web/src/routers/cloud-agent-next-router.test.ts @@ -146,8 +146,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 5ef263fcd3..bcc31460a3 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 { @@ -123,6 +124,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) * @@ -156,7 +166,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, @@ -237,7 +247,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, @@ -264,7 +274,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 @@ -330,7 +340,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); }), @@ -339,7 +350,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,8 @@ export const cloudAgentNextRouter = createTRPCRouter({ .output(getWorktreeFileOutputSchema) .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.getWorktreeFile(input); }), @@ -359,7 +372,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({ @@ -398,7 +411,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) { @@ -413,7 +426,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) { @@ -508,7 +521,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); @@ -524,7 +537,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); @@ -535,7 +548,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); }), @@ -545,7 +558,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); }), @@ -555,7 +568,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); }), @@ -569,7 +582,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); @@ -585,9 +598,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 @@ -595,7 +608,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 6f39ad9989..4a68159a01 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 @@ -172,8 +172,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 e939ee63ad..9425cd6e39 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, @@ -65,6 +66,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'; @@ -272,7 +287,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, @@ -377,7 +396,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, @@ -408,7 +431,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 @@ -482,7 +509,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); }), @@ -495,7 +527,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); }), @@ -508,7 +545,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.getWorktreeFile({ cloudAgentSessionId: input.cloudAgentSessionId, path: input.path, @@ -527,7 +569,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, @@ -580,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.resizeTerminal({ cloudAgentSessionId: input.cloudAgentSessionId, @@ -604,7 +654,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, @@ -692,7 +746,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); @@ -712,7 +770,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); @@ -727,7 +789,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, @@ -745,7 +811,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, @@ -762,7 +832,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, @@ -784,7 +858,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); @@ -810,9 +888,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 @@ -824,7 +902,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 07be7b485e..dcb43c52e8 100644 --- a/services/cloud-agent-next/.dev.vars.example +++ b/services/cloud-agent-next/.dev.vars.example @@ -47,6 +47,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=true + 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/kilo-targets.test.ts b/services/cloud-agent-next/src/kilo/kilo-targets.test.ts index 6873bb3953..4cdc00643f 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,52 @@ import { deriveKiloSandboxTargets, providerBaseUrlEncodedInToken, } from './kilo-targets.js'; +import { + inferRuntimeCredentialProxyRoute, + resolveRuntimeCredentialProxyRoute, +} from './runtime-credential-proxy-routes.js'; + +// 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([ + ['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); + }); +}); 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/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/persistence/CloudAgentSession.ts b/services/cloud-agent-next/src/persistence/CloudAgentSession.ts index b567e6aae9..e548ebdefb 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,23 @@ 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, + 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'; +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 +157,7 @@ import { import { createQueuedSessionMessageState, getSessionMessageState, + hasNonTerminalSessionMessage, listNonTerminalAcceptedMessages, markAgentActivityObserved, markMessageAccepted, @@ -251,6 +275,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 +1692,390 @@ 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 = 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 { + this.inspectRuntimeAuthorizationRecovery(); + 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) { + 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, { + 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; + } + 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() + : this.orchestrator + ? { status: 'absent' as const } + : await createAgentSandbox( + this.env, + metadata, + this.getAgentSandboxRuntimeContext() + ).observeWrappersWithoutWaking(); + 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) { + 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()); + 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') { + 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) + ); + 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 transaction.delete(RUNTIME_AUTHORIZATION_RECOVERY_DIAGNOSTICS_KEY); + }); + failureReason = 'wrapper_identity_clear_failed'; + await clearWrapperRuntimeIdentity(this.ctx.storage, {}, { incrementGeneration: true }); + return { status: 'recovered' }; + } catch { + diagnostic(failureReason); + 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 +2566,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 +2598,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 +2912,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 +3020,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 +4097,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 +4167,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 a81af3db0e..12efb305aa 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' || input.operation === 'session.git.snapshot') { @@ -705,6 +786,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 @@ -761,6 +850,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 () => { @@ -1877,6 +1975,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); @@ -2365,6 +2547,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 b1da297637..76af29f386 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, sessionGitSnapshotPayloadSchema, @@ -61,6 +62,7 @@ const REQUEST_PAYLOAD_SCHEMAS: Record = { 'session.git.summary': sessionGitSummaryPayloadSchema, 'session.git.snapshot': sessionGitSnapshotPayloadSchema, '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 e00c4abbca..579d4d05a2 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 d4df6143f9..ea1ad8acdc 100644 --- a/services/cloud-agent-next/src/sandbox-control/socket.test.ts +++ b/services/cloud-agent-next/src/sandbox-control/socket.test.ts @@ -66,6 +66,8 @@ function helloFrame( requestId = 'req_hello', capabilities?: { nativeRuntimeRetirement?: boolean; + runtimeIsolation?: true; + runtimeRecovery?: true; scopedCleanupResult?: boolean; workingBranches?: boolean; } @@ -105,6 +107,29 @@ describe('sandbox control socket handler', () => { expect(JSON.stringify(parsed)).not.toContain('private'); } }); + it('retains runtime isolation and recovery alongside scoped cleanup negotiation', 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, + runtimeRecovery: true, + scopedCleanupResult: true, + }) + ); + + expect(handler.getConnectionIdentity()).toMatchObject({ + providerInstanceId: 'inst_1', + runtimeIsolation: true, + runtimeRecovery: true, + }); + expect(handler.supportsScopedCleanupResult?.()).toBe(true); + expect(incoming.send).toHaveBeenCalledWith( + expect.stringContaining('"scopedCleanupResult":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 91d23e150f..a2da4764c7 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 }; @@ -227,6 +230,8 @@ function readConnectionIdentity( providerInstanceId: attachment.providerInstanceId, ...(attachment.recoveryCapable ? { recoveryCapable: true } : {}), ...(attachment.wrapperInstanceId ? { wrapperInstanceId: attachment.wrapperInstanceId } : {}), + ...(attachment.runtimeIsolation ? { runtimeIsolation: true } : {}), + ...(attachment.runtimeRecovery ? { runtimeRecovery: true } : {}), }; } @@ -550,6 +555,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, @@ -561,6 +568,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 8a114ee4b2..4c1f49a74d 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,25 @@ import type { import { generateBranchSlug } from '@kilocode/worker-utils/deployment-slug'; 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'; +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 { @@ -87,6 +107,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 { @@ -122,12 +151,14 @@ import { sessionOperationExpiresAt, sessionOperationResultHash, sessionPromptResultSchema, + sessionRuntimeRetireResultSchema, sessionSyncResultSchema, sessionPermissionResolveResultSchema, sessionQuestionResolveResultSchema, sessionAbortResultSchema, sameSessionOperation, wrapperInstanceIdSchema, + type SessionAttachPayload, type SessionOperationAck, type SessionOperationAuthorization, type SessionRequestIdentity, @@ -275,6 +306,7 @@ type ControlEventEvaluationRequest = type SandboxSessionRegistrationInput = { identity: SessionMetadata['identity']; auth: SessionMetadata['auth']; + runtimeAuthorizationSeal?: string; agent: SessionMetadata['agent']; repository?: SessionMetadata['repository']; workspace?: SessionMetadata['workspace']; @@ -765,11 +797,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, @@ -778,8 +845,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); @@ -1037,11 +1109,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(); @@ -1500,6 +1860,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)); @@ -1510,6 +1873,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)); } @@ -1868,6 +2234,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) { @@ -1902,6 +2290,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, @@ -2264,6 +2655,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) { @@ -2305,10 +2703,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', @@ -2330,6 +2755,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' }; @@ -2435,6 +2867,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(); @@ -2553,7 +2986,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 = { @@ -2567,7 +3005,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'), @@ -2807,10 +3245,73 @@ export class SandboxSession extends DurableObject { attachInPreparation = needsPreparation; 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 = 'preparing'; await wait(() => control.attachSession({ @@ -2830,7 +3331,21 @@ export class SandboxSession extends DurableObject { } phase = 'attach'; 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; } @@ -2844,7 +3359,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-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/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 650510c356..4561b21189 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,10 @@ +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'; import { createServiceState } from '../../../../packages/cloud-agent-sdk/src/service-state'; @@ -32,6 +39,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 { PENDING_SESSION_MESSAGE_LIMIT } from '../session/pending-messages.js'; import { createControlStopRequest } from '../shared/control-plane-session.js'; import type { @@ -1109,6 +1117,7 @@ function sessionFixture( 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) => { @@ -1161,6 +1170,7 @@ function sessionFixture( blockConcurrencyWhile: async (callback: () => Promise) => callback(), getWebSockets: () => [], waitUntil: (promise: Promise) => { + background.push(promise); void promise.catch(error => { errors.push(error); }); @@ -1184,6 +1194,13 @@ function sessionFixture( }); 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, @@ -1191,6 +1208,7 @@ function sessionFixture( }) ), attachSession: vi.fn(async () => ({})), + bindRuntimeCredentialProxyHandle: vi.fn(async () => ({ bound: true as const })), detachSession: vi.fn(async () => ({ existed: true })), quarantineRuntime: vi.fn( async ( @@ -1208,6 +1226,8 @@ function sessionFixture( } satisfies Control; const env = { SANDBOX_CONTROL: { getByName: () => sharedControl ?? control }, + WORKER_URL: 'https://worker.example.test', + NEXTAUTH_SECRET: 'test-secret', CALLBACK_QUEUE: callbackQueue, CLOUD_AGENT_CONTAINER_BILLING_ENABLED: 'true', CLOUD_AGENT_CONTAINER_BILLING_ORG_IDS: 'org_1', @@ -1219,6 +1239,8 @@ function sessionFixture( return session; }, control, + env, + settleBackground: () => Promise.all(background), metadata, storage, values, @@ -1298,6 +1320,38 @@ function sessionFixture( }; } +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'], @@ -1838,6 +1892,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'); @@ -4562,6 +4729,102 @@ 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('fences admissions and snapshots callbacks before deletion waits on an interrupt', async () => { const send = vi.fn(async (_job: CallbackJob) => ({}) as QueueSendResponse); const fixture = sessionFixture( @@ -5830,6 +6093,22 @@ 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(); + installModernRuntimeAuthorization(fixture); + 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 a5d88e53a1..abd7a0d8e7 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 { @@ -152,7 +155,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 { @@ -228,6 +235,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 8f9e8da381..68f9c1b4b2 100644 --- a/services/cloud-agent-next/src/sandbox-session/terminal-lifecycle.ts +++ b/services/cloud-agent-next/src/sandbox-session/terminal-lifecycle.ts @@ -816,6 +816,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); @@ -828,6 +845,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/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, 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..b8bcc317cf --- /dev/null +++ b/services/cloud-agent-next/src/session/runtime-authorization-persistence.test.ts @@ -0,0 +1,334 @@ +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, + inspectRuntimeAuthorizationRecoveryLock, + runtimeAuthorizationRecoveryLockSchema, + RUNTIME_AUTHORIZATION_RECOVERY_WARNING_MS, + 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); + }); +}); + +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 new file mode 100644 index 0000000000..e43fbad6ad --- /dev/null +++ b/services/cloud-agent-next/src/session/runtime-authorization-persistence.ts @@ -0,0 +1,185 @@ +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(); +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 { + 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 b8258d9df4..3a17b56098 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, @@ -384,6 +422,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 => { @@ -2163,6 +2217,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 => { @@ -3661,6 +3804,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 d17a252b02..d6c86384f8 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, @@ -176,6 +182,7 @@ type NewSessionAllocation = SessionRegistrationResult & { credentialContainment: CredentialContainment; sessionService: SessionService; rollbackCliSession: () => Promise; + runtimeAuthorization?: { token: string; seal: string }; }; // ----- operation-ledger boundary (P1-A-08b) ----------------------------------- @@ -510,6 +517,11 @@ function sessionPlaneForCreate( ); } +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) { @@ -531,6 +543,55 @@ function effectiveSessionRegistrationInput( }; } +async function issueSessionRuntimeAuthorization( + input: SessionRegistrationInput, + ctx: SessionRegistrationContext, + cloudAgentSessionId: string, + initialTurn: AcceptedExecutionTurn | undefined +): Promise { + const orgId = input.options?.kilocodeOrganizationId; + 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 && + 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', + }); + } + } + + return runtimeAuthorization; +} + async function allocateNewSession( input: SessionRegistrationInput, ctx: SessionRegistrationContext, @@ -553,6 +614,12 @@ async function allocateNewSession( ) : undefined; const createdOnPlatform = input.options?.createdOnPlatform ?? 'cloud-agent'; + const runtimeAuthorization = await issueSessionRuntimeAuthorization( + input, + ctx, + cloudAgentSessionId, + initialTurn + ); try { if (ledger) { @@ -779,6 +846,7 @@ async function allocateNewSession( .error('Failed to rollback cli_sessions_v2 record'); } }, + ...(runtimeAuthorization ? { runtimeAuthorization } : {}), }; } @@ -914,8 +982,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, @@ -1754,6 +1825,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 @@ -1821,6 +1898,12 @@ async function resumeFirstWorktreeCreate( } } + allocation.runtimeAuthorization = await issueSessionRuntimeAuthorization( + input, + ctx, + allocation.cloudAgentSessionId, + allocation.initialTurn + ); const result = await registerAndAdmitInitialTurn( input, ctx, 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 a66f5ac4b2..371bbc7a3c 100644 --- a/services/cloud-agent-next/src/shared/sandbox-control-protocol.ts +++ b/services/cloud-agent-next/src/shared/sandbox-control-protocol.ts @@ -72,6 +72,7 @@ export const SESSION_OPERATIONS = [ 'session.git.summary', 'session.git.snapshot', 'session.detach', + 'session.runtime.retire', 'session.terminal.create', 'session.terminal.resize', 'session.terminal.close', @@ -189,6 +190,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(), scopedCleanupResult: z.boolean().optional(), workingBranches: z.boolean().optional(), }) @@ -206,6 +209,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(), scopedCleanupResult: z.boolean().optional(), }) .optional(), @@ -382,6 +387,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), @@ -604,6 +610,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), @@ -754,6 +768,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; @@ -949,6 +965,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 bc5cfbe16e..719c44a8da 100644 --- a/services/cloud-agent-next/src/types.ts +++ b/services/cloud-agent-next/src/types.ts @@ -604,6 +604,7 @@ export type Env = { /** Comma-separated user or org IDs admitted to the call-home control plane for interactive web creates. `*` 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 012bb82178..372ff52958 100644 --- a/services/cloud-agent-next/test/e2e/README.md +++ b/services/cloud-agent-next/test/e2e/README.md @@ -171,10 +171,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` @@ -233,7 +241,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 @@ -273,7 +282,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. | @@ -331,19 +340,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 852e4c72c1..3ad2482aba 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; @@ -206,6 +208,7 @@ export async function startSession( api === 'legacy' ? await startSessionLegacy(config, started) : await startSessionUnified(config, started); + 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` @@ -302,6 +305,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 16fe642b68..def28fbbc8 100644 --- a/services/cloud-agent-next/test/e2e/fake-llm-server.ts +++ b/services/cloud-agent-next/test/e2e/fake-llm-server.ts @@ -759,6 +759,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..f9c0e56f79 --- /dev/null +++ b/services/cloud-agent-next/test/integration/runtime-authorization-recovery.test.ts @@ -0,0 +1,631 @@ +import { env, listDurableObjectIds, runInDurableObject } from 'cloudflare:test'; +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'; +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'; + +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('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'; + 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 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, + recoveryId: '00000000-0000-4000-8000-000000000253', + 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, + 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 f389758ba4..0b9e0b877b 100644 --- a/services/cloud-agent-next/test/integration/sandbox-control.test.ts +++ b/services/cloud-agent-next/test/integration/sandbox-control.test.ts @@ -54,6 +54,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 { @@ -1100,6 +1104,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, @@ -3713,6 +3718,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; @@ -6451,6 +6529,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); @@ -6469,6 +6549,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({ @@ -6481,7 +6574,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, @@ -6510,6 +6604,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 }), @@ -6522,6 +6619,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 51b75ebddf..ff8a836313 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'); @@ -1096,6 +1105,7 @@ type ProcFixture = { port: number; inode: string; directory: string; + home: string; roots: string[]; address?: string; state?: string; @@ -1152,6 +1162,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'); @@ -1243,6 +1255,7 @@ function directoryProcesses(): ProcFixture[] { port: 41001, inode: '501', directory: '/workspace/worktrees/worktree-a', + home: '/tmp/kilo-worktrees/a1b2c3d4', roots: ['ses_a', 'ses_sibling'], }, { @@ -1250,6 +1263,7 @@ function directoryProcesses(): ProcFixture[] { port: 41002, inode: '502', directory: '/workspace/worktrees/worktree-b', + home: '/tmp/kilo-worktrees/e5f6a7b8', roots: ['ses_b'], address: '00000000000000000000000001000000', }, @@ -1267,13 +1281,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 12c93dc1bd..7fd57040ae 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 @@ -835,8 +835,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 dd2ef41c98..b488f3b0d7 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: c8a767d5a742c2b932c5d8038b7dec5e) +// Generated by Wrangler by running `wrangler types` (hash: f489c5c6e23b7870ee56b12b9167fc99) // 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?: "*"; + RUNTIME_ISOLATION_ENABLED: "true"; 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: "*"; + RUNTIME_ISOLATION_ENABLED: "true"; 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 21cdd2c3e5..97d29c855b 100644 --- a/services/cloud-agent-next/wrangler.jsonc +++ b/services/cloud-agent-next/wrangler.jsonc @@ -64,6 +64,7 @@ "CLOUD_AGENT_CONTAINER_BILLING_USER_IDS": "f1a848ca-bade-48d8-a5ad-1042d08651e6", "CLOUD_AGENT_CONTAINER_BILLING_ORG_IDS": "9d278969-5453-4ae3-a51f-a8d2274a7b56", "CONTAINER_BILLING_HEARTBEAT_SECONDS": "300", + "RUNTIME_ISOLATION_ENABLED": "true", "VERCEL_SANDBOX_ORG_IDS": "", "VERCEL_PROJECT_ID": "", "VERCEL_TEAM_ID": "", @@ -455,6 +456,7 @@ "CONTAINER_BILLING_HEARTBEAT_SECONDS": "60", "CONTROL_PLANE_IDS": "*", "WORKTREE_CREATION_ENABLED_IDS": "*", + "RUNTIME_ISOLATION_ENABLED": "true", "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 b872ffb8ba..381d5cf742 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: () => {}, }; @@ -1520,7 +1534,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( @@ -1603,7 +1617,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 }, @@ -1621,7 +1635,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 }); @@ -1633,7 +1647,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, @@ -1652,11 +1666,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 }; @@ -1667,7 +1681,7 @@ describe('applySessionAttach', () => { const grant = { ...kilo, targets: { ...kilo.targets } }; const attaching = applySessionAttach( identity, - { kilo: grant }, + { kilo: grant, runtimeIsolation: 'per-session' }, { ...noFs, kiloRuntimes: runtimes, @@ -1687,22 +1701,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`, @@ -1714,12 +1729,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; @@ -1802,7 +1817,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, @@ -2071,7 +2086,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 7d2e3f015c..bb97252455 100644 --- a/services/cloud-agent-next/wrapper/src/control/apply-attach.ts +++ b/services/cloud-agent-next/wrapper/src/control/apply-attach.ts @@ -367,6 +367,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 9b17e08582..fe647ceed4 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 @@ -336,6 +336,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() })); @@ -387,7 +413,7 @@ describe('native-scoped control event failures', () => { handleFailure(failure); - expect(getRuntime).toHaveBeenCalledWith('/root'); + expect(getRuntime).toHaveBeenCalledWith('/root', runtime.runtimeId); expect(onFailure).toHaveBeenCalledWith(failure, runtime); }); @@ -416,7 +442,7 @@ describe('native-scoped control event failures', () => { handleFailure(failure); - expect(getRuntime).toHaveBeenCalledWith('/root'); + expect(getRuntime).toHaveBeenCalledWith('/root', runtime.runtimeId); expect(onFailure).toHaveBeenCalledWith(failure, runtime); }); 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 8409b0dc34..d2f45b8c28 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 @@ -9,7 +9,7 @@ import { ownerDirectoryForSession } from './session-directories.js'; 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) => unknown; }) { const inFlight = new WeakMap>(); @@ -22,7 +22,7 @@ export function createControlEventFailureHandler(); diff --git a/services/cloud-agent-next/wrapper/src/control/control-test-fixtures.ts b/services/cloud-agent-next/wrapper/src/control/control-test-fixtures.ts index 331b8784ba..342de32321 100644 --- a/services/cloud-agent-next/wrapper/src/control/control-test-fixtures.ts +++ b/services/cloud-agent-next/wrapper/src/control/control-test-fixtures.ts @@ -91,6 +91,7 @@ export function createHandlerFixture( const nativeLifetime = new AbortController(); const runtime: WorktreeKiloRuntime | undefined = client ? { + identity: { ...identity }, scopeId: kilo.scopeId, runtimeId: 'native_1', directory: identity.directory, @@ -119,8 +120,10 @@ export function createHandlerFixture( }; }, detach: () => 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 f68b7410ac..499ffaab19 100644 --- a/services/cloud-agent-next/wrapper/src/control/main.ts +++ b/services/cloud-agent-next/wrapper/src/control/main.ts @@ -72,10 +72,10 @@ function main( deps.operations.settleRootPublication(retirement); if (!retirement.reportToService || retirement.result !== 'retired' || !retirement.retirementId) return; - const current = kiloRuntimes.get(retirement.directory); + const current = kiloRuntimes.getRetained?.(retirement.directory, retirement.nativeRuntimeId); if ( !isRetirementReportCurrent( - kiloRuntimes.getEntryRuntimeId?.(retirement.directory), + kiloRuntimes.getEntryRuntimeId?.(retirement.directory, retirement.root), retirement.nativeRuntimeId ) || (current && @@ -102,7 +102,7 @@ function main( if ( !reported && isRetirementReportCurrent( - kiloRuntimes.getEntryRuntimeId?.(retirement.directory), + kiloRuntimes.getEntryRuntimeId?.(retirement.directory, retirement.root), retirement.nativeRuntimeId ) ) @@ -111,7 +111,7 @@ function main( () => { if ( isRetirementReportCurrent( - kiloRuntimes.getEntryRuntimeId?.(retirement.directory), + kiloRuntimes.getEntryRuntimeId?.(retirement.directory, retirement.root), retirement.nativeRuntimeId ) ) @@ -137,7 +137,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, @@ -167,7 +172,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) { @@ -197,7 +202,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({ @@ -314,7 +319,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), }); @@ -493,7 +502,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'); const attempt = beginPublicationFailure( 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 cbe26058a1..49ef30d7bd 100644 --- a/services/cloud-agent-next/wrapper/src/control/operation-registry.ts +++ b/services/cloud-agent-next/wrapper/src/control/operation-registry.ts @@ -9,7 +9,7 @@ import { } from '../../../src/shared/sandbox-control-protocol.js'; import { rejectBeforeAdmission } from './control-handler-result.js'; import { rootForSession } from './session-directories.js'; -import type { WorktreeKiloRuntimes } from './worktree-runtime.js'; +import type { WorktreeKiloRuntime, WorktreeKiloRuntimes } from './worktree-runtime.js'; import type { NativeOperationTarget, NativeRetirement, @@ -24,12 +24,9 @@ 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; + getEntryRuntimeId?(directory: string, root: string): string | undefined; + getRetained(directory: string, runtimeId?: string): WorktreeKiloRuntime | undefined; prepareForNewWork?(directory: string): boolean; retireRuntime( directory: string, @@ -199,8 +196,8 @@ export function createOperationRegistry(deps: OperationRegistryDependencies) { return JSON.stringify([root, nativeRuntimeId]); } - function currentRuntime(directory: string) { - return deps.native.getRetained(directory) ?? deps.native.get(directory); + function currentRuntime(directory: string, nativeRuntimeId: string) { + return deps.native.getRetained(directory, nativeRuntimeId); } function installPhysical( @@ -263,8 +260,15 @@ export function createOperationRegistry(deps: OperationRegistryDependencies) { function clearStaleScopedFailures(): void { for (const [id, failure] of scopedFailures) { - const runtime = currentRuntime(failure.directory); - if (runtime && runtime.runtimeId !== failure.nativeRuntimeId) scopedFailures.delete(id); + const currentId = deps.native.getEntryRuntimeId + ? deps.native.getEntryRuntimeId(failure.directory, failure.root) + : currentRuntime(failure.directory, failure.nativeRuntimeId)?.runtimeId; + if ( + deps.native.getEntryRuntimeId + ? currentId !== failure.nativeRuntimeId + : currentId !== undefined && currentId !== failure.nativeRuntimeId + ) + scopedFailures.delete(id); } removeOrphanedArchivedClaims(); } @@ -610,7 +614,7 @@ export function createOperationRegistry(deps: OperationRegistryDependencies) { clearStaleScopedFailures(); const root = rootForSession(session.kiloSessionId, session.directory); if (!root) return false; - const runtime = currentRuntime(session.directory); + const runtime = deps.native.get(session); for (const failure of scopedFailures.values()) { if ( failure.directory === session.directory && @@ -788,7 +792,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 ab39351675..9e428c14db 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 @@ -138,8 +138,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() {}, }, @@ -386,6 +388,8 @@ describe('createSandboxControlClient', () => { nativeRuntimeRetirement?: boolean; connectionRecovery?: boolean; eventReceipts?: boolean; + runtimeIsolation?: boolean; + runtimeRecovery?: boolean; scopedCleanupResult?: boolean; workingBranches?: boolean; }; @@ -402,6 +406,8 @@ describe('createSandboxControlClient', () => { nativeRuntimeRetirement: true, connectionRecovery: true, eventReceipts: true, + runtimeIsolation: true, + runtimeRecovery: true, scopedCleanupResult: 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 6ca24d1a95..cc44b9575e 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 @@ -524,6 +524,8 @@ export function createSandboxControlClient( nativeRuntimeRetirement: true, connectionRecovery: true, eventReceipts: true, + runtimeIsolation: true, + runtimeRecovery: true, scopedCleanupResult: true, workingBranches: true, }, 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 09190df041..e90d9238c3 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 @@ -11,6 +11,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'; @@ -123,6 +124,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, @@ -148,6 +150,7 @@ function deps( release: () => {}, }), detach: () => true, + retireForRecovery: async () => 'retired', deleteDirectory: async () => {}, getRetained: directory => (directory === runtime.directory ? runtime : undefined), retireRuntime: async (directory, _deadlineAt, target) => @@ -160,7 +163,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) @@ -258,6 +270,9 @@ function fakeTerminalRuntime( rememberAttachedSession: () => {}, detachSession: async () => {}, detachDirectory: async () => {}, + hasActivePty: () => false, + beginRecoveryRetirement: () => {}, + endRecoveryRetirement: () => {}, create: async () => ({ pty }), resize: async () => ({ pty }), close: async () => ({ success: true }), @@ -967,6 +982,7 @@ describe('handleControlRequest', () => { }, }); runtimes.set(directory, { + identity: { ...session }, directory, scopeId: directory, runtimeId: crypto.randomUUID(), @@ -987,6 +1003,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) => { @@ -1006,7 +1023,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, @@ -1262,7 +1284,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, @@ -1369,7 +1391,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({ @@ -1382,7 +1408,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]); @@ -1793,10 +1819,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'); @@ -1805,6 +1833,7 @@ describe('production worktree deletion routes', () => { calls.push('detach'); return true; }, + retireForRecovery: async () => 'retired', deleteDirectory: async dir => { calls.push(`delete:${dir}`); }, @@ -1906,7 +1935,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({ @@ -1928,7 +1957,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, @@ -1946,9 +1975,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; @@ -2036,7 +2070,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', @@ -2095,11 +2131,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); @@ -2131,7 +2168,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([ @@ -2190,12 +2227,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'); @@ -2204,6 +2244,7 @@ describe('production worktree deletion routes', () => { forbidden.push('detach'); return false; }, + retireForRecovery: async () => 'retired', deleteDirectory: async dir => { retirements.push(dir); }, @@ -2230,7 +2271,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([ @@ -3978,7 +4019,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( { @@ -4003,12 +4044,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); @@ -4197,6 +4246,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, @@ -4273,6 +4328,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, @@ -4835,7 +4896,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 () => { @@ -5309,7 +5370,7 @@ describe('refreshHeartbeatPayload', () => { } ); } - const runtime = handlerDeps.kiloRuntimes?.get(session.directory); + const runtime = handlerDeps.kiloRuntimes?.get(session); expect(runtime).toBeDefined(); const refresh = refreshHeartbeatPayload(handlerDeps); try { @@ -5323,7 +5384,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', @@ -5430,16 +5491,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); @@ -5539,16 +5614,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' }, @@ -5642,16 +5723,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 15dc69f886..c6d1d1e741 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 @@ -17,6 +17,7 @@ import { sessionPermissionResolvePayloadSchema, sessionPromptPayloadSchema, sessionQuestionResolvePayloadSchema, + sessionRuntimeRetirePayloadSchema, sessionSyncPayloadSchema, sessionTerminalClosePayloadSchema, sessionTerminalCloseResultSchema, @@ -382,8 +383,15 @@ 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), + ...(input.kiloRuntimes?.getEntryRuntimeId + ? { + getEntryRuntimeId: (directory: string, root: string) => + input.kiloRuntimes?.getEntryRuntimeId?.(directory, root), + } + : {}), + 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) ?? @@ -435,7 +443,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), }); } @@ -522,13 +534,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); @@ -636,6 +653,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': @@ -704,7 +723,7 @@ function sessionKiloRuntime( rootForSession(session.kiloSessionId) !== session.kiloSessionId ) return undefined; - return deps.kiloRuntimes?.get(session.directory); + return deps.kiloRuntimes?.get(session); } function currentRuntimeMatchesTarget( @@ -714,8 +733,8 @@ function currentRuntimeMatchesTarget( ): boolean { if (!target) return false; const runtime = - deps.kiloRuntimes?.getRetained?.(session.directory) ?? - deps.kiloRuntimes?.get(session.directory); + deps.kiloRuntimes?.getRetained?.(session.directory, target.runtimeId) ?? + deps.kiloRuntimes?.get(session); return ( runtime !== undefined && runtime.runtimeId === target.runtimeId && @@ -831,6 +850,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 }; }; @@ -1031,7 +1093,10 @@ async function handleAbort( if (!parsed.success) return fail('protocol_error', 'Invalid payload', false); const scopedCleanupResultGranted = supportsScopedCleanupResult(deps); 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 6c0314a5f4..c3fdde846b 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 @@ -314,11 +314,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: () => {}, }; @@ -425,7 +427,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( @@ -464,7 +466,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 c9f08a094b..f6e965a28b 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; }, @@ -588,12 +590,42 @@ 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 }, + isolation: 'per-session', + 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, @@ -618,7 +650,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); @@ -631,7 +663,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([ @@ -640,7 +672,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' }); @@ -806,7 +838,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, @@ -821,7 +854,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 3de535e405..5ed845074d 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,9 +411,40 @@ 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) { throw new ControlTerminalRuntimeError( 'unauthorized', @@ -412,6 +452,16 @@ export function createControlTerminalRuntime(options: { false ); } + if ( + kiloRuntime.isolation === 'per-session' && + (!kiloRuntime.identity || !sameSession(kiloRuntime.identity, identity)) + ) { + throw new ControlTerminalRuntimeError( + 'unauthorized', + 'Terminal session ownership mismatch', + false + ); + } if ( directoryForSession(identity.kiloSessionId) !== identity.directory || rootForSession(identity.kiloSessionId) !== identity.kiloSessionId @@ -615,6 +665,9 @@ export function createControlTerminalRuntime(options: { return connection; }, + hasActivePty, + beginRecoveryRetirement, + endRecoveryRetirement, shutdown() { if (shutDown) return; shutDown = true; @@ -628,6 +681,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 a3b851773b..afb10b2979 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 9c14ac3f55..68fee14961 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 @@ -108,10 +108,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; @@ -253,7 +256,10 @@ function proveOwnedProcesses( } as unknown as OwnedProcessScope); } -function createRegistry(overrides: Partial[0]> = {}) { +function createRegistry( + overrides: Partial[0]> = {}, + isolation: 'per-session' | 'directory-shared' = 'per-session' +) { const launches: Array[0]> = []; let closes = 0; let unexpectedCloses = 0; @@ -281,11 +287,25 @@ function createRegistry(overrides: Partial, + canRefreshCredentials?: () => boolean + ) { + return registry.attach(identity, kilo, env, canRefreshCredentials, isolation); + }, 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, + isolation + ); try { const runtime = await attachment.ready; attachment.commit(); @@ -295,6 +315,7 @@ function createRegistry(overrides: Partial[0]> = {} +) { + return createRegistry(overrides, 'directory-shared'); +} + function createHandlerDeps(registry: WorktreeKiloRuntimes): HandlerDeps { const terminalRuntime = createControlTerminalRuntime({ controlUrl: 'ws://127.0.0.1:1/sandbox-control/test', wrapperInstanceId: crypto.randomUUID(), - getKiloRuntime: directory => registry.get(directory), + getKiloRuntime: identity => registry.get(identity), }); terminalRuntimes.push(terminalRuntime); return createControlHandlerDeps({ @@ -328,7 +355,7 @@ function createIntegratedRegistry( ) { const context: { handlerDeps?: HandlerDeps } = {}; const settlements: RootRuntimeRetirement[] = []; - const harness = createRegistry({ + const harness = createSharedRegistry({ onRootRetirementStarted: attempt => { context.handlerDeps?.operations.markRootRetirementStarted(attempt); }, @@ -453,7 +480,7 @@ describe('retirement report ownership', () => { const releaseReplacement = Promise.withResolvers(); const retirements: RootRuntimeRetirement[] = []; let launches = 0; - const harness = createRegistry({ + const harness = createSharedRegistry({ startServer: async options => { launches += 1; const server = createKiloStub(); @@ -652,11 +679,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), @@ -666,7 +804,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); @@ -768,7 +906,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]) => @@ -782,7 +920,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 => { @@ -803,6 +941,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 @@ -810,6 +949,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( @@ -823,23 +967,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); }); @@ -887,7 +1034,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', @@ -923,7 +1070,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, @@ -960,11 +1107,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; @@ -1108,13 +1268,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 = { @@ -1139,15 +1300,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( @@ -1190,7 +1351,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( @@ -1203,14 +1364,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); @@ -1224,18 +1385,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'); @@ -1248,9 +1411,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 () => { @@ -1268,7 +1431,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); @@ -1299,36 +1462,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: () => { @@ -1340,23 +1509,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]); @@ -1419,7 +1589,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') @@ -1427,12 +1597,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(); @@ -1441,7 +1612,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); @@ -1461,10 +1632,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 () => { @@ -1483,7 +1653,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'); @@ -1509,7 +1679,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', }); @@ -1555,7 +1726,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); }); @@ -1578,13 +1750,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 () => { @@ -1639,7 +1812,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( @@ -1689,7 +1862,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 { @@ -1734,7 +1907,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'); @@ -1819,7 +1992,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; @@ -1837,7 +2010,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); @@ -1845,25 +2018,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); }); @@ -1876,7 +2050,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([ @@ -1903,7 +2077,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); @@ -1976,7 +2150,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); @@ -2026,7 +2200,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); @@ -2043,7 +2217,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( @@ -2060,7 +2234,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'); @@ -2068,7 +2242,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( @@ -2092,16 +2266,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; @@ -2132,10 +2308,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); @@ -2155,7 +2331,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( @@ -2186,14 +2362,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(); @@ -2210,8 +2388,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; @@ -2260,7 +2438,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; @@ -2439,12 +2617,74 @@ 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('keeps positive process-death proof authoritative after the cleanup deadline', async () => { - const harness = createRegistry(); + const harness = createSharedRegistry(); const directory = path.join(tmpDir, 'deadline'); const runtime = await harness.registry.ensure(directory, auth); @@ -2495,7 +2735,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' ); @@ -2506,7 +2746,7 @@ setInterval(() => {}, 1000); const exited = Promise.withResolvers(); let launches = 0; let observations = 0; - const harness = createRegistry({ + const harness = createSharedRegistry({ startServer: async options => { const server = createKiloStub(); servers.push(server); @@ -2555,7 +2795,7 @@ setInterval(() => {}, 1000); it('awaits the same bounded observation when deleting an unconfirmed runtime', async () => { const absent = Promise.withResolvers(); let observations = 0; - const harness = createRegistry({ + const harness = createSharedRegistry({ startServer: async options => { const server = createKiloStub(); servers.push(server); @@ -2590,7 +2830,7 @@ setInterval(() => {}, 1000); it('deletes homes after concurrent attach observation already proved death', async () => { const absent = Promise.withResolvers(); let observations = 0; - const harness = createRegistry({ + const harness = createSharedRegistry({ startServer: async options => { const server = createKiloStub(); servers.push(server); @@ -2706,7 +2946,7 @@ setInterval(() => {}, 1000); }); it('rejects a stale target without retiring the current runtime', async () => { - const harness = createRegistry(); + const harness = createSharedRegistry(); const { registry } = harness; const directory = path.join(tmpDir, 'stale-target'); const attachment = registry.attach(rootIdentity(directory), auth); @@ -2730,7 +2970,7 @@ setInterval(() => {}, 1000); it('reports immediate root retirement with its captured incarnation even without a deferred intent', async () => { const reports: string[] = []; const attempts: string[] = []; - const harness = createRegistry({ + const harness = createSharedRegistry({ onRootRetirementStarted: retirement => attempts.push(retirement.retirementId), onRootRetirement: retirement => { reports.push(`${retirement.root}:${retirement.nativeRuntimeId}:${retirement.result}`); @@ -2759,7 +2999,7 @@ setInterval(() => {}, 1000); }); it('does not settle a current N2 deferred failure from a stale N1 cleanup callback', async () => { - const harness = createRegistry(); + const harness = createSharedRegistry(); const { registry } = harness; const directory = path.join(tmpDir, 'stale-deferred-target'); const first = registry.attach(rootIdentity(directory, 'first'), auth); @@ -2799,7 +3039,7 @@ setInterval(() => {}, 1000); const disappearedRoots: string[] = []; const server = createKiloStub(); servers.push(server); - const harness = createRegistry({ + const harness = createSharedRegistry({ startServer: async options => { options.onProcessScope?.({ stop: async (_deadlineAt: number) => false, @@ -2834,7 +3074,7 @@ setInterval(() => {}, 1000); }); it('does not retire a healthy survivor when multiple failed roots disappear', async () => { - const harness = createRegistry(); + const harness = createSharedRegistry(); const { registry } = harness; const directory = path.join(tmpDir, 'multiple-failures'); const roots = ['first', 'second', 'healthy'].map(name => rootIdentity(directory, name)); @@ -2877,7 +3117,7 @@ setInterval(() => {}, 1000); it.each([{ order: ['first', 'second'] as const }, { order: ['second', 'first'] as const }])( 'retires either failed root when the other failed root detaches first', async ({ order }) => { - const harness = createRegistry(); + const harness = createSharedRegistry(); const { registry } = harness; const directory = path.join(tmpDir, `failed-order-${order[0]}`); const first = registry.attach(rootIdentity(directory, 'first'), auth); @@ -2916,8 +3156,92 @@ setInterval(() => {}, 1000); }); describe('runtime-to-registry root settlement', () => { + it('observes and replaces only the failed isolated runtime and clears its stale publication failure', async () => { + let absent = false; + let observations = 0; + let launches = 0; + const harness = createRegistry({ + startServer: async options => { + const first = launches++ === 0; + const server = createKiloStub(); + servers.push(server); + options.onProcessScope?.({ + stop: async () => !first, + verify: async () => { + observations += 1; + return absent; + }, + } as unknown as OwnedProcessScope); + return { url: server.url, close: () => {} }; + }, + }); + const directory = path.join(tmpDir, 'isolated-observation'); + const identityA = rootIdentity(directory, 'a'); + const identityB = rootIdentity(directory, 'b'); + const attachmentA = harness.registry.attach(identityA, auth); + const runtimeA = await attachmentA.ready; + attachmentA.commit(); + const attachmentB = harness.registry.attach(identityB, auth); + const runtimeB = await attachmentB.ready; + attachmentB.commit(); + const dependencies = createHandlerDeps(harness.registry); + const failure = { + directory, + root: identityA.kiloSessionId, + nativeRuntimeId: runtimeA.runtimeId, + target: { runtimeId: runtimeA.runtimeId, client: runtimeA.kiloClient }, + reason: 'isolated publication failure', + deadlineAt: Date.now() + 1_000, + }; + const originalCleanup = dependencies.operations.retireRootPublication(failure); + expect(await originalCleanup).toBe('unconfirmed'); + expect( + harness.registry.rootRetirementScope?.(directory, failure.target, identityA.kiloSessionId) + ).toBe('sole'); + expect( + await harness.registry.retireRuntimeIfUnshared?.( + directory, + failure.target, + identityA.kiloSessionId, + Date.now() + 1_000 + ) + ).toBe('unconfirmed'); + dependencies.operations.prune(); + expect(dependencies.operations.retireRootPublication(failure)).toBe(originalCleanup); + expect(harness.registry.get(identityB)).toBe(runtimeB); + expect(runtimeB.signal.aborted).toBe(false); + expect(() => harness.registry.attach(identityA, { ...auth, token: 'unauthorized' })).toThrow( + 'Kilo worktree auth context mismatch' + ); + expect(observations).toBe(0); + absent = true; + expect(() => harness.registry.attach(identityA, auth)).toThrow( + 'Native runtime retirement is unconfirmed' + ); + await waitUntil( + () => harness.registry.getRetained?.(directory, runtimeA.runtimeId) === undefined + ); + const replacement = harness.registry.attach(identityA, auth); + const replacementRuntime = await replacement.ready; + replacement.commit(); + expect(harness.registry.getEntryRuntimeId?.(directory, identityA.kiloSessionId)).toBe( + replacementRuntime.runtimeId + ); + expect(harness.registry.getEntryRuntimeId?.(directory, identityB.kiloSessionId)).toBe( + runtimeB.runtimeId + ); + dependencies.operations.prune(); + // A repeated failure can no longer reuse the old runtime's retained claim. + expect(dependencies.operations.retireRootPublication(failure)).not.toBe(originalCleanup); + expect(dependencies.operations.admission('session.prompt', identityA, undefined).kind).toBe( + 'continue' + ); + expect(harness.registry.get(identityB)).toBe(runtimeB); + expect(runtimeB.signal.aborted).toBe(false); + }); + it('allows A to reattach after a retained completed operation retires its runtime', async () => { - const harness = createRegistry(); + const harness = createSharedRegistry(); const handlerDeps = createHandlerDeps(harness.registry); const identity = rootIdentity(path.join(tmpDir, 'retained-abort-reattach'), 'a'); const prompt = { @@ -3000,7 +3324,7 @@ describe('runtime-to-registry root settlement', () => { }); it('allows A to reattach after a sole active abort replaces its runtime', async () => { - const harness = createRegistry(); + const harness = createSharedRegistry(); const handlerDeps = createHandlerDeps(harness.registry); const identity = rootIdentity(path.join(tmpDir, 'abort-reattach'), 'a'); expect( @@ -3075,7 +3399,7 @@ describe('runtime-to-registry root settlement', () => { }); it('detaches A before returning a failed non-scoped abort so A can reattach', async () => { - const harness = createRegistry({ + const harness = createSharedRegistry({ startServer: async options => { const server = createKiloStub(); servers.push(server); @@ -3161,7 +3485,7 @@ describe('runtime-to-registry root settlement', () => { }); it('keeps B attached while A aborts and reattaches on a shared runtime', async () => { - const harness = createRegistry(); + const harness = createSharedRegistry(); const handlerDeps = createHandlerDeps(harness.registry); handlerDeps.scopedCleanupResult = true; const directory = path.join(tmpDir, 'shared-abort-reattach'); @@ -3392,7 +3716,7 @@ describe('runtime-to-registry root settlement', () => { }); it('does not detach a replacement attachment for a stale native runtime abort', async () => { - const harness = createRegistry(); + const harness = createSharedRegistry(); const handlerDeps = createHandlerDeps(harness.registry); const identity = rootIdentity(path.join(tmpDir, 'stale-abort-attachment'), 'a'); expect( 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 9924ec8e7d..2db2fa3f17 100644 --- a/services/cloud-agent-next/wrapper/src/control/worktree-runtime.ts +++ b/services/cloud-agent-next/wrapper/src/control/worktree-runtime.ts @@ -33,8 +33,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'; @@ -44,6 +46,8 @@ type WorktreeKiloFailure = { }; export type WorktreeKiloRuntime = { + readonly identity?: SessionRequestIdentity; + readonly isolation?: RuntimeIsolation; readonly scopeId: string; readonly runtimeId: string; readonly directory: string; @@ -60,6 +64,7 @@ export type WorktreeKiloAttachment = { release(): void; }; +type RecoveryRetirement = 'retired' | 'absent' | 'acknowledged'; export type RootRuntimeRetirement = { directory: string; root: string; @@ -98,10 +103,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, @@ -125,9 +136,11 @@ export type WorktreeKiloRuntimes = { target: NativeOperationTarget, deadlineAt: number ): Promise; - getRetained?(directory: string): WorktreeKiloRuntime | undefined; - getEntryRuntimeId?(directory: string): string | 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; + getEntryRuntimeId?(directory: string, root?: string): string | undefined; prepareForNewWork?(directory: string): boolean; isHealthy(): boolean; shutdown(): void; @@ -152,6 +165,8 @@ type WorktreeKiloServerHandle = Omit & { }; type RuntimeEntry = { + identity: SessionRequestIdentity; + isolation: RuntimeIsolation; kilo: WorktreeKiloAuth; directory: string; env: Record; @@ -397,6 +412,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(); const deferredRetirements = new Map< @@ -420,24 +437,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 { @@ -562,8 +569,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); if (!root.entry.retiring) options.onRootDisappeared?.({ @@ -607,8 +614,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); } @@ -644,7 +651,9 @@ export function createWorktreeKiloRuntimes(options: { deadlineAt: number, reason = 'Native runtime retirement requested' ): Promise { - const entry = entries.get(directory); + const entry = [...entries.values()].find( + entry => entry.directory === directory && runtimeTargetMatches(entry, target) + ); if (!entry || !runtimeTargetMatches(entry, target)) { const intent = { directory, @@ -695,7 +704,9 @@ export function createWorktreeKiloRuntimes(options: { target: NativeOperationTarget, retiringRoot: string ): RootRetirementScope { - const entry = entries.get(directory); + const entry = [...entries.values()].find( + entry => entry.directory === directory && runtimeTargetMatches(entry, target) + ); if (!entry || entry.retiring || !runtimeTargetMatches(entry, target)) return 'stale'; const failedRoot = [...entry.roots].find(root => root.identity.kiloSessionId === retiringRoot); if (!failedRoot || !(failedRoot.attached || failedRoot.pending.size > 0)) return 'stale'; @@ -734,14 +745,14 @@ export function createWorktreeKiloRuntimes(options: { } if (Date.now() >= deadlineAt) return false; if ( - entries.get(entry.directory) !== entry || + entries.get(entryKey(entry.identity, entry.isolation)) !== entry || entry.runtimeId !== runtimeId || entry.runtime !== runtime || entry.retirementResult !== 'unconfirmed' || !entry.abort.signal.aborted ) return false; - entries.delete(entry.directory); + entries.delete(entryKey(entry.identity, entry.isolation)); if (directoriesByScope.get(entry.kilo.scopeId) === entry.directory) directoriesByScope.delete(entry.kilo.scopeId); failedDirectories.delete(entry.directory); @@ -762,7 +773,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); @@ -778,6 +789,7 @@ export function createWorktreeKiloRuntimes(options: { }).then(result => { if (result === 'unconfirmed') failedDirectories.add(entry.directory); options.onUnexpectedClose({ + identity: entry.identity, retirementId: crypto.randomUUID(), directory: entry.directory, reason, @@ -875,6 +887,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; @@ -899,7 +913,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, @@ -1018,11 +1032,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); } @@ -1030,8 +1057,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); let entry = previous?.retiring ? undefined : previous; let cleanupRequired = false; if ( @@ -1111,7 +1139,7 @@ export function createWorktreeKiloRuntimes(options: { if (!entry) { if (previous) settleEntryDeferredRetirements(previous, undefined, 'stale'); const homeId = createHash('sha256') - .update(kilo.scopeId) + .update(isolation === 'per-session' ? key : kilo.scopeId) .update('\0') .update(directory) .digest('hex'); @@ -1120,6 +1148,8 @@ export function createWorktreeKiloRuntimes(options: { homeId ); entry = { + identity: { ...identity }, + isolation, kilo: { ...kilo, targets: { ...kilo.targets } }, directory, env: buildWorktreeKiloEnvironment( @@ -1137,10 +1167,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 = { @@ -1150,7 +1180,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); } @@ -1192,20 +1222,90 @@ 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; let quiescent = await withTimeoutAndAbort(retire(entry), { timeoutMs: KILO_STARTUP_TIMEOUT_MS, timeoutMessage: 'Kilo worktree retirement timed out', abortMessage: 'Kilo worktree retirement cancelled', }); if (quiescent !== 'retired' && entry.retirementResult === 'unconfirmed') { - if ((await observeRetained(entry)) || entries.get(directory) !== entry) + if ( + (await observeRetained(entry)) || + entries.get(entryKey(entry.identity, entry.isolation)) !== entry + ) quiescent = 'retired'; } if (quiescent !== 'retired') throw new Error('Native worktree cleanup is unconfirmed'); @@ -1216,7 +1316,12 @@ 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); }, @@ -1225,7 +1330,9 @@ export function createWorktreeKiloRuntimes(options: { }, rootRetirementScope, 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 || @@ -1237,29 +1344,63 @@ 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; }, - getEntryRuntimeId(directory) { - return entries.get(directory)?.runtimeId; + getEntryRuntimeId(directory, root) { + const entry = + [...entries.values()].find( + entry => + entry.directory === directory && + entry.isolation === 'per-session' && + (root === undefined || entry.identity.kiloSessionId === root) + ) ?? entries.get(directory); + return entry?.runtimeId; }, - 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() { 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 3f77b361ff..84a9af6042 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 = { @@ -3302,3 +3303,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); + }); +}); 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'], + }, +});