@meistrari/agent-core 0.0.0 → 0.1.1

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

@@ -0,0 +1,301 @@
+import type { Logger } from '../logger'
+import type {
+    BootstrapProvider,
+    CommandSource,
+    ConnectionAuthorityStore,
+    ControlAuthoritySigner,
+    DurableEventSink,
+    EphemeralEventSink,
+    RpcDispatcher,
+    RuntimeNotifications,
+    SandboxConnectionRuntimeTimings,
+    SessionFailureSink,
+    SupervisorConnectionOpener,
+} from './ports'
+import {
+    ConnectionAuthorityLostError,
+    ConnectionClosedError,
+    DeterministicEstablishmentError,
+    SessionSandboxConnectionAttempt,
+} from './connection-attempt'
+import { TokenDecryptionError } from './token-crypto'
+
+const defaultTimings: SandboxConnectionRuntimeTimings = {
+    leaseDurationMs: 30_000,
+    leaseRenewIntervalMs: 10_000,
+    establishmentBudgetMs: 30_000,
+    dialTimeoutMs: 10_000,
+    retryInitialDelayMs: 250,
+    retryMaxDelayMs: 5_000,
+    candidateLimit: 100,
+}
+
+export function createSandboxConnectionRuntime(input: {
+    replicaId: string
+    trafficTokenEncryptionKey: string
+    signControlAuthority: ControlAuthoritySigner
+    openSupervisorConnection: SupervisorConnectionOpener
+    authorityStore: ConnectionAuthorityStore
+    commands: CommandSource
+    durableEvents: DurableEventSink
+    ephemeralEvents: EphemeralEventSink
+    rpc: RpcDispatcher
+    bootstrap: BootstrapProvider
+    failures: SessionFailureSink
+    notifications: RuntimeNotifications
+    timings?: Partial<SandboxConnectionRuntimeTimings>
+    logger: Logger
+}): {
+    start: () => Promise<void>
+    stop: () => Promise<void>
+    wake: (sessionSandboxId: string) => void
+    wakeCommands: (sessionId: string) => void
+} {
+    const timings = validatedTimings({ ...defaultTimings, ...input.timings })
+    const owned = new Map<string, OwnedConnection>()
+    const sessionIndex = new Map<string, string>()
+    let started = false
+    let stopping = false
+
+    const wake = (sessionSandboxId: string) => {
+        if (!started || stopping || owned.has(sessionSandboxId))
+            return
+        const connection = new OwnedConnection({ ...input, timings, sessionSandboxId })
+        owned.set(sessionSandboxId, connection)
+        void connection.run().finally(() => {
+            if (owned.get(sessionSandboxId) === connection)
+                owned.delete(sessionSandboxId)
+            const sessionId = connection.sessionId()
+            if (sessionId && sessionIndex.get(sessionId) === sessionSandboxId)
+                sessionIndex.delete(sessionId)
+        })
+        connection.onClaimed((sessionId) => {
+            sessionIndex.set(sessionId, sessionSandboxId)
+        })
+    }
+
+    return {
+        async start() {
+            if (started)
+                return
+            started = true
+            try {
+                await input.notifications.start({
+                    onConnectionDemand: wake,
+                    onCommandsReady: sessionId => owned.get(sessionIndex.get(sessionId) ?? '')?.wakeCommands(),
+                    onOwnerChanged: change => owned.get(change.sessionSandboxId)?.ownerChanged(change.generation),
+                })
+                const candidates = await input.authorityStore.listConnectionCandidates({ limit: timings.candidateLimit })
+                for (const candidate of candidates)
+                    wake(candidate.sessionSandboxId)
+            }
+            catch (error) {
+                started = false
+                await input.notifications.stop().catch(() => undefined)
+                throw error
+            }
+        },
+        async stop() {
+            if (stopping)
+                return
+            stopping = true
+            await input.notifications.stop()
+            const connections = [...owned.values()]
+            for (const connection of connections)
+                connection.stop()
+            await Promise.allSettled(connections.map(connection => connection.finished()))
+            owned.clear()
+            sessionIndex.clear()
+        },
+        wake,
+        wakeCommands(sessionId) {
+            owned.get(sessionIndex.get(sessionId) ?? '')?.wakeCommands()
+        },
+    }
+}
+
+class OwnedConnection {
+    private readonly controller = new AbortController()
+    private readonly completion = Promise.withResolvers<void>()
+    private attempt: SessionSandboxConnectionAttempt | undefined
+    private claimSessionId: string | undefined
+    private claimGeneration: number | undefined
+    private claimed: ((sessionId: string) => void) | undefined
+
+    constructor(private readonly dependencies: Parameters<typeof createSandboxConnectionRuntime>[0] & {
+        sessionSandboxId: string
+        timings: SandboxConnectionRuntimeTimings
+    }) {}
+
+    onClaimed(callback: (sessionId: string) => void): void {
+        this.claimed = callback
+        if (this.claimSessionId)
+            callback(this.claimSessionId)
+    }
+
+    sessionId(): string | undefined {
+        return this.claimSessionId
+    }
+
+    async run(): Promise<void> {
+        try {
+            const claim = await this.dependencies.authorityStore.claimConnection({
+                sessionSandboxId: this.dependencies.sessionSandboxId,
+                replicaId: this.dependencies.replicaId,
+                leaseDurationMs: this.dependencies.timings.leaseDurationMs,
+            })
+            if (!claim)
+                return
+            this.claimSessionId = claim.sessionId
+            this.claimGeneration = claim.runtimeConnectionGeneration
+            this.claimed?.(claim.sessionId)
+            let establishmentDeadline = performance.now() + this.dependencies.timings.establishmentBudgetMs
+            let backoff = this.dependencies.timings.retryInitialDelayMs
+            while (!this.controller.signal.aborted) {
+                const attempt = new SessionSandboxConnectionAttempt({
+                    claim,
+                    establishmentDeadline,
+                    replicaId: this.dependencies.replicaId,
+                    trafficTokenEncryptionKey: this.dependencies.trafficTokenEncryptionKey,
+                    signControlAuthority: this.dependencies.signControlAuthority,
+                    openSupervisorConnection: this.dependencies.openSupervisorConnection,
+                    authorityStore: this.dependencies.authorityStore,
+                    commands: this.dependencies.commands,
+                    durableEvents: this.dependencies.durableEvents,
+                    ephemeralEvents: this.dependencies.ephemeralEvents,
+                    rpc: this.dependencies.rpc,
+                    bootstrap: this.dependencies.bootstrap,
+                    failures: this.dependencies.failures,
+                    timings: this.dependencies.timings,
+                    logger: this.dependencies.logger,
+                    parentSignal: this.controller.signal,
+                })
+                this.attempt = attempt
+                try {
+                    await attempt.run()
+                }
+                catch (error) {
+                    attempt.stop(error)
+                    await attempt.drain()
+                    if (this.controller.signal.aborted)
+                        return
+                    if (error instanceof DeterministicEstablishmentError) {
+                        if (!error.reported) {
+                            await this.dependencies.failures.failDeterministicEstablishment({
+                                sessionId: claim.sessionId,
+                                sessionSandboxId: claim.sessionSandboxId,
+                                providerSandboxId: claim.providerSandboxId,
+                                generation: claim.runtimeConnectionGeneration,
+                                code: error.code,
+                                detail: error.detail,
+                            })
+                        }
+                        return
+                    }
+                    if (error instanceof ConnectionAuthorityLostError)
+                        return
+                    if (error instanceof TokenDecryptionError) {
+                        await this.dependencies.failures.markSessionFailed({
+                            sessionId: claim.sessionId,
+                            sessionSandboxId: claim.sessionSandboxId,
+                            providerSandboxId: claim.providerSandboxId,
+                            code: 'agent-core.traffic-token-decryption-failed',
+                        })
+                        return
+                    }
+                    if (error instanceof ConnectionClosedError && error.ready) {
+                        establishmentDeadline = performance.now() + this.dependencies.timings.establishmentBudgetMs
+                        backoff = this.dependencies.timings.retryInitialDelayMs
+                    }
+                    await this.dependencies.authorityStore.markConnectionReconnecting({
+                        sessionSandboxId: claim.sessionSandboxId,
+                        generation: claim.runtimeConnectionGeneration,
+                    })
+                    if (performance.now() + backoff >= establishmentDeadline)
+                        break
+                    await abortableDelay(backoff, this.controller.signal)
+                    backoff = Math.min(backoff * 2, this.dependencies.timings.retryMaxDelayMs)
+                }
+            }
+            if (!this.controller.signal.aborted) {
+                const exhausted = await this.dependencies.authorityStore.markConnectionRetryExhausted({
+                    sessionSandboxId: claim.sessionSandboxId,
+                    generation: claim.runtimeConnectionGeneration,
+                    replicaId: this.dependencies.replicaId,
+                })
+                if (exhausted) {
+                    await this.dependencies.failures.markEstablishmentExhausted({
+                        sessionId: claim.sessionId,
+                        sessionSandboxId: claim.sessionSandboxId,
+                        providerSandboxId: claim.providerSandboxId,
+                        generation: claim.runtimeConnectionGeneration,
+                    })
+                }
+            }
+        }
+        catch (error) {
+            this.dependencies.logger.error({ err: error, sessionSandboxId: this.dependencies.sessionSandboxId }, 'sandbox connection owner failed')
+        }
+        finally {
+            this.attempt?.stop()
+            if (this.claimGeneration !== undefined) {
+                await this.dependencies.authorityStore.releaseConnection({
+                    sessionSandboxId: this.dependencies.sessionSandboxId,
+                    generation: this.claimGeneration,
+                    reason: this.controller.signal.aborted ? 'runtime_stopped' : 'connection_owner_finished',
+                }).catch(error => this.dependencies.logger.warn({ err: error }, 'sandbox connection release failed'))
+            }
+            this.completion.resolve()
+        }
+    }
+
+    wakeCommands(): void {
+        this.attempt?.wakeCommands()
+    }
+
+    ownerChanged(generation: number): void {
+        if (!this.attempt || generation !== this.claimGeneration)
+            return
+        this.stop(new ConnectionAuthorityLostError())
+    }
+
+    stop(reason?: unknown): void {
+        if (!this.controller.signal.aborted)
+            this.controller.abort(reason ?? new DOMException('Sandbox connection runtime stopped.', 'AbortError'))
+        this.attempt?.stop(reason)
+    }
+
+    async finished(): Promise<void> {
+        await this.completion.promise
+    }
+}
+
+async function abortableDelay(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 validatedTimings(timings: SandboxConnectionRuntimeTimings): SandboxConnectionRuntimeTimings {
+    for (const [name, value] of Object.entries(timings)) {
+        if (!Number.isSafeInteger(value) || value <= 0)
+            throw new RangeError(`${name} must be a positive safe integer.`)
+    }
+    if (timings.leaseRenewIntervalMs >= timings.leaseDurationMs)
+        throw new RangeError('leaseRenewIntervalMs must be lower than leaseDurationMs.')
+    if (timings.retryInitialDelayMs > timings.retryMaxDelayMs)
+        throw new RangeError('retryInitialDelayMs must not exceed retryMaxDelayMs.')
+    return timings
+}

package/src/worker-runtime-client/token-crypto.ts

@@ -0,0 +1,46 @@
+import { Buffer } from 'node:buffer'
+import { createCipheriv, createDecipheriv, createHash, randomBytes } from 'node:crypto'
+
+const algorithm = 'aes-256-gcm'
+const ivBytes = 12
+
+function encryptionKey(secret: string): Buffer {
+    if (!secret)
+        throw new TypeError('Token encryption key must not be empty.')
+    return createHash('sha256').update(secret, 'utf8').digest()
+}
+
+export function encryptToken(input: { token: string, secret: string }): string {
+    const iv = randomBytes(ivBytes)
+    const cipher = createCipheriv(algorithm, encryptionKey(input.secret), iv)
+    const ciphertext = Buffer.concat([cipher.update(input.token, 'utf8'), cipher.final()])
+    return [iv, cipher.getAuthTag(), ciphertext].map(value => value.toString('base64url')).join('.')
+}
+
+export function decryptToken(input: { encryptedToken: string, secret: string }): string {
+    try {
+        const parts = input.encryptedToken.split('.')
+        if (parts.length !== 3)
+            throw new Error('Encrypted token has an invalid shape.')
+        const [encodedIv, encodedTag, encodedCiphertext] = parts
+        if (!encodedIv || !encodedTag || !encodedCiphertext)
+            throw new Error('Encrypted token contains an empty segment.')
+        const decipher = createDecipheriv(algorithm, encryptionKey(input.secret), Buffer.from(encodedIv, 'base64url'))
+        decipher.setAuthTag(Buffer.from(encodedTag, 'base64url'))
+        return Buffer.concat([
+            decipher.update(Buffer.from(encodedCiphertext, 'base64url')),
+            decipher.final(),
+        ]).toString('utf8')
+    }
+    catch (cause) {
+        throw new TokenDecryptionError({ cause })
+    }
+}
+
+export class TokenDecryptionError extends Error {
+    override readonly name = 'TokenDecryptionError'
+
+    constructor(options: { cause: unknown }) {
+        super('Stored sandbox traffic credential could not be decrypted.', options)
+    }
+}