Merge branch 'fix_event_idempotency' of https://github.com/vr-varad/mosoo into fix_event_idempotency

Author: vr-varad

apps/api/package.json

@@ -16,6 +16,7 @@
     "@cloudflare/sandbox": "0.12.3",
     "@larksuiteoapi/node-sdk": "^1.66.0",
     "@mosoo/ag-ui-session": "workspace:*",
+    "@mosoo/agent-driver": "workspace:*",
     "@mosoo/agent-package": "workspace:*",
     "@mosoo/contracts": "workspace:*",
     "@mosoo/db": "workspace:*",
@@ -28,7 +29,6 @@
     "@mosoo/session-policy": "workspace:*",
     "@mosoo/skill-package": "workspace:*",
     "@orpc/server": "^1.14.3",
-    "agent-driver": "workspace:*",
     "arktype": "^2.2.0",
     "better-auth": "1.6.23",
     "cloudflare": "7.0.0",

apps/api/src/modules/agents/application/agent-versioned-config.service.ts

@@ -1,3 +1,5 @@
+import { isSupportedDriverRuntime } from "@mosoo/agent-driver/runtime";
+import type { DriverRuntime } from "@mosoo/agent-driver/runtime";
 import type { JsonObject } from "@mosoo/contracts";
 import type { AgentBuiltInToolConfig, AgentEnvironmentConfig } from "@mosoo/contracts/agent";
 import { classifyAgentConfigChanges } from "@mosoo/contracts/agent-config-change-plan";
@@ -7,8 +9,6 @@ import type {
 } from "@mosoo/contracts/agent-config-change-plan";
 import type { AgentId, McpServerId, SkillId } from "@mosoo/id";
 import { getRuntimeCatalogEntry } from "@mosoo/runtime-catalog";
-import { isSupportedDriverRuntime } from "agent-driver/runtime";
-import type { DriverRuntime } from "agent-driver/runtime";
 
 import { listEditableAgentSkillReferences } from "./agent-deployment-version.service";
 import type { AgentRow } from "./agent-types";

apps/api/src/modules/api-command/application/api-command-enqueue.ts

@@ -2,7 +2,7 @@ import type { AppDeploymentRunId } from "@mosoo/id";
 
 import type { ApiBindings } from "../../../platform/cloudflare/worker-types";
 import { admitApiCommand, enqueueApiCommand } from "./api-command-ledger";
