stagent 0.9.5 → 0.10.0

package/src/lib/schedules/config.ts

@@ -0,0 +1,53 @@
+import { getSettingSync } from "@/lib/settings/helpers";
+import { SETTINGS_KEYS } from "@/lib/constants/settings";
+
+const DEFAULT_MAX_CONCURRENT = 2;
+const DEFAULT_MAX_RUN_DURATION_SEC = 1200; // 20 minutes
+const DEFAULT_CHAT_PRESSURE_DELAY_SEC = 30;
+
+function readIntConfig(
+  envVar: string,
+  settingKey: string,
+  defaultValue: number,
+): number {
+  const envRaw = process.env[envVar];
+  if (envRaw !== undefined) {
+    const parsed = parseInt(envRaw, 10);
+    if (Number.isFinite(parsed) && parsed > 0) return parsed;
+    console.warn(
+      `[schedule-config] ${envVar}="${envRaw}" is not a positive integer; using default ${defaultValue}`,
+    );
+  }
+
+  const settingRaw = getSettingSync(settingKey);
+  if (settingRaw !== null) {
+    const parsed = parseInt(settingRaw, 10);
+    if (Number.isFinite(parsed) && parsed > 0) return parsed;
+  }
+
+  return defaultValue;
+}
+
+export function getScheduleMaxConcurrent(): number {
+  return readIntConfig(
+    "SCHEDULE_MAX_CONCURRENT",
+    SETTINGS_KEYS.SCHEDULE_MAX_CONCURRENT,
+    DEFAULT_MAX_CONCURRENT,
+  );
+}
+
+export function getScheduleMaxRunDurationSec(): number {
+  return readIntConfig(
+    "SCHEDULE_MAX_RUN_DURATION_SEC",
+    SETTINGS_KEYS.SCHEDULE_MAX_RUN_DURATION_SEC,
+    DEFAULT_MAX_RUN_DURATION_SEC,
+  );
+}
+
+export function getScheduleChatPressureDelaySec(): number {
+  return readIntConfig(
+    "SCHEDULE_CHAT_PRESSURE_DELAY_SEC",
+    SETTINGS_KEYS.SCHEDULE_CHAT_PRESSURE_DELAY_SEC,
+    DEFAULT_CHAT_PRESSURE_DELAY_SEC,
+  );
+}

package/src/lib/schedules/scheduler.ts

@@ -12,8 +12,9 @@
  */
 
 import { db } from "@/lib/db";
-import { schedules, tasks, agentLogs, scheduleDocumentInputs, documents } from "@/lib/db/schema";
-import { eq, and, lte, inArray, sql, asc } from "drizzle-orm";
+import { schedules, tasks, agentLogs, scheduleDocumentInputs, documents, workflows, scheduleFiringMetrics } from "@/lib/db/schema";
+import { eq, and, lte, inArray, sql, asc, isNotNull } from "drizzle-orm";
+import { resumeWorkflow } from "@/lib/workflows/engine";
 import { computeNextFireTime } from "./interval-parser";
 import { executeTaskWithRuntime } from "@/lib/agents/runtime";
 import { getSetting } from "@/lib/settings/helpers";
