Harden task-pod claim with atomic reservation and rollback

src/__tests__/integration/api/agent.test.ts

@@ -10,7 +10,10 @@ import {
 import { createTestUser } from "@/__tests__/support/factories/user.factory";
 import { createTestTask } from "@/__tests__/support/factories/task.factory";
 import { createTestWorkspace } from "@/__tests__/support/factories/workspace.factory";
+import { createTestSwarm } from "@/__tests__/support/factories/swarm.factory";
+import { createTestPod } from "@/__tests__/support/factories/pod.factory";
 import { db } from "@/lib/db";
+import { PodUsageStatus } from "@prisma/client";
 
 // Mock the gooseWeb provider
 const mockStreamText = vi.fn();
@@ -182,6 +185,85 @@ describe("POST /api/agent Integration Tests", () => {
     });
   });
 
+  describe("Pod claim rollback", () => {
+    test("releases the pod and clears task state when claim setup fails after reservation", async () => {
+      const user = await createTestUser();
+      const workspace = await createTestWorkspace({ ownerId: user.id });
+      const swarm = await createTestSwarm({
+        workspaceId: workspace.id,
+        status: "ACTIVE",
+      });
+
+      await db.swarm.update({
+        where: { id: swarm.id },
+        data: {
+          poolApiKey: "test-pool-api-key",
+        },
+      });
+
+      const pod = await createTestPod({
+        swarmId: swarm.id,
+        portMappings: [3000],
+        password: "plain-password",
+      });
+
+      const task = await createTestTask({
+        workspaceId: workspace.id,
+        createdById: user.id,
+        title: "Rollback task",
+      });
+
+      await db.task.update({
+        where: { id: task.id },
+        data: { mode: "agent" },
+      });
+
+      getMockedSession().mockResolvedValue(createAuthenticatedSession(user));
+
+      const request = createPostRequest("http://localhost/api/agent", {
+        message: "Start the agent",
+        taskId: task.id,
+      });
+
+      const response = await POST(request);
+      const data = await response.json();
+
+      expect(response.status).toBe(503);
+      expect(data.error).toBe("No pods available");
+      expect(mockFetch).not.toHaveBeenCalled();
+
+      const [updatedTask, updatedPod] = await Promise.all([
+        db.task.findUnique({
+          where: { id: task.id },
+          select: {
+            podId: true,
+            agentUrl: true,
+            agentPassword: true,
+          },
+        }),
+        db.pod.findUnique({
+          where: { id: pod.id },
+          select: {
+            usageStatus: true,
+            usageStatusMarkedAt: true,
+            usageStatusMarkedBy: true,
+          },
+        }),
+      ]);
+
+      expect(updatedTask).toEqual({
+        podId: null,
+        agentUrl: null,
+        agentPassword: null,
+      });
+      expect(updatedPod).toEqual({
+        usageStatus: PodUsageStatus.UNUSED,
+        usageStatusMarkedAt: null,
+        usageStatusMarkedBy: null,
+      });
+    });
+  });
+
   // NOTE: Most tests commented out due to significant implementation gaps:
   // 1. Production code does NOT validate message field (empty, missing, whitespace)
   // 2. Task/ChatMessage persistence only works with taskId (no standalone messages saved without task)

src/__tests__/integration/api/chat-response-notifications.test.ts

@@ -9,7 +9,8 @@ import { db } from "@/lib/db";
 import { POST } from "@/app/api/chat/response/route";
 import { NextRequest } from "next/server";
 import { resetDatabase } from "@/__tests__/support/utilities/database";
