@meistrari/agent-core 0.2.0-experimental.0 → 0.2.0-experimental.2

package/src/testing/index.ts

@@ -1,6 +1,6 @@
 export { createTestControlAuthority } from './es256-test-keys'
 export type { TestControlAuthority } from './es256-test-keys'
-export { InMemoryRuntimeControlPlane } from './in-memory-runtime-control-plane'
+export { inMemoryJournalIdentityChangedCode, 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/worker-runtime-client/command-pump.ts

@@ -3,11 +3,19 @@ import type { CommandAckEnvelope } from '../supervisor-protocol/envelopes/superv
 import type { CommandSource, ConnectionAuthorityStore, SupervisorConnection } from './ports'
 import { encodeWireFrame } from '../supervisor-protocol/wire-codec'
 
+export class CommandAckTimeoutError extends Error {
+    constructor() {
+        super('agent-core.command-ack-timeout')
+        this.name = 'CommandAckTimeoutError'
+    }
+}
+
 export class SessionCommandPump {
     private afterSequence: number
     private running = false
     private stopped = false
     private wakePending = false
+    private readonly idleWaiters = new Set<() => void>()
     private pending: {
         commandId: string
         commandSeq: number
@@ -19,13 +27,20 @@ export class SessionCommandPump {
         sessionSandboxId: string
         runtimeConnectionId: string
         afterSequence: number
+        individualReceipts?: boolean
+        ackTimeoutMs?: number
+        onFailure?: (error: unknown) => void
         commands: CommandSource
-        authorityStore: ConnectionAuthorityStore
+        authorityStore?: ConnectionAuthorityStore
         connection: SupervisorConnection
         logger: Logger
         onRejected: (ack: Extract<CommandAckEnvelope, { status: 'rejected' }>) => Promise<void>
     }) {
         this.afterSequence = dependencies.afterSequence
+        if (dependencies.ackTimeoutMs !== undefined && (!Number.isSafeInteger(dependencies.ackTimeoutMs) || dependencies.ackTimeoutMs < 1))
+            throw new RangeError('ackTimeoutMs must be a positive safe integer')
+        if (dependencies.individualReceipts && (!dependencies.commands.pending || !dependencies.commands.received || !dependencies.commands.rejected))
+            throw new Error('wss-v1 requires individual command selection, receipt persistence and rejection persistence')
     }
 
     start(): void {
@@ -43,7 +58,7 @@ export class SessionCommandPump {
     acknowledge(ack: CommandAckEnvelope): void {
         const pending = this.pending
         if (!pending) {
-            if (ack.status === 'received' && ack.commandSeq <= this.afterSequence)
+            if (ack.status === 'received' && (this.dependencies.individualReceipts || 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.')
@@ -75,6 +90,13 @@ export class SessionCommandPump {
         return this.afterSequence
     }
 
+    /** Stop dispatch and wait for any in-flight selection/send/receipt to settle. */
+    async close(): Promise<void> {
+        this.stop()
+        if (this.running)
+            await new Promise<void>(resolve => this.idleWaiters.add(resolve))
+    }
+
     private async run(): Promise<void> {
         if (this.running)
             return
@@ -83,26 +105,58 @@ export class SessionCommandPump {
             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,
-                    })
+                    const command = this.dependencies.individualReceipts
+                        ? await this.dependencies.commands.pending!({ sessionId: this.dependencies.sessionId, runtimeConnectionId: this.dependencies.runtimeConnectionId })
+                        : await this.dependencies.commands.next({
+                                sessionId: this.dependencies.sessionId,
+                                runtimeConnectionId: this.dependencies.runtimeConnectionId,
+                                afterSequence: this.afterSequence,
+                            })
+                    if (this.stopped)
+                        return
                     if (!command)
                         break
                     const ack = await this.sendAndWait(command)
                     if (this.stopped)
                         return
                     if (ack.status === 'rejected') {
+                        if (!this.dependencies.individualReceipts) {
+                            await this.dependencies.onRejected(ack)
+                            return
+                        }
+                        // Terminal for this command only: persist it so `pending` never offers it again,
+                        // then keep delivering the rest of the session's work.
+                        await this.dependencies.commands.rejected!({
+                            sessionId: this.dependencies.sessionId,
+                            runtimeConnectionId: this.dependencies.runtimeConnectionId,
+                            commandId: ack.commandId,
+                            commandSeq: ack.commandSeq,
+                            errorCode: ack.errorCode,
+                            detail: ack.detail,
+                        })
                         await this.dependencies.onRejected(ack)
-                        return
+                        continue
                     }
-                    this.afterSequence = ack.commandSeq
+                    if (this.dependencies.individualReceipts) {
+                        await this.dependencies.commands.received!({
+                            sessionId: this.dependencies.sessionId,
+                            runtimeConnectionId: this.dependencies.runtimeConnectionId,
+                            commandId: ack.commandId,
+                            commandSeq: ack.commandSeq,
+                        })
+                    }
+                    this.afterSequence = Math.max(this.afterSequence, ack.commandSeq)
                 }
             }
         }
+        catch (error) {
+            this.fail(error)
+        }
         finally {
             this.running = false
+            for (const resolve of this.idleWaiters)
+                resolve()
+            this.idleWaiters.clear()
             if (!this.stopped && this.wakePending)
                 void this.run()
         }
@@ -110,18 +164,39 @@ export class SessionCommandPump {
 
     private async sendAndWait(command: NonNullable<Awaited<ReturnType<CommandSource['next']>>>): Promise<CommandAckEnvelope> {
         const waiter = Promise.withResolvers<CommandAckEnvelope>()
+        const timeout = this.dependencies.ackTimeoutMs === undefined ? undefined : setTimeout(() => this.fail(new CommandAckTimeoutError()), this.dependencies.ackTimeoutMs)
         this.pending = {
             commandId: command.commandId,
             commandSeq: command.commandSeq,
-            settle: waiter.resolve,
+            settle: (ack) => {
+                clearTimeout(timeout)
+                waiter.resolve(ack)
+            },
         }
         try {
             await this.dependencies.connection.send(encodeWireFrame(command))
             return await waiter.promise
         }
         finally {
+            clearTimeout(timeout)
             if (this.pending?.commandId === command.commandId)
                 this.pending = undefined
         }
     }
+
+    private fail(error: unknown): void {
+        if (this.stopped)
+            return
+        this.dependencies.logger.warn({ err: error }, 'command delivery failed; reconnect and reconcile durable attempts')
+        try {
+            this.dependencies.onFailure?.(error)
+        }
+        catch (callbackError) {
+            this.dependencies.logger.warn({ err: callbackError }, 'command failure callback failed')
+        }
+        finally {
+            this.stop()
+            this.dependencies.connection.close(1011, 'Command delivery failed.')
+        }
+    }
 }

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

@@ -1,4 +1,6 @@
 import type { Logger } from '../logger'
+import type { SessionBootstrapBody } from '../supervisor-protocol/bootstrap'
+import type { SessionBootstrapEnvelope } from '../supervisor-protocol/envelopes/control-plane-to-supervisor'
 import type { CommandAckEnvelope } from '../supervisor-protocol/envelopes/supervisor-to-control-plane'
 import type {
     BootstrapProvider,
@@ -16,6 +18,7 @@ import type {
 } from './ports'
 import { randomUUID } from 'node:crypto'
 import { wireCommandBodySchema } from '../supervisor-protocol/command-body'
+import { sessionBootstrapEnvelopeSchema } from '../supervisor-protocol/envelopes/control-plane-to-supervisor'
 import { encodeWireFrame } from '../supervisor-protocol/wire-codec'
 import { SessionCommandPump } from './command-pump'
 import { BootstrapRejectedAckError, SupervisorFrameProcessor } from './frame-processor'
@@ -85,7 +88,7 @@ export class SessionSandboxConnectionAttempt {
         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({
+        const connection = await openSupervisorConnectionBeforeDeadline(this.dependencies.openSupervisorConnection, {
             providerSandboxId: claim.providerSandboxId,
             trafficAccessToken,
             sessionId: claim.sessionId,
@@ -97,7 +100,7 @@ export class SessionSandboxConnectionAttempt {
             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({
+        }, Math.min(this.dependencies.establishmentDeadline, performance.now() + this.dependencies.timings.dialTimeoutMs), () => new ConnectionClosedError({
             ready: false,
             detail: 'Supervisor connection dial timed out.',
         }))
@@ -136,30 +139,25 @@ export class SessionSandboxConnectionAttempt {
             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.',
-            )
+        let initialCommandBody: SessionBootstrapBody['initialCommand']
+        if (body.commandTransport !== 'wss-v1') {
+            if (!initialCommand || initialCommand.commandSeq !== 1)
+                throw new DeterministicEstablishmentError('agent-core.initial-command-invalid', 'Session command 1 is missing or invalid.')
+            const parsed = wireCommandBodySchema.safeParse(initialCommand.body)
+            if (!parsed.success)
+                throw new DeterministicEstablishmentError('agent-core.initial-command-invalid', 'Session command 1 has an invalid body.')
+            initialCommandBody = { commandId: initialCommand.commandId, commandSeq: 1, body: parsed.data }
         }
-        const initialBody = parsedInitialBody.data
-        await connection.send(encodeWireFrame({
+        await sendSupervisorBootstrap(connection, {
             kind: 'session.bootstrap',
             connectionGeneration: claim.runtimeConnectionGeneration,
             body: {
                 ...body,
                 sessionId: claim.sessionId,
-                initialCommand: {
-                    commandId: initialCommand.commandId,
-                    commandSeq: initialCommand.commandSeq,
-                    body: initialBody,
-                },
+                ...(initialCommandBody ? { initialCommand: initialCommandBody } : {}),
             },
             ...secrets,
-        }))
+        }, this.controller.signal)
 
         const readiness = await beforeDeadline(Promise.race([
             processor.waitForReadiness({ signal: this.controller.signal }),
@@ -176,15 +174,24 @@ export class SessionSandboxConnectionAttempt {
             }
             throw error
         })
+        const wssV1 = body.commandTransport === 'wss-v1'
+        if (wssV1 && readiness.sourceId === undefined)
+            throw new DeterministicEstablishmentError('agent-core.journal-identity-missing', 'Supervisor did not announce its journal identity.')
+        // The store pins the journal identity before any command is dispatched or event persisted;
+        // a replaced journal is refused here instead of being resumed from a stale cursor.
         const markedReady = await this.dependencies.authorityStore.markConnectionReady({
             sessionSandboxId: claim.sessionSandboxId,
             generation: claim.runtimeConnectionGeneration,
             runtimeConnectionId: claim.runtimeConnectionId,
             lastReceivedCommandSeq: readiness.lastReceivedCommandSeq,
             agentRunSnapshot: readiness.agentRun,
+            ...(readiness.sourceId === undefined ? {} : { journalSourceId: readiness.sourceId }),
         })
-        if (!markedReady)
+        if (markedReady === false)
             throw new ConnectionAuthorityLostError()
+        if (markedReady !== true)
+            throw new DeterministicEstablishmentError(markedReady.code, markedReady.detail)
+        processor.confirmReadiness()
         this.ready = true
 
         const pump = new SessionCommandPump({
@@ -192,11 +199,17 @@ export class SessionSandboxConnectionAttempt {
             sessionSandboxId: claim.sessionSandboxId,
             runtimeConnectionId: claim.runtimeConnectionId,
             afterSequence: readiness.lastReceivedCommandSeq,
+            individualReceipts: wssV1,
             commands: this.dependencies.commands,
             authorityStore: this.dependencies.authorityStore,
             connection,
             logger: this.dependencies.logger,
             onRejected: async (ack) => {
+                if (wssV1) {
+                    // The pump already persisted the terminal rejection; the session keeps running.
+                    this.dependencies.logger.warn({ commandSeq: ack.commandSeq, errorCode: ack.errorCode }, 'supervisor rejected a command; it will not be redispatched')
+                    return
+                }
                 await this.dependencies.failures.markSessionFailed({
                     sessionId: claim.sessionId,
                     sessionSandboxId: claim.sessionSandboxId,
@@ -233,7 +246,10 @@ export class SessionSandboxConnectionAttempt {
     }
 
     async drain(): Promise<void> {
-        await this.processor?.drain()
+        await drainSupervisorConnection([
+            () => this.processor?.drain(),
+            () => this.commandPump?.close(),
+        ])
     }
 
     private async renewLease(): Promise<never> {
@@ -347,3 +363,62 @@ async function beforeDeadline<T>(
         signal.removeEventListener('abort', abort!)
     }
 }
+
+/** A cancelled/timed-out dial must never leave a late socket alive. */
+export async function openSupervisorConnectionBeforeDeadline(
+    open: SupervisorConnectionOpener,
+    input: Parameters<SupervisorConnectionOpener>[0],
+    deadline: number,
+    timeoutError: () => Error,
+): Promise<SupervisorConnection> {
+    input.signal.throwIfAborted()
+    const stop = new AbortController()
+    const signal = AbortSignal.any([input.signal, stop.signal])
+    let opened: SupervisorConnection | undefined
+    let retired = false
+    const retire = () => {
+        if (opened && !retired) {
+            retired = true
+            opened.close(1001, 'Connection dial cancelled.')
+        }
+    }
+    const task = Promise.resolve().then(() => {
+        signal.throwIfAborted()
+        return open({ ...input, signal })
+    }).then((connection) => {
+        opened = connection
+        if (signal.aborted) {
+            retire()
+            throw signal.reason
+        }
+        return connection
+    })
+    try {
+        return await beforeDeadline(task, deadline, signal, timeoutError)
+    }
+    catch (error) {
+        stop.abort(error)
+        retire()
+        // Observe late rejection even if the deadline elapsed before beforeDeadline raced it.
+        void task.catch(() => undefined)
+        throw error
+    }
+}
+
+/** Release ownership only after every pending connection task has settled. */
+export async function drainSupervisorConnection(tasks: readonly (() => void | Promise<void>)[]): Promise<void> {
+    const results = await Promise.allSettled(tasks.map(task => Promise.resolve().then(task)))
+    const errors = results.flatMap(result => result.status === 'rejected' ? [result.reason] : [])
+    if (errors.length === 1)
+        throw errors[0]
+    if (errors.length > 1)
+        throw new AggregateError(errors, 'Supervisor connection drain failed.')
+}
+
+/** Validate the wire envelope and reject a bootstrap whose owner was cancelled while loading it. */
+export async function sendSupervisorBootstrap(connection: SupervisorConnection, envelope: SessionBootstrapEnvelope, signal: AbortSignal): Promise<void> {
+    signal.throwIfAborted()
+    const encoded = encodeWireFrame(sessionBootstrapEnvelopeSchema.parse(envelope))
+    signal.throwIfAborted()
+    await connection.send(encoded)
+}