@hassangameryt/freeflow 0.5.0 → 0.7.0

package/pi-extension/src/conversation-history/ranking.ts

@@ -0,0 +1,155 @@
+import type { ConversationSearchEntry } from "./types.js";
+import {
+  normalizeText,
+  splitTextIntoPassages,
+  throwIfAborted,
+  tokenize,
+  unique,
+  type TextPassage,
+} from "./passages.js";
+
+export const BM25_K1 = 1.2;
+export const BM25_B = 0.75;
+
+export type RankedPassage = {
+  entry: ConversationSearchEntry;
+  passage: TextPassage;
+  exactPhrase: boolean;
+  matchedTerms: string[];
+  score: number;
+};
+
+type RankedTextPassage = Omit<RankedPassage, "entry">;
+
+type PassageInput = {
+  entry?: ConversationSearchEntry;
+  passage: TextPassage;
+  exactPhrase: boolean;
+  matchedTerms: string[];
+};
+
+function bm25Score(
+  tokens: readonly string[],
+  terms: readonly string[],
+  documentFrequency: ReadonlyMap<string, number>,
+  documentCount: number,
+  averageLength: number,
+): number {
+  const counts = new Map<string, number>();
+  for (const token of tokens) counts.set(token, (counts.get(token) ?? 0) + 1);
+  const length = Math.max(1, tokens.length);
+  let score = 0;
+  for (const term of terms) {
+    const termFrequency = counts.get(term) ?? 0;
+    if (termFrequency === 0) continue;
+    const df = documentFrequency.get(term) ?? 0;
+    const inverseDocumentFrequency = Math.log(1 + (documentCount - df + 0.5) / (df + 0.5));
+    const normalization = BM25_K1 * (1 - BM25_B + BM25_B * (length / Math.max(1, averageLength)));
+    score += inverseDocumentFrequency * ((termFrequency * (BM25_K1 + 1)) / (termFrequency + normalization));
+  }
+  return score;
+}
+
+function comparePassages(left: RankedPassage, right: RankedPassage): number {
+  if (left.exactPhrase !== right.exactPhrase) return left.exactPhrase ? -1 : 1;
+  if (left.score !== right.score) return right.score - left.score;
+  if (left.matchedTerms.length !== right.matchedTerms.length) {
+    return right.matchedTerms.length - left.matchedTerms.length;
+  }
+  if (left.entry.position !== right.entry.position) return right.entry.position - left.entry.position;
+  return left.passage.start - right.passage.start;
+}
+
+function compareTextPassages(left: RankedTextPassage, right: RankedTextPassage): number {
+  if (left.exactPhrase !== right.exactPhrase) return left.exactPhrase ? -1 : 1;
+  if (left.score !== right.score) return right.score - left.score;
+  if (left.matchedTerms.length !== right.matchedTerms.length) {
+    return right.matchedTerms.length - left.matchedTerms.length;
+  }
+  return left.passage.start - right.passage.start;
+}
+
+function rankInputs<T extends PassageInput>(
+  inputs: T[],
+  terms: string[],
+  signal?: AbortSignal,
+): Array<T & { score: number }> {
+  const documentFrequency = new Map<string, number>();
+  for (const input of inputs) {
+    throwIfAborted(signal);
+    for (const term of unique(input.passage.tokens)) {
+      documentFrequency.set(term, (documentFrequency.get(term) ?? 0) + 1);
+    }
+  }
+  const averageLength =
+    inputs.length === 0 ? 1 : inputs.reduce((total, input) => total + input.passage.length, 0) / inputs.length;
+  const matching: T[] = [];
+  for (const input of inputs) {
+    throwIfAborted(signal);
+    if (input.matchedTerms.length > 0 || input.exactPhrase) matching.push(input);
+  }
+  return matching
+    .map((input) => ({
+      ...input,
+      score: bm25Score(input.passage.tokens, terms, documentFrequency, inputs.length, averageLength),
+    }))
+    .sort((left, right) => {
+      if (left.entry && right.entry) {
+        return comparePassages(
+          { ...left, entry: left.entry } as RankedPassage,
+          { ...right, entry: right.entry } as RankedPassage,
+        );
+      }
+      return compareTextPassages(left as RankedTextPassage & T, right as RankedTextPassage & T);
+    });
+}
+
+export function rankMatchingPassages(
+  entries: readonly ConversationSearchEntry[],
+  query: string,
+  signal?: AbortSignal,
+): RankedPassage[] {
+  throwIfAborted(signal);
+  const normalizedQuery = normalizeText(query);
+  const terms = unique(tokenize(query));
+  if (!normalizedQuery || terms.length === 0) return [];
+  const inputs: PassageInput[] = [];
+  for (const entry of entries) {
+    throwIfAborted(signal);
+    for (const passage of splitTextIntoPassages(entry.text, signal)) {
+      const normalizedPassage = normalizeText(passage.text);
+      inputs.push({
+        entry,
+        passage,
+        exactPhrase: normalizedPassage.includes(normalizedQuery),
+        matchedTerms: terms.filter((term) => passage.tokens.includes(term)),
+      });
+    }
+  }
+  return rankInputs(inputs, terms, signal) as RankedPassage[];
+}
+
+export function strongestTextPassage(text: string, query: string): RankedTextPassage | undefined {
+  const normalizedQuery = normalizeText(query);
+  const terms = unique(tokenize(query));
+  if (!normalizedQuery || terms.length === 0) return undefined;
+  const inputs: PassageInput[] = splitTextIntoPassages(text).map((passage) => ({
+    passage,
+    exactPhrase: normalizeText(passage.text).includes(normalizedQuery),
+    matchedTerms: terms.filter((term) => passage.tokens.includes(term)),
+  }));
+  return rankInputs(inputs, terms)[0] as RankedTextPassage | undefined;
+}
+
+export function focusedWindow(text: string, focus: string, maxCharacters: number): string | undefined {
+  const strongest = strongestTextPassage(text, focus);
+  if (!strongest) return undefined;
+  const center = Math.floor((strongest.passage.start + strongest.passage.end) / 2);
+  const proposedStart = Math.max(0, Math.min(center - Math.floor(maxCharacters / 2), text.length - maxCharacters));
+  const previousLineBreak = proposedStart > 0 ? text.lastIndexOf("\n", proposedStart - 1) : -1;
+  const start = previousLineBreak >= 0 ? previousLineBreak + 1 : proposedStart;
+  const proposedEnd = Math.min(text.length, start + maxCharacters);
+  const previousEndLineBreak = proposedEnd < text.length ? text.lastIndexOf("\n", proposedEnd) : -1;
+  const end = previousEndLineBreak > start ? previousEndLineBreak + 1 : proposedEnd;
+  return text.slice(start, end);
+}

