agent-relay 2.0.29 → 2.0.33

package/packages/storage/src/batched-sqlite-adapter.test.ts

@@ -0,0 +1,240 @@
+import { describe, it, expect, beforeEach, afterEach, vi } from 'vitest';
+import fs from 'node:fs';
+import os from 'node:os';
+import path from 'node:path';
+import { BatchedSqliteAdapter } from './batched-sqlite-adapter.js';
+import type { StoredMessage } from './adapter.js';
+
+const makeMessage = (overrides: Partial<StoredMessage> = {}): StoredMessage => ({
+  id: overrides.id ?? `msg-${Math.random().toString(36).slice(2, 10)}`,
+  ts: overrides.ts ?? Date.now(),
+  from: overrides.from ?? 'AgentA',
+  to: overrides.to ?? 'AgentB',
+  topic: overrides.topic,
+  kind: overrides.kind ?? 'message',
+  body: overrides.body ?? 'hello',
+  data: overrides.data,
+  payloadMeta: overrides.payloadMeta,
+  thread: overrides.thread,
+  deliverySeq: overrides.deliverySeq,
+  deliverySessionId: overrides.deliverySessionId,
+  sessionId: overrides.sessionId,
+  status: overrides.status ?? 'unread',
+  is_urgent: overrides.is_urgent ?? false,
+  is_broadcast: overrides.is_broadcast ?? false,
+});
+
+describe('BatchedSqliteAdapter', () => {
+  let dbPath: string;
+  let adapter: BatchedSqliteAdapter;
+
+  beforeEach(async () => {
+    const tmpDir = fs.mkdtempSync(path.join(os.tmpdir(), 'relay-batch-sqlite-'));
+    dbPath = path.join(tmpDir, 'messages.sqlite');
+    adapter = new BatchedSqliteAdapter({
+      dbPath,
+      batch: {
+        maxBatchSize: 5,
+        maxBatchDelayMs: 50,
+        maxBatchBytes: 10 * 1024,
+        logBatches: false,
+      },
+    });
+    await adapter.init();
+  });
+
+  afterEach(async () => {
+    await adapter.close();
+    try {
+      fs.rmSync(path.dirname(dbPath), { recursive: true, force: true });
+    } catch {
+      // ignore cleanup errors
+    }
+  });
+
+  describe('saveMessage batching', () => {
+    it('queues messages and flushes when batch size is reached', async () => {
+      // Save 5 messages (batch size threshold)
+      for (let i = 0; i < 5; i++) {
+        await adapter.saveMessage(makeMessage({ id: `batch-msg-${i}` }));
+      }
+
+      // Messages should be flushed after reaching batch size
+      const messages = await adapter.getMessages();
+      expect(messages).toHaveLength(5);
+    });
+
+    it('flushes on time threshold when batch size not reached', async () => {
+      vi.useFakeTimers();
+      try {
+        // Save fewer messages than batch size
+        await adapter.saveMessage(makeMessage({ id: 'time-msg-1' }));
+        await adapter.saveMessage(makeMessage({ id: 'time-msg-2' }));
+
+        // Advance past the delay threshold
+        await vi.advanceTimersByTimeAsync(100);
+
+        // Force flush to ensure completion
+        await adapter.flush();
+
+        const messages = await adapter.getMessages();
+        expect(messages).toHaveLength(2);
+      } finally {
+        vi.useRealTimers();
+      }
+    });
+
+    it('flushes when memory threshold is exceeded', async () => {
+      // Create adapter with small memory threshold
+      await adapter.close();
+
+      adapter = new BatchedSqliteAdapter({
+        dbPath,
+        batch: {
+          maxBatchSize: 100, // Large batch size so it won't trigger
+          maxBatchDelayMs: 5000, // Long delay so it won't trigger
+          maxBatchBytes: 100, // Small byte threshold
+          logBatches: false,
+        },
+      });
+      await adapter.init();
+
+      // Create a message with large body to exceed byte threshold
+      const largeBody = 'x'.repeat(200);
+      await adapter.saveMessage(makeMessage({ id: 'large-msg', body: largeBody }));
+
+      // Should have been flushed due to byte threshold
+      const messages = await adapter.getMessages();
+      expect(messages).toHaveLength(1);
+      expect(messages[0].body).toBe(largeBody);
+    });
+
+    it('tracks metrics for flush reasons', async () => {
+      // Save batch size worth of messages
+      for (let i = 0; i < 5; i++) {
+        await adapter.saveMessage(makeMessage({ id: `metric-msg-${i}` }));
+      }
+
+      const metrics = adapter.getBatchMetrics();
+      expect(metrics.batchesWritten).toBe(1);
+      expect(metrics.messagesWritten).toBe(5);
+      expect(metrics.flushDueToSize).toBe(1);
+    });
+  });
+
+  describe('flush behavior', () => {
+    it('flush is idempotent when queue is empty', async () => {
+      // Multiple flushes on empty queue should not error
+      await adapter.flush();
+      await adapter.flush();
+      await adapter.flush();
+
+      const metrics = adapter.getBatchMetrics();
+      expect(metrics.batchesWritten).toBe(0);
+    });
+
+    it('does not flush concurrently', async () => {
+      // Queue some messages
+      for (let i = 0; i < 3; i++) {
+        adapter.saveMessage(makeMessage({ id: `concurrent-msg-${i}` }));
+      }
+
+      // Call flush multiple times concurrently
+      await Promise.all([
+        adapter.flush(),
+        adapter.flush(),
+        adapter.flush(),
+      ]);
+
+      const metrics = adapter.getBatchMetrics();
+      // Only one batch should have been written
+      expect(metrics.batchesWritten).toBe(1);
+    });
+
+    it('re-queues messages on write failure', async () => {
+      // We can't easily test this without mocking internals,
+      // but we can verify the basic happy path
+      await adapter.saveMessage(makeMessage({ id: 'requeue-msg' }));
+      await adapter.flush();
+
+      const messages = await adapter.getMessages();
+      expect(messages).toHaveLength(1);
+    });
+  });
+
+  describe('close behavior', () => {
+    it('flushes pending messages on close', async () => {
+      // Save messages without triggering batch threshold
+      await adapter.saveMessage(makeMessage({ id: 'close-msg-1' }));
+      await adapter.saveMessage(makeMessage({ id: 'close-msg-2' }));
+
+      // Close should flush pending messages
+      await adapter.close();
+
+      // Re-open to verify messages were saved
+      adapter = new BatchedSqliteAdapter({ dbPath });
+      await adapter.init();
+
+      const messages = await adapter.getMessages();
+      expect(messages).toHaveLength(2);
+    });
+  });
+
+  describe('metrics', () => {
+    it('tracks pending count and bytes', async () => {
+      // Create adapter with large batch to keep messages pending
+      await adapter.close();
+      adapter = new BatchedSqliteAdapter({
+        dbPath,
+        batch: {
+          maxBatchSize: 100,
+          maxBatchDelayMs: 10000,
+          maxBatchBytes: 1024 * 1024,
+          logBatches: false,
+        },
+      });
+      await adapter.init();
+
+      await adapter.saveMessage(makeMessage({ id: 'pending-1', body: 'hello' }));
+      await adapter.saveMessage(makeMessage({ id: 'pending-2', body: 'world' }));
+
+      const metrics = adapter.getBatchMetrics();
+      expect(metrics.pendingCount).toBe(2);
+      expect(metrics.pendingBytes).toBeGreaterThan(0);
+
+      // Flush and check metrics are cleared
+      await adapter.flush();
+      const afterFlush = adapter.getBatchMetrics();
+      expect(afterFlush.pendingCount).toBe(0);
+      expect(afterFlush.pendingBytes).toBe(0);
+    });
+
+    it('resetMetrics clears counters', async () => {
+      for (let i = 0; i < 5; i++) {
+        await adapter.saveMessage(makeMessage({ id: `reset-msg-${i}` }));
+      }
+
+      let metrics = adapter.getBatchMetrics();
+      expect(metrics.batchesWritten).toBe(1);
+
+      adapter.resetMetrics();
+      metrics = adapter.getBatchMetrics();
+      expect(metrics.batchesWritten).toBe(0);
+      expect(metrics.messagesWritten).toBe(0);
+    });
+  });
+
+  describe('updateMessageStatus', () => {
+    it('updates status immediately without batching', async () => {
+      // Save and flush a message
+      await adapter.saveMessage(makeMessage({ id: 'status-msg', status: 'unread' }));
+      await adapter.flush();
+
+      // Update status - should be immediate
+      await adapter.updateMessageStatus('status-msg', 'read');
+
+      const messages = await adapter.getMessages();
+      expect(messages[0].status).toBe('read');
+    });
+  });
+});

