@groeponline/pi-wishcraft 0.17.3

package/src/extension/segment-context.ts

@@ -0,0 +1,212 @@
+import type {
+  ReadonlyFooterDataProvider,
+  Theme,
+} from "@earendil-works/pi-coding-agent";
+
+import type { SegmentContext, ColorScheme } from "../config/types.ts";
+import { getPreset } from "../config/presets.ts";
+import {
+  collectHiddenExtensionStatusKeys,
+  mergeSegmentOptions,
+} from "../config/powerline-config.ts";
+import { getDefaultColors } from "../theme/theme.ts";
+import { getGitStatus } from "../git/status.ts";
+import { getQueueContext } from "./queue-context.ts";
+import { getUsageTokenTotal } from "../usage/ledger.ts";
+import {
+  CUSTOM_COMPACTION_STATUS_KEY,
+  EDITOR_STATUS_DEFER_MS,
+} from "./constants.ts";
+import {
+  config,
+  customCompactionEnabled,
+} from "./state.ts";
+import type { RuntimeState } from "./types.ts";
+
+/**
+ * Reset the cached layout + underlying session-scan caches. Called on
+ * session boundaries and whenever config that affects layout changes.
+ */
+export function resetLayoutCache(rt: RuntimeState): void {
+  rt.lastLayoutResult = null;
+  rt.layoutDirty = true;
+  rt.sessionBranchCache.reset();
+  rt.tokenStatsCache.reset();
+  rt.coreContextUsageCache.reset();
+}
+
+export function requestStatusRender(rt: RuntimeState, delayMs?: number): void {
+  rt.layoutDirty = true;
+  rt.statusRenderScheduler.schedule(delayMs);
+}
+
+export function requestImmediateStatusRender(
+  rt: RuntimeState,
+  options: { deferDuringTyping?: boolean } = {},
+): void {
+  rt.layoutDirty = true;
+  if (
+    options.deferDuringTyping !== false &&
+    Date.now() - rt.lastEditorInputAt < EDITOR_STATUS_DEFER_MS
+  ) {
+    rt.statusRenderScheduler.schedule();
+    return;
+  }
+
+  rt.forceNextLayoutRecompute = true;
+  rt.statusRenderScheduler.cancel();
+  rt.statusRenderScheduler.schedule(0);
+}
+
+export function installFooterStatusRepaintHook(
+  rt: RuntimeState,
+  footerData: ReadonlyFooterDataProvider,
+): void {
+  rt.restoreFooterStatusRepaintHook?.();
+  rt.restoreFooterStatusRepaintHook = null;
+
+  const writableFooterData = footerData as ReadonlyFooterDataProvider & {
+    setExtensionStatus?: (key: string, text: string | undefined) => void;
+    clearExtensionStatuses?: () => void;
+  };
+  if (typeof writableFooterData.setExtensionStatus !== "function") return;
+
+  const originalSetExtensionStatus = writableFooterData.setExtensionStatus;
+  const originalClearExtensionStatuses =
+    writableFooterData.clearExtensionStatuses;
+  const setExtensionStatusAndRepaint = function setExtensionStatusAndRepaint(
+    this: unknown,
+    key: string,
+    text: string | undefined,
+  ) {
+    originalSetExtensionStatus.call(this, key, text);
+    requestImmediateStatusRender(rt);
+  };
+  writableFooterData.setExtensionStatus = setExtensionStatusAndRepaint;
+
+  let clearExtensionStatusesAndRepaint: (() => void) | null = null;
+  if (typeof originalClearExtensionStatuses === "function") {
+    clearExtensionStatusesAndRepaint =
+      function clearExtensionStatusesAndRepaint(this: unknown) {
+        originalClearExtensionStatuses.call(this);
+        requestImmediateStatusRender(rt);
+      };
+    writableFooterData.clearExtensionStatuses =
+      clearExtensionStatusesAndRepaint;
+  }
+
+  rt.restoreFooterStatusRepaintHook = () => {
+    if (
+      writableFooterData.setExtensionStatus === setExtensionStatusAndRepaint
+    ) {
+      writableFooterData.setExtensionStatus = originalSetExtensionStatus;
+    }
+    if (
+      clearExtensionStatusesAndRepaint &&
+      writableFooterData.clearExtensionStatuses ===
+        clearExtensionStatusesAndRepaint
+    ) {
+      writableFooterData.clearExtensionStatuses =
+        originalClearExtensionStatuses;
+    }
+  };
+}
+
+export function buildSegmentContext(
+  rt: RuntimeState,
+  ctx: any,
+  theme: Theme,
+): SegmentContext {
+  const presetDef = getPreset(config.preset);
+  const colors: ColorScheme = presetDef.colors ?? getDefaultColors();
+
+  // Build usage stats and get thinking level from session (cached; the full
+  // event list is only re-scanned when events are appended or the trailing
+  // event's stats-relevant fields change, e.g. in-place streaming updates)
+  const sessionEvents = rt.sessionBranchCache.get(ctx.sessionManager);
+  const tokenStats = rt.tokenStatsCache.get(sessionEvents);
+  const { input, output, cacheRead, cacheWrite, cost, subagentCost } =
+    tokenStats;
+  const lastAssistant = tokenStats.lastAssistant;
+  const thinkingLevelFromSession = tokenStats.thinkingLevelFromSession;
+
+  // Calculate context percentage.
+  const latestUsage = rt.isStreaming
+    ? (rt.liveAssistantUsage ?? lastAssistant?.usage)
+    : lastAssistant?.usage;
+  const coreContextUsage =
+    rt.isStreaming && rt.liveAssistantUsage
+      ? null
+      : rt.coreContextUsageCache.get(ctx);
+  const contextTokens =
+    coreContextUsage?.contextTokens ??
+    (latestUsage ? getUsageTokenTotal(latestUsage) : 0);
+  const contextWindow =
+    coreContextUsage?.contextWindow ?? ctx.model?.contextWindow ?? 0;
+  const contextPercent =
+    coreContextUsage?.contextPercent ??
+    (contextWindow > 0 ? (contextTokens / contextWindow) * 100 : 0);
+
+  const segmentOptions = mergeSegmentOptions(
+    presetDef.segmentOptions,
+    config.segmentOptions,
+  );
+
+  // Get git status (cached)
+  const gitBranch = rt.footerDataRef?.getGitBranch() ?? null;
+  const gitStatus = getGitStatus(gitBranch, segmentOptions.git?.polling);
+  const extensionStatuses =
+    rt.footerDataRef?.getExtensionStatuses() ?? new Map();
+  const customItemsById = new Map(
+    config.customItems.map((item) => [item.id, item]),
+  );
+  const hiddenExtensionStatusKeys = collectHiddenExtensionStatusKeys(
+    config.customItems,
+  );
+
+  // Check if using OAuth subscription
+  const usingSubscription = ctx.model
+    ? (ctx.modelRegistry?.isUsingOAuth?.(ctx.model) ?? false)
+    : false;
+
+  const thinkingLevel =
+    rt.currentThinkingLevel ??
+    thinkingLevelFromSession ??
+    rt.getThinkingLevelFn?.() ??
+    "off";
+  const queueSummary = rt.queueStore.summarize(
+    getQueueContext(ctx),
+    rt.powerlineCompacting,
+  );
+
+  return {
+    model: ctx.model,
+    thinkingLevel,
+    sessionId: ctx.sessionManager?.getSessionId?.(),
+    cwd: ctx.cwd,
+    usageStats: { input, output, cacheRead, cacheWrite, cost, subagentCost },
+    contextTokens,
+    contextPercent,
+    contextWindow,
+    autoCompactEnabled:
+      ctx.settingsManager?.getCompactionSettings?.()?.enabled ?? true,
+    customCompactionEnabled:
+      customCompactionEnabled ||
+      extensionStatuses.has(CUSTOM_COMPACTION_STATUS_KEY),
+    usingSubscription,
+    queueSummary,
+    sessionStartTime: rt.sessionStartTime,
+    shellModeActive: rt.bashModeActive,
+    shellRunning: rt.shellSession?.state.running ?? false,
+    shellName: rt.shellSession?.state.shellName ?? null,
+    shellCwd: rt.shellSession?.state.cwd ?? null,
+    git: gitStatus,
+    extensionStatuses,
+    hiddenExtensionStatusKeys,
+    customItemsById,
+    options: segmentOptions,
+    segmentLabels: new Map(Object.entries(config.segmentLabels)),
+    theme,
+    colors,
+  };
+}