package/pi-extension/src/conversation-history/runtime.ts

@@ -0,0 +1,250 @@
+import { FreeflowContextRuntime } from "../freeflow-context/runtime.js";
+import { contextRefForEntry, entryIdFromContextRef } from "../freeflow-context/types.js";
+import type { FreeflowContextSnapshot, ProjectedContextSource } from "../freeflow-context/types.js";
+import { focusedWindow } from "./ranking.js";
+import { searchConversationEntries, type ConversationSearchOptions } from "./search.js";
+import { tokenize, unique } from "./passages.js";
+import type { ConversationSearchEntry } from "./types.js";
+
+const MAX_QUERY_CHARACTERS = 500;
+const MAX_QUERY_TERMS = 32;
+const MAX_TOOL_NAMES = 8;
+const MAX_TOOL_NAME_CHARACTERS = 128;
+const MAX_REF_CHARACTERS = 256;
+const MAX_FOCUS_CHARACTERS = 500;
+const MAX_RETRIEVE_REFS = 3;
+const MAX_SOURCE_CHARACTERS = 8_000;
+const MAX_TOTAL_RETRIEVED_CHARACTERS = 24_000;
+const VALID_KINDS = new Set(["user", "assistant", "toolResult", "summary"]);
+
+function normalizedTerms(value: string): string[] {
+  return unique(tokenize(value));
+}
+
+function failure(operation: "search" | "retrieve", reason: string, message?: string) {
+  return {
+    status: "rejected" as const,
+    operation,
+    reason: reason.slice(0, 64),
+    ...(message ? { message: message.slice(0, 500) } : {}),
+  };
+}
+
+function unavailable(operation: "search" | "retrieve", reason = "conversation_history_unavailable") {
+  return { status: "unavailable" as const, operation, reason: reason.slice(0, 64) };
+}
+
+function sourceMetadata(source: ProjectedContextSource) {
+  return {
+    ref: source.ref,
+    kind: source.kind,
+    ...(source.toolNames
+      ? {
+          toolNames: source.toolNames.slice(0, MAX_TOOL_NAMES),
+          ...(source.toolNames.length > MAX_TOOL_NAMES ? { toolNamesTruncated: true } : {}),
+        }
+      : {}),
+    ...(source.isError === undefined ? {} : { isError: source.isError }),
+    timestamp: source.timestamp.slice(0, 64),
+  };
+}
+
+export class ConversationHistoryRuntime {
+  private ctx: any;
+  private readonly freeflowContext: FreeflowContextRuntime;
+  private readonly isSourceFullyProjected: (entryId: string) => boolean;
+  private snapshot: FreeflowContextSnapshot | undefined;
+  private available = true;
+  private unavailableReason: string | undefined;
+
+  constructor(
+    ctx: any,
+    isSourceFullyProjected: (entryId: string) => boolean = () => true,
+    freeflowContext = new FreeflowContextRuntime(ctx),
+  ) {
+    this.ctx = ctx;
+    this.isSourceFullyProjected = isSourceFullyProjected;
+    this.freeflowContext = freeflowContext;
+  }
+
+  setContext(ctx: any): void {
+    this.ctx = ctx;
+    this.freeflowContext.setContext(ctx);
+  }
+
+  capture(contextVirtualizationEnabled = false): boolean {
+    try {
+      this.freeflowContext.setContext(this.ctx);
+      this.snapshot = this.freeflowContext.captureSnapshot({
+        contextVirtualizationEnabled,
+        isSourceFullyProjected: this.isSourceFullyProjected,
+      });
+      this.available = true;
+      this.unavailableReason = undefined;
+      return true;
+    } catch (error) {
+      this.snapshot = undefined;
+      this.available = false;
+      this.unavailableReason = error instanceof Error ? error.message : String(error);
+      return false;
+    }
+  }
+
+  async search(params: Record<string, unknown>, signal?: AbortSignal): Promise<any> {
+    if (signal?.aborted) return unavailable("search", "conversation_history_search_cancelled");
+    if (!this.available || !this.snapshot) return unavailable("search", this.unavailableReason);
+    const validation = this.validateSearch(params);
+    if (validation) return validation;
+
+    const query = (params.query as string).trim();
+    const hidden = [...this.snapshot.entries.values()].filter(
+      (entry) =>
+        !this.snapshot?.visibleSourceIds.has(entry.source.source.entryId) &&
+        !this.snapshot?.materializedSourceIds.has(entry.source.source.entryId),
+    );
+    if (hidden.length === 0) {
+      if (this.snapshot.skippedEntries === 0) return failure("search", "no_hidden_conversation_history");
+      return {
+        status: "ok",
+        operation: "search",
+        query,
+        coverage: "partial",
+        skippedEntries: this.snapshot.skippedEntries,
+        returned: 0,
+        truncated: false,
+        hits: [],
+      };
+    }
+
+    let result;
+    try {
+      result = searchConversationEntries(
+        hidden as ConversationSearchEntry[],
+        {
+          query,
+          kinds: params.kinds as ConversationSearchOptions["kinds"],
+          toolNames: Array.isArray(params.toolNames)
+            ? params.toolNames.map((name) => (typeof name === "string" ? name.trim() : ""))
+            : undefined,
+          limit: params.limit as number | undefined,
+        },
+        signal,
+      );
+    } catch (error) {
+      if (signal?.aborted || (error instanceof Error && error.name === "AbortError")) {
+        return unavailable("search", "conversation_history_search_cancelled");
+      }
+      return unavailable("search");
+    }
+    return {
+      status: "ok",
+      operation: "search",
+      query,
+      coverage: this.snapshot.skippedEntries > 0 ? "partial" : "complete",
+      ...(this.snapshot.skippedEntries > 0 ? { skippedEntries: this.snapshot.skippedEntries } : {}),
+      returned: result.returned,
+      truncated: result.truncated,
+      hits: result.hits,
+    };
+  }
+
+  async retrieve(params: Record<string, unknown>, signal?: AbortSignal): Promise<any> {
+    if (signal?.aborted) return unavailable("retrieve", "conversation_history_retrieve_cancelled");
+    if (!this.available || !this.snapshot) return unavailable("retrieve", this.unavailableReason);
+    const refs = params.refs;
+    if (!Array.isArray(refs) || refs.length < 1 || refs.length > MAX_RETRIEVE_REFS) {
+      return failure("retrieve", "refs_count_invalid");
+    }
+    if (refs.some((ref) => typeof ref !== "string" || ref.length > MAX_REF_CHARACTERS)) {
+      return failure("retrieve", "ref_invalid");
+    }
+    const canonicalRefs = refs.map((ref) => {
+      const entryId = entryIdFromContextRef(ref);
+      return entryId ? contextRefForEntry(entryId) : undefined;
+    });
+    if (canonicalRefs.some((ref) => ref === undefined)) return failure("retrieve", "ref_invalid");
+    const uniqueRefs = canonicalRefs as string[];
+    if (new Set(uniqueRefs).size !== uniqueRefs.length) return failure("retrieve", "duplicate_ref");
+    const focus = params.focus;
+    if (
+      focus !== undefined &&
+      (typeof focus !== "string" || focus.trim().length === 0 || focus.length > MAX_FOCUS_CHARACTERS)
+    ) {
+      return failure("retrieve", "focus_invalid");
+    }
+
+    const items = [];
+    let total = 0;
+    for (const ref of uniqueRefs) {
+      if (signal?.aborted) return unavailable("retrieve", "conversation_history_retrieve_cancelled");
+      const entryId = entryIdFromContextRef(ref);
+      const source = entryId ? this.snapshot.entries.get(contextRefForEntry(entryId)) : undefined;
+      if (!source) return failure("retrieve", "reference_unresolved");
+      if (this.snapshot.visibleSourceIds.has(entryId) || this.snapshot.materializedSourceIds.has(entryId)) {
+        return failure("retrieve", "source_visible");
+      }
+      const sourceCharacters = source.text.length;
+      let content = source.text;
+      let completeness: "complete" | "partial" = "complete";
+      if (sourceCharacters > MAX_SOURCE_CHARACTERS) {
+        if (typeof focus !== "string") return failure("retrieve", "focus_required_for_oversized_source");
+        const focused = focusedWindow(source.text, focus.trim(), MAX_SOURCE_CHARACTERS);
+        if (focused === undefined) return failure("retrieve", "focus_no_match");
+        content = focused;
+        completeness = "partial";
+      }
+      if (total + content.length > MAX_TOTAL_RETRIEVED_CHARACTERS)
+        return failure("retrieve", "retrieval_budget_exceeded");
+      total += content.length;
+      items.push({
+        ...sourceMetadata(source),
+        content,
+        completeness,
+        sourceCharacters,
+        returnedCharacters: content.length,
+      });
+    }
+    return {
+      status: "ok",
+      operation: "retrieve",
+      ...(typeof focus === "string" ? { focus: focus.trim() } : {}),
+      items,
+    };
+  }
+
+  private validateSearch(params: Record<string, unknown>): any | undefined {
+    if (typeof params.query !== "string" || params.query.trim().length === 0)
+      return failure("search", "query_required");
+    if (params.query.length > MAX_QUERY_CHARACTERS) return failure("search", "query_too_long");
+    if (normalizedTerms(params.query).length > MAX_QUERY_TERMS) return failure("search", "query_terms_too_many");
+    if (params.kinds !== undefined) {
+      if (!Array.isArray(params.kinds) || params.kinds.length < 1 || params.kinds.length > 4)
+        return failure("search", "kinds_invalid");
+      if (
+        new Set(params.kinds).size !== params.kinds.length ||
+        params.kinds.some((kind) => typeof kind !== "string" || !VALID_KINDS.has(kind))
+      ) {
+        return failure("search", "kinds_invalid");
+      }
+    }
+    if (params.toolNames !== undefined) {
+      if (!Array.isArray(params.toolNames) || params.toolNames.length < 1 || params.toolNames.length > MAX_TOOL_NAMES)
+        return failure("search", "tool_names_invalid");
+      const normalizedNames = params.toolNames.map((name) => (typeof name === "string" ? name.trim() : ""));
+      if (
+        params.toolNames.some((name) => typeof name !== "string") ||
+        new Set(normalizedNames.map((name) => name.toLocaleLowerCase())).size !== normalizedNames.length ||
+        normalizedNames.some((name) => name.length === 0 || name.length > MAX_TOOL_NAME_CHARACTERS)
+      ) {
+        return failure("search", "tool_names_invalid");
+      }
+    }
+    if (
+      params.limit !== undefined &&
+      (!Number.isInteger(params.limit) || (params.limit as number) < 1 || (params.limit as number) > 20)
+    ) {
+      return failure("search", "limit_invalid");
+    }
+    return undefined;
+  }
+}

