ds4-context-engine 0.1.0

package/src/pi-adapter/message-converter.ts

@@ -0,0 +1,181 @@
+import type {
+  CanonicalBlock,
+  CanonicalMessage,
+  CanonicalRole,
+  ImageBlock,
+} from "ds4-context-core/core/canonical-message";
+import { estimateMessageTokens } from "ds4-context-core/core/token-estimator";
+
+export interface MessageConversionInput {
+  sessionId: string;
+  entryId: string;
+  entryTimestamp?: string;
+  message: unknown;
+}
+
+function isRecord(value: unknown): value is Record<string, unknown> {
+  return value !== null && typeof value === "object" && !Array.isArray(value);
+}
+
+function stringValue(value: unknown): string | undefined {
+  return typeof value === "string" ? value : undefined;
+}
+
+function safeJson(value: unknown): string {
+  try {
+    return JSON.stringify(value) ?? "";
+  } catch {
+    return String(value);
+  }
+}
+
+function imageBlock(block: Record<string, unknown>): ImageBlock {
+  const source = isRecord(block.source) ? block.source : undefined;
+  return {
+    type: "image",
+    ...(stringValue(block.mimeType) || stringValue(source?.mediaType)
+      ? { mimeType: stringValue(block.mimeType) ?? stringValue(source?.mediaType) }
+      : {}),
+    ...(stringValue(block.data) || stringValue(source?.data)
+      ? { data: stringValue(block.data) ?? stringValue(source?.data) }
+      : {}),
+  };
+}
+
+function convertContentBlock(block: unknown): CanonicalBlock {
+  if (!isRecord(block)) return { type: "opaqueProvider", value: block };
+
+  if (block.type === "text" && typeof block.text === "string") {
+    return { type: "text", text: block.text };
+  }
+  if (block.type === "thinking" && typeof block.thinking === "string") {
+    return { type: "thinking", thinking: block.thinking };
+  }
+  if (block.type === "toolCall") {
+    return {
+      type: "toolCall",
+      id: stringValue(block.id) ?? "",
+      name: stringValue(block.name) ?? "",
+      arguments: block.arguments,
+    };
+  }
+  if (block.type === "image") return imageBlock(block);
+
+  return {
+    type: "opaqueProvider",
+    ...(typeof block.type === "string" ? { originalType: block.type } : {}),
+    value: block,
+  };
+}
+
+function contentBlocks(content: unknown): CanonicalBlock[] {
+  if (typeof content === "string") return [{ type: "text", text: content }];
+  if (!Array.isArray(content)) return content == null ? [] : [{ type: "opaqueProvider", value: content }];
+  return content.map(convertContentBlock);
+}
+
+function canonicalRole(originalRole: string | undefined): CanonicalRole {
+  switch (originalRole) {
+    case "system":
+      return "system";
+    case "user":
+      return "user";
+    case "assistant":
+      return "assistant";
+    case "toolResult":
+    case "bashExecution":
+      return "tool";
+    default:
+      return "custom";
+  }
+}
+
+function timestamp(message: Record<string, unknown>, entryTimestamp: string | undefined): number | undefined {
+  if (typeof message.timestamp === "number" && Number.isFinite(message.timestamp)) return message.timestamp;
+  if (!entryTimestamp) return undefined;
+  const parsed = Date.parse(entryTimestamp);
+  return Number.isNaN(parsed) ? undefined : parsed;
+}
+
+export function toCanonicalMessage(input: MessageConversionInput): CanonicalMessage {
+  const message = isRecord(input.message) ? input.message : {};
+  const originalRole = stringValue(message.role);
+  let blocks = contentBlocks(message.content);
+
+  if (originalRole === "toolResult") {
+    const text = blocks
+      .filter((block): block is Extract<CanonicalBlock, { type: "text" }> => block.type === "text")
+      .map((block) => block.text)
+      .join("\n");
+    const images = blocks.filter((block): block is ImageBlock => block.type === "image");
+    blocks = [
+      {
+        type: "toolResult",
+        toolCallId: stringValue(message.toolCallId) ?? "",
+        toolName: stringValue(message.toolName) ?? "",
+        content: text,
+        isError: message.isError === true,
+      },
+      ...images,
+    ];
+  } else if (originalRole === "bashExecution") {
+    blocks = [
+      { type: "text", text: `Command: ${stringValue(message.command) ?? ""}` },
+      { type: "text", text: stringValue(message.output) ?? "" },
+    ];
+  }
+
+  const createdAt = timestamp(message, input.entryTimestamp);
+  return {
+    id: `${input.sessionId}:${input.entryId}`,
+    sourceEntryId: input.entryId,
+    role: canonicalRole(originalRole),
+    blocks,
+    ...(createdAt !== undefined ? { createdAt } : {}),
+    ...(stringValue(message.provider) ? { provider: stringValue(message.provider) } : {}),
+    ...(stringValue(message.model) ? { model: stringValue(message.model) } : {}),
+    provenance: {
+      source: "pi-session",
+      sessionId: input.sessionId,
+      entryId: input.entryId,
+      ...(originalRole ? { originalRole } : {}),
+    },
+    tokenEstimate: estimateMessageTokens(message),
+    flags: {
+      atomic: originalRole === "toolResult" || blocks.some((block) => block.type === "toolCall"),
+    },
+  };
+}
+
+export function canonicalMessageSearchText(message: CanonicalMessage): string {
+  const parts: string[] = [];
+  for (const block of message.blocks) {
+    switch (block.type) {
+      case "text":
+        parts.push(block.text);
+        break;
+      case "toolCall":
+        parts.push(block.name, safeJson(block.arguments));
+        break;
+      case "toolResult":
+        parts.push(block.toolName, block.content);
+        break;
+      case "fileReference":
+        parts.push(block.path);
+        break;
+      case "artifactReference":
+        parts.push(block.artifactId);
+        break;
+      case "summaryReference":
+        parts.push(block.summaryId);
+        break;
+      // Thinking, image data, and opaque provider payloads are preserved but
+      // intentionally excluded from the lexical index.
+      case "thinking":
+      case "image":
+      case "opaqueProvider":
+        break;
+    }
+  }
+  return parts.filter(Boolean).join("\n");
+}

