Refactoring and cleanup.

README.md

@@ -54,23 +54,21 @@ docker run -d -p 5672:5672 -p 15672:15672 cloudamqp/lavinmq:latest
 ### 1. HTTP Server (enqueues tasks)
 
 ```typescript
-import { AMQPAgentBackend, QueueingRequestHandler } from "@84codes/a2a-amqp";
+import { AMQPAgentBackend, QueuingRequestHandler } from "@84codes/a2a-amqp";
 import { A2AExpressApp } from "@a2a-js/sdk/server/express";
 import express from "express";
 
+// Create AMQP backend
 const backend = await AMQPAgentBackend.create({
   url: "amqp://localhost:5672",
   agentName: "my-agent",
-  enableWorkQueue: true,  // Enable task queuing
 });
 
-const requestHandler = new QueueingRequestHandler(
-  agentCard,
-  backend.taskStore,
-  backend.eventBusManager,
-  backend.workQueue
-);
+// Create request handler (handles task queuing + event projection)
+const requestHandler = new QueuingRequestHandler(agentCard, backend);
+await requestHandler.initialize();
 
+// Setup Express with A2A routes
 const app = express();
 new A2AExpressApp(requestHandler).setupRoutes(app, "/");
 app.listen(3000);
@@ -79,53 +77,65 @@ app.listen(3000);
 ### 2. Worker Process (processes tasks)
 
 ```typescript
-import { AMQPAgentBackend } from "@84codes/a2a-amqp";
+import { AMQPAgentBackend, WorkerEventBus } from "@84codes/a2a-amqp";
 import { AgentExecutor, RequestContext } from "@a2a-js/sdk/server";
 
+// Create backend with same agent name as server
 const backend = await AMQPAgentBackend.create({
   url: "amqp://localhost:5672",
   agentName: "my-agent",
-  enableWorkQueue: true,
 });
 
+// Initialize work queue
+await backend.workQueue.initialize();
+
 class MyExecutor implements AgentExecutor {
   async execute(context: RequestContext, eventBus: ExecutionEventBus) {
     // Your long-running task logic here
-    // Events are automatically published to AMQP stream
     eventBus.publish({
       kind: "status-update",
       taskId: context.taskId,
-      status: { state: "working" },
+      contextId: context.contextId,
+      status: { state: "working", timestamp: new Date().toISOString() },
+      final: false,
     });
 
     // ... do work ...
 
     eventBus.publish({
       kind: "status-update",
       taskId: context.taskId,
-      status: { state: "completed" },
+      contextId: context.contextId,
+      status: { state: "completed", timestamp: new Date().toISOString() },
       final: true,
     });
     eventBus.finished();
   }
 }
 
 const executor = new MyExecutor();
-await backend.workQueue.startConsumer(async (taskMsg) => {
+
+// Start consuming with async generator pattern
+const messages = backend.workQueue.start();
+
+for await (const taskMessage of messages) {
+  const { taskId, contextId, requestContext } = taskMessage;
+
+  // Create request context
   const context = new RequestContext(
-    taskMsg.requestContext.userMessage,
-    taskMsg.taskId,
-    taskMsg.contextId
+    requestContext.userMessage,
+    taskId,
+    contextId,
+    requestContext.task,
+    requestContext.referenceTasks
   );
 
-  const eventBus = backend.eventBusManager.createOrGetByTaskIdWithContext(
-    taskMsg.taskId,
-    taskMsg.contextId
-  );
+  // Create event bus for publishing task events
+  const eventBus = new WorkerEventBus(backend.amqpConnection, taskId, contextId);
 
+  // Execute task
   await executor.execute(context, eventBus);
-  backend.eventBusManager.cleanupByTaskId(taskMsg.taskId);
-});
+}
 ```
 
 ### 3. Scale horizontally
@@ -145,33 +155,41 @@ bun run worker  # Add more workers for higher throughput
 
 - **Work Queue**: Distribute tasks across multiple worker processes
 - **Event Sourcing**: All task events stored in AMQP streams for replay
-- **Task Store**: In-memory task projection with automatic recovery from streams
+- **In-Memory Projection**: Fast task lookups with automatic recovery from streams
 - **SSE Streaming**: Automatic streaming of task events back to clients
 - **Horizontal Scaling**: Add more workers to increase throughput
-- **Fault Tolerance**: Task replay from event streams on crashes
-- **Type-Safe**: Full TypeScript support
+- **Graceful Shutdown**: Clean consumer and connection handling
+- **Type-Safe**: Full TypeScript support with Zod validation
 
 ## Configuration
 
 ```typescript
 interface AMQPAgentBackendConfig {
   url: string;                    // AMQP broker URL
   agentName: string;              // Agent identifier
-  enableWorkQueue?: boolean;      // Enable task queuing (default: false)
   streamRetention?: string;       // Event retention (default: "7d")
   streamMaxBytes?: number;        // Max stream size (default: 1GB)
-  workQueueName?: string;         // Custom queue name
-  taskQueueName?: string;         // Custom task store name
+  workQueueName?: string;         // Custom work queue name
   exchangeName?: string;          // Custom exchange name
   logger?: Logger;                // Custom logger
+  connection?: {
+    heartbeat?: number;           // Heartbeat interval in seconds
+    reconnectDelay?: number;      // Reconnection delay in ms
+    maxReconnectAttempts?: number;// Max reconnection attempts
+  };
+  publishing?: {
+    persistent?: boolean;         // Persistent messages (default: true)
+    confirmMode?: boolean;        // Publisher confirms (default: true)
+    messageTtl?: number;          // Message TTL in ms (0 = no expiration)
+  };
 }
 ```
 
 ## Examples
 
 See complete working examples:
 
-- `src/examples/http-server-sdk.ts` - HTTP server with queuing
+- `src/examples/http-server.ts` - HTTP server with queuing
 - `src/examples/worker.ts` - Worker process
 
 ```bash