-import { NotificationTriggerType, NotificationTriggerStatus } from "@prisma/client";
+import { NotificationTriggerType, NotificationTriggerStatus, PodUsageStatus } from "@prisma/client";
+import { EncryptionService } from "@/lib/encryption";
 
 // Mock Sphinx delivery
 vi.mock("@/lib/sphinx/direct-message", () => ({
@@ -201,4 +202,168 @@ describe("POST /api/chat/response — plan artifact notifications", () => {
     expect(record!.sendAfter!.getTime()).toBeGreaterThan(Date.now() + 4 * 60 * 1000);
     expect(record!.message).toBeTruthy();
   });
+
+  it("does not mutate task pod fields when an artifact references a nonexistent pod", async () => {
+    const task = await db.task.create({
+      data: {
+        title: "Missing pod task",
+        workspaceId: workspace.id,
+        createdById: owner.id,
+        updatedById: owner.id,
+      },
+    });
+
+    const req = makeRequest({
+      taskId: task.id,
+      message: "IDE opened",
+      artifacts: [
+        {
+          type: "IDE",
+          content: {
+            url: "https://ide.test",
+            podId: "missing-pod-id",
+          },
+        },
+      ],
+    });
+
+    const res = await POST(req);
+    expect(res.status).toBe(201);
+
+    const [updatedTask, claimedPodCount] = await Promise.all([
+      db.task.findUnique({
+        where: { id: task.id },
+        select: {
+          podId: true,
+          agentUrl: true,
+          agentPassword: true,
+        },
+      }),
+      db.pod.count({
+        where: {
+          usageStatusMarkedBy: task.id,
+        },
+      }),
+    ]);
+
+    expect(updatedTask).toEqual({
+      podId: null,
+      agentUrl: null,
+      agentPassword: null,
+    });
+    expect(claimedPodCount).toBe(0);
+  });
+
+  it("attaches a valid artifact pod to the task and mirrors pod usage state", async () => {
+    const { createTestSwarm } = await import("@/__tests__/support/factories/swarm.factory");
+    const { createTestPod } = await import("@/__tests__/support/factories/pod.factory");
+    const encryptionService = EncryptionService.getInstance();
+
+    const swarm = await createTestSwarm({
+      workspaceId: workspace.id,
+      status: "ACTIVE",
+    });
+
+    const pod = await createTestPod({
+      swarmId: swarm.id,
+      password: "plain-pod-password",
+      portMappings: [3000, 15552],
+    });
+
+    const task = await db.task.create({
+      data: {
+        title: "Attach artifact pod task",
+        workspaceId: workspace.id,
+        createdById: owner.id,
+        updatedById: owner.id,
+      },
+    });
+
+    const req = makeRequest({
+      taskId: task.id,
+      message: "IDE opened",
+      artifacts: [
+        {
+          type: "IDE",
+          content: {
+            url: "https://ide.test",
+            podId: pod.podId,
+            agentPassword: "artifact-secret",
+          },
+        },
+      ],
+    });
+
+    const res = await POST(req);
+    expect(res.status).toBe(201);
+
+    const [updatedTask, updatedPod] = await Promise.all([
+      db.task.findUnique({
+        where: { id: task.id },
+        select: {
+          podId: true,
+          agentUrl: true,
+          agentPassword: true,
+        },
+      }),
+      db.pod.findUnique({
+        where: { id: pod.id },
+        select: {
+          usageStatus: true,
+          usageStatusMarkedBy: true,
+        },
+      }),
+    ]);
+
+    expect(updatedTask?.podId).toBe(pod.podId);
+    expect(updatedTask?.agentUrl).toBeNull();
+    expect(updatedTask?.agentPassword).toBeTruthy();
+    expect(encryptionService.decryptField("agentPassword", updatedTask!.agentPassword!)).toBe("artifact-secret");
+    expect(updatedPod).toEqual({
+      usageStatus: PodUsageStatus.USED,
+      usageStatusMarkedBy: task.id,
+    });
+  });
+
+  it("does not store agentPassword when the task podId points to a missing pod", async () => {
+    const task = await db.task.create({
+      data: {
+        title: "Stale pod task",
+        workspaceId: workspace.id,
+        createdById: owner.id,
+        updatedById: owner.id,
+        podId: "stale-pod-id",
+      },
+    });
+
+    const req = makeRequest({
+      taskId: task.id,
+      message: "IDE reopened",
+      artifacts: [
+        {
+          type: "IDE",
+          content: {
+            url: "https://ide.test",
+            agentPassword: "artifact-secret",
+          },
+        },
+      ],
+    });
+
+    const res = await POST(req);
+    expect(res.status).toBe(201);
+
+    const updatedTask = await db.task.findUnique({
+      where: { id: task.id },
+      select: {
+        podId: true,
+        agentPassword: true,
+      },
+    });
+
+    expect(updatedTask).toEqual({
+      podId: "stale-pod-id",
+      agentPassword: null,
+    });
+  });
 });