package/src/pi-adapter/openai-responses-stream.ts

@@ -0,0 +1,226 @@
+import {
+  createAssistantMessageEventStream,
+  type Api,
+  type AssistantMessage,
+  type AssistantMessageEvent,
+  type AssistantMessageEventStream,
+  type Context,
+  type Model,
+  type SimpleStreamOptions,
+} from "@earendil-works/pi-ai";
+import { createGrammarToolInputProperties } from "@earendil-works/pi-ai/api/constrained-sampling";
+import { openAIResponsesApi } from "@earendil-works/pi-ai/api/openai-responses.lazy";
+import { convertResponsesMessages } from "@earendil-works/pi-ai/api/openai-responses-shared";
+import {
+  continuationItemHashes,
+  type NativeContinuationAttempt,
+  type PreparedNativeContinuation,
+} from "ds4-context-core/continuation/native-continuation";
+
+const OPENAI_TOOL_CALL_PROVIDERS = new Set(["openai", "openai-codex", "opencode"]);
+
+export type OpenAIResponsesStreamDelegate = (
+  model: Model<Api>,
+  context: Context,
+  options?: SimpleStreamOptions,
+) => AssistantMessageEventStream;
+
+export interface NativeContinuationStreamController {
+  prepare(
+    payload: unknown,
+    model: Model<Api>,
+    requestSessionId: string | undefined,
+    options: {
+      forceManagedReplay: boolean;
+      retryOfContinuation: boolean;
+      fallbackReason?: string;
+    },
+  ): PreparedNativeContinuation;
+  beginManagedReplayRetry(attempt: NativeContinuationAttempt): void;
+  complete(
+    attempt: NativeContinuationAttempt,
+    message: AssistantMessage,
+    responseItemHashes: readonly string[],
+  ): void;
+  fail(attempt: NativeContinuationAttempt, reason: string): void;
+  shouldRetryManagedReplay(): boolean;
+}
+
+interface AttemptRun {
+  stream: AssistantMessageEventStream;
+  attempt: () => NativeContinuationAttempt | undefined;
+}
+
+function terminalMessage(event: AssistantMessageEvent): AssistantMessage | undefined {
+  if (event.type === "done") return event.message;
+  if (event.type === "error") return event.error;
+  return undefined;
+}
+
+function retryableContinuationError(event: AssistantMessageEvent): boolean {
+  if (event.type !== "error") return false;
+  const message = event.error.errorMessage ?? "";
+  return /previous[ _-]?response|previous_response_id|response(?:\s+with)?\s+id[^\n]{0,80}(?:not found|expired|invalid|unknown)|conversation(?:\s+with)?\s+id[^\n]{0,80}(?:not found|expired|invalid|unknown)/iu.test(message);
+}
+
+function outputItemHashes(
+  model: Model<Api>,
+  context: Context,
+  message: AssistantMessage,
+): string[] {
+  try {
+    const supportsGrammar = Boolean(
+      model.compat
+      && "supportsOpenAIGrammarTools" in model.compat
+      && model.compat.supportsOpenAIGrammarTools === true,
+    );
+    const grammarToolInputProperties = createGrammarToolInputProperties(
+      context.tools,
+      supportsGrammar,
+    );
+    const items = convertResponsesMessages(
+      model,
+      { messages: [message] },
+      OPENAI_TOOL_CALL_PROVIDERS,
+      {
+        includeSystemPrompt: false,
+        grammarToolInputProperties,
+      },
+    ).filter((item) => item.type !== "function_call_output"
+      && item.type !== "custom_tool_call_output");
+    return continuationItemHashes(items);
+  } catch {
+    return [];
+  }
+}
+
+function syntheticError(model: Model<Api>, error: unknown): AssistantMessage {
+  return {
+    role: "assistant",
+    content: [],
+    api: model.api,
+    provider: model.provider,
+    model: model.id,
+    usage: {
+      input: 0,
+      output: 0,
+      cacheRead: 0,
+      cacheWrite: 0,
+      totalTokens: 0,
+      cost: { input: 0, output: 0, cacheRead: 0, cacheWrite: 0, total: 0 },
+    },
+    stopReason: "error",
+    errorMessage: error instanceof Error ? error.message : String(error),
+    timestamp: Date.now(),
+  };
+}
+
+export function createOpenAIResponsesContinuationStream(
+  controller: NativeContinuationStreamController,
+  delegate: OpenAIResponsesStreamDelegate = (model, context, options) =>
+    openAIResponsesApi().streamSimple(model, context, options),
+): OpenAIResponsesStreamDelegate {
+  return (model, context, options) => {
+    const output = createAssistantMessageEventStream();
+
+    const startAttempt = (
+      forceManagedReplay: boolean,
+      retryOfContinuation: boolean,
+    ): AttemptRun => {
+      let preparedAttempt: NativeContinuationAttempt | undefined;
+      const stream = delegate(model, context, {
+        ...options,
+        onPayload: async (payload, payloadModel) => {
+          const replacement = await options?.onPayload?.(payload, payloadModel);
+          const finalPayload = replacement === undefined ? payload : replacement;
+          const prepared = controller.prepare(
+            finalPayload,
+            model,
+            options?.sessionId,
+            {
+              forceManagedReplay,
+              retryOfContinuation,
+              ...(retryOfContinuation
+                ? { fallbackReason: "provider-rejected-continuation" }
+                : {}),
+            },
+          );
+          preparedAttempt = prepared.attempt;
+          return prepared.payload;
+        },
+      });
+      return { stream, attempt: () => preparedAttempt };
+    };
+
+    const consume = async (
+      run: AttemptRun,
+      allowRetry: boolean,
+    ): Promise<"complete" | "retry"> => {
+      let started = false;
+      for await (const event of run.stream) {
+        if (event.type === "start") started = true;
+        const attempt = run.attempt();
+        let retryEnabled = false;
+        try {
+          retryEnabled = controller.shouldRetryManagedReplay();
+        } catch {
+          retryEnabled = false;
+        }
+        if (event.type === "error"
+          && allowRetry
+          && !started
+          && attempt?.usedContinuation
+          && !options?.signal?.aborted
+          && retryEnabled
+          && retryableContinuationError(event)) {
+          try {
+            controller.beginManagedReplayRetry(attempt);
+            return "retry";
+          } catch {
+            // Continuation bookkeeping must not replace the provider's original error.
+          }
+        }
+
+        if (event.type === "done" && attempt) {
+          try {
+            controller.complete(attempt, event.message, outputItemHashes(model, context, event.message));
+          } catch {
+            // Provider output remains authoritative if optimization bookkeeping fails.
+          }
+        } else if (event.type === "error" && attempt) {
+          try {
+            controller.fail(attempt, started
+              ? "provider-stream-failed-after-start"
+              : "provider-request-failed");
+          } catch {
+            // Provider output remains authoritative if optimization bookkeeping fails.
+          }
+        }
+
+        output.push(event);
+        if (terminalMessage(event)) {
+          output.end();
+          return "complete";
+        }
+      }
+      throw new Error("OpenAI Responses provider stream ended without a terminal event");
+    };
+
+    void (async () => {
+      try {
+        const first = startAttempt(false, false);
+        const result = await consume(first, true);
+        if (result === "retry") {
+          const replay = startAttempt(true, true);
+          await consume(replay, false);
+        }
+      } catch (error) {
+        const message = syntheticError(model, error);
+        output.push({ type: "error", reason: "error", error: message });
+        output.end();
+      }
+    })();
+
+    return output;
+  };
+}

