[Admin Governance] Settle repair side effects

Author: philipperolet

front/migrations/20260922_restore_editorless_agent_authors.test.ts

@@ -3,11 +3,13 @@ import { GroupPermissionResource } from "@app/lib/resources/group_permission_res
 import { withTransaction } from "@app/lib/utils/sql_utils";
 import baseLogger from "@app/logger/logger";
 import { restoreEditorlessAgentAuthors } from "@app/migrations/20260922_restore_editorless_agent_authors";
+import * as searchIndexationClient from "@app/temporal/es_indexation/client";
 import { AgentConfigurationFactory } from "@app/tests/utils/AgentConfigurationFactory";
 import { createResourceTest } from "@app/tests/utils/generic_resource_tests";
 import { MembershipFactory } from "@app/tests/utils/MembershipFactory";
 import { UserFactory } from "@app/tests/utils/UserFactory";
 import type { AgentConfigurationType } from "@app/types/assistant/agent";
+import { Err, Ok } from "@app/types/shared/result";
 import type { UserType } from "@app/types/user";
 import assert from "assert";
 import { afterEach, describe, expect, it, vi } from "vitest";
@@ -50,8 +52,8 @@ async function editorIds(
 describe("restoreEditorlessAgentAuthors", () => {
   it("restores only editorless agents, with a dry run and idempotent rerun", async () => {
     const launchSearchIndexation = vi
-      .spyOn(AgentResource, "launchSearchIndexation")
-      .mockResolvedValue();
+      .spyOn(searchIndexationClient, "launchIndexAgentSearchWorkflow")
+      .mockResolvedValue(new Ok(undefined));
     const {
       authenticator: auth,
       user,
@@ -117,10 +119,10 @@ describe("restoreEditorlessAgentAuthors", () => {
       otherEditor.sId,
     ]);
     expect(await editorIds(auth, resourceWithGlobalEditor)).toEqual([]);
-    expect(launchSearchIndexation).toHaveBeenCalledExactlyOnceWith(
-      expect.anything(),
-      [editorlessAgent.sId]
-    );
+    expect(launchSearchIndexation).toHaveBeenCalledExactlyOnceWith({
+      workspaceId: workspace.sId,
+      agentId: editorlessAgent.sId,
+    });
 
     await expect(
       restoreEditorlessAgentAuthors({
@@ -135,6 +137,95 @@ describe("restoreEditorlessAgentAuthors", () => {
     });
   });
 
+  it("retries failed search workflow launches", async () => {
+    const launchSearchIndexation = vi
+      .spyOn(searchIndexationClient, "launchIndexAgentSearchWorkflow")
+      .mockResolvedValue(new Ok(undefined));
+    const { authenticator: auth, workspace } = await createResourceTest({
+      role: "admin",
+    });
+    const agent = await AgentConfigurationFactory.createTestAgent(auth, {
+      name: "Editorless agent",
+    });
+    await replaceEditors(auth, agent, []);
+    launchSearchIndexation.mockReset();
+    launchSearchIndexation
+      .mockResolvedValueOnce(new Err(new Error("Temporal unavailable")))
+      .mockResolvedValueOnce(new Err(new Error("Temporal unavailable")))
+      .mockResolvedValueOnce(new Ok(undefined));
+
+    await restoreEditorlessAgentAuthors({
+      execute: true,
+      logger,
+      wId: workspace.sId,
+    });
+
+    expect(launchSearchIndexation).toHaveBeenCalledTimes(3);
+  });
+
+  it("waits for every in-flight repair before propagating a failure", async () => {
+    vi.spyOn(
+      searchIndexationClient,
+      "launchIndexAgentSearchWorkflow"
+    ).mockResolvedValue(new Ok(undefined));
+    const {
+      authenticator: auth,
+      user,
+      workspace,
+    } = await createResourceTest({ role: "admin" });
+    const firstAgent = await AgentConfigurationFactory.createTestAgent(auth, {
+      name: "First editorless agent",
+    });
+    const secondAgent = await AgentConfigurationFactory.createTestAgent(auth, {
+      name: "Second editorless agent",
+    });
+    const firstResource = await replaceEditors(auth, firstAgent, []);
+    const secondResource = await replaceEditors(auth, secondAgent, []);
+
+    const originalGrantToUser = GroupPermissionResource.grantToUser;
+    let releaseSecondGrant: () => void = () => undefined;
+    const secondGrantGate = new Promise<void>((resolve) => {
+      releaseSecondGrant = resolve;
+    });
+    let markSecondGrantStarted: () => void = () => undefined;
+    const secondGrantStarted = new Promise<void>((resolve) => {
+      markSecondGrantStarted = resolve;
+    });
+    vi.spyOn(GroupPermissionResource, "grantToUser").mockImplementation(
+      async (grantAuth, grant) => {
+        if (grant.resourceId === firstResource.id) {
+          throw new Error("Failed first repair");
+        }
+        if (grant.resourceId === secondResource.id) {
+          markSecondGrantStarted();
+          await secondGrantGate;
+        }
+        return originalGrantToUser.call(
+          GroupPermissionResource,
+          grantAuth,
+          grant
+        );
+      }
+    );
+
+    let settled = false;
+    const run = restoreEditorlessAgentAuthors({
+      execute: true,
+      logger,
+      wId: workspace.sId,
+    }).finally(() => {
+      settled = true;
+    });
+    await secondGrantStarted;
+    await new Promise((resolve) => setTimeout(resolve, 0));
+    expect(settled).toBe(false);
+
+    releaseSecondGrant();
+    await expect(run).rejects.toThrow("Failed first repair");
+    expect(await editorIds(auth, firstResource)).toEqual([]);
+    expect(await editorIds(auth, secondResource)).toEqual([user.sId]);
+  });
+
   it("rejects an unknown workspace scope", async () => {
     await expect(
       restoreEditorlessAgentAuthors({

front/migrations/20260922_restore_editorless_agent_authors.ts

@@ -3,7 +3,6 @@ import {
   emitAuditLogEventDirect,
 } from "@app/lib/api/audit/workos_audit";
 import { Authenticator } from "@app/lib/auth";
-import { AgentResource } from "@app/lib/resources/agent_resource";
 import { GroupPermissionResource } from "@app/lib/resources/group_permission_resource";
 import { frontSequelize } from "@app/lib/resources/storage";
 import { UserResource } from "@app/lib/resources/user_resource";
@@ -12,8 +11,11 @@ import { concurrentExecutor, withRetry } from "@app/lib/utils/async_utils";
 import { withTransaction } from "@app/lib/utils/sql_utils";
 import type { Logger } from "@app/logger/logger";
 import { makeScript } from "@app/scripts/helpers";
+import { launchIndexAgentSearchWorkflow } from "@app/temporal/es_indexation/client";
 import type { AgentConfigurationScope } from "@app/types/assistant/agent";
 import type { ModelId } from "@app/types/shared/model_id";
+import { Err, Ok } from "@app/types/shared/result";
+import { normalizeError } from "@app/types/shared/utils/error_utils";
 import assert from "assert";
 import type { Transaction } from "sequelize";
 import { QueryTypes } from "sequelize";
@@ -220,11 +222,17 @@ async function restoreAuthor(
       },
     });
 
-    // Keep indexing outside the write transaction so a Temporal connection failure cannot skip the
-    // audit event after the repair has committed. Retry connection failures, then log and continue.
-    const indexation = await withRetry(() =>
-      AgentResource.launchSearchIndexation(auth, [restored.agentId])
-    );
+    // Keep indexing outside the write transaction so a Temporal failure cannot skip the audit event
+    // after the repair has committed. Retry launch failures, then log and continue.
+    const indexation = await withRetry(async () => {
+      const result = await launchIndexAgentSearchWorkflow({
+        workspaceId: restored.workspaceId,
+        agentId: restored.agentId,
+      });
+      if (result.isErr()) {
+        throw result.error;
+      }
+    });
     if (indexation.isErr()) {
       logger.error(
         {
@@ -329,14 +337,35 @@ export async function restoreEditorlessAgentAuthors({
           const auth = authsByWorkspaceId.get(agent.workspaceId);
           assert(auth);
           // Each agent needs its own transaction and advisory lock to serialize with live editor
-          // changes. Concurrency is capped at four to bound database pressure.
-          return restoreAuthor(auth, agent, logger);
+          // changes. Catch failures here so every in-flight repair finishes its post-commit audit
+          // attempt before the batch propagates an error to makeScript.
+          try {
+            return new Ok(await restoreAuthor(auth, agent, logger));
+          } catch (error) {
+            return new Err(normalizeError(error));
+          }
         },
         { concurrency: CONCURRENCY }
       );
 
-      for (const [index, restored] of results.entries()) {
+      let firstError: Error | undefined;
+      for (const [index, result] of results.entries()) {
         const discovered = agents[index];
+        if (result.isErr()) {
+          firstError ??= result.error;
+          logger.error(
+            {
+              error: result.error,
+              workspaceId: discovered.workspaceId,
+              agentId: discovered.agentId,
+              agentModelId: discovered.agentModelId,
+            },
+            "Failed to restore author as agent editor"
+          );
+          continue;
+        }
+
+        const restored = result.value;
         if (!restored) {
           stats.agentsSkipped += 1;
           logger.info(
@@ -361,6 +390,9 @@ export async function restoreEditorlessAgentAuthors({
           "Restored author as agent editor"
         );
       }
+      if (firstError) {
+        throw firstError;
+      }
     }
 
     cursor = agents[agents.length - 1].agentModelId;