@meistrari/agent-core 0.0.0 → 0.1.0

package/src/testing/in-memory-runtime-control-plane.ts

@@ -0,0 +1,205 @@
+import type { SessionBootstrapBody } from '../supervisor-protocol/bootstrap'
+import type { WireCommandBody } from '../supervisor-protocol/command-body'
+import type { CommandEnvelope } from '../supervisor-protocol/envelopes/control-plane-to-supervisor'
+import type { DurableEventEnvelope } from '../supervisor-protocol/envelopes/supervisor-to-control-plane'
+import type { SupervisorRpcRequestEnvelope } from '../supervisor-protocol/rpc'
+import type { SupervisorAgentRunSnapshot } from '../supervisor-protocol/supervisor-agent-run-snapshot'
+import type {
+    BootstrapProvider,
+    CommandSource,
+    ConnectionAuthorityStore,
+    ConnectionClaim,
+    DurableEventSink,
+    EphemeralEventSink,
+    RpcDispatcher,
+    RuntimeNotifications,
+    SessionFailureSink,
+} from '../worker-runtime-client/ports'
+import { ulid } from 'ulid'
+
+type NotificationCallbacks = Parameters<RuntimeNotifications['start']>[0]
+
+export class InMemoryRuntimeControlPlane implements ConnectionAuthorityStore, CommandSource, DurableEventSink,
+EphemeralEventSink, RpcDispatcher, BootstrapProvider, SessionFailureSink, RuntimeNotifications {
+    readonly durableEventLog: DurableEventEnvelope[] = []
+    readonly durableRecordAttempts: number[] = []
+    readonly ephemeralEventLog: Array<Parameters<EphemeralEventSink['publish']>[0]> = []
+    readonly readySnapshots: SupervisorAgentRunSnapshot[] = []
+    readonly projectedCommandSequences: number[] = []
+    readonly failureLog: Array<{ kind: string, code?: string, detail?: string }> = []
+    readonly rpcCalls: SupervisorRpcRequestEnvelope[] = []
+    onDurableRecord: ((event: DurableEventEnvelope) => void | Promise<void>) | undefined
+    private readonly commandLog: CommandEnvelope[]
+    private notifications: NotificationCallbacks | undefined
+    private claimed = false
+    private generation = 0
+    private connectionId = ''
+    private ackFloor = 0
+
+    constructor(private readonly input: {
+        ref: { sessionId: string, sessionSandboxId: string, providerSandboxId: string }
+        trafficAccessTokenEncrypted: string
+        bootstrap: Omit<SessionBootstrapBody, 'sessionId' | 'initialCommand'>
+        initialCommand: { commandId: string, body: WireCommandBody }
+        secrets?: Awaited<ReturnType<BootstrapProvider['issueEphemeralSecrets']>>
+        rpcHandlers?: Readonly<Record<string, (body: unknown) => unknown | Promise<unknown>>>
+    }) {
+        this.commandLog = [{
+            kind: 'command',
+            commandId: input.initialCommand.commandId,
+            commandSeq: 1,
+            body: input.initialCommand.body,
+        }]
+    }
+
+    enqueueCommand(input: { commandId?: string, body: WireCommandBody }): CommandEnvelope {
+        const command: CommandEnvelope = {
+            kind: 'command',
+            commandId: input.commandId ?? ulid(),
+            commandSeq: this.commandLog.length + 1,
+            body: input.body,
+        }
+        this.commandLog.push(command)
+        this.notifications?.onCommandsReady(this.input.ref.sessionId)
+        return command
+    }
+
+    demandConnection(): void {
+        this.notifications?.onConnectionDemand(this.input.ref.sessionSandboxId)
+    }
+
+    notifyOwnerChanged(generation = this.generation): void {
+        this.notifications?.onOwnerChanged({ sessionSandboxId: this.input.ref.sessionSandboxId, generation })
+    }
+
+    async listConnectionCandidates(): Promise<{ sessionSandboxId: string }[]> {
+        return this.claimed ? [] : [{ sessionSandboxId: this.input.ref.sessionSandboxId }]
+    }
+
+    async notifyConnectionDemand(input: { sessionSandboxIds: string[] }): Promise<void> {
+        for (const sessionSandboxId of input.sessionSandboxIds)
+            this.notifications?.onConnectionDemand(sessionSandboxId)
+    }
+
+    async claimConnection(input: { sessionSandboxId: string }): Promise<ConnectionClaim | undefined> {
+        if (this.claimed || input.sessionSandboxId !== this.input.ref.sessionSandboxId)
+            return undefined
+        this.claimed = true
+        this.generation += 1
+        this.connectionId = ulid()
+        return {
+            ...this.input.ref,
+            runtimeConnectionGeneration: this.generation,
+            runtimeConnectionId: this.connectionId,
+            trafficAccessTokenEncrypted: this.input.trafficAccessTokenEncrypted,
+            leaseExpiresAt: new Date(Date.now() + 30_000),
+        }
+    }
+
+    async renewConnectionLease(input: { generation: number }): Promise<{
+        renewed: boolean
+        leaseExpiresAt?: Date
+        hasOutstandingCommand: boolean
+    }> {
+        const renewed = this.claimed && input.generation === this.generation
+        return {
+            renewed,
+            ...(renewed ? { leaseExpiresAt: new Date(Date.now() + 30_000) } : {}),
+            hasOutstandingCommand: this.projectedCommandSequences.length < this.commandLog.length - 1,
+        }
+    }
+
+    async confirmConnectionAuthority(input: { generation: number }): Promise<boolean> {
+        return this.claimed && input.generation === this.generation
+    }
+
+    async markConnectionReady(input: { generation: number, agentRunSnapshot: SupervisorAgentRunSnapshot }): Promise<boolean> {
+        if (!this.claimed || input.generation !== this.generation)
+            return false
+        this.readySnapshots.push(input.agentRunSnapshot)
+        return true
+    }
+
+    async markConnectionReconnecting(): Promise<void> {}
+
+    async releaseConnection(input: { generation: number }): Promise<void> {
+        if (input.generation === this.generation)
+            this.claimed = false
+    }
+
+    async markConnectionRetryExhausted(input: { generation: number }): Promise<boolean> {
+        return this.claimed && input.generation === this.generation
+    }
+
+    async projectAppliedCommand(input: { commandSeq: number }): Promise<void> {
+        this.projectedCommandSequences.push(input.commandSeq)
+    }
+
+    async initial(): Promise<CommandEnvelope | undefined> {
+        return this.commandLog[0]
+    }
+
+    async next(input: { afterSequence: number }): Promise<CommandEnvelope | undefined> {
+        return this.commandLog.find(command => command.commandSeq > input.afterSequence)
+    }
+
+    async getAckFloor(): Promise<number> {
+        return this.ackFloor
+    }
+
+    async record(input: { event: DurableEventEnvelope }): Promise<{
+        status: 'recorded' | 'duplicate'
+        highWaterMark: number
+    } | { status: 'gap' }> {
+        const { event } = input
+        this.durableRecordAttempts.push(event.seq)
+        if (event.seq <= this.ackFloor)
+            return { status: 'duplicate', highWaterMark: this.ackFloor }
+        if (event.seq !== this.ackFloor + 1)
+            return { status: 'gap' }
+        this.durableEventLog.push(event)
+        this.ackFloor = event.seq
+        await this.onDurableRecord?.(event)
+        return { status: 'recorded', highWaterMark: this.ackFloor }
+    }
+
+    publish(input: Parameters<EphemeralEventSink['publish']>[0]): void {
+        this.ephemeralEventLog.push(input)
+    }
+
+    async handle(input: { request: SupervisorRpcRequestEnvelope }): Promise<unknown> {
+        this.rpcCalls.push(input.request)
+        const handler = this.input.rpcHandlers?.[input.request.method]
+        if (!handler)
+            throw new Error(`No in-memory RPC handler is registered for ${input.request.method}.`)
+        return await handler(input.request.body)
+    }
+
+    async loadBootstrapBody(): Promise<Omit<SessionBootstrapBody, 'sessionId' | 'initialCommand'>> {
+        return this.input.bootstrap
+    }
+
+    async issueEphemeralSecrets(): Promise<Awaited<ReturnType<BootstrapProvider['issueEphemeralSecrets']>>> {
+        return this.input.secrets ?? {}
+    }
+
+    async markSessionFailed(input: { code: string, detail?: string }): Promise<void> {
+        this.failureLog.push({ kind: 'session', code: input.code, detail: input.detail })
+    }
+
+    async failDeterministicEstablishment(input: { code: string, detail?: string }): Promise<void> {
+        this.failureLog.push({ kind: 'deterministic', code: input.code, detail: input.detail })
+    }
+
+    async markEstablishmentExhausted(): Promise<void> {
+        this.failureLog.push({ kind: 'exhausted' })
+    }
+
+    async start(callbacks: NotificationCallbacks): Promise<void> {
+        this.notifications = callbacks
+    }
+
+    async stop(): Promise<void> {
+        this.notifications = undefined
+    }
+}

