agent-relay 1.6.0 → 2.0.0

package/packages/daemon/dist/connection.d.ts

@@ -0,0 +1,130 @@
+/**
+ * Connection state machine for agent connections.
+ *
+ * States:
+ *   CONNECTING -> HANDSHAKING -> ACTIVE -> CLOSING -> CLOSED
+ *                     |            |
+ *                     v            v
+ *                   ERROR -------> CLOSED
+ */
+import net from 'node:net';
+import { type Envelope, type AckPayload, type EntityType } from '@agent-relay/protocol/types';
+export type ConnectionState = 'CONNECTING' | 'HANDSHAKING' | 'ACTIVE' | 'CLOSING' | 'CLOSED' | 'ERROR';
+export interface ConnectionConfig {
+    maxFrameBytes: number;
+    heartbeatMs: number;
+    /** Multiplier for heartbeat timeout (timeout = heartbeatMs * multiplier). Default: 6 (30s with 5s heartbeat) */
+    heartbeatTimeoutMultiplier: number;
+    /** Optional handler to validate resume tokens and provide session state */
+    resumeHandler?: (params: {
+        agent: string;
+        resumeToken: string;
+    }) => Promise<{
+        sessionId: string;
+        resumeToken?: string;
+        seedSequences?: Array<{
+            topic?: string;
+            peer: string;
+            seq: number;
+        }>;
+    } | null>;
+    /** Optional callback to check if agent is currently processing (exempts from heartbeat timeout) */
+    isProcessing?: (agentName: string) => boolean;
+    /** Maximum messages in write queue before dropping (default: 2000) */
+    maxWriteQueueSize?: number;
+    /** High water mark - emit backpressure signal when queue exceeds this (default: 1500) */
+    writeQueueHighWaterMark?: number;
+    /** Low water mark - release backpressure when queue drops below this (default: 500) */
+    writeQueueLowWaterMark?: number;
+}
+export declare const DEFAULT_CONFIG: ConnectionConfig;
+export declare class Connection {
+    readonly id: string;
+    private socket;
+    private parser;
+    private config;
+    private _state;
+    private _agentName?;
+    private _entityType?;
+    private _cli?;
+    private _program?;
+    private _model?;
+    private _task?;
+    private _workingDirectory?;
+    private _displayName?;
+    private _avatarUrl?;
+    private _sessionId;
+    private _resumeToken;
+    private _isResumed;
+    private heartbeatTimer?;
+    private lastPongReceived?;
+    private sequences;
+    private writeQueue;
+    private draining;
+    private _backpressured;
+    private socketDrainHandler?;
+    onMessage?: (envelope: Envelope) => void;
+    onClose?: () => void;
+    onError?: (error: Error) => void;
+    onActive?: () => void;
+    onAck?: (envelope: Envelope<AckPayload>) => void;
+    onPong?: () => void;
+    /** Fires when write queue crosses high/low water marks */
+    onBackpressure?: (backpressured: boolean) => void;
+    constructor(socket: net.Socket, config?: Partial<ConnectionConfig>);
+    get state(): ConnectionState;
+    get agentName(): string | undefined;
+    get entityType(): EntityType | undefined;
+    get cli(): string | undefined;
+    get program(): string | undefined;
+    get model(): string | undefined;
+    get task(): string | undefined;
+    get workingDirectory(): string | undefined;
+    get displayName(): string | undefined;
+    get avatarUrl(): string | undefined;
+    get sessionId(): string;
+    get resumeToken(): string;
+    get isResumed(): boolean;
+    /** Whether this connection is currently backpressured (write queue above high water mark) */
+    get backpressured(): boolean;
+    /** Current number of messages queued for writing */
+    get writeQueueLength(): number;
+    private setupSocketHandlers;
+    private handleData;
+    private processFrame;
+    private handleHello;
+    private handleSend;
+    private handlePong;
+    private startHeartbeat;
+    private stopHeartbeat;
+    /**
+     * Seed a sequence counter for a stream so that the next value continues from the provided seq.
+     */
+    seedSequence(topic: string | undefined, peer: string, seq: number): void;
+    /**
+     * Get next sequence number for a stream.
+     */
+    getNextSeq(topic: string, peer: string): number;
+    /**
+     * Send an envelope to this connection.
+     *
+     * Uses a write queue to prevent blocking on slow consumers.
+     * Returns false if the connection is closed or the queue is full.
+     */
+    send(envelope: Envelope): boolean;
+    /**
+     * Schedule the drain loop to run on next tick if not already running.
+     */
+    private scheduleDrain;
+    /**
+     * Drain the write queue to the socket.
+     * Respects socket backpressure by waiting for 'drain' events.
+     */
+    private drain;
+    private sendError;
+    private handleClose;
+    private handleError;
+    private cleanup;
+    close(): void;
+}
+//# sourceMappingURL=connection.d.ts.map

