pi-oracle 0.1.12 → 0.2.0

package/extensions/oracle/lib/locks.ts

@@ -1,8 +1,8 @@
 import { createHash } from "node:crypto";
 import { existsSync } from "node:fs";
-import { mkdirSync, readdirSync, readFileSync } from "node:fs";
-import { mkdir, readFile, rm, writeFile } from "node:fs/promises";
-import { join } from "node:path";
+import { mkdirSync, readdirSync, readFileSync, statSync } from "node:fs";
+import { mkdir, readFile, rename, rm, writeFile } from "node:fs/promises";
+import { basename, join } from "node:path";
 
 export const DEFAULT_ORACLE_STATE_DIR = "/tmp/pi-oracle-state";
 export const ORACLE_STATE_DIR_ENV = "PI_ORACLE_STATE_DIR";
@@ -11,6 +11,7 @@ const LOCKS_DIR = join(ORACLE_STATE_DIR, "locks");
 const LEASES_DIR = join(ORACLE_STATE_DIR, "leases");
 const DEFAULT_WAIT_MS = 30_000;
 const POLL_MS = 200;
+export const ORACLE_METADATA_WRITE_GRACE_MS = 1_000;
 
 export interface OracleLockHandle {
   path: string;
@@ -52,11 +53,51 @@ async function sleep(ms: number): Promise<void> {
 }
 
 async function writeMetadata(path: string, metadata: unknown): Promise<void> {
-  await writeFile(join(path, "metadata.json"), `${JSON.stringify(metadata, null, 2)}\n`, { encoding: "utf8", mode: 0o600 });
+  const targetPath = join(path, "metadata.json");
+  const tempPath = join(path, `metadata.${process.pid}.${Date.now()}.${Math.random().toString(36).slice(2)}.tmp`);
+  await writeFile(tempPath, `${JSON.stringify(metadata, null, 2)}\n`, { encoding: "utf8", mode: 0o600 });
+  await rename(tempPath, targetPath);
+}
+
+async function createStateDirAtomically(parentDir: string, finalPath: string, metadata: unknown): Promise<void> {
+  const tempPath = join(parentDir, `.tmp-${basename(finalPath)}.${process.pid}.${Date.now()}.${Math.random().toString(36).slice(2)}`);
+  await mkdir(tempPath, { recursive: false, mode: 0o700 });
+  try {
+    await writeMetadata(tempPath, metadata);
+    await rename(tempPath, finalPath);
+  } catch (error) {
+    await rm(tempPath, { recursive: true, force: true }).catch(() => undefined);
+    throw error;
+  }
+}
+
+function getMetadataPath(path: string): string {
+  return join(path, "metadata.json");
+}
+
+function getMetadataState(path: string): "present" | "missing" | "invalid" {
+  const metadataPath = getMetadataPath(path);
+  if (!existsSync(metadataPath)) return "missing";
+  try {
+    JSON.parse(readFileSync(metadataPath, "utf8"));
+    return "present";
+  } catch {
+    return "invalid";
+  }
+}
+
+function isIncompleteStateDirStale(path: string, now = Date.now()): boolean {
+  try {
+    const stats = statSync(path);
+    const baselineMs = Math.max(stats.mtimeMs, stats.ctimeMs);
+    return now - baselineMs >= ORACLE_METADATA_WRITE_GRACE_MS;
+  } catch {
+    return false;
+  }
 }
 
 function readLockProcessPid(path: string): number | undefined {
-  const metadataPath = join(path, "metadata.json");
+  const metadataPath = getMetadataPath(path);
   if (!existsSync(metadataPath)) return undefined;
   try {
     const metadata = JSON.parse(readFileSync(metadataPath, "utf8")) as { processPid?: unknown };
@@ -78,7 +119,19 @@ function isProcessAlive(pid: number): boolean {
   }
 }
 
-async function maybeReclaimStaleLock(path: string): Promise<boolean> {
+function isStateDirExistsError(error: unknown): boolean {
+  return Boolean(error && typeof error === "object" && "code" in error && (error.code === "EEXIST" || error.code === "ENOTEMPTY"));
+}
+
+async function maybeReclaimIncompleteStateDir(path: string, now = Date.now()): Promise<boolean> {
+  if (getMetadataState(path) === "present") return false;
+  if (!isIncompleteStateDirStale(path, now)) return false;
+  await rm(path, { recursive: true, force: true }).catch(() => undefined);
+  return true;
+}
+
+async function maybeReclaimStaleLock(path: string, now = Date.now()): Promise<boolean> {
+  if (await maybeReclaimIncompleteStateDir(path, now)) return true;
   const processPid = readLockProcessPid(path);
   if (!processPid || isProcessAlive(processPid)) return false;
   await rm(path, { recursive: true, force: true }).catch(() => undefined);
@@ -105,17 +158,17 @@ export async function acquireLock(
   metadata: unknown,
   options?: { timeoutMs?: number },
 ): Promise<OracleLockHandle> {
-  const path = lockPath(kind, key);
+  const parentDir = getLocksDir();
+  const path = join(parentDir, leaseKey(kind, key));
   const timeoutMs = options?.timeoutMs ?? DEFAULT_WAIT_MS;
   const deadline = Date.now() + timeoutMs;
 
   while (Date.now() < deadline) {
     try {
-      await mkdir(path, { recursive: false, mode: 0o700 });
-      await writeMetadata(path, metadata);
+      await createStateDirAtomically(parentDir, path, metadata);
       return { path };
     } catch (error) {
-      if (!(error && typeof error === "object" && "code" in error && error.code === "EEXIST")) throw error;
+      if (!isStateDirExistsError(error)) throw error;
       if (await maybeReclaimStaleLock(path)) continue;
     }
     await sleep(POLL_MS);
@@ -176,9 +229,42 @@ export async function withJobLock<T>(jobId: string, metadata: unknown, fn: () =>
 }
 
 export async function createLease(kind: string, key: string, metadata: unknown): Promise<string> {
-  const path = leasePath(kind, key);
-  await mkdir(path, { recursive: false, mode: 0o700 });
-  await writeMetadata(path, metadata);
+  const parentDir = getLeasesDir();
+  const path = join(parentDir, leaseKey(kind, key));
+  const deadline = Date.now() + DEFAULT_WAIT_MS;
+
+  while (Date.now() < deadline) {
+    try {
+      await createStateDirAtomically(parentDir, path, metadata);
+      return path;
+    } catch (error) {
+      if (!isStateDirExistsError(error)) throw error;
+      if (await maybeReclaimIncompleteStateDir(path)) continue;
+      if (getMetadataState(path) === "present") throw error;
+    }
+    await sleep(POLL_MS);
+  }
+
+  throw new Error(`Timed out waiting for oracle ${kind} lease: ${key}`);
+}
+
+export async function writeLeaseMetadata(kind: string, key: string, metadata: unknown): Promise<string> {
+  const parentDir = getLeasesDir();
+  const path = join(parentDir, leaseKey(kind, key));
+  if (existsSync(path)) {
+    await writeMetadata(path, metadata);
+    return path;
+  }
+  try {
+    await createStateDirAtomically(parentDir, path, metadata);
+  } catch (error) {
+    if (!isStateDirExistsError(error)) throw error;
+    if (await maybeReclaimIncompleteStateDir(path)) {
+      await createStateDirAtomically(parentDir, path, metadata);
+    } else {
+      await writeMetadata(path, metadata);
+    }
+  }
   return path;
 }
 

package/extensions/oracle/lib/poller.ts

@@ -1,21 +1,56 @@
+import { execFileSync } from "node:child_process";
 import type { ExtensionAPI, ExtensionContext } from "@mariozechner/pi-coding-agent";
-import { isLockTimeoutError, withGlobalReconcileLock } from "./locks.js";
+import { isLockTimeoutError, listLeaseMetadata, releaseLease, withGlobalReconcileLock, writeLeaseMetadata } from "./locks.js";
 import {
+  getJobDir,
   getSessionFile,
   getStaleOracleJobReason,
+  hasPersistedOriginSession,
   isActiveOracleJob,
   listOracleJobDirs,
-  markJobNotified,
+  noteWakeupRequested,
   readJob,
+  recordNotificationTarget,
   reconcileStaleOracleJobs,
   releaseNotificationClaim,
+  shouldPruneTerminalJob,
+  shouldRequestWakeup,
   tryClaimNotification,
 } from "./jobs.js";
+import { promoteQueuedJobs } from "./queue.js";
 import { getProjectId, getSessionId } from "./runtime.js";
 
 const activePollers = new Map<string, NodeJS.Timeout>();
 const scansInFlight = new Set<string>();
 const POLLER_LOCK_TIMEOUT_MS = 50;
+const WAKEUP_TARGET_LEASE_KIND = "wakeup-target";
+const WAKEUP_TARGET_STALE_MS = 2 * 60 * 1000;
+const ORACLE_WAKEUP_REMINDER_CUSTOM_TYPE = "oracle-job-wakeup";
+
+interface OracleWakeupTargetLeaseMetadata {
+  leaseKey: string;
+  projectId: string;
+  sessionId: string;
+  processPid: number;
+  processStartedAt?: string;
+  updatedAt: string;
+}
+
+
+type OraclePollerJob = NonNullable<ReturnType<typeof readJob>>;
+
+export interface OraclePollerHooks {
+  collectLiveWakeupTargets?: (now?: number) => Promise<Set<string>>;
+  beforeNotificationClaim?: (jobId: string) => Promise<void> | void;
+  afterNotificationClaim?: (job: OraclePollerJob) => Promise<void> | void;
+  beforeNotificationPersist?: (job: OraclePollerJob) => Promise<void> | void;
+  afterNotificationPersisted?: (job: OraclePollerJob) => Promise<void> | void;
+  beforeMarkJobNotified?: (job: OraclePollerJob) => Promise<void> | void;
+}
+
+export interface OraclePollerOptions {
+  hooks?: OraclePollerHooks;
+}
 
 export function getPollerSessionKey(sessionFile: string | undefined, cwd: string): string {
   const projectId = getProjectId(cwd);
@@ -29,53 +64,161 @@ function jobMatchesContext(job: { projectId: string; sessionId: string }, sessio
   return job.projectId === projectId && job.sessionId === sessionId;
 }
 
-function getActiveJobCount(ctx: ExtensionContext): number {
+function readProcessStartedAt(pid: number | undefined): string | undefined {
+  if (!pid || pid <= 0) return undefined;
+  try {
+    const startedAt = execFileSync("ps", ["-o", "lstart=", "-p", String(pid)], { encoding: "utf8" }).trim();
+    return startedAt || undefined;
+  } catch {
+    return undefined;
+  }
+}
+
+function isProcessAlive(pid: number): boolean {
+  try {
+    process.kill(pid, 0);
+    return true;
+  } catch (error) {
+    if (error && typeof error === "object" && "code" in error && error.code === "ESRCH") return false;
+    return true;
+  }
+}
+
+function parseTimestamp(value: string | undefined): number | undefined {
+  if (!value) return undefined;
+  const parsed = Date.parse(value);
+  return Number.isFinite(parsed) ? parsed : undefined;
+}
+
+function getWakeupTargetLeaseKey(sessionKey: string, processPid = process.pid, processStartedAt = readProcessStartedAt(process.pid) || "unknown"): string {
+  return `${sessionKey}::${processPid}::${processStartedAt}`;
+}
+
+async function collectLiveWakeupTargetLeases(now = Date.now()): Promise<Array<OracleWakeupTargetLeaseMetadata & { sessionKey: string }>> {
+  const liveTargets: Array<OracleWakeupTargetLeaseMetadata & { sessionKey: string }> = [];
+  for (const lease of listLeaseMetadata<OracleWakeupTargetLeaseMetadata>(WAKEUP_TARGET_LEASE_KIND)) {
+    const sessionKey = `${lease?.projectId ?? ""}::${lease?.sessionId ?? ""}`;
+    const leaseKey = lease?.leaseKey;
+    const currentStartedAt = readProcessStartedAt(lease?.processPid);
+    const updatedAtMs = parseTimestamp(lease?.updatedAt);
+    const stale = updatedAtMs !== undefined && now - updatedAtMs > WAKEUP_TARGET_STALE_MS;
+    const missingIdentity = !lease?.projectId || !lease?.sessionId || !lease?.processPid || !leaseKey;
+    const deadProcess = !missingIdentity && (!isProcessAlive(lease.processPid) || !currentStartedAt || (lease.processStartedAt && currentStartedAt !== lease.processStartedAt));
+    if (missingIdentity || deadProcess || stale) {
+      if (leaseKey) {
+        await releaseLease(WAKEUP_TARGET_LEASE_KIND, leaseKey).catch(() => undefined);
+      }
+      continue;
+    }
+    liveTargets.push({ ...lease, sessionKey } as OracleWakeupTargetLeaseMetadata & { sessionKey: string });
+  }
+  return liveTargets;
+}
+
+async function collectLiveWakeupTargets(now = Date.now()): Promise<Set<string>> {
+  return new Set((await collectLiveWakeupTargetLeases(now)).map((lease) => lease.sessionKey));
+}
+
+function jobHasLiveWakeupTarget(job: { projectId: string; sessionId: string }, liveWakeupTargets: Set<string>): boolean {
+  return liveWakeupTargets.has(`${job.projectId}::${job.sessionId}`);
+}
+
+function jobCanNotifyContext(
+  job: { projectId: string; sessionId: string; originSessionFile?: string },
+  sessionFile: string | undefined,
+  cwd: string,
+  liveWakeupTargets: Set<string>,
+): boolean {
+  if (!hasPersistedOriginSession(job)) return false;
+  if (jobMatchesContext(job, sessionFile, cwd)) return true;
+  return job.projectId === getProjectId(cwd) && !jobHasLiveWakeupTarget(job, liveWakeupTargets);
+}
+
+function getJobCounts(ctx: ExtensionContext): { active: number; queued: number } {
   const currentSessionFile = getSessionFile(ctx);
+  if (!currentSessionFile) return { active: 0, queued: 0 };
   return listOracleJobDirs()
     .map((jobDir) => readJob(jobDir))
     .filter((job): job is NonNullable<typeof job> => Boolean(job))
-    .filter((job) => {
-      if (!isActiveOracleJob(job)) return false;
-      if (getStaleOracleJobReason(job)) return false;
-      return jobMatchesContext(job, currentSessionFile, ctx.cwd);
-    }).length;
+    .filter((job) => jobMatchesContext(job, currentSessionFile, ctx.cwd))
+    .reduce(
+      (counts, job) => {
+        if (job.status === "queued") counts.queued += 1;
+        else if (isActiveOracleJob(job) && !getStaleOracleJobReason(job)) counts.active += 1;
+        return counts;
+      },
+      { active: 0, queued: 0 },
+    );
 }
 
 export function refreshOracleStatus(ctx: ExtensionContext): void {
-  const activeJobCount = getActiveJobCount(ctx);
-  if (activeJobCount > 0) {
-    const suffix = activeJobCount > 1 ? ` (${activeJobCount})` : "";
+  if (!getSessionFile(ctx)) {
+    ctx.ui.setStatus("oracle", ctx.ui.theme.fg("accent", "oracle: unavailable"));
+    return;
+  }
+  const counts = getJobCounts(ctx);
+  if (counts.active > 0 && counts.queued > 0) {
+    ctx.ui.setStatus("oracle", ctx.ui.theme.fg("success", `oracle: running (${counts.active}), queued (${counts.queued})`));
+    return;
+  }
+  if (counts.active > 0) {
+    const suffix = counts.active > 1 ? ` (${counts.active})` : "";
     ctx.ui.setStatus("oracle", ctx.ui.theme.fg("success", `oracle: running${suffix}`));
     return;
   }
+  if (counts.queued > 0) {
+    const suffix = counts.queued > 1 ? ` (${counts.queued})` : "";
+    ctx.ui.setStatus("oracle", ctx.ui.theme.fg("accent", `oracle: queued${suffix}`));
+    return;
+  }
 
   ctx.ui.setStatus("oracle", ctx.ui.theme.fg("accent", "oracle: ready"));
 }
 
-function notifyForJob(pi: ExtensionAPI, job: NonNullable<ReturnType<typeof readJob>>): void {
-  const responsePath = job.responsePath || `${job.id}/response.md`;
-  const artifactsPath = `/tmp/oracle-${job.id}/artifacts`;
+function buildNotificationContent(job: OraclePollerJob): string {
+  const responsePath = job.responsePath || `${getJobDir(job.id)}/response.md`;
+  const artifactsPath = `${getJobDir(job.id)}/artifacts`;
+  return [
+    `Oracle job ${job.id} is ${job.status}.`,
+    `Read response: ${responsePath}`,
+    `Artifacts: ${artifactsPath}`,
+    job.error ? `Error: ${job.error}` : "Continue from the oracle output.",
+  ].join("\n");
+}
+
+
+function requestWakeupTurn(pi: ExtensionAPI, job: OraclePollerJob): void {
   pi.sendMessage(
     {
-      customType: "oracle-job-complete",
-      display: true,
-      content: [
-        `Oracle job ${job.id} is ${job.status}.`,
-        `Read response: ${responsePath}`,
-        `Artifacts: ${artifactsPath}`,
-        job.error ? `Error: ${job.error}` : "Continue from the oracle output.",
-      ].join("\n"),
+      customType: ORACLE_WAKEUP_REMINDER_CUSTOM_TYPE,
+      display: false,
+      content: buildNotificationContent(job),
       details: { jobId: job.id, status: job.status },
     },
-    { triggerTurn: true },
+    { triggerTurn: true, deliverAs: "followUp" },
   );
 }
 
-async function scan(pi: ExtensionAPI, ctx: ExtensionContext): Promise<void> {
+async function scan(pi: ExtensionAPI, ctx: ExtensionContext, workerPath: string, hooks: OraclePollerHooks = {}): Promise<void> {
   const currentSessionFile = getSessionFile(ctx);
   const pollerKey = getPollerSessionKey(currentSessionFile, ctx.cwd);
   const notificationClaimant = `${pollerKey}:${process.pid}`;
 
+  const projectId = getProjectId(ctx.cwd);
+  const sessionId = getSessionId(currentSessionFile, projectId);
+  const processStartedAt = readProcessStartedAt(process.pid);
+  const wakeupTargetLeaseKey = getWakeupTargetLeaseKey(pollerKey, process.pid, processStartedAt || "unknown");
+  const resolveLiveWakeupTargets = hooks.collectLiveWakeupTargets ?? collectLiveWakeupTargets;
+  await writeLeaseMetadata(WAKEUP_TARGET_LEASE_KIND, wakeupTargetLeaseKey, {
+    leaseKey: wakeupTargetLeaseKey,
+    projectId,
+    sessionId,
+    processPid: process.pid,
+    processStartedAt,
+    updatedAt: new Date().toISOString(),
+  }).catch(() => undefined);
+  const liveWakeupTargets = await resolveLiveWakeupTargets();
+
   try {
     await withGlobalReconcileLock(
       { processPid: process.pid, cwd: ctx.cwd, sessionFile: currentSessionFile, source: "poller" },
@@ -88,27 +231,60 @@ async function scan(pi: ExtensionAPI, ctx: ExtensionContext): Promise<void> {
     if (!isLockTimeoutError(error, "reconcile", "global")) throw error;
   }
 
-  const candidateJobIds = listOracleJobDirs()
+  await promoteQueuedJobs({ workerPath, source: "poller" });
+
+  const terminalJobs = listOracleJobDirs()
     .map((jobDir) => readJob(jobDir))
     .filter((job): job is NonNullable<typeof job> => Boolean(job))
+    .filter((job) => job.status === "complete" || job.status === "failed" || job.status === "cancelled");
+
+  const now = Date.now();
+  const candidateJobIds = terminalJobs
     .filter((job) => {
-      if (job.status !== "complete" && job.status !== "failed" && job.status !== "cancelled") return false;
-      if (!jobMatchesContext(job, currentSessionFile, ctx.cwd)) return false;
-      return !job.notifiedAt;
+      if (!jobCanNotifyContext(job, currentSessionFile, ctx.cwd, liveWakeupTargets)) return false;
+      if (job.notifiedAt) return false;
+      if (shouldPruneTerminalJob(job, now)) return false;
+      return shouldRequestWakeup(job, now);
     })
     .map((job) => job.id);
 
   for (const jobId of candidateJobIds) {
+    await hooks.beforeNotificationClaim?.(jobId);
     const claimed = await tryClaimNotification(jobId, notificationClaimant);
     if (!claimed) continue;
-    if (!jobMatchesContext(claimed, currentSessionFile, ctx.cwd)) {
+
+    await hooks.afterNotificationClaim?.(claimed);
+    const preNotifyLiveWakeupTargets = await resolveLiveWakeupTargets();
+    if (!jobCanNotifyContext(claimed, currentSessionFile, ctx.cwd, preNotifyLiveWakeupTargets)) {
       await releaseNotificationClaim(jobId, notificationClaimant).catch(() => undefined);
       continue;
     }
 
     try {
-      notifyForJob(pi, claimed);
-      await markJobNotified(jobId, notificationClaimant);
+      if (currentSessionFile) {
+        await recordNotificationTarget(jobId, notificationClaimant, {
+          notificationSessionKey: pollerKey,
+          notificationSessionFile: currentSessionFile,
+        });
+      }
+      await hooks.beforeNotificationPersist?.(claimed);
+      const preWakeupLiveWakeupTargets = await resolveLiveWakeupTargets();
+      if (!jobCanNotifyContext(claimed, currentSessionFile, ctx.cwd, preWakeupLiveWakeupTargets)) {
+        await releaseNotificationClaim(jobId, notificationClaimant).catch(() => undefined);
+        continue;
+      }
+      const deliverable = readJob(jobId);
+      if (!deliverable || shouldPruneTerminalJob(deliverable, Date.now())) {
+        await releaseNotificationClaim(jobId, notificationClaimant).catch(() => undefined);
+        continue;
+      }
+
+      requestWakeupTurn(pi, deliverable);
+      await noteWakeupRequested(jobId).catch(() => undefined);
+      if (ctx.hasUI) {
+        ctx.ui.notify(`Oracle job ${claimed.id} is ${claimed.status}.`, "info");
+      }
+      await releaseNotificationClaim(jobId, notificationClaimant).catch(() => undefined);
     } catch (error) {
       await releaseNotificationClaim(jobId, notificationClaimant).catch(() => undefined);
       throw error;
@@ -116,7 +292,11 @@ async function scan(pi: ExtensionAPI, ctx: ExtensionContext): Promise<void> {
   }
 }
 
-export function startPoller(pi: ExtensionAPI, ctx: ExtensionContext, intervalMs: number): void {
+export async function scanOracleJobsOnce(pi: ExtensionAPI, ctx: ExtensionContext, workerPath: string, options: OraclePollerOptions = {}): Promise<void> {
+  await scan(pi, ctx, workerPath, options.hooks);
+}
+
+export function startPoller(pi: ExtensionAPI, ctx: ExtensionContext, intervalMs: number, workerPath: string, options: OraclePollerOptions = {}): void {
   const sessionKey = getPollerSessionKey(getSessionFile(ctx), ctx.cwd);
   const existing = activePollers.get(sessionKey);
   if (existing) clearInterval(existing);
@@ -125,7 +305,7 @@ export function startPoller(pi: ExtensionAPI, ctx: ExtensionContext, intervalMs:
     if (scansInFlight.has(sessionKey)) return;
     scansInFlight.add(sessionKey);
     try {
-      await scan(pi, ctx);
+      await scanOracleJobsOnce(pi, ctx, workerPath, options);
     } catch (error) {
       console.error(`Oracle poller scan failed (${sessionKey}):`, error);
     } finally {
@@ -145,12 +325,17 @@ export function startPoller(pi: ExtensionAPI, ctx: ExtensionContext, intervalMs:
 export function stopPollerForSession(sessionFile: string | undefined, cwd: string): void {
   const sessionKey = getPollerSessionKey(sessionFile, cwd);
   const timer = activePollers.get(sessionKey);
-  if (!timer) return;
-  clearInterval(timer);
-  activePollers.delete(sessionKey);
-  scansInFlight.delete(sessionKey);
+  if (timer) {
+    clearInterval(timer);
+    activePollers.delete(sessionKey);
+    scansInFlight.delete(sessionKey);
+  }
+  const wakeupTargetLeaseKey = getWakeupTargetLeaseKey(sessionKey);
+  void releaseLease(WAKEUP_TARGET_LEASE_KIND, wakeupTargetLeaseKey).catch(() => undefined);
 }
 
 export function stopPoller(ctx: ExtensionContext): void {
-  stopPollerForSession(getSessionFile(ctx), ctx.cwd);
+  const sessionFile = getSessionFile(ctx);
+  if (!sessionFile) return;
+  stopPollerForSession(sessionFile, ctx.cwd);
 }