package/src/testing/index.ts

@@ -0,0 +1,6 @@
+export { createTestControlAuthority } from './es256-test-keys'
+export type { TestControlAuthority } from './es256-test-keys'
+export { InMemoryRuntimeControlPlane } from './in-memory-runtime-control-plane'
+export { createLoopbackSupervisorConnectionOpener } from './loopback-supervisor-connection'
+export { createScriptedProvider } from './scripted-provider'
+export type { ScriptedCommandContext } from './scripted-provider'

package/src/testing/loopback-supervisor-connection.ts

@@ -0,0 +1,71 @@
+import type { SupervisorConnection, SupervisorConnectionOpener } from '../worker-runtime-client/ports'
+import { sandboxControlAuthorityHeaderName } from '../supervisor-protocol/control-authority'
+
+export function createLoopbackSupervisorConnectionOpener(input: {
+    url: string
+    supervisorPort?: number
+    dialTimeoutMs?: number
+}): SupervisorConnectionOpener {
+    const dialTimeoutMs = input.dialTimeoutMs ?? 2_000
+    return async (request) => {
+        request.signal.throwIfAborted()
+        const socket = new WebSocket(input.url, {
+            headers: {
+                'E2B-Traffic-Access-Token': request.trafficAccessToken,
+                'E2b-Sandbox-Id': request.providerSandboxId,
+                'E2b-Sandbox-Port': String(input.supervisorPort ?? 8080),
+                '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),
+            },
+        })
+        socket.binaryType = 'arraybuffer'
+        const connection: SupervisorConnection = {
+            send(data) {
+                if (socket.readyState !== WebSocket.OPEN)
+                    throw new Error('Loopback supervisor connection is not open.')
+                socket.send(data)
+            },
+            close: (code, reason) => socket.close(code, reason),
+        }
+        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)
+            }
+            timeout = setTimeout(() => {
+                request.signal.removeEventListener('abort', abort)
+                socket.close()
+                reject(new Error('Timed out opening the loopback supervisor connection.'))
+            }, dialTimeoutMs)
+            request.signal.addEventListener('abort', abort, { once: true })
+            socket.onopen = () => {
+                clearTimeout(timeout)
+                request.signal.removeEventListener('abort', abort)
+                socket.onclose = event => request.onClose({ code: event.code, reason: event.reason })
+                socket.onerror = () => socket.close()
+                resolve()
+            }
+            socket.onerror = () => undefined
+            socket.onclose = (event) => {
+                clearTimeout(timeout)
+                request.signal.removeEventListener('abort', abort)
+                reject(new Error(`Loopback supervisor connection closed during upgrade (${event.code}).`))
+            }
+        })
+        return connection
+    }
+}