-import type { ApiCommandAdmission } from "./api-command-ledger";
+import type { ApiCommandAdmission, EnqueueApiCommandInput } from "./api-command-ledger";
 import type {
   AppDeploymentRunDispatchCommandPayload,
   ChannelWorkTriggerCommandPayload,
@@ -74,9 +74,15 @@ export async function admitSessionRunDispatchCommand(
   bindings: Pick<ApiBindings, "API_COMMAND_QUEUE" | "DB">,
   payload: SessionRunDispatchCommandPayload,
 ): Promise<ApiCommandAdmission> {
-  return admitApiCommand(bindings, {
+  return admitApiCommand(bindings, createSessionRunDispatchApiCommandInput(payload));
+}
+
+export function createSessionRunDispatchApiCommandInput(
+  payload: SessionRunDispatchCommandPayload,
+): EnqueueApiCommandInput {
+  return {
     dedupeKey: `session_run_dispatch:${payload.sessionRunId}`,
     kind: "session_run_dispatch",
     payload,
-  });
+  };
 }

apps/api/src/modules/api-command/application/api-command-ledger.ts

@@ -30,6 +30,11 @@ export interface EnqueueApiCommandInput {
   retryTerminal?: boolean;
 }
 
+export interface PreparedApiCommand {
+  commandId: ApiCommandId;
+  record: ApiCommandRow;
+}
+
 export interface ApiCommandClaim {
   attemptCount: number;
   commandId: ApiCommandId;
@@ -188,40 +193,51 @@ export async function redriveFailedApiCommandEnqueues(
   }
 }
 
-export async function admitApiCommand(
-  bindings: ApiCommandDeliveryBindings,
+export function prepareApiCommand(
   input: EnqueueApiCommandInput,
-): Promise<ApiCommandAdmission> {
-  const nowMs = currentTimestampMs();
-  const commandId = createPlatformId<ApiCommandId>();
-  const dedupeKey = normalizeDedupeKey(input.dedupeKey);
-  const payloadJson = JSON.stringify(input.payload);
-  const database = getAppDatabase(bindings.DB);
-  const insertResult = await database
-    .insert(apiCommandsTable)
-    .values({
+  options: { commandId?: ApiCommandId; timestampMs?: number } = {},
+): PreparedApiCommand {
+  const timestampMs = options.timestampMs ?? currentTimestampMs();
+  const commandId = options.commandId ?? createPlatformId<ApiCommandId>();
+
+  return {
+    commandId,
+    record: {
       attemptCount: 0,
       claimExpiresAt: null,
       claimOwner: null,
       completedAt: null,
-      createdAt: nowMs,
-      dedupeKey,
+      createdAt: timestampMs,
+      dedupeKey: normalizeDedupeKey(input.dedupeKey),
       id: commandId,
       kind: input.kind,
       lastErrorCode: API_COMMAND_QUEUE_DELIVERY_PENDING_CODE,
       lastErrorMessage: API_COMMAND_QUEUE_DELIVERY_PENDING_MESSAGE,
-      payloadJson,
+      payloadJson: JSON.stringify(input.payload),
       status: "queued",
-      updatedAt: nowMs,
-    })
+      updatedAt: timestampMs,
+    },
+  };
+}
+
+export async function admitApiCommand(
+  bindings: ApiCommandDeliveryBindings,
+  input: EnqueueApiCommandInput,
+): Promise<ApiCommandAdmission> {
+  const prepared = prepareApiCommand(input);
+  const database = getAppDatabase(bindings.DB);
+
+  const insertResult = await database
+    .insert(apiCommandsTable)
+    .values(prepared.record)
     .onConflictDoNothing()
     .run();
 
   if (getD1ChangeCount(insertResult) > 0) {
-    return { commandId, kind: input.kind, shouldDeliver: true };
+    return { commandId: prepared.commandId, kind: input.kind, shouldDeliver: true };
   }
 
-  const current = await findApiCommandByDedupeKey(bindings.DB, dedupeKey);
+  const current = await findApiCommandByDedupeKey(bindings.DB, prepared.record.dedupeKey);
 
   if (current === null) {
     throw new Error("API command enqueue could not confirm the ledger row.");
@@ -237,9 +253,9 @@ export async function admitApiCommand(
         completedAt: null,
         lastErrorCode: API_COMMAND_QUEUE_DELIVERY_PENDING_CODE,
         lastErrorMessage: API_COMMAND_QUEUE_DELIVERY_PENDING_MESSAGE,
-        payloadJson,
+        payloadJson: prepared.record.payloadJson,
         status: "queued",
-        updatedAt: nowMs,
+        updatedAt: prepared.record.updatedAt,
       })
       .where(
         and(

apps/api/src/modules/public-api/public-api-errors.ts

@@ -39,6 +39,8 @@ export function toPublicApiError(error: unknown): PublicApiError | null {
       return publicForbidden(error.message);
     case API_ERROR_CODE.notFound:
       return publicNotFound(error.message);
+    case API_ERROR_CODE.sessionRunClientRequestDuplicate:
+      return publicIdempotencyConflict(error.message);
     case API_ERROR_CODE.unauthorized:
       return publicUnauthenticated(error.message);
     default:

apps/api/src/modules/public-api/public-thread-events.ts

@@ -1,3 +1,4 @@
+import { OpenAiPrivateCitationStreamFilter } from "@mosoo/agent-driver/provider-output";
 import type {
   PublicThreadApiListThreadEventsResponse,
   PublicThreadEventLogEntry,
@@ -12,7 +13,6 @@ import type { SessionProcessEvent } from "@mosoo/contracts/session";
 import { sessionEventsTable, sessionMessagesTable } from "@mosoo/db";
 import { parsePlatformId } from "@mosoo/id";
 import type { RuntimeEventId, SessionId, SessionRunId } from "@mosoo/id";
-import { OpenAiPrivateCitationStreamFilter } from "agent-driver/provider-output";
 import { and, asc, desc, eq, gt, lt } from "drizzle-orm";
 import type { SQL } from "drizzle-orm";
 
@@ -58,15 +58,64 @@ interface LiveMessageState {
 
 class PublicLiveEventRowProjector {
   readonly #messages = new Map<string, LiveMessageState>();
+  readonly #messageKeysByRun = new Map<string, Set<string>>();
+
+  #getMessageRunKey(row: PublicThreadEventProcessRow): string {
+    return row.run_id ?? "";
+  }
+
+  #getMessageKey(row: PublicThreadEventProcessRow): string {
+    return `${this.#getMessageRunKey(row)}:${row.process_type}`;
+  }
+
+  #setMessage(row: PublicThreadEventProcessRow, message: LiveMessageState): void {
+    const key = this.#getMessageKey(row);
+    const runKey = this.#getMessageRunKey(row);
+    const runKeys = this.#messageKeysByRun.get(runKey);
+
+    this.#messages.set(key, message);
+    if (runKeys === undefined) {
+      this.#messageKeysByRun.set(runKey, new Set([key]));
+    } else {
+      runKeys.add(key);
+    }
+  }
+
+  #deleteMessage(row: PublicThreadEventProcessRow): void {
+    const key = this.#getMessageKey(row);
+    const runKey = this.#getMessageRunKey(row);
+    const runKeys = this.#messageKeysByRun.get(runKey);
+
+    this.#messages.delete(key);
+    runKeys?.delete(key);
+    if (runKeys?.size === 0) {
+      this.#messageKeysByRun.delete(runKey);
+    }
+  }
+
+  #deleteRunMessages(runId: SessionRunId | null): void {
+    const runKey = runId ?? "";
+    const messageKeys = this.#messageKeysByRun.get(runKey);
+
+    if (messageKeys === undefined) {
+      return;
+    }
+
+    for (const messageKey of messageKeys) {
+      this.#messages.delete(messageKey);
+    }
+
+    this.#messageKeysByRun.delete(runKey);
+  }
 
   project(rows: readonly PublicThreadEventProcessRow[]): PublicThreadEventProcessRow[] {
     const output: PublicThreadEventProcessRow[] = [];
 
     for (const row of rows) {
-      const key = `${row.run_id ?? ""}:${row.process_type}`;
+      const key = this.#getMessageKey(row);
 
       if (row.event_type === "message.started") {
-        this.#messages.set(key, {
+        this.#setMessage(row, {
           filter: new OpenAiPrivateCitationStreamFilter(),
           text: "",
         });
@@ -78,7 +127,7 @@ class PublicLiveEventRowProjector {
 
         if (message === undefined) {
           message = { filter: new OpenAiPrivateCitationStreamFilter(), text: "" };
-          this.#messages.set(key, message);
+          this.#setMessage(row, message);
         }
 
         const contentText = message.filter.push(row.content_text).text;
@@ -105,7 +154,7 @@ class PublicLiveEventRowProjector {
 
       if (row.event_type === "message.added") {
         const message = this.#messages.get(key);
-        this.#messages.delete(key);
+        this.#deleteMessage(row);
 
         if (message !== undefined) {
           const snapshotText = sanitizePublicOutput(row.content_text).text;
@@ -132,13 +181,7 @@ class PublicLiveEventRowProjector {
         row.event_type === "run.completed" ||
         row.event_type === "run.failed"
       ) {
-        const runPrefix = `${row.run_id ?? ""}:`;
-
-        for (const messageKey of this.#messages.keys()) {
-          if (messageKey.startsWith(runPrefix)) {
-            this.#messages.delete(messageKey);
-          }
-        }
+        this.#deleteRunMessages(row.run_id);
       }
 
       output.push(row);

apps/api/src/modules/runtime/application/agent-runtime-profile.ts

@@ -1,7 +1,7 @@
+import { getSessionOrganizationPath, getSessionRuntimeStatePath } from "@mosoo/agent-driver/paths";
 import type { JsonObject } from "@mosoo/contracts";
 import type { AgentKind, AgentReadiness } from "@mosoo/contracts/agent";
 import type { AccountId, AgentId, SandboxId, SandboxSessionId, SessionId } from "@mosoo/id";
-import { getSessionOrganizationPath, getSessionRuntimeStatePath } from "agent-driver/paths";
 
 import type {
   DriverConfigRevision,

apps/api/src/modules/runtime/application/execution-plane/driver-boot-payload-prepared.ts

@@ -1,4 +1,4 @@
-import type { DriverBootPayload } from "agent-driver/boot";
+import type { DriverBootPayload } from "@mosoo/agent-driver/boot";
 
 export interface DriverBootPayloadPreparedInput {
   readonly bootPayload: DriverBootPayload;

apps/api/src/modules/runtime/application/session-definition/hydrate-run-context.service.ts

@@ -1,3 +1,4 @@
+import { getSessionOrganizationPath } from "@mosoo/agent-driver/paths";
 import type { SessionSummary } from "@mosoo/contracts/session";
 import type { UserWarning } from "@mosoo/contracts/session-run";
 import type { ResolvedRunSkill } from "@mosoo/contracts/skill";
@@ -6,7 +7,6 @@ import type { AgentId, PlatformId, AppId, SandboxId, SandboxSessionId, SessionId
 import { getRuntimeCatalogEntry, getRuntimeCatalogVendorForProvider } from "@mosoo/runtime-catalog";
 import type { RuntimeCatalogVendor } from "@mosoo/runtime-catalog";
 import { RUNTIME_DIAGNOSTIC_EVENT } from "@mosoo/runtime-events";
-import { getSessionOrganizationPath } from "agent-driver/paths";
 
 import type { ApiBindings } from "../../../../platform/cloudflare/worker-types";
 import { validationError } from "../../../../platform/errors";

apps/api/src/modules/runtime/application/session-definition/session-config-trace-event.ts

@@ -2,7 +2,7 @@ import type {
   MosooSessionConfigTraceMcpServer,
   MosooSessionConfigTraceValue,
 } from "@mosoo/ag-ui-session";
-import type { DriverBootMcpServer, DriverBootPayload } from "agent-driver/boot";
+import type { DriverBootMcpServer, DriverBootPayload } from "@mosoo/agent-driver/boot";
 
 function summarizeMcpServer(server: DriverBootMcpServer): MosooSessionConfigTraceMcpServer {
   return {