@meistrari/agent-core 0.0.0 → 0.1.1

package/src/worker-runtime-client/connection-attempt.ts

@@ -0,0 +1,349 @@
+import type { Logger } from '../logger'
+import type { CommandAckEnvelope } from '../supervisor-protocol/envelopes/supervisor-to-control-plane'
+import type {
+    BootstrapProvider,
+    CommandSource,
+    ConnectionAuthorityStore,
+    ConnectionClaim,
+    ControlAuthoritySigner,
+    DurableEventSink,
+    EphemeralEventSink,
+    RpcDispatcher,
+    SandboxConnectionRuntimeTimings,
+    SessionFailureSink,
+    SupervisorConnection,
+    SupervisorConnectionOpener,
+} from './ports'
+import { randomUUID } from 'node:crypto'
+import { wireCommandBodySchema } from '../supervisor-protocol/command-body'
+import { encodeWireFrame } from '../supervisor-protocol/wire-codec'
+import { SessionCommandPump } from './command-pump'
+import { BootstrapRejectedAckError, SupervisorFrameProcessor } from './frame-processor'
+import { decryptToken } from './token-crypto'
+
+export class SessionSandboxConnectionAttempt {
+    private readonly controller = new AbortController()
+    private commandPump: SessionCommandPump | undefined
+    private processor: SupervisorFrameProcessor | undefined
+    private connection: SupervisorConnection | undefined
+    private ready = false
+    private readonly parentAbort: () => void
+
+    constructor(private readonly dependencies: {
+        claim: ConnectionClaim
+        establishmentDeadline: number
+        replicaId: string
+        trafficTokenEncryptionKey: string
+        signControlAuthority: ControlAuthoritySigner
+        openSupervisorConnection: SupervisorConnectionOpener
+        authorityStore: ConnectionAuthorityStore
+        commands: CommandSource
+        durableEvents: DurableEventSink
+        ephemeralEvents: EphemeralEventSink
+        rpc: RpcDispatcher
+        bootstrap: BootstrapProvider
+        failures: SessionFailureSink
+        timings: SandboxConnectionRuntimeTimings
+        logger: Logger
+        parentSignal: AbortSignal
+    }) {
+        this.parentAbort = () => this.stop(dependencies.parentSignal.reason)
+        dependencies.parentSignal.addEventListener('abort', this.parentAbort, { once: true })
+    }
+
+    async run(): Promise<never> {
+        const { claim } = this.dependencies
+        const renewalFailure = Promise.withResolvers<unknown>()
+        void this.renewLease().catch((error) => {
+            renewalFailure.resolve(error)
+            this.stop(error)
+        })
+        const runtimeConnectionAttemptId = randomUUID()
+        const eventAckFloor = await this.dependencies.durableEvents.getAckFloor({
+            sessionSandboxId: claim.sessionSandboxId,
+        })
+        const initialCommandPromise = this.dependencies.commands.initial({ sessionId: claim.sessionId })
+        const bootstrapBodyPromise = this.dependencies.bootstrap.loadBootstrapBody({ sessionId: claim.sessionId })
+        const secretsPromise = this.dependencies.bootstrap.issueEphemeralSecrets({
+            sessionId: claim.sessionId,
+            eventAckFloor,
+            signal: this.controller.signal,
+        })
+        const authorityAssertionPromise = this.dependencies.signControlAuthority({
+            providerSandboxId: claim.providerSandboxId,
+            sessionId: claim.sessionId,
+            sessionSandboxId: claim.sessionSandboxId,
+            connectionGeneration: claim.runtimeConnectionGeneration,
+            runtimeConnectionAttemptId,
+            eventAckFloor,
+        })
+        const trafficAccessToken = decryptToken({
+            encryptedToken: claim.trafficAccessTokenEncrypted,
+            secret: this.dependencies.trafficTokenEncryptionKey,
+        })
+
+        const queuedFrames: Array<string | ArrayBuffer | Uint8Array> = []
+        const closed = Promise.withResolvers<{ code: number, reason: string }>()
+        const authorityAssertion = await authorityAssertionPromise
+        const connection = await beforeDeadline(this.dependencies.openSupervisorConnection({
+            providerSandboxId: claim.providerSandboxId,
+            trafficAccessToken,
+            sessionId: claim.sessionId,
+            sessionSandboxId: claim.sessionSandboxId,
+            runtimeConnectionAttemptId,
+            authorityAssertion,
+            connectionGeneration: claim.runtimeConnectionGeneration,
+            eventAckFloor,
+            signal: this.controller.signal,
+            onMessage: frame => this.processor ? this.processor.receive(frame) : queuedFrames.push(frame),
+            onClose: close => closed.resolve(close),
+        }), Math.min(this.dependencies.establishmentDeadline, performance.now() + this.dependencies.timings.dialTimeoutMs), this.controller.signal, () => new ConnectionClosedError({
+            ready: false,
+            detail: 'Supervisor connection dial timed out.',
+        }))
+        this.connection = connection
+        let bootstrapAck: CommandAckEnvelope | undefined
+        const processor = new SupervisorFrameProcessor({
+            ref: {
+                sessionId: claim.sessionId,
+                sessionSandboxId: claim.sessionSandboxId,
+                providerSandboxId: claim.providerSandboxId,
+            },
+            generation: claim.runtimeConnectionGeneration,
+            runtimeConnectionId: claim.runtimeConnectionId,
+            runtimeConnectionAttemptId,
+            connection,
+            authorityStore: this.dependencies.authorityStore,
+            durableEvents: this.dependencies.durableEvents,
+            ephemeralEvents: this.dependencies.ephemeralEvents,
+            rpcDispatcher: this.dependencies.rpc,
+            failures: this.dependencies.failures,
+            logger: this.dependencies.logger,
+            onCommandAck: (ack) => {
+                if (this.commandPump)
+                    this.commandPump.acknowledge(ack)
+                else
+                    bootstrapAck = ack
+            },
+            onFatal: error => this.stop(error),
+        })
+        this.processor = processor
+        for (const frame of queuedFrames)
+            processor.receive(frame)
+
+        const [initialCommand, body, secrets] = await Promise.all([
+            initialCommandPromise,
+            bootstrapBodyPromise,
+            secretsPromise,
+        ])
+        if (!initialCommand || initialCommand.commandSeq !== 1)
+            throw new DeterministicEstablishmentError('agent-core.initial-command-invalid', 'Session command 1 is missing or invalid.')
+        const parsedInitialBody = wireCommandBodySchema.safeParse(initialCommand.body)
+        if (!parsedInitialBody.success) {
+            throw new DeterministicEstablishmentError(
+                'agent-core.initial-command-invalid',
+                'Session command 1 has an invalid body.',
+            )
+        }
+        const initialBody = parsedInitialBody.data
+        await connection.send(encodeWireFrame({
+            kind: 'session.bootstrap',
+            connectionGeneration: claim.runtimeConnectionGeneration,
+            body: {
+                ...body,
+                sessionId: claim.sessionId,
+                initialCommand: {
+                    commandId: initialCommand.commandId,
+                    commandSeq: initialCommand.commandSeq,
+                    body: initialBody,
+                },
+            },
+            ...secrets,
+        }))
+
+        const readiness = await beforeDeadline(Promise.race([
+            processor.waitForReadiness({ signal: this.controller.signal }),
+            closed.promise.then((close) => {
+                throw new ConnectionClosedError({ ready: false, detail: `${close.code}:${close.reason}` })
+            }),
+        ]), Math.min(this.dependencies.establishmentDeadline, performance.now() + 10_000), this.controller.signal, () => new ConnectionClosedError({
+            ready: false,
+            detail: 'Supervisor bootstrap timed out.',
+        })).catch(async (error: unknown) => {
+            if (error instanceof BootstrapRejectedAckError) {
+                await processor.acknowledgeBootstrapRejection(error.ack)
+                throw new DeterministicEstablishmentError(error.ack.errorCode, error.ack.detail, true)
+            }
+            throw error
+        })
+        const markedReady = await this.dependencies.authorityStore.markConnectionReady({
+            sessionSandboxId: claim.sessionSandboxId,
+            generation: claim.runtimeConnectionGeneration,
+            runtimeConnectionId: claim.runtimeConnectionId,
+            lastReceivedCommandSeq: readiness.lastReceivedCommandSeq,
+            agentRunSnapshot: readiness.agentRun,
+        })
+        if (!markedReady)
+            throw new ConnectionAuthorityLostError()
+        this.ready = true
+
+        const pump = new SessionCommandPump({
+            sessionId: claim.sessionId,
+            sessionSandboxId: claim.sessionSandboxId,
+            runtimeConnectionId: claim.runtimeConnectionId,
+            afterSequence: readiness.lastReceivedCommandSeq,
+            commands: this.dependencies.commands,
+            authorityStore: this.dependencies.authorityStore,
+            connection,
+            logger: this.dependencies.logger,
+            onRejected: async (ack) => {
+                await this.dependencies.failures.markSessionFailed({
+                    sessionId: claim.sessionId,
+                    sessionSandboxId: claim.sessionSandboxId,
+                    providerSandboxId: claim.providerSandboxId,
+                    code: ack.errorCode,
+                    detail: ack.detail,
+                })
+                connection.close(1008, 'Supervisor command was rejected.')
+            },
+        })
+        this.commandPump = pump
+        if (bootstrapAck)
+            pump.acknowledge(bootstrapAck)
+        pump.start()
+        const close = await Promise.race([
+            closed.promise,
+            aborted(this.controller.signal),
+            renewalFailure.promise.then((error) => { throw error }),
+        ])
+        throw new ConnectionClosedError({ ready: this.ready, detail: 'code' in close ? `${close.code}:${close.reason}` : undefined })
+    }
+
+    wakeCommands(): void {
+        this.commandPump?.wake()
+    }
+
+    stop(reason?: unknown): void {
+        this.dependencies.parentSignal.removeEventListener('abort', this.parentAbort)
+        if (!this.controller.signal.aborted)
+            this.controller.abort(reason ?? new DOMException('Connection attempt stopped.', 'AbortError'))
+        this.commandPump?.stop()
+        this.processor?.stop(reason)
+        this.connection?.close(1001, 'Connection attempt stopped.')
+    }
+
+    async drain(): Promise<void> {
+        await this.processor?.drain()
+    }
+
+    private async renewLease(): Promise<never> {
+        const { claim, timings } = this.dependencies
+        let leaseSafetyDeadline = monotonicDeadline(claim.leaseExpiresAt)
+        while (true) {
+            const untilSafetyDeadline = leaseSafetyDeadline - performance.now()
+            if (untilSafetyDeadline <= 0)
+                throw new ConnectionAuthorityLostError()
+            await delay(Math.min(timings.leaseRenewIntervalMs, untilSafetyDeadline), this.controller.signal)
+            if (performance.now() >= leaseSafetyDeadline)
+                throw new ConnectionAuthorityLostError()
+            const result = await beforeDeadline(this.dependencies.authorityStore.renewConnectionLease({
+                sessionSandboxId: claim.sessionSandboxId,
+                replicaId: this.dependencies.replicaId,
+                generation: claim.runtimeConnectionGeneration,
+                leaseDurationMs: timings.leaseDurationMs,
+                supervisorLivenessObserved: this.processor?.takeSupervisorLivenessObservation() ?? false,
+            }), leaseSafetyDeadline, this.controller.signal, () => new ConnectionAuthorityLostError())
+            if (!result.renewed)
+                throw new ConnectionAuthorityLostError()
+            leaseSafetyDeadline = monotonicDeadline(
+                result.leaseExpiresAt ?? new Date(Date.now() + timings.leaseDurationMs),
+            )
+            if (result.hasOutstandingCommand)
+                this.commandPump?.wake()
+        }
+    }
+}
+
+export class DeterministicEstablishmentError extends Error {
+    override readonly name = 'DeterministicEstablishmentError'
+    readonly code: string
+    readonly detail: string | undefined
+    readonly reported: boolean
+
+    constructor(code: string, detail?: string, reported = false) {
+        super(detail ?? code)
+        this.code = code
+        this.detail = detail
+        this.reported = reported
+    }
+}
+
+export class ConnectionAuthorityLostError extends Error {
+    override readonly name = 'ConnectionAuthorityLostError'
+}
+
+export class ConnectionClosedError extends Error {
+    override readonly name = 'ConnectionClosedError'
+    readonly ready: boolean
+
+    constructor(input: { ready: boolean, detail?: string }) {
+        super(input.detail ?? 'Supervisor connection closed.')
+        this.ready = input.ready
+    }
+}
+
+async function aborted(signal: AbortSignal): Promise<never> {
+    if (signal.aborted)
+        throw signal.reason
+    return await new Promise<never>((_resolve, reject) => {
+        signal.addEventListener('abort', () => reject(signal.reason), { once: true })
+    })
+}
+
+async function delay(milliseconds: number, signal: AbortSignal): Promise<void> {
+    if (signal.aborted)
+        throw signal.reason
+    await new Promise<void>((resolve, reject) => {
+        let timeout: ReturnType<typeof setTimeout>
+        const abort = () => {
+            clearTimeout(timeout)
+            signal.removeEventListener('abort', abort)
+            reject(signal.reason)
+        }
+        timeout = setTimeout(() => {
+            signal.removeEventListener('abort', abort)
+            resolve()
+        }, milliseconds)
+        signal.addEventListener('abort', abort, { once: true })
+    })
+}
+
+function monotonicDeadline(expiresAt: Date): number {
+    return performance.now() + Math.max(0, expiresAt.getTime() - Date.now())
+}
+
+async function beforeDeadline<T>(
+    task: Promise<T>,
+    deadline: number,
+    signal: AbortSignal,
+    timeoutError: () => Error,
+): Promise<T> {
+    signal.throwIfAborted()
+    const remaining = deadline - performance.now()
+    if (remaining <= 0)
+        throw timeoutError()
+    let timer: ReturnType<typeof setTimeout>
+    let abort: () => void
+    const guard = new Promise<never>((_resolve, reject) => {
+        abort = () => reject(signal.reason)
+        timer = setTimeout(() => reject(timeoutError()), remaining)
+        signal.addEventListener('abort', abort, { once: true })
+    })
+    try {
+        return await Promise.race([task, guard])
+    }
+    finally {
+        clearTimeout(timer!)
+        signal.removeEventListener('abort', abort!)
+    }
+}

