@duckmind/dm-darwin-arm64 0.13.6 → 0.13.8

package/extensions/dm-phone/src/session-pool/session-worker.ts

@@ -0,0 +1,734 @@
+import { calculateContextTokens, estimateTokens } from "@mariozechner/pi-coding-agent";
+import { spawn, type ChildProcessWithoutNullStreams } from "node:child_process";
+import { StringDecoder } from "node:string_decoder";
+import type {
+  PendingClientResponse,
+  PendingRequest,
+  SessionSnapshot,
+  SessionSummary,
+  SessionWorkerOptions,
+} from "./types";
+import { contentToPreviewText, shortId } from "./utils";
+
+let workerCounter = 0;
+
+export class PhoneSessionWorker {
+  id = `active-session-${++workerCounter}`;
+  kind = "parallel" as const;
+  cwd: string;
+  previousCwd: string | null = null;
+  currentSessionFile: string | null;
+  child: ChildProcessWithoutNullStreams | null = null;
+  lastError = "";
+  lastState: any = null;
+  lastMessages: any[] = [];
+  lastCommands: any[] = [];
+  isStreaming = false;
+  lastActivityAt = Date.now();
+  pendingUiRequest: any = null;
+  liveAssistantMessage: any = null;
+  liveTools = new Map<string, any>();
+
+  private readonly options: SessionWorkerOptions<PhoneSessionWorker>;
+  private readonly decoder = new StringDecoder("utf8");
+  private stdoutBuffer = "";
+  private startPromise: Promise<void> | null = null;
+  private requestCounter = 0;
+  private pendingRequests = new Map<string, PendingRequest>();
+  private pendingClientResponses = new Map<string, PendingClientResponse>();
+  private snapshotRefreshTimer: NodeJS.Timeout | null = null;
+  private reloadPromise: Promise<void> | null = null;
+  private isRestarting = false;
+  private disposed = false;
+  private firstUserPreview = "";
+  private lastUserPreview = "";
+
+  constructor(options: SessionWorkerOptions<PhoneSessionWorker>, sessionFile: string | null = null) {
+    this.options = options;
+    this.cwd = options.cwd;
+    this.currentSessionFile = sessionFile;
+  }
+
+  private touch() {
+    this.lastActivityAt = Date.now();
+    this.options.onActivity();
+    this.options.onStateChange();
+  }
+
+  private updateMessagePreviews() {
+    const firstUser = this.lastMessages.find((message) => message?.role === "user");
+    const lastUser = [...this.lastMessages].reverse().find((message) => message?.role === "user");
+    this.firstUserPreview = firstUser ? contentToPreviewText(firstUser.content) : "";
+    this.lastUserPreview = lastUser ? contentToPreviewText(lastUser.content) : "";
+  }
+
+  private setMessages(messages: any[]) {
+    this.lastMessages = Array.isArray(messages) ? messages : [];
+    this.updateMessagePreviews();
+    this.syncDerivedState();
+    this.options.onStateChange();
+  }
+
+  private estimateContextTokens(messages: any[]) {
+    let lastUsageIndex = -1;
+    let usageTokens = 0;
+
+    for (let index = messages.length - 1; index >= 0; index -= 1) {
+      const message = messages[index];
+      if (message?.role !== "assistant") continue;
+      if (message.stopReason === "aborted" || message.stopReason === "error") continue;
+      if (!message.usage) continue;
+
+      usageTokens = calculateContextTokens(message.usage);
+      lastUsageIndex = index;
+      break;
+    }
+
+    if (lastUsageIndex === -1) {
+      let estimated = 0;
+      for (const message of messages) {
+        estimated += estimateTokens(message);
+      }
+      return estimated;
+    }
+
+    let trailingTokens = 0;
+    for (let index = lastUsageIndex + 1; index < messages.length; index += 1) {
+      trailingTokens += estimateTokens(messages[index]);
+    }
+
+    return usageTokens + trailingTokens;
+  }
+
+  private buildContextUsage() {
+    const contextWindow = Number(this.lastState?.model?.contextWindow);
+    if (!Number.isFinite(contextWindow) || contextWindow <= 0) {
+      return null;
+    }
+
+    let latestCompactionIndex = -1;
+    for (let index = this.lastMessages.length - 1; index >= 0; index -= 1) {
+      if (this.lastMessages[index]?.role === "compactionSummary") {
+        latestCompactionIndex = index;
+        break;
+      }
+    }
+
+    if (latestCompactionIndex !== -1) {
+      let hasPostCompactionUsage = false;
+
+      for (let index = this.lastMessages.length - 1; index > latestCompactionIndex; index -= 1) {
+        const message = this.lastMessages[index];
+        if (message?.role !== "assistant") continue;
+        if (message.stopReason === "aborted" || message.stopReason === "error") continue;
+
+        if (message.usage && calculateContextTokens(message.usage) > 0) {
+          hasPostCompactionUsage = true;
+        }
+        break;
+      }
+
+      if (!hasPostCompactionUsage) {
+        return {
+          tokens: null,
+          contextWindow,
+          percent: null,
+        };
+      }
+    }
+
+    const tokens = this.estimateContextTokens(this.lastMessages);
+    return {
+      tokens,
+      contextWindow,
+      percent: (tokens / contextWindow) * 100,
+    };
+  }
+
+  private syncDerivedState() {
+    if (!this.lastState || typeof this.lastState !== "object") return;
+
+    const contextUsage = this.buildContextUsage();
+    if (!contextUsage) {
+      const { contextUsage: _ignored, ...nextState } = this.lastState;
+      this.lastState = nextState;
+      return;
+    }
+
+    this.lastState = {
+      ...this.lastState,
+      contextUsage,
+    };
+  }
+
+  private rememberState(state: any) {
+    if (!state || typeof state !== "object") return;
+    this.lastState = state;
+    if (typeof state.isStreaming === "boolean") {
+      this.isStreaming = state.isStreaming;
+    }
+    if (typeof state.sessionFile === "string" && state.sessionFile.trim()) {
+      this.currentSessionFile = state.sessionFile;
+    }
+    this.syncDerivedState();
+    this.options.onStateChange();
+  }
+
+  private buildSpawnArgs(sessionFile = this.currentSessionFile) {
+    const args = ["--mode", "rpc"];
+    if (sessionFile) {
+      args.push("--session", sessionFile);
+    }
+    return args;
+  }
+
+  getStatus() {
+    return {
+      childRunning: Boolean(this.child),
+      cwd: this.cwd,
+      previousCwd: this.previousCwd,
+      isStreaming: this.isStreaming,
+      isCompacting: Boolean(this.lastState?.isCompacting),
+      lastError: this.lastError,
+      childPid: this.child?.pid ?? null,
+      sessionWorkerId: this.id,
+      sessionKind: "parallel" as const,
+    };
+  }
+
+  setTrackedCwd(cwd: string, previousCwd: string | null = this.cwd) {
+    this.previousCwd = previousCwd;
+    this.cwd = cwd;
+    this.options.onStateChange();
+  }
+
+  getSummary(): SessionSummary {
+    const sessionId = this.lastState?.sessionId || null;
+    const sessionName = this.lastState?.sessionName || null;
+    const label = sessionName || this.firstUserPreview || (sessionId ? `Session ${shortId(sessionId)}` : `Session ${shortId(this.id)}`);
+    const secondaryLabel = sessionName ? this.firstUserPreview || shortId(sessionId) || "" : shortId(sessionId) || "";
+
+    return {
+      id: this.id,
+      kind: "parallel",
+      sessionId,
+      sessionFile: this.currentSessionFile || this.lastState?.sessionFile || null,
+      sessionName,
+      label,
+      secondaryLabel,
+      firstUserPreview: this.firstUserPreview || null,
+      lastUserPreview: this.lastUserPreview || null,
+      model: this.lastState?.model
+        ? {
+            id: this.lastState.model.id,
+            name: this.lastState.model.name,
+            provider: this.lastState.model.provider,
+          }
+        : null,
+      isRunning: Boolean(this.child),
+      isStreaming: this.isStreaming,
+      isCompacting: Boolean(this.lastState?.isCompacting),
+      messageCount: this.lastState?.messageCount ?? this.lastMessages.length,
+      pendingMessageCount: this.lastState?.pendingMessageCount ?? 0,
+      hasPendingUiRequest: Boolean(this.pendingUiRequest),
+      lastError: this.lastError,
+      lastActivityAt: this.lastActivityAt,
+      childPid: this.child?.pid ?? null,
+      cwd: this.cwd,
+    };
+  }
+
+  private cachedSnapshot(): SessionSnapshot {
+    return {
+      state: this.lastState,
+      messages: this.lastMessages,
+      commands: this.lastCommands,
+      liveAssistantMessage: this.liveAssistantMessage,
+      liveTools: [...this.liveTools.values()],
+    };
+  }
+
+  getCachedSnapshot(): SessionSnapshot {
+    return this.cachedSnapshot();
+  }
+
+  async ensureStarted(startOptions: { sessionFile?: string | null } = {}) {
+    if (this.disposed) {
+      throw new Error("Session worker disposed.");
+    }
+
+    if (this.child) return;
+    if (this.startPromise) return this.startPromise;
+
+    const sessionFile = startOptions.sessionFile ?? this.currentSessionFile;
+    this.stdoutBuffer = "";
+
+    this.startPromise = new Promise<void>((resolvePromise, rejectPromise) => {
+      const spawned = spawn("dm", this.buildSpawnArgs(sessionFile), {
+        cwd: this.cwd,
+        env: {
+          ...process.env,
+          DM_PHONE_CHILD: "1",
+          PI_PHONE_CHILD: "1",
+        },
+        stdio: ["pipe", "pipe", "pipe"],
+      });
+
+      let settled = false;
+
+      const failStart = (error: Error) => {
+        if (settled) return;
+        settled = true;
+        this.lastError = error.message;
+        this.child = null;
+        this.options.onStateChange();
+        rejectPromise(error);
+      };
+
+      spawned.once("error", (error) => {
+        failStart(error instanceof Error ? error : new Error(String(error)));
+      });
+
+      spawned.stdout.on("data", (chunk) => {
+        this.handleStdoutChunk(chunk);
+      });
+
+      spawned.stderr.on("data", (chunk) => {
+        const text = chunk.toString();
+        this.lastError = text.trim() || this.lastError;
+        this.options.onEnvelope(this, { channel: "server", event: "stderr", data: { text } });
+        this.touch();
+      });
+
+      spawned.once("exit", (code, signal) => {
+        const message = `dm rpc exited${code !== null ? ` with code ${code}` : ""}${signal ? ` (${signal})` : ""}`;
+        const restarting = this.isRestarting;
+
+        if (!settled) {
+          failStart(new Error(message));
+          return;
+        }
+
+        this.child = null;
+        this.isStreaming = false;
+        this.lastState = this.lastState ? { ...this.lastState, isStreaming: false } : this.lastState;
+        this.pendingUiRequest = null;
+        this.liveAssistantMessage = null;
+        this.liveTools.clear();
+        this.rejectAllPending(new Error(restarting ? "DM rpc is reloading." : message));
+
+        if (restarting || this.disposed) {
+          this.lastError = "";
+          this.options.onStateChange();
+          return;
+        }
+
+        this.lastError = message;
+        this.options.onEnvelope(this, { channel: "server", event: "agent-exit", data: { code, signal, message } });
+        this.touch();
+
+        if (this.options.shouldAutoRestart(this)) {
+          setTimeout(() => {
+            if (this.disposed) return;
+            this.ensureStarted({ sessionFile: this.currentSessionFile })
+              .then(() => this.refreshCachedSnapshot().catch(() => {}))
+              .catch((error) => {
+                this.lastError = error instanceof Error ? error.message : String(error);
+                this.options.onStateChange();
+              });
+          }, 1500);
+        }
+      });
+
+      this.child = spawned;
+      this.lastError = "";
+      this.touch();
+
+      setTimeout(() => {
+        if (settled) return;
+        settled = true;
+        resolvePromise();
+      }, 300);
+    }).finally(() => {
+      this.startPromise = null;
+    });
+
+    return this.startPromise;
+  }
+
+  private rejectAllPending(error: Error) {
+    for (const pending of this.pendingRequests.values()) {
+      clearTimeout(pending.timer);
+      pending.reject(error);
+    }
+    this.pendingRequests.clear();
+
+    for (const [id, meta] of this.pendingClientResponses.entries()) {
+      this.options.send(meta.ws, {
+        channel: "rpc",
+        payload: {
+          type: "response",
+          id,
+          command: meta.responseCommand || "unknown",
+          success: false,
+          error: error.message,
+        },
+      });
+    }
+    this.pendingClientResponses.clear();
+  }
+
+  private handleStdoutChunk(chunk: Buffer | string) {
+    this.stdoutBuffer += typeof chunk === "string" ? chunk : this.decoder.write(chunk);
+
+    while (true) {
+      const newlineIndex = this.stdoutBuffer.indexOf("\n");
+      if (newlineIndex === -1) break;
+
+      let line = this.stdoutBuffer.slice(0, newlineIndex);
+      this.stdoutBuffer = this.stdoutBuffer.slice(newlineIndex + 1);
+      if (line.endsWith("\r")) line = line.slice(0, -1);
+      if (!line.length) continue;
+      this.handleRpcLine(line);
+    }
+  }
+
+  private scheduleSnapshotRefresh(delayMs = 80) {
+    if (this.snapshotRefreshTimer || this.disposed) return;
+
+    this.snapshotRefreshTimer = setTimeout(() => {
+      this.snapshotRefreshTimer = null;
+      this.refreshCachedSnapshot().catch(() => {});
+    }, delayMs);
+  }
+
+  private handleRpcLine(line: string) {
+    let payload: any;
+    try {
+      payload = JSON.parse(line);
+    } catch (error) {
+      this.lastError = `Failed to parse child rpc output: ${line.slice(0, 200)}`;
+      this.options.onEnvelope(this, { channel: "server", event: "parse-error", data: { line, error: String(error) } });
+      this.touch();
+      return;
+    }
+
+    this.touch();
+
+    if (payload.type === "response" && typeof payload.id === "string") {
+      if (payload.success && payload.command === "get_state") {
+        this.rememberState(payload.data);
+      }
+
+      if (payload.success && payload.command === "get_messages") {
+        this.setMessages(payload.data?.messages || []);
+      }
+
+      if (payload.success && payload.command === "get_commands") {
+        this.lastCommands = payload.data?.commands || [];
+        this.options.onStateChange();
+      }
+
+      const pending = this.pendingRequests.get(payload.id);
+      if (pending) {
+        clearTimeout(pending.timer);
+        this.pendingRequests.delete(payload.id);
+        pending.resolve(payload);
+      }
+
+      const clientMeta = this.pendingClientResponses.get(payload.id);
+      if (clientMeta) {
+        this.pendingClientResponses.delete(payload.id);
+
+        try {
+          if (payload.success) clientMeta.onSuccess?.(payload);
+          else clientMeta.onError?.(payload);
+        } catch {
+          // Ignore local response side effects and still forward the payload.
+        }
+
+        const normalizedPayload = payload.success && payload.command === "get_state"
+          ? { ...payload, data: this.lastState || payload.data }
+          : payload;
+
+        const nextPayload = {
+          ...normalizedPayload,
+          ...(clientMeta.responseCommand ? { command: clientMeta.responseCommand } : {}),
+          ...(normalizedPayload.success && clientMeta.responseData
+            ? { data: { ...(normalizedPayload.data || {}), ...clientMeta.responseData } }
+            : {}),
+        };
+        this.options.send(clientMeta.ws, { channel: "rpc", payload: nextPayload });
+      }
+
+      if (payload.success && !payload.data?.cancelled && ["new_session", "switch_session", "set_session_name", "reload"].includes(payload.command)) {
+        this.pendingUiRequest = null;
+        this.liveAssistantMessage = null;
+        this.liveTools.clear();
+        this.scheduleSnapshotRefresh(40);
+      }
+
+      return;
+    }
+
+    if (payload.type === "agent_start") {
+      this.isStreaming = true;
+      this.lastState = this.lastState ? { ...this.lastState, isStreaming: true } : this.lastState;
+      this.options.onStateChange();
+    }
+
+    if (payload.type === "agent_end") {
+      this.isStreaming = false;
+      this.lastState = this.lastState ? { ...this.lastState, isStreaming: false } : this.lastState;
+      this.liveAssistantMessage = null;
+      this.liveTools.clear();
+      this.options.onStateChange();
+      this.scheduleSnapshotRefresh(30);
+    }
+
+    if (payload.type === "message_start" && payload.message?.role === "assistant") {
+      this.liveAssistantMessage = payload.message;
+      this.options.onStateChange();
+    }
+
+    if (payload.type === "message_update" && payload.message?.role === "assistant") {
+      this.liveAssistantMessage = payload.message;
+      this.options.onStateChange();
+    }
+
+    if (payload.type === "message_end" && payload.message?.role === "assistant") {
+      this.liveAssistantMessage = null;
+      this.options.onStateChange();
+    }
+
+    if (payload.type === "tool_execution_start") {
+      this.liveTools.set(payload.toolCallId, {
+        toolCallId: payload.toolCallId,
+        toolName: payload.toolName || "tool",
+        args: payload.args || {},
+        partialResult: null,
+        result: null,
+        isError: false,
+      });
+      this.options.onStateChange();
+    }
+
+    if (payload.type === "tool_execution_update") {
+      const current = this.liveTools.get(payload.toolCallId) || {};
+      this.liveTools.set(payload.toolCallId, {
+        ...current,
+        toolCallId: payload.toolCallId,
+        toolName: payload.toolName || current.toolName || "tool",
+        args: payload.args || current.args || {},
+        partialResult: payload.partialResult || current.partialResult || null,
+        result: current.result || null,
+        isError: current.isError || false,
+      });
+      this.options.onStateChange();
+    }
+
+    if (payload.type === "tool_execution_end") {
+      const current = this.liveTools.get(payload.toolCallId) || {};
+      this.liveTools.set(payload.toolCallId, {
+        ...current,
+        toolCallId: payload.toolCallId,
+        toolName: payload.toolName || current.toolName || "tool",
+        args: payload.args || current.args || {},
+        partialResult: current.partialResult || null,
+        result: payload.result || null,
+        isError: Boolean(payload.isError),
+      });
+      this.options.onStateChange();
+    }
+
+    if (payload.type === "extension_ui_request" && ["select", "confirm", "input", "editor"].includes(payload.method)) {
+      this.pendingUiRequest = payload;
+      this.options.onStateChange();
+    }
+
+    this.options.onEnvelope(this, { channel: "rpc", payload });
+  }
+
+  async request(command: Record<string, unknown>, timeoutMs = 30000) {
+    await this.ensureStarted();
+    if (!this.child) throw new Error("dm rpc child is not running");
+
+    const id = `srv-${++this.requestCounter}`;
+    const payload = { ...command, id };
+
+    return new Promise<any>((resolvePromise, rejectPromise) => {
+      const timer = setTimeout(() => {
+        this.pendingRequests.delete(id);
+        rejectPromise(new Error(`Timed out waiting for child response to ${String(command.type)}`));
+      }, timeoutMs);
+
+      this.pendingRequests.set(id, {
+        resolve: resolvePromise,
+        reject: rejectPromise,
+        timer,
+      });
+
+      this.child!.stdin.write(`${JSON.stringify(payload)}\n`);
+    });
+  }
+
+  async refreshCachedSnapshot(timeoutMs = 4500): Promise<SessionSnapshot> {
+    await this.ensureStarted();
+
+    const [stateResponse, messagesResponse, commandsResponse] = await Promise.all([
+      this.request({ type: "get_state" }, timeoutMs),
+      this.request({ type: "get_messages" }, timeoutMs),
+      this.request({ type: "get_commands" }, timeoutMs),
+    ]);
+
+    if (stateResponse?.success) {
+      this.rememberState(stateResponse.data || null);
+    }
+
+    if (messagesResponse?.success) {
+      this.setMessages(messagesResponse.data?.messages || []);
+    }
+
+    if (commandsResponse?.success) {
+      this.lastCommands = commandsResponse.data?.commands || [];
+      this.options.onStateChange();
+    }
+
+    return this.cachedSnapshot();
+  }
+
+  async getSnapshot(): Promise<SessionSnapshot> {
+    const hasCache = Boolean(this.lastState) || this.lastMessages.length > 0 || this.lastCommands.length > 0;
+
+    if (this.isStreaming || this.pendingUiRequest) {
+      if (hasCache) {
+        return this.cachedSnapshot();
+      }
+
+      try {
+        return await this.refreshCachedSnapshot(2500);
+      } catch {
+        return this.cachedSnapshot();
+      }
+    }
+
+    try {
+      return await this.refreshCachedSnapshot(4500);
+    } catch (error) {
+      if (hasCache || this.pendingUiRequest) {
+        return this.cachedSnapshot();
+      }
+      throw error;
+    }
+  }
+
+  async sendClientCommand(command: Record<string, unknown>, meta?: PendingClientResponse) {
+    await this.ensureStarted();
+    if (!this.child) throw new Error("dm rpc child is not running");
+
+    const nextCommand = { ...command } as Record<string, any>;
+
+    if (nextCommand.type === "extension_ui_response") {
+      this.pendingUiRequest = null;
+      this.options.onStateChange();
+    } else if (!nextCommand.id) {
+      nextCommand.id = `cli-${++this.requestCounter}`;
+    }
+
+    if (nextCommand.type !== "extension_ui_response" && nextCommand.id && meta?.ws) {
+      this.pendingClientResponses.set(String(nextCommand.id), meta);
+    }
+
+    this.touch();
+    this.child.stdin.write(`${JSON.stringify(nextCommand)}\n`);
+    return nextCommand.id as string | undefined;
+  }
+
+  private async stopChildForRestart() {
+    const runningChild = this.child;
+    if (!runningChild) return;
+
+    await new Promise<void>((resolvePromise) => {
+      let settled = false;
+      const finish = () => {
+        if (settled) return;
+        settled = true;
+        clearTimeout(forceKillTimer);
+        resolvePromise();
+      };
+
+      const forceKillTimer = setTimeout(() => {
+        try {
+          runningChild.kill("SIGKILL");
+        } catch {
+          finish();
+        }
+      }, 2000);
+
+      runningChild.once("exit", finish);
+
+      try {
+        runningChild.kill("SIGTERM");
+      } catch {
+        finish();
+      }
+    });
+  }
+
+  async reload() {
+    if (this.reloadPromise) return this.reloadPromise;
+
+    this.reloadPromise = (async () => {
+      await this.ensureStarted();
+
+      const stateResponse = await this.request({ type: "get_state" });
+      if (!stateResponse?.success) {
+        throw new Error(stateResponse?.error || "Failed to read DM state before reload.");
+      }
+
+      const nextState = stateResponse.data || {};
+      this.rememberState(nextState);
+
+      if (nextState.isStreaming) {
+        throw new Error("Wait for the current response to finish before reloading.");
+      }
+
+      if (nextState.isCompacting) {
+        throw new Error("Wait for compaction to finish before reloading.");
+      }
+
+      this.isRestarting = true;
+      this.options.onEnvelope(this, {
+        channel: "server",
+        event: "reloading",
+        data: { message: "Reloading extensions, skills, prompts, and themes…" },
+      });
+
+      await this.stopChildForRestart();
+      await this.ensureStarted({ sessionFile: this.currentSessionFile });
+      await this.refreshCachedSnapshot(5000);
+    })().finally(() => {
+      this.isRestarting = false;
+      this.options.onEnvelope(this, { channel: "server", event: "reloading", data: { message: "" } });
+      this.reloadPromise = null;
+    });
+
+    return this.reloadPromise;
+  }
+
+  async dispose() {
+    this.disposed = true;
+    if (this.snapshotRefreshTimer) {
+      clearTimeout(this.snapshotRefreshTimer);
+      this.snapshotRefreshTimer = null;
+    }
+    this.rejectAllPending(new Error("dm phone session stopped"));
+    await this.stopChildForRestart();
+    this.child = null;
+    this.isStreaming = false;
+    this.pendingUiRequest = null;
+    this.liveAssistantMessage = null;
+    this.liveTools.clear();
+    this.options.onStateChange();
+  }
+}