package/src/pi-adapter/session-indexer.ts

@@ -0,0 +1,255 @@
+import { existsSync, statSync } from "node:fs";
+import type { Logger } from "ds4-context-core/shared/logging";
+import { silentLogger } from "ds4-context-core/shared/logging";
+import {
+  AppendOnlyEntryChangedError,
+  type SessionIdentity,
+  type SessionIndexCheckpointInput,
+  type SessionIndexRepository,
+  type SessionIndexState,
+} from "ds4-context-core/persistence/repositories/session-index-repository";
+import type { PiSessionSnapshot } from "./session-reader.ts";
+import { isPiSessionEntryRecord, toIndexedSessionEntry } from "./indexed-entry.ts";
+import {
+  hashFileRange,
+  readJsonlRecords,
+  readSessionHeaderRecord,
+  type JsonlReadResult,
+  type SessionHeaderRecord,
+} from "./session-jsonl.ts";
+
+const SUPPORTED_SESSION_VERSION = 3;
+
+export type SessionIndexMode = "noop" | "incremental" | "rebuild" | "skipped";
+
+export interface SessionIndexResult {
+  mode: SessionIndexMode;
+  reason: string;
+  processedEntries: number;
+  insertedEntries: number;
+  totalEntries: number;
+  malformedLines: number;
+  durationMs: number;
+}
+
+export interface SessionIndexerOptions {
+  logger?: Logger;
+  now?: () => number;
+  monotonicNow?: () => number;
+}
+
+function validateHeader(header: SessionHeaderRecord, expectedSessionId: string): void {
+  if (header.value.id !== expectedSessionId) {
+    throw new Error(
+      `Session header id ${header.value.id} does not match Pi session id ${expectedSessionId}`,
+    );
+  }
+  const version = header.value.version;
+  if (typeof version === "number" && version > SUPPORTED_SESSION_VERSION) {
+    throw new Error(`Pi session format v${version} is newer than supported v${SUPPORTED_SESSION_VERSION}`);
+  }
+}
+
+function identity(session: PiSessionSnapshot, indexedAt: number): SessionIdentity {
+  if (!session.sessionFile) throw new Error("Cannot index an ephemeral Pi session");
+  return {
+    sessionId: session.sessionId,
+    sessionFile: session.sessionFile,
+    projectPath: session.projectPath,
+    indexedAt,
+  };
+}
+
+function checkpoint(
+  session: PiSessionSnapshot,
+  header: SessionHeaderRecord,
+  read: JsonlReadResult,
+  indexedAt: number,
+  previous?: SessionIndexState,
+): SessionIndexCheckpointInput {
+  const hasNewTerminatedLine = read.checkpointHash !== undefined;
+  return {
+    sessionId: session.sessionId,
+    sessionFile: session.sessionFile ?? "",
+    headerHash: header.rawHash,
+    fileSize: read.fileSize,
+    fileMtimeMs: read.fileMtimeMs,
+    checkpointOffset: hasNewTerminatedLine
+      ? read.safeCheckpointOffset
+      : previous?.checkpointOffset ?? read.safeCheckpointOffset,
+    checkpointHashStart: hasNewTerminatedLine
+      ? read.checkpointHashStart
+      : previous?.checkpointHashStart ?? read.checkpointHashStart,
+    ...(hasNewTerminatedLine
+      ? { checkpointHash: read.checkpointHash }
+      : previous?.checkpointHash
+        ? { checkpointHash: previous.checkpointHash }
+        : {}),
+    malformedLines: (previous?.malformedLines ?? 0) + read.malformedLines,
+    indexedAt,
+  };
+}
+
+export class PiSessionIndexer {
+  private readonly logger: Logger;
+  private readonly now: () => number;
+  private readonly monotonicNow: () => number;
+
+  constructor(
+    private readonly repository: SessionIndexRepository,
+    options: SessionIndexerOptions = {},
+  ) {
+    this.logger = options.logger ?? silentLogger;
+    this.now = options.now ?? Date.now;
+    this.monotonicNow = options.monotonicNow ?? (() => performance.now());
+  }
+
+  sync(session: PiSessionSnapshot, forceRebuild = false): SessionIndexResult {
+    const started = this.monotonicNow();
+    if (!session.sessionFile || !existsSync(session.sessionFile)) {
+      return this.result("skipped", "session file is not persisted yet", 0, 0, 0, 0, started);
+    }
+
+    const header = readSessionHeaderRecord(session.sessionFile);
+    validateHeader(header, session.sessionId);
+    const state = this.repository.getState(session.sessionId);
+    const stat = statSync(session.sessionFile);
+
+    const rebuildReason = this.rebuildReason(session, header, state, stat.size, stat.mtimeMs, forceRebuild);
+    if (rebuildReason) return this.rebuild(session, header, rebuildReason, started);
+
+    if (!state) return this.rebuild(session, header, "missing index checkpoint", started);
+    if (stat.size === state.fileSize && stat.mtimeMs === state.fileMtimeMs) {
+      return this.result(
+        "noop",
+        "session file unchanged",
+        0,
+        0,
+        state.indexedEntries,
+        state.malformedLines,
+        started,
+      );
+    }
+
+    const read = readJsonlRecords(session.sessionFile, state.checkpointOffset);
+    const indexedAt = this.now();
+    const entries = read.records
+      .filter((record) => record.value.type !== "session")
+      .map((record) => toIndexedSessionEntry(session.sessionId, record, indexedAt));
+
+    try {
+      const write = this.repository.append(
+        identity(session, indexedAt),
+        entries,
+        checkpoint(session, header, read, indexedAt, state),
+      );
+      const result = this.result(
+        "incremental",
+        "appended JSONL records",
+        entries.length,
+        write.inserted,
+        write.totalEntries,
+        state.malformedLines + read.malformedLines,
+        started,
+      );
+      this.logger.debug("session_index.incremental", { sessionId: session.sessionId, ...result });
+      return result;
+    } catch (error) {
+      if (error instanceof AppendOnlyEntryChangedError) {
+        return this.rebuild(session, header, `append-only validation failed for ${error.entryId}`, started);
+      }
+      throw error;
+    }
+  }
+
+  private rebuildReason(
+    session: PiSessionSnapshot,
+    header: SessionHeaderRecord,
+    state: SessionIndexState | undefined,
+    fileSize: number,
+    fileMtimeMs: number,
+    force: boolean,
+  ): string | undefined {
+    if (force) return "forced by user";
+    if (!state) return "initial index";
+    if (state.sessionFile !== session.sessionFile) return "session file path changed";
+    if (state.headerHash !== header.rawHash) return "session header changed";
+    if (fileSize < state.fileSize) return "session file was truncated";
+    if (fileSize === state.fileSize && fileMtimeMs !== state.fileMtimeMs) return "session file changed in place";
+    if (fileSize > state.fileSize && state.checkpointOffset < state.fileSize) {
+      return "session grew after an unterminated JSONL tail";
+    }
+    if (state.checkpointOffset > 0) {
+      if (!state.checkpointHash) return "index checkpoint hash is missing";
+      const currentHash = hashFileRange(
+        session.sessionFile ?? "",
+        state.checkpointHashStart,
+        state.checkpointOffset,
+      );
+      if (currentHash !== state.checkpointHash) return "index checkpoint no longer matches the JSONL file";
+    }
+    return undefined;
+  }
+
+  private rebuild(
+    session: PiSessionSnapshot,
+    header: SessionHeaderRecord,
+    reason: string,
+    started: number,
+  ): SessionIndexResult {
+    if (!session.sessionFile) throw new Error("Cannot rebuild an ephemeral Pi session index");
+    const read = readJsonlRecords(session.sessionFile, 0);
+    const first = read.records[0];
+    if (!first || first.value.type !== "session" || first.rawHash !== header.rawHash) {
+      throw new Error("Pi session header changed while rebuilding the index");
+    }
+
+    const indexedAt = this.now();
+    const seen = new Set<string>();
+    const entries = read.records.slice(1).map((record) => {
+      if (!isPiSessionEntryRecord(record.value)) {
+        throw new Error(`Invalid Pi session entry at byte ${record.startOffset}`);
+      }
+      if (seen.has(record.value.id)) throw new Error(`Duplicate Pi session entry id: ${record.value.id}`);
+      seen.add(record.value.id);
+      return toIndexedSessionEntry(session.sessionId, record, indexedAt);
+    });
+
+    const write = this.repository.rebuild(
+      identity(session, indexedAt),
+      entries,
+      checkpoint(session, header, read, indexedAt),
+    );
+    const result = this.result(
+      "rebuild",
+      reason,
+      entries.length,
+      write.inserted,
+      write.totalEntries,
+      read.malformedLines,
+      started,
+    );
+    this.logger.info("session_index.rebuilt", { sessionId: session.sessionId, ...result });
+    return result;
+  }
+
+  private result(
+    mode: SessionIndexMode,
+    reason: string,
+    processedEntries: number,
+    insertedEntries: number,
+    totalEntries: number,
+    malformedLines: number,
+    started: number,
+  ): SessionIndexResult {
+    return {
+      mode,
+      reason,
+      processedEntries,
+      insertedEntries,
+      totalEntries,
+      malformedLines,
+      durationMs: Math.max(0, this.monotonicNow() - started),
+    };
+  }
+}