ds4-context-engine 0.3.0-alpha.4 → 0.3.0-beta.1

package/src/pi-adapter/compaction-adapter.ts

@@ -19,14 +19,22 @@ interface SourceCandidate {
   used: boolean;
 }
 
-export interface PreparedCompactionSource {
-  messages: SessionBeforeCompactEvent["preparation"]["messagesToSummarize"];
+type CompactionMessage = SessionBeforeCompactEvent["preparation"]["messagesToSummarize"][number];
+
+export interface PreparedCompactionSourceSlice {
+  messages: CompactionMessage[];
   conversationText: string;
   sourceText: string;
   sourceEntryIds: string[];
   sourceHash: string;
   segmentReadFiles: string[];
   segmentModifiedFiles: string[];
+  isSplitTurn: boolean;
+}
+
+export interface PreparedCompactionSource extends PreparedCompactionSourceSlice {
+  messageEntryIds: string[];
+  turnPrefixStartIndex: number;
   readFiles: string[];
   modifiedFiles: string[];
   previousSummary?: string;
@@ -53,7 +61,7 @@ function mapSourceEntryIds(messages: readonly unknown[], entries: readonly Sessi
     const match = available.find((candidate) => !candidate.used && candidate.fingerprint === fingerprint(message));
     if (!match) throw new Error("Compaction source message has no exact canonical Pi session entry");
     match.used = true;
-    if (!ids.includes(match.entryId)) ids.push(match.entryId);
+    ids.push(match.entryId);
   }
 
   return ids;
@@ -69,8 +77,93 @@ function previousFiles(entries: readonly SessionEntry[]): { read: string[]; modi
   return { read: [], modified: [] };
 }
 
+function unique(values: Iterable<string>): string[] {
+  return [...new Set(values)];
+}
+
 function sortedUnique(values: Iterable<string>): string[] {
-  return [...new Set(values)].sort((left, right) => left.localeCompare(right));
+  return unique(values).sort((left, right) => left.localeCompare(right));
+}
+
+function segmentFiles(messages: readonly CompactionMessage[]): { readFiles: string[]; modifiedFiles: string[] } {
+  const read = new Set<string>();
+  const modified = new Set<string>();
+  for (const message of messages) {
+    if (message.role !== "assistant" || !Array.isArray(message.content)) continue;
+    for (const block of message.content) {
+      if (block.type !== "toolCall" || !block.arguments || typeof block.arguments !== "object") continue;
+      const path = "path" in block.arguments && typeof block.arguments.path === "string"
+        ? block.arguments.path
+        : undefined;
+      if (!path) continue;
+      if (block.name === "read") read.add(path);
+      if (block.name === "write" || block.name === "edit") modified.add(path);
+    }
+  }
+  return {
+    readFiles: sortedUnique([...read].filter((path) => !modified.has(path))),
+    modifiedFiles: sortedUnique(modified),
+  };
+}
+
+function createSourceSlice(input: {
+  messages: CompactionMessage[];
+  sourceEntryIds: string[];
+  segmentReadFiles: string[];
+  segmentModifiedFiles: string[];
+  isSplitTurn: boolean;
+}): PreparedCompactionSourceSlice {
+  const conversationText = serializeConversation(convertToLlm(input.messages));
+  const canonicalSourceText = stableStringify(input.messages);
+  const sourceEntryIds = unique(input.sourceEntryIds);
+  const sourceText = [
+    canonicalSourceText,
+    conversationText,
+    input.segmentReadFiles.join("\n"),
+    input.segmentModifiedFiles.join("\n"),
+  ].filter(Boolean).join("\n\n");
+  return {
+    messages: input.messages,
+    conversationText,
+    sourceText,
+    sourceEntryIds,
+    sourceHash: computeSummarySourceHash({
+      conversationText: canonicalSourceText,
+      sourceEntryIds,
+      readFiles: input.segmentReadFiles,
+      modifiedFiles: input.segmentModifiedFiles,
+    }),
+    segmentReadFiles: input.segmentReadFiles,
+    segmentModifiedFiles: input.segmentModifiedFiles,
+    isSplitTurn: input.isSplitTurn,
+  };
+}
+
+export function sliceCompactionSource(
+  source: PreparedCompactionSource,
+  messageIndices: readonly number[],
+): PreparedCompactionSourceSlice {
+  const indices = [...new Set(messageIndices)].sort((left, right) => left - right);
+  if (indices.length === 0) throw new Error("Compaction segment contains no source messages");
+  for (let offset = 0; offset < indices.length; offset++) {
+    const index = indices[offset];
+    if (index === undefined || !source.messages[index] || !source.messageEntryIds[index]) {
+      throw new Error("Compaction segment references an unavailable source message");
+    }
+    if (offset > 0 && index !== (indices[offset - 1] ?? -1) + 1) {
+      throw new Error("Compaction segment source messages must be contiguous");
+    }
+  }
+  const messages = indices.map((index) => source.messages[index] as CompactionMessage);
+  const files = segmentFiles(messages);
+  return createSourceSlice({
+    messages,
+    sourceEntryIds: indices.map((index) => source.messageEntryIds[index] as string),
+    segmentReadFiles: files.readFiles,
+    segmentModifiedFiles: files.modifiedFiles,
+    isSplitTurn: source.isSplitTurn
+      && indices.some((index) => index >= source.turnPrefixStartIndex),
+  });
 }
 
 export function findActiveBranchSummary(
@@ -109,7 +202,8 @@ export function prepareCompactionSource(event: SessionBeforeCompactEvent): Prepa
   ];
   if (messages.length === 0) throw new Error("Pi compaction preparation contains no source messages");
 
-  const sourceEntryIds = mapSourceEntryIds(messages, event.branchEntries);
+  const messageEntryIds = mapSourceEntryIds(messages, event.branchEntries);
+  const sourceEntryIds = unique(messageEntryIds);
   const priorFiles = previousFiles(event.branchEntries);
   const segmentReadFiles = sortedUnique(event.preparation.fileOps.read);
   const segmentModifiedFiles = sortedUnique([
@@ -118,30 +212,20 @@ export function prepareCompactionSource(event: SessionBeforeCompactEvent): Prepa
   ]);
   const readFiles = sortedUnique([...priorFiles.read, ...segmentReadFiles]);
   const modifiedFiles = sortedUnique([...priorFiles.modified, ...segmentModifiedFiles]);
-  const conversationText = serializeConversation(convertToLlm(messages));
-  const canonicalSourceText = stableStringify(messages);
   const previousSummary = event.preparation.previousSummary;
   const previousNode = findActiveBranchSummary(event.branchEntries, previousSummary);
-  const sourceText = [
-    canonicalSourceText,
-    conversationText,
-    segmentReadFiles.join("\n"),
-    segmentModifiedFiles.join("\n"),
-  ].filter(Boolean).join("\n\n");
-
-  return {
+  const prepared = createSourceSlice({
     messages,
-    conversationText,
-    sourceText,
     sourceEntryIds,
-    sourceHash: computeSummarySourceHash({
-      conversationText: canonicalSourceText,
-      sourceEntryIds,
-      readFiles: segmentReadFiles,
-      modifiedFiles: segmentModifiedFiles,
-    }),
     segmentReadFiles,
     segmentModifiedFiles,
+    isSplitTurn: event.preparation.isSplitTurn,
+  });
+
+  return {
+    ...prepared,
+    messageEntryIds,
+    turnPrefixStartIndex: event.preparation.messagesToSummarize.length,
     readFiles,
     modifiedFiles,
     ...(previousSummary ? { previousSummary } : {}),