@@ -188,9 +206,11 @@ curl -X POST http://localhost:3000/ \
 ## Testing
 
 ```bash
-bun test              # Run all tests
-bun run test:unit     # Unit tests only
-bun run test:watch    # Watch mode
+bun test                 # Run all tests (unit + integration)
+bun run test:unit        # Unit tests only
+bun run test:integration # Integration tests only
+bun run test:watch       # Watch mode
+bun run test:coverage    # With coverage
 ```
 
 ## License

package.json

@@ -18,11 +18,11 @@
   ],
   "scripts": {
     "build": "tsc",
-    "type-check": "tsc --noEmit",
+    "type-check": "tsc --project tsconfig.typecheck.json",
     "prepare": "bun run build",
     "server": "bun run src/examples/http-server.ts",
     "worker": "bun run src/examples/worker.ts",
-    "test": "vitest run",
+    "test": "bun run test:unit && bun run test:integration",
     "test:unit": "vitest run --config vitest.config.ts",
     "test:integration": "vitest run --config vitest.integration.config.ts",
     "test:watch": "vitest watch --config vitest.config.ts",
@@ -46,7 +46,8 @@
   "author": "84codes",
   "license": "MIT",
   "dependencies": {
-    "@cloudamqp/amqp-client": "^3.1.1"
+    "@cloudamqp/amqp-client": "^3.4.0",
+    "zod": "^4.1.12"
   },
   "devDependencies": {
     "@types/bun": "latest",

src/examples/http-server.ts

@@ -1,7 +1,7 @@
 /**
  * HTTP Server Example (Using A2A SDK) - Queuing Architecture
  *
- * This server uses the A2A SDK's A2AExpressApp with a custom QueueingRequestHandler:
+ * This server uses the A2A SDK's A2AExpressApp with a custom QueuingRequestHandler:
  * 1. Accepts incoming A2A requests via HTTP (handled by SDK)
  * 2. Creates initial Task event and publishes to exchange
  * 3. Enqueues task to work queue for worker processing
@@ -20,7 +20,7 @@
 import { AgentCard } from "@a2a-js/sdk";
 import { A2AExpressApp } from "@a2a-js/sdk/server/express";
 import express from "express";
-import { AMQPAgentBackend, QueueingRequestHandler } from "../index.js";
+import { AMQPAgentBackend, QueuingRequestHandler } from "../index.js";
 
 // AMQP configuration
 const AMQP_URL = process.env.AMQP_URL || "amqp://localhost:5672";
@@ -54,16 +54,20 @@ async function main() {
 
   console.log("AMQP backend initialized");
 
-  // Log projection stats
-  const stats = backend.taskStore.getProjectionStats();
-  console.log(`TaskStore projection: ${stats.taskCount} tasks loaded into memory`);
-
   // Create queuing request handler
-  const requestHandler = new QueueingRequestHandler(
+  const requestHandler = new QueuingRequestHandler(
     agentCard,
     backend
   );
 
+  // Initialize request handler (sets up event stream)
+  await requestHandler.initialize();
+  console.log("Request handler initialized");
+
+  // Log projection stats
+  const stats = requestHandler.getProjectionStats();
+  console.log(`Event projection: ${stats.taskCount} tasks loaded into memory`);
+
   // Create A2A Express app
   const a2aExpressApp = new A2AExpressApp(requestHandler);
 
@@ -151,7 +155,14 @@ Make sure to run worker.ts in another process!
   process.on("SIGINT", async () => {
     console.log("\nShutting down...");
     server.close();
+
+    // Close request handler first (stops projection consumer)
+    await requestHandler.close();
+
+    // Then close backend
     await backend.close();
+
+    console.log("Server stopped gracefully");
     process.exit(0);
   });
 }

src/examples/worker.ts

@@ -18,12 +18,12 @@ import {
   RequestContext,
   type ExecutionEventBus,
 } from "@a2a-js/sdk/server";
-import { AMQPAgentBackend, type QueuedTaskMessage } from "../index.js";
+import { AMQPAgentBackend, WorkerEventBus, type QueuedTaskMessage } from "../index.js";
 
 // AMQP configuration
 const AMQP_URL = process.env.AMQP_URL || "amqp://localhost:5672";
 const AGENT_NAME = "greeting-agent-http";
-const WORKER_ID = process.env.WORKER_ID || `worker-${Math.random().toString(36).substr(2, 9)}`;
+const WORKER_ID = process.env.WORKER_ID || `worker-${Math.random().toString(36).slice(2, 11)}`;
 
 // Implement the agent executor
 class GreetingExecutor implements AgentExecutor {
@@ -151,48 +151,18 @@ async function main() {
     agentName: AGENT_NAME,
   });
 
-  if (!backend.workQueue) {
-    throw new Error("Work queue not initialized");
-  }
-
   const workQueue = backend.workQueue;
   const executor = new GreetingExecutor();
 
   console.log(`[${WORKER_ID}] AMQP backend initialized, starting consumer...`);
 
   // Start consuming from work queue
-  const consumerTag = await workQueue.startConsumer(async (taskMessage: QueuedTaskMessage) => {
-    const { taskId, contextId, requestContext } = taskMessage;
-
-    console.log(`[${WORKER_ID}] Received task ${taskId} from queue`);
-
-    try {
-      // Create RequestContext
-      const reqContext = new RequestContext(
-        requestContext.userMessage,
-        taskId,
-        contextId,
-        requestContext.task,
-        requestContext.referenceTasks
-      );
-
-      // Create EventBus for this task (publisher mode)
-      const eventBus = backend.createEventBus(taskId, contextId);
-
-      // Execute the task
-      await executor.execute(reqContext, eventBus);
-    } catch (error) {
-      console.error(`[${WORKER_ID}] Failed to process task ${taskId}:`, error);
-      // Error handling: task will be NACK'd and requeued by AMQPWorkQueue
-      throw error;
-    }
-  });
+  const messages = workQueue.start();
 
   console.log(`
 🔧 A2A Worker is running!
 
 Worker ID: ${WORKER_ID}
-Consumer Tag: ${consumerTag}
 Queue: ${workQueue.getQueueName()}
 
 Waiting for tasks from work queue...
@@ -210,8 +180,10 @@ Press Ctrl+C to stop.
     console.log(`\n[${WORKER_ID}] Shutting down...`);
 
     try {
-      // Stop consuming
-      await workQueue.stopConsumer(consumerTag);
+      // Close generator first to stop consuming messages
+      if (messages) {
+        await messages.return();
+      }
 
       // Close backend
       await backend.close();
@@ -226,6 +198,33 @@ Press Ctrl+C to stop.
 
   process.on("SIGINT", shutdown);
   process.on("SIGTERM", shutdown);
+
+  for await (const taskMessage of messages) {
+    const { taskId, contextId, requestContext } = taskMessage;
+
+    console.log(`[${WORKER_ID}] Received task ${taskId} from queue`);
+
+    try {
+      // Create RequestContext
+      const reqContext = new RequestContext(
+        requestContext.userMessage,
+        taskId,
+        contextId,
+        requestContext.task,
+        requestContext.referenceTasks
+      );
+
+      // Create WorkerEventBus for this task to publish events
+      const eventBus = new WorkerEventBus(backend.amqpConnection, taskId, contextId);
+
+      // Execute the task
+      await executor.execute(reqContext, eventBus);
+    } catch (error) {
+      console.error(`[${WORKER_ID}] Failed to process task ${taskId}:`, error);
+      // Error handling: task will be NACK'd and requeued by AMQPWorkQueue
+      throw error;
+    }
+  }
 }
 
 // Run the worker