package/src/extension/session-lifecycle.ts

@@ -0,0 +1,373 @@
+import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";
+
+import { BashTranscriptStore } from "../../bash-mode/transcript.ts";
+import { BashCompletionEngine } from "../../bash-mode/completion.ts";
+import { parsePowerlineConfig } from "../config/powerline-config.ts";
+import { registerCustomSegments } from "../segments/index.ts";
+import { registerCustomPresets } from "../config/presets.ts";
+import { invalidateGitStatus } from "../git/status.ts";
+import { invalidateGitForCommand } from "./git-invalidation.ts";
+import {
+  initVibeManager,
+  onVibeAgentEnd,
+  onVibeAgentStart,
+  onVibeBeforeAgentStart,
+  onVibeToolCall,
+} from "../working-vibes/index.ts";
+import {
+  getUsageTokenTotal,
+  hasSessionAssistantUsage,
+  isSessionAssistantMessage,
+} from "../usage/ledger.ts";
+import {
+  detectCustomCompactionEnabled,
+  readSettings,
+} from "./settings-io.ts";
+import {
+  resolveShortcutConfig,
+  parseBashModeSettings,
+} from "./shortcuts-config.ts";
+import { warnInvalidSegmentSettings } from "./layout.ts";
+import { readPersistedStashHistory } from "./stash-history.ts";
+import {
+  requestImmediateStatusRender,
+  requestStatusRender,
+  resetLayoutCache,
+} from "./segment-context.ts";
+import { getQueueContext } from "./queue-context.ts";
+import {
+  requestQueueRender,
+  schedulePostCompactionDelivery,
+  finishFailedCompaction,
+} from "./queue-integration.ts";
+import { setupCustomEditor } from "./custom-editor.ts";
+import {
+  setupWelcomeHeader,
+  setupWelcomeOverlay,
+} from "./welcome-integration.ts";
+import {
+  config,
+  PRESET_NAMES,
+  setConfig,
+  setCustomCompactionEnabled,
+} from "./state.ts";
+import { CONTEXT_STATUS_RENDER_MS } from "./constants.ts";
+import type { RuntimeState } from "./types.ts";
+import { isStaleExtensionContextError } from "./stale-context.ts";
+import {
+  dismissWelcome,
+  scheduleDismissWelcome,
+} from "./welcome-control.ts";
+
+// Helper to extract recent agent response text (skipping thinking blocks)
+function getRecentAgentContext(ctx: any): string | undefined {
+  const sessionEvents = ctx.sessionManager?.getBranch?.() ?? [];
+
+  // Find the most recent assistant message
+  for (let i = sessionEvents.length - 1; i >= 0; i--) {
+    const e = sessionEvents[i];
+    if (e.type === "message" && e.message?.role === "assistant") {
+      const content = e.message.content;
+      if (!Array.isArray(content)) continue;
+
+      // Extract text content, skip thinking blocks
+      for (const block of content) {
+        if (block.type === "text" && block.text) {
+          // Return first ~200 chars of non-empty text
+          const text = block.text.trim();
+          if (text.length > 0) {
+            return text.slice(0, 200);
+          }
+        }
+      }
+    }
+  }
+  return undefined;
+}
+
+export function shouldShowStartupWelcome(reason: unknown, welcomeEnabled: boolean): boolean {
+  return reason === "startup" && welcomeEnabled;
+}
+
+export function registerSessionLifecycle(
+  pi: ExtensionAPI,
+  rt: RuntimeState,
+): void {
+  // Track session start
+  pi.on("session_start", async (event, ctx) => {
+    rt.shellSession?.dispose();
+    rt.shellSession = null;
+    rt.sessionGeneration++;
+    rt.sessionStartTime = Date.now();
+    rt.currentCtx = ctx;
+    setCustomCompactionEnabled(detectCustomCompactionEnabled(ctx.cwd));
+    rt.lastUserPrompt = "";
+    rt.isStreaming = false;
+    rt.liveAssistantUsage = null;
+    rt.powerlineCompacting = false;
+    rt.deliverAfterRetrySettles = false;
+    rt.stashedEditorText = null;
+
+    const settings = readSettings(ctx.cwd);
+    rt.resolvedShortcuts = resolveShortcutConfig(settings);
+    rt.bashModeSettings = parseBashModeSettings(settings, rt.resolvedShortcuts);
+    rt.showLastPrompt = settings.showLastPrompt !== false;
+    setConfig(parsePowerlineConfig(settings.powerline, PRESET_NAMES));
+    registerCustomSegments(config.segments);
+    registerCustomPresets(config.presets);
+    warnInvalidSegmentSettings(ctx);
+    rt.stashedPromptHistory = readPersistedStashHistory();
+    rt.bashModeActive = false;
+    rt.bashTranscript = new BashTranscriptStore(rt.bashModeSettings);
+    rt.bashCompletionEngine = new BashCompletionEngine();
+
+    rt.getThinkingLevelFn = () => ctx.thinkingLevel ?? "off";
+    rt.currentThinkingLevel = rt.getThinkingLevelFn();
+
+    if (ctx.hasUI) {
+      ctx.ui.setStatus("stash", undefined);
+      const pendingIdeas = rt.queueStore
+        .activeItems(getQueueContext(ctx))
+        .filter((item) => item.intent === "idea").length;
+      if (pendingIdeas > 0) {
+        ctx.ui.notify(
+          `${pendingIdeas} idea${pendingIdeas === 1 ? "" : "s"} waiting — /ideas`,
+          "info",
+        );
+      }
+    }
+
+    // Initialize vibe manager (needs modelRegistry from ctx)
+    initVibeManager(ctx);
+
+    if (rt.enabled && ctx.hasUI) {
+      setupCustomEditor(pi, rt, ctx);
+      if (shouldShowStartupWelcome(event.reason, config.welcome)) {
+        if (settings.quietStartup === true) {
+          setupWelcomeHeader(rt, ctx);
+        } else {
+          setupWelcomeOverlay(rt, ctx);
+        }
+      } else {
+        dismissWelcome(rt, ctx);
+      }
+    }
+  });
+
+  pi.on("session_shutdown", async (_event, ctx) => {
+    rt.sessionGeneration++;
+    rt.dismissWelcomeOverlay?.();
+    rt.dismissWelcomeOverlay = null;
+    rt.welcomeHeaderActive = false;
+    rt.welcomeOverlayShouldDismiss = false;
+    rt.welcomeDismissScheduler.cancel();
+    rt.statusRenderScheduler.cancel();
+    rt.restoreFooterStatusRepaintHook?.();
+    rt.restoreFooterStatusRepaintHook = null;
+    rt.stashShortcutInputUnsubscribe?.();
+    rt.stashShortcutInputUnsubscribe = null;
+    rt.shellSession?.dispose();
+    rt.shellSession = null;
+    if (rt.queueDeliveryTimer) {
+      clearTimeout(rt.queueDeliveryTimer);
+      rt.queueDeliveryTimer = null;
+    }
+    rt.powerlineCompacting = false;
+    rt.deliverAfterRetrySettles = false;
+    rt.bashModeActive = false;
+    rt.currentCtx = null;
+    rt.footerDataRef = null;
+    rt.getThinkingLevelFn = null;
+    rt.currentThinkingLevel = null;
+    rt.liveAssistantUsage = null;
+    rt.tuiRef = null;
+    rt.currentEditor = null;
+    resetLayoutCache(rt);
+  });
+
+  // Invalidate git status on file changes, trigger re-render on potential branch changes
+  pi.on("tool_result", async (event) => {
+    if (event.toolName === "write" || event.toolName === "edit") {
+      invalidateGitStatus();
+      requestStatusRender(rt);
+    }
+    // Check for bash commands that might change git branch
+    if (event.toolName === "bash" && event.input?.command) {
+      invalidateGitForCommand(rt, String(event.input.command));
+    }
+  });
+
+  // Also catch user escape commands (! prefix)
+  // Note: This fires BEFORE execution, so we use a longer delay and multiple re-renders
+  // to ensure we catch the update after the command completes.
+  pi.on("user_bash", async (event) => {
+    invalidateGitForCommand(rt, event.command, { stagger: true });
+  });
+
+  pi.on("model_select", async (_event, ctx) => {
+    rt.currentCtx = ctx;
+    rt.coreContextUsageCache.reset();
+    requestStatusRender(rt);
+  });
+
+  pi.on("thinking_level_select", async (event, ctx) => {
+    rt.currentCtx = ctx;
+    rt.currentThinkingLevel =
+      rt.getThinkingLevelFn?.() ??
+      (typeof event.level === "string" ? event.level : null);
+    requestImmediateStatusRender(rt, { deferDuringTyping: false });
+  });
+
+  pi.on("session_tree", async (_event, ctx) => {
+    rt.currentCtx = ctx;
+    rt.currentThinkingLevel = null;
+    rt.liveAssistantUsage = null;
+    requestImmediateStatusRender(rt, { deferDuringTyping: false });
+  });
+
+  // Generate themed working message before agent starts (has access to user's prompt)
+  pi.on("before_agent_start", async (event, ctx) => {
+    rt.lastUserPrompt = event.prompt;
+    if (ctx.hasUI) {
+      onVibeBeforeAgentStart(event.prompt, ctx.ui.setWorkingMessage);
+    }
+  });
+
+  // Track streaming state (footer only shows status during streaming)
+  // Also dismiss welcome when agent starts responding (handles `p "command"` case)
+  pi.on("agent_start", async (_event, ctx) => {
+    rt.isStreaming = true;
+    rt.liveAssistantUsage = null;
+    onVibeAgentStart();
+    dismissWelcome(rt, ctx);
+    rt.currentCtx = ctx;
+  });
+
+  pi.on("message_update", async (event, ctx) => {
+    if (
+      isSessionAssistantMessage(event.message) &&
+      event.message.stopReason !== "error" &&
+      event.message.stopReason !== "aborted" &&
+      getUsageTokenTotal(event.message.usage) > 0
+    ) {
+      rt.liveAssistantUsage = event.message.usage;
+      rt.currentCtx = ctx;
+      rt.layoutDirty = true;
+      rt.statusRenderScheduler.schedule(CONTEXT_STATUS_RENDER_MS);
+    }
+  });
+
+  pi.on("message_end", async (event, ctx) => {
+    rt.currentCtx = ctx;
+    rt.coreContextUsageCache.reset();
+    if (isSessionAssistantMessage(event.message)) {
+      if (
+        event.message.stopReason === "error" ||
+        event.message.stopReason === "aborted"
+      ) {
+        rt.liveAssistantUsage = null;
+      } else if (getUsageTokenTotal(event.message.usage) > 0) {
+        rt.liveAssistantUsage = event.message.usage;
+      }
+    }
+    requestImmediateStatusRender(rt, { deferDuringTyping: false });
+  });
+
+  pi.on("turn_end", async (_event, ctx) => {
+    rt.currentCtx = ctx;
+    rt.coreContextUsageCache.reset();
+    requestImmediateStatusRender(rt, { deferDuringTyping: false });
+  });
+
+  pi.on("session_before_compact", async (_event, ctx) => {
+    rt.powerlineCompacting = true;
+    rt.currentCtx = ctx;
+    requestQueueRender(rt);
+  });
+
+  pi.on("session_compact", async (event, ctx) => {
+    rt.powerlineCompacting = false;
+    rt.currentCtx = ctx;
+    // Compaction rewrites the conversation, so the cached context-usage (tokens/window/percent)
+    // is stale. Reset it and force a redraw — otherwise the bar keeps showing the pre-compact fill.
+    rt.coreContextUsageCache.reset();
+    requestImmediateStatusRender(rt, { deferDuringTyping: false });
+    if (event.willRetry) {
+      rt.deliverAfterRetrySettles = true;
+    } else {
+      rt.deliverAfterRetrySettles = false;
+      schedulePostCompactionDelivery(pi, rt, ctx);
+    }
+    requestQueueRender(rt);
+  });
+
+  pi.on("agent_settled", async (_event, ctx) => {
+    if (rt.powerlineCompacting) {
+      finishFailedCompaction(rt, ctx, "Compaction did not complete");
+      return;
+    }
+    if (rt.deliverAfterRetrySettles) {
+      rt.deliverAfterRetrySettles = false;
+      schedulePostCompactionDelivery(pi, rt, ctx);
+    }
+  });
+
+  // Also dismiss on tool calls (agent is working) + refresh vibe if rate limit allows
+  pi.on("tool_call", async (event, ctx) => {
+    dismissWelcome(rt, ctx);
+    if (ctx.hasUI) {
+      // Extract recent agent context from session for richer vibe generation
+      const agentContext = getRecentAgentContext(ctx);
+      onVibeToolCall(
+        event.toolName,
+        event.input,
+        ctx.ui.setWorkingMessage,
+        agentContext,
+      );
+    }
+  });
+
+  pi.on("agent_end", async (_event, ctx) => {
+    rt.isStreaming = false;
+    rt.liveAssistantUsage = null;
+    rt.coreContextUsageCache.reset();
+
+    let hasUI = false;
+    try {
+      hasUI = Boolean(ctx.hasUI);
+    } catch (error) {
+      if (!isStaleExtensionContextError(error)) throw error;
+      rt.currentCtx = null;
+      return;
+    }
+
+    rt.currentCtx = ctx;
+    try {
+      if (hasUI) {
+        onVibeAgentEnd(ctx.ui.setWorkingMessage); // working-vibes internal state + reset message
+        if (rt.stashedEditorText !== null) {
+          if (ctx.ui.getEditorText().trim() === "") {
+            ctx.ui.setEditorText(rt.stashedEditorText);
+            rt.stashedEditorText = null;
+            ctx.ui.setStatus("stash", undefined);
+            ctx.ui.notify("Stash restored", "info");
+          } else {
+            ctx.ui.notify(
+              "Stash preserved — clear editor then Alt+S to restore",
+              "info",
+            );
+          }
+        }
+      }
+    } catch (error) {
+      if (!isStaleExtensionContextError(error)) throw error;
+      rt.currentCtx = null;
+      return;
+    }
+
+    requestStatusRender(rt);
+    if (!rt.powerlineCompacting && !rt.deliverAfterRetrySettles) {
+      schedulePostCompactionDelivery(pi, rt, ctx);
+    }
+  });
+}