package/src/worker-runtime-client/control-authority-signer.ts

@@ -0,0 +1,100 @@
+import type { SandboxControlAuthorityClaims, SandboxControlAuthorityPublicJwks } from '../supervisor-protocol/control-authority'
+import type { ControlAuthoritySigner } from './ports'
+import { SignJWT } from 'jose/jwt/sign'
+import { importJWK } from 'jose/key/import'
+import { z } from 'zod'
+import { InfrastructureError } from '../errors'
+import {
+    sandboxControlAuthorityAlgorithm,
+    sandboxControlAuthorityAssertionLifetimeSeconds,
+    sandboxControlAuthorityAudience,
+    sandboxControlAuthorityClaimsSchema,
+    sandboxControlAuthorityPublicJwksSchema,
+} from '../supervisor-protocol/control-authority'
+
+const base64UrlCoordinateSchema = z.string().regex(/^[\w-]+$/u)
+
+const privateJwkSchema = z.object({
+    kty: z.literal('EC'),
+    crv: z.literal('P-256'),
+    x: base64UrlCoordinateSchema,
+    y: base64UrlCoordinateSchema,
+    d: base64UrlCoordinateSchema,
+    kid: z.string().regex(/^[\w.-]{1,128}$/u),
+    alg: z.literal(sandboxControlAuthorityAlgorithm),
+    use: z.literal('sig'),
+    key_ops: z.tuple([z.literal('sign')]),
+}).strict()
+
+export class ControlAuthorityConfigurationError extends InfrastructureError {
+    constructor(input: { message: string, cause?: unknown }) {
+        super({
+            code: 'agent-core.control-authority-configuration-invalid',
+            message: input.message,
+            publicMessage: 'Sandbox control authority is not configured.',
+            retryable: false,
+            cause: input.cause,
+        })
+    }
+}
+
+export async function createControlAuthoritySigner(input: {
+    issuer: string
+    activeKeyId: string
+    signingJwk: string | Record<string, unknown>
+    publicJwks: string | SandboxControlAuthorityPublicJwks
+}): Promise<{ sign: ControlAuthoritySigner, publicJwks: SandboxControlAuthorityPublicJwks }> {
+    try {
+        const rawPrivateJwk = typeof input.signingJwk === 'string' ? JSON.parse(input.signingJwk) : input.signingJwk
+        const rawPublicJwks = typeof input.publicJwks === 'string' ? JSON.parse(input.publicJwks) : input.publicJwks
+        const privateJwk = privateJwkSchema.parse(rawPrivateJwk)
+        const publicJwks = sandboxControlAuthorityPublicJwksSchema.parse(rawPublicJwks)
+        if (privateJwk.kid !== input.activeKeyId)
+            configurationFailure('The active key ID does not match the signing key.')
+        const activePublicKey = publicJwks.keys.find(key => key.kid === input.activeKeyId)
+        if (!activePublicKey)
+            configurationFailure('The active signing key is absent from the verification keyset.')
+        if (activePublicKey.x !== privateJwk.x || activePublicKey.y !== privateJwk.y)
+            configurationFailure('The active signing key does not match its verification key.')
+        const signingKey = await importJWK(privateJwk, sandboxControlAuthorityAlgorithm)
+        if (!(signingKey instanceof CryptoKey))
+            configurationFailure('The control-authority signing key is not asymmetric.')
+
+        const sign: ControlAuthoritySigner = async (assertionInput) => {
+            const now = Math.floor(Date.now() / 1_000)
+            const claims = sandboxControlAuthorityClaimsSchema.parse({
+                iss: input.issuer,
+                aud: sandboxControlAuthorityAudience,
+                jti: assertionInput.runtimeConnectionAttemptId,
+                iat: now,
+                nbf: now,
+                exp: now + sandboxControlAuthorityAssertionLifetimeSeconds,
+                ...assertionInput,
+            })
+            return await signClaims({ claims, keyId: input.activeKeyId, signingKey })
+        }
+        return { sign, publicJwks }
+    }
+    catch (cause) {
+        if (cause instanceof ControlAuthorityConfigurationError)
+            throw cause
+        throw new ControlAuthorityConfigurationError({
+            message: 'Sandbox control-authority signing configuration is invalid.',
+            cause,
+        })
+    }
+}
+
+function configurationFailure(message: string): never {
+    throw new ControlAuthorityConfigurationError({ message })
+}
+
+async function signClaims(input: {
+    claims: SandboxControlAuthorityClaims
+    keyId: string
+    signingKey: CryptoKey
+}): Promise<string> {
+    return await new SignJWT(input.claims)
+        .setProtectedHeader({ alg: sandboxControlAuthorityAlgorithm, kid: input.keyId, typ: 'JWT' })
+        .sign(input.signingKey)
+}

