From 033002e172e89954621790797f05b70c8fe36bf7 Mon Sep 17 00:00:00 2001 From: Evgeny Shurakov Date: Mon, 6 Jul 2026 14:28:11 +0200 Subject: [PATCH 1/2] feat(cloud-agent-next): consume Kilo opaque session capabilities --- .../cloud-agent-next/MessageBubble.test.ts | 46 +++ .../cloud-agent-next/MessageBubble.tsx | 6 +- services/cloud-agent-next/.dev.vars.example | 4 + services/cloud-agent-next/src/auth.test.ts | 122 +++++++- services/cloud-agent-next/src/auth.ts | 116 +++++++ .../src/kilo/kilo-targets.test.ts | 107 +++++++ .../cloud-agent-next/src/kilo/kilo-targets.ts | 88 ++++++ .../src/model-validation.test.ts | 3 +- .../cloud-agent-next/src/model-validation.ts | 7 +- .../cloud-agent-next/src/persistence/types.ts | 2 - .../src/sandbox-outbound.test.ts | 193 +++++++++++- .../cloud-agent-next/src/sandbox-outbound.ts | 217 ++++++++++++- services/cloud-agent-next/src/server.test.ts | 257 ++++++++++++++- services/cloud-agent-next/src/server.ts | 114 ++++++- .../services/git-token-service-client.test.ts | 2 + .../src/services/git-token-service-client.ts | 52 ++++ .../src/session-service.test.ts | 294 +++++++++++++++++- .../cloud-agent-next/src/session-service.ts | 183 +++++++++-- .../src/session/agent-runtime.test.ts | 27 +- services/cloud-agent-next/src/types.ts | 51 +++ .../worker-configuration.d.ts | 8 +- services/cloud-agent-next/wrangler.jsonc | 2 +- .../wrapper/src/log-uploader.test.ts | 140 +++++++++ .../wrapper/src/log-uploader.ts | 26 +- services/cloud-agent-next/wrapper/src/main.ts | 3 +- .../wrapper/src/server.test.ts | 40 +++ .../cloud-agent-next/wrapper/src/server.ts | 19 +- .../wrapper/src/session-bootstrap.test.ts | 25 +- .../wrapper/src/session-bootstrap.ts | 10 +- 29 files changed, 2062 insertions(+), 102 deletions(-) create mode 100644 apps/web/src/components/cloud-agent-next/MessageBubble.test.ts create mode 100644 services/cloud-agent-next/src/kilo/kilo-targets.test.ts create mode 100644 services/cloud-agent-next/src/kilo/kilo-targets.ts create mode 100644 services/cloud-agent-next/wrapper/src/log-uploader.test.ts diff --git a/apps/web/src/components/cloud-agent-next/MessageBubble.test.ts b/apps/web/src/components/cloud-agent-next/MessageBubble.test.ts new file mode 100644 index 0000000000..0b8484b8bb --- /dev/null +++ b/apps/web/src/components/cloud-agent-next/MessageBubble.test.ts @@ -0,0 +1,46 @@ +import React from 'react'; +import { renderToStaticMarkup } from 'react-dom/server'; +import type { AssistantMessage } from '@/types/opencode.gen'; +import type { StoredMessage } from './types'; + +jest.mock('./PartRenderer', () => ({ PartRenderer: () => null })); +jest.mock('@/components/shared/TimeAgo', () => ({ TimeAgo: () => null })); + +import { MessageBubble } from './MessageBubble'; + +describe('MessageBubble', () => { + it('renders a sanitized string assistant error', () => { + const info: AssistantMessage = { + id: 'msg-1', + sessionID: 'ses-1', + role: 'assistant', + time: { created: 1, completed: 2 }, + parentID: 'msg-parent', + modelID: 'test-model', + providerID: 'test-provider', + mode: 'code', + agent: 'test-agent', + path: { cwd: '/', root: '/' }, + cost: 0, + tokens: { + input: 0, + output: 0, + reasoning: 0, + cache: { read: 0, write: 0 }, + }, + }; + Object.defineProperty(info, 'error', { + value: 'Assistant request was rate limited', + enumerable: true, + }); + const message: StoredMessage = { + info, + parts: [], + }; + + const html = renderToStaticMarkup(React.createElement(MessageBubble, { message })); + + expect(html).toContain('Assistant request was rate limited'); + expect(html).toContain('Failed'); + }); +}); diff --git a/apps/web/src/components/cloud-agent-next/MessageBubble.tsx b/apps/web/src/components/cloud-agent-next/MessageBubble.tsx index 8589c43320..3cd4f7508c 100644 --- a/apps/web/src/components/cloud-agent-next/MessageBubble.tsx +++ b/apps/web/src/components/cloud-agent-next/MessageBubble.tsx @@ -1,6 +1,6 @@ 'use client'; -import { useCallback } from 'react'; +import React, { useCallback } from 'react'; import { Scissors, Image, FileText, AlertCircle, Clock } from 'lucide-react'; import { TimeAgo } from '@/components/shared/TimeAgo'; import type { AssistantMessage } from '@/types/opencode.gen'; @@ -140,7 +140,9 @@ function getAssistantTextContent(parts: Part[]): string { /** * Extract a human-readable error message from an AssistantMessage error field. */ -function getAssistantErrorMessage(error: NonNullable): string { +function getAssistantErrorMessage(error: NonNullable | string): string { + if (typeof error === 'string') return error; + if ('data' in error && 'message' in error.data && typeof error.data.message === 'string') { return error.data.message; } diff --git a/services/cloud-agent-next/.dev.vars.example b/services/cloud-agent-next/.dev.vars.example index 58f7512a65..685166fac6 100644 --- a/services/cloud-agent-next/.dev.vars.example +++ b/services/cloud-agent-next/.dev.vars.example @@ -88,6 +88,10 @@ R2_SECRET_ACCESS_KEY="" # @url cloudflare-session-ingest KILO_SESSION_INGEST_URL=http://host.docker.internal:8800 +# Local debugging only. Set to 1 to log full URLs rejected during Kilo capability +# redemption. Those URLs can contain tokens; never enable this in a deployed Worker. +LOG_REJECTED_KILO_URLS=0 + # Local dev worker origins allowed to connect to /stream # @url nextjs,cloud-agent-next WS_ALLOWED_ORIGINS=http://localhost:3000,http://localhost:8794 diff --git a/services/cloud-agent-next/src/auth.test.ts b/services/cloud-agent-next/src/auth.test.ts index 7afd749988..a7319d329a 100644 --- a/services/cloud-agent-next/src/auth.test.ts +++ b/services/cloud-agent-next/src/auth.test.ts @@ -1,6 +1,11 @@ import { describe, expect, it } from 'vitest'; import jwt from 'jsonwebtoken'; -import { validateStreamTicket } from './auth.js'; +import { + mintWrapperDispatchTicket, + validateStreamTicket, + validateWrapperDispatchTicket, + type WrapperDispatchTicketClaims, +} from './auth.js'; const secret = 'test-secret'; @@ -50,3 +55,118 @@ describe('validateStreamTicket', () => { }); }); }); + +const wrapperDispatchTicketClaims: WrapperDispatchTicketClaims = { + type: 'wrapper_dispatch_ticket', + userId: 'user-1', + cloudAgentSessionId: 'session-1', + kiloSessionId: 'kilo-session-1', + wrapperRunId: 'run-1', + wrapperGeneration: 2, + wrapperConnectionId: 'connection-1', +}; + +describe('validateWrapperDispatchTicket', () => { + it('returns a configuration error when NEXTAUTH_SECRET is missing', async () => { + await expect(validateWrapperDispatchTicket('Bearer ticket', null)).resolves.toEqual({ + success: false, + error: 'NEXTAUTH_SECRET is not configured on the worker', + }); + }); + + it('returns an error when the Authorization header is missing or malformed', async () => { + await expect(validateWrapperDispatchTicket(null, secret)).resolves.toEqual({ + success: false, + error: 'Missing or malformed Authorization header', + }); + await expect(validateWrapperDispatchTicket('NotBearer ticket', secret)).resolves.toEqual({ + success: false, + error: 'Missing or malformed Authorization header', + }); + }); + + it('round-trips a minted ticket, preserving claims exactly', async () => { + const ticket = mintWrapperDispatchTicket(wrapperDispatchTicketClaims, secret); + + await expect(validateWrapperDispatchTicket(`Bearer ${ticket}`, secret)).resolves.toEqual({ + success: true, + claims: wrapperDispatchTicketClaims, + }); + }); + + it('rejects a raw Kilo JWT / stream ticket by its type discriminant', async () => { + const streamTicket = jwt.sign( + { type: 'stream_ticket', userId: 'user-1', cloudAgentSessionId: 'session-1' }, + secret, + { algorithm: 'HS256', expiresIn: '1 minute' } + ); + const kiloJwt = jwt.sign({ kiloUserId: 'user-1' }, secret, { + algorithm: 'HS256', + expiresIn: '1 minute', + }); + + await expect(validateWrapperDispatchTicket(`Bearer ${streamTicket}`, secret)).resolves.toEqual({ + success: false, + error: 'Invalid ticket type', + }); + await expect(validateWrapperDispatchTicket(`Bearer ${kiloJwt}`, secret)).resolves.toEqual({ + success: false, + error: 'Invalid ticket type', + }); + }); + + it('accepts a legacy raw Kilo JWT so wrapper processes bound before ticket support shipped keep working', async () => { + const legacyToken = jwt.sign({ version: 3, kiloUserId: 'user-1' }, secret, { + algorithm: 'HS256', + expiresIn: '1 minute', + }); + + await expect(validateWrapperDispatchTicket(`Bearer ${legacyToken}`, secret)).resolves.toEqual({ + success: true, + claims: { type: 'legacy_kilo_token', userId: 'user-1' }, + }); + }); + + it('rejects an audience-scoped Kilo JWT on the legacy path', async () => { + const audienceScopedToken = jwt.sign({ version: 3, kiloUserId: 'user-1' }, secret, { + algorithm: 'HS256', + expiresIn: '1 minute', + audience: 'git-token-service', + }); + + await expect( + validateWrapperDispatchTicket(`Bearer ${audienceScopedToken}`, secret) + ).resolves.toEqual({ + success: false, + error: 'Invalid ticket type', + }); + }); + + it('rejects an expired ticket', async () => { + const ticket = jwt.sign(wrapperDispatchTicketClaims, secret, { + algorithm: 'HS256', + expiresIn: -1, + }); + + await expect(validateWrapperDispatchTicket(`Bearer ${ticket}`, secret)).resolves.toEqual({ + success: false, + error: 'Ticket expired', + }); + }); + + it('rejects a malformed ticket', async () => { + await expect(validateWrapperDispatchTicket('Bearer not-a-jwt', secret)).resolves.toEqual({ + success: false, + error: 'Invalid ticket signature', + }); + }); + + it('rejects a ticket signed with the wrong secret', async () => { + const ticket = mintWrapperDispatchTicket(wrapperDispatchTicketClaims, 'other-secret'); + + await expect(validateWrapperDispatchTicket(`Bearer ${ticket}`, secret)).resolves.toEqual({ + success: false, + error: 'Invalid ticket signature', + }); + }); +}); diff --git a/services/cloud-agent-next/src/auth.ts b/services/cloud-agent-next/src/auth.ts index 5f2b77034e..ab1b7687c3 100644 --- a/services/cloud-agent-next/src/auth.ts +++ b/services/cloud-agent-next/src/auth.ts @@ -13,6 +13,32 @@ type StreamTicketPayload = { nonce?: string; }; +export type WrapperDispatchTicketClaims = { + type: 'wrapper_dispatch_ticket'; + userId: string; + cloudAgentSessionId: string; + kiloSessionId: string; + wrapperRunId: string; + wrapperGeneration: number; + wrapperConnectionId: string; +}; + +/** + * A pre-migration raw Kilo user JWT presented as a workerAuthToken. Accepted + * only so wrapper processes already bound before wrapper dispatch tickets + * shipped can keep talking to the Worker until their next session/ready + * dispatch mints a real ticket. Carries no fence claims — callers must treat + * it as exempt from ticketClaimsMismatchRequestFence. + */ +export type LegacyKiloTokenClaims = { + type: 'legacy_kilo_token'; + userId: string; +}; + +export type WrapperAuthClaims = WrapperDispatchTicketClaims | LegacyKiloTokenClaims; + +const WRAPPER_DISPATCH_TICKET_MAX_LIFETIME_SECONDS = 4 * 60 * 60; + export type SecretBinding = string | { get(): Promise }; export async function resolveSecret( @@ -88,3 +114,93 @@ export function validateStreamTicket( return { success: false, error: 'Ticket validation failed' }; } } + +export function mintWrapperDispatchTicket( + claims: WrapperDispatchTicketClaims, + secret: string +): string { + return jwt.sign(claims, secret, { + algorithm: 'HS256', + expiresIn: WRAPPER_DISPATCH_TICKET_MAX_LIFETIME_SECONDS, + }); +} + +function isWrapperDispatchTicketClaims(payload: unknown): payload is WrapperDispatchTicketClaims { + if (typeof payload !== 'object' || payload === null) return false; + const claims = payload as Record; + return ( + claims.type === 'wrapper_dispatch_ticket' && + typeof claims.userId === 'string' && + typeof claims.cloudAgentSessionId === 'string' && + typeof claims.kiloSessionId === 'string' && + typeof claims.wrapperRunId === 'string' && + typeof claims.wrapperGeneration === 'number' && + typeof claims.wrapperConnectionId === 'string' + ); +} + +/** + * Re-verifies the token through the canonical verifyKiloToken so the legacy + * path enforces the same version/audience/shape guards as validateKiloToken — + * in particular, audience-scoped tokens (e.g. internal service tokens) must + * not be accepted as wrapper auth. + */ +async function legacyKiloTokenClaims( + ticket: string, + secret: string +): Promise { + try { + const payload = await verifyKiloToken(ticket, secret); + return { type: 'legacy_kilo_token', userId: payload.kiloUserId }; + } catch { + return undefined; + } +} + +export async function validateWrapperDispatchTicket( + authHeader: string | null, + secret: string | null | undefined +): Promise<{ success: true; claims: WrapperAuthClaims } | { success: false; error: string }> { + if (!secret) { + return { success: false, error: 'NEXTAUTH_SECRET is not configured on the worker' }; + } + + const ticket = extractBearerToken(authHeader); + if (!ticket) { + return { success: false, error: 'Missing or malformed Authorization header' }; + } + + let payload: string | jwt.JwtPayload; + try { + payload = jwt.verify(ticket, secret, { algorithms: ['HS256'] }); + } catch (error) { + if (error instanceof jwt.TokenExpiredError) { + return { success: false, error: 'Ticket expired' }; + } + if (error instanceof jwt.JsonWebTokenError) { + return { success: false, error: 'Invalid ticket signature' }; + } + return { success: false, error: 'Ticket validation failed' }; + } + + if (!isWrapperDispatchTicketClaims(payload)) { + const legacyClaims = await legacyKiloTokenClaims(ticket, secret); + if (legacyClaims) { + return { success: true, claims: legacyClaims }; + } + return { success: false, error: 'Invalid ticket type' }; + } + + return { + success: true, + claims: { + type: payload.type, + userId: payload.userId, + cloudAgentSessionId: payload.cloudAgentSessionId, + kiloSessionId: payload.kiloSessionId, + wrapperRunId: payload.wrapperRunId, + wrapperGeneration: payload.wrapperGeneration, + wrapperConnectionId: payload.wrapperConnectionId, + }, + }; +} diff --git a/services/cloud-agent-next/src/kilo/kilo-targets.test.ts b/services/cloud-agent-next/src/kilo/kilo-targets.test.ts new file mode 100644 index 0000000000..f948597102 --- /dev/null +++ b/services/cloud-agent-next/src/kilo/kilo-targets.test.ts @@ -0,0 +1,107 @@ +import { describe, expect, it } from 'vitest'; +import { + backendUrlForSandbox, + deriveKiloSandboxTargets, + providerBaseUrlEncodedInToken, +} from './kilo-targets.js'; + +describe('providerBaseUrlEncodedInToken', () => { + it('extracts and normalizes a provider base while preserving the full token separately', () => { + expect( + providerBaseUrlEncodedInToken('http://localhost:9911/api/openrouter/:provider-token') + ).toBe('http://localhost:9911/api/openrouter'); + }); + + it.each([undefined, '', 'ordinary-token', 'ftp://localhost/path:token', 'not a url:token'])( + 'does not infer a target from %s', + token => { + expect(providerBaseUrlEncodedInToken(token)).toBeUndefined(); + } + ); +}); + +describe('backendUrlForSandbox', () => { + it.each([ + ['http://localhost:3000/api/', 'http://host.docker.internal:3000/api'], + ['http://127.0.0.1:8800/', 'http://host.docker.internal:8800'], + ['https://api.kilo.ai/base/', 'https://api.kilo.ai/base'], + ])('normalizes %s to %s', (input, expected) => { + expect(backendUrlForSandbox(input)).toBe(expected); + }); +}); + +describe('deriveKiloSandboxTargets', () => { + it('uses the approved defaults', () => { + expect(deriveKiloSandboxTargets({}, 'user-token')).toEqual({ + success: true, + targets: { + backendBaseUrl: 'https://api.kilo.ai', + providerBaseUrl: 'https://api.kilo.ai', + sessionIngestBaseUrl: 'https://ingest.kilosessions.ai', + }, + }); + }); + + it('applies token URL, configured provider, then backend precedence', () => { + expect( + deriveKiloSandboxTargets( + { + KILOCODE_BACKEND_BASE_URL: 'https://backend.example.com/base', + KILO_OPENROUTER_BASE: 'https://configured.example.com/api', + }, + 'http://localhost:9911/api/openrouter:raw-provider-token' + ) + ).toEqual({ + success: true, + targets: { + backendBaseUrl: 'https://backend.example.com/base', + providerBaseUrl: 'http://host.docker.internal:9911/api/openrouter', + sessionIngestBaseUrl: 'https://ingest.kilosessions.ai', + }, + }); + + expect( + deriveKiloSandboxTargets( + { + KILOCODE_BACKEND_BASE_URL: 'https://backend.example.com/base', + KILO_OPENROUTER_BASE: 'https://configured.example.com/api', + }, + 'raw-user-token' + ) + ).toMatchObject({ + success: true, + targets: { providerBaseUrl: 'https://configured.example.com/api' }, + }); + }); + + it('rewrites explicit localhost backend and ingest targets for the sandbox', () => { + expect( + deriveKiloSandboxTargets( + { + KILOCODE_BACKEND_BASE_URL: 'http://localhost:3000/root/', + KILO_SESSION_INGEST_URL: 'http://127.0.0.1:8800/ingest/', + }, + 'user-token' + ) + ).toMatchObject({ + success: true, + targets: { + backendBaseUrl: 'http://host.docker.internal:3000/root', + providerBaseUrl: 'http://host.docker.internal:3000/root', + sessionIngestBaseUrl: 'http://host.docker.internal:8800/ingest', + }, + }); + }); + + it.each([ + ['production HTTP', { KILOCODE_BACKEND_BASE_URL: 'http://api.kilo.ai' }], + ['userinfo', { KILO_OPENROUTER_BASE: 'https://user@example.com/api' }], + ['query', { KILO_SESSION_INGEST_URL: 'https://ingest.example.com/root?target=other' }], + ['encoded separator', { KILOCODE_BACKEND_BASE_URL: 'https://api.example.com/base%2fescape' }], + ] as const)('rejects %s targets', (_description, env) => { + expect(deriveKiloSandboxTargets(env, 'user-token')).toEqual({ + success: false, + reason: 'invalid_target', + }); + }); +}); diff --git a/services/cloud-agent-next/src/kilo/kilo-targets.ts b/services/cloud-agent-next/src/kilo/kilo-targets.ts new file mode 100644 index 0000000000..656dcd8172 --- /dev/null +++ b/services/cloud-agent-next/src/kilo/kilo-targets.ts @@ -0,0 +1,88 @@ +import { DEFAULT_BACKEND_URL } from '../constants.js'; + +const DEFAULT_SESSION_INGEST_URL = 'https://ingest.kilosessions.ai'; +const LOCAL_SANDBOX_HOSTNAME = 'host.docker.internal'; + +type KiloTargetEnv = { + KILOCODE_BACKEND_BASE_URL?: string; + KILO_OPENROUTER_BASE?: string; + KILO_SESSION_INGEST_URL?: string; +}; + +export type KiloSandboxTargets = { + backendBaseUrl: string; + providerBaseUrl: string; + sessionIngestBaseUrl: string; +}; + +export type DerivedKiloSandboxTargets = + | { + success: true; + targets: KiloSandboxTargets; + } + | { success: false; reason: 'invalid_target' }; + +export function providerBaseUrlEncodedInToken(token: string | undefined): string | undefined { + if (!token) return undefined; + const match = token.match(/^(https?:\/\/[^:]+(?::\d+)?(?:\/[^:]*)?):/); + if (!match?.[1]) return undefined; + try { + return new URL(match[1]).toString().replace(/\/+$/, ''); + } catch { + return undefined; + } +} + +export function backendUrlForSandbox(workerBackendUrl: string): string { + try { + const url = new URL(workerBackendUrl); + if (url.hostname === 'localhost' || url.hostname === '127.0.0.1') { + url.hostname = LOCAL_SANDBOX_HOSTNAME; + } + return url.toString().replace(/\/+$/, ''); + } catch { + return workerBackendUrl; + } +} + +function normalizeSandboxTarget(value: string): string | null { + if (/%(?:2f|5c)/i.test(value)) return null; + let url: URL; + try { + url = new URL(backendUrlForSandbox(value)); + } catch { + return null; + } + if (url.username || url.password || url.search || url.hash) return null; + const localHttp = + url.protocol === 'http:' && + ['localhost', '127.0.0.1', LOCAL_SANDBOX_HOSTNAME].includes(url.hostname); + if (url.protocol !== 'https:' && !localHttp) return null; + url.pathname = url.pathname === '/' ? '' : url.pathname.replace(/\/+$/, ''); + return url.toString().replace(/\/+$/, ''); +} + +export function deriveKiloSandboxTargets( + env: KiloTargetEnv, + userToken: string +): DerivedKiloSandboxTargets { + const backendBaseUrl = normalizeSandboxTarget( + env.KILOCODE_BACKEND_BASE_URL ?? DEFAULT_BACKEND_URL + ); + if (!backendBaseUrl) return { success: false, reason: 'invalid_target' }; + + const providerSource = + providerBaseUrlEncodedInToken(userToken) ?? env.KILO_OPENROUTER_BASE ?? backendBaseUrl; + const providerBaseUrl = normalizeSandboxTarget(providerSource); + const sessionIngestBaseUrl = normalizeSandboxTarget( + env.KILO_SESSION_INGEST_URL ?? DEFAULT_SESSION_INGEST_URL + ); + if (!providerBaseUrl || !sessionIngestBaseUrl) { + return { success: false, reason: 'invalid_target' }; + } + + return { + success: true, + targets: { backendBaseUrl, providerBaseUrl, sessionIngestBaseUrl }, + }; +} diff --git a/services/cloud-agent-next/src/model-validation.test.ts b/services/cloud-agent-next/src/model-validation.test.ts index 973236be34..c9016e11fa 100644 --- a/services/cloud-agent-next/src/model-validation.test.ts +++ b/services/cloud-agent-next/src/model-validation.test.ts @@ -15,7 +15,6 @@ describe('model validation', () => { const officialEnv = { KILOCODE_BACKEND_BASE_URL: 'https://api.kilo.test', - KILOCODE_TOKEN_OVERRIDE: 'override-token', KILOCODE_ORG_ID_OVERRIDE: 'override-org', } as unknown as Env; @@ -50,7 +49,7 @@ describe('model validation', () => { if (typeof init.body !== 'string') throw new Error('Expected JSON request body'); expect(JSON.parse(init.body)).toEqual({ modelId: 'anthropic/claude-sonnet' }); const headers = init.headers as Headers; - expect(headers.get('Authorization')).toBe('Bearer override-token'); + expect(headers.get('Authorization')).toBe('Bearer stored-token'); expect(headers.get('X-KiloCode-OrganizationId')).toBe('override-org'); expect(headers.get('X-KiloCode-Feature')).toBe('cloud-agent-web'); }); diff --git a/services/cloud-agent-next/src/model-validation.ts b/services/cloud-agent-next/src/model-validation.ts index e15ebfba4a..5b4e9d8dcc 100644 --- a/services/cloud-agent-next/src/model-validation.ts +++ b/services/cloud-agent-next/src/model-validation.ts @@ -13,10 +13,7 @@ const MODEL_VALIDATION_UNAVAILABLE_MESSAGE = 'Model availability could not be ve type ModelValidationEnv = Pick< PersistenceEnv, - | 'KILOCODE_BACKEND_BASE_URL' - | 'KILO_OPENROUTER_BASE' - | 'KILOCODE_TOKEN_OVERRIDE' - | 'KILOCODE_ORG_ID_OVERRIDE' + 'KILOCODE_BACKEND_BASE_URL' | 'KILO_OPENROUTER_BASE' | 'KILOCODE_ORG_ID_OVERRIDE' >; type EffectiveCatalogContext = { @@ -53,7 +50,7 @@ type EndpointValidationResult = function effectiveCatalogContext(input: AssertKiloModelAvailableInput): EffectiveCatalogContext { return { - token: input.env.KILOCODE_TOKEN_OVERRIDE ?? input.originalToken, + token: input.originalToken, organizationId: input.env.KILOCODE_ORG_ID_OVERRIDE ?? input.originalOrganizationId, feature: input.createdOnPlatform ?? 'cloud-agent', }; diff --git a/services/cloud-agent-next/src/persistence/types.ts b/services/cloud-agent-next/src/persistence/types.ts index 6996afb440..8d41d52ded 100644 --- a/services/cloud-agent-next/src/persistence/types.ts +++ b/services/cloud-agent-next/src/persistence/types.ts @@ -134,8 +134,6 @@ export type PersistenceEnv = { NEXTAUTH_SECRET: SecretBinding; /** Comma-separated list of allowed Origins for /stream WebSocket connections */ WS_ALLOWED_ORIGINS?: string; - /** Optional override for Kilocode token injected into session environment (does not affect authentication) */ - KILOCODE_TOKEN_OVERRIDE?: string; /** Optional override for Kilocode org ID injected into session environment (does not affect authentication) */ KILOCODE_ORG_ID_OVERRIDE?: string; /** Backend base URL for API calls and session environment variables (defaults to https://kilo.ai) */ diff --git a/services/cloud-agent-next/src/sandbox-outbound.test.ts b/services/cloud-agent-next/src/sandbox-outbound.test.ts index 7685979f66..e6332c2d2b 100644 --- a/services/cloud-agent-next/src/sandbox-outbound.test.ts +++ b/services/cloud-agent-next/src/sandbox-outbound.test.ts @@ -40,9 +40,11 @@ const CAPABILITY = 'kgh2.opaque'; const LEGACY_CAPABILITY = 'kgh1.opaque'; const GITLAB_CAPABILITY = 'kgl2.opaque'; const LEGACY_GITLAB_CAPABILITY = 'kgl1.opaque'; +const KILO_CAPABILITY = 'kka1.opaque'; const OUTBOUND_CONTEXT = { containerId: 'container-test', className: 'SandboxContainment' }; const REDEEMED_GIT_AUTHORIZATION = `Basic ${Buffer.from('x-access-token:upstream-token').toString('base64')}`; const REDEEMED_GITLAB_AUTHORIZATION = `Basic ${Buffer.from('oauth2:upstream-token').toString('base64')}`; +const REDEEMED_KILO_AUTHORIZATION = 'Bearer upstream-kilo-token'; function basicCredential(password: string, scheme = 'Basic', username = 'x-access-token'): string { return `${scheme} ${Buffer.from(`${username}:${password}`).toString('base64')}`; @@ -50,10 +52,17 @@ function basicCredential(password: string, scheme = 'Basic', username = 'x-acces function createEnv( redeemGitHubSessionCapability: ReturnType = vi.fn(), - redeemGitLabSessionCapability: ReturnType = vi.fn() + redeemGitLabSessionCapability: ReturnType = vi.fn(), + redeemKiloSessionCapability: ReturnType = vi.fn(), + logRejectedKiloUrls?: string ) { return { - GIT_TOKEN_SERVICE: { redeemGitHubSessionCapability, redeemGitLabSessionCapability }, + GIT_TOKEN_SERVICE: { + redeemGitHubSessionCapability, + redeemGitLabSessionCapability, + redeemKiloSessionCapability, + }, + LOG_REJECTED_KILO_URLS: logRejectedKiloUrls, } as never; } @@ -990,3 +999,183 @@ describe('handleManagedScmOutbound API authorization', () => { expect(forward).not.toHaveBeenCalled(); }); }); + +describe('handleManagedScmOutbound Kilo authorization', () => { + beforeEach(() => { + vi.unstubAllGlobals(); + }); + + it('redeems a managed Kilo capability, rewrites authorization and uses manual redirects', async () => { + const redeemKiloSessionCapability = vi.fn().mockResolvedValue({ + success: true, + authorization: REDEEMED_KILO_AUTHORIZATION, + routeClass: 'backend_api', + }); + const forward = vi.fn().mockResolvedValue(new Response('forwarded')); + vi.stubGlobal('fetch', forward); + + await handleOutbound( + new Request('https://api.kilo.ai/api/users/me', { + headers: { Authorization: `Bearer ${KILO_CAPABILITY}` }, + }), + createEnv(vi.fn(), vi.fn(), redeemKiloSessionCapability) + ); + + expect(redeemKiloSessionCapability).toHaveBeenCalledWith({ + capability: KILO_CAPABILITY, + outboundContainerId: OUTBOUND_CONTEXT.containerId, + requestMethod: 'GET', + requestUrl: 'https://api.kilo.ai/api/users/me', + bootstrapKiloSessionId: undefined, + }); + const forwarded = forward.mock.calls[0]?.[0] as Request; + expect(forwarded.headers.get('Authorization')).toBe(REDEEMED_KILO_AUTHORIZATION); + expect(forwarded.redirect).toBe('manual'); + }); + + it('passes a matching session identity when redeeming Kilo session-ingest bootstrap', async () => { + const redeemKiloSessionCapability = vi.fn().mockResolvedValue({ + success: true, + authorization: REDEEMED_KILO_AUTHORIZATION, + routeClass: 'session_ingest', + }); + const forward = vi.fn().mockResolvedValue(new Response('forwarded')); + vi.stubGlobal('fetch', forward); + + await handleOutbound( + new Request('https://ingest.kilosessions.ai/api/session', { + method: 'POST', + headers: { + Authorization: `Bearer ${KILO_CAPABILITY}`, + 'Content-Type': 'application/json', + }, + body: JSON.stringify({ sessionId: 'kilo-session-1' }), + }), + createEnv(vi.fn(), vi.fn(), redeemKiloSessionCapability) + ); + + expect(redeemKiloSessionCapability).toHaveBeenCalledWith({ + capability: KILO_CAPABILITY, + outboundContainerId: OUTBOUND_CONTEXT.containerId, + requestMethod: 'POST', + requestUrl: 'https://ingest.kilosessions.ai/api/session', + bootstrapKiloSessionId: 'kilo-session-1', + }); + }); + + it('fails closed without logging a rejected Kilo URL by default', async () => { + const forward = vi.fn(); + vi.stubGlobal('fetch', forward); + const rejectedUrl = + 'https://api.kilo.ai/api/openrouter/v1/chat/completions?access_token=query-secret'; + + const response = await handleOutbound( + new Request(rejectedUrl, { + headers: { Authorization: `Bearer ${KILO_CAPABILITY}` }, + }), + createEnv( + vi.fn(), + vi.fn(), + vi.fn().mockResolvedValue({ success: false, reason: 'upstream_not_allowed' }) + ) + ); + + expect(response.status).toBe(502); + expect(forward).not.toHaveBeenCalled(); + expect(serializedLogCalls()).not.toContain(rejectedUrl); + expect(serializedLogCalls()).not.toContain('query-secret'); + }); + + it('logs a rejected Kilo URL without its query string when local diagnostics are enabled', async () => { + const forward = vi.fn(); + vi.stubGlobal('fetch', forward); + const rejectedUrl = + 'https://api.kilo.ai/api/openrouter/v1/chat/completions?access_token=query-secret'; + + const response = await handleOutbound( + new Request(rejectedUrl, { + headers: { Authorization: `Bearer ${KILO_CAPABILITY}` }, + }), + createEnv( + vi.fn(), + vi.fn(), + vi.fn().mockResolvedValue({ success: false, reason: 'upstream_not_allowed' }), + '1' + ) + ); + + expect(response.status).toBe(502); + expect(forward).not.toHaveBeenCalled(); + expect(logging.logger.withFields).toHaveBeenCalledWith( + expect.objectContaining({ + rejectedRequestUrl: 'https://api.kilo.ai/api/openrouter/v1/chat/completions', + }) + ); + expect(serializedLogCalls()).not.toContain('query-secret'); + expect(logging.logger.debug).toHaveBeenCalledWith( + 'Managed Kilo outbound redemption rejected with local URL diagnostics' + ); + expect(logging.logger.warn).toHaveBeenCalledWith('Managed Kilo outbound redemption rejected'); + }); + + it('fails closed without forwarding when Kilo redemption throws', async () => { + const forward = vi.fn(); + vi.stubGlobal('fetch', forward); + + const response = await handleOutbound( + new Request('https://api.kilo.ai/api/users/me', { + headers: { Authorization: `Bearer ${KILO_CAPABILITY}` }, + }), + createEnv(vi.fn(), vi.fn(), vi.fn().mockRejectedValue(new Error('rpc unavailable'))) + ); + + expect(response.status).toBe(502); + expect(forward).not.toHaveBeenCalled(); + }); + + it('fails closed when the Kilo redemption binding is unavailable', async () => { + const forward = vi.fn(); + vi.stubGlobal('fetch', forward); + + const response = await handleOutbound( + new Request('https://api.kilo.ai/api/users/me', { + headers: { Authorization: `Bearer ${KILO_CAPABILITY}` }, + }), + { GIT_TOKEN_SERVICE: {} } as never + ); + + expect(response.status).toBe(502); + expect(forward).not.toHaveBeenCalled(); + }); + + it('fails closed for an unsupported Kilo capability version', async () => { + const redeemKiloSessionCapability = vi.fn(); + const forward = vi.fn(); + vi.stubGlobal('fetch', forward); + + const response = await handleOutbound( + new Request('https://api.kilo.ai/api/users/me', { + headers: { Authorization: 'Bearer kka2.opaque' }, + }), + createEnv(vi.fn(), vi.fn(), redeemKiloSessionCapability) + ); + + expect(response.status).toBe(502); + expect(redeemKiloSessionCapability).not.toHaveBeenCalled(); + expect(forward).not.toHaveBeenCalled(); + }); + + it('passes an ordinary unrelated outbound request through unchanged', async () => { + const redeemKiloSessionCapability = vi.fn(); + const forward = vi.fn().mockResolvedValue(new Response('forwarded')); + vi.stubGlobal('fetch', forward); + + await handleOutbound( + new Request('https://example.com/unrelated'), + createEnv(vi.fn(), vi.fn(), redeemKiloSessionCapability) + ); + + expect(redeemKiloSessionCapability).not.toHaveBeenCalled(); + expect(forward).toHaveBeenCalledOnce(); + }); +}); diff --git a/services/cloud-agent-next/src/sandbox-outbound.ts b/services/cloud-agent-next/src/sandbox-outbound.ts index 5c29283633..0609f1678a 100644 --- a/services/cloud-agent-next/src/sandbox-outbound.ts +++ b/services/cloud-agent-next/src/sandbox-outbound.ts @@ -8,11 +8,14 @@ export { MANAGED_SCM_OUTBOUND_HANDLER } from './sandbox-id.js'; const GITHUB_CAPABILITY_PREFIXES = ['kgh1.', 'kgh2.']; const GITLAB_CAPABILITY_PREFIXES = ['kgl1.', 'kgl2.']; +const KILO_CAPABILITY_PREFIXES = ['kka1.']; +const MAX_KILO_SESSION_BOOTSTRAP_BYTES = 16_000; type GitHubTokenRedemptionBinding = Pick; type GitLabTokenRedemptionBinding = Pick; +type KiloTokenRedemptionBinding = Pick; type ManagedScmOutboundContext = { containerId: string }; -type RedeemableAuthorization = { provider: 'github' | 'gitlab'; capability: string }; +type RedeemableAuthorization = { provider: 'github' | 'gitlab' | 'kilo'; capability: string }; type AuthorizationExtraction = | { type: 'none' } | { type: 'capability'; value: RedeemableAuthorization } @@ -31,6 +34,13 @@ type AuthorizationClass = | 'unmanaged' | 'none'; type DiagnosticLevel = 'debug' | 'info' | 'warn'; +type LocalKiloUrlDiagnosticsEnvironment = { LOG_REJECTED_KILO_URLS?: string }; + +function isLocalKiloUrlDiagnosticsEnabled( + env: Cloudflare.Env & LocalKiloUrlDiagnosticsEnvironment +): boolean { + return env.LOG_REJECTED_KILO_URLS === '1'; +} function logDiagnostic( level: DiagnosticLevel, @@ -110,10 +120,97 @@ function getCapabilityVersion(capability: string): string { return separator === -1 ? 'unknown' : capability.slice(0, separator); } +function isKiloSessionBootstrapRequest(request: Request): boolean { + if (request.method.toUpperCase() !== 'POST') return false; + try { + return new URL(request.url).pathname.endsWith('/api/session'); + } catch { + return false; + } +} + +/** + * Reads a clone of the request body, stopping as soon as it exceeds maxBytes + * so an unbounded (or under-declared) body is never fully buffered. + */ +async function readBoundedRequestBody( + request: Request, + maxBytes: number +): Promise { + const body: ReadableStream | null = request.clone().body; + if (!body) return undefined; + const reader = body.getReader(); + const chunks: Uint8Array[] = []; + let total = 0; + try { + for (;;) { + const { done, value } = await reader.read(); + if (done) break; + total += value.byteLength; + if (total > maxBytes) return undefined; + chunks.push(value); + } + } finally { + reader.cancel().catch(() => {}); + } + const buffer = new Uint8Array(total); + let offset = 0; + for (const chunk of chunks) { + buffer.set(chunk, offset); + offset += chunk.byteLength; + } + return new TextDecoder().decode(buffer); +} + +async function getKiloSessionBootstrapId(request: Request): Promise { + if (!isKiloSessionBootstrapRequest(request)) return undefined; + const contentType = request.headers.get('Content-Type')?.split(';', 1)[0]?.trim().toLowerCase(); + if (contentType !== 'application/json') return undefined; + + const contentLength = request.headers.get('Content-Length'); + if ( + contentLength && + (!/^[0-9]+$/.test(contentLength) || Number(contentLength) > MAX_KILO_SESSION_BOOTSTRAP_BYTES) + ) { + return undefined; + } + + try { + const body = await readBoundedRequestBody(request, MAX_KILO_SESSION_BOOTSTRAP_BYTES); + if (body === undefined) { + return undefined; + } + const parsed: unknown = JSON.parse(body); + if ( + typeof parsed === 'object' && + parsed !== null && + !Array.isArray(parsed) && + 'sessionId' in parsed && + typeof parsed.sessionId === 'string' && + parsed.sessionId.length > 0 + ) { + return parsed.sessionId; + } + } catch { + // The broker rejects a bootstrap route without a matching session identity. + } + return undefined; +} + function classifyDiagnosticError(error: unknown): 'error' | 'unknown' { return error instanceof Error ? 'error' : 'unknown'; } +/** Query strings can carry credentials — even local diagnostics log only origin + path. */ +function redactUrlQuery(url: string): string { + try { + const parsed = new URL(url); + return `${parsed.origin}${parsed.pathname}`; + } catch { + return 'invalid-url'; + } +} + function supportsGitHubSessionCapabilityRedemption( service: unknown ): service is GitHubTokenRedemptionBinding { @@ -136,6 +233,17 @@ function supportsGitLabSessionCapabilityRedemption( ); } +function supportsKiloSessionCapabilityRedemption( + service: unknown +): service is KiloTokenRedemptionBinding { + return ( + typeof service === 'object' && + service !== null && + 'redeemKiloSessionCapability' in service && + typeof service.redeemKiloSessionCapability === 'function' + ); +} + function classifyCapability(capability: string): AuthorizationExtraction { if (GITHUB_CAPABILITY_PREFIXES.some(prefix => capability.startsWith(prefix))) { return { type: 'capability', value: { provider: 'github', capability } }; @@ -143,7 +251,10 @@ function classifyCapability(capability: string): AuthorizationExtraction { if (GITLAB_CAPABILITY_PREFIXES.some(prefix => capability.startsWith(prefix))) { return { type: 'capability', value: { provider: 'gitlab', capability } }; } - return /^(?:kgh|kgl)\d+\./.test(capability) + if (KILO_CAPABILITY_PREFIXES.some(prefix => capability.startsWith(prefix))) { + return { type: 'capability', value: { provider: 'kilo', capability } }; + } + return /^(?:kgh|kgl|kka)\d+\./.test(capability) ? { type: 'unsupported_capability' } : NO_AUTHORIZATION_CAPABILITY; } @@ -313,6 +424,98 @@ async function handleManagedGitHubOutbound( return response; } +async function handleManagedKiloOutbound( + request: Request, + env: Cloudflare.Env, + capability: { capability: string }, + outboundContainerId: string +): Promise { + const logFields = { + ...getSafeRequestLogFields(request), + provider: 'kilo', + capabilityVersion: getCapabilityVersion(capability.capability), + outboundContainerId, + }; + logDiagnostic('debug', logFields, 'Redeeming managed Kilo outbound request'); + + const tokenService = env.GIT_TOKEN_SERVICE; + if (!supportsKiloSessionCapabilityRedemption(tokenService)) { + logDiagnostic( + 'warn', + { ...logFields, failureStage: 'redemption-binding' }, + 'Managed Kilo outbound redemption unavailable' + ); + return new Response('Kilo authorization unavailable', { status: 502 }); + } + + let result: Awaited>; + try { + const bootstrapKiloSessionId = await getKiloSessionBootstrapId(request); + result = await tokenService.redeemKiloSessionCapability({ + capability: capability.capability, + outboundContainerId, + requestMethod: request.method, + requestUrl: request.url, + bootstrapKiloSessionId, + }); + } catch (error) { + logDiagnostic( + 'warn', + { + ...logFields, + failureStage: 'redemption-rpc', + errorClass: classifyDiagnosticError(error), + }, + 'Managed Kilo outbound redemption failed' + ); + return new Response('Kilo authorization unavailable', { status: 502 }); + } + + if (!result.success) { + if (isLocalKiloUrlDiagnosticsEnabled(env)) { + logDiagnostic( + 'debug', + { + ...logFields, + failureStage: 'redemption-policy', + reason: result.reason, + rejectedRequestUrl: redactUrlQuery(request.url), + }, + 'Managed Kilo outbound redemption rejected with local URL diagnostics' + ); + } + logDiagnostic( + 'warn', + { ...logFields, failureStage: 'redemption-policy', reason: result.reason }, + 'Managed Kilo outbound redemption rejected' + ); + return new Response('Kilo authorization unavailable', { status: 502 }); + } + + let response: Response; + try { + response = await forwardRedeemedRequest(request, { authorization: result.authorization }); + } catch (error) { + logDiagnostic( + 'warn', + { + ...logFields, + failureStage: 'upstream-forward', + errorClass: classifyDiagnosticError(error), + }, + 'Managed Kilo outbound forwarding failed' + ); + return new Response('Kilo authorization unavailable', { status: 502 }); + } + + logDiagnostic( + 'info', + { ...logFields, upstreamStatus: response.status, routeClass: result.routeClass }, + 'Managed Kilo outbound request forwarded' + ); + return response; +} + async function handleManagedGitLabOutbound( request: Request, env: Cloudflare.Env, @@ -391,9 +594,13 @@ export function handleManagedScmOutbound( } const capability = authorizationCapability ?? gitLabPrivateTokenCapability; if (!capability) return fetch(request); - return capability.provider === 'github' - ? handleManagedGitHubOutbound(request, env, capability, ctx.containerId) - : handleManagedGitLabOutbound(request, env, capability, ctx.containerId); + if (capability.provider === 'github') { + return handleManagedGitHubOutbound(request, env, capability, ctx.containerId); + } + if (capability.provider === 'kilo') { + return handleManagedKiloOutbound(request, env, capability, ctx.containerId); + } + return handleManagedGitLabOutbound(request, env, capability, ctx.containerId); } const managedScmOutboundHandlers = { diff --git a/services/cloud-agent-next/src/server.test.ts b/services/cloud-agent-next/src/server.test.ts index 4c68a30bf0..faf63a11cf 100644 --- a/services/cloud-agent-next/src/server.test.ts +++ b/services/cloud-agent-next/src/server.test.ts @@ -1,6 +1,7 @@ import { beforeEach, describe, expect, it, vi } from 'vitest'; import jwt from 'jsonwebtoken'; import type { Env } from './types.js'; +import { mintWrapperDispatchTicket, type WrapperDispatchTicketClaims } from './auth.js'; const { getRunningTerminalClientMock, @@ -142,6 +143,22 @@ function signKiloToken(userId = 'usr_1'): string { ); } +function signWrapperDispatchTicket(overrides: Partial = {}): string { + return mintWrapperDispatchTicket( + { + type: 'wrapper_dispatch_ticket', + userId: 'usr_feed', + cloudAgentSessionId: 'agent_live', + kiloSessionId: 'ses_12345678901234567890123456', + wrapperRunId: 'wr_1', + wrapperGeneration: 2, + wrapperConnectionId: 'conn_1', + ...overrides, + }, + secret + ); +} + beforeEach(() => { getRunningTerminalClientMock.mockReset(); consumeCloudAgentReportBatchMock.mockClear(); @@ -513,7 +530,7 @@ describe('server raw global feed route', () => { ); env.USER_KILO_FACADE.idFromName.mockReturnValue('facade-id'); env.USER_KILO_FACADE.get.mockReturnValue({ fetch: facadeFetch }); - const token = signKiloToken('usr_feed'); + const ticket = signWrapperDispatchTicket(); const response = await fetchWorker( new Request( @@ -521,7 +538,7 @@ describe('server raw global feed route', () => { { headers: { Upgrade: 'websocket', - Authorization: `Bearer ${token}`, + Authorization: `Bearer ${ticket}`, Cookie: 'session=secret', 'x-kilo-facade-user-id': 'usr_attacker', 'x-kilo-facade-auth-token': 'attacker-token', @@ -565,7 +582,7 @@ describe('server raw global feed route', () => { }); const facadeFetch = vi.fn(); env.USER_KILO_FACADE.get.mockReturnValue({ fetch: facadeFetch }); - const token = signKiloToken('usr_feed'); + const ticket = signWrapperDispatchTicket(); const response = await fetchWorker( new Request( @@ -573,7 +590,7 @@ describe('server raw global feed route', () => { { headers: { Upgrade: 'websocket', - Authorization: `Bearer ${token}`, + Authorization: `Bearer ${ticket}`, }, } ), @@ -586,7 +603,7 @@ describe('server raw global feed route', () => { it('rejects repeated producer identity parameters before session validation', async () => { const env = createEnv(); - const token = signKiloToken('usr_feed'); + const ticket = signWrapperDispatchTicket(); const response = await fetchWorker( new Request( @@ -594,7 +611,7 @@ describe('server raw global feed route', () => { { headers: { Upgrade: 'websocket', - Authorization: `Bearer ${token}`, + Authorization: `Bearer ${ticket}`, }, } ), @@ -607,7 +624,7 @@ describe('server raw global feed route', () => { it('rejects malformed producer generation before session validation', async () => { const env = createEnv(); - const token = signKiloToken('usr_feed'); + const ticket = signWrapperDispatchTicket(); const response = await fetchWorker( new Request( @@ -615,7 +632,7 @@ describe('server raw global feed route', () => { { headers: { Upgrade: 'websocket', - Authorization: `Bearer ${token}`, + Authorization: `Bearer ${ticket}`, }, } ), @@ -625,4 +642,228 @@ describe('server raw global feed route', () => { expect(response.status).toBe(400); expect(env.CLOUD_AGENT_SESSION.idFromName).not.toHaveBeenCalled(); }); + + it('accepts a legacy raw Kilo JWT for wrapper processes bound before dispatch tickets shipped', async () => { + 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 }); + const facadeFetch = vi.fn<(request: Request) => Promise>( + async () => new Response('accepted', { status: 200 }) + ); + env.USER_KILO_FACADE.idFromName.mockReturnValue('facade-id'); + env.USER_KILO_FACADE.get.mockReturnValue({ fetch: facadeFetch }); + const token = signKiloToken('usr_feed'); + + 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 ${token}`, + }, + } + ), + env + ); + + expect(response.status).toBe(200); + expect(env.CLOUD_AGENT_SESSION.idFromName).toHaveBeenCalledWith('usr_feed:agent_live'); + expect(validateKiloGlobalFeedProducer).toHaveBeenCalledWith({ + kiloSessionId: 'ses_12345678901234567890123456', + wrapperRunId: 'wr_1', + wrapperGeneration: 2, + wrapperConnectionId: 'conn_1', + }); + }); + + it('rejects a ticket whose fence claims disagree with the request query', async () => { + const env = createEnv(); + const ticket = signWrapperDispatchTicket({ wrapperRunId: 'wr_stale' }); + + 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 ${ticket}`, + }, + } + ), + env + ); + + expect(response.status).toBe(403); + await expect(response.text()).resolves.toBe('Ticket does not match dispatch fence'); + expect(env.CLOUD_AGENT_SESSION.idFromName).not.toHaveBeenCalled(); + }); +}); + +describe('server wrapper ingest route', () => { + 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 }); + const ticket = signWrapperDispatchTicket(); + + const response = await fetchWorker( + new Request( + 'http://worker.test/sessions/usr_feed/agent_live/ingest?kiloSessionId=ses_12345678901234567890123456&wrapperRunId=wr_1&wrapperGeneration=2&wrapperConnectionId=conn_1', + { + headers: { Upgrade: 'websocket', Authorization: `Bearer ${ticket}` }, + } + ), + env + ); + + expect(response.status).toBe(200); + expect(env.CLOUD_AGENT_SESSION.idFromName).toHaveBeenCalledWith('usr_feed:agent_live'); + expect(doFetch).toHaveBeenCalledOnce(); + expect(new URL(doFetch.mock.calls[0][0].url).pathname).toBe('/ingest'); + }); + + it('accepts a legacy raw Kilo JWT for wrapper processes bound before dispatch tickets shipped', 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 }); + const token = signKiloToken('usr_feed'); + + const response = await fetchWorker( + new Request('http://worker.test/sessions/usr_feed/agent_live/ingest', { + headers: { Upgrade: 'websocket', Authorization: `Bearer ${token}` }, + }), + env + ); + + expect(response.status).toBe(200); + expect(env.CLOUD_AGENT_SESSION.idFromName).toHaveBeenCalledWith('usr_feed:agent_live'); + expect(doFetch).toHaveBeenCalledOnce(); + }); + + it('rejects a ticket minted for a different user', async () => { + const env = createEnv(); + const ticket = signWrapperDispatchTicket(); + + const response = await fetchWorker( + new Request('http://worker.test/sessions/usr_other/agent_live/ingest', { + headers: { Upgrade: 'websocket', Authorization: `Bearer ${ticket}` }, + }), + env + ); + + expect(response.status).toBe(403); + await expect(response.text()).resolves.toBe('Token does not match session user'); + expect(env.CLOUD_AGENT_SESSION.idFromName).not.toHaveBeenCalled(); + }); + + it('rejects a ticket whose fence claims disagree with the request query', async () => { + const env = createEnv(); + const ticket = signWrapperDispatchTicket(); + + const response = await fetchWorker( + new Request( + 'http://worker.test/sessions/usr_feed/agent_live/ingest?kiloSessionId=ses_12345678901234567890123456&wrapperRunId=wr_1&wrapperGeneration=99&wrapperConnectionId=conn_1', + { + headers: { Upgrade: 'websocket', Authorization: `Bearer ${ticket}` }, + } + ), + env + ); + + expect(response.status).toBe(403); + await expect(response.text()).resolves.toBe('Ticket does not match dispatch fence'); + expect(env.CLOUD_AGENT_SESSION.idFromName).not.toHaveBeenCalled(); + }); +}); + +describe('server wrapper log upload route', () => { + function createLogEnv() { + const env = createEnv(); + return Object.assign(env, { + R2_BUCKET: { put: vi.fn().mockResolvedValue(undefined) }, + }); + } + + it('accepts a valid wrapper dispatch ticket and stores the upload in R2', async () => { + const env = createLogEnv(); + const ticket = signWrapperDispatchTicket(); + + const response = await fetchWorker( + new Request( + 'http://worker.test/sessions/usr_feed/agent_live/logs/session/logs.tar.gz?kiloSessionId=ses_12345678901234567890123456', + { + method: 'PUT', + headers: { Authorization: `Bearer ${ticket}` }, + body: new Uint8Array([1, 2, 3]), + } + ), + env + ); + + expect(response.status).toBe(204); + expect(env.R2_BUCKET.put).toHaveBeenCalledOnce(); + expect(env.R2_BUCKET.put.mock.calls[0][0]).toBe('logs/usr_feed/agent_live/session/logs.tar.gz'); + }); + + it('accepts a legacy raw Kilo JWT for wrapper processes bound before dispatch tickets shipped', async () => { + const env = createLogEnv(); + const token = signKiloToken('usr_feed'); + + const response = await fetchWorker( + new Request( + 'http://worker.test/sessions/usr_feed/agent_live/logs/session/logs.tar.gz?kiloSessionId=ses_12345678901234567890123456', + { + method: 'PUT', + headers: { Authorization: `Bearer ${token}` }, + body: new Uint8Array([1, 2, 3]), + } + ), + env + ); + + expect(response.status).toBe(204); + expect(env.R2_BUCKET.put).toHaveBeenCalledOnce(); + }); + + it('rejects an upload missing the kiloSessionId parameter', async () => { + const env = createLogEnv(); + const ticket = signWrapperDispatchTicket(); + + const response = await fetchWorker( + new Request('http://worker.test/sessions/usr_feed/agent_live/logs/session/logs.tar.gz', { + method: 'PUT', + headers: { Authorization: `Bearer ${ticket}` }, + body: new Uint8Array([1, 2, 3]), + }), + env + ); + + expect(response.status).toBe(400); + expect(env.R2_BUCKET.put).not.toHaveBeenCalled(); + }); + + it('rejects a ticket whose kiloSessionId disagrees with the request query', async () => { + const env = createLogEnv(); + const ticket = signWrapperDispatchTicket({ kiloSessionId: 'ses_other0000000000000000000' }); + + const response = await fetchWorker( + new Request( + 'http://worker.test/sessions/usr_feed/agent_live/logs/session/logs.tar.gz?kiloSessionId=ses_12345678901234567890123456', + { + method: 'PUT', + headers: { Authorization: `Bearer ${ticket}` }, + body: new Uint8Array([1, 2, 3]), + } + ), + env + ); + + expect(response.status).toBe(403); + await expect(response.text()).resolves.toBe('Ticket does not match dispatch fence'); + expect(env.R2_BUCKET.put).not.toHaveBeenCalled(); + }); }); diff --git a/services/cloud-agent-next/src/server.ts b/services/cloud-agent-next/src/server.ts index 7fd4a0a3cf..30edcd46d2 100644 --- a/services/cloud-agent-next/src/server.ts +++ b/services/cloud-agent-next/src/server.ts @@ -5,7 +5,13 @@ import { appRouter } from './router.js'; import type { Env } from './types.js'; import type { HonoContext } from './hono-context.js'; import { logger, withLogTags } from './logger.js'; -import { resolveSecret, validateStreamTicket, validateKiloToken } from './auth.js'; +import { + resolveSecret, + validateStreamTicket, + validateKiloToken, + validateWrapperDispatchTicket, + type WrapperAuthClaims, +} from './auth.js'; import { createErrorHandler, createNotFoundHandler } from '@kilocode/worker-utils'; import { createCallbackQueueConsumer } from './callbacks/index.js'; import type { CallbackJob } from './callbacks/index.js'; @@ -163,6 +169,52 @@ function createSanitizedForwardRequest( return new Request(url, init); } +function parseOptionalWrapperGeneration(raw: string | null): number | undefined { + if (raw === null) return undefined; + const parsed = Number(raw); + return Number.isInteger(parsed) ? parsed : undefined; +} + +/** + * Defense-in-depth on top of the DO's own fencing checks: reject a wrapper + * dispatch ticket whose claims disagree with the fence tuple carried on the + * request itself. Only compares fields the caller supplies — a route that + * doesn't parse a given fence field yet is not forced to require it. + * + * A legacy raw Kilo JWT (see auth.ts) carries no fence claims to compare, so + * it is exempt — wrapper processes bound before ticket support shipped rely + * on requireCurrentSessionAccess/the DO's own checks until their next dispatch. + */ +function ticketClaimsMismatchRequestFence( + claims: WrapperAuthClaims, + expected: { + cloudAgentSessionId: string; + kiloSessionId?: string | null; + wrapperRunId?: string | null; + wrapperGeneration?: number; + wrapperConnectionId?: string | null; + } +): boolean { + if (claims.type !== 'wrapper_dispatch_ticket') return false; + if (claims.cloudAgentSessionId !== expected.cloudAgentSessionId) return true; + if (expected.kiloSessionId != null && claims.kiloSessionId !== expected.kiloSessionId) + return true; + if (expected.wrapperRunId != null && claims.wrapperRunId !== expected.wrapperRunId) return true; + if ( + expected.wrapperGeneration !== undefined && + claims.wrapperGeneration !== expected.wrapperGeneration + ) { + return true; + } + if ( + expected.wrapperConnectionId != null && + claims.wrapperConnectionId !== expected.wrapperConnectionId + ) { + return true; + } + return false; +} + function stripPublicCredentialHeaders(headers: Headers): Headers { const sanitized = new Headers(headers); sanitized.delete('Authorization'); @@ -290,11 +342,14 @@ app.all('/sessions/:userId/:sessionId/kilo-global-ingest', async (c: Context) = const authHeader = c.req.header('Authorization'); const nextAuthSecret = await resolveSecret(c.env.NEXTAUTH_SECRET); - const authResult = await validateKiloToken(authHeader ?? null, nextAuthSecret); + const authResult = await validateWrapperDispatchTicket(authHeader ?? null, nextAuthSecret); if (!authResult.success) { return c.text(authResult.error, 401); } - if (authResult.userId !== userId) { + if (authResult.claims.userId !== userId) { return c.text('Token does not match session user', 403); } + const url = new URL(c.req.url); + const wrapperGenerationParam = url.searchParams.get('wrapperGeneration'); + const wrapperGeneration = parseOptionalWrapperGeneration(wrapperGenerationParam); + if (wrapperGenerationParam !== null && wrapperGeneration === undefined) { + return c.text('Invalid wrapperGeneration parameter', 400); + } + if ( + ticketClaimsMismatchRequestFence(authResult.claims, { + cloudAgentSessionId: sessionId, + kiloSessionId: url.searchParams.get('kiloSessionId'), + wrapperRunId: url.searchParams.get('wrapperRunId'), + wrapperGeneration, + wrapperConnectionId: url.searchParams.get('wrapperConnectionId'), + }) + ) { + return c.text('Ticket does not match dispatch fence', 403); + } + try { await requireCurrentSessionAccess({ env: c.env, @@ -432,19 +517,34 @@ app.put( const authHeader = c.req.header('Authorization'); const nextAuthSecret = await resolveSecret(c.env.NEXTAUTH_SECRET); - const authResult = await validateKiloToken(authHeader ?? null, nextAuthSecret); + const authResult = await validateWrapperDispatchTicket(authHeader ?? null, nextAuthSecret); if (!authResult.success) { return c.text(authResult.error, 401); } - if (authResult.userId !== userId) { + if (authResult.claims.userId !== userId) { return c.text('Token does not match session user', 403); } + const kiloSessionId = new URL(c.req.url).searchParams.get('kiloSessionId'); + if (!kiloSessionId) { + return c.text('Missing kiloSessionId parameter', 400); + } + + if ( + ticketClaimsMismatchRequestFence(authResult.claims, { + cloudAgentSessionId: sessionId, + kiloSessionId, + }) + ) { + return c.text('Ticket does not match dispatch fence', 403); + } + try { await requireCurrentSessionAccess({ env: c.env, kiloUserId: userId, cloudAgentSessionId: sessionId, + expectedKiloSessionId: kiloSessionId, }); } catch (error) { return projectSessionAccessHttpError(error); diff --git a/services/cloud-agent-next/src/services/git-token-service-client.test.ts b/services/cloud-agent-next/src/services/git-token-service-client.test.ts index 35878fa79a..06d0a77ac5 100644 --- a/services/cloud-agent-next/src/services/git-token-service-client.test.ts +++ b/services/cloud-agent-next/src/services/git-token-service-client.test.ts @@ -27,6 +27,8 @@ function createGitTokenService() { redeemGitHubSessionCapability: vi.fn(), issueGitLabSessionCapability: vi.fn(), redeemGitLabSessionCapability: vi.fn(), + issueKiloSessionCapability: vi.fn(), + redeemKiloSessionCapability: vi.fn(), } satisfies GitTokenService; } diff --git a/services/cloud-agent-next/src/services/git-token-service-client.ts b/services/cloud-agent-next/src/services/git-token-service-client.ts index 6d0ac36c7c..0e34d40284 100644 --- a/services/cloud-agent-next/src/services/git-token-service-client.ts +++ b/services/cloud-agent-next/src/services/git-token-service-client.ts @@ -3,6 +3,7 @@ import type { BitbucketTokenFailureReason, GitAuthorConfig, GitTokenService, + KiloSessionCapabilityTargets, ManagedGitHubFallbackReason, } from '../types.js'; @@ -424,3 +425,54 @@ export async function resolveManagedGitLabToken( return { success: false, reason: 'rpc_error' }; } } + +export type ResolvedCloudAgentKiloCapability = { + capability: string; +}; + +export async function issueCloudAgentKiloSessionCapability( + env: GitTokenServiceEnv, + params: { + userId: string; + cloudAgentSessionId: string; + kiloSessionId: string; + outboundContainerId: string; + userToken: string; + targets: KiloSessionCapabilityTargets; + } +): Promise< + | { success: true; value: ResolvedCloudAgentKiloCapability } + | { success: false; error: ResolveGitHubTokenError } +> { + if (!env.GIT_TOKEN_SERVICE) { + return { + success: false, + error: { + reason: 'service_not_configured', + message: 'git-token-service capability issuance is not configured', + }, + }; + } + + try { + const result = await env.GIT_TOKEN_SERVICE.issueKiloSessionCapability(params); + if (!result.success) { + return { + success: false, + error: { + reason: result.reason, + message: `Kilo session capability issuance failed (${result.reason})`, + }, + }; + } + logger.info('Issued Kilo session capability via git-token-service'); + return { success: true, value: { capability: result.capability } }; + } catch (error) { + const message = error instanceof Error ? error.message : String(error); + logger.withFields({ error: message }).error('Failed to issue Kilo session capability'); + return { + success: false, + error: { reason: 'rpc_error', message: `git-token-service RPC failed: ${message}` }, + }; + } +} diff --git a/services/cloud-agent-next/src/session-service.test.ts b/services/cloud-agent-next/src/session-service.test.ts index beb80d78cb..555a36f81d 100644 --- a/services/cloud-agent-next/src/session-service.test.ts +++ b/services/cloud-agent-next/src/session-service.test.ts @@ -2,6 +2,8 @@ import { dirname, relative } from 'node:path'; import { beforeEach, describe, expect, it, vi } from 'vitest'; 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'; vi.mock('./logger.js', () => ({ logger: { @@ -112,7 +114,7 @@ describe('SessionService.buildRuntimeEnv', () => { const runtimeEnv = service.buildRuntimeEnv({ context, env: createEnv(), - originalToken: 'kilo-token', + kiloCapability: 'kilo-token', }); expect(runtimeEnv.HOME).toBe('/home/agent_test'); @@ -382,7 +384,13 @@ function createEnv(metadata?: CloudAgentSessionState | null): PersistenceEnv { }), issueGitLabSessionCapability: vi.fn(), redeemGitLabSessionCapability: vi.fn(), + issueKiloSessionCapability: vi.fn().mockResolvedValue({ + success: true, + capability: 'kka1.default', + }), + redeemKiloSessionCapability: vi.fn(), }, + MANAGED_SCM_CONTAINMENT_ORG_IDS: '*', NOTIFICATIONS: {} as unknown as PersistenceEnv['NOTIFICATIONS'], } satisfies PersistenceEnv; } @@ -742,6 +750,102 @@ describe('SessionService.prepareWorkspace', () => { ); }); + it('writes the opaque Kilo capability to the sandbox auth file, never the raw token', async () => { + const session = createSession(false); + const writeFile = vi.fn().mockResolvedValue(undefined); + const sandbox = createSandbox(session, false, writeFile); + const metadata = createMetadata({ managedScmContainment: true }); + const env = createEnv(); + const issueKiloSessionCapability = vi.fn().mockResolvedValue({ + success: true, + capability: 'kka1.workspace-issued', + }); + if (!env.GIT_TOKEN_SERVICE) throw new Error('Expected GIT_TOKEN_SERVICE in test env'); + env.GIT_TOKEN_SERVICE.issueKiloSessionCapability = issueKiloSessionCapability; + + await new SessionService().prepareWorkspace({ + sandbox, + sandboxId: 'ses-abcdef', + userId: 'user_test', + sessionId: 'agent_test' as SessionId, + env, + metadata, + kilocodeModel: 'test-model', + }); + + expect(issueKiloSessionCapability).toHaveBeenCalledWith( + expect.objectContaining({ + userToken: 'kilo-token', + outboundContainerId: 'containment-small-sandbox-do-id', + }) + ); + const authFileCall = writeFile.mock.calls.find( + ([path]) => path === '/home/agent_test/.local/share/kilo/auth.json' + ); + expect(authFileCall).toBeDefined(); + const authFileContent = authFileCall?.[1] as string; + expect(authFileContent).toContain('kka1.workspace-issued'); + expect(authFileContent).not.toContain('kilo-token'); + for (const [, content] of writeFile.mock.calls) { + if (typeof content === 'string') { + expect(content).not.toContain('kilo-token'); + } + } + }); + + it('binds the restore-token file to the opaque Kilo capability for a standard devcontainer session, never the raw token', async () => { + const session = createSession(false); + const writeFile = vi.fn().mockResolvedValue(undefined); + const sandbox = createSandbox(session, false, writeFile); + const metadata = { + ...createMetadata({ preparedAt: 1 }), + workspace: { + sandboxId: 'ses-abcdef' as const, + devcontainerRequested: true, + managedScmContainment: true, + }, + } satisfies CloudAgentSessionState; + const env = createEnv(); + const issueKiloSessionCapability = vi.fn().mockResolvedValue({ + success: true, + capability: 'kka1.restore-issued', + }); + if (!env.GIT_TOKEN_SERVICE) throw new Error('Expected GIT_TOKEN_SERVICE in test env'); + env.GIT_TOKEN_SERVICE.issueKiloSessionCapability = issueKiloSessionCapability; + const devcontainerHandle = { + containerId: 'container-dev', + innerWorkspaceFolder: '/workspaces/repo', + workspacePath: '/workspace/user/sessions/agent_test', + agentSessionId: 'agent_test', + overrideConfigPath: '/tmp/devcontainer-override-agent_test/devcontainer.json', + teardown: vi.fn().mockResolvedValue(undefined), + }; + devcontainerMocks.detectDevContainer.mockResolvedValue({ + configPath: '.devcontainer/devcontainer.json', + }); + devcontainerMocks.bringUpDevContainer.mockResolvedValue(devcontainerHandle); + + await new SessionService().prepareWorkspace({ + sandbox, + sandboxId: 'ses-abcdef', + userId: 'user_test', + sessionId: 'agent_test' as SessionId, + env, + metadata, + kilocodeModel: 'test-model', + }); + + const restoreTokenCall = writeFile.mock.calls.find( + ([path]) => path === '/home/agent_test/.local/share/kilo/session-restore-token' + ); + expect(restoreTokenCall).toBeDefined(); + expect(restoreTokenCall?.[1]).toBe('kka1.restore-issued'); + const restoreCall = session.exec.mock.calls.find( + ([command]) => typeof command === 'string' && command.includes('kilo-restore-session.js') + ); + expect(restoreCall?.[0]).not.toContain('kilo-token'); + }); + it('types ENOSPC during the cold devcontainer probe before provisioning', async () => { const session = createSession(false); const sandbox = createSandbox(session); @@ -1898,15 +2002,21 @@ describe('SessionService.buildWrapperSessionReadyAndPromptRequests', () => { expect(result.readyRequest.devcontainer).toEqual({ requested: true, resolved: devcontainer }); }); - it('materializes workspace setup and prompt delivery without using provider token overrides for ingest auth', async () => { + it('materializes workspace setup and prompt delivery behind an opaque Kilo capability, never the raw tokens', async () => { const service = new SessionService(); const env = createEnv(); env.WORKER_URL = 'https://cloud-agent.example.com'; - env.KILOCODE_TOKEN_OVERRIDE = 'provider-override-token'; + const issueKiloSessionCapability = vi.fn().mockResolvedValue({ + success: true, + capability: 'kka1.issued', + }); + if (!env.GIT_TOKEN_SERVICE) throw new Error('Expected GIT_TOKEN_SERVICE in test env'); + env.GIT_TOKEN_SERVICE.issueKiloSessionCapability = issueKiloSessionCapability; const metadata = createMetadata({ setupCommands: ['pnpm install'], envVars: { PUBLIC_VALUE: 'visible' }, upstreamBranch: 'main', + managedScmContainment: true, }); const result = await service.buildWrapperSessionReadyAndPromptRequests({ @@ -1952,11 +2062,11 @@ describe('SessionService.buildWrapperSessionReadyAndPromptRequests', () => { sessionHome: '/home/agent_test', branchName: 'main', kiloSessionId: 'kilo-session', - gitToken: 'resolved-gitlab-token', + gitToken: 'kgl2.default', gitlabTokenManaged: true, }); - expect(tokenMocks.resolveManagedGitLabToken).toHaveBeenCalled(); - expect(tokenMocks.issueCloudAgentGitLabSessionCapability).not.toHaveBeenCalled(); + expect(tokenMocks.resolveManagedGitLabToken).not.toHaveBeenCalled(); + expect(tokenMocks.issueCloudAgentGitLabSessionCapability).toHaveBeenCalled(); expect(result.readyRequest).toMatchObject({ agentSessionId: 'agent_test', userId: 'user_test', @@ -1971,7 +2081,7 @@ describe('SessionService.buildWrapperSessionReadyAndPromptRequests', () => { repo: { kind: 'git', url: 'https://gitlab.com/acme/repo.git', - token: 'resolved-gitlab-token', + token: 'kgl2.default', platform: 'gitlab', }, materialized: { @@ -1984,15 +2094,41 @@ describe('SessionService.buildWrapperSessionReadyAndPromptRequests', () => { expect(result.promptRequest).not.toHaveProperty('workspace'); expect(result.promptRequest).not.toHaveProperty('materialized'); expect(result.readyRequest.materialized.env.PUBLIC_VALUE).toBe('visible'); - expect(result.readyRequest.materialized.env.KILOCODE_TOKEN).toBe('provider-override-token'); + expect(result.readyRequest.materialized.env.KILOCODE_TOKEN).toBe('kka1.issued'); expect(JSON.parse(result.readyRequest.materialized.env.KILO_AUTH_CONTENT)).toEqual({ - kilo: { type: 'api', key: 'kilo-token' }, + kilo: { type: 'api', key: 'kka1.issued' }, }); - expect(result.readyRequest.materialized.env.GITLAB_TOKEN).toBe('resolved-gitlab-token'); - expect(JSON.stringify(result.readyRequest)).not.toContain('kgl2.default'); + expect(issueKiloSessionCapability).toHaveBeenCalledWith( + expect.objectContaining({ + userId: 'user_test', + cloudAgentSessionId: 'agent_test', + kiloSessionId: 'kilo-session', + userToken: 'kilo-token', + outboundContainerId: 'containment-small-sandbox-do-id', + }) + ); + expect(result.readyRequest.materialized.env.GITLAB_TOKEN).toBe('kgl2.default'); + expect(JSON.stringify(result.readyRequest)).not.toContain('resolved-gitlab-token'); + expect(JSON.stringify(result.readyRequest)).not.toContain('kilo-token'); expect(result.readyRequest.materialized.env.GITLAB_HOST).toBe('gitlab.com'); expect(result.readyRequest.materialized.env.GLAB_IS_OAUTH2).toBe('true'); - expect(result.readyRequest.session.workerAuthToken).toBe('kilo-token'); + expect(result.readyRequest.session.workerAuthToken).not.toBe('kilo-token'); + const ticketResult = await validateWrapperDispatchTicket( + `Bearer ${result.readyRequest.session.workerAuthToken}`, + 'secret' + ); + expect(ticketResult).toMatchObject({ + success: true, + claims: { + type: 'wrapper_dispatch_ticket', + userId: 'user_test', + cloudAgentSessionId: 'agent_test', + kiloSessionId: 'kilo-session', + wrapperRunId: 'wr_test', + wrapperGeneration: 2, + wrapperConnectionId: 'conn_test', + }, + }); expect(result.readyRequest.session.wrapperRunId).toBe('wr_test'); expect(result.readyRequest).not.toHaveProperty('message'); expect(result.readyRequest).not.toHaveProperty('agent'); @@ -2018,6 +2154,140 @@ describe('SessionService.buildWrapperSessionReadyAndPromptRequests', () => { expect(result.promptRequest.session).toEqual(result.readyRequest.session); }); + it('pins the provider baseURL to the target baked into the issued capability, not env.KILO_OPENROUTER_BASE independently', async () => { + const service = new SessionService(); + const env = createEnv(); + env.WORKER_URL = 'https://cloud-agent.example.com'; + env.KILO_OPENROUTER_BASE = 'https://openrouter.wrong-base.example.com'; + const issueKiloSessionCapability = vi.fn().mockResolvedValue({ + success: true, + capability: 'kka1.issued', + }); + if (!env.GIT_TOKEN_SERVICE) throw new Error('Expected GIT_TOKEN_SERVICE in test env'); + env.GIT_TOKEN_SERVICE.issueKiloSessionCapability = issueKiloSessionCapability; + // The stored Kilo token can encode its own provider base URL (e.g. a + // self-hosted proxy) — that must win over the worker's static + // KILO_OPENROUTER_BASE — the capability's baked-in target and the + // sandbox's actual outbound baseURL must agree, or the outbound + // interceptor rejects the request as upstream_not_allowed. + const metadata = createMetadata({ + upstreamBranch: 'main', + kilocodeToken: 'http://localhost:9911/api/openrouter/:provider-token', + managedScmContainment: true, + }); + const derivedTargets = deriveKiloSandboxTargets(env, metadata.auth.kilocodeToken ?? ''); + if (!derivedTargets.success) throw new Error('Expected valid derived Kilo sandbox targets'); + + const result = await service.buildWrapperSessionReadyAndPromptRequests({ + env, + plan: { + scope: { sessionId: 'agent_test', userId: 'user_test' }, + turn: { + type: 'prompt', + messageId: 'msg_018f1e2d3c4bPayloadTestAAAA', + prompt: 'Do the work', + }, + agent: { mode: 'code', model: 'test-model', variant: 'thinking' }, + finalization: { autoCommit: true, condenseOnComplete: false }, + workspace: { sandboxId: 'ses-abcdef', metadata }, + wrapper: { + fence: { + wrapperRunId: 'wr_test', + wrapperGeneration: 2, + wrapperConnectionId: 'conn_test', + }, + }, + } satisfies FencedWrapperDispatchRequest, + }); + + const configContent = JSON.parse(result.readyRequest.materialized.env.KILO_CONFIG_CONTENT) as { + provider: { kilo: { options: { baseURL?: string } } }; + }; + expect(configContent.provider.kilo.options.baseURL).toBe( + derivedTargets.targets.providerBaseUrl + ); + expect(issueKiloSessionCapability).toHaveBeenCalledWith( + expect.objectContaining({ targets: derivedTargets.targets }) + ); + expect(configContent.provider.kilo.options.baseURL).not.toBe( + 'https://openrouter.wrong-base.example.com' + ); + }); + + it('falls back to the raw Kilo token for DIND sandboxes, which have no outbound interceptor to redeem a capability against', async () => { + const service = new SessionService(); + const env = createEnv(); + env.WORKER_URL = 'https://cloud-agent.example.com'; + const issueKiloSessionCapability = vi.fn(); + if (!env.GIT_TOKEN_SERVICE) throw new Error('Expected GIT_TOKEN_SERVICE in test env'); + env.GIT_TOKEN_SERVICE.issueKiloSessionCapability = issueKiloSessionCapability; + const metadata = createMetadata({ sandboxId: 'dind-abcdef' }); + + const result = await service.buildWrapperSessionReadyAndPromptRequests({ + env, + plan: { + scope: { sessionId: 'agent_test', userId: 'user_test' }, + turn: { + type: 'prompt', + messageId: 'msg_018f1e2d3c4bDindFallbackAA', + prompt: 'Do the work', + }, + agent: { mode: 'code', model: 'test-model' }, + workspace: { sandboxId: 'dind-abcdef', metadata }, + wrapper: { + fence: { + wrapperRunId: 'wr_dind', + wrapperGeneration: 1, + wrapperConnectionId: 'conn_dind', + }, + }, + } satisfies FencedWrapperDispatchRequest, + }); + + expect(issueKiloSessionCapability).not.toHaveBeenCalled(); + expect(result.readyRequest.materialized.env.KILOCODE_TOKEN).toBe('kilo-token'); + expect(JSON.parse(result.readyRequest.materialized.env.KILO_AUTH_CONTENT)).toEqual({ + kilo: { type: 'api', key: 'kilo-token' }, + }); + }); + + it('falls back to the raw Kilo token for an uncontained session', async () => { + const service = new SessionService(); + const env = createEnv(); + env.WORKER_URL = 'https://cloud-agent.example.com'; + const issueKiloSessionCapability = vi.fn(); + if (!env.GIT_TOKEN_SERVICE) throw new Error('Expected GIT_TOKEN_SERVICE in test env'); + env.GIT_TOKEN_SERVICE.issueKiloSessionCapability = issueKiloSessionCapability; + const metadata = createMetadata({ managedScmContainment: false }); + + const result = await service.buildWrapperSessionReadyAndPromptRequests({ + env, + plan: { + scope: { sessionId: 'agent_test', userId: 'user_test' }, + turn: { + type: 'prompt', + messageId: 'msg_018f1e2d3c4bContainmentFallbackAA', + prompt: 'Do the work', + }, + agent: { mode: 'code', model: 'test-model' }, + workspace: { sandboxId: 'ses-abcdef', metadata }, + wrapper: { + fence: { + wrapperRunId: 'wr_containment_fallback', + wrapperGeneration: 1, + wrapperConnectionId: 'conn_cf', + }, + }, + } satisfies FencedWrapperDispatchRequest, + }); + + expect(issueKiloSessionCapability).not.toHaveBeenCalled(); + expect(result.readyRequest.materialized.env.KILOCODE_TOKEN).toBe('kilo-token'); + expect(JSON.parse(result.readyRequest.materialized.env.KILO_AUTH_CONTENT)).toEqual({ + kilo: { type: 'api', key: 'kilo-token' }, + }); + }); + it('allowlists only the active session attachment directory for Kilo file access', async () => { const result = await buildPromptWrapperRequests(createMetadata()); const config: unknown = JSON.parse(result.readyRequest.materialized.env.KILO_CONFIG_CONTENT); diff --git a/services/cloud-agent-next/src/session-service.ts b/services/cloud-agent-next/src/session-service.ts index e29bb0db87..5e25ab09dc 100644 --- a/services/cloud-agent-next/src/session-service.ts +++ b/services/cloud-agent-next/src/session-service.ts @@ -10,15 +10,18 @@ import { type ManagedGitHubFallbackReason, } from './types.js'; import { generateSandboxId, getOutboundContainerId } from './sandbox-id.js'; +import { mintWrapperDispatchTicket, resolveSecret } from './auth.js'; import { normalizeKilocodeModel } from './persistence/model-utils.js'; import { isTemporaryManagedBitbucketTokenFailure, issueCloudAgentGitHubSessionCapability, issueCloudAgentGitLabSessionCapability, + issueCloudAgentKiloSessionCapability, resolveCloudAgentGitHubAuthForRepo, resolveManagedBitbucketToken, resolveManagedGitLabToken, } from './services/git-token-service-client.js'; +import { deriveKiloSandboxTargets } from './kilo/kilo-targets.js'; import { ExecutionError } from './execution/errors.js'; import { checkDiskAndCleanBeforeSetup, @@ -1255,7 +1258,8 @@ export class SessionService { sessionId, workspacePath, env: opts.env, - originalToken: opts.originalToken, + kiloCapability: opts.kiloCapability, + kiloProviderBaseUrl: opts.kiloProviderBaseUrl, kilocodeModel: opts.kilocodeModel, originalOrgId: opts.originalOrgId, githubToken: context.githubToken, @@ -1281,7 +1285,8 @@ export class SessionService { sessionId, workspacePath, env, - originalToken, + kiloCapability, + kiloProviderBaseUrl, kilocodeModel, originalOrgId, githubToken, @@ -1310,8 +1315,6 @@ export class SessionService { const runtimeAgents = profile?.runtimeAgents; const kiloCommands = profile?.kiloCommands; - // Use override if available, otherwise use original values from API - const kilocodeToken = env.KILOCODE_TOKEN_OVERRIDE ?? originalToken; const kilocodeOrganizationId = env.KILOCODE_ORG_ID_OVERRIDE ?? originalOrgId; // Bitbucket Code Reviewer sessions use only trusted worker-owned environment values. @@ -1343,9 +1346,10 @@ export class SessionService { SESSION_ID: sessionId, SESSION_HOME: sessionHome, [PNPM_STORE_ENV_VAR]: PNPM_STORE_DIR, - // Inject Kilocode credentials (with override support) - KILOCODE_TOKEN: kilocodeToken, - KILO_AUTH_CONTENT: JSON.stringify({ kilo: { type: 'api', key: originalToken } }), + // Opaque Kilo capability — redeemed for the real credential at the outbound interceptor + KILOCODE_TOKEN: kiloCapability, + // Backend auth surface (session restore/import). + KILO_AUTH_CONTENT: JSON.stringify({ kilo: { type: 'api', key: kiloCapability } }), // Platform identifier - defaults to 'cloud-agent' if not specified KILO_PLATFORM: createdOnPlatform ?? 'cloud-agent', KILO_DISABLE_AUTOUPDATE: 'true', @@ -1354,13 +1358,18 @@ export class SessionService { }; const providerOptions: Record = { - apiKey: kilocodeToken, - kilocodeToken: kilocodeToken, + apiKey: kiloCapability, + kilocodeToken: kiloCapability, }; if (kilocodeOrganizationId) { providerOptions.kilocodeOrganizationId = kilocodeOrganizationId; } - if (env.KILO_OPENROUTER_BASE) { + if (kiloProviderBaseUrl) { + // Must match the provider target baked into the issued capability's claims — + // otherwise the outbound interceptor's route classification rejects the + // request as upstream_not_allowed even though the capability is valid. + providerOptions.baseURL = kiloProviderBaseUrl; + } else if (env.KILO_OPENROUTER_BASE) { providerOptions.baseURL = backendUrlForSandbox(env.KILO_OPENROUTER_BASE); } const isInteractive = createdOnPlatform == 'cloud-agent-web'; @@ -1612,7 +1621,8 @@ export class SessionService { sandbox, context, env, - originalToken, + kiloCapability, + kiloProviderBaseUrl, kilocodeModel, originalOrgId, createdOnPlatform, @@ -1630,7 +1640,8 @@ export class SessionService { const saferEnvVars = this.buildRuntimeEnv({ context, env, - originalToken, + kiloCapability, + kiloProviderBaseUrl, kilocodeModel, originalOrgId, createdOnPlatform, @@ -1811,6 +1822,68 @@ export class SessionService { }; } + /** + * Issue one Kilo capability for this dispatch when its sandbox uses managed + * credential containment, replacing every sandbox-visible carrier of the raw + * Kilo token (env vars, auth files, provider config). The outbound interceptor + * redeems it per-request; the sandbox never sees `userToken`. + * + * For uncontained sessions, or DIND sandboxes (`SandboxDIND`) which don't set + * `interceptHttps`/`outboundHandlers`, fall back to the real token so the + * sandbox can authenticate directly. + */ + private async issueKiloSessionCapability( + env: PersistenceEnv, + params: { + userId: string; + cloudAgentSessionId: string; + kiloSessionId: string; + sandboxId: string; + managedScmContainment: boolean; + userToken: string; + } + ): Promise<{ capability: string; providerBaseUrl?: string }> { + if (params.sandboxId.startsWith('dind-')) { + return { capability: params.userToken }; + } + + if (!params.managedScmContainment) { + logger + .withFields({ sandboxId: params.sandboxId }) + .info('Using raw Kilo token; session does not use managed containment'); + return { capability: params.userToken }; + } + + const derivedTargets = deriveKiloSandboxTargets(env, params.userToken); + if (!derivedTargets.success) { + throw ExecutionError.invalidRequest('Unable to derive Kilo sandbox targets'); + } + const issued = await issueCloudAgentKiloSessionCapability(env, { + userId: params.userId, + cloudAgentSessionId: params.cloudAgentSessionId, + kiloSessionId: params.kiloSessionId, + outboundContainerId: getOutboundContainerId(env, params.sandboxId, { + managedScmContainment: true, + }), + userToken: params.userToken, + targets: derivedTargets.targets, + }); + if (!issued.success) { + throw ExecutionError.invalidRequest( + `Kilo session capability issuance failed (${issued.error.reason})` + ); + } + // The provider base URL baked into the capability's targets MUST match the + // baseURL the sandbox actually calls — otherwise the outbound interceptor's + // route classification rejects the request as `upstream_not_allowed`, even + // though the capability itself is valid. Never let getSaferEnvVars derive + // this independently from env.KILO_OPENROUTER_BASE. + return { + capability: issued.value.capability, + providerBaseUrl: derivedTargets.targets.providerBaseUrl, + }; + } + async buildWrapperSessionReadyAndPromptRequests( options: BuildWrapperSessionReadyAndPromptRequestsOptions ): Promise< @@ -1834,6 +1907,19 @@ export class SessionService { if (!metadata.auth.kiloSessionId) { throw ExecutionError.invalidRequest('Missing kiloSessionId in session metadata'); } + const nextAuthSecret = await resolveSecret(env.NEXTAUTH_SECRET); + if (!nextAuthSecret) { + throw ExecutionError.invalidRequest('NEXTAUTH_SECRET is not configured on the worker'); + } + const { capability: kiloCapability, providerBaseUrl: kiloProviderBaseUrl } = + await this.issueKiloSessionCapability(env, { + userId, + cloudAgentSessionId: sessionId, + kiloSessionId: metadata.auth.kiloSessionId, + sandboxId, + managedScmContainment: metadata.workspace?.managedScmContainment === true, + userToken: metadata.auth.kilocodeToken, + }); const devcontainerRequested = metadata.workspace?.devcontainerRequested === true || metadata.devcontainer !== undefined; @@ -1876,7 +1962,8 @@ export class SessionService { sessionId, workspacePath, env, - originalToken: metadata.auth.kilocodeToken, + kiloCapability, + kiloProviderBaseUrl, kilocodeModel: agent.model, originalOrgId: orgId, githubToken: resolvedTokens.githubToken, @@ -1912,9 +1999,10 @@ export class SessionService { const repo = this.buildWrapperRepoSource(metadata, resolvedTokens); const session = buildWrapperSessionBinding({ workerUrl: env.WORKER_URL, - kilocodeToken: metadata.auth.kilocodeToken, + nextAuthSecret, userId, sessionId, + kiloSessionId: metadata.auth.kiloSessionId, wrapper, upstreamBranch: metadata.repository?.upstreamBranch, }); @@ -2068,6 +2156,15 @@ export class SessionService { if (!metadata.auth.kiloSessionId) { throw ExecutionError.invalidRequest('Missing kiloSessionId in session metadata'); } + const { capability: kiloCapability, providerBaseUrl: kiloProviderBaseUrl } = + await this.issueKiloSessionCapability(env, { + userId, + cloudAgentSessionId: sessionId, + kiloSessionId: metadata.auth.kiloSessionId, + sandboxId, + managedScmContainment: metadata.workspace?.managedScmContainment === true, + userToken: metadata.auth.kilocodeToken, + }); const resolvedTokens = await this.resolveWorkspaceTokens(env, metadata, sandboxId); const github = githubRepository(metadata); @@ -2122,7 +2219,8 @@ export class SessionService { const runtimeEnv = this.buildRuntimeEnv({ context, env, - originalToken: metadata.auth.kilocodeToken, + kiloCapability, + kiloProviderBaseUrl, kilocodeModel: options.kilocodeModel, originalOrgId: orgId, createdOnPlatform: metadata.identity.createdOnPlatform, @@ -2145,7 +2243,9 @@ export class SessionService { env, metadata, options.kilocodeModel, - orgId + orgId, + kiloCapability, + kiloProviderBaseUrl ); if ( !(await this.sanitizeBitbucketCodeReviewRemote(session, context.workspacePath, metadata)) @@ -2214,7 +2314,9 @@ export class SessionService { env, metadata, options.kilocodeModel, - orgId + orgId, + kiloCapability, + kiloProviderBaseUrl ); let devcontainer: DevContainerHandle | undefined; @@ -2227,7 +2329,7 @@ export class SessionService { await this.prepareBranch(session, workspacePath, branchName, metadata); await this.sanitizeBitbucketCodeReviewRemote(session, workspacePath, metadata); - await writeAuthFile(sandbox, sessionHome, metadata.auth.kilocodeToken); + await writeAuthFile(sandbox, sessionHome, kiloCapability); await writeGlobalRules(sandbox, sessionHome, sessionId); const detectedDevcontainer = metadata.workspace?.devcontainerRequested @@ -2282,7 +2384,7 @@ export class SessionService { devcontainer, dockerEnv, env, - kilocodeToken: metadata.auth.kilocodeToken, + kiloCapability, runtimeEnv, sessionHome, } @@ -2325,16 +2427,16 @@ export class SessionService { env: PersistenceEnv, metadata: CloudAgentSessionState, kilocodeModel: string | undefined, - orgId: string | undefined + orgId: string | undefined, + kiloCapability: string, + kiloProviderBaseUrl: string | undefined ): Promise { - if (!metadata.auth.kilocodeToken) { - throw ExecutionError.invalidRequest('Missing kilocodeToken in session metadata'); - } return this.getOrCreateSession({ sandbox, context, env, - originalToken: metadata.auth.kilocodeToken, + kiloCapability, + kiloProviderBaseUrl, kilocodeModel, originalOrgId: orgId, createdOnPlatform: metadata.identity.createdOnPlatform, @@ -2554,7 +2656,7 @@ export class SessionService { options: RestoreRuntimeOptions ): Promise { const restoreTokenFilePath = options.devcontainer - ? await writeRestoreTokenFile(sandbox, session, options.sessionHome, options.kilocodeToken) + ? await writeRestoreTokenFile(sandbox, session, options.sessionHome, options.kiloCapability) : undefined; const restoreCommand = buildRestoreCommand({ kiloSessionId, @@ -2631,7 +2733,7 @@ export class SessionService { : `/tmp/kilo-empty-session-${kiloSessionId}.json`; await sandbox.writeFile(importFilePath, minimalSessionJson); const restoreTokenFilePath = options.devcontainer - ? await writeRestoreTokenFile(sandbox, session, options.sessionHome, options.kilocodeToken) + ? await writeRestoreTokenFile(sandbox, session, options.sessionHome, options.kiloCapability) : undefined; const restoreCommand = buildRestoreCommand({ kiloSessionId, @@ -2747,7 +2849,8 @@ export type GetOrCreateSessionOptions = { sandbox: SandboxInstance; context: SessionContext; env: PersistenceEnv; - originalToken: string; + kiloCapability: string; + kiloProviderBaseUrl?: string; kilocodeModel?: string; originalOrgId?: string; createdOnPlatform?: string; @@ -2762,7 +2865,7 @@ type RestoreRuntimeOptions = { devcontainer?: DevContainerHandle; dockerEnv?: Record; env: PersistenceEnv; - kilocodeToken: string; + kiloCapability: string; runtimeEnv: Record; sessionHome: string; }; @@ -2772,7 +2875,8 @@ type GetSaferEnvVarsOptions = { sessionId: string; workspacePath: string; env: PersistenceEnv; - originalToken: string; + kiloCapability: string; + kiloProviderBaseUrl?: string; kilocodeModel?: string; originalOrgId?: string; githubToken?: string; @@ -2828,13 +2932,15 @@ export type BuildWrapperSessionReadyAndPromptRequestsOptions = { function buildWrapperSessionBinding(options: { workerUrl?: string; - kilocodeToken: string; + nextAuthSecret: string; userId: string; sessionId: string; + kiloSessionId: string; wrapper: FencedWrapperDispatchRequest['wrapper']; upstreamBranch?: string; }): WrapperSessionReadyRequest['session'] { - const { workerUrl, kilocodeToken, userId, sessionId, wrapper, upstreamBranch } = options; + const { workerUrl, nextAuthSecret, userId, sessionId, kiloSessionId, wrapper, upstreamBranch } = + options; if (!workerUrl) { throw ExecutionError.invalidRequest('WORKER_URL is required for wrapper bootstrap'); } @@ -2843,9 +2949,22 @@ function buildWrapperSessionBinding(options: { wsUrl.protocol = wsUrl.protocol === 'https:' ? 'wss:' : 'ws:'; wsUrl.pathname = `/sessions/${encodeURIComponent(userId)}/${sessionId}/ingest`; + const workerAuthToken = mintWrapperDispatchTicket( + { + type: 'wrapper_dispatch_ticket', + userId, + cloudAgentSessionId: sessionId, + kiloSessionId, + wrapperRunId: wrapper.fence.wrapperRunId, + wrapperGeneration: wrapper.fence.wrapperGeneration, + wrapperConnectionId: wrapper.fence.wrapperConnectionId, + }, + nextAuthSecret + ); + return { ingestUrl: wsUrl.toString(), - workerAuthToken: kilocodeToken, + workerAuthToken, wrapperRunId: wrapper.fence.wrapperRunId, wrapperGeneration: wrapper.fence.wrapperGeneration, wrapperConnectionId: wrapper.fence.wrapperConnectionId, 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 64370edf9c..30e7daddf0 100644 --- a/services/cloud-agent-next/src/session/agent-runtime.test.ts +++ b/services/cloud-agent-next/src/session/agent-runtime.test.ts @@ -329,7 +329,14 @@ describe('AgentRuntime', () => { } as unknown as AgentSandbox; const runtime = createAgentRuntime({ storage, - env: { WORKER_URL: 'https://worker.example.com' } as Env, + env: { + WORKER_URL: 'https://worker.example.com', + NEXTAUTH_SECRET: 'test-secret', + SandboxSmall: { idFromName: () => ({ toString: () => 'container-id' }) }, + GIT_TOKEN_SERVICE: { + issueKiloSessionCapability: async () => ({ success: true, capability: 'kka1.test' }), + }, + } as unknown as Env, getMetadata: async () => createMetadata(), getSessionIdForLogs: () => 'agent_runtime', sendToWrapper: () => false, @@ -388,7 +395,14 @@ describe('AgentRuntime', () => { } as unknown as AgentSandbox; const runtime = createAgentRuntime({ storage, - env: { WORKER_URL: 'https://worker.example.com' } as Env, + env: { + WORKER_URL: 'https://worker.example.com', + NEXTAUTH_SECRET: 'test-secret', + SandboxSmall: { idFromName: () => ({ toString: () => 'container-id' }) }, + GIT_TOKEN_SERVICE: { + issueKiloSessionCapability: async () => ({ success: true, capability: 'kka1.test' }), + }, + } as unknown as Env, getMetadata: async () => createMetadata(), getSessionIdForLogs: () => 'agent_runtime', sendToWrapper: () => false, @@ -445,7 +459,14 @@ describe('AgentRuntime', () => { } as unknown as AgentSandbox; const runtime = createAgentRuntime({ storage, - env: { WORKER_URL: 'https://worker.example.com' } as Env, + env: { + WORKER_URL: 'https://worker.example.com', + NEXTAUTH_SECRET: 'test-secret', + SandboxSmall: { idFromName: () => ({ toString: () => 'container-id' }) }, + GIT_TOKEN_SERVICE: { + issueKiloSessionCapability: async () => ({ success: true, capability: 'kka1.test' }), + }, + } as unknown as Env, getMetadata: async () => createMetadata(), getSessionIdForLogs: () => 'agent_runtime', sendToWrapper: () => false, diff --git a/services/cloud-agent-next/src/types.ts b/services/cloud-agent-next/src/types.ts index 90d6424a1b..0a97ba67ce 100644 --- a/services/cloud-agent-next/src/types.ts +++ b/services/cloud-agent-next/src/types.ts @@ -350,6 +350,42 @@ type GetBitbucketTokenResult = | { success: true; token: string } | { success: false; reason: BitbucketTokenFailureReason }; +export type KiloSessionCapabilityTargets = { + backendBaseUrl: string; + providerBaseUrl: string; + sessionIngestBaseUrl: string; +}; + +export type KiloCapabilityRouteClass = + | 'provider_model' + | 'organization_models' + | 'backend_api' + | 'session_ingest'; + +type IssueKiloSessionCapabilityResult = + | { success: true; capability: string } + | { + success: false; + reason: + | 'invalid_targets' + | 'invalid_capability' + | 'expired_capability' + | 'capability_configuration_error'; + }; + +type RedeemKiloSessionCapabilityResult = + | { success: true; authorization: string; routeClass: KiloCapabilityRouteClass } + | { + success: false; + reason: + | 'invalid_capability' + | 'expired_capability' + | 'capability_configuration_error' + | 'container_mismatch' + | 'invalid_upstream_url' + | 'upstream_not_allowed'; + }; + export type GitTokenService = { getTokenForRepo(params: { githubRepo: string; @@ -396,6 +432,21 @@ export type GitTokenService = { requestMethod: string; requestUrl: string; }): Promise; + issueKiloSessionCapability(params: { + userId: string; + cloudAgentSessionId: string; + kiloSessionId: string; + outboundContainerId: string; + userToken: string; + targets: KiloSessionCapabilityTargets; + }): Promise; + redeemKiloSessionCapability(params: { + capability: string; + outboundContainerId: string; + requestMethod: string; + requestUrl: string; + bootstrapKiloSessionId?: string; + }): Promise; }; export type Env = { diff --git a/services/cloud-agent-next/worker-configuration.d.ts b/services/cloud-agent-next/worker-configuration.d.ts index 6fffe251f1..400834613b 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: a412e8ebdbae3b406ab482df8ffb606e) +// Generated by Wrangler by running `wrangler types` (hash: 34a9a79845e310ce4d99ebec7c8db838) // Runtime types generated with workerd@1.20260603.1 2026-06-03 nodejs_compat interface __BaseEnv_Env { SHARED_SANDBOX_OVERRIDES: KVNamespace; @@ -16,8 +16,8 @@ interface __BaseEnv_Env { BACKUP_BUCKET_NAME: "kilocode-sessions-dev" | "kilocode-sessions"; CLOUDFLARE_R2_ACCOUNT_ID: "e115e769bcdd4c3d66af59d3332cb394"; PER_SESSION_SANDBOX_ORG_IDS?: "*"; - MANAGED_SCM_CONTAINMENT_ORG_IDS: ""; - REPO_SNAPSHOT_ORG_IDS: ""; + MANAGED_SCM_CONTAINMENT_ORG_IDS?: "*"; + REPO_SNAPSHOT_ORG_IDS?: ""; NEXTAUTH_SECRET: string; KILO_SESSION_INGEST_URL: string; WORKER_URL: string; @@ -75,7 +75,7 @@ declare namespace Cloudflare { BACKUP_BUCKET_NAME: "kilocode-sessions-dev"; CLOUDFLARE_R2_ACCOUNT_ID: "e115e769bcdd4c3d66af59d3332cb394"; PER_SESSION_SANDBOX_ORG_IDS: "*"; - MANAGED_SCM_CONTAINMENT_ORG_IDS: ""; + MANAGED_SCM_CONTAINMENT_ORG_IDS: "*"; REPO_SNAPSHOT_ORG_IDS: ""; NEXTAUTH_SECRET: string; KILO_SESSION_INGEST_URL: string; diff --git a/services/cloud-agent-next/wrangler.jsonc b/services/cloud-agent-next/wrangler.jsonc index 4680452e6b..86a672052c 100644 --- a/services/cloud-agent-next/wrangler.jsonc +++ b/services/cloud-agent-next/wrangler.jsonc @@ -406,7 +406,7 @@ "WS_ALLOWED_ORIGINS": "http://localhost:3000,http://host.docker.internal:3000", "KILO_SESSION_INGEST_URL": "http://localhost:8800", "PER_SESSION_SANDBOX_ORG_IDS": "*", - "MANAGED_SCM_CONTAINMENT_ORG_IDS": "", + "MANAGED_SCM_CONTAINMENT_ORG_IDS": "*", "REPO_SNAPSHOT_ORG_IDS": "", "TOOL_CGROUP_ORG_IDS": "", }, diff --git a/services/cloud-agent-next/wrapper/src/log-uploader.test.ts b/services/cloud-agent-next/wrapper/src/log-uploader.test.ts new file mode 100644 index 0000000000..e7c2af03dd --- /dev/null +++ b/services/cloud-agent-next/wrapper/src/log-uploader.test.ts @@ -0,0 +1,140 @@ +import { afterEach, describe, expect, it } from 'bun:test'; +import fs from 'node:fs'; +import fsp from 'node:fs/promises'; +import os from 'node:os'; +import path from 'node:path'; +import { createLogUploader } from './log-uploader'; + +const originalFetch = globalThis.fetch; +const temporaryDirectories: string[] = []; + +afterEach(async () => { + globalThis.fetch = originalFetch; + await Promise.all( + temporaryDirectories + .splice(0) + .map(directory => fsp.rm(directory, { recursive: true, force: true })) + ); +}); + +describe('createLogUploader', () => { + it('aborts an active upload when stopped', async () => { + const directory = await fsp.mkdtemp(path.join(os.tmpdir(), 'log-uploader-stop-test-')); + temporaryDirectories.push(directory); + const wrapperLogPath = path.join(directory, 'wrapper.log'); + await fsp.writeFile(wrapperLogPath, 'wrapper log'); + + let requestSignal: AbortSignal | undefined; + globalThis.fetch = Object.assign( + async (_input: string | URL | Request, init?: RequestInit) => { + requestSignal = init?.signal ?? undefined; + return new Promise((_resolve, reject) => { + requestSignal?.addEventListener('abort', () => reject(requestSignal?.reason), { + once: true, + }); + }); + }, + { preconnect: originalFetch.preconnect } + ); + const uploader = createLogUploader({ + workerBaseUrl: 'https://worker.example.com', + sessionId: 'agent-session', + getKiloSessionId: () => 'kilo-session', + executionId: 'session', + userId: 'user', + getWorkerAuthToken: () => 'kka1.opaque', + cliLogDir: path.join(directory, 'missing-cli-logs'), + wrapperLogPath, + }); + + const upload = uploader.uploadNow(); + while (!requestSignal) await Bun.sleep(1); + uploader.stop(); + const settled = await Promise.race([upload.then(() => true), Bun.sleep(100).then(() => false)]); + + expect(requestSignal.aborted).toBe(true); + expect(settled).toBe(true); + }); + + it('binds log uploads to the Kilo session and sends the opaque worker credential', async () => { + const directory = await fsp.mkdtemp(path.join(os.tmpdir(), 'log-uploader-test-')); + temporaryDirectories.push(directory); + const cliLogDir = path.join(directory, 'cli-logs'); + const wrapperLogPath = path.join(directory, 'wrapper.log'); + await fsp.mkdir(cliLogDir); + await fsp.writeFile(path.join(cliLogDir, 'kilo.log'), 'kilo log'); + await fsp.writeFile(wrapperLogPath, 'wrapper log'); + + let capturedUrl: URL | undefined; + let capturedInit: RequestInit | undefined; + globalThis.fetch = Object.assign( + async (input: string | URL | Request, init?: RequestInit) => { + capturedUrl = new URL(input instanceof Request ? input.url : input.toString()); + capturedInit = init; + return new Response(null, { status: 204 }); + }, + { preconnect: originalFetch.preconnect } + ); + + const uploader = createLogUploader({ + workerBaseUrl: 'https://worker.example.com', + sessionId: 'agent/session', + getKiloSessionId: () => 'kilo/session?one', + executionId: 'session', + userId: 'user@example.com', + getWorkerAuthToken: () => 'kka1.opaque', + cliLogDir, + wrapperLogPath, + }); + + await uploader.uploadNow(); + + expect(capturedUrl?.pathname).toBe( + '/sessions/user%40example.com/agent%2Fsession/logs/session/logs.tar.gz' + ); + expect(capturedUrl?.searchParams.get('kiloSessionId')).toBe('kilo/session?one'); + expect(new Headers(capturedInit?.headers).get('Authorization')).toBe('Bearer kka1.opaque'); + expect(capturedInit?.body).toBeInstanceOf(ReadableStream); + expect(fs.existsSync(wrapperLogPath)).toBe(true); + }); + + it('reads the worker auth token fresh on every upload instead of a value captured at creation', async () => { + const directory = await fsp.mkdtemp(path.join(os.tmpdir(), 'log-uploader-refresh-test-')); + temporaryDirectories.push(directory); + const cliLogDir = path.join(directory, 'cli-logs'); + const wrapperLogPath = path.join(directory, 'wrapper.log'); + await fsp.mkdir(cliLogDir); + await fsp.writeFile(path.join(cliLogDir, 'kilo.log'), 'kilo log'); + await fsp.writeFile(wrapperLogPath, 'wrapper log'); + + const capturedAuthHeaders: Array = []; + globalThis.fetch = Object.assign( + async (_input: string | URL | Request, init?: RequestInit) => { + capturedAuthHeaders.push(new Headers(init?.headers).get('Authorization')); + return new Response(null, { status: 204 }); + }, + { preconnect: originalFetch.preconnect } + ); + + let currentToken = 'kka1.first-ticket'; + const uploader = createLogUploader({ + workerBaseUrl: 'https://worker.example.com', + sessionId: 'agent-session', + getKiloSessionId: () => 'kilo-session', + executionId: 'session', + userId: 'user', + getWorkerAuthToken: () => currentToken, + cliLogDir, + wrapperLogPath, + }); + + await uploader.uploadNow(); + currentToken = 'kka1.refreshed-ticket'; + await uploader.uploadNow(); + + expect(capturedAuthHeaders).toEqual([ + 'Bearer kka1.first-ticket', + 'Bearer kka1.refreshed-ticket', + ]); + }); +}); diff --git a/services/cloud-agent-next/wrapper/src/log-uploader.ts b/services/cloud-agent-next/wrapper/src/log-uploader.ts index 5f415fa37e..8c7b0b1966 100644 --- a/services/cloud-agent-next/wrapper/src/log-uploader.ts +++ b/services/cloud-agent-next/wrapper/src/log-uploader.ts @@ -6,9 +6,21 @@ import { logToFile } from './utils.js'; type LogUploaderOpts = { workerBaseUrl: string; sessionId: string; + /** + * Read at upload time rather than captured once at creation: a later session + * bind can carry a different kiloSessionId, and the Worker rejects uploads + * whose query param disagrees with the ticket's kiloSessionId claim. + */ + getKiloSessionId: () => string; executionId: string; userId: string; - workerAuthToken: string; + /** + * Reads the current wrapper dispatch ticket at upload time rather than a + * value captured once at creation. The uploader is created once per wrapper + * process and outlives the ticket's own (shorter) lifetime, which is + * refreshed on every subsequent session bind. + */ + getWorkerAuthToken: () => string; /** Directory containing CLI log files (e.g. ~/.local/share/kilo/log/) */ cliLogDir: string; wrapperLogPath: string; @@ -82,6 +94,7 @@ function createTarStream(paths: Array): TarStream | undefined { export function createLogUploader(opts: LogUploaderOpts): LogUploader { let intervalId: ReturnType | undefined; let isUploading = false; + let activeAbort: AbortController | undefined; async function uploadNow(): Promise { if (isUploading) return; @@ -93,12 +106,16 @@ export function createLogUploader(opts: LogUploaderOpts): LogUploader { } const abort = new AbortController(); + activeAbort = abort; const timer = setTimeout(() => abort.abort(), UPLOAD_TIMEOUT_MS); try { - const url = `${opts.workerBaseUrl}/sessions/${encodeURIComponent(opts.userId)}/${encodeURIComponent(opts.sessionId)}/logs/${encodeURIComponent(opts.executionId)}/logs.tar.gz`; + const url = new URL( + `${opts.workerBaseUrl}/sessions/${encodeURIComponent(opts.userId)}/${encodeURIComponent(opts.sessionId)}/logs/${encodeURIComponent(opts.executionId)}/logs.tar.gz` + ); + url.searchParams.set('kiloSessionId', opts.getKiloSessionId()); const response = await fetch(url, { method: 'PUT', - headers: { Authorization: `Bearer ${opts.workerAuthToken}` }, + headers: { Authorization: `Bearer ${opts.getWorkerAuthToken()}` }, body: tar.stream, // @ts-expect-error -- Node/Bun fetch supports duplex for streaming request bodies duplex: 'half', @@ -117,6 +134,7 @@ export function createLogUploader(opts: LogUploaderOpts): LogUploader { clearTimeout(timer); tar.kill(); isUploading = false; + if (activeAbort === abort) activeAbort = undefined; } } @@ -132,6 +150,8 @@ export function createLogUploader(opts: LogUploaderOpts): LogUploader { clearInterval(intervalId); intervalId = undefined; } + activeAbort?.abort(); + activeAbort = undefined; } return { start, uploadNow, stop }; diff --git a/services/cloud-agent-next/wrapper/src/main.ts b/services/cloud-agent-next/wrapper/src/main.ts index d75edcf0a7..39bc1d49ae 100644 --- a/services/cloud-agent-next/wrapper/src/main.ts +++ b/services/cloud-agent-next/wrapper/src/main.ts @@ -594,7 +594,8 @@ async function main() { request.session, serverConfig, serverDeps, - 'close-until-runtime-ready' + 'close-until-runtime-ready', + request.kiloSessionId ); if (bindError) { const error = (await bindError.json()) as { diff --git a/services/cloud-agent-next/wrapper/src/server.test.ts b/services/cloud-agent-next/wrapper/src/server.test.ts index 197f9fbb75..acdd135cdf 100644 --- a/services/cloud-agent-next/wrapper/src/server.test.ts +++ b/services/cloud-agent-next/wrapper/src/server.test.ts @@ -549,6 +549,46 @@ describe('wrapper Kilo proxy route', () => { }); describe('wrapper session binding', () => { + it('binds the kiloSessionId supplied by the caller even when config has not been updated yet', async () => { + // A freshly bootstrapped wrapper's ServerConfig.sessionId starts out empty + // and is only set by the caller after the ready request's binding is + // processed. bindSessionContext must use the id the caller already knows + // (from the ready request), not a stale config value, since that id also + // seeds the log uploader's kiloSessionId for the wrapper's entire life. + const state = new WrapperState(); + + const response = await bindSessionContext( + { + ingestUrl: 'ws://worker.test/ingest', + workerAuthToken: 'worker-token', + wrapperRunId: 'run_1', + wrapperGeneration: 1, + wrapperConnectionId: 'conn_1', + }, + { + port: 5000, + workspacePath: '/workspace/repo', + version: 'test', + sessionId: '', + agentSessionId: 'agent_00000000-0000-0000-0000-000000000000', + userId: 'user_test', + }, + { + state, + kiloClient: {} as WrapperKiloClient, + openConnection: async () => {}, + closeConnection: async () => {}, + setAborted: () => {}, + resetLifecycle: () => {}, + }, + 'close-until-runtime-ready', + 'kilo_sess_real' + ); + + expect(response).toBeNull(); + expect(state.currentSession?.kiloSessionId).toBe('kilo_sess_real'); + }); + it('rejects even the current binding while the wrapper is finalizing', async () => { const state = new WrapperState(); const sessionBinding = { diff --git a/services/cloud-agent-next/wrapper/src/server.ts b/services/cloud-agent-next/wrapper/src/server.ts index fbc2f12765..c933657018 100644 --- a/services/cloud-agent-next/wrapper/src/server.ts +++ b/services/cloud-agent-next/wrapper/src/server.ts @@ -270,9 +270,17 @@ export async function bindSessionContext( binding: SessionBinding | undefined, config: ServerConfig, deps: ServerDependencies, - feedPolicy: SessionBoundFeedPolicy = 'restart' + feedPolicy: SessionBoundFeedPolicy = 'restart', + /** + * The kiloSessionId to bind, when the caller already knows it (e.g. from a + * fresh session/ready request) but hasn't yet written it back to `config`. + * Falls back to `config.sessionId` for callers that rebind an + * already-bootstrapped wrapper, where `config.sessionId` is already current. + */ + kiloSessionIdOverride?: string ): Promise { const { state } = deps; + const kiloSessionId = kiloSessionIdOverride ?? config.sessionId; const blockedWrapperRunId = state.finalizingWrapperRunId; const isFreshRunAfterFinalization = state.admissionsBlocked && @@ -325,7 +333,7 @@ export async function bindSessionContext( deps.resetLifecycle(); const sessionContext: SessionContext = { - kiloSessionId: config.sessionId, + kiloSessionId, ingestUrl: binding.ingestUrl, ingestToken: binding.ingestToken, workerAuthToken: binding.workerAuthToken, @@ -342,21 +350,22 @@ export async function bindSessionContext( const logUploader = createLogUploader({ workerBaseUrl, sessionId: config.agentSessionId, + getKiloSessionId: () => state.currentSession?.kiloSessionId ?? kiloSessionId, executionId: 'session', userId: config.userId, - workerAuthToken: binding.workerAuthToken, + getWorkerAuthToken: () => state.currentSession?.workerAuthToken ?? binding.workerAuthToken, cliLogDir, wrapperLogPath, }); state.setLogUploader(logUploader); logUploader.start(); - logToFile(`session bound: sessionId=${config.sessionId}`); + logToFile(`session bound: sessionId=${kiloSessionId}`); await notifySessionBound(deps, feedPolicy); return null; } const sessionContext: SessionContext = { - kiloSessionId: config.sessionId, + kiloSessionId, ingestUrl: binding.ingestUrl, ingestToken: binding.ingestToken, workerAuthToken: binding.workerAuthToken, 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 2751dfbe75..0fa97b1d53 100644 --- a/services/cloud-agent-next/wrapper/src/session-bootstrap.test.ts +++ b/services/cloud-agent-next/wrapper/src/session-bootstrap.test.ts @@ -40,7 +40,7 @@ function makeRequest(tmpDir: string, overrides: Partial { expect( fs.existsSync(path.join(request.workspace.workspacePath, '.git', 'kilo-bootstrap-complete')) ).toBe(true); + const authFile = await fsp.readFile( + path.join(request.workspace.sessionHome, '.local/share/kilo/auth.json'), + 'utf8' + ); + expect(JSON.parse(authFile)).toEqual({ kilo: { type: 'api', key: 'kilo-capability' } }); + expect(authFile).not.toContain('wrapper-dispatch-ticket'); }); it('uses activity watchdogs and reports sanitized progress for long git operations', async () => { @@ -1051,10 +1057,16 @@ describe('prepareWrapperBootstrapWorkspace', () => { refreshRemote: true, }, materialized: { - env: { KILO_PLATFORM: 'code-review' }, + env: { KILO_PLATFORM: 'code-review', KILOCODE_TOKEN: 'kilo-capability' }, }, }); await createCompleteGitWorkspace(request.workspace.workspacePath); + const authPath = path.join(request.workspace.sessionHome, '.local/share/kilo/auth.json'); + await fsp.mkdir(path.dirname(authPath), { recursive: true }); + await fsp.writeFile( + authPath, + JSON.stringify({ kilo: { type: 'api', key: 'stale-capability' } }) + ); const gitCalls: string[][] = []; await prepareWrapperBootstrapWorkspace(request, undefined, { @@ -1068,6 +1080,9 @@ describe('prepareWrapperBootstrapWorkspace', () => { ['remote', 'set-url', 'origin', 'https://bitbucket.org/acme/repo.git'], ]); expect(gitCalls.at(-1)?.join(' ')).not.toContain('bitbucket-token'); + expect(JSON.parse(await fsp.readFile(authPath, 'utf8'))).toEqual({ + kilo: { type: 'api', key: 'kilo-capability' }, + }); }); it('refreshes a warm GitHub remote, author, and selected CLI credential', async () => { @@ -1086,7 +1101,7 @@ describe('prepareWrapperBootstrapWorkspace', () => { refreshRemote: true, }, materialized: { - env: { GH_TOKEN: 'user-token' }, + env: { GH_TOKEN: 'user-token', KILOCODE_TOKEN: 'kilo-capability' }, }, }); await createCompleteGitWorkspace(request.workspace.workspacePath); @@ -1124,7 +1139,7 @@ describe('prepareWrapperBootstrapWorkspace', () => { runProcess: async (command, args) => { events.push(`process:${command} ${args.join(' ')}`); expect(process.env.HOME).toBe(request.workspace.sessionHome); - expect(process.env.KILOCODE_TOKEN).toBe('kilo-token'); + expect(process.env.KILOCODE_TOKEN).toBe('kilo-capability'); expect(process.env[PNPM_STORE_ENV_VAR]).toBe(PNPM_STORE_DIR); expect( fs.existsSync(path.join(request.workspace.sessionHome, '.local/share/kilo/auth.json')) diff --git a/services/cloud-agent-next/wrapper/src/session-bootstrap.ts b/services/cloud-agent-next/wrapper/src/session-bootstrap.ts index 494fe25bc8..a51bac5708 100644 --- a/services/cloud-agent-next/wrapper/src/session-bootstrap.ts +++ b/services/cloud-agent-next/wrapper/src/session-bootstrap.ts @@ -464,11 +464,16 @@ async function refreshGitRemoteToken( } async function writeSessionAuthFile(request: WrapperSessionReadyRequest): Promise { + const kilocodeToken = request.materialized.env.KILOCODE_TOKEN; + if (!kilocodeToken) { + throw new Error('KILOCODE_TOKEN is required to write the Kilo auth file'); + } + const authFilePath = sessionAuthFilePath(request.workspace.sessionHome); await fs.mkdir(path.dirname(authFilePath), { recursive: true }); await fs.writeFile( authFilePath, - JSON.stringify({ kilo: { type: 'api', key: request.session.workerAuthToken } }, null, 2) + JSON.stringify({ kilo: { type: 'api', key: kilocodeToken } }, null, 2) ); } @@ -786,6 +791,8 @@ async function prepareWrapperBootstrapWorkspaceWithinDeadline( logToFile(`bootstrap cold workspace clone ready kiloSessionId=${request.kiloSessionId}`); } + await writeSessionAuthFile(request); + if (workspaceNeedsBootstrap) { progress?.('branch', 'Setting up branch...'); logToFile( @@ -808,7 +815,6 @@ async function prepareWrapperBootstrapWorkspaceWithinDeadline( ); await sanitizeBitbucketCodeReviewRemote(request, runGit); - await writeSessionAuthFile(request); await writeRuntimeSkills(request); progress?.( From 775c382e37bed18f872e6a6674a8ef024dc0d1d2 Mon Sep 17 00:00:00 2001 From: Evgeny Shurakov Date: Thu, 9 Jul 2026 22:28:14 +0200 Subject: [PATCH 2/2] fix(cloud-agent-next): address review feedback on Kilo capability consumption - Materialize the derived sessionIngestBaseUrl into normal and devcontainer runtime env so contained restore/import/ingest calls match the interceptor origin baked into the capability, not the raw localhost value - Accept the legacy log-upload URL without kiloSessionId for legacy_kilo_token claims; resolve the authoritative kiloSessionId from current session access, keeping the query/claim requirement for dispatch tickets - Map transient Kilo capability issuance failures (rpc_error, service_not_configured, temporarily_unavailable) to retryable workspaceSetupFailed instead of terminal invalidRequest --- services/cloud-agent-next/src/server.test.ts | 23 +++++ services/cloud-agent-next/src/server.ts | 10 +- .../src/session-service.test.ts | 22 +++++ .../cloud-agent-next/src/session-service.ts | 94 ++++++++++++------- 4 files changed, 114 insertions(+), 35 deletions(-) diff --git a/services/cloud-agent-next/src/server.test.ts b/services/cloud-agent-next/src/server.test.ts index faf63a11cf..3380cdba1b 100644 --- a/services/cloud-agent-next/src/server.test.ts +++ b/services/cloud-agent-next/src/server.test.ts @@ -829,6 +829,29 @@ describe('server wrapper log upload route', () => { expect(env.R2_BUCKET.put).toHaveBeenCalledOnce(); }); + it('accepts the old legacy log upload URL without a kiloSessionId query parameter', async () => { + const env = createLogEnv(); + const token = signKiloToken('usr_feed'); + + const response = await fetchWorker( + new Request('http://worker.test/sessions/usr_feed/agent_live/logs/session/logs.tar.gz', { + method: 'PUT', + headers: { Authorization: `Bearer ${token}` }, + body: new Uint8Array([1, 2, 3]), + }), + env + ); + + expect(response.status).toBe(204); + expect(requireCurrentSessionAccessMock).toHaveBeenCalledWith({ + env, + kiloUserId: 'usr_feed', + cloudAgentSessionId: 'agent_live', + }); + expect(env.R2_BUCKET.put).toHaveBeenCalledOnce(); + expect(env.R2_BUCKET.put.mock.calls[0][0]).toBe('logs/usr_feed/agent_live/session/logs.tar.gz'); + }); + it('rejects an upload missing the kiloSessionId parameter', async () => { const env = createLogEnv(); const ticket = signWrapperDispatchTicket(); diff --git a/services/cloud-agent-next/src/server.ts b/services/cloud-agent-next/src/server.ts index 30edcd46d2..a98e2a6f55 100644 --- a/services/cloud-agent-next/src/server.ts +++ b/services/cloud-agent-next/src/server.ts @@ -526,7 +526,7 @@ app.put( } const kiloSessionId = new URL(c.req.url).searchParams.get('kiloSessionId'); - if (!kiloSessionId) { + if (!kiloSessionId && authResult.claims.type === 'wrapper_dispatch_ticket') { return c.text('Missing kiloSessionId parameter', 400); } @@ -540,12 +540,16 @@ app.put( } try { - await requireCurrentSessionAccess({ + const sessionAccess = await requireCurrentSessionAccess({ env: c.env, kiloUserId: userId, cloudAgentSessionId: sessionId, - expectedKiloSessionId: kiloSessionId, + expectedKiloSessionId: kiloSessionId ?? undefined, }); + const authoritativeKiloSessionId = kiloSessionId ?? sessionAccess.kiloSessionId; + if (!authoritativeKiloSessionId) { + return c.text('Missing kiloSessionId parameter', 400); + } } catch (error) { return projectSessionAccessHttpError(error); } diff --git a/services/cloud-agent-next/src/session-service.test.ts b/services/cloud-agent-next/src/session-service.test.ts index 555a36f81d..0131a23221 100644 --- a/services/cloud-agent-next/src/session-service.test.ts +++ b/services/cloud-agent-next/src/session-service.test.ts @@ -2159,6 +2159,7 @@ describe('SessionService.buildWrapperSessionReadyAndPromptRequests', () => { const env = createEnv(); env.WORKER_URL = 'https://cloud-agent.example.com'; env.KILO_OPENROUTER_BASE = 'https://openrouter.wrong-base.example.com'; + env.KILO_SESSION_INGEST_URL = 'http://localhost:8800/ingest/'; const issueKiloSessionCapability = vi.fn().mockResolvedValue({ success: true, capability: 'kka1.issued', @@ -2206,12 +2207,33 @@ describe('SessionService.buildWrapperSessionReadyAndPromptRequests', () => { expect(configContent.provider.kilo.options.baseURL).toBe( derivedTargets.targets.providerBaseUrl ); + expect(result.readyRequest.materialized.env.KILO_SESSION_INGEST_URL).toBe( + derivedTargets.targets.sessionIngestBaseUrl + ); expect(issueKiloSessionCapability).toHaveBeenCalledWith( expect.objectContaining({ targets: derivedTargets.targets }) ); expect(configContent.provider.kilo.options.baseURL).not.toBe( 'https://openrouter.wrong-base.example.com' ); + expect(result.readyRequest.materialized.env.KILO_SESSION_INGEST_URL).not.toBe( + 'http://localhost:8800/ingest/' + ); + }); + + it('treats transient Kilo capability issuance failures as retryable workspace setup failures', async () => { + await expect( + buildPromptWrapperRequests(createMetadata({ managedScmContainment: true }), env => { + if (!env.GIT_TOKEN_SERVICE) throw new Error('Expected GIT_TOKEN_SERVICE in test env'); + env.GIT_TOKEN_SERVICE.issueKiloSessionCapability = vi + .fn() + .mockRejectedValue(new Error('binding unavailable')); + }) + ).rejects.toMatchObject({ + code: 'WORKSPACE_SETUP_FAILED', + retryable: true, + message: 'Kilo session capability issuance failed (rpc_error)', + }); }); it('falls back to the raw Kilo token for DIND sandboxes, which have no outbound interceptor to redeem a capability against', async () => { diff --git a/services/cloud-agent-next/src/session-service.ts b/services/cloud-agent-next/src/session-service.ts index 5e25ab09dc..baf1e49f43 100644 --- a/services/cloud-agent-next/src/session-service.ts +++ b/services/cloud-agent-next/src/session-service.ts @@ -133,6 +133,14 @@ function gitLabTokenLookupFailureMessage(reason: string): string { } } +function isRetryableKiloCapabilityIssuanceFailure(reason: string): boolean { + return ( + reason === 'rpc_error' || + reason === 'service_not_configured' || + reason === 'temporarily_unavailable' + ); +} + // Keep in sync with: cloudflare-code-review-infra/src/code-review-orchestrator.ts // mkdir and touch are intentionally allowed for agent scratch space during analysis const CODE_REVIEW_ALLOWED_COMMANDS = [ @@ -1260,6 +1268,7 @@ export class SessionService { env: opts.env, kiloCapability: opts.kiloCapability, kiloProviderBaseUrl: opts.kiloProviderBaseUrl, + kiloSessionIngestBaseUrl: opts.kiloSessionIngestBaseUrl, kilocodeModel: opts.kilocodeModel, originalOrgId: opts.originalOrgId, githubToken: context.githubToken, @@ -1287,6 +1296,7 @@ export class SessionService { env, kiloCapability, kiloProviderBaseUrl, + kiloSessionIngestBaseUrl, kilocodeModel, originalOrgId, githubToken, @@ -1603,7 +1613,9 @@ export class SessionService { envVars.KILO_API_URL = sandboxUrl; } - if (env.KILO_SESSION_INGEST_URL) { + if (kiloSessionIngestBaseUrl) { + envVars.KILO_SESSION_INGEST_URL = kiloSessionIngestBaseUrl; + } else if (env.KILO_SESSION_INGEST_URL) { envVars.KILO_SESSION_INGEST_URL = env.KILO_SESSION_INGEST_URL; } @@ -1842,7 +1854,7 @@ export class SessionService { managedScmContainment: boolean; userToken: string; } - ): Promise<{ capability: string; providerBaseUrl?: string }> { + ): Promise<{ capability: string; providerBaseUrl?: string; sessionIngestBaseUrl?: string }> { if (params.sandboxId.startsWith('dind-')) { return { capability: params.userToken }; } @@ -1869,18 +1881,19 @@ export class SessionService { targets: derivedTargets.targets, }); if (!issued.success) { - throw ExecutionError.invalidRequest( - `Kilo session capability issuance failed (${issued.error.reason})` - ); + const message = `Kilo session capability issuance failed (${issued.error.reason})`; + if (isRetryableKiloCapabilityIssuanceFailure(issued.error.reason)) { + throw ExecutionError.workspaceSetupFailed(message); + } + throw ExecutionError.invalidRequest(message); } - // The provider base URL baked into the capability's targets MUST match the - // baseURL the sandbox actually calls — otherwise the outbound interceptor's - // route classification rejects the request as `upstream_not_allowed`, even - // though the capability itself is valid. Never let getSaferEnvVars derive - // this independently from env.KILO_OPENROUTER_BASE. + // Capability targets MUST match the URLs the sandbox actually calls — + // otherwise the outbound interceptor rejects the request as + // `upstream_not_allowed`, even though the capability itself is valid. return { capability: issued.value.capability, providerBaseUrl: derivedTargets.targets.providerBaseUrl, + sessionIngestBaseUrl: derivedTargets.targets.sessionIngestBaseUrl, }; } @@ -1911,15 +1924,18 @@ export class SessionService { if (!nextAuthSecret) { throw ExecutionError.invalidRequest('NEXTAUTH_SECRET is not configured on the worker'); } - const { capability: kiloCapability, providerBaseUrl: kiloProviderBaseUrl } = - await this.issueKiloSessionCapability(env, { - userId, - cloudAgentSessionId: sessionId, - kiloSessionId: metadata.auth.kiloSessionId, - sandboxId, - managedScmContainment: metadata.workspace?.managedScmContainment === true, - userToken: metadata.auth.kilocodeToken, - }); + const { + capability: kiloCapability, + providerBaseUrl: kiloProviderBaseUrl, + sessionIngestBaseUrl: kiloSessionIngestBaseUrl, + } = await this.issueKiloSessionCapability(env, { + userId, + cloudAgentSessionId: sessionId, + kiloSessionId: metadata.auth.kiloSessionId, + sandboxId, + managedScmContainment: metadata.workspace?.managedScmContainment === true, + userToken: metadata.auth.kilocodeToken, + }); const devcontainerRequested = metadata.workspace?.devcontainerRequested === true || metadata.devcontainer !== undefined; @@ -1964,6 +1980,7 @@ export class SessionService { env, kiloCapability, kiloProviderBaseUrl, + kiloSessionIngestBaseUrl, kilocodeModel: agent.model, originalOrgId: orgId, githubToken: resolvedTokens.githubToken, @@ -2156,15 +2173,18 @@ export class SessionService { if (!metadata.auth.kiloSessionId) { throw ExecutionError.invalidRequest('Missing kiloSessionId in session metadata'); } - const { capability: kiloCapability, providerBaseUrl: kiloProviderBaseUrl } = - await this.issueKiloSessionCapability(env, { - userId, - cloudAgentSessionId: sessionId, - kiloSessionId: metadata.auth.kiloSessionId, - sandboxId, - managedScmContainment: metadata.workspace?.managedScmContainment === true, - userToken: metadata.auth.kilocodeToken, - }); + const { + capability: kiloCapability, + providerBaseUrl: kiloProviderBaseUrl, + sessionIngestBaseUrl: kiloSessionIngestBaseUrl, + } = await this.issueKiloSessionCapability(env, { + userId, + cloudAgentSessionId: sessionId, + kiloSessionId: metadata.auth.kiloSessionId, + sandboxId, + managedScmContainment: metadata.workspace?.managedScmContainment === true, + userToken: metadata.auth.kilocodeToken, + }); const resolvedTokens = await this.resolveWorkspaceTokens(env, metadata, sandboxId); const github = githubRepository(metadata); @@ -2221,6 +2241,7 @@ export class SessionService { env, kiloCapability, kiloProviderBaseUrl, + kiloSessionIngestBaseUrl, kilocodeModel: options.kilocodeModel, originalOrgId: orgId, createdOnPlatform: metadata.identity.createdOnPlatform, @@ -2245,7 +2266,8 @@ export class SessionService { options.kilocodeModel, orgId, kiloCapability, - kiloProviderBaseUrl + kiloProviderBaseUrl, + kiloSessionIngestBaseUrl ); if ( !(await this.sanitizeBitbucketCodeReviewRemote(session, context.workspacePath, metadata)) @@ -2316,7 +2338,8 @@ export class SessionService { options.kilocodeModel, orgId, kiloCapability, - kiloProviderBaseUrl + kiloProviderBaseUrl, + kiloSessionIngestBaseUrl ); let devcontainer: DevContainerHandle | undefined; @@ -2385,6 +2408,7 @@ export class SessionService { dockerEnv, env, kiloCapability, + kiloSessionIngestBaseUrl, runtimeEnv, sessionHome, } @@ -2429,7 +2453,8 @@ export class SessionService { kilocodeModel: string | undefined, orgId: string | undefined, kiloCapability: string, - kiloProviderBaseUrl: string | undefined + kiloProviderBaseUrl: string | undefined, + kiloSessionIngestBaseUrl: string | undefined ): Promise { return this.getOrCreateSession({ sandbox, @@ -2437,6 +2462,7 @@ export class SessionService { env, kiloCapability, kiloProviderBaseUrl, + kiloSessionIngestBaseUrl, kilocodeModel, originalOrgId: orgId, createdOnPlatform: metadata.identity.createdOnPlatform, @@ -2642,7 +2668,8 @@ export class SessionService { : undefined; return { KILOCODE_TOKEN_FILE: restoreTokenFilePath, - KILO_SESSION_INGEST_URL: options.env.KILO_SESSION_INGEST_URL, + KILO_SESSION_INGEST_URL: + options.kiloSessionIngestBaseUrl ?? options.env.KILO_SESSION_INGEST_URL, ...buildKiloSessionXdgEnv(options.sessionHome), ...(backendUrl ? { KILOCODE_BACKEND_BASE_URL: backendUrl, KILO_API_URL: backendUrl } : {}), }; @@ -2851,6 +2878,7 @@ export type GetOrCreateSessionOptions = { env: PersistenceEnv; kiloCapability: string; kiloProviderBaseUrl?: string; + kiloSessionIngestBaseUrl?: string; kilocodeModel?: string; originalOrgId?: string; createdOnPlatform?: string; @@ -2866,6 +2894,7 @@ type RestoreRuntimeOptions = { dockerEnv?: Record; env: PersistenceEnv; kiloCapability: string; + kiloSessionIngestBaseUrl?: string; runtimeEnv: Record; sessionHome: string; }; @@ -2877,6 +2906,7 @@ type GetSaferEnvVarsOptions = { env: PersistenceEnv; kiloCapability: string; kiloProviderBaseUrl?: string; + kiloSessionIngestBaseUrl?: string; kilocodeModel?: string; originalOrgId?: string; githubToken?: string;