package/pi-extension/src/conversation-history/search.ts

@@ -0,0 +1,79 @@
+import { normalizeText, snippet, throwIfAborted, tokenize, unique } from "./passages.js";
+import { rankMatchingPassages } from "./ranking.js";
+import type {
+  ConversationSearchEntry,
+  ConversationSearchHit,
+  ConversationSearchOptions,
+  ConversationSearchResult,
+} from "./types.js";
+
+export type {
+  ConversationSearchEntry,
+  ConversationSearchHit,
+  ConversationSearchOptions,
+  ConversationSearchResult,
+} from "./types.js";
+
+const MAX_DISPLAYED_TOOL_NAMES = 8;
+
+function matchesFilters(entry: ConversationSearchEntry, options: ConversationSearchOptions): boolean {
+  if (options.kinds && !options.kinds.includes(entry.kind)) return false;
+  if (options.toolNames) {
+    const names = new Set((entry.toolNames ?? []).map((name) => name.toLocaleLowerCase()));
+    if (!options.toolNames.some((name) => names.has(name.toLocaleLowerCase()))) return false;
+  }
+  return true;
+}
+
+export function searchConversationEntries(
+  entries: readonly ConversationSearchEntry[],
+  options: ConversationSearchOptions,
+  signal?: AbortSignal,
+): ConversationSearchResult {
+  throwIfAborted(signal);
+  const normalizedQuery = normalizeText(options.query);
+  const terms = unique(tokenize(options.query));
+  const limit = Math.min(20, Math.max(1, Math.trunc(options.limit ?? 8)));
+  if (!normalizedQuery || terms.length === 0) return { returned: 0, truncated: false, hits: [] };
+
+  const candidates = entries.filter((entry) => matchesFilters(entry, options));
+  const rankedPassages = rankMatchingPassages(candidates, options.query, signal);
+  const byEntry = new Map<string, (typeof rankedPassages)[number]>();
+  for (const passage of rankedPassages) {
+    throwIfAborted(signal);
+    if (!byEntry.has(passage.entry.ref)) byEntry.set(passage.entry.ref, passage);
+  }
+
+  const ranked = [...byEntry.values()];
+  const hits = ranked.slice(0, limit).map((passage) => {
+    const matchType = passage.exactPhrase
+      ? "exact-phrase"
+      : passage.matchedTerms.length === terms.length
+        ? "all-terms"
+        : "partial-terms";
+    return {
+      ref: passage.entry.ref,
+      kind: passage.entry.kind,
+      ...(passage.entry.toolNames
+        ? {
+            toolNames: passage.entry.toolNames.slice(0, MAX_DISPLAYED_TOOL_NAMES),
+            ...(passage.entry.toolNames.length > MAX_DISPLAYED_TOOL_NAMES ? { toolNamesTruncated: true } : {}),
+          }
+        : {}),
+      ...(passage.entry.isError === undefined ? {} : { isError: passage.entry.isError }),
+      timestamp: passage.entry.timestamp.slice(0, 64),
+      snippet: snippet(passage.passage.text, normalizedQuery, terms),
+      match: {
+        type: matchType,
+        matchedTerms: passage.matchedTerms,
+        queryTermCount: terms.length,
+      },
+    } satisfies ConversationSearchHit;
+  });
+
+  return {
+    returned: hits.length,
+    truncated: ranked.length > limit,
+    hits,
+  };
+}

package/pi-extension/src/conversation-history/types.ts

@@ -0,0 +1,31 @@
+import type { ProjectedContextSource } from "../freeflow-context/types.js";
+
+export type ConversationSearchEntry = ProjectedContextSource;
+
+export type ConversationSearchOptions = {
+  query: string;
+  kinds?: readonly ConversationSearchEntry["kind"][];
+  toolNames?: readonly string[];
+  limit?: number;
+};
+
+export type ConversationSearchHit = {
+  ref: string;
+  kind: ConversationSearchEntry["kind"];
+  toolNames?: string[];
+  toolNamesTruncated?: boolean;
+  isError?: boolean;
+  timestamp: string;
+  snippet: string;
+  match: {
+    type: "exact-phrase" | "all-terms" | "partial-terms";
+    matchedTerms: string[];
+    queryTermCount: number;
+  };
+};
+
+export type ConversationSearchResult = {
+  returned: number;
+  truncated: boolean;
+  hits: ConversationSearchHit[];
+};

package/pi-extension/src/freeflow-context/resolver.ts

@@ -0,0 +1,139 @@
+import { entryIdFromContextRef } from "./types.js";
+import type { ContextSourceKind, ResolvedContextEntry, ResolvedContextSource } from "./types.js";
+
+export function isToolResultEntry(entry: any): boolean {
+  return entry?.type === "message" && entry?.message?.role === "toolResult";
+}
+
+export class ContextSourceResolver {
+  private ctx: any;
+
+  constructor(ctx: any) {
+    this.ctx = ctx;
+  }
+
+  setContext(ctx: any): void {
+    this.ctx = ctx;
+  }
+
+  sessionId(): string {
+    const value = this.ctx?.sessionManager?.getSessionId?.() ?? this.ctx?.sessionId ?? this.ctx?.cwd ?? "current";
+    return String(value);
+  }
+
+  branchLeafId(): string | null | undefined {
+    try {
+      return this.ctx?.sessionManager?.getLeafId?.();
+    } catch {
+      return undefined;
+    }
+  }
+
+  branchEntries(): any[] {
+    const entries = this.ctx?.sessionManager?.getBranch?.() ?? this.ctx?.sessionManager?.getEntries?.() ?? [];
+    return Array.isArray(entries) ? entries : [];
+  }
+
+  activeContextEntries(): any[] {
+    try {
+      const entries = this.ctx?.sessionManager?.buildContextEntries?.() ?? this.branchEntries();
+      return Array.isArray(entries) ? entries : [];
+    } catch {
+      return this.branchEntries();
+    }
+  }
+
+  activeContextEntriesStrict(): any[] {
+    const manager = this.ctx?.sessionManager;
+    if (typeof manager?.buildContextEntries !== "function") return this.branchEntries();
+    const entries = manager.buildContextEntries();
+    if (!Array.isArray(entries)) throw new Error("active_context_entries_unavailable");
+    return entries;
+  }
+
+  resolvedEntriesById(entries = this.branchEntries(), sessionId = this.sessionId()): Map<string, ResolvedContextEntry> {
+    const result = new Map<string, ResolvedContextEntry>();
+    for (const entry of entries) {
+      const resolved = this.entryFromSessionEntry(entry, sessionId);
+      if (resolved) result.set(resolved.source.entryId, resolved);
+    }
+    return result;
+  }
+
+  toolResultSources(entries = this.branchEntries(), sessionId = this.sessionId()): Map<string, ResolvedContextSource> {
+    const grouped = new Map<string, ResolvedContextSource[]>();
+    for (const entry of entries) {
+      const source = this.sourceFromEntry(entry, sessionId);
+      if (!source) continue;
+      const group = grouped.get(source.source.toolCallId!) ?? [];
+      group.push(source);
+      grouped.set(source.source.toolCallId!, group);
+    }
+
+    const unique = new Map<string, ResolvedContextSource>();
+    for (const [toolCallId, sources] of grouped) {
+      if (sources.length === 1) unique.set(toolCallId, sources[0]);
+    }
+    return unique;
+  }
+
+  sourcesByEntryId(entries = this.branchEntries(), sessionId = this.sessionId()): Map<string, ResolvedContextSource> {
+    const result = new Map<string, ResolvedContextSource>();
+    for (const source of this.toolResultSources(entries, sessionId).values()) {
+      result.set(source.source.entryId, source);
+    }
+    return result;
+  }
+
+  activeToolResultIds(): Set<string> {
+    const ids = new Set<string>();
+    for (const entry of this.activeContextEntries()) {
+      if (isToolResultEntry(entry) && typeof entry.id === "string") ids.add(entry.id);
+    }
+    return ids;
+  }
+
+  resolveCurrent(ref: unknown): ResolvedContextEntry | undefined {
+    const entryId = entryIdFromContextRef(ref);
+    if (!entryId) return undefined;
+    const entry = this.branchEntries().find((candidate) => candidate?.id === entryId);
+    return entry ? this.entryFromSessionEntry(entry, this.sessionId()) : undefined;
+  }
+
+  private sourceFromEntry(entry: any, sessionId: string): ResolvedContextSource | undefined {
+    const resolved = this.entryFromSessionEntry(entry, sessionId);
+    return resolved?.kind === "toolResult" ? (resolved as ResolvedContextSource) : undefined;
+  }
+
+  private entryFromSessionEntry(entry: any, sessionId: string): ResolvedContextEntry | undefined {
+    if (!entry || typeof entry.id !== "string") return undefined;
+    const kind = this.kindForEntry(entry);
+    if (!kind) return undefined;
+    const message = entry.message;
+    const toolCallId =
+      kind === "toolResult" && typeof message?.toolCallId === "string" ? message.toolCallId : undefined;
+    return {
+      kind,
+      entry,
+      ...(message ? { message } : {}),
+      source: {
+        sessionId,
+        entryId: entry.id,
+        ...(toolCallId ? { toolCallId } : {}),
+        ...(typeof message?.toolName === "string" ? { toolName: message.toolName } : {}),
+      },
+    };
+  }
+
+  private kindForEntry(entry: any): ContextSourceKind | undefined {
+    if (entry.type === "custom") return "custom";
+    if (entry.type === "compaction" || entry.type === "branch_summary") return "summary";
+    if (entry.type !== "message") return undefined;
+    if (entry.message?.role === "user") return "user";
+    if (entry.message?.role === "assistant") return "assistant";
+    if (entry.message?.role === "toolResult") return "toolResult";
+    if (entry.message?.role === "custom") return "custom";
+    if (entry.message?.role === "branchSummary" || entry.message?.role === "compactionSummary") return "summary";
+    return undefined;
+  }
+}