pi-sessions 0.3.2 → 0.4.1

package/extensions/session-handoff/extract.ts

@@ -1,18 +1,21 @@
 import type { ThinkingLevel } from "@earendil-works/pi-agent-core";
-import {
-  type AssistantMessage,
-  completeSimple,
-  type Message,
-  type SimpleStreamOptions,
-  type TextContent,
-  type ThinkingContent,
-  type Tool,
-  type ToolCall,
+import type {
+  Api,
+  AssistantMessage,
+  Model,
+  TextContent,
+  ThinkingContent,
+  ToolCall,
 } from "@earendil-works/pi-ai";
 import type { ExtensionContext } from "@earendil-works/pi-coding-agent";
 import {
   buildSessionContext,
   convertToLlm,
+  createAgentSession,
+  DefaultResourceLoader,
+  defineTool,
+  getAgentDir,
+  SessionManager,
   serializeConversation,
 } from "@earendil-works/pi-coding-agent";
 import { type Static, Type } from "typebox";
@@ -56,12 +59,7 @@ const HANDOFF_EXTRACTION_PARAMETERS = Type.Object({
   ),
 });
 
-const HANDOFF_EXTRACTION_TOOL: Tool<typeof HANDOFF_EXTRACTION_PARAMETERS> = {
-  name: "create_handoff_context",
-  description: "Extract the structured handoff context for the next session.",
-  parameters: HANDOFF_EXTRACTION_PARAMETERS,
-};
-
+type HandoffExtractionArgs = Static<typeof HANDOFF_EXTRACTION_PARAMETERS>;
 type RequiredHandoffExtractionArgs = Static<typeof REQUIRED_HANDOFF_EXTRACTION_PARAMETERS>;
 
 const REQUIRED_HANDOFF_EXTRACTION_PARAMETERS = Type.Object({
@@ -95,11 +93,7 @@ export async function generateHandoffDraft(
     throw new Error("No model is available for handoff.");
   }
 
-  const auth = await ctx.modelRegistry.getApiKeyAndHeaders(ctx.model);
-  if (!auth.ok || !auth.apiKey) {
-    throw new Error(`No API key is available for ${ctx.model.provider}/${ctx.model.id}.`);
-  }
-
+  const model = ctx.model;
   const sessionContext = buildSessionContext(
     ctx.sessionManager.getEntries(),
     ctx.sessionManager.getLeafId(),
@@ -109,51 +103,18 @@ export async function generateHandoffDraft(
   }
 
   const conversationText = serializeConversation(convertToLlm(sessionContext.messages));
-  const userMessage: Message = {
-    role: "user",
-    content: [
-      {
-        type: "text",
-        text: buildExtractionPrompt(conversationText, goal),
-      },
-    ],
-    timestamp: Date.now(),
-  };
-
-  const requestOptions: SimpleStreamOptions = {
-    apiKey: auth.apiKey,
-    ...(auth.headers ? { headers: auth.headers } : {}),
-    ...(signal ? { signal } : {}),
-    ...(ctx.model.reasoning && thinkingLevel && thinkingLevel !== "off"
-      ? { reasoning: thinkingLevel }
-      : {}),
-  };
-
-  const response = await completeSimple(
-    ctx.model,
-    {
-      systemPrompt: HANDOFF_SYSTEM_PROMPT,
-      messages: [userMessage],
-      tools: [HANDOFF_EXTRACTION_TOOL],
-    },
-    requestOptions,
+  const handoffContext = await runHandoffExtractionAgent(
+    ctx,
+    model,
+    conversationText,
+    goal,
+    thinkingLevel,
+    signal,
   );
-
-  if (response.stopReason === "aborted") {
+  if (!handoffContext) {
     return undefined;
   }
 
-  if (response.stopReason === "error") {
-    throw new Error(response.errorMessage ?? "Handoff generation failed.");
-  }
-
-  const extraction = extractHandoffContext(response, goal);
-  if (!extraction.context) {
-    throw new Error(extraction.error);
-  }
-
-  const handoffContext = extraction.context;
-
   const sessionId = ctx.sessionManager.getSessionId();
   const sessionPath = ctx.sessionManager.getSessionFile();
 
@@ -177,6 +138,84 @@ export function buildExtractionPrompt(conversationText: string, goal: string): s
   ].join("\n");
 }
 
+async function runHandoffExtractionAgent(
+  ctx: ExtensionContext,
+  model: Model<Api>,
+  conversationText: string,
+  goal: string,
+  thinkingLevel: ThinkingLevel | undefined,
+  signal?: AbortSignal,
+): Promise<HandoffContext | undefined> {
+  let capturedArguments: HandoffExtractionArgs | undefined;
+  const createHandoffContextTool = defineTool({
+    name: "create_handoff_context",
+    label: "Create handoff context",
+    description: "Extract the structured handoff context for the next session.",
+    parameters: HANDOFF_EXTRACTION_PARAMETERS,
+    execute: async (_toolCallId, params) => {
+      capturedArguments = params;
+      return {
+        content: [{ type: "text", text: "Handoff context captured. Stopping." }],
+        details: {},
+        terminate: true,
+      };
+    },
+  });
+
+  const resourceLoader = new DefaultResourceLoader({
+    cwd: ctx.cwd,
+    agentDir: getAgentDir(),
+    noExtensions: true,
+    noPromptTemplates: true,
+    noSkills: true,
+    appendSystemPromptOverride: (base) => [...base, HANDOFF_SYSTEM_PROMPT],
+  });
+  await resourceLoader.reload();
+
+  const { session } = await createAgentSession({
+    cwd: ctx.cwd,
+    model,
+    modelRegistry: ctx.modelRegistry,
+    ...(thinkingLevel ? { thinkingLevel } : {}),
+    tools: ["read", "grep", "find", "ls", "create_handoff_context"],
+    customTools: [createHandoffContextTool],
+    resourceLoader,
+    sessionManager: SessionManager.inMemory(ctx.cwd),
+  });
+
+  const abortHandler = (): void => {
+    void session.abort();
+  };
+
+  try {
+    signal?.addEventListener("abort", abortHandler, { once: true });
+    if (signal?.aborted) {
+      await session.abort();
+      return undefined;
+    }
+
+    await session.prompt(buildExtractionPrompt(conversationText, goal));
+  } finally {
+    signal?.removeEventListener("abort", abortHandler);
+    session.dispose();
+  }
+
+  if (signal?.aborted) {
+    return undefined;
+  }
+
+  if (!capturedArguments) {
+    throw new Error("Handoff extraction did not return structured context.");
+  }
+
+  const extraction = extractHandoffContextFromArguments(capturedArguments, goal);
+  if (!extraction.context) {
+    throw new Error(extraction.error);
+  }
+
+  return extraction.context;
+}
+
 export function assembleHandoffDraft(
   sessionId: string,
   sessionPath: string | undefined,
@@ -224,11 +263,18 @@ export function extractHandoffContext(
     return { error: "Handoff extraction did not return structured context." };
   }
 
+  return extractHandoffContextFromArguments(toolCall.arguments, goal);
+}
+
+function extractHandoffContextFromArguments(
+  argumentsValue: unknown,
+  goal: string,
+): { context: HandoffContext; error?: undefined } | { context?: undefined; error: string } {
   let requiredArguments: RequiredHandoffExtractionArgs;
   try {
     requiredArguments = parseTypeBoxValue(
       REQUIRED_HANDOFF_EXTRACTION_PARAMETERS,
-      toolCall.arguments,
+      argumentsValue,
       "Invalid create_handoff_context arguments",
     );
   } catch {
@@ -241,9 +287,9 @@ export function extractHandoffContext(
   }
 
   const summary = normalizeText(requiredArguments.summary);
-  const relevantFiles = normalizeStringArray(toolCall.arguments.relevantFiles, MAX_RELEVANT_FILES);
+  const relevantFiles = getRelevantFiles(argumentsValue);
   const nextTask = normalizeText(requiredArguments.nextTask) || goal.trim();
-  const openQuestions = normalizeStringArray(toolCall.arguments.openQuestions, MAX_OPEN_QUESTIONS);
+  const openQuestions = getOpenQuestions(argumentsValue);
 
   if (!summary || !nextTask || !title) {
     return { error: "Handoff extraction did not return structured context." };
@@ -285,10 +331,30 @@ function normalizeStringArray(value: unknown, limit: number): string[] {
   return [...uniqueValues];
 }
 
+function getRelevantFiles(argumentsValue: unknown): string[] {
+  if (!isRecord(argumentsValue)) {
+    return [];
+  }
+
+  return normalizeStringArray(argumentsValue.relevantFiles, MAX_RELEVANT_FILES);
+}
+
+function getOpenQuestions(argumentsValue: unknown): string[] {
+  if (!isRecord(argumentsValue)) {
+    return [];
+  }
+
+  return normalizeStringArray(argumentsValue.openQuestions, MAX_OPEN_QUESTIONS);
+}
+
 function normalizeText(value: unknown): string {
   return typeof value === "string" ? value.trim() : "";
 }
 
+function isRecord(value: unknown): value is Record<string, unknown> {
+  return typeof value === "object" && value !== null;
+}
+
 function isCreateHandoffContextToolCall(
   content: TextContent | ThinkingContent | ToolCall,
 ): content is ToolCall {

package/extensions/session-search/extract.ts

@@ -1,4 +1,13 @@
-import { existsSync, readdirSync, readFileSync } from "node:fs";
+import {
+  closeSync,
+  existsSync,
+  fstatSync,
+  openSync,
+  readdirSync,
+  readFileSync,
+  readSync,
+  statSync,
+} from "node:fs";
 import path from "node:path";
 import type { AgentMessage } from "@earendil-works/pi-agent-core";
 import type { AssistantMessage, ToolCall, ToolResultMessage } from "@earendil-works/pi-ai";
@@ -34,7 +43,7 @@ export interface SearchTextChunk {
   text: string;
 }
 
-interface DurableHandoffMetadataRecord {
+export interface DurableHandoffMetadataRecord {
   entryId?: string | undefined;
   ts: string;
   metadata: HandoffSessionMetadata;
@@ -70,10 +79,39 @@ export interface ExtractedSessionRecord {
   sessionOrigin?: SessionOrigin | undefined;
   handoffGoal?: string | undefined;
   handoffNextTask?: string | undefined;
+  indexedFileSize: number;
+  indexedFileMtimeMs: number;
+  indexedFileAnchor: string;
   chunks: SearchTextChunk[];
   fileTouches: SessionFileTouch[];
 }
 
+export interface SessionEntryScan {
+  chunks: SearchTextChunk[];
+  fileTouches: SessionFileTouch[];
+  messageCount: number;
+  entryCount: number;
+  maxEntryTs?: string | undefined;
+  firstUserPrompt?: string | undefined;
+  sessionName?: string | undefined;
+  handoffMetadata?: DurableHandoffMetadataRecord | undefined;
+}
+
+export interface SessionTailBaseline {
+  sessionId: string;
+  cwd: string;
+  startedAt: string;
+  indexedFileSize: number;
+  indexedFileAnchor: string;
+}
+
+export interface ExtractedSessionTail {
+  scan: SessionEntryScan;
+  indexedFileSize: number;
+  indexedFileMtimeMs: number;
+  indexedFileAnchor: string;
+}
+
 export interface ParsedSessionFile {
   header: SessionHeader;
   entries: SessionEntry[];
@@ -82,6 +120,8 @@ export interface ParsedSessionFile {
 
 const TOOL_RESULT_TEXT_LIMIT = 500;
 const BASH_OUTPUT_TEXT_LIMIT = 500;
+const FILE_ANCHOR_BYTES = 64;
+const NEWLINE_BYTE = 0x0a;
 const SUMMARY_DETAILS_SCHEMA = Type.Object({
   readFiles: Type.Optional(Type.Array(Type.String())),
   modifiedFiles: Type.Optional(Type.Array(Type.String())),
@@ -115,30 +155,124 @@ export function listSessionFiles(sessionsDir: string): string[] {
 }
 
 export function extractSessionRecord(sessionPath: string): ExtractedSessionRecord | undefined {
-  const parsed = parseSessionFile(sessionPath);
+  const indexedFileMtimeMs = Math.trunc(statSync(sessionPath).mtimeMs);
+  const buffer = readFileSync(sessionPath);
+  const slice = sliceCompleteJsonlPrefix(buffer);
+  const parsed = parseSessionContent(slice.content);
   if (!parsed) return undefined;
 
   const fallbackTs = parsed.header.timestamp;
+  const scan = scanSessionEntries(parsed.entries, fallbackTs, parsed.header.cwd);
+
+  const chunks: SearchTextChunk[] = [];
+  if (parsed.sessionName) {
+    chunks.push(createSessionNameChunk(parsed.sessionName, fallbackTs));
+  }
+  chunks.push(...scan.chunks);
+  appendDurableHandoffMetadataChunks(chunks, scan.handoffMetadata);
+
+  const parentSessionPath = normalizeParentSessionPath(parsed.header.parentSession);
+
+  return {
+    sessionId: parsed.header.id,
+    sessionPath,
+    sessionName: parsed.sessionName,
+    firstUserPrompt: trimmedText(scan.firstUserPrompt),
+    cwd: parsed.header.cwd,
+    repoRoots: deriveSessionRepoRoots(parsed.header.cwd, scan.fileTouches),
+    startedAt: fallbackTs,
+    modifiedAt: maxTimestamp(fallbackTs, scan.maxEntryTs),
+    messageCount: scan.messageCount,
+    entryCount: scan.entryCount,
+    parentSessionPath,
+    parentSessionId: parentSessionPath ? readSessionIdFromPath(parentSessionPath) : undefined,
+    sessionOrigin: inferSessionOrigin(parentSessionPath, scan.handoffMetadata?.metadata),
+    handoffGoal: scan.handoffMetadata?.metadata.goal,
+    handoffNextTask: scan.handoffMetadata?.metadata.nextTask,
+    indexedFileSize: slice.consumedBytes,
+    indexedFileMtimeMs,
+    indexedFileAnchor: buildFileAnchor(buffer, slice.consumedBytes),
+    chunks: chunks,
+    fileTouches: scan.fileTouches,
+  };
+}
+
+// Reads only the bytes appended since the last sync. The stored anchor (the
+// final bytes of the previously indexed prefix) proves the prefix is unchanged;
+// any in-place rewrite invalidates it and the caller falls back to a full
+// re-extract. Session files are append-only in normal pi operation.
+export function extractSessionTail(
+  sessionPath: string,
+  baseline: SessionTailBaseline,
+): ExtractedSessionTail | undefined {
+  const anchorBytes = Buffer.from(baseline.indexedFileAnchor, "base64");
+  const anchorStart = baseline.indexedFileSize - anchorBytes.length;
+  if (anchorBytes.length === 0 || anchorStart < 0) {
+    return undefined;
+  }
+
+  const indexedFileMtimeMs = Math.trunc(statSync(sessionPath).mtimeMs);
+  const fd = openSync(sessionPath, "r");
+  try {
+    const fileSize = fstatSync(fd).size;
+    if (fileSize < baseline.indexedFileSize) {
+      return undefined;
+    }
+
+    const window = Buffer.alloc(fileSize - anchorStart);
+    const bytesRead = readSync(fd, window, 0, window.length, anchorStart);
+    if (
+      bytesRead < anchorBytes.length ||
+      !window.subarray(0, anchorBytes.length).equals(anchorBytes)
+    ) {
+      return undefined;
+    }
+
+    const tailBuffer = window.subarray(anchorBytes.length, bytesRead);
+    const slice = sliceCompleteJsonlPrefix(tailBuffer);
+    const entries = parseSessionEntries(slice.content).filter(isSessionEntry);
+    const consumedWindow = Buffer.concat([
+      anchorBytes,
+      tailBuffer.subarray(0, slice.consumedBytes),
+    ]);
+
+    return {
+      scan: scanSessionEntries(entries, baseline.startedAt, baseline.cwd),
+      indexedFileSize: baseline.indexedFileSize + slice.consumedBytes,
+      indexedFileMtimeMs,
+      indexedFileAnchor: buildFileAnchor(consumedWindow, consumedWindow.length),
+    };
+  } finally {
+    closeSync(fd);
+  }
+}
+
+export function createSessionNameChunk(sessionName: string, ts: string): SearchTextChunk {
+  return {
+    entryType: "session_info",
+    ts,
+    sourceKind: "session_name",
+    text: sessionName,
+  };
+}
+
+export function scanSessionEntries(
+  entries: SessionEntry[],
+  fallbackTs: string,
+  cwd: string,
+): SessionEntryScan {
   const chunks: SearchTextChunk[] = [];
   const fileTouches: SessionFileTouch[] = [];
-  let modifiedAt = fallbackTs;
   let messageCount = 0;
+  let maxEntryTs: string | undefined;
   let firstUserPrompt: string | undefined;
-  let durableHandoffMetadata: DurableHandoffMetadataRecord | undefined;
+  let sessionName: string | undefined;
+  let handoffMetadata: DurableHandoffMetadataRecord | undefined;
 
-  if (parsed.sessionName) {
-    chunks.push({
-      entryType: "session_info",
-      ts: parsed.header.timestamp,
-      sourceKind: "session_name",
-      text: parsed.sessionName,
-    });
-  }
-
-  for (const entry of parsed.entries) {
+  for (const entry of entries) {
     const entryTs = getEntryTimestamp(entry, fallbackTs);
-    if (entryTs > modifiedAt) {
-      modifiedAt = entryTs;
+    if (maxEntryTs === undefined || entryTs > maxEntryTs) {
+      maxEntryTs = entryTs;
     }
 
     switch (entry.type) {
@@ -150,15 +284,19 @@ export function extractSessionRecord(sessionPath: string): ExtractedSessionRecor
           firstUserPrompt = contentToText(message.content);
         }
         chunks.push(...extractMessageChunks(entry.id, entryTs, message));
-        fileTouches.push(
-          ...extractMessageFileTouches(entry.id, entryTs, message, parsed.header.cwd),
-        );
+        fileTouches.push(...extractMessageFileTouches(entry.id, entryTs, message, cwd));
+        continue;
+      }
+      case "session_info": {
+        if (typeof entry.name === "string") {
+          sessionName = entry.name.trim();
+        }
         continue;
       }
       case "custom": {
         const nextHandoffMetadata = extractHandoffMetadata(entry, entryTs);
-        if (nextHandoffMetadata && !durableHandoffMetadata) {
-          durableHandoffMetadata = nextHandoffMetadata;
+        if (nextHandoffMetadata && !handoffMetadata) {
+          handoffMetadata = nextHandoffMetadata;
         }
         continue;
       }
@@ -179,7 +317,7 @@ export function extractSessionRecord(sessionPath: string): ExtractedSessionRecor
             entryTs,
             "branch_summary_details",
             entry.details,
-            parsed.header.cwd,
+            cwd,
           ),
         );
         continue;
@@ -193,13 +331,7 @@ export function extractSessionRecord(sessionPath: string): ExtractedSessionRecor
           trimmedText(entry.summary),
         );
         fileTouches.push(
-          ...extractDetailFileTouches(
-            entry.id,
-            entryTs,
-            "compaction_details",
-            entry.details,
-            parsed.header.cwd,
-          ),
+          ...extractDetailFileTouches(entry.id, entryTs, "compaction_details", entry.details, cwd),
         );
         continue;
       }
@@ -208,33 +340,69 @@ export function extractSessionRecord(sessionPath: string): ExtractedSessionRecor
     }
   }
 
-  appendDurableHandoffMetadataChunks(chunks, durableHandoffMetadata);
-
-  const parentSessionPath = normalizeParentSessionPath(parsed.header.parentSession);
-
   return {
-    sessionId: parsed.header.id,
-    sessionPath,
-    sessionName: parsed.sessionName,
-    firstUserPrompt: trimmedText(firstUserPrompt),
-    cwd: parsed.header.cwd,
-    repoRoots: deriveSessionRepoRoots(parsed.header.cwd, fileTouches),
-    startedAt: parsed.header.timestamp,
-    modifiedAt,
-    messageCount,
-    entryCount: parsed.entries.length,
-    parentSessionPath,
-    parentSessionId: parentSessionPath ? readSessionIdFromPath(parentSessionPath) : undefined,
-    sessionOrigin: inferSessionOrigin(parentSessionPath, durableHandoffMetadata?.metadata),
-    handoffGoal: durableHandoffMetadata?.metadata.goal,
-    handoffNextTask: durableHandoffMetadata?.metadata.nextTask,
     chunks,
     fileTouches,
+    messageCount,
+    entryCount: entries.length,
+    maxEntryTs,
+    firstUserPrompt,
+    sessionName,
+    handoffMetadata,
   };
 }
 
+function maxTimestamp(fallbackTs: string, entryTs: string | undefined): string {
+  return entryTs !== undefined && entryTs > fallbackTs ? entryTs : fallbackTs;
+}
+
+// Consumes only whole JSONL lines so the recorded byte offset always lands on
+// an entry boundary. A final unterminated line is consumed only if it parses as
+// complete JSON: a strict prefix of a JSON object can never parse, so a torn
+// in-progress write is reliably excluded.
+function sliceCompleteJsonlPrefix(buffer: Buffer): { content: string; consumedBytes: number } {
+  if (buffer.length === 0) {
+    return { content: "", consumedBytes: 0 };
+  }
+
+  if (buffer[buffer.length - 1] === NEWLINE_BYTE) {
+    return { content: buffer.toString("utf8"), consumedBytes: buffer.length };
+  }
+
+  const lastNewline = buffer.lastIndexOf(NEWLINE_BYTE);
+  const finalLine = buffer.subarray(lastNewline + 1).toString("utf8");
+  if (isCompleteJsonLine(finalLine)) {
+    return { content: buffer.toString("utf8"), consumedBytes: buffer.length };
+  }
+
+  const consumedBytes = lastNewline + 1;
+  return { content: buffer.subarray(0, consumedBytes).toString("utf8"), consumedBytes };
+}
+
+function isCompleteJsonLine(line: string): boolean {
+  const trimmed = line.trim();
+  if (!trimmed) {
+    return true;
+  }
+
+  try {
+    JSON.parse(trimmed);
+    return true;
+  } catch {
+    return false;
+  }
+}
+
+function buildFileAnchor(buffer: Buffer, end: number): string {
+  const start = Math.max(0, end - FILE_ANCHOR_BYTES);
+  return buffer.subarray(start, end).toString("base64");
+}
+
 export function parseSessionFile(sessionPath: string): ParsedSessionFile | undefined {
-  const raw = readFileSync(sessionPath, "utf8");
+  return parseSessionContent(readFileSync(sessionPath, "utf8"));
+}
+
+export function parseSessionContent(raw: string): ParsedSessionFile | undefined {
   const fileEntries = parseSessionEntries(raw);
   if (fileEntries.length === 0) {
     return undefined;