fix: Refactor run handling to use run object instead of runId

control-plane/src/modules/router.ts

@@ -2,38 +2,34 @@ import { initServer } from "@ts-rest/fastify";
 import { generateOpenApi } from "@ts-rest/open-api";
 import fs from "fs";
 import path from "path";
+import { ulid } from "ulid";
 import util from "util";
-import { contract } from "./contract";
-import * as data from "./data";
-import * as management from "./management";
-import * as events from "./observability/events";
-import {
-  assertMessageOfType,
-  editHumanMessage,
-  getRunMessagesForDisplayWithPolling,
-} from "./runs/messages";
-import { addMessageAndResume, assertRunReady } from "./runs";
-import { runsRouter } from "./runs/router";
-import { machineRouter } from "./machines/router";
+import { BadRequestError } from "../utilities/errors";
+import { deleteAgent, getAgent, listAgents, upsertAgent, validateSchema } from "./agents";
 import { authRouter } from "./auth/router";
-import { ulid } from "ulid";
 import { getBlobData } from "./blobs";
-import { posthog } from "./posthog";
-import { BadRequestError } from "../utilities/errors";
-import { upsertAgent, getAgent, deleteAgent, listAgents, validateSchema } from "./agents";
+import { contract } from "./contract";
+import * as data from "./data";
+import { integrationsRouter } from "./integrations/router";
+import { jobsRouter } from "./jobs/router";
 import {
   createClusterKnowledgeArtifact,
-  getKnowledge,
-  upsertKnowledgeArtifact,
   deleteKnowledgeArtifact,
-  getKnowledgeArtifact,
   getAllKnowledgeArtifacts,
+  getKnowledge,
+  getKnowledgeArtifact,
+  upsertKnowledgeArtifact,
 } from "./knowledge/knowledgebase";
-import { jobsRouter } from "./jobs/router";
+import { machineRouter } from "./machines/router";
+import * as management from "./management";
 import { buildModel } from "./models";
-import { getServiceDefinitions, getStandardLibraryToolsMeta } from "./service-definitions";
-import { integrationsRouter } from "./integrations/router";
+import * as events from "./observability/events";
+import { posthog } from "./posthog";
+import { addMessageAndResume, assertRunReady, getRun } from "./runs";
+import { editHumanMessage, getRunMessagesForDisplayWithPolling } from "./runs/messages";
+import { runsRouter } from "./runs/router";
 import { getServerStats } from "./server-stats";
+import { getServiceDefinitions, getStandardLibraryToolsMeta } from "./service-definitions";
 
 const readFile = util.promisify(fs.readFile);
 
@@ -289,7 +285,8 @@ export const router = initServer().router(contract, {
     const { clusterId, runId, messageId } = request.params;
     const { message } = request.body;
 
-    await assertRunReady({ clusterId, runId });
+    const run = await getRun({ clusterId, runId });
+    await assertRunReady({ clusterId, run });
 
     const auth = request.request.getAuth();
     await auth.canManage({ run: { clusterId, runId } });

control-plane/src/modules/runs/index.test.ts

@@ -17,9 +17,9 @@ describe("assertRunReady", () => {
 
     await expect(
       assertRunReady({
-        runId: run.id,
+        run,
         clusterId: owner.clusterId,
-      }),
+      })
     ).resolves.not.toThrow();
   });
 
@@ -35,9 +35,9 @@ describe("assertRunReady", () => {
 
     await expect(
       assertRunReady({
-        runId: run.id,
+        run,
         clusterId: owner.clusterId,
-      }),
+      })
     ).rejects.toThrow(RunBusyError);
   });
 
@@ -49,9 +49,9 @@ describe("assertRunReady", () => {
 
     await expect(
       assertRunReady({
-        runId: run.id,
+        run,
         clusterId: owner.clusterId,
-      }),
+      })
     ).rejects.toThrow(BadRequestError);
   });
 
@@ -87,42 +87,39 @@ describe("assertRunReady", () => {
 
     await expect(
       assertRunReady({
-        runId: run.id,
+        run,
         clusterId: owner.clusterId,
-      }),
+      })
     ).resolves.not.toThrow();
   });
 
-  it.each(["human", "template"] as const)(
-    "should throw if last message is %s",
-    async (type) => {
-      const run = await createRun({
-        clusterId: owner.clusterId,
-      });
+  it.each(["human", "template"] as const)("should throw if last message is %s", async type => {
+    const run = await createRun({
+      clusterId: owner.clusterId,
+    });
 
-      await insertRunMessage({
-        id: ulid(),
-        data: {
-          message: "Some request",
-        },
-        type,
+    await insertRunMessage({
+      id: ulid(),
+      data: {
+        message: "Some request",
+      },
+      type,
+      clusterId: owner.clusterId,
+      runId: run.id,
+    });
+
+    await updateRun({
+      ...run,
+      status: "done",
+    });
+
+    await expect(
+      assertRunReady({
+        run,
         clusterId: owner.clusterId,
-        runId: run.id,
-      });
-
-      await updateRun({
-        ...run,
-        status: "done",
-      });
-
-      await expect(
-        assertRunReady({
-          runId: run.id,
-          clusterId: owner.clusterId,
-        }),
-      ).rejects.toThrow(RunBusyError);
-    },
-  );
+      })
+    ).rejects.toThrow(RunBusyError);
+  });
 
   it.each([
     {
@@ -153,7 +150,7 @@ describe("assertRunReady", () => {
       },
       type: "template" as const,
     },
-  ])("messages should throw unless AI with no tool calls", async (message) => {
+  ])("messages should throw unless AI with no tool calls", async message => {
     const run = await createRun({
       clusterId: owner.clusterId,
     });
@@ -185,16 +182,16 @@ describe("assertRunReady", () => {
       clusterId: owner.clusterId,
     });
 
-    await updateRun({
+    const updatedRun = await updateRun({
       ...run,
       status: "done",
     });
 
     await expect(
       assertRunReady({
-        runId: run.id,
+        run: updatedRun,
         clusterId: owner.clusterId,
-      }),
+      })
     ).rejects.toThrow(RunBusyError);
   });
 });