diff --git a/services/cloud-agent-next/src/session/agent-runtime-runtime-proxy-grant.test.ts b/services/cloud-agent-next/src/session/agent-runtime-runtime-proxy-grant.test.ts new file mode 100644 index 0000000000..120bf4edbc --- /dev/null +++ b/services/cloud-agent-next/src/session/agent-runtime-runtime-proxy-grant.test.ts @@ -0,0 +1,258 @@ +import { describe, expect, it, vi } from 'vitest'; +import jwt from 'jsonwebtoken'; +import type { RuntimeAuthorization } from '@kilocode/worker-utils/runtime-authorization-contract'; +import type { AgentSandbox } from '../agent-sandbox/protocol.js'; +import type { Env } from '../types.js'; +import type { FencedWrapperDispatchRequest } from '../execution/types.js'; +import type { SessionMetadata } from '../persistence/session-metadata.js'; +import { RUNTIME_AUTHORIZATION_KEY } from './runtime-authorization-persistence.js'; +import { createAgentRuntime } from './agent-runtime.js'; +import { getWrapperLease, getWrapperRuntimeState } from './wrapper-runtime-state.js'; + +vi.mock('@cloudflare/sandbox', () => ({ + Sandbox: class Sandbox {}, + getSandbox: vi.fn(), + ContainerProxy: class ContainerProxy {}, +})); + +vi.mock('@cloudflare/containers', () => ({})); + +vi.mock('cloudflare:workers', () => ({ + DurableObject: class DurableObject { + ctx: unknown; + env: unknown; + constructor(ctx: unknown, env: unknown) { + this.ctx = ctx; + this.env = env; + } + }, +})); + +vi.mock('../logger.js', () => { + const logger = { + setTags: vi.fn(), + info: vi.fn(), + warn: vi.fn(), + error: vi.fn(), + debug: vi.fn(), + withFields: vi.fn(), + }; + logger.withFields.mockReturnValue(logger); + return { + logger, + withLogTags: async (_tags: unknown, fn: () => Promise) => fn(), + WithLogTags: () => (_target: unknown, _propertyKey: string, descriptor: PropertyDescriptor) => + descriptor, + }; +}); + +vi.mock('drizzle-orm/durable-sqlite', () => ({ + drizzle: vi.fn(() => ({})), +})); + +vi.mock('drizzle-orm/durable-sqlite/migrator', () => ({ + migrate: vi.fn(), +})); + +vi.mock('../../drizzle/migrations', () => ({ + default: { journal: {}, migrations: {} }, +})); + +vi.mock('./queries/index.js', () => ({ + createExecutionQueries: vi.fn(() => ({})), + createEventQueries: vi.fn(() => ({})), + createLeaseQueries: vi.fn(() => ({})), +})); + +vi.mock('../websocket/stream.js', () => ({ + createStreamHandler: vi.fn(), + getConnectedStreamClientCount: vi.fn(() => 0), +})); + +vi.mock('@kilocode/db/client', () => ({ + getWorkerDb: () => ({ + select: () => ({ + from: () => ({ + where: () => ({ + limit: async () => [{ api_token_pepper: null, blocked_reason: null }], + }), + }), + }), + }), +})); + +const { CloudAgentSession } = await import('../persistence/CloudAgentSession.js'); + +const secret = 'test-secret'; +const authorizationId = '11111111-1111-4111-8111-111111111111'; + +type MemoryStorage = Pick & DurableObjectStorage; + +function createMemoryStorage(initialEntries?: Array<[string, unknown]>): MemoryStorage { + const store = new Map(initialEntries ?? []); + return { + async get(key: string) { + return store.get(key) as T | undefined; + }, + async put(key: string, value: unknown) { + store.set(key, value); + }, + async delete(keys: string | string[]) { + let deleted = false; + for (const key of Array.isArray(keys) ? keys : [keys]) { + deleted = store.delete(key) || deleted; + } + return deleted; + }, + } as MemoryStorage; +} + +function authorization(): RuntimeAuthorization { + const issuedAt = new Date(); + return { + version: 1, + id: authorizationId, + resourceKind: 'cloud-agent-next', + resourceId: 'agent_runtime', + userId: 'user_runtime', + authorizationUserId: 'user_runtime', + organizationId: 'org_runtime', + issuedAt: issuedAt.toISOString(), + delegationExpiresAt: new Date(issuedAt.getTime() + 24 * 60 * 60_000).toISOString(), + state: 'active', + bindings: { + userPepperDigest: 'a'.repeat(64), + authorizationPepperDigest: 'b'.repeat(64), + userMembershipId: 'membership_1', + authorizationUserMembershipId: 'membership_1', + }, + source: { admissionSource: 'user' }, + }; +} + +function metadata(token: string): SessionMetadata { + return { + metadataSchemaVersion: 2, + identity: { + sessionId: 'agent_runtime', + userId: 'user_runtime', + orgId: 'org_runtime', + }, + auth: { + kiloSessionId: 'kilo_runtime', + kilocodeToken: token, + }, + lifecycle: { + version: 1, + timestamp: 1, + }, + workspace: { + sandboxId: 'ses-abcdef', + sandboxProvider: 'cloudflare', + workspacePath: '/workspace/runtime', + sessionHome: '/home/agent_runtime', + branchName: 'main', + }, + }; +} + +describe('AgentRuntime restart credential proxy grant', () => { + it('issues a runtime credential proxy grant after allocating a replacement wrapper', async () => { + const token = jwt.sign( + { + runtimeAuthorization: { id: authorizationId }, + exp: Math.floor(Date.now() / 1000) + 3600, + }, + secret, + { algorithm: 'HS256' } + ); + const storage = createMemoryStorage([ + ['metadata', metadata(token)], + [RUNTIME_AUTHORIZATION_KEY, authorization()], + [ + 'wrapper_runtime_state', + { + wrapperGeneration: 7, + wrapperConnectionId: 'conn_stale', + wrapperRunId: 'wr_stale', + }, + ], + [ + 'wrapper_lease', + { + state: 'owns_wrapper', + nextInstanceGeneration: 2, + instance: { instanceId: 'instance_old', instanceGeneration: 1 }, + }, + ], + ]); + const deliveredPlans: FencedWrapperDispatchRequest[] = []; + const runtime = createAgentRuntime({ + storage, + env: { NEXTAUTH_SECRET: secret } as Env, + getMetadata: async () => metadata(token), + getOrchestratorOverride: () => ({ + execute: async (plan: FencedWrapperDispatchRequest) => { + deliveredPlans.push(plan); + return { kiloSessionId: 'kilo_runtime' }; + }, + }), + getSessionIdForLogs: () => 'agent_runtime', + sendToWrapper: () => false, + createAgentSandbox: () => + ({ + discoverSessionWrappers: async () => ({ status: 'absent' }), + }) as unknown as AgentSandbox, + }); + + await expect( + runtime.send({ + scope: { sessionId: 'agent_runtime', userId: 'user_runtime' }, + turn: { + type: 'prompt', + messageId: 'msg_018f1e2d3c4bRuntimeGrant01', + prompt: 'Continue after wrapper restart', + }, + agent: { mode: 'code', model: 'runtime-model' }, + workspace: { + sandboxId: 'ses-abcdef', + metadata: metadata(token), + }, + wrapper: { kiloSessionId: 'kilo_runtime' }, + }) + ).resolves.toMatchObject({ success: true, outcome: 'accepted' }); + + const runtimeState = await getWrapperRuntimeState(storage); + const physicalLease = await getWrapperLease(storage); + if (physicalLease.state !== 'owns_wrapper') { + throw new Error('expected owned wrapper after restart allocation'); + } + expect(runtimeState.wrapperGeneration).toBe(physicalLease.instance.instanceGeneration); + const fence = deliveredPlans[0]?.wrapper.fence; + expect(fence).toEqual({ + wrapperRunId: runtimeState.wrapperRunId, + wrapperGeneration: physicalLease.instance.instanceGeneration, + wrapperConnectionId: runtimeState.wrapperConnectionId, + }); + + const session = new CloudAgentSession( + { + id: { name: 'user_runtime:agent_runtime' }, + storage: { + ...storage, + sql: {}, + getAlarm: async () => null, + setAlarm: async () => undefined, + list: async () => new Map(), + }, + blockConcurrencyWhile: async () => undefined, + getWebSockets: () => [], + } as never, + { NEXTAUTH_SECRET: secret } as never + ); + if (!fence) throw new Error('expected dispatch fence'); + await expect(session.issueRuntimeCredentialProxyGrant(fence)).resolves.toEqual( + expect.any(String) + ); + }); +}); diff --git a/services/cloud-agent-next/src/session/agent-runtime.test.ts b/services/cloud-agent-next/src/session/agent-runtime.test.ts index 11e20c5548..fbec7cd1db 100644 --- a/services/cloud-agent-next/src/session/agent-runtime.test.ts +++ b/services/cloud-agent-next/src/session/agent-runtime.test.ts @@ -764,7 +764,13 @@ describe('AgentRuntime', () => { nextInstanceGeneration: 3, instance: { instanceGeneration: 2 }, }); - await expect(getWrapperRuntimeState(storage)).resolves.not.toMatchObject({ + const runtimeState = await getWrapperRuntimeState(storage); + const physicalLease = await getWrapperLease(storage); + if (physicalLease.state !== 'owns_wrapper') { + throw new Error('expected owned wrapper after restart allocation'); + } + expect(runtimeState.wrapperGeneration).toBe(physicalLease.instance.instanceGeneration); + expect(runtimeState).not.toMatchObject({ wrapperRunId: 'wr_stale', wrapperConnectionId: 'conn_stale', }); @@ -796,7 +802,7 @@ describe('AgentRuntime', () => { target: { kind: 'instance', instance: { instanceGeneration: 1 } }, reason: 'startup-failed', }); - await expect(getWrapperRuntimeState(storage)).resolves.toEqual({ wrapperGeneration: 2 }); + await expect(getWrapperRuntimeState(storage)).resolves.toEqual({ wrapperGeneration: 1 }); }); it('stores cleanup obligation when a newly leased wrapper readies but its initial dispatch fails', async () => { diff --git a/services/cloud-agent-next/src/session/agent-runtime.ts b/services/cloud-agent-next/src/session/agent-runtime.ts index 9ab6afe235..43ba1436bd 100644 --- a/services/cloud-agent-next/src/session/agent-runtime.ts +++ b/services/cloud-agent-next/src/session/agent-runtime.ts @@ -357,10 +357,13 @@ export function createAgentRuntime(dependencies: AgentRuntimeDependencies): Agen (allocatedPhysicalInstance || requiresFreshRunFence) && (previousRuntimeState.wrapperConnectionId || previousRuntimeState.wrapperRunId) ) { - await clearWrapperRuntimeIdentity(storage, {}, { incrementGeneration: true }); + await clearWrapperRuntimeIdentity(storage); } - const { state: wrapperRuntimeState, allocatedNewIdentity } = - await allocateWrapperRuntimeState(storage); + const { state: wrapperRuntimeState, allocatedNewIdentity } = await allocateWrapperRuntimeState( + storage, + Date.now(), + leasedInstance.instanceGeneration + ); logger .withFields({ sessionId, diff --git a/services/cloud-agent-next/src/session/wrapper-runtime-state.test.ts b/services/cloud-agent-next/src/session/wrapper-runtime-state.test.ts index 8bb93b513a..cb80be7969 100644 --- a/services/cloud-agent-next/src/session/wrapper-runtime-state.test.ts +++ b/services/cloud-agent-next/src/session/wrapper-runtime-state.test.ts @@ -394,6 +394,17 @@ describe('WrapperLease', () => { expect(nextSandboxRecoveryDeadline(retrying)).toBe(3_000); }); + it('adopts an explicit wrapper generation instead of incrementing the runtime counter', async () => { + const storage = createMemoryStorage(); + await storage.put('wrapper_runtime_state', { wrapperGeneration: 7 }); + + const { state, allocatedNewIdentity } = await allocateWrapperRuntimeState(storage, 1_000, 2); + + expect(allocatedNewIdentity).toBe(true); + expect(state.wrapperGeneration).toBe(2); + await expect(getWrapperRuntimeState(storage)).resolves.toMatchObject({ wrapperGeneration: 2 }); + }); + it('marks a newly allocated wrapper run as maintaining its message index', async () => { const storage = createMemoryStorage(); diff --git a/services/cloud-agent-next/src/session/wrapper-runtime-state.ts b/services/cloud-agent-next/src/session/wrapper-runtime-state.ts index 245a861ff2..beee621ddb 100644 --- a/services/cloud-agent-next/src/session/wrapper-runtime-state.ts +++ b/services/cloud-agent-next/src/session/wrapper-runtime-state.ts @@ -599,7 +599,8 @@ export type AllocatedWrapperRuntimeState = { export async function allocateWrapperRuntimeState( storage: DurableObjectStorage, - now = Date.now() + now = Date.now(), + wrapperGeneration?: number ): Promise { const current = await getWrapperRuntimeState(storage); if (isActiveWrapperRuntimeState(current)) { @@ -617,7 +618,7 @@ export async function allocateWrapperRuntimeState( // Obsolete grace cleanup is best-effort; fresh fenced work must proceed. } const next = { - wrapperGeneration: current.wrapperGeneration + 1, + wrapperGeneration: wrapperGeneration ?? current.wrapperGeneration + 1, wrapperConnectionId: crypto.randomUUID(), wrapperRunId: `wr_${crypto.randomUUID().replace(/-/g, '')}`, messageIndexVersion: WRAPPER_RUN_MESSAGE_INDEX_VERSION, @@ -666,14 +667,10 @@ export async function clearAllocatedWrapperRuntimeState( ): Promise { if (!allocated.wrapperConnectionId) return; - await clearWrapperRuntimeIdentity( - storage, - { - wrapperGeneration: allocated.wrapperGeneration, - wrapperConnectionId: allocated.wrapperConnectionId, - }, - { incrementGeneration: true } - ); + await clearWrapperRuntimeIdentity(storage, { + wrapperGeneration: allocated.wrapperGeneration, + wrapperConnectionId: allocated.wrapperConnectionId, + }); } export async function isCurrentWrapperConnection( diff --git a/services/cloud-agent-next/test/integration/session/execute-directly-failure.test.ts b/services/cloud-agent-next/test/integration/session/execute-directly-failure.test.ts index 8fd99da7db..b1f9322ddc 100644 --- a/services/cloud-agent-next/test/integration/session/execute-directly-failure.test.ts +++ b/services/cloud-agent-next/test/integration/session/execute-directly-failure.test.ts @@ -421,7 +421,7 @@ describe('executeDirectly failure handling', () => { expect(result.pendingAfterAlarm[0]?.executionId).toBeUndefined(); expect(result.pendingAfterAlarm[0]?.lastFlushError).toBe('Sandbox connect failed'); expect(result.executionsAfterFirstAlarm).toEqual([]); - expect(result.wrapperRuntimeState.wrapperGeneration).toBe(2); + expect(result.wrapperRuntimeState.wrapperGeneration).toBe(1); expect(result.wrapperRuntimeState.wrapperConnectionId).toBeUndefined(); expect(result.wrapperLeaseAfterFailure).toMatchObject({ state: 'stop_needed',