self-bench 0.3.5 → 0.3.6

package/src/cli.ts

@@ -11,13 +11,23 @@ import { fileURLToPath } from "node:url";
 import { parseArgs } from "node:util";
 import { buildCommit } from "./build-metadata.js";
 import { loadWorkerConfig } from "./config.js";
+import { MAX_CANDIDATES_PER_RUN } from "./contracts.js";
 import { runCommand } from "./process.js";
-import { collectGitHubPullRequestProvenance, collectRepositoryProvenance } from "./provenance.js";
+import {
+  collectRepositoryProvenance,
+  collectRepositoryProvenanceWithMetadata,
+} from "./provenance.js";
+import {
+  applyProvenanceAssociationManifests,
+  associationSessionSummaries,
+  createProvenanceAssociationManifest,
+  resolveMergedPullRequest,
+  writeProvenanceAssociationManifest,
+} from "./provenance-associations.js";
 import { isExecutionBackend, isHarborEnvironment, matchingHarborEnvironment } from "./providers.js";
 import { type PolledRunStatus, waitForRun } from "./run-wait.js";
+import { applyVercelProfile, setupVercel } from "./setup/vercel/index.js";
 import { SetupCanceledError } from "./terminal-prompts.js";
-import { applyVercelProfile } from "./vercel-profile.js";
-import { setupVercel } from "./vercel-setup.js";
 
 const [command, ...rest] = process.argv.slice(2);
 switch (command) {
@@ -30,6 +40,9 @@ switch (command) {
   case "run":
     await run(rest);
     break;
+  case "associate":
+    await associate(rest);
+    break;
   case "down":
     await down();
     break;
@@ -162,6 +175,57 @@ async function down(): Promise<void> {
   await runCommand("docker", ["compose", "--file", resolve(projectRoot, "compose.yaml"), "down"]);
 }
 
+async function associate(args: string[]): Promise<void> {
+  const parsed = parseArgs({
+    args,
+    options: {
+      repo: { type: "string", short: "r" },
+      pr: { type: "string" },
+      session: { type: "string", multiple: true },
+      output: { type: "string", short: "o" },
+      "list-sessions": { type: "boolean", default: false },
+    },
+    strict: true,
+  });
+  const repositoryPath = resolve(parsed.values.repo ?? fail("--repo is required"));
+  const [repository, collected] = await Promise.all([
+    resolveRepository(repositoryPath),
+    collectRepositoryProvenanceWithMetadata(repositoryPath, process.env.HOME ?? homedir()),
+  ]);
+  const localMessages = collected.messages;
+  const sessions = associationSessionSummaries(localMessages, collected.sessions);
+  if (parsed.values["list-sessions"]) {
+    console.log(JSON.stringify({ repository: repository.url, sessions }, null, 2));
+    return;
+  }
+
+  const sourcePr = positiveInteger(parsed.values.pr ?? fail("--pr is required"), "--pr");
+  const output = resolve(parsed.values.output ?? fail("--output is required"));
+  const selectors = parsed.values.session ?? [];
+  const pullRequest = await resolveMergedPullRequest(repository.url, sourcePr);
+  const manifest = createProvenanceAssociationManifest({
+    repositoryUrl: repository.url,
+    pullRequest,
+    messages: localMessages,
+    sessionSelectors: selectors,
+  });
+  await writeProvenanceAssociationManifest(output, manifest);
+  console.log(
+    JSON.stringify(
+      {
+        output,
+        repository: manifest.repository,
+        sourcePr: manifest.sourcePr,
+        sourceUrl: manifest.sourceUrl,
+        sessions: selectors.length,
+        messages: manifest.messages.length,
+      },
+      null,
+      2,
+    ),
+  );
+}
+
 async function run(args: string[]): Promise<void> {
   const parsed = parseArgs({
     args,
@@ -174,6 +238,7 @@ async function run(args: string[]): Promise<void> {
       model: { type: "string", default: "gpt-5.6-sol" },
       wait: { type: "boolean", default: false },
       output: { type: "string", short: "o" },
+      association: { type: "string", multiple: true },
     },
     strict: true,
   });
@@ -184,21 +249,25 @@ async function run(args: string[]): Promise<void> {
     hard: nonnegativeInteger(parsed.values["hard-count"] ?? "0", "--hard-count"),
   };
   const totalCandidates = candidateCounts.easy + candidateCounts.medium + candidateCounts.hard;
-  if (totalCandidates < 1 || totalCandidates > 100) {
-    fail("the total candidate count must be between 1 and 100");
+  if (totalCandidates < 1 || totalCandidates > MAX_CANDIDATES_PER_RUN) {
+    fail(`the total candidate count must be between 1 and ${MAX_CANDIDATES_PER_RUN}`);
   }
   const runId = parsed.values["run-id"] ?? defaultRunId();
-  const [repository, localMessages, selfbenchCommit] = await Promise.all([
+  const [repository, collectedMessages, selfbenchCommit] = await Promise.all([
     resolveRepository(repositoryPath),
     collectRepositoryProvenance(repositoryPath, process.env.HOME ?? homedir()),
     resolveSelfBenchCommit(),
   ]);
-  const githubMessages = await collectGitHubPullRequestProvenance(repository.url);
-  const messages = [...localMessages, ...githubMessages];
-  if (messages.length === 0) {
-    throw new Error("no sanitized local-session or GitHub pull-request provenance was found");
-  }
-  const corpus = Buffer.from(`${messages.map((message) => JSON.stringify(message)).join("\n")}\n`);
+  const localMessages = await applyProvenanceAssociationManifests(
+    collectedMessages,
+    repository.url,
+    parsed.values.association ?? [],
+  );
+  const corpus = Buffer.from(
+    localMessages.length > 0
+      ? `${localMessages.map((message) => JSON.stringify(message)).join("\n")}\n`
+      : "",
+  );
   const provenance = await requestJson(`/v1/provenance?runId=${encodeURIComponent(runId)}`, {
     method: "POST",
     body: corpus,
@@ -222,9 +291,8 @@ async function run(args: string[]): Promise<void> {
     JSON.stringify(
       {
         ...response,
-        provenanceMessages: messages.length,
         localProvenanceMessages: localMessages.length,
-        githubPullRequestMessages: githubMessages.length,
+        associationManifests: parsed.values.association?.length ?? 0,
       },
       null,
       2,
@@ -377,6 +445,14 @@ function asPolledRunStatus(value: Record<string, unknown>): PolledRunStatus {
   return { ...value, phase: value.phase };
 }
 
+function positiveInteger(value: string, label: string): number {
+  const parsed = Number(value);
+  if (!Number.isInteger(parsed) || parsed < 1) {
+    throw new Error(`${label} must be a positive integer`);
+  }
+  return parsed;
+}
+
 function nonnegativeInteger(value: string, label: string): number {
   const parsed = Number(value);
   if (!Number.isInteger(parsed) || parsed < 0) {
@@ -401,8 +477,12 @@ Usage:
   self-bench up [--backend docker|modal|vercel] [--harbor-environment docker|modal]
                 [--modal-config PATH] [--vercel-profile NAME]
   self-bench down
+  self-bench associate --repo PATH --list-sessions
+  self-bench associate --repo PATH --pr NUMBER --session TYPE:SESSION_ID [...]
+                       --output ASSOCIATION.json
   self-bench run --repo PATH [--easy-count N] [--medium-count N] [--hard-count N]
-                  [--model MODEL] [--run-id ID] [--wait] [--output OUTPUT.tar.gz]
+                  [--model MODEL] [--association ASSOCIATION.json ...]
+                  [--run-id ID] [--wait] [--output OUTPUT.tar.gz]
   self-bench status RUN_ID
   self-bench cancel RUN_ID
   self-bench download RUN_ID OUTPUT.tar.gz
@@ -413,6 +493,10 @@ requires --harbor-environment because Harbor does not support Vercel. Modal gene
 ~/.modal.toml unless --modal-config overrides it. Run self-bench setup vercel once to create or select a
 project, publish the pinned runtime image, verify access, and save an owner-only local profile.
 
+The associate command runs locally. --list-sessions prints selectors, counts, local paths, and modification
+times, but never request text. Association writes a create-only, text-free manifest; it never uploads or
+starts a run. Pass the manifest to run with --association (repeatable). No LLM participates in association.
+
 The tier counts are candidate authoring budgets, not accepted-task targets. Rejected candidates are not
 replaced, and the export contains only accepted tasks. The run command performs only repository metadata
 and sanitized provenance upload locally; discovery, authoring, validation, review, and audit run remotely.

package/src/config.ts

@@ -6,7 +6,7 @@ import {
   type HarborEnvironment,
   matchingHarborEnvironment,
 } from "./providers.js";
-import { parseSandboxTimeoutCapText, STANDARD_VERCEL_TIMEOUT_CAP_MS } from "./sandbox-timeout.js";
+import { parseSandboxTimeoutCapText, STANDARD_VERCEL_TIMEOUT_CAP_MS } from "./sandbox/timeout.js";
 
 const emptyStringAsUndefined = (value: unknown): unknown =>
   typeof value === "string" && value.trim() === "" ? undefined : value;

package/src/contracts.ts

@@ -22,17 +22,19 @@ export type ArtifactRef = z.infer<typeof artifactRefSchema>;
 export const difficultySchema = z.enum(["easy", "medium", "hard"]);
 export type Difficulty = z.infer<typeof difficultySchema>;
 
+export const MAX_CANDIDATES_PER_RUN = 10_000;
+
 const candidateCountsSchema = z
   .object({
-    easy: z.number().int().min(0).max(100),
-    medium: z.number().int().min(0).max(100),
-    hard: z.number().int().min(0).max(100),
+    easy: z.number().int().min(0).max(MAX_CANDIDATES_PER_RUN),
+    medium: z.number().int().min(0).max(MAX_CANDIDATES_PER_RUN),
+    hard: z.number().int().min(0).max(MAX_CANDIDATES_PER_RUN),
   })
   .refine(({ easy, medium, hard }) => easy + medium + hard >= 1, {
     message: "at least one candidate must be requested",
   })
-  .refine(({ easy, medium, hard }) => easy + medium + hard <= 100, {
-    message: "at most 100 candidates may be requested",
+  .refine(({ easy, medium, hard }) => easy + medium + hard <= MAX_CANDIDATES_PER_RUN, {
+    message: `at most ${MAX_CANDIDATES_PER_RUN} candidates may be requested`,
   });
 
 const runVersionSchema = z

package/src/extensions/discovery.ts

@@ -1,4 +1,4 @@
-import { writeFileSync } from "node:fs";
+import { readFileSync, writeFileSync } from "node:fs";
 import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";
 import { Type } from "@sinclair/typebox";
 
@@ -22,7 +22,6 @@ export default function discoveryExtension(pi: ExtensionAPI): void {
               sourceUrl: Type.String({ minLength: 1 }),
               baseCommit: Type.String({ pattern: "^[0-9a-fA-F]{40}$" }),
               completedCommit: Type.String({ pattern: "^[0-9a-fA-F]{40}$" }),
-              request: Type.String({ minLength: 1 }),
               provenance: Type.Object({
                 sourceType: Type.Union([
                   Type.Literal("pi"),
@@ -51,10 +50,20 @@ export default function discoveryExtension(pi: ExtensionAPI): void {
       if (!candidates) {
         throw new Error("submit_discovery received no candidates");
       }
-      writeFileSync(output, `${JSON.stringify({ candidates }, null, 2)}\n`, { flag: "wx" });
+      const exclusionsPath = process.env.SELFBENCH_DISCOVERY_EXCLUSIONS;
+      if (!exclusionsPath) {
+        throw new Error("SELFBENCH_DISCOVERY_EXCLUSIONS is not configured");
+      }
+      const excluded = new Set<number>(JSON.parse(readFileSync(exclusionsPath, "utf8")));
+      const accepted = candidates.filter(
+        ({ sourcePr }) => sourcePr !== undefined && !excluded.has(sourcePr),
+      );
+      writeFileSync(output, `${JSON.stringify({ candidates: accepted }, null, 2)}\n`, {
+        flag: "wx",
+      });
       return {
-        content: [{ type: "text", text: `Submitted ${candidates.length} candidates.` }],
-        details: { candidates },
+        content: [{ type: "text", text: `Submitted ${accepted.length} new candidates.` }],
+        details: { candidates: accepted },
       };
     },
   });

package/src/harbor-task.ts

@@ -6,7 +6,7 @@ import { runCommand } from "./process.js";
 import { patchPaths } from "./repair.js";
 
 const HARBOR_SCHEMA_VERSION = "1.4";
-const COMPILER_REVISION = 23;
+const COMPILER_REVISION = 24;
 
 export interface AuthoredTaskFiles {
   readonly definition: TaskDefinition;
@@ -70,6 +70,10 @@ export async function compileHarborTask(
   await Promise.all([
     cp(snapshot, join(environment, "repo.tar.gz")),
     cp(snapshot, join(tests, "repo.tar.gz")),
+    writeFile(
+      join(outputDirectory, "definition.json"),
+      `${JSON.stringify(task.definition, null, 2)}\n`,
+    ),
     writeFile(join(outputDirectory, "instruction.md"), `${task.definition.prompt.trim()}\n`),
     writeFile(join(solution, "gold.patch"), task.goldPatch),
     writeFile(join(solution, "solve.sh"), solutionScript()),
@@ -127,6 +131,7 @@ export async function refreshHarborTask(
   const preinstallGoldDependencies = dependencySetupPatch.length > 0;
   await writeFile(join(outputDirectory, "tests/test.sh"), testScript(definition, testPatch));
   await Promise.all([
+    writeFile(join(outputDirectory, "definition.json"), `${JSON.stringify(definition, null, 2)}\n`),
     writeFile(join(outputDirectory, "task.toml"), taskToml(definition)),
     writeFile(join(outputDirectory, "environment/Dockerfile"), agentDockerfile(definition)),
     writeFile(

package/src/project-paths.ts

@@ -0,0 +1,15 @@
+import { existsSync } from "node:fs";
+import { dirname, join, parse } from "node:path";
+import { fileURLToPath } from "node:url";
+
+export function projectRoot(moduleUrl: string): string {
+  let directory = dirname(fileURLToPath(moduleUrl));
+  const root = parse(directory).root;
+  while (directory !== root) {
+    if (existsSync(join(directory, "package.json"))) {
+      return directory;
+    }
+    directory = dirname(directory);
+  }
+  throw new Error(`could not find package.json above ${fileURLToPath(moduleUrl)}`);
+}

package/src/provenance-associations.ts

@@ -0,0 +1,311 @@
+import { readFile, writeFile } from "node:fs/promises";
+import { z } from "zod";
+import { assertPullRequestBelongsToRepository, githubRepository } from "./github.js";
+import { sha256 } from "./hash.js";
+import { runCommand } from "./process.js";
+import type {
+  LocalSessionMetadata,
+  ProvenanceMessage,
+  SessionProvenanceFormat,
+} from "./provenance.js";
+
+const localSourceTypeSchema = z.enum(["pi", "claude-code", "codex"]);
+type LocalSourceType = z.infer<typeof localSourceTypeSchema>;
+
+const associationMessageSchema = z
+  .object({
+    sourceType: localSourceTypeSchema,
+    sessionId: z.string().min(1),
+    messageIndex: z.number().int().nonnegative(),
+    contentSha256: z.string().regex(/^[0-9a-f]{64}$/),
+  })
+  .strict();
+
+export const provenanceAssociationManifestSchema = z
+  .object({
+    schemaVersion: z.literal(1),
+    repository: z.string().regex(/^[^/]+\/[^/]+$/),
+    sourcePr: z.number().int().positive(),
+    sourceUrl: z.string().url(),
+    messages: z.array(associationMessageSchema).min(1),
+  })
+  .strict();
+
+export type ProvenanceAssociationManifest = z.infer<typeof provenanceAssociationManifestSchema>;
+
+export interface AssociationSessionSummary {
+  readonly selector: string;
+  readonly sourceType: LocalSourceType;
+  readonly sessionId: string;
+  readonly messageCount: number;
+  readonly modifiedAt?: string;
+  readonly paths?: readonly string[];
+}
+
+export interface MergedPullRequest {
+  readonly sourcePr: number;
+  readonly sourceUrl: string;
+}
+
+export function associationSessionSummaries(
+  messages: readonly ProvenanceMessage[],
+  metadata: readonly LocalSessionMetadata[] = [],
+): AssociationSessionSummary[] {
+  const metadataBySelector = new Map<string, { modifiedAt: string; paths: Set<string> }>();
+  for (const session of metadata) {
+    if (!isLocalSourceType(session.sourceType)) {
+      continue;
+    }
+    const selector = sessionSelector(session.sourceType, session.sessionId);
+    const existing = metadataBySelector.get(selector);
+    metadataBySelector.set(selector, {
+      modifiedAt:
+        !existing || session.modifiedAt > existing.modifiedAt
+          ? session.modifiedAt
+          : existing.modifiedAt,
+      paths: new Set([...(existing?.paths ?? []), session.path]),
+    });
+  }
+
+  const counts = new Map<string, AssociationSessionSummary>();
+  for (const message of messages) {
+    if (!isLocalSourceType(message.sourceType)) {
+      continue;
+    }
+    const selector = sessionSelector(message.sourceType, message.sessionId);
+    const existing = counts.get(selector);
+    const sessionMetadata = metadataBySelector.get(selector);
+    counts.set(selector, {
+      selector,
+      sourceType: message.sourceType,
+      sessionId: message.sessionId,
+      messageCount: (existing?.messageCount ?? 0) + 1,
+      ...(sessionMetadata
+        ? {
+            modifiedAt: sessionMetadata.modifiedAt,
+            paths: [...sessionMetadata.paths].sort(),
+          }
+        : {}),
+    });
+  }
+  return [...counts.values()].sort(
+    (left, right) =>
+      (right.modifiedAt ?? "").localeCompare(left.modifiedAt ?? "") ||
+      left.selector.localeCompare(right.selector),
+  );
+}
+
+export function createProvenanceAssociationManifest(input: {
+  readonly repositoryUrl: string;
+  readonly pullRequest: MergedPullRequest;
+  readonly messages: readonly ProvenanceMessage[];
+  readonly sessionSelectors: readonly string[];
+}): ProvenanceAssociationManifest {
+  const repository = githubRepository(input.repositoryUrl);
+  assertPullRequestBelongsToRepository(
+    input.repositoryUrl,
+    input.pullRequest.sourceUrl,
+    input.pullRequest.sourcePr,
+  );
+  if (input.sessionSelectors.length === 0) {
+    throw new Error("at least one --session is required to create an association manifest");
+  }
+
+  const selected = new Set(input.sessionSelectors.map(parseSessionSelector));
+  if (selected.size !== input.sessionSelectors.length) {
+    throw new Error("each --session selector must be unique");
+  }
+  const available = new Set(
+    associationSessionSummaries(input.messages).map(({ selector }) => selector),
+  );
+  for (const selector of selected) {
+    if (!available.has(selector)) {
+      throw new Error(`local session ${selector} was not found for this repository`);
+    }
+  }
+
+  const messageKeys = new Set<string>();
+  const messages = input.messages
+    .filter(
+      (message): message is ProvenanceMessage & { sourceType: LocalSourceType } =>
+        isLocalSourceType(message.sourceType) &&
+        selected.has(sessionSelector(message.sourceType, message.sessionId)),
+    )
+    .map((message) => {
+      const key = messageKey(message.sourceType, message.sessionId, message.messageIndex);
+      if (messageKeys.has(key)) {
+        throw new Error(`local provenance contains duplicate message identity ${key}`);
+      }
+      messageKeys.add(key);
+      return {
+        sourceType: message.sourceType,
+        sessionId: message.sessionId,
+        messageIndex: message.messageIndex,
+        contentSha256: sha256(message.content),
+      };
+    })
+    .sort(compareAssociationMessages);
+
+  return provenanceAssociationManifestSchema.parse({
+    schemaVersion: 1,
+    repository,
+    sourcePr: input.pullRequest.sourcePr,
+    sourceUrl: input.pullRequest.sourceUrl,
+    messages,
+  });
+}
+
+export async function writeProvenanceAssociationManifest(
+  path: string,
+  manifest: ProvenanceAssociationManifest,
+): Promise<void> {
+  await writeFile(path, `${JSON.stringify(manifest, null, 2)}\n`, { flag: "wx", mode: 0o600 });
+}
+
+export async function applyProvenanceAssociationManifests(
+  messages: readonly ProvenanceMessage[],
+  repositoryUrl: string,
+  paths: readonly string[],
+): Promise<ProvenanceMessage[]> {
+  if (paths.length === 0) {
+    return [...messages];
+  }
+  const repository = githubRepository(repositoryUrl);
+  const manifests = await Promise.all(
+    paths.map(async (path) => {
+      const value = JSON.parse(await readFile(path, "utf8"));
+      return { path, manifest: provenanceAssociationManifestSchema.parse(value) };
+    }),
+  );
+  const messageByKey = new Map<string, ProvenanceMessage>();
+  for (const message of messages) {
+    const key = messageKey(message.sourceType, message.sessionId, message.messageIndex);
+    if (messageByKey.has(key)) {
+      throw new Error(`local provenance contains duplicate message identity ${key}`);
+    }
+    messageByKey.set(key, message);
+  }
+
+  const associations = new Map<string, MergedPullRequest>();
+  for (const { path, manifest } of manifests) {
+    if (manifest.repository.toLowerCase() !== repository) {
+      throw new Error(
+        `association manifest ${path} belongs to ${manifest.repository}, not ${repository}`,
+      );
+    }
+    assertPullRequestBelongsToRepository(repositoryUrl, manifest.sourceUrl, manifest.sourcePr);
+    for (const selector of new Set(
+      manifest.messages.map((message) => sessionSelector(message.sourceType, message.sessionId)),
+    )) {
+      const expected = manifest.messages
+        .filter((message) => sessionSelector(message.sourceType, message.sessionId) === selector)
+        .sort(compareAssociationMessages);
+      const current = messages
+        .filter(
+          (message): message is ProvenanceMessage & { sourceType: LocalSourceType } =>
+            isLocalSourceType(message.sourceType) &&
+            sessionSelector(message.sourceType, message.sessionId) === selector,
+        )
+        .map((message) => ({
+          sourceType: message.sourceType,
+          sessionId: message.sessionId,
+          messageIndex: message.messageIndex,
+          contentSha256: sha256(message.content),
+        }))
+        .sort(compareAssociationMessages);
+      if (JSON.stringify(current) !== JSON.stringify(expected)) {
+        throw new Error(`association manifest ${path} does not match local session ${selector}`);
+      }
+    }
+    for (const reference of manifest.messages) {
+      const key = messageKey(reference.sourceType, reference.sessionId, reference.messageIndex);
+      const message = messageByKey.get(key);
+      if (!message || sha256(message.content) !== reference.contentSha256) {
+        throw new Error(`association manifest ${path} does not match local message ${key}`);
+      }
+      if (associations.has(key)) {
+        throw new Error(`local message ${key} is associated more than once`);
+      }
+      associations.set(key, {
+        sourcePr: manifest.sourcePr,
+        sourceUrl: manifest.sourceUrl,
+      });
+    }
+  }
+
+  return messages.map((message) => {
+    const association = associations.get(
+      messageKey(message.sourceType, message.sessionId, message.messageIndex),
+    );
+    return association ? { ...message, ...association } : message;
+  });
+}
+
+export async function resolveMergedPullRequest(
+  repositoryUrl: string,
+  sourcePr: number,
+): Promise<MergedPullRequest> {
+  const repository = githubRepository(repositoryUrl);
+  const result = await runCommand("gh", [
+    "pr",
+    "view",
+    String(sourcePr),
+    "--repo",
+    repository,
+    "--json",
+    "number,url,state,mergedAt",
+  ]);
+  const parsed = z
+    .object({
+      number: z.number().int().positive(),
+      url: z.string().url(),
+      state: z.string(),
+      mergedAt: z.string().nullable(),
+    })
+    .parse(JSON.parse(result.stdout));
+  if (parsed.number !== sourcePr) {
+    throw new Error(`GitHub returned pull request ${parsed.number}; expected ${sourcePr}`);
+  }
+  assertPullRequestBelongsToRepository(repositoryUrl, parsed.url, sourcePr);
+  if (parsed.state !== "MERGED" || !parsed.mergedAt) {
+    throw new Error(`pull request ${repository}#${sourcePr} is not merged`);
+  }
+  return { sourcePr, sourceUrl: parsed.url };
+}
+
+function parseSessionSelector(value: string): string {
+  const separator = value.indexOf(":");
+  if (separator < 1 || separator === value.length - 1) {
+    throw new Error(`invalid session selector ${JSON.stringify(value)}; expected TYPE:SESSION_ID`);
+  }
+  const sourceType = value.slice(0, separator);
+  localSourceTypeSchema.parse(sourceType);
+  return sessionSelector(sourceType as LocalSourceType, value.slice(separator + 1));
+}
+
+function sessionSelector(sourceType: LocalSourceType, sessionId: string): string {
+  return `${sourceType}:${sessionId}`;
+}
+
+function messageKey(
+  sourceType: SessionProvenanceFormat | "github-pull-request",
+  sessionId: string,
+  messageIndex: number,
+): string {
+  return `${sourceType}:${sessionId}:${messageIndex}`;
+}
+
+function compareAssociationMessages(
+  left: z.infer<typeof associationMessageSchema>,
+  right: z.infer<typeof associationMessageSchema>,
+): number {
+  return (
+    left.sourceType.localeCompare(right.sourceType) ||
+    left.sessionId.localeCompare(right.sessionId) ||
+    left.messageIndex - right.messageIndex
+  );
+}
+
+function isLocalSourceType(value: string): value is LocalSourceType {
+  return value === "pi" || value === "claude-code" || value === "codex";
+}