package/packages/daemon/dist/connection.js

@@ -0,0 +1,438 @@
+/**
+ * Connection state machine for agent connections.
+ *
+ * States:
+ *   CONNECTING -> HANDSHAKING -> ACTIVE -> CLOSING -> CLOSED
+ *                     |            |
+ *                     v            v
+ *                   ERROR -------> CLOSED
+ */
+import { generateId } from '@agent-relay/wrapper';
+import { PROTOCOL_VERSION, } from '@agent-relay/protocol/types';
+import { encodeFrameLegacy as encodeFrame, FrameParser } from '@agent-relay/protocol/framing';
+import { DEFAULT_CONNECTION_CONFIG } from '@agent-relay/config/relay-config';
+export const DEFAULT_CONFIG = {
+    ...DEFAULT_CONNECTION_CONFIG,
+};
+export class Connection {
+    id;
+    socket;
+    parser;
+    config;
+    _state = 'CONNECTING';
+    _agentName;
+    _entityType;
+    _cli;
+    _program;
+    _model;
+    _task;
+    _workingDirectory;
+    _displayName;
+    _avatarUrl;
+    _sessionId;
+    _resumeToken;
+    _isResumed = false;
+    heartbeatTimer;
+    lastPongReceived;
+    // Sequence numbers per (topic, peer) stream
+    sequences = new Map();
+    // Write queue for backpressure handling
+    writeQueue = [];
+    draining = false;
+    _backpressured = false;
+    socketDrainHandler;
+    // Event handlers
+    onMessage;
+    onClose;
+    onError;
+    onActive; // Fires when connection transitions to ACTIVE state
+    onAck;
+    onPong; // Fires on successful heartbeat response
+    /** Fires when write queue crosses high/low water marks */
+    onBackpressure;
+    constructor(socket, config = {}) {
+        this.id = generateId();
+        this.socket = socket;
+        this.config = { ...DEFAULT_CONFIG, ...config };
+        this.parser = new FrameParser(this.config.maxFrameBytes);
+        this.parser.setLegacyMode(true); // Use 4-byte header for backwards compatibility
+        this._sessionId = generateId();
+        this._resumeToken = generateId();
+        this.setupSocketHandlers();
+        this._state = 'HANDSHAKING';
+    }
+    get state() {
+        return this._state;
+    }
+    get agentName() {
+        return this._agentName;
+    }
+    get entityType() {
+        return this._entityType;
+    }
+    get cli() {
+        return this._cli;
+    }
+    get program() {
+        return this._program;
+    }
+    get model() {
+        return this._model;
+    }
+    get task() {
+        return this._task;
+    }
+    get workingDirectory() {
+        return this._workingDirectory;
+    }
+    get displayName() {
+        return this._displayName;
+    }
+    get avatarUrl() {
+        return this._avatarUrl;
+    }
+    get sessionId() {
+        return this._sessionId;
+    }
+    get resumeToken() {
+        return this._resumeToken;
+    }
+    get isResumed() {
+        return this._isResumed;
+    }
+    /** Whether this connection is currently backpressured (write queue above high water mark) */
+    get backpressured() {
+        return this._backpressured;
+    }
+    /** Current number of messages queued for writing */
+    get writeQueueLength() {
+        return this.writeQueue.length;
+    }
+    setupSocketHandlers() {
+        this.socket.on('data', (data) => this.handleData(data));
+        this.socket.on('close', () => this.handleClose());
+        this.socket.on('error', (err) => this.handleError(err));
+    }
+    handleData(data) {
+        try {
+            const frames = this.parser.push(data);
+            for (const frame of frames) {
+                this.processFrame(frame).catch((err) => {
+                    this.sendError('BAD_REQUEST', `Frame error: ${err}`, true);
+                    this.close();
+                });
+            }
+        }
+        catch (err) {
+            this.sendError('BAD_REQUEST', `Frame error: ${err}`, true);
+            this.close();
+        }
+    }
+    async processFrame(envelope) {
+        switch (envelope.type) {
+            case 'HELLO':
+                await this.handleHello(envelope);
+                break;
+            case 'SEND':
+                this.handleSend(envelope);
+                break;
+            case 'ACK':
+                if (this.onAck) {
+                    this.onAck(envelope);
+                }
+                break;
+            case 'PONG':
+                this.handlePong(envelope);
+                break;
+            case 'BYE':
+                this.close();
+                break;
+            default:
+                if (this.onMessage) {
+                    this.onMessage(envelope);
+                }
+        }
+    }
+    async handleHello(envelope) {
+        if (this._state !== 'HANDSHAKING') {
+            this.sendError('BAD_REQUEST', 'Unexpected HELLO', false);
+            return;
+        }
+        this._agentName = envelope.payload.agent;
+        this._entityType = envelope.payload.entityType;
+        this._cli = envelope.payload.cli;
+        this._program = envelope.payload.program;
+        this._model = envelope.payload.model;
+        this._task = envelope.payload.task;
+        this._workingDirectory = envelope.payload.workingDirectory;
+        this._displayName = envelope.payload.displayName;
+        this._avatarUrl = envelope.payload.avatarUrl;
+        // Check for session resume
+        const resumeToken = envelope.payload.session?.resume_token;
+        if (resumeToken) {
+            if (this.config.resumeHandler) {
+                try {
+                    const resumeState = await this.config.resumeHandler({
+                        agent: this._agentName,
+                        resumeToken,
+                    });
+                    if (resumeState) {
+                        this._sessionId = resumeState.sessionId;
+                        this._resumeToken = resumeState.resumeToken ?? resumeToken;
+                        this._isResumed = true;
+                        // Seed sequence counters so new deliveries continue from last known seq per stream
+                        for (const seed of resumeState.seedSequences ?? []) {
+                            this.seedSequence(seed.topic ?? 'default', seed.peer, seed.seq);
+                        }
+                    }
+                    else {
+                        this.sendError('RESUME_TOO_OLD', 'Resume token rejected; starting new session', false);
+                    }
+                }
+                catch (err) {
+                    this.sendError('RESUME_TOO_OLD', `Resume validation failed: ${err?.message ?? err}`, false);
+                }
+            }
+            else {
+                this.sendError('RESUME_TOO_OLD', 'Session resume not configured; starting new session', false);
+            }
+        }
+        // Send WELCOME
+        const welcome = {
+            v: PROTOCOL_VERSION,
+            type: 'WELCOME',
+            id: generateId(),
+            ts: Date.now(),
+            payload: {
+                session_id: this._sessionId,
+                resume_token: this._resumeToken,
+                server: {
+                    max_frame_bytes: this.config.maxFrameBytes,
+                    heartbeat_ms: this.config.heartbeatMs,
+                },
+            },
+        };
+        this.send(welcome);
+        this._state = 'ACTIVE';
+        this.lastPongReceived = Date.now();
+        this.startHeartbeat();
+        // Notify that connection is now active (for registration)
+        if (this.onActive) {
+            this.onActive();
+        }
+    }
+    handleSend(envelope) {
+        if (this._state !== 'ACTIVE') {
+            this.sendError('BAD_REQUEST', 'Not in ACTIVE state', false);
+            return;
+        }
+        // Forward to router via callback
+        if (this.onMessage) {
+            this.onMessage(envelope);
+        }
+    }
+    handlePong(_envelope) {
+        // Note: envelope.payload.nonce could be used for RTT calculation in the future
+        this.lastPongReceived = Date.now();
+        if (this.onPong) {
+            this.onPong();
+        }
+    }
+    startHeartbeat() {
+        this.heartbeatTimer = setInterval(() => {
+            if (this._state !== 'ACTIVE')
+                return;
+            const now = Date.now();
+            // Check for missed pong - use configurable timeout multiplier
+            const timeoutMs = this.config.heartbeatMs * this.config.heartbeatTimeoutMultiplier;
+            if (this.lastPongReceived && now - this.lastPongReceived > timeoutMs) {
+                // Exempt agents that are actively processing (long tool calls, thinking, etc.)
+                if (this._agentName && this.config.isProcessing?.(this._agentName)) {
+                    // Agent is processing - reset the pong timer to avoid timeout
+                    // but don't kill the connection
+                    console.log(`[connection] Heartbeat timeout exemption for ${this._agentName} (processing)`);
+                }
+                else {
+                    this.handleError(new Error(`Heartbeat timeout (no pong in ${timeoutMs}ms)`));
+                    return;
+                }
+            }
+            // Send ping
+            const nonce = generateId();
+            this.send({
+                v: PROTOCOL_VERSION,
+                type: 'PING',
+                id: generateId(),
+                ts: now,
+                payload: { nonce },
+            });
+        }, this.config.heartbeatMs);
+    }
+    stopHeartbeat() {
+        if (this.heartbeatTimer) {
+            clearInterval(this.heartbeatTimer);
+            this.heartbeatTimer = undefined;
+        }
+    }
+    /**
+     * Seed a sequence counter for a stream so that the next value continues from the provided seq.
+     */
+    seedSequence(topic, peer, seq) {
+        const key = `${topic ?? 'default'}:${peer}`;
+        const current = this.sequences.get(key) ?? 0;
+        if (seq > current) {
+            this.sequences.set(key, seq);
+        }
+    }
+    /**
+     * Get next sequence number for a stream.
+     */
+    getNextSeq(topic, peer) {
+        const key = `${topic}:${peer}`;
+        const seq = (this.sequences.get(key) ?? 0) + 1;
+        this.sequences.set(key, seq);
+        return seq;
+    }
+    /**
+     * Send an envelope to this connection.
+     *
+     * Uses a write queue to prevent blocking on slow consumers.
+     * Returns false if the connection is closed or the queue is full.
+     */
+    send(envelope) {
+        if (this._state === 'CLOSED' || this._state === 'ERROR' || this._state === 'CLOSING') {
+            console.log(`[connection] Send to ${this._agentName} blocked - state: ${this._state}`);
+            return false;
+        }
+        const maxQueueSize = this.config.maxWriteQueueSize ?? 2000;
+        const highWaterMark = this.config.writeQueueHighWaterMark ?? 1500;
+        // Check queue capacity
+        if (this.writeQueue.length >= maxQueueSize) {
+            // Queue full - this is a serious condition, log it
+            console.warn(`[connection] Write queue full for ${this._agentName ?? this.id}, dropping message`);
+            return false;
+        }
+        try {
+            const frame = encodeFrame(envelope);
+            this.writeQueue.push(frame);
+            // Check if we should signal backpressure
+            if (!this._backpressured && this.writeQueue.length >= highWaterMark) {
+                this._backpressured = true;
+                this.onBackpressure?.(true);
+            }
+            // Schedule drain if not already draining
+            this.scheduleDrain();
+            return true;
+        }
+        catch (err) {
+            this.handleError(err);
+            return false;
+        }
+    }
+    /**
+     * Schedule the drain loop to run on next tick if not already running.
+     */
+    scheduleDrain() {
+        if (this.draining)
+            return;
+        this.draining = true;
+        setImmediate(() => this.drain());
+    }
+    /**
+     * Drain the write queue to the socket.
+     * Respects socket backpressure by waiting for 'drain' events.
+     */
+    drain() {
+        while (this.writeQueue.length > 0) {
+            if (this._state === 'CLOSED' || this._state === 'ERROR' || this._state === 'CLOSING') {
+                this.draining = false;
+                return;
+            }
+            const frame = this.writeQueue[0];
+            const canWrite = this.socket.write(frame);
+            this.writeQueue.shift();
+            if (!canWrite) {
+                // Socket buffer full - wait for drain event
+                if (!this.socketDrainHandler) {
+                    this.socketDrainHandler = () => {
+                        this.socketDrainHandler = undefined;
+                        this.drain();
+                    };
+                    this.socket.once('drain', this.socketDrainHandler);
+                }
+                return;
+            }
+        }
+        this.draining = false;
+        // Check if we should release backpressure
+        const lowWaterMark = this.config.writeQueueLowWaterMark ?? 500;
+        if (this._backpressured && this.writeQueue.length <= lowWaterMark) {
+            this._backpressured = false;
+            this.onBackpressure?.(false);
+        }
+    }
+    sendError(code, message, fatal) {
+        const error = {
+            v: PROTOCOL_VERSION,
+            type: 'ERROR',
+            id: generateId(),
+            ts: Date.now(),
+            payload: {
+                code: code,
+                message,
+                fatal,
+            },
+        };
+        this.send(error);
+    }
+    handleClose() {
+        this._state = 'CLOSED';
+        this.stopHeartbeat();
+        this.cleanup();
+        if (this.onClose) {
+            this.onClose();
+        }
+    }
+    handleError(err) {
+        this._state = 'ERROR';
+        this.stopHeartbeat();
+        this.cleanup();
+        if (this.onError) {
+            this.onError(err);
+        }
+        this.socket.destroy();
+    }
+    cleanup() {
+        this.parser.reset();
+        // Clear write queue
+        this.writeQueue = [];
+        this.draining = false;
+        this._backpressured = false;
+        // Remove drain handler if registered
+        if (this.socketDrainHandler) {
+            this.socket.removeListener('drain', this.socketDrainHandler);
+            this.socketDrainHandler = undefined;
+        }
+    }
+    close() {
+        if (this._state === 'CLOSED' || this._state === 'CLOSING')
+            return;
+        this._state = 'CLOSING';
+        // Write BYE message directly (not via queue) to avoid race with socket.end()
+        try {
+            const byeFrame = encodeFrame({
+                v: PROTOCOL_VERSION,
+                type: 'BYE',
+                id: generateId(),
+                ts: Date.now(),
+                payload: {},
+            });
+            this.socket.write(byeFrame);
+        }
+        catch {
+            // Ignore write errors during close
+        }
+        this.socket.end();
+    }
+}
+//# sourceMappingURL=connection.js.map