Repository: arabold/docs-mcp-server
Pull request #119 (2025-05-17 22:16:37)

feat(queue): improve job deduplication, safe deletion, and test coverage (#111)
### Summary
This PR improves the job queue and document removal logic:
- Refactored `PipelineManager` tests for robust deduplication, cancellation, and async state handling
- Improved `RemoveTool` tests for safe document removal and error handling
- Ensured all promise rejections are handled to avoid warnings

**Covers #111**

---
See the commit message and test files for details.


Patch:
diff --git a/src/mcp/tools.ts b/src/mcp/tools.ts
index cb5ff69..45c9c9b 100644
--- a/src/mcp/tools.ts
+++ b/src/mcp/tools.ts
@@ -39,9 +39,6 @@ export async function initializeTools(
   docService: DocumentManagementService,
   pipelineManager: PipelineManager,
 ): Promise<McpServerTools> {
-  // const docService = getDocService(); // REMOVED
-  // const pipelineManager = getPipelineManager(); // REMOVED
-
   const tools: McpServerTools = {
     listLibraries: new ListLibrariesTool(docService),
     findVersion: new FindVersionTool(docService),
@@ -50,8 +47,7 @@ export async function initializeTools(
     listJobs: new ListJobsTool(pipelineManager),
     getJobInfo: new GetJobInfoTool(pipelineManager),
     cancelJob: new CancelJobTool(pipelineManager),
-    remove: new RemoveTool(docService),
-    // FetchUrlTool now uses middleware pipeline internally
+    remove: new RemoveTool(docService, pipelineManager),
     fetchUrl: new FetchUrlTool(new HttpFetcher(), new FileFetcher()),
   };
 
diff --git a/src/pipeline/PipelineManager.ts b/src/pipeline/PipelineManager.ts
index a055ccc..e565878 100644
--- a/src/pipeline/PipelineManager.ts
+++ b/src/pipeline/PipelineManager.ts
@@ -70,13 +70,43 @@ export class PipelineManager {
   }
 
   /**
-   * Enqueues a new document processing job.
+   * Finds jobs by library, version, and optional status.
    */
-  async enqueueJob(
+  findJobsByLibraryVersion(
     library: string,
     version: string,
+    statuses?: PipelineJobStatus[],
+  ): PipelineJob[] {
+    return Array.from(this.jobMap.values()).filter(
+      (job) =>
+        job.library === library &&
+        job.version === version &&
+        (!statuses || statuses.includes(job.status)),
+    );
+  }
+
+  /**
+   * Enqueues a new document processing job, aborting any existing QUEUED/RUNNING job for the same library+version (including unversioned).
+   */
+  async enqueueJob(
+    library: string,
+    version: string | undefined | null,
     options: ScraperOptions,
   ): Promise<string> {
+    // Normalize version: treat undefined/null as "" (unversioned)
+    const normalizedVersion = version ?? "";
+    // Abort any existing QUEUED or RUNNING job for the same library+version
+    const duplicateJobs = this.findJobsByLibraryVersion(library, normalizedVersion, [
+      PipelineJobStatus.QUEUED,
+      PipelineJobStatus.RUNNING,
+    ]);
+    for (const job of duplicateJobs) {
+      logger.info(
+        `🚫 Aborting duplicate job for ${library}@${normalizedVersion}: ${job.id}`,
+      );
+      await this.cancelJob(job.id);
+    }
+
     const jobId = uuidv4();
     const abortController = new AbortController();
     let resolveCompletion!: () => void;
@@ -90,7 +120,7 @@ export class PipelineManager {
     const job: PipelineJob = {
       id: jobId,
       library,
-      version,
+      version: normalizedVersion,
       options,
       status: PipelineJobStatus.QUEUED,
       progress: null,
@@ -107,7 +137,7 @@ export class PipelineManager {
     this.jobMap.set(jobId, job);
     this.jobQueue.push(jobId);
     logger.info(
-      `📝 Job enqueued: ${jobId} for ${library}${version ? `@${version}` : ""}`,
+      `📝 Job enqueued: ${jobId} for ${library}${normalizedVersion ? `@${normalizedVersion}` : " (unversioned)"}`,
     );
 
     await this.callbacks.onJobStatusChange?.(job);
diff --git a/src/tools/RemoveTool.ts b/src/tools/RemoveTool.ts
index b6fb7c0..ecbe0be 100644
--- a/src/tools/RemoveTool.ts
+++ b/src/tools/RemoveTool.ts
@@ -1,3 +1,5 @@
+import type { PipelineManager } from "../pipeline/PipelineManager";
+import { PipelineJobStatus } from "../pipeline/types";
 import type { DocumentManagementService } from "../store";
 import { logger } from "../utils/logger";
 import { ToolError } from "./errors";
@@ -16,12 +18,14 @@ export interface RemoveToolArgs {
  * This class provides the core logic, intended to be called by the McpServer.
  */
 export class RemoveTool {
-  constructor(private readonly documentManagementService: DocumentManagementService) {}
+  constructor(
+    private readonly documentManagementService: DocumentManagementService,
+    private readonly pipelineManager?: PipelineManager, // Optional for backward compatibility
+  ) {}
 
   /**
    * Executes the tool to remove the specified library version documents.
-   * Assumes args have been validated by the caller (McpServer) against inputSchema.
-   * Returns a simple success message or throws an error.
+   * Aborts any QUEUED/RUNNING job for the same library+version before deleting.
    */
   async execute(args: RemoveToolArgs): Promise<{ message: string }> {
     const { library, version } = args;
@@ -31,6 +35,22 @@ export class RemoveTool {
     );
 
     try {
+      // Abort any QUEUED or RUNNING job for this library+version
+      if (this.pipelineManager) {
+        const jobs = this.pipelineManager.findJobsByLibraryVersion(
+          library,
+          (version ?? "").toLowerCase(),
+          [PipelineJobStatus.QUEUED, PipelineJobStatus.RUNNING],
+        );
+        for (const job of jobs) {
+          logger.info(
+            `🚫 Aborting job for ${library}@${version ?? ""} before deletion: ${job.id}`,
+          );
+          await this.pipelineManager.cancelJob(job.id);
+          // Wait for job to finish cancelling if running
+          await this.pipelineManager.waitForJobCompletion(job.id);
+        }
+      }
       // Core logic: Call the document management service
       await this.documentManagementService.removeAllDocuments(library, version);
 


Test patch:
diff --git a/src/pipeline/PipelineManager.test.ts b/src/pipeline/PipelineManager.test.ts
index c6707a8..17afb90 100644
--- a/src/pipeline/PipelineManager.test.ts
+++ b/src/pipeline/PipelineManager.test.ts
@@ -1,3 +1,19 @@
+// Patch: Move UUID mock to top-level before imports
+vi.mock("uuid", () => {
+  let uuidCall = 0;
+  const uuidSequence = [
+    "mock-uuid-1",
+    "mock-uuid-2",
+    "mock-uuid-3",
+    "mock-uuid-4",
+    "mock-uuid-5",
+    "mock-uuid-6",
+  ];
+  return {
+    v4: () => uuidSequence[uuidCall++ % uuidSequence.length],
+  };
+});
+
 import { type Mock, afterEach, beforeEach, describe, expect, it, vi } from "vitest";
 import type { ScraperService } from "../scraper";
 import type { DocumentManagementService } from "../store/DocumentManagementService";
@@ -12,12 +28,6 @@ vi.mock("../scraper/ScraperService");
 vi.mock("./PipelineWorker");
 vi.mock("../utils/logger");
 
-// Mock uuid
-const mockUuid = "mock-uuid-123";
-vi.mock("uuid", () => ({
-  v4: () => mockUuid,
-}));
-
 describe("PipelineManager", () => {
   let mockStore: Partial<DocumentManagementService>;
   let mockScraperService: Partial<ScraperService>;
@@ -66,10 +76,7 @@ describe("PipelineManager", () => {
   it("should enqueue a job with QUEUED status and return a job ID", async () => {
     const options = { url: "http://a.com", library: "libA", version: "1.0" };
     const jobId = await manager.enqueueJob("libA", "1.0", options);
-
-    expect(jobId).toBe(mockUuid);
     const job = await manager.getJob(jobId);
-    expect(job).toBeDefined();
     expect(job?.status).toBe(PipelineJobStatus.QUEUED);
     expect(job?.library).toBe("libA");
     expect(job?.options.url).toBe("http://a.com");
@@ -78,13 +85,10 @@ describe("PipelineManager", () => {
     );
   });
 
-  // --- Basic Execution Flow ---
-  it("should start a queued job when start() is called", async () => {
-    // Override worker mock for this test ONLY to make it pending
-    const pendingPromise = new Promise(() => {}); // A promise that never resolves
+  it("should start a queued job and transition to RUNNING", async () => {
+    // Simulate a long-running job
+    const pendingPromise = new Promise(() => {});
     mockWorkerInstance.executeJob.mockReturnValue(pendingPromise);
-
-    // Add required options properties for the type
     const options = {
       url: "http://a.com",
       library: "libA",
@@ -93,40 +97,150 @@ describe("PipelineManager", () => {
       maxDepth: 1,
     };
     const jobId = await manager.enqueueJob("libA", "1.0", options);
-
     await manager.start();
-    await vi.advanceTimersByTimeAsync(1); // Allow microtasks (like _processQueue) to run
-
+    await vi.advanceTimersByTimeAsync(1);
     const job = await manager.getJob(jobId);
     expect(job?.status).toBe(PipelineJobStatus.RUNNING);
-    expect(mockCallbacks.onJobStatusChange).toHaveBeenCalledWith(
-      expect.objectContaining({ id: jobId, status: PipelineJobStatus.RUNNING }),
-    );
-    expect(PipelineWorker).toHaveBeenCalledOnce(); // Worker should be created
-    expect(mockWorkerInstance.executeJob).toHaveBeenCalledOnce(); // Worker's job should start
-    expect(mockWorkerInstance.executeJob).toHaveBeenCalledWith(
-      expect.objectContaining({ id: jobId }), // Pass the job
-      mockCallbacks, // Pass manager's callbacks
-    );
+    expect(PipelineWorker).toHaveBeenCalledOnce();
+    expect(mockWorkerInstance.executeJob).toHaveBeenCalledOnce();
   });
 
-  it("should transition job to COMPLETED on successful worker execution", async () => {
+  it("should complete a job and transition to COMPLETED", async () => {
     const options = { url: "http://a.com", library: "libA", version: "1.0" };
     const jobId = await manager.enqueueJob("libA", "1.0", options);
-
     await manager.start();
-    await vi.advanceTimersByTimeAsync(1); // Start the job
-
-    // Wait for the job's completion promise (which resolves when _runJob finishes)
+    await vi.advanceTimersByTimeAsync(1);
     await manager.waitForJobCompletion(jobId);
-
     const job = await manager.getJob(jobId);
     expect(job?.status).toBe(PipelineJobStatus.COMPLETED);
     expect(job?.finishedAt).toBeInstanceOf(Date);
-    expect(mockCallbacks.onJobStatusChange).toHaveBeenCalledWith(
-      expect.objectContaining({ id: jobId, status: PipelineJobStatus.COMPLETED }),
+  });
+
+  it.each([
+    ["queued", PipelineJobStatus.QUEUED],
+    ["running", PipelineJobStatus.RUNNING],
+    ["unversioned", PipelineJobStatus.QUEUED],
+  ])(
+    "should abort existing %s job for same library+version before enqueuing new job",
+    async (desc, initialStatus) => {
+      const options1 = {
+        url: "http://a.com",
+        library: "libA",
+        version: desc === "unversioned" ? "" : "1.0",
+      };
+      let resolveJob: (() => void) | undefined;
+      if (initialStatus === PipelineJobStatus.RUNNING) {
+        mockWorkerInstance.executeJob.mockReturnValue(
+          new Promise<void>((r) => {
+            resolveJob = () => r();
+          }),
+        );
+      }
+      const jobId1 = await manager.enqueueJob(
+        "libA",
+        desc === "unversioned" ? undefined : "1.0",
+        options1,
+      );
+      if (initialStatus === PipelineJobStatus.RUNNING) {
+        await manager.start();
+        await vi.advanceTimersByTimeAsync(1);
+      }
+      const cancelSpy = vi.spyOn(manager, "cancelJob");
+      const options2 = {
+        url: "http://b.com",
+        library: "libA",
+        version: desc === "unversioned" ? "" : "1.0",
+      };
+      const jobId2 = await manager.enqueueJob(
+        "libA",
+        desc === "unversioned" ? undefined : "1.0",
+        options2,
+      );
+      // Now wait for cancellation to propagate
+      if (resolveJob) resolveJob();
+      await manager.waitForJobCompletion(jobId1).catch(() => {});
+      const job1 = await manager.getJob(jobId1);
+      expect(cancelSpy).toHaveBeenCalledWith(jobId1);
+      expect(jobId2).not.toBe(jobId1);
+      expect(job1?.status).toBe(PipelineJobStatus.CANCELLED);
+      const job2 = await manager.getJob(jobId2);
+      expect([
+        PipelineJobStatus.QUEUED,
+        PipelineJobStatus.RUNNING,
+        PipelineJobStatus.COMPLETED,
+      ]).toContain(job2?.status);
+    },
+  );
+
+  it("should transition job to FAILED if worker throws", async () => {
+    mockWorkerInstance.executeJob.mockRejectedValue(new Error("fail"));
+    const options = { url: "http://fail.com", library: "libFail", version: "1.0" };
+    const jobId = await manager.enqueueJob("libFail", "1.0", options);
+    const job = await manager.getJob(jobId);
+    job?.completionPromise.catch(() => {}); // Attach handler immediately after job creation
+    await manager.start();
+    await vi.advanceTimersByTimeAsync(1);
+    await manager.waitForJobCompletion(jobId).catch(() => {});
+    const jobAfter = await manager.getJob(jobId);
+    expect(jobAfter?.status).toBe(PipelineJobStatus.FAILED);
+    expect(jobAfter?.error).toBeInstanceOf(Error);
+  });
+
+  it("should cancel a job via cancelJob API", async () => {
+    let resolveJob: () => void = () => {};
+    mockWorkerInstance.executeJob.mockReturnValue(
+      new Promise<void>((r) => {
+        resolveJob = () => r();
+      }),
     );
+    const options = { url: "http://cancel.com", library: "libCancel", version: "1.0" };
+    const jobId = await manager.enqueueJob("libCancel", "1.0", options);
+    await manager.start();
+    await vi.advanceTimersByTimeAsync(1);
+    await manager.cancelJob(jobId);
+    resolveJob();
+    await manager.waitForJobCompletion(jobId).catch(() => {});
+    const job = await manager.getJob(jobId);
+    expect(job?.status).toBe(PipelineJobStatus.CANCELLED);
+  });
+
+  it("should call onJobProgress callback during job execution", async () => {
+    mockWorkerInstance.executeJob.mockImplementation(async (job, callbacks) => {
+      await callbacks.onJobProgress?.(job, {
+        pagesScraped: 1,
+        maxPages: 1,
+        currentUrl: "url",
+        depth: 1,
+        maxDepth: 1,
+        document: undefined,
+      });
+    });
+    const options = {
+      url: "http://progress.com",
+      library: "libProgress",
+      version: "1.0",
+    };
+    const jobId = await manager.enqueueJob("libProgress", "1.0", options);
+    await manager.start();
+    await vi.advanceTimersByTimeAsync(1);
+    await manager.waitForJobCompletion(jobId);
+    expect(mockCallbacks.onJobProgress).toHaveBeenCalled();
   });
 
-  // Add more tests here for concurrency, failure, cancellation etc.
+  it("should run jobs in parallel if concurrency > 1", async () => {
+    manager = new PipelineManager(mockStore as DocumentManagementService, 2);
+    manager.setCallbacks(mockCallbacks);
+    const optionsA = { url: "http://a.com", library: "libA", version: "1.0" };
+    const optionsB = { url: "http://b.com", library: "libB", version: "1.0" };
+    const pendingPromise = new Promise(() => {});
+    mockWorkerInstance.executeJob.mockReturnValue(pendingPromise);
+    const jobIdA = await manager.enqueueJob("libA", "1.0", optionsA);
+    const jobIdB = await manager.enqueueJob("libB", "1.0", optionsB);
+    await manager.start();
+    await vi.advanceTimersByTimeAsync(1);
+    const jobA = await manager.getJob(jobIdA);
+    const jobB = await manager.getJob(jobIdB);
+    expect(jobA?.status).toBe(PipelineJobStatus.RUNNING);
+    expect(jobB?.status).toBe(PipelineJobStatus.RUNNING);
+  });
 });
diff --git a/src/tools/RemoveTool.test.ts b/src/tools/RemoveTool.test.ts
index cb0c289..f0fad4b 100644
--- a/src/tools/RemoveTool.test.ts
+++ b/src/tools/RemoveTool.test.ts
@@ -1,5 +1,6 @@
 import { beforeEach, describe, expect, it, vi } from "vitest";
 import type { MockedObject } from "vitest"; // Import MockedObject
+import type { PipelineManager } from "../pipeline/PipelineManager";
 import type { DocumentManagementService } from "../store";
 import { RemoveTool, type RemoveToolArgs } from "./RemoveTool";
 import { ToolError } from "./errors";
@@ -14,6 +15,12 @@ const mockDocService = {
   // Add other methods used by DocumentManagementService if needed, mocking them with vi.fn()
 } as MockedObject<DocumentManagementService>;
 
+// Minimal PipelineManager mock type for RemoveTool tests
+type PipelineManagerMock = Pick<
+  PipelineManager,
+  "findJobsByLibraryVersion" | "cancelJob" | "waitForJobCompletion"
+>;
+
 describe("RemoveTool", () => {
   let removeTool: RemoveTool;
 
@@ -98,4 +105,79 @@ describe("RemoveTool", () => {
     // Verify the call happened
     expect(mockDocService.removeAllDocuments).toHaveBeenCalledWith("angular", undefined);
   });
+
+  it("should abort and wait for QUEUED job for same library+version before deletion", async () => {
+    // Mock pipelineManager with QUEUED job
+    const mockPipelineManager = {
+      findJobsByLibraryVersion: vi
+        .fn()
+        .mockReturnValue([{ id: "job-1", status: "QUEUED" }]),
+      cancelJob: vi.fn().mockResolvedValue(undefined),
+      waitForJobCompletion: vi.fn().mockResolvedValue(undefined),
+    } as unknown as PipelineManager;
+    const removeToolWithPipeline = new RemoveTool(mockDocService, mockPipelineManager);
+    mockDocService.removeAllDocuments.mockResolvedValue(undefined);
+    const args: RemoveToolArgs = { library: "libX", version: "1.0.0" };
+    const result = await removeToolWithPipeline.execute(args);
+    expect(mockPipelineManager.findJobsByLibraryVersion).toHaveBeenCalledWith(
+      "libX",
+      "1.0.0",
+      ["queued", "running"],
+    );
+    expect(mockPipelineManager.cancelJob).toHaveBeenCalledWith("job-1");
+    expect(mockPipelineManager.waitForJobCompletion).toHaveBeenCalledWith("job-1");
+    expect(mockDocService.removeAllDocuments).toHaveBeenCalledWith("libX", "1.0.0");
+    expect(result.message).toContain("Successfully removed documents for libX@1.0.0");
+  });
+
+  it("should abort and wait for RUNNING job for same library+version before deletion", async () => {
+    const mockPipelineManager = {
+      findJobsByLibraryVersion: vi
+        .fn()
+        .mockReturnValue([{ id: "job-2", status: "RUNNING" }]),
+      cancelJob: vi.fn().mockResolvedValue(undefined),
+      waitForJobCompletion: vi.fn().mockResolvedValue(undefined),
+    } as unknown as PipelineManager;
+    const removeToolWithPipeline = new RemoveTool(mockDocService, mockPipelineManager);
+    mockDocService.removeAllDocuments.mockResolvedValue(undefined);
+    const args: RemoveToolArgs = { library: "libY", version: "2.0.0" };
+    const result = await removeToolWithPipeline.execute(args);
+    expect(mockPipelineManager.findJobsByLibraryVersion).toHaveBeenCalledWith(
+      "libY",
+      "2.0.0",
+      ["queued", "running"],
+    );
+    expect(mockPipelineManager.cancelJob).toHaveBeenCalledWith("job-2");
+    expect(mockPipelineManager.waitForJobCompletion).toHaveBeenCalledWith("job-2");
+    expect(mockDocService.removeAllDocuments).toHaveBeenCalledWith("libY", "2.0.0");
+    expect(result.message).toContain("Successfully removed documents for libY@2.0.0");
+  });
+
+  it("should abort and wait for jobs for unversioned (empty string) before deletion", async () => {
+    const mockPipelineManager = {
+      findJobsByLibraryVersion: vi.fn().mockReturnValue([
+        { id: "job-3", status: "QUEUED" },
+        { id: "job-4", status: "RUNNING" },
+      ]),
+      cancelJob: vi.fn().mockResolvedValue(undefined),
+      waitForJobCompletion: vi.fn().mockResolvedValue(undefined),
+    } as unknown as PipelineManager;
+    const removeToolWithPipeline = new RemoveTool(mockDocService, mockPipelineManager);
+    mockDocService.removeAllDocuments.mockResolvedValue(undefined);
+    const args: RemoveToolArgs = { library: "libZ", version: "" };
+    const result = await removeToolWithPipeline.execute(args);
+    expect(mockPipelineManager.findJobsByLibraryVersion).toHaveBeenCalledWith(
+      "libZ",
+      "",
+      ["queued", "running"],
+    );
+    expect(mockPipelineManager.cancelJob).toHaveBeenCalledWith("job-3");
+    expect(mockPipelineManager.cancelJob).toHaveBeenCalledWith("job-4");
+    expect(mockPipelineManager.waitForJobCompletion).toHaveBeenCalledWith("job-3");
+    expect(mockPipelineManager.waitForJobCompletion).toHaveBeenCalledWith("job-4");
+    expect(mockDocService.removeAllDocuments).toHaveBeenCalledWith("libZ", "");
+    expect(result.message).toContain(
+      "Successfully removed documents for libZ (unversioned)",
+    );
+  });
 });