@@ -27,6 +28,13 @@ import {
 import { sendToChannels } from "@/lib/channels/registry";
 import type { ChannelMessage } from "@/lib/channels/types";
 import { processHandoffs } from "@/lib/agents/handoff/bus";
+import { claimSlot, reapExpiredLeases, countRunningScheduledSlots } from "./slot-claim";
+import { isAnyChatStreaming } from "@/lib/chat/active-streams";
+import {
+  getScheduleMaxConcurrent,
+  getScheduleMaxRunDurationSec,
+  getScheduleChatPressureDelaySec,
+} from "./config";
 
 const POLL_INTERVAL_MS = 60_000; // 60 seconds
 
@@ -54,6 +62,10 @@ export async function drainQueue(): Promise<void> {
     // Loop until the queue is empty so a single drain cycle clears all
     // collided tasks rather than only the next one.
     while (true) {
+      // Respect the global cap — stop draining if we're already at capacity
+      const cap = getScheduleMaxConcurrent();
+      if (countRunningScheduledSlots() >= cap) return;
+
       const [nextQueued] = await db
         .select({ id: tasks.id })
         .from(tasks)
@@ -68,7 +80,17 @@ export async function drainQueue(): Promise<void> {
 
       if (!nextQueued) return;
 
-      console.log(`[scheduler] draining queue → executing task ${nextQueued.id}`);
+      // Atomic claim — could lose the race if a concurrent tick already took
+      // this specific task, OR the cap filled between the select and the claim.
+      // On a lost-race (task-level) we should try the next queued task; on a
+      // cap-full the next iteration's cap check at the top of the loop will
+      // return. Continue rather than return so we don't strand other queued
+      // tasks that could still claim.
+      const leaseSec = getScheduleMaxRunDurationSec();
+      const { claimed } = claimSlot(nextQueued.id, cap, leaseSec);
+      if (!claimed) continue;
+
+      console.log(`[scheduler] draining queue → running task ${nextQueued.id}`);
       try {
         await executeTaskWithRuntime(nextQueued.id);
       } catch (err) {
@@ -131,17 +153,31 @@ function detectFailureReason(result: string | null): string {
   return "error";
 }
 
+const TURN_BUDGET_BREACH_AUTO_PAUSE_THRESHOLD = 5;
+const GRACE_PERIOD_MULTIPLIER = 2; // grace window = 2 × cron interval
+
 /**
  * Record per-firing health metrics on a schedule and auto-pause after
- * 3 consecutive failures. Uses an exponential moving average for turn count
- * so the metric reflects recent behavior more than ancient firings.
+ * 3 consecutive generic failures or 5 consecutive turn-budget breaches.
+ * Uses an exponential moving average for turn count so the metric reflects
+ * recent behavior more than ancient firings.
+ *
+ * Turn-budget breaches are tracked separately (turnBudgetBreachStreak) so a
+ * misconfigured maxTurns doesn't auto-pause via the generic threshold of 3.
+ * A first-breach grace window (2× cron interval after maxTurnsSetAt) forgives
+ * the first firing that hits a newly-lowered cap.
  */
 export async function recordFiringMetrics(
   scheduleId: string,
-  taskId: string
+  taskId: string,
 ): Promise<void> {
   const [task] = await db
-    .select({ status: tasks.status, result: tasks.result })
+    .select({
+      status: tasks.status,
+      result: tasks.result,
+      failureReason: tasks.failureReason,
+      updatedAt: tasks.updatedAt,
+    })
     .from(tasks)
     .where(eq(tasks.id, taskId));
   if (!task) return;
@@ -162,26 +198,124 @@ export async function recordFiringMetrics(
   const newAvg = Math.round(prevAvg * 0.7 + turns * 0.3);
 
   const isFailure = task.status === "failed";
-  const newStreak = isFailure ? (schedule.failureStreak ?? 0) + 1 : 0;
-  const shouldAutoPause = isFailure && newStreak >= 3 && schedule.status === "active";
+  const failureReason =
+    task.failureReason ?? (isFailure ? detectFailureReason(task.result) : null);
+  const isTurnBudgetBreach = failureReason === "turn_limit_exceeded";
+  const isGenericFailure = isFailure && !isTurnBudgetBreach;
+
+  // First-breach grace: if this is the first firing after maxTurns was edited,
+  // don't count the breach toward the auto-pause streak.
+  let turnBudgetStreakDelta = 0;
+  if (isTurnBudgetBreach) {
+    const graceApplies = shouldApplyGrace(
+      schedule.maxTurnsSetAt,
+      schedule.cronExpression,
+      task.updatedAt,
+    );
+    if (!graceApplies) turnBudgetStreakDelta = 1;
+  }
+
+  const newFailureStreak = isGenericFailure ? (schedule.failureStreak ?? 0) + 1 : 0;
+  const newBudgetStreak =
+    turnBudgetStreakDelta > 0
+      ? (schedule.turnBudgetBreachStreak ?? 0) + 1
+      : isTurnBudgetBreach
+      ? (schedule.turnBudgetBreachStreak ?? 0)  // hold-steady but coerce null→0
+      : 0;
+  const shouldAutoPauseGeneric =
+    isGenericFailure && newFailureStreak >= 3 && schedule.status === "active";
+  const shouldAutoPauseBudget =
+    newBudgetStreak >= TURN_BUDGET_BREACH_AUTO_PAUSE_THRESHOLD &&
+    schedule.status === "active";
+  const shouldAutoPause = shouldAutoPauseGeneric || shouldAutoPauseBudget;
 
   await db
     .update(schedules)
     .set({
       lastTurnCount: turns,
       avgTurnsPerFiring: newAvg,
-      failureStreak: newStreak,
-      lastFailureReason: isFailure ? detectFailureReason(task.result) : null,
+      failureStreak: newFailureStreak,
+      turnBudgetBreachStreak: newBudgetStreak,
+      lastFailureReason: failureReason,
       status: shouldAutoPause ? "paused" : schedule.status,
       updatedAt: new Date(),
     })
     .where(eq(schedules.id, scheduleId));
 
-  if (shouldAutoPause) {
+  if (shouldAutoPauseGeneric) {
+    console.warn(
+      `[scheduler] auto-paused "${schedule.name}" after 3 consecutive failures`,
+    );
+  }
+  if (shouldAutoPauseBudget) {
     console.warn(
-      `[scheduler] auto-paused "${schedule.name}" after 3 consecutive failures`
+      `[scheduler] auto-paused "${schedule.name}" after 5 consecutive turn-budget breaches (avg: ${newAvg} steps, cap: ${schedule.maxTurns})`,
     );
   }
+
+  try {
+    const [taskRow] = await db
+      .select()
+      .from(tasks)
+      .where(eq(tasks.id, taskId));
+    if (taskRow) {
+      const firedAtDate = taskRow.createdAt;
+      const slotClaimedAt = taskRow.slotClaimedAt;
+      const completedAt = taskRow.updatedAt;
+      const slotWaitMs =
+        slotClaimedAt && firedAtDate
+          ? slotClaimedAt.getTime() - firedAtDate.getTime()
+          : null;
+      const durationMs =
+        slotClaimedAt && completedAt
+          ? completedAt.getTime() - slotClaimedAt.getTime()
+          : null;
+
+      await db.insert(scheduleFiringMetrics).values({
+        id: crypto.randomUUID(),
+        scheduleId,
+        taskId,
+        firedAt: firedAtDate,
+        slotClaimedAt,
+        completedAt,
+        slotWaitMs,
+        durationMs,
+        turnCount: turns,
+        maxTurnsAtFiring: schedule.maxTurns,
+        eventLoopLagMs: null,
+        peakRssMb: null,
+        chatStreamsActive: null,
+        concurrentSchedules: null,
+        failureReason,
+      });
+    }
+  } catch (err) {
+    console.error(`[scheduler] failed to insert firing metrics for ${taskId}:`, err);
+  }
+}
+
+/**
+ * First-breach grace: if maxTurnsSetAt was recent enough that this is the
+ * first-or-second firing after the edit, don't count the breach toward the
+ * auto-pause streak.
+ */
+function shouldApplyGrace(
+  maxTurnsSetAt: Date | null,
+  cronExpression: string,
+  completedAt: Date | null,
+): boolean {
+  if (!maxTurnsSetAt || !completedAt) return false;
+  try {
+    const t1 = computeNextFireTime(cronExpression, maxTurnsSetAt);
+    const t2 = computeNextFireTime(cronExpression, t1);
+    const cronIntervalMs = t2.getTime() - t1.getTime();
+    const graceWindowEnd = new Date(
+      maxTurnsSetAt.getTime() + GRACE_PERIOD_MULTIPLIER * cronIntervalMs,
+    );
+    return completedAt <= graceWindowEnd;
+  } catch {
+    return false;
+  }
 }
 
 /**
@@ -220,6 +354,18 @@ export function stopScheduler(): void {
 export async function tickScheduler(): Promise<void> {
   const now = new Date();
 
+  // Reap any running tasks whose lease has expired before claiming new slots.
+  try {
+    const reaped = reapExpiredLeases();
+    if (reaped.length > 0) {
+      console.warn(
+        `[scheduler] reaped ${reaped.length} expired lease(s): ${reaped.join(", ")}`,
+      );
+    }
+  } catch (err) {
+    console.error("[scheduler] lease reaper error:", err);
+  }
+
   const dueSchedules = await db
     .select()
     .from(schedules)
@@ -230,6 +376,34 @@ export async function tickScheduler(): Promise<void> {
       )
     );
 
+  // Chat soft pressure: defer new firings by N seconds when any chat stream
+  // is in flight. In-flight scheduled runs are NOT affected — this only gates
+  // new claims. Per-iteration try/catch so a single failed deferral doesn't
+  // silently skip the remaining schedules in this tick.
+  if (isAnyChatStreaming() && dueSchedules.length > 0) {
+    const delayMs = getScheduleChatPressureDelaySec() * 1000;
+    const deferredUntil = new Date(now.getTime() + delayMs);
+    let deferredCount = 0;
+    for (const schedule of dueSchedules) {
+      try {
+        await db
+          .update(schedules)
+          .set({ nextFireAt: deferredUntil, updatedAt: now })
+          .where(eq(schedules.id, schedule.id));
+        deferredCount++;
+      } catch (err) {
+        console.error(
+          `[scheduler] failed to defer schedule ${schedule.id} under chat pressure:`,
+          err,
+        );
+      }
+    }
+    console.warn(
+      `[scheduler] chat streaming — deferred ${deferredCount}/${dueSchedules.length} firings by ${delayMs}ms`,
+    );
+    return;
+  }
+
   for (const schedule of dueSchedules) {
     try {
       // Atomic claim: attempt to update nextFireAt to null as a lock.
@@ -268,6 +442,32 @@ export async function tickScheduler(): Promise<void> {
   } catch (err) {
     console.error("[scheduler] handoff processing error:", err);
   }
+
+  // Resume delayed workflows whose resume_at has passed. Uses the partial index
+  // idx_workflows_resume_at (WHERE resume_at IS NOT NULL) for efficiency.
+  // resumeWorkflow is idempotent via atomic status transition, so even if the
+  // scheduler tick races a user's "Resume Now" click, exactly one resume wins.
+  try {
+    const nowMs = now.getTime();
+    const dueDelayedWorkflows = await db
+      .select({ id: workflows.id })
+      .from(workflows)
+      .where(
+        and(
+          eq(workflows.status, "paused"),
+          isNotNull(workflows.resumeAt),
+          lte(workflows.resumeAt, nowMs),
+        ),
+      );
+
+    for (const wf of dueDelayedWorkflows) {
+      resumeWorkflow(wf.id).catch((err) => {
+        console.error(`[scheduler] failed to resume workflow ${wf.id}:`, err);
+      });
+    }
+  } catch (err) {
+    console.error("[scheduler] delayed-workflow check error:", err);
+  }
 }
 
 async function fireSchedule(
@@ -332,6 +532,7 @@ async function fireSchedule(
     agentProfile: schedule.agentProfile,
     priority: 2,
     sourceType: "scheduled",
+    maxTurns: schedule.maxTurns, // per-schedule override, NULL = inherit global
     createdAt: now,
     updatedAt: now,
   });
@@ -378,6 +579,23 @@ async function fireSchedule(
     })
     .where(eq(schedules.id, schedule.id));
 
+  // Atomic slot claim — if the global cap is full, leave the task in queued
+  // state. The task will be picked up by drainQueue when a currently-running
+  // task completes (its .then(drainQueue) chain runs the drain loop), OR by
+  // the next tickScheduler pass up to POLL_INTERVAL_MS (60s) later — whichever
+  // comes first. In a saturated-cap scenario where no running task completes
+  // before the next poll, expect up to a 60s drain latency.
+  const cap = getScheduleMaxConcurrent();
+  const leaseSec = schedule.maxRunDurationSec ?? getScheduleMaxRunDurationSec();
+  const { claimed } = claimSlot(taskId, cap, leaseSec);
+
+  if (!claimed) {
+    console.warn(
+      `[scheduler] schedule "${schedule.name}" queued — cap full (${countRunningScheduledSlots()}/${cap})`,
+    );
+    return;
+  }
+
   // Drain-aware task execution. We still don't await in fireSchedule (the
   // poll loop must keep claiming other due schedules), but on completion we
   // record metrics and trigger drainQueue() so any tasks queued by colliding
@@ -510,6 +728,7 @@ async function fireHeartbeat(
     agentProfile: schedule.agentProfile,
     priority: 2,
     sourceType: "heartbeat",
+    maxTurns: schedule.maxTurns, // per-schedule override, NULL = inherit global
     createdAt: now,
     updatedAt: now,
   });

package/src/lib/schedules/slot-claim.ts

@@ -0,0 +1,105 @@
+import { sqlite } from "@/lib/db";
+import { getExecution } from "@/lib/agents/execution-manager";
+
+export interface ClaimResult {
+  claimed: boolean;
+}
+
+// Module-level prepared statements. These are hot-path primitives called on
+// every scheduler tick and drain pass, so we pay the SQL compilation cost once
+// at module load rather than on every invocation.
+const claimStmt = sqlite.prepare(
+  "UPDATE tasks SET status = 'running', slot_claimed_at = ?, lease_expires_at = ?, updated_at = ? WHERE id = ? AND status = 'queued' AND source_type IN ('scheduled', 'heartbeat') AND (SELECT COUNT(*) FROM tasks WHERE status = 'running' AND source_type IN ('scheduled', 'heartbeat')) < ?",
+);
+
+const countRunningStmt = sqlite.prepare(
+  "SELECT COUNT(*) AS n FROM tasks WHERE status = 'running' AND source_type IN ('scheduled', 'heartbeat')",
+);
+
+/**
+ * Atomic slot claim: transitions a queued scheduled task to running IFF the
+ * global cap of concurrent running scheduled tasks is not exceeded.
+ *
+ * Must be a single SQL statement — check-then-act would race between the
+ * scheduler tick loop and the drain loop that scheduler.ts currently dispatches
+ * concurrently. Using a subquery inside the WHERE clause guarantees SQLite
+ * serializes the count and update under its write lock, so two concurrent
+ * claim attempts cannot both succeed against the same cap.
+ *
+ * Returns `{ claimed: true }` when the task transitioned; `{ claimed: false }`
+ * when either (a) the task is no longer in queued state (already claimed) or
+ * (b) the global cap is full.
+ *
+ * @param cap — must be ≥ 0. Negative values refuse to claim (the SQL COUNT
+ *   cannot be less than a negative number) — treat as an input error upstream.
+ * @param leaseSec — must be > 0 for the lease to be meaningful. A lease of 0
+ *   expires immediately and will be reaped on the next scheduler tick.
+ */
+export function claimSlot(
+  taskId: string,
+  cap: number,
+  leaseSec: number,
+): ClaimResult {
+  // Drizzle integer({ mode: "timestamp" }) stores Unix seconds and deserializes
+  // to Date(seconds * 1000). Pass seconds here so the round-trip is correct.
+  // Use Math.ceil so slotClaimedAt.getTime() >= Date.now() captured just before
+  // the call (sub-second precision would cause test assertions to fail with floor).
+  const nowSec = Math.ceil(Date.now() / 1000);
+  const leaseExpiresSec = nowSec + leaseSec;
+
+  const result = claimStmt.run(nowSec, leaseExpiresSec, nowSec, taskId, cap);
+  return { claimed: result.changes === 1 };
+}
+
+/**
+ * Count currently running scheduled/heartbeat tasks — used by the drain loop,
+ * manual-execute endpoint, and telemetry.
+ */
+export function countRunningScheduledSlots(): number {
+  const row = countRunningStmt.get() as { n: number } | undefined;
+  return row?.n ?? 0;
+}
+
+// Module-level prepared statements for the reaper hot path
+const selectExpiredStmt = sqlite.prepare(
+  "SELECT id FROM tasks WHERE status = 'running' AND source_type IN ('scheduled', 'heartbeat') AND lease_expires_at IS NOT NULL AND lease_expires_at < ?",
+);
+
+const reapUpdateStmt = sqlite.prepare(
+  "UPDATE tasks SET status = 'failed', failure_reason = 'lease_expired', updated_at = ? WHERE id = ? AND status = 'running'",
+);
+
+/**
+ * Reap running scheduled tasks whose lease has expired. For each expired
+ * task: (1) abort the in-memory execution via AbortController, (2) mark
+ * the DB row as failed with failure_reason='lease_expired'. Returns the
+ * list of reaped task IDs for logging.
+ *
+ * Idempotent — safe to call on every scheduler tick.
+ */
+export function reapExpiredLeases(): string[] {
+  // Drizzle mode: "timestamp" stores seconds; raw SQL comparisons must use
+  // seconds. Use Math.floor (strict < comparison, so floor catches everything
+  // already past).
+  const nowSec = Math.floor(Date.now() / 1000);
+
+  const expiredRows = selectExpiredStmt.all(nowSec) as Array<{ id: string }>;
+
+  const reaped: string[] = [];
+  for (const { id } of expiredRows) {
+    // Abort the in-process execution so the SDK stops immediately
+    const execution = getExecution(id);
+    if (execution) {
+      try {
+        execution.abortController.abort();
+      } catch {
+        // Already aborted — safe to ignore
+      }
+    }
+
+    const updateResult = reapUpdateStmt.run(nowSec, id);
+    if (updateResult.changes === 1) reaped.push(id);
+  }
+
+  return reaped;
+}

package/src/lib/settings/__tests__/openai-auth.test.ts

@@ -0,0 +1,101 @@
+import { beforeEach, describe, expect, it, vi } from "vitest";
+
+const mockFrom = vi.fn();
+const mockWhere = vi.fn();
+const mockValues = vi.fn();
+const mockSet = vi.fn();
+const mockRun = vi.fn();
+
+vi.mock("@/lib/db", () => ({
+  db: {
+    select: () => ({ from: mockFrom }),
+    insert: () => ({ values: mockValues }),
+    update: () => ({ set: mockSet }),
+  },
+}));
+
+vi.mock("@/lib/db/schema", () => ({
+  settings: { key: "key" },
+}));
+
+vi.mock("@/lib/utils/crypto", () => ({
+  encrypt: vi.fn((v: string) => `encrypted:${v}`),
+  decrypt: vi.fn((v: string) => v.replace("encrypted:", "")),
+}));
+
+vi.mock("@/lib/utils/stagent-paths", () => ({
+  getStagentCodexAuthPath: () => "/tmp/stagent-codex/auth.json",
+}));
+
+mockFrom.mockReturnValue({ where: mockWhere });
+mockValues.mockReturnValue({ run: mockRun });
+mockSet.mockReturnValue({ where: vi.fn().mockReturnValue({ run: mockRun }) });
+
+function mockGetSettingSequence(values: (string | null)[]) {
+  let callIndex = 0;
+  mockWhere.mockImplementation(() => {
+    const val = values[callIndex] ?? null;
+    callIndex++;
+    return val !== null ? [{ value: val }] : [];
+  });
+}
+
+describe("openai auth settings", () => {
+  beforeEach(() => {
+    vi.clearAllMocks();
+    vi.resetModules();
+    vi.unstubAllEnvs();
+    vi.stubEnv("OPENAI_API_KEY", "");
+    mockWhere.mockReturnValue([]);
+  });
+
+  it("defaults to api_key mode with no key", async () => {
+    const { getOpenAIAuthSettings } = await import("../openai-auth");
+    const result = await getOpenAIAuthSettings();
+    expect(result.method).toBe("api_key");
+    expect(result.hasKey).toBe(false);
+    expect(result.oauthConnected).toBe(false);
+  });
+
+  it("detects env-backed API key", async () => {
+    vi.stubEnv("OPENAI_API_KEY", "sk-openai");
+    mockGetSettingSequence([null, null, null, null, null, null]);
+    const { getOpenAIAuthSettings } = await import("../openai-auth");
+    const result = await getOpenAIAuthSettings();
+    expect(result.hasKey).toBe(true);
+    expect(result.apiKeySource).toBe("env");
+  });
+
+  it("returns oauth connection metadata when stored", async () => {
+    mockGetSettingSequence([
+      "oauth",
+      null,
+      null,
+      "true",
+      JSON.stringify({
+        account: { type: "chatgpt", email: "dev@example.com", planType: "pro" },
+        authMode: "chatgpt",
+      }),
+      JSON.stringify({
+        limitId: "codex",
+        limitName: null,
+        primary: { usedPercent: 25, windowDurationMins: 15, resetsAt: 1730947200 },
+        secondary: null,
+      }),
+    ]);
+
+    const { getOpenAIAuthSettings } = await import("../openai-auth");
+    const result = await getOpenAIAuthSettings();
+
+    expect(result.method).toBe("oauth");
+    expect(result.oauthConnected).toBe(true);
+    expect(result.account?.email).toBe("dev@example.com");
+    expect(result.rateLimits?.primary?.usedPercent).toBe(25);
+  });
+
+  it("stores method changes without clearing existing key data", async () => {
+    const { setOpenAIAuthSettings } = await import("../openai-auth");
+    await setOpenAIAuthSettings({ method: "oauth" });
+    expect(mockValues).toHaveBeenCalled();
+  });
+});

package/src/lib/settings/__tests__/openai-login-manager.test.ts

@@ -0,0 +1,64 @@
+import { beforeEach, describe, expect, it, vi } from "vitest";
+
+const fakeClient = {
+  request: vi.fn(),
+  close: vi.fn(async () => {}),
+  onProcessError: undefined as ((error: Error) => void) | undefined,
+  onNotification: undefined as
+    | ((notification: { method: string; params?: unknown }) => void)
+    | undefined,
+};
+
+vi.mock("@/lib/agents/runtime/openai-codex-auth", () => ({
+  connectStagentCodexClient: vi.fn(async () => fakeClient),
+  initializeCodexClient: vi.fn(async () => {}),
+  readStagentCodexAuthState: vi.fn(async () => ({
+    connected: false,
+    account: null,
+    rateLimits: null,
+  })),
+}));
+
+vi.mock("@/lib/settings/openai-auth", () => ({
+  clearOpenAIOAuthStatus: vi.fn(async () => {}),
+}));
+
+describe("openai login manager", () => {
+  beforeEach(() => {
+    vi.resetModules();
+    vi.clearAllMocks();
+    fakeClient.request.mockImplementation(async (method: string) => {
+      if (method === "account/login/start") {
+        return {
+          type: "chatgpt",
+          loginId: "login-1",
+          authUrl: "https://auth.openai.com/log-in",
+        };
+      }
+
+      if (method === "account/login/cancel") {
+        return {};
+      }
+
+      throw new Error(`Unexpected method: ${method}`);
+    });
+  });
+
+  it("returns a cancelled state when the active ChatGPT login is cancelled", async () => {
+    const {
+      startOpenAIChatGPTLogin,
+      cancelOpenAIChatGPTLogin,
+      getOpenAILoginState,
+    } = await import("@/lib/settings/openai-login-manager");
+
+    const started = await startOpenAIChatGPTLogin();
+    expect(started.phase).toBe("pending");
+
+    const cancelled = await cancelOpenAIChatGPTLogin();
+
+    expect(cancelled.phase).toBe("cancelled");
+    expect(cancelled.error).toBeNull();
+    expect(fakeClient.close).toHaveBeenCalledTimes(1);
+    expect(getOpenAILoginState().phase).toBe("cancelled");
+  });
+});