package/src/worker-runtime-client/e2b-supervisor-connection.ts

@@ -0,0 +1,102 @@
+import type { SupervisorConnection, SupervisorConnectionOpener } from './ports'
+import { ProviderError } from '../errors'
+import { sandboxControlAuthorityHeaderName } from '../supervisor-protocol/control-authority'
+
+const defaultIngressUrl = 'wss://sandbox.e2b.app/control/ws'
+const defaultSupervisorPort = 8080
+
+export function createE2bSupervisorConnectionOpener(input: {
+    ingressUrl?: string
+    supervisorPort?: number
+    dialTimeoutMs?: number
+} = {}): SupervisorConnectionOpener {
+    const ingressUrl = input.ingressUrl ?? defaultIngressUrl
+    const supervisorPort = input.supervisorPort ?? defaultSupervisorPort
+    const dialTimeoutMs = input.dialTimeoutMs ?? 10_000
+
+    return async (request) => {
+        request.signal.throwIfAborted()
+        const socket = new WebSocket(ingressUrl, {
+            headers: {
+                'E2B-Traffic-Access-Token': request.trafficAccessToken,
+                'E2b-Sandbox-Id': request.providerSandboxId,
+                'E2b-Sandbox-Port': String(supervisorPort),
+                'X-Coding-Agent-Session-Id': request.sessionId,
+                'X-Coding-Agent-Session-Sandbox-Id': request.sessionSandboxId,
+                'X-Coding-Agent-Runtime-Connection-Attempt-Id': request.runtimeConnectionAttemptId,
+                [sandboxControlAuthorityHeaderName]: request.authorityAssertion,
+                'X-Coding-Agent-Connection-Generation': String(request.connectionGeneration),
+                'X-Coding-Agent-Event-Ack': String(request.eventAckFloor),
+            },
+        })
+        const connection: SupervisorConnection = {
+            send(frame) {
+                if (socket.readyState !== WebSocket.OPEN) {
+                    throw new E2bSupervisorConnectionError({
+                        message: 'E2B supervisor connection is not open.',
+                    })
+                }
+                socket.send(frame)
+            },
+            close: (code, reason) => socket.close(code, reason),
+        }
+        socket.binaryType = 'arraybuffer'
+        socket.onmessage = (event) => {
+            if (typeof event.data === 'string')
+                request.onMessage(event.data)
+            else if (event.data instanceof ArrayBuffer)
+                request.onMessage(event.data)
+        }
+
+        await new Promise<void>((resolve, reject) => {
+            let timeout: ReturnType<typeof setTimeout>
+            const abort = () => {
+                clearTimeout(timeout)
+                request.signal.removeEventListener('abort', abort)
+                socket.close()
+                reject(request.signal.reason)
+            }
+            const cleanup = () => {
+                clearTimeout(timeout)
+                request.signal.removeEventListener('abort', abort)
+            }
+            timeout = setTimeout(() => {
+                cleanup()
+                socket.close()
+                reject(new E2bSupervisorConnectionError({ message: 'Timed out opening the E2B supervisor connection.' }))
+            }, dialTimeoutMs)
+            request.signal.addEventListener('abort', abort, { once: true })
+            socket.onopen = () => {
+                cleanup()
+                socket.onclose = event => request.onClose({ code: event.code, reason: event.reason })
+                socket.onerror = () => {
+                    if (socket.readyState === WebSocket.OPEN)
+                        socket.close()
+                }
+                resolve()
+            }
+            socket.onerror = () => undefined
+            socket.onclose = (event) => {
+                cleanup()
+                reject(new E2bSupervisorConnectionError({
+                    message: `E2B supervisor connection closed during upgrade (${event.code}).`,
+                }))
+            }
+        })
+        return connection
+    }
+}
+
+export const openE2bSupervisorConnection = createE2bSupervisorConnectionOpener()
+
+class E2bSupervisorConnectionError extends ProviderError {
+    constructor(input: { message: string, cause?: unknown }) {
+        super({
+            code: 'agent-core.e2b-supervisor-connection-failed',
+            message: input.message,
+            publicMessage: 'The sandbox control connection is unavailable.',
+            retryable: true,
+            cause: input.cause,
+        })
+    }
+}