refactor(agent): dedupe harness and command workflows

src/agents/pi-embedded-runner.run-embedded-pi-agent.auth-profile-rotation.e2e.test.ts

@@ -120,55 +120,154 @@ const writeAuthStore = async (
   await fs.writeFile(authPath, JSON.stringify(payload));
 };
 
+const mockFailedThenSuccessfulAttempt = (errorMessage = "rate limit") => {
+  runEmbeddedAttemptMock
+    .mockResolvedValueOnce(
+      makeAttempt({
+        assistantTexts: [],
+        lastAssistant: buildAssistant({
+          stopReason: "error",
+          errorMessage,
+        }),
+      }),
+    )
+    .mockResolvedValueOnce(
+      makeAttempt({
+        assistantTexts: ["ok"],
+        lastAssistant: buildAssistant({
+          stopReason: "stop",
+          content: [{ type: "text", text: "ok" }],
+        }),
+      }),
+    );
+};
+
+async function runAutoPinnedOpenAiTurn(params: {
+  agentDir: string;
+  workspaceDir: string;
+  sessionKey: string;
+  runId: string;
+  authProfileId?: string;
+}) {
+  await runEmbeddedPiAgent({
+    sessionId: "session:test",
+    sessionKey: params.sessionKey,
+    sessionFile: path.join(params.workspaceDir, "session.jsonl"),
+    workspaceDir: params.workspaceDir,
+    agentDir: params.agentDir,
+    config: makeConfig(),
+    prompt: "hello",
+    provider: "openai",
+    model: "mock-1",
+    authProfileId: params.authProfileId ?? "openai:p1",
+    authProfileIdSource: "auto",
+    timeoutMs: 5_000,
+    runId: params.runId,
+  });
+}
+
+async function readUsageStats(agentDir: string) {
+  const stored = JSON.parse(
+    await fs.readFile(path.join(agentDir, "auth-profiles.json"), "utf-8"),
+  ) as { usageStats?: Record<string, { lastUsed?: number; cooldownUntil?: number }> };
+  return stored.usageStats ?? {};
+}
+
+async function expectProfileP2UsageUpdated(agentDir: string) {
+  const usageStats = await readUsageStats(agentDir);
+  expect(typeof usageStats["openai:p2"]?.lastUsed).toBe("number");
+}
+
+async function expectProfileP2UsageUnchanged(agentDir: string) {
+  const usageStats = await readUsageStats(agentDir);
+  expect(usageStats["openai:p2"]?.lastUsed).toBe(2);
+}
+
+function mockSingleSuccessfulAttempt() {
+  runEmbeddedAttemptMock.mockResolvedValueOnce(
+    makeAttempt({
+      assistantTexts: ["ok"],
+      lastAssistant: buildAssistant({
+        stopReason: "stop",
+        content: [{ type: "text", text: "ok" }],
+      }),
+    }),
+  );
+}
+
+async function withTimedAgentWorkspace<T>(
+  run: (ctx: { agentDir: string; workspaceDir: string; now: number }) => Promise<T>,
+) {
+  vi.useFakeTimers();
+  try {
+    const agentDir = await fs.mkdtemp(path.join(os.tmpdir(), "openclaw-agent-"));
+    const workspaceDir = await fs.mkdtemp(path.join(os.tmpdir(), "openclaw-workspace-"));
+    const now = Date.now();
+    vi.setSystemTime(now);
+
+    try {
+      return await run({ agentDir, workspaceDir, now });
+    } finally {
+      await fs.rm(agentDir, { recursive: true, force: true });
+      await fs.rm(workspaceDir, { recursive: true, force: true });
+    }
+  } finally {
+    vi.useRealTimers();
+  }
+}
+
+async function runTurnWithCooldownSeed(params: {
+  sessionKey: string;
+  runId: string;
+  authProfileId: string | undefined;
+  authProfileIdSource: "auto" | "user";
+}) {
+  return await withTimedAgentWorkspace(async ({ agentDir, workspaceDir, now }) => {
+    await writeAuthStore(agentDir, {
+      usageStats: {
+        "openai:p1": { lastUsed: 1, cooldownUntil: now + 60 * 60 * 1000 },
+        "openai:p2": { lastUsed: 2 },
+      },
+    });
+    mockSingleSuccessfulAttempt();
+
+    await runEmbeddedPiAgent({
+      sessionId: "session:test",
+      sessionKey: params.sessionKey,
+      sessionFile: path.join(workspaceDir, "session.jsonl"),
+      workspaceDir,
+      agentDir,
+      config: makeConfig(),
+      prompt: "hello",
+      provider: "openai",
+      model: "mock-1",
+      authProfileId: params.authProfileId,
+      authProfileIdSource: params.authProfileIdSource,
+      timeoutMs: 5_000,
+      runId: params.runId,
+    });
+
+    expect(runEmbeddedAttemptMock).toHaveBeenCalledTimes(1);
+    return { usageStats: await readUsageStats(agentDir), now };
+  });
+}
+
 describe("runEmbeddedPiAgent auth profile rotation", () => {
   it("rotates for auto-pinned profiles", async () => {
     const agentDir = await fs.mkdtemp(path.join(os.tmpdir(), "openclaw-agent-"));
     const workspaceDir = await fs.mkdtemp(path.join(os.tmpdir(), "openclaw-workspace-"));
     try {
       await writeAuthStore(agentDir);
-
-      runEmbeddedAttemptMock
-        .mockResolvedValueOnce(
-          makeAttempt({
-            assistantTexts: [],
-            lastAssistant: buildAssistant({
-              stopReason: "error",
-              errorMessage: "rate limit",
-            }),
-          }),
-        )
-        .mockResolvedValueOnce(
-          makeAttempt({
-            assistantTexts: ["ok"],
-            lastAssistant: buildAssistant({
-              stopReason: "stop",
-              content: [{ type: "text", text: "ok" }],
-            }),
-          }),
-        );
-
-      await runEmbeddedPiAgent({
-        sessionId: "session:test",
-        sessionKey: "agent:test:auto",
-        sessionFile: path.join(workspaceDir, "session.jsonl"),
-        workspaceDir,
+      mockFailedThenSuccessfulAttempt("rate limit");
+      await runAutoPinnedOpenAiTurn({
         agentDir,
-        config: makeConfig(),
-        prompt: "hello",
-        provider: "openai",
-        model: "mock-1",
-        authProfileId: "openai:p1",
-        authProfileIdSource: "auto",
-        timeoutMs: 5_000,
+        workspaceDir,
+        sessionKey: "agent:test:auto",
         runId: "run:auto",
       });
 
       expect(runEmbeddedAttemptMock).toHaveBeenCalledTimes(2);
-
-      const stored = JSON.parse(
-        await fs.readFile(path.join(agentDir, "auth-profiles.json"), "utf-8"),
-      ) as { usageStats?: Record<string, { lastUsed?: number }> };
-      expect(typeof stored.usageStats?.["openai:p2"]?.lastUsed).toBe("number");
+      await expectProfileP2UsageUpdated(agentDir);
     } finally {
       await fs.rm(agentDir, { recursive: true, force: true });
       await fs.rm(workspaceDir, { recursive: true, force: true });
@@ -180,49 +279,16 @@ describe("runEmbeddedPiAgent auth profile rotation", () => {
     const workspaceDir = await fs.mkdtemp(path.join(os.tmpdir(), "openclaw-workspace-"));
     try {
       await writeAuthStore(agentDir);
-
-      runEmbeddedAttemptMock
-        .mockResolvedValueOnce(
-          makeAttempt({
-            assistantTexts: [],
-            lastAssistant: buildAssistant({
-              stopReason: "error",
-              errorMessage: "request ended without sending any chunks",
-            }),
-          }),
-        )
-        .mockResolvedValueOnce(
-          makeAttempt({
-            assistantTexts: ["ok"],
-            lastAssistant: buildAssistant({
-              stopReason: "stop",
-              content: [{ type: "text", text: "ok" }],
-            }),
-          }),
-        );
-
-      await runEmbeddedPiAgent({
-        sessionId: "session:test",
-        sessionKey: "agent:test:empty-chunk-stream",
-        sessionFile: path.join(workspaceDir, "session.jsonl"),
-        workspaceDir,
+      mockFailedThenSuccessfulAttempt("request ended without sending any chunks");
+      await runAutoPinnedOpenAiTurn({
         agentDir,
-        config: makeConfig(),
-        prompt: "hello",
-        provider: "openai",
-        model: "mock-1",
-        authProfileId: "openai:p1",
-        authProfileIdSource: "auto",
-        timeoutMs: 5_000,
+        workspaceDir,
+        sessionKey: "agent:test:empty-chunk-stream",
         runId: "run:empty-chunk-stream",
       });
 
       expect(runEmbeddedAttemptMock).toHaveBeenCalledTimes(2);
-
-      const stored = JSON.parse(
-        await fs.readFile(path.join(agentDir, "auth-profiles.json"), "utf-8"),
-      ) as { usageStats?: Record<string, { lastUsed?: number }> };
-      expect(typeof stored.usageStats?.["openai:p2"]?.lastUsed).toBe("number");
+      await expectProfileP2UsageUpdated(agentDir);
     } finally {
       await fs.rm(agentDir, { recursive: true, force: true });
       await fs.rm(workspaceDir, { recursive: true, force: true });
@@ -267,10 +333,7 @@ describe("runEmbeddedPiAgent auth profile rotation", () => {
       expect(runEmbeddedAttemptMock).toHaveBeenCalledTimes(1);
       expect(result.meta.aborted).toBe(true);
 
-      const stored = JSON.parse(
-        await fs.readFile(path.join(agentDir, "auth-profiles.json"), "utf-8"),
-      ) as { usageStats?: Record<string, { lastUsed?: number }> };
-      expect(stored.usageStats?.["openai:p2"]?.lastUsed).toBe(2);
+      await expectProfileP2UsageUnchanged(agentDir);
     } finally {
       await fs.rm(agentDir, { recursive: true, force: true });
       await fs.rm(workspaceDir, { recursive: true, force: true });
@@ -310,11 +373,7 @@ describe("runEmbeddedPiAgent auth profile rotation", () => {
       });
 
       expect(runEmbeddedAttemptMock).toHaveBeenCalledTimes(1);
-
-      const stored = JSON.parse(
-        await fs.readFile(path.join(agentDir, "auth-profiles.json"), "utf-8"),
-      ) as { usageStats?: Record<string, { lastUsed?: number }> };
-      expect(stored.usageStats?.["openai:p2"]?.lastUsed).toBe(2);
+      await expectProfileP2UsageUnchanged(agentDir);
     } finally {
       await fs.rm(agentDir, { recursive: true, force: true });
       await fs.rm(workspaceDir, { recursive: true, force: true });
@@ -322,71 +381,16 @@ describe("runEmbeddedPiAgent auth profile rotation", () => {
   });
 
   it("honors user-pinned profiles even when in cooldown", async () => {
-    vi.useFakeTimers();
-    try {
-      const agentDir = await fs.mkdtemp(path.join(os.tmpdir(), "openclaw-agent-"));
-      const workspaceDir = await fs.mkdtemp(path.join(os.tmpdir(), "openclaw-workspace-"));
-      const now = Date.now();
-      vi.setSystemTime(now);
+    const { usageStats } = await runTurnWithCooldownSeed({
+      sessionKey: "agent:test:user-cooldown",
+      runId: "run:user-cooldown",
+      authProfileId: "openai:p1",
+      authProfileIdSource: "user",
+    });
 
-      try {
-        const authPath = path.join(agentDir, "auth-profiles.json");
-        const payload = {
-          version: 1,
-          profiles: {
-            "openai:p1": { type: "api_key", provider: "openai", key: "sk-one" },
-            "openai:p2": { type: "api_key", provider: "openai", key: "sk-two" },
-          },
-          usageStats: {
-            "openai:p1": { lastUsed: 1, cooldownUntil: now + 60 * 60 * 1000 },
-            "openai:p2": { lastUsed: 2 },
-          },
-        };
-        await fs.writeFile(authPath, JSON.stringify(payload));
-
-        runEmbeddedAttemptMock.mockResolvedValueOnce(
-          makeAttempt({
-            assistantTexts: ["ok"],
-            lastAssistant: buildAssistant({
-              stopReason: "stop",
-              content: [{ type: "text", text: "ok" }],
-            }),
-          }),
-        );
-
-        await runEmbeddedPiAgent({
-          sessionId: "session:test",
-          sessionKey: "agent:test:user-cooldown",
-          sessionFile: path.join(workspaceDir, "session.jsonl"),
-          workspaceDir,
-          agentDir,
-          config: makeConfig(),
-          prompt: "hello",
-          provider: "openai",
-          model: "mock-1",
-          authProfileId: "openai:p1",
-          authProfileIdSource: "user",
-          timeoutMs: 5_000,
-          runId: "run:user-cooldown",
-        });
-
-        expect(runEmbeddedAttemptMock).toHaveBeenCalledTimes(1);
-
-        const stored = JSON.parse(
-          await fs.readFile(path.join(agentDir, "auth-profiles.json"), "utf-8"),
-        ) as {
-          usageStats?: Record<string, { lastUsed?: number; cooldownUntil?: number }>;
-        };
-        expect(stored.usageStats?.["openai:p1"]?.cooldownUntil).toBeUndefined();
-        expect(stored.usageStats?.["openai:p1"]?.lastUsed).not.toBe(1);
-        expect(stored.usageStats?.["openai:p2"]?.lastUsed).toBe(2);
-      } finally {
-        await fs.rm(agentDir, { recursive: true, force: true });
-        await fs.rm(workspaceDir, { recursive: true, force: true });
-      }
-    } finally {
-      vi.useRealTimers();
-    }
+    expect(usageStats["openai:p1"]?.cooldownUntil).toBeUndefined();
+    expect(usageStats["openai:p1"]?.lastUsed).not.toBe(1);
+    expect(usageStats["openai:p2"]?.lastUsed).toBe(2);
   });
 
   it("ignores user-locked profile when provider mismatches", async () => {
@@ -429,116 +433,50 @@ describe("runEmbeddedPiAgent auth profile rotation", () => {
   });
 
   it("skips profiles in cooldown during initial selection", async () => {
-    vi.useFakeTimers();
-    try {
-      const agentDir = await fs.mkdtemp(path.join(os.tmpdir(), "openclaw-agent-"));
-      const workspaceDir = await fs.mkdtemp(path.join(os.tmpdir(), "openclaw-workspace-"));
-      const now = Date.now();
-      vi.setSystemTime(now);
+    const { usageStats, now } = await runTurnWithCooldownSeed({
+      sessionKey: "agent:test:skip-cooldown",
+      runId: "run:skip-cooldown",
+      authProfileId: undefined,
+      authProfileIdSource: "auto",
+    });
 
-      try {
-        const authPath = path.join(agentDir, "auth-profiles.json");
-        const payload = {
-          version: 1,
-          profiles: {
-            "openai:p1": { type: "api_key", provider: "openai", key: "sk-one" },
-            "openai:p2": { type: "api_key", provider: "openai", key: "sk-two" },
-          },
-          usageStats: {
-            "openai:p1": { lastUsed: 1, cooldownUntil: now + 60 * 60 * 1000 }, // p1 in cooldown for 1 hour
-            "openai:p2": { lastUsed: 2 },
-          },
-        };
-        await fs.writeFile(authPath, JSON.stringify(payload));
-
-        runEmbeddedAttemptMock.mockResolvedValueOnce(
-          makeAttempt({
-            assistantTexts: ["ok"],
-            lastAssistant: buildAssistant({
-              stopReason: "stop",
-              content: [{ type: "text", text: "ok" }],
-            }),
-          }),
-        );
-
-        await runEmbeddedPiAgent({
-          sessionId: "session:test",
-          sessionKey: "agent:test:skip-cooldown",
-          sessionFile: path.join(workspaceDir, "session.jsonl"),
-          workspaceDir,
-          agentDir,
-          config: makeConfig(),
-          prompt: "hello",
-          provider: "openai",
-          model: "mock-1",
-          authProfileId: undefined,
-          authProfileIdSource: "auto",
-          timeoutMs: 5_000,
-          runId: "run:skip-cooldown",
-        });
-
-        expect(runEmbeddedAttemptMock).toHaveBeenCalledTimes(1);
-
-        const stored = JSON.parse(
-          await fs.readFile(path.join(agentDir, "auth-profiles.json"), "utf-8"),
-        ) as { usageStats?: Record<string, { lastUsed?: number; cooldownUntil?: number }> };
-        expect(stored.usageStats?.["openai:p1"]?.cooldownUntil).toBe(now + 60 * 60 * 1000);
-        expect(typeof stored.usageStats?.["openai:p2"]?.lastUsed).toBe("number");
-      } finally {
-        await fs.rm(agentDir, { recursive: true, force: true });
-        await fs.rm(workspaceDir, { recursive: true, force: true });
-      }
-    } finally {
-      vi.useRealTimers();
-    }
+    expect(usageStats["openai:p1"]?.cooldownUntil).toBe(now + 60 * 60 * 1000);
+    expect(typeof usageStats["openai:p2"]?.lastUsed).toBe("number");
   });
 
   it("fails over when all profiles are in cooldown and fallbacks are configured", async () => {
-    vi.useFakeTimers();
-    try {
-      const agentDir = await fs.mkdtemp(path.join(os.tmpdir(), "openclaw-agent-"));
-      const workspaceDir = await fs.mkdtemp(path.join(os.tmpdir(), "openclaw-workspace-"));
-      const now = Date.now();
-      vi.setSystemTime(now);
+    await withTimedAgentWorkspace(async ({ agentDir, workspaceDir, now }) => {
+      await writeAuthStore(agentDir, {
+        usageStats: {
+          "openai:p1": { lastUsed: 1, cooldownUntil: now + 60 * 60 * 1000 },
+          "openai:p2": { lastUsed: 2, cooldownUntil: now + 60 * 60 * 1000 },
+        },
+      });
 
-      try {
-        await writeAuthStore(agentDir, {
-          usageStats: {
-            "openai:p1": { lastUsed: 1, cooldownUntil: now + 60 * 60 * 1000 },
-            "openai:p2": { lastUsed: 2, cooldownUntil: now + 60 * 60 * 1000 },
-          },
-        });
-
-        await expect(
-          runEmbeddedPiAgent({
-            sessionId: "session:test",
-            sessionKey: "agent:test:cooldown-failover",
-            sessionFile: path.join(workspaceDir, "session.jsonl"),
-            workspaceDir,
-            agentDir,
-            config: makeConfig({ fallbacks: ["openai/mock-2"] }),
-            prompt: "hello",
-            provider: "openai",
-            model: "mock-1",
-            authProfileIdSource: "auto",
-            timeoutMs: 5_000,
-            runId: "run:cooldown-failover",
-          }),
-        ).rejects.toMatchObject({
-          name: "FailoverError",
-          reason: "rate_limit",
+      await expect(
+        runEmbeddedPiAgent({
+          sessionId: "session:test",
+          sessionKey: "agent:test:cooldown-failover",
+          sessionFile: path.join(workspaceDir, "session.jsonl"),
+          workspaceDir,
+          agentDir,
+          config: makeConfig({ fallbacks: ["openai/mock-2"] }),
+          prompt: "hello",
           provider: "openai",
           model: "mock-1",
-        });
+          authProfileIdSource: "auto",
+          timeoutMs: 5_000,
+          runId: "run:cooldown-failover",
+        }),
+      ).rejects.toMatchObject({
+        name: "FailoverError",
+        reason: "rate_limit",
+        provider: "openai",
+        model: "mock-1",
+      });
 
-        expect(runEmbeddedAttemptMock).not.toHaveBeenCalled();
-      } finally {
-        await fs.rm(agentDir, { recursive: true, force: true });
-        await fs.rm(workspaceDir, { recursive: true, force: true });
-      }
-    } finally {
-      vi.useRealTimers();
-    }
+      expect(runEmbeddedAttemptMock).not.toHaveBeenCalled();
+    });
   });
 
   it("fails over when auth is unavailable and fallbacks are configured", async () => {
@@ -604,52 +542,19 @@ describe("runEmbeddedPiAgent auth profile rotation", () => {
         };
         await fs.writeFile(authPath, JSON.stringify(payload));
 
-        runEmbeddedAttemptMock
-          .mockResolvedValueOnce(
-            makeAttempt({
-              assistantTexts: [],
-              lastAssistant: buildAssistant({
-                stopReason: "error",
-                errorMessage: "rate limit",
-              }),
-            }),
-          )
-          .mockResolvedValueOnce(
-            makeAttempt({
-              assistantTexts: ["ok"],
-              lastAssistant: buildAssistant({
-                stopReason: "stop",
-                content: [{ type: "text", text: "ok" }],
-              }),
-            }),
-          );
-
-        await runEmbeddedPiAgent({
-          sessionId: "session:test",
-          sessionKey: "agent:test:rotate-skip-cooldown",
-          sessionFile: path.join(workspaceDir, "session.jsonl"),
-          workspaceDir,
+        mockFailedThenSuccessfulAttempt("rate limit");
+        await runAutoPinnedOpenAiTurn({
           agentDir,
-          config: makeConfig(),
-          prompt: "hello",
-          provider: "openai",
-          model: "mock-1",
-          authProfileId: "openai:p1",
-          authProfileIdSource: "auto",
-          timeoutMs: 5_000,
+          workspaceDir,
+          sessionKey: "agent:test:rotate-skip-cooldown",
           runId: "run:rotate-skip-cooldown",
         });
 
         expect(runEmbeddedAttemptMock).toHaveBeenCalledTimes(2);
-
-        const stored = JSON.parse(
-          await fs.readFile(path.join(agentDir, "auth-profiles.json"), "utf-8"),
-        ) as {
-          usageStats?: Record<string, { lastUsed?: number; cooldownUntil?: number }>;
-        };
-        expect(typeof stored.usageStats?.["openai:p1"]?.lastUsed).toBe("number");
-        expect(typeof stored.usageStats?.["openai:p3"]?.lastUsed).toBe("number");
-        expect(stored.usageStats?.["openai:p2"]?.cooldownUntil).toBe(now + 60 * 60 * 1000);
+        const usageStats = await readUsageStats(agentDir);
+        expect(typeof usageStats["openai:p1"]?.lastUsed).toBe("number");
+        expect(typeof usageStats["openai:p3"]?.lastUsed).toBe("number");
+        expect(usageStats["openai:p2"]?.cooldownUntil).toBe(now + 60 * 60 * 1000);
       } finally {
         await fs.rm(agentDir, { recursive: true, force: true });
         await fs.rm(workspaceDir, { recursive: true, force: true });