package/packages/storage/src/batched-sqlite-adapter.ts

@@ -0,0 +1,239 @@
+/**
+ * Batched SQLite Storage Adapter
+ *
+ * Wraps SqliteStorageAdapter to provide batched writes for improved throughput.
+ * Messages are buffered and flushed either when:
+ *   - Batch size is reached
+ *   - Time threshold is exceeded
+ *   - Memory threshold is exceeded
+ *   - close() is called
+ *
+ * All reads still go directly to SQLite for consistency.
+ */
+
+import {
+  SqliteStorageAdapter,
+  type SqliteAdapterOptions,
+} from './sqlite-adapter.js';
+import type { StoredMessage, MessageStatus } from './adapter.js';
+
+export interface BatchConfig {
+  /** Maximum messages in a batch before flush (default: 50) */
+  maxBatchSize: number;
+  /** Maximum time to wait before flush in ms (default: 100) */
+  maxBatchDelayMs: number;
+  /** Maximum bytes in memory before flush (default: 1MB) */
+  maxBatchBytes: number;
+  /** Whether to log batch operations (default: false) */
+  logBatches: boolean;
+}
+
+export const DEFAULT_BATCH_CONFIG: BatchConfig = {
+  maxBatchSize: 50,
+  maxBatchDelayMs: 100,
+  maxBatchBytes: 1024 * 1024, // 1MB
+  logBatches: false,
+};
+
+export interface BatchedSqliteAdapterOptions extends SqliteAdapterOptions {
+  batch?: Partial<BatchConfig>;
+}
+
+interface PendingMessage {
+  message: StoredMessage;
+  sizeBytes: number;
+}
+
+export class BatchedSqliteAdapter extends SqliteStorageAdapter {
+  private batchConfig: BatchConfig;
+  private pending: PendingMessage[] = [];
+  private pendingBytes = 0;
+  private flushTimer?: NodeJS.Timeout;
+  private flushing = false;
+  private flushPromise?: Promise<void>;
+
+  // Metrics
+  private metrics = {
+    batchesWritten: 0,
+    messagesWritten: 0,
+    bytesWritten: 0,
+    flushDueToSize: 0,
+    flushDueToTime: 0,
+    flushDueToBytes: 0,
+  };
+
+  constructor(options: BatchedSqliteAdapterOptions) {
+    super(options);
+    this.batchConfig = { ...DEFAULT_BATCH_CONFIG, ...options.batch };
+  }
+
+  /**
+   * Queue a message for batched writing.
+   * May trigger an immediate flush if thresholds are exceeded.
+   */
+  async saveMessage(message: StoredMessage): Promise<void> {
+    // Ensure any pending flush completes first
+    if (this.flushPromise) {
+      await this.flushPromise;
+    }
+
+    const msgJson = JSON.stringify(message);
+    const sizeBytes = Buffer.byteLength(msgJson, 'utf-8');
+
+    this.pending.push({ message, sizeBytes });
+    this.pendingBytes += sizeBytes;
+
+    // Check flush conditions
+    let flushReason: 'size' | 'bytes' | null = null;
+
+    if (this.pending.length >= this.batchConfig.maxBatchSize) {
+      flushReason = 'size';
+      this.metrics.flushDueToSize++;
+    } else if (this.pendingBytes >= this.batchConfig.maxBatchBytes) {
+      flushReason = 'bytes';
+      this.metrics.flushDueToBytes++;
+    }
+
+    if (flushReason) {
+      await this.flush();
+    } else if (!this.flushTimer) {
+      // Schedule time-based flush
+      this.flushTimer = setTimeout(() => {
+        this.metrics.flushDueToTime++;
+        this.flush().catch((err) => {
+          console.error('[batched-sqlite] Timer flush failed:', err);
+        });
+      }, this.batchConfig.maxBatchDelayMs);
+    }
+  }
+
+  /**
+   * Flush all pending messages to SQLite.
+   */
+  async flush(): Promise<void> {
+    // Clear timer if set
+    if (this.flushTimer) {
+      clearTimeout(this.flushTimer);
+      this.flushTimer = undefined;
+    }
+
+    // Skip if nothing to flush or already flushing
+    if (this.pending.length === 0 || this.flushing) {
+      return;
+    }
+
+    this.flushing = true;
+
+    // Take current batch
+    const batch = this.pending;
+    const batchBytes = this.pendingBytes;
+    this.pending = [];
+    this.pendingBytes = 0;
+
+    this.flushPromise = this.writeBatch(batch, batchBytes);
+
+    try {
+      await this.flushPromise;
+    } finally {
+      this.flushing = false;
+      this.flushPromise = undefined;
+    }
+  }
+
+  /**
+   * Write a batch of messages within a transaction.
+   */
+  private async writeBatch(batch: PendingMessage[], batchBytes: number): Promise<void> {
+    const startTime = Date.now();
+
+    try {
+      // Use transaction for atomicity and performance
+      // The parent class's db is private, so we call saveMessage in a loop
+      // but the SQLite driver will batch them efficiently due to WAL mode
+      for (const { message } of batch) {
+        await super.saveMessage(message);
+      }
+
+      // Update metrics
+      this.metrics.batchesWritten++;
+      this.metrics.messagesWritten += batch.length;
+      this.metrics.bytesWritten += batchBytes;
+
+      if (this.batchConfig.logBatches) {
+        const elapsed = Date.now() - startTime;
+        console.log(
+          `[batched-sqlite] Wrote ${batch.length} messages (${(batchBytes / 1024).toFixed(1)}KB) in ${elapsed}ms`
+        );
+      }
+    } catch (err) {
+      // On failure, re-queue messages for retry
+      console.error('[batched-sqlite] Batch write failed, re-queuing:', err);
+      this.pending = [...batch, ...this.pending];
+      this.pendingBytes += batchBytes;
+      throw err;
+    }
+  }
+
+  /**
+   * Get batch metrics for monitoring.
+   */
+  getBatchMetrics(): typeof this.metrics & { pendingCount: number; pendingBytes: number } {
+    return {
+      ...this.metrics,
+      pendingCount: this.pending.length,
+      pendingBytes: this.pendingBytes,
+    };
+  }
+
+  /**
+   * Reset metrics (useful for testing or periodic reporting).
+   */
+  resetMetrics(): void {
+    this.metrics = {
+      batchesWritten: 0,
+      messagesWritten: 0,
+      bytesWritten: 0,
+      flushDueToSize: 0,
+      flushDueToTime: 0,
+      flushDueToBytes: 0,
+    };
+  }
+
+  /**
+   * Close the adapter, flushing any pending messages first.
+   */
+  async close(): Promise<void> {
+    // Ensure all pending messages are written
+    if (this.flushTimer) {
+      clearTimeout(this.flushTimer);
+      this.flushTimer = undefined;
+    }
+
+    // Wait for any in-progress flush
+    if (this.flushPromise) {
+      await this.flushPromise;
+    }
+
+    // Flush remaining
+    if (this.pending.length > 0) {
+      await this.flush();
+    }
+
+    await super.close();
+  }
+
+  /**
+   * Update message status - goes directly to SQLite (not batched).
+   */
+  async updateMessageStatus(id: string, status: MessageStatus): Promise<void> {
+    // Status updates should be immediate, not batched
+    return super.updateMessageStatus(id, status);
+  }
+
+  /**
+   * Get channel memberships for an agent - read-only, delegates to parent.
+   */
+  async getChannelMembershipsForAgent(memberName: string): Promise<string[]> {
+    return super.getChannelMembershipsForAgent(memberName);
+  }
+}