package/src/testing/scripted-provider.ts

@@ -0,0 +1,64 @@
+import type { EphemeralCredentials, SessionBootstrapGitToken } from '../supervisor-protocol/bootstrap'
+import type { WireCommandBody } from '../supervisor-protocol/command-body'
+import type { WireEventBody } from '../supervisor-protocol/event-body'
+import type { SupervisorAgentRunSnapshot } from '../supervisor-protocol/supervisor-agent-run-snapshot'
+import type { SupervisorAgentRuntime, SupervisorProviderFactory } from '../supervisor/provider-factory'
+import type { SupervisorRpcClient } from '../supervisor/rpc-client'
+import { createSupervisorProviderFactory } from '../supervisor/provider-factory'
+
+export interface ScriptedCommandContext {
+    commandId: string
+    commandSeq: number
+    body: WireCommandBody
+    signal: AbortSignal
+    rpc: SupervisorRpcClient
+    emit: (body: WireEventBody) => void
+}
+
+export function createScriptedProvider(input: {
+    snapshot?: SupervisorAgentRunSnapshot
+    onCommand?: (context: ScriptedCommandContext) => Promise<void>
+    onRefreshSecrets?: (input: { credentials?: EphemeralCredentials, gitToken?: SessionBootstrapGitToken }) => void | Promise<void>
+    onClose?: () => Promise<void>
+} = {}): {
+    factory: SupervisorProviderFactory
+    handledCommands: Array<{ commandId: string, commandSeq: number, body: WireCommandBody }>
+    refreshedSecrets: Array<{ credentials?: EphemeralCredentials, gitToken?: SessionBootstrapGitToken }>
+    createdCount: () => number
+    closedCount: () => number
+} {
+    const handledCommands: Array<{ commandId: string, commandSeq: number, body: WireCommandBody }> = []
+    const refreshedSecrets: Array<{ credentials?: EphemeralCredentials, gitToken?: SessionBootstrapGitToken }> = []
+    let created = 0
+    let closed = 0
+    const factory = createSupervisorProviderFactory(async ({ rpc }) => {
+        created += 1
+        const runtime: SupervisorAgentRuntime = {
+            snapshot: () => input.snapshot ?? { status: 'not_attached' },
+            async handle(command) {
+                handledCommands.push({
+                    commandId: command.commandId,
+                    commandSeq: command.commandSeq,
+                    body: command.body,
+                })
+                await input.onCommand?.({ ...command, rpc })
+            },
+            async refreshSecrets(secrets) {
+                refreshedSecrets.push(secrets)
+                await input.onRefreshSecrets?.(secrets)
+            },
+            async close() {
+                closed += 1
+                await input.onClose?.()
+            },
+        }
+        return runtime
+    })
+    return {
+        factory,
+        handledCommands,
+        refreshedSecrets,
+        createdCount: () => created,
+        closedCount: () => closed,
+    }
+}

package/src/worker-runtime-client/command-pump.ts

@@ -0,0 +1,132 @@
+import type { Logger } from '../logger'
+import type { CommandAckEnvelope } from '../supervisor-protocol/envelopes/supervisor-to-control-plane'
+import type { CommandSource, ConnectionAuthorityStore, SupervisorConnection } from './ports'
+import { encodeWireFrame } from '../supervisor-protocol/wire-codec'
+
+export class SessionCommandPump {
+    private afterSequence: number
+    private running = false
+    private stopped = false
+    private wakePending = false
+    private pending: {
+        commandId: string
+        commandSeq: number
+        settle: (ack: CommandAckEnvelope) => void
+    } | undefined
+
+    constructor(private readonly dependencies: {
+        sessionId: string
+        sessionSandboxId: string
+        runtimeConnectionId: string
+        afterSequence: number
+        commands: CommandSource
+        authorityStore: ConnectionAuthorityStore
+        connection: SupervisorConnection
+        logger: Logger
+        onRejected: (ack: Extract<CommandAckEnvelope, { status: 'rejected' }>) => Promise<void>
+    }) {
+        this.afterSequence = dependencies.afterSequence
+    }
+
+    start(): void {
+        this.wake()
+    }
+
+    wake(): void {
+        if (this.stopped)
+            return
+        this.wakePending = true
+        if (!this.running)
+            void this.run()
+    }
+
+    acknowledge(ack: CommandAckEnvelope): void {
+        const pending = this.pending
+        if (!pending) {
+            if (ack.status === 'received' && ack.commandSeq <= this.afterSequence)
+                return
+            this.dependencies.logger.warn({ commandSeq: ack.commandSeq }, 'closing after an unsolicited command acknowledgement')
+            this.dependencies.connection.close(1008, 'Command ACK has no pending command.')
+            this.stop()
+            return
+        }
+        if (ack.commandId !== pending.commandId || ack.commandSeq !== pending.commandSeq) {
+            this.dependencies.logger.warn({ commandSeq: ack.commandSeq }, 'closing after a mismatched command acknowledgement')
+            this.dependencies.connection.close(1008, 'Command ACK did not match the pending command.')
+            this.stop()
+            return
+        }
+        pending.settle(ack)
+    }
+
+    stop(): void {
+        this.stopped = true
+        this.pending?.settle({
+            kind: 'command.ack',
+            commandId: this.pending.commandId,
+            commandSeq: this.pending.commandSeq,
+            status: 'rejected',
+            errorCode: 'agent-core.connection-stopped',
+        })
+        this.pending = undefined
+    }
+
+    cursor(): number {
+        return this.afterSequence
+    }
+
+    private async run(): Promise<void> {
+        if (this.running)
+            return
+        this.running = true
+        try {
+            while (!this.stopped && this.wakePending) {
+                this.wakePending = false
+                while (!this.stopped) {
+                    const command = await this.dependencies.commands.next({
+                        sessionId: this.dependencies.sessionId,
+                        runtimeConnectionId: this.dependencies.runtimeConnectionId,
+                        afterSequence: this.afterSequence,
+                    })
+                    if (!command)
+                        break
+                    const ack = await this.sendAndWait(command)
+                    if (this.stopped)
+                        return
+                    if (ack.status === 'rejected') {
+                        await this.dependencies.onRejected(ack)
+                        return
+                    }
+                    this.afterSequence = ack.commandSeq
+                    await this.dependencies.authorityStore.projectAppliedCommand({
+                        sessionSandboxId: this.dependencies.sessionSandboxId,
+                        runtimeConnectionId: this.dependencies.runtimeConnectionId,
+                        commandSeq: ack.commandSeq,
+                    })
+                }
+            }
+        }
+        finally {
+            this.running = false
+            if (!this.stopped && this.wakePending)
+                void this.run()
+        }
+    }
+
+    private async sendAndWait(command: NonNullable<Awaited<ReturnType<CommandSource['next']>>>): Promise<CommandAckEnvelope> {
+        const waiter = Promise.withResolvers<CommandAckEnvelope>()
+        this.pending = {
+            commandId: command.commandId,
+            commandSeq: command.commandSeq,
+            settle: waiter.resolve,
+        }
+        try {
+            await this.dependencies.connection.send(encodeWireFrame(command))
+            return await waiter.promise
+        }
+        finally {
+            if (this.pending?.commandId === command.commandId)
+                this.pending = undefined
+        }
+    }
+}