@librechat/agents 3.1.86 → 3.1.88

package/src/scripts/compare_pi_vs_ours.ts

@@ -1,7 +1,7 @@
 /**
  * src/scripts/compare_pi_vs_ours.ts
  *
- * Side-by-side runs: pi-mono's `pi` CLI vs our local engine, same
+ * Side-by-side runs: pi-mono's `pi` CLI vs our AgentSession facade, same
  * task, same model, two parallel temp workspaces. We track:
  *
  *   - tool calls (name + args length, ordered)
@@ -10,7 +10,8 @@
  *   - whether the final on-disk state matches the expected outcome
  *
  * The tasks intentionally probe areas where we expect the local
- * engine to behave differently:
+ * engine to behave differently, while the preflight probes compare the
+ * programmatic session DX now exposed by the SDK:
  *
  *   T1 simple-edit       — both should one-shot
  *   T2 fuzzy-edit        — model emits an `oldText` with off-by-
@@ -32,25 +33,35 @@ config();
 import { spawn } from 'child_process';
 import { homedir, tmpdir } from 'os';
 import { join, resolve } from 'path';
-import { mkdtemp, readFile, rm, writeFile } from 'fs/promises';
+import {
+  copyFile,
+  mkdir,
+  mkdtemp,
+  readFile,
+  rm,
+  stat,
+  symlink,
+  writeFile,
+} from 'fs/promises';
 import { performance } from 'perf_hooks';
-import { HumanMessage, ToolMessage } from '@langchain/core/messages';
+import { AIMessage, HumanMessage, ToolMessage } from '@langchain/core/messages';
+import { MemorySaver } from '@langchain/langgraph';
 import type { BaseMessage } from '@langchain/core/messages';
+import type {
+  AgentSessionCheckpointing,
+  AgentSessionConfig,
+  AgentSessionRunResult,
+} from '@/session';
 import type * as t from '@/types';
-import { ChatModelStreamHandler, createContentAggregator } from '@/stream';
-import { ToolEndHandler, ModelEndHandler } from '@/events';
 import { getLLMConfig } from '@/utils/llmConfig';
-import { GraphEvents, Providers } from '@/common';
-import { Run } from '@/run';
+import { Providers, StepTypes } from '@/common';
+import { createAgentSession } from '@/session';
 
 const PROVIDER = Providers.ANTHROPIC;
 const MODEL = 'claude-sonnet-4-5';
 const PI_BIN =
   process.env.PI_BIN ??
-  resolve(
-    homedir(),
-    'Projects/pi-mono/packages/coding-agent/dist/cli.js'
-  );
+  resolve(homedir(), 'Projects/pi-mono/packages/coding-agent/dist/cli.js');
 
 interface Task {
   name: string;
@@ -95,6 +106,21 @@ interface RunOutcome {
   finalAssistant: string;
   errored: boolean;
   errorMessage?: string;
+  sessionEvents?: number;
+  sessionEntries?: number;
+  sessionPath?: string;
+}
+
+interface DxProbeResult {
+  ok: boolean;
+  detail: string;
+}
+
+interface DxProbe {
+  feature: string;
+  pi: string;
+  ours: string;
+  run: () => Promise<DxProbeResult>;
 }
 
 const TASKS: Task[] = [
@@ -102,8 +128,7 @@ const TASKS: Task[] = [
     name: 'T1 simple-edit',
     description: 'Single literal substitution in an existing file.',
     seed: {
-      'greet.py':
-        'def greet(name):\n    return f"Hello, {name}!"\n',
+      'greet.py': 'def greet(name):\n    return f"Hello, {name}!"\n',
     },
     prompt:
       'Edit greet.py: change the greeting from "Hello" to "Hi". ' +
@@ -163,12 +188,11 @@ const TASKS: Task[] = [
         null,
         2
       ),
-      'broken.ts':
-        'export const port: number = "not a number";\n',
+      'broken.ts': 'export const port: number = "not a number";\n',
     },
     prompt:
       'broken.ts has a type error. Fix it so the project typechecks cleanly. ' +
-      'After fixing, verify by running the project\'s typecheck (or `compile_check` if available). ' +
+      "After fixing, verify by running the project's typecheck (or `compile_check` if available). " +
       'Reply with "done".',
     verify: async (cwd) => {
       const text = await readFile(join(cwd, 'broken.ts'), 'utf8').catch(
@@ -224,8 +248,12 @@ const TASKS: Task[] = [
       'Rename the exported function `calc_total` to `calculateTotal` across src/lib.ts, ' +
       'src/index.ts, and src/index.test.ts. Update every reference. Reply "done" when finished.',
     verify: async (cwd) => {
-      const lib = await readFile(join(cwd, 'src/lib.ts'), 'utf8').catch(() => '');
-      const idx = await readFile(join(cwd, 'src/index.ts'), 'utf8').catch(() => '');
+      const lib = await readFile(join(cwd, 'src/lib.ts'), 'utf8').catch(
+        () => ''
+      );
+      const idx = await readFile(join(cwd, 'src/index.ts'), 'utf8').catch(
+        () => ''
+      );
       const tst = await readFile(join(cwd, 'src/index.test.ts'), 'utf8').catch(
         () => ''
       );
@@ -240,9 +268,7 @@ const TASKS: Task[] = [
       const ok = allRenamed && noOldName;
       return {
         ok,
-        detail: ok
-          ? ''
-          : `lib:\n${lib}\nindex:\n${idx}\ntest:\n${tst}`,
+        detail: ok ? '' : `lib:\n${lib}\nindex:\n${idx}\ntest:\n${tst}`,
       };
     },
   },
@@ -252,7 +278,6 @@ const TASKS: Task[] = [
       'Reads a PNG and describes it. Ours embeds via attachReadAttachments + image_url block; pi has no equivalent and is skipped.',
     seed: {},
     setup: async (cwd) => {
-      const { copyFile } = await import('fs/promises');
       // Use a real PNG (Anthropic refuses tiny 1x1 PNGs with "Could not
       // process image"). Try a few well-known macOS app icons; fall back to
       // any *.png we can find under /System.
@@ -277,7 +302,6 @@ const TASKS: Task[] = [
       // The verify step is soft — we just check the file is still on disk
       // (the agent shouldn't have deleted it) and the script-level error
       // tracking will fail this task if Anthropic refused the request.
-      const { stat } = await import('fs/promises');
       try {
         await stat(join(cwd, 'sample.png'));
         return { ok: true, detail: '' };
@@ -430,58 +454,23 @@ async function runPi(task: Task, cwd: string): Promise<RunOutcome> {
 /* Our local-engine runner                                             */
 /* ------------------------------------------------------------------ */
 
-async function runOurs(
-  task: Task,
-  cwd: string,
-  overrides: Partial<t.LocalExecutionConfig> = {}
-): Promise<RunOutcome> {
-  const start = performance.now();
-  const conversation: BaseMessage[] = [];
-  const observedToolCalls: ToolCallObservation[] = [];
-  let inputTokens = 0;
-  let outputTokens = 0;
-  let cacheReadTokens = 0;
-  let cacheWriteTokens = 0;
-
-  const { aggregateContent } = createContentAggregator();
-  const customHandlers = {
-    [GraphEvents.TOOL_END]: new ToolEndHandler(),
-    [GraphEvents.CHAT_MODEL_END]: new ModelEndHandler(),
-    [GraphEvents.CHAT_MODEL_STREAM]: new ChatModelStreamHandler(),
-    // ON_RUN_STEP must be forwarded too — without it the aggregator's
-    // `stepMap` is empty when ON_RUN_STEP_COMPLETED arrives and you
-    // get a "No run step or runId found for completed step event"
-    // warn for every tool call. The harness doesn't actually use the
-    // aggregated content, but feeding both events keeps logs clean.
-    [GraphEvents.ON_RUN_STEP]: {
-      handle: (
-        event: GraphEvents.ON_RUN_STEP,
-        data: t.StreamEventData
-      ): void => {
-        aggregateContent({ event, data: data as t.RunStep });
-      },
-    },
-    [GraphEvents.ON_RUN_STEP_COMPLETED]: {
-      handle: (
-        event: GraphEvents.ON_RUN_STEP_COMPLETED,
-        data: t.StreamEventData
-      ): void => {
-        aggregateContent({
-          event,
-          data: data as unknown as { result: t.ToolEndEvent },
-        });
-      },
-    },
-  };
-
+function createOursSessionConfig(params: {
+  cwd: string;
+  sessionPath?: string;
+  overrides?: Partial<t.LocalExecutionConfig>;
+  checkpointing?: AgentSessionCheckpointing;
+  ephemeral?: boolean;
+  humanInTheLoop?: t.HumanInTheLoopConfig;
+  graphConfig?: t.RunConfig['graphConfig'];
+}): AgentSessionConfig {
   const llmConfig = getLLMConfig(PROVIDER);
-  const runConfig: t.RunConfig = {
-    runId: `compare-${Date.now()}`,
-    graphConfig: {
+  return {
+    cwd: params.cwd,
+    sessionPath: params.sessionPath,
+    ephemeral: params.ephemeral,
+    checkpointing: params.checkpointing ?? false,
+    graphConfig: params.graphConfig ?? {
       type: 'standard',
-      // NB: in the legacy path Run.createLegacyGraph rebuilds
-      // `clientOptions` from llmConfig (it ignores graphConfig.clientOptions),
-      // so promptCache lives here and not on a separate clientOptions field.
       llmConfig: { ...llmConfig, model: MODEL, promptCache: true },
       instructions:
         'You are a coding assistant with local file tools. Use read_file, ' +
@@ -490,106 +479,199 @@ async function runOurs(
     toolExecution: {
       engine: 'local',
       local: {
-        cwd,
+        cwd: params.cwd,
         postEditSyntaxCheck: 'auto',
         timeoutMs: 30_000,
-        ...overrides,
+        ...params.overrides,
       },
     },
-    returnContent: true,
+    humanInTheLoop: params.humanInTheLoop,
     skipCleanup: true,
-    customHandlers,
   };
+}
 
-  let errored = false;
-  let errorMessage: string | undefined;
-  try {
-    const run = await Run.create<t.IState>(runConfig);
-    conversation.push(new HumanMessage(task.prompt));
-    const streamConfig = {
-      configurable: { provider: PROVIDER, thread_id: `compare-${Date.now()}` },
-      streamMode: 'values',
-      version: 'v2' as const,
-    };
-    await run.processStream(
-      { messages: conversation },
-      streamConfig as Parameters<typeof run.processStream>[1]
-    );
-    const finalMessages = run.getRunMessages();
-    if (finalMessages) {
-      conversation.push(...finalMessages);
+function getToolCallName(toolCall: t.AgentToolCall): string {
+  return 'function' in toolCall
+    ? toolCall.function.name
+    : (toolCall.name ?? '?');
+}
+
+function getToolCallArgs(toolCall: t.AgentToolCall): unknown {
+  return 'function' in toolCall ? toolCall.function.arguments : toolCall.args;
+}
+
+function collectToolCallsFromSteps(steps: t.RunStep[]): ToolCallObservation[] {
+  const calls: ToolCallObservation[] = [];
+  for (const step of steps) {
+    if (step.stepDetails.type !== StepTypes.TOOL_CALLS) {
+      continue;
+    }
+    for (const toolCall of step.stepDetails.tool_calls ?? []) {
+      calls.push({
+        name: getToolCallName(toolCall),
+        argsBytes: JSON.stringify(getToolCallArgs(toolCall) ?? {}).length,
+        isError: false,
+      });
     }
-  } catch (err) {
-    errored = true;
-    errorMessage = (err as Error).message.slice(0, 500);
   }
+  return calls;
+}
 
-  // Walk the conversation: tool calls live on AIMessage as `tool_calls`,
-  // tool results are ToolMessage entries (already chronologically next to them).
-  for (const msg of conversation) {
+function collectToolCallsFromMessages(
+  messages: BaseMessage[]
+): ToolCallObservation[] {
+  const calls: ToolCallObservation[] = [];
+  for (const msg of messages) {
     if (msg._getType() === 'ai') {
       const ai = msg as unknown as {
         tool_calls?: Array<{ name?: string; args?: unknown }>;
-        usage_metadata?: { input_tokens?: number; output_tokens?: number };
       };
-      if (ai.tool_calls != null) {
-        for (const tc of ai.tool_calls) {
-          observedToolCalls.push({
-            name: tc.name ?? '?',
-            argsBytes: JSON.stringify(tc.args ?? {}).length,
-            isError: false,
-          });
-        }
-      }
-      if (ai.usage_metadata != null) {
-        const reportedInput = ai.usage_metadata.input_tokens ?? 0;
-        outputTokens += ai.usage_metadata.output_tokens ?? 0;
-        const idu =
-          (ai.usage_metadata as unknown as {
-            input_token_details?: {
-              cache_read?: number;
-              cache_creation?: number;
-            };
-          }).input_token_details;
-        const cacheRead = idu?.cache_read ?? 0;
-        const cacheCreate = idu?.cache_creation ?? 0;
-        cacheReadTokens += cacheRead;
-        cacheWriteTokens += cacheCreate;
-        // The Anthropic adapter at src/llm/anthropic/utils/message_outputs.ts:31
-        // reports usage_metadata.input_tokens as the TOTAL prompt
-        // (input + cache_creation + cache_read), not just the uncached
-        // portion. Subtract cached fields so `inputTokens` here is
-        // apples-to-apples with pi's `input` field (uncached only).
-        const trulyUncached = Math.max(
-          0,
-          reportedInput - cacheRead - cacheCreate
-        );
-        inputTokens += trulyUncached;
+      for (const toolCall of ai.tool_calls ?? []) {
+        calls.push({
+          name: toolCall.name ?? '?',
+          argsBytes: JSON.stringify(toolCall.args ?? {}).length,
+          isError: false,
+        });
       }
+      continue;
     }
-    if (msg instanceof ToolMessage) {
-      if (msg.status === 'error' && observedToolCalls.length > 0) {
-        observedToolCalls[observedToolCalls.length - 1].isError = true;
-      }
+    if (
+      msg instanceof ToolMessage &&
+      msg.status === 'error' &&
+      calls.length > 0
+    ) {
+      calls[calls.length - 1].isError = true;
     }
   }
+  return calls;
+}
 
-  const lastAssistant = [...conversation]
+function collectTokenUsage(messages: BaseMessage[]): {
+  inputTokens: number;
+  outputTokens: number;
+  cacheReadTokens: number;
+  cacheWriteTokens: number;
+} {
+  let inputTokens = 0;
+  let outputTokens = 0;
+  let cacheReadTokens = 0;
+  let cacheWriteTokens = 0;
+  for (const msg of messages) {
+    if (msg._getType() !== 'ai') {
+      continue;
+    }
+    const ai = msg as unknown as {
+      usage_metadata?: { input_tokens?: number; output_tokens?: number };
+    };
+    if (ai.usage_metadata == null) {
+      continue;
+    }
+    const reportedInput = ai.usage_metadata.input_tokens ?? 0;
+    outputTokens += ai.usage_metadata.output_tokens ?? 0;
+    const inputTokenDetails = (
+      ai.usage_metadata as unknown as {
+        input_token_details?: {
+          cache_read?: number;
+          cache_creation?: number;
+        };
+      }
+    ).input_token_details;
+    const cacheRead = inputTokenDetails?.cache_read ?? 0;
+    const cacheCreate = inputTokenDetails?.cache_creation ?? 0;
+    cacheReadTokens += cacheRead;
+    cacheWriteTokens += cacheCreate;
+    inputTokens += Math.max(0, reportedInput - cacheRead - cacheCreate);
+  }
+  return {
+    inputTokens,
+    outputTokens,
+    cacheReadTokens,
+    cacheWriteTokens,
+  };
+}
+
+function getFinalAssistant(messages: BaseMessage[], fallback: string): string {
+  const lastAssistant = [...messages]
     .reverse()
-    .find((m) => m._getType() === 'ai');
-  let finalAssistant = '';
-  if (lastAssistant) {
-    const c = lastAssistant.content;
-    finalAssistant =
-      typeof c === 'string'
-        ? c
-        : Array.isArray(c)
-          ? c
-            .map((b) => ('text' in b ? b.text : ''))
-            .filter(Boolean)
-            .join(' ')
-          : '';
+    .find((message) => message._getType() === 'ai');
+  if (!lastAssistant) {
+    return fallback;
+  }
+  const content = lastAssistant.content;
+  if (typeof content === 'string') {
+    return content;
+  }
+  if (!Array.isArray(content)) {
+    return fallback;
+  }
+  return content
+    .map((block) => ('text' in block ? block.text : ''))
+    .filter(Boolean)
+    .join(' ');
+}
+
+async function runOurs(
+  task: Task,
+  cwd: string,
+  overrides: Partial<t.LocalExecutionConfig> = {}
+): Promise<RunOutcome> {
+  const start = performance.now();
+  let errored = false;
+  let errorMessage: string | undefined;
+  let result: AgentSessionRunResult | undefined;
+  let finalResultPromise: Promise<AgentSessionRunResult> | undefined;
+  let sessionEvents = 0;
+  let sessionEntries = 0;
+  const sessionPath = join(cwd, '.librechat-agent-session.jsonl');
+  try {
+    const session = await createAgentSession(
+      createOursSessionConfig({
+        cwd,
+        sessionPath,
+        overrides,
+        checkpointing: false,
+      })
+    );
+    const stream = session.stream(task.prompt, {
+      runId: `compare-${Date.now()}`,
+      threadId: `compare-${Date.now()}`,
+      config: {
+        configurable: { provider: PROVIDER },
+      },
+    });
+    finalResultPromise = stream.finalResult();
+    for await (const event of stream) {
+      sessionEvents = Math.max(sessionEvents, event.sequence + 1);
+    }
+    result = await finalResultPromise;
+    sessionEntries = session.getSessionStore()?.getEntries().length ?? 0;
+  } catch (err) {
+    await finalResultPromise?.catch(() => undefined);
+    errored = true;
+    errorMessage = (err as Error).message.slice(0, 500);
+  }
+
+  const messages = result?.messages ?? [];
+  const usage = collectTokenUsage(messages);
+  let observedToolCalls = collectToolCallsFromSteps(result?.steps ?? []);
+  const messageToolCalls = collectToolCallsFromMessages(messages);
+  if (observedToolCalls.length === 0) {
+    observedToolCalls = messageToolCalls;
+  }
+  for (let i = 0; i < observedToolCalls.length; i++) {
+    observedToolCalls[i].isError = messageToolCalls[i]?.isError ?? false;
   }
+  const inputTokens =
+    usage.inputTokens === 0 && result != null
+      ? result.usage.inputTokens
+      : usage.inputTokens;
+  const outputTokens =
+    usage.outputTokens === 0 && result != null
+      ? result.usage.outputTokens
+      : usage.outputTokens;
+  const cacheReadTokens = usage.cacheReadTokens;
+  const cacheWriteTokens = usage.cacheWriteTokens;
+  const finalAssistant = getFinalAssistant(messages, result?.text ?? '');
 
   // Sonnet 4.5 pricing (USD per 1M tokens). Pi computes its own cost; we
   // compute ours from the same per-turn breakdown so the cost columns are
@@ -615,15 +697,292 @@ async function runOurs(
     finalAssistant: finalAssistant.slice(0, 500),
     errored,
     errorMessage,
+    sessionEvents,
+    sessionEntries,
+    sessionPath,
   };
 }
 
+/* ------------------------------------------------------------------ */
+/* Programmatic DX probes                                              */
+/* ------------------------------------------------------------------ */
+
+async function withDxWorkspace<T>(
+  name: string,
+  run: (cwd: string) => Promise<T>
+): Promise<T> {
+  const cwd = await mkdtemp(join(tmpdir(), `lc-dx-${name}-`));
+  try {
+    return await run(cwd);
+  } finally {
+    await rm(cwd, { recursive: true, force: true });
+  }
+}
+
+async function probeSessionFacade(): Promise<DxProbeResult> {
+  return withDxWorkspace('facade', async (cwd) => {
+    const session = await createAgentSession(
+      createOursSessionConfig({
+        cwd,
+        sessionPath: join(cwd, 'facade.jsonl'),
+        checkpointing: false,
+      })
+    );
+    const methods: Array<keyof typeof session> = [
+      'run',
+      'stream',
+      'clone',
+      'fork',
+      'branch',
+      'compact',
+      'resumeInterrupt',
+    ];
+    const missing = methods.filter(
+      (method) => typeof session[method] !== 'function'
+    );
+    return {
+      ok: missing.length === 0,
+      detail:
+        missing.length === 0
+          ? 'session exposes run/stream plus lifecycle methods'
+          : `missing: ${missing.join(', ')}`,
+    };
+  });
+}
+
+async function probeJsonlTree(): Promise<DxProbeResult> {
+  return withDxWorkspace('jsonl', async (cwd) => {
+    const session = await createAgentSession(
+      createOursSessionConfig({
+        cwd,
+        sessionPath: join(cwd, 'tree.jsonl'),
+        checkpointing: false,
+      })
+    );
+    const store = session.getSessionStore();
+    if (!store) {
+      return { ok: false, detail: 'store was not created' };
+    }
+    const prompt = await store.appendMessage(
+      new HumanMessage('rename calc_total')
+    );
+    const reply = await store.appendMessage(new AIMessage('done'));
+    await store.setLabel(prompt.id, 'coding prompt');
+    const path = store.getPath(reply.id);
+    const entries = store.getEntries();
+    const ok =
+      path.length === 2 &&
+      store.getTree().length === 1 &&
+      store.getLabel(prompt.id) === 'coding prompt' &&
+      entries.some((entry) => entry.type === 'session_state');
+    return {
+      ok,
+      detail: `${entries.length} JSONL entries, active path length ${path.length}`,
+    };
+  });
+}
+
+async function probeForkCloneBranch(): Promise<DxProbeResult> {
+  return withDxWorkspace('branch', async (cwd) => {
+    const session = await createAgentSession(
+      createOursSessionConfig({
+        cwd,
+        sessionPath: join(cwd, 'branch.jsonl'),
+        checkpointing: false,
+      })
+    );
+    const store = session.getSessionStore();
+    if (!store) {
+      return { ok: false, detail: 'store was not created' };
+    }
+    const prompt = await store.appendMessage(new HumanMessage('turn one'));
+    const reply = await store.appendMessage(new AIMessage('reply one'));
+    const clone = await session.clone({ name: 'clone' });
+    const fork = await session.fork(reply.id, {
+      position: 'before',
+      name: 'fork-before-reply',
+    });
+    await session.branch(prompt.id, { position: 'at' });
+    const clonePathLength = clone.getSessionStore()?.getPath().length ?? 0;
+    const forkLeafId = fork.getSessionStore()?.getLeafEntry()?.id;
+    const activeLeafId = store.getLeafEntry()?.id;
+    const ok =
+      clonePathLength === 2 &&
+      forkLeafId === prompt.id &&
+      activeLeafId === prompt.id;
+    return {
+      ok,
+      detail:
+        `clone path ${clonePathLength}, fork leaf ${forkLeafId ?? 'none'}, ` +
+        `active leaf ${activeLeafId ?? 'none'}`,
+    };
+  });
+}
+
+async function probeResumeByPath(): Promise<DxProbeResult> {
+  return withDxWorkspace('resume', async (cwd) => {
+    const sessionPath = join(cwd, 'resume.jsonl');
+    const session = await createAgentSession(
+      createOursSessionConfig({ cwd, sessionPath, checkpointing: false })
+    );
+    const store = session.getSessionStore();
+    if (!store) {
+      return { ok: false, detail: 'store was not created' };
+    }
+    await store.appendMessage(new HumanMessage('persist me'));
+    const resumed = await createAgentSession(
+      createOursSessionConfig({ cwd, sessionPath, checkpointing: false })
+    );
+    const resumedMessages = resumed.getSessionStore()?.getMessages() ?? [];
+    const ok =
+      resumed.threadId === session.threadId && resumedMessages.length === 1;
+    return {
+      ok,
+      detail: `thread ${resumed.threadId}, messages ${resumedMessages.length}`,
+    };
+  });
+}
+
+async function probeCheckpointComposition(): Promise<DxProbeResult> {
+  return withDxWorkspace('checkpointing', async (cwd) => {
+    const checkpointer = new MemorySaver();
+    const injected = await createAgentSession(
+      createOursSessionConfig({
+        cwd,
+        ephemeral: true,
+        checkpointing: { checkpointer },
+      })
+    );
+    const disabled = await createAgentSession(
+      createOursSessionConfig({
+        cwd,
+        ephemeral: true,
+        checkpointing: false,
+        humanInTheLoop: { enabled: true },
+      })
+    );
+    const ok =
+      injected.getCheckpointer() === checkpointer &&
+      injected.getSessionStore() == null &&
+      disabled.getCheckpointer() == null &&
+      disabled.getSessionStore() == null;
+    return {
+      ok,
+      detail: ok
+        ? 'custom checkpointer injected; JSONL/checkpointing can both be disabled'
+        : 'composition check failed',
+    };
+  });
+}
+
+async function probeMultiAgentWrapping(): Promise<DxProbeResult> {
+  return withDxWorkspace('multi', async (cwd) => {
+    const clientOptions = {
+      modelName: MODEL,
+      apiKey: process.env.ANTHROPIC_API_KEY,
+    };
+    const session = await createAgentSession(
+      createOursSessionConfig({
+        cwd,
+        sessionPath: join(cwd, 'multi.jsonl'),
+        checkpointing: false,
+        graphConfig: {
+          type: 'multi-agent',
+          agents: [
+            {
+              agentId: 'supervisor',
+              provider: PROVIDER,
+              clientOptions,
+              instructions: 'Route coding work to the specialist.',
+            },
+            {
+              agentId: 'coder',
+              provider: PROVIDER,
+              clientOptions,
+              instructions: 'Make precise code changes.',
+            },
+          ],
+          edges: [
+            {
+              from: 'supervisor',
+              to: 'coder',
+              description: 'Delegate implementation work',
+              edgeType: 'handoff',
+            },
+          ],
+        },
+      })
+    );
+    const ok =
+      session.getSessionStore() != null &&
+      typeof session.stream === 'function' &&
+      typeof session.run === 'function';
+    return {
+      ok,
+      detail:
+        'AgentSession accepted a multi-agent handoff graph without special harness code',
+    };
+  });
+}
+
+const DX_PROBES: DxProbe[] = [
+  {
+    feature: 'Session facade',
+    pi: 'SDK sessions expose a high-level run loop',
+    ours: 'createAgentSession().run/stream wraps existing Run',
+    run: probeSessionFacade,
+  },
+  {
+    feature: 'Append-only JSONL tree',
+    pi: 'Native session JSONL tree',
+    ours: 'JsonlSessionStore v1 with entries, labels, state',
+    run: probeJsonlTree,
+  },
+  {
+    feature: 'Clone/fork/branch',
+    pi: 'Clone, fork, and branch from tree positions',
+    ours: 'clone(), fork(before/at), branch(before/at)',
+    run: probeForkCloneBranch,
+  },
+  {
+    feature: 'Resume',
+    pi: 'Resume by session id/path',
+    ours: 'resume/open by exact JSONL path with stable threadId',
+    run: probeResumeByPath,
+  },
+  {
+    feature: 'Composable state',
+    pi: 'Session log is separate from execution provider state',
+    ours: 'JSONL optional; LangGraph checkpointer injectable or off',
+    run: probeCheckpointComposition,
+  },
+  {
+    feature: 'Multi-agent/subagent surface',
+    pi: 'Coding-agent session loop',
+    ours: 'Same AgentSession facade wraps multi-agent graphs and tools',
+    run: probeMultiAgentWrapping,
+  },
+];
+
+async function runDxProbes(): Promise<boolean> {
+  console.log('\n================ DX SESSION PROBES ================');
+  let allPassed = true;
+  for (const probe of DX_PROBES) {
+    const result = await probe.run();
+    allPassed &&= result.ok;
+    console.log(`\n[${result.ok ? 'ok' : 'fail'}] ${probe.feature}`);
+    console.log(`  pi:   ${probe.pi}`);
+    console.log(`  ours: ${probe.ours}`);
+    console.log(`  live: ${result.detail}`);
+  }
+  return allPassed;
+}
+
 /* ------------------------------------------------------------------ */
 /* Harness                                                             */
 /* ------------------------------------------------------------------ */
 
 async function setupWorkspace(task: Task): Promise<string> {
-  const { mkdir } = await import('fs/promises');
   const dir = await mkdtemp(join(tmpdir(), 'lc-compare-'));
   for (const [relPath, content] of Object.entries(task.seed)) {
     const abs = join(dir, relPath);
@@ -642,7 +1001,6 @@ async function setupWorkspace(task: Task): Promise<string> {
 }
 
 async function symlinkRepoNodeModules(cwd: string): Promise<void> {
-  const { symlink } = await import('fs/promises');
   const repo = resolve(process.cwd(), 'node_modules');
   await symlink(repo, join(cwd, 'node_modules'), 'dir').catch(() => {
     /* fall through; tsc just won't be available */
@@ -655,9 +1013,7 @@ function summariseToolCalls(calls: ToolCallObservation[]): string {
   for (const c of calls) {
     grouped.set(c.name, (grouped.get(c.name) ?? 0) + 1);
   }
-  const inline = [...grouped.entries()]
-    .map(([n, c]) => `${n}×${c}`)
-    .join(', ');
+  const inline = [...grouped.entries()].map(([n, c]) => `${n}×${c}`).join(', ');
   const errors = calls.filter((c) => c.isError).length;
   return `${calls.length} call(s) [${inline}]${errors > 0 ? ` (${errors} errored)` : ''}`;
 }
@@ -677,10 +1033,28 @@ function avg(xs: number[]): number {
   return xs.length === 0 ? 0 : xs.reduce((a, b) => a + b, 0) / xs.length;
 }
 
+function selectTasks(tasks: Task[]): Task[] {
+  const raw = process.env.COMPARE_TASKS;
+  if (raw == null || raw.trim() === '') {
+    return tasks;
+  }
+  const requested = raw
+    .split(',')
+    .map((token) => token.trim().toLowerCase())
+    .filter(Boolean);
+  return tasks.filter((task) => {
+    const name = task.name.toLowerCase();
+    return requested.some((token) => name === token || name.startsWith(token));
+  });
+}
+
 async function runOnce(
   task: Task,
   side: 'pi' | 'ours'
-): Promise<{ outcome: RunOutcome; verify: { ok: boolean; detail: string } } | null> {
+): Promise<{
+  outcome: RunOutcome;
+  verify: { ok: boolean; detail: string };
+} | null> {
   if (task.skip === side) return null;
   const cwd = await setupWorkspace(task);
   const outcome =
@@ -702,10 +1076,28 @@ async function runOnce(
 
 async function main(): Promise<void> {
   const ITERS = Math.max(1, Number(process.env.COMPARE_ITERS ?? '1'));
+  const DX_ONLY = process.env.COMPARE_DX_ONLY === '1';
+  const selectedTasks = selectTasks(TASKS);
   console.log(`pi binary: ${PI_BIN}`);
   console.log(`model:     ${MODEL}`);
   console.log(`provider:  ${PROVIDER}`);
   console.log(`iters:     ${ITERS}`);
+  console.log(`dx only:   ${DX_ONLY ? 'yes' : 'no'}`);
+  console.log(
+    `tasks:     ${
+      DX_ONLY ? 'DX probes' : selectedTasks.map((task) => task.name).join(', ')
+    }`
+  );
+  const dxPassed = await runDxProbes();
+  if (DX_ONLY) {
+    process.exitCode = dxPassed ? 0 : 1;
+    return;
+  }
+  if (selectedTasks.length === 0) {
+    throw new Error(
+      `No tasks matched COMPARE_TASKS=${process.env.COMPARE_TASKS}`
+    );
+  }
 
   const results: Array<{
     task: Task;
@@ -713,7 +1105,7 @@ async function main(): Promise<void> {
     ours: AggregatedSide;
   }> = [];
 
-  for (const task of TASKS) {
+  for (const task of selectedTasks) {
     console.log(`\n========== ${task.name} ==========`);
     console.log(task.description);
 
@@ -733,7 +1125,9 @@ async function main(): Promise<void> {
             `cacheR=${piRes.outcome.cacheReadTokens} cacheW=${piRes.outcome.cacheWriteTokens} ` +
             `$${piRes.outcome.cost.toFixed(4)}`
         );
-        if (piRes.outcome.errored) console.log(`  err: ${piRes.outcome.errorMessage}`);
+        if (piRes.outcome.errored) {
+          console.log(`  err: ${piRes.outcome.errorMessage}`);
+        }
       } else {
         console.log(`[pi]${tag} (skipped)`);
       }
@@ -746,9 +1140,12 @@ async function main(): Promise<void> {
           `[ours]${tag} ${oursRes.outcome.errored ? 'ERROR' : oursRes.verify.ok ? 'ok' : 'fail'} ` +
             `${fmtMs(oursRes.outcome.wallMs)} ${summariseToolCalls(oursRes.outcome.toolCalls)} ` +
             `in=${oursRes.outcome.inputTokens} out=${oursRes.outcome.outputTokens} ` +
-            `cacheR=${oursRes.outcome.cacheReadTokens} cacheW=${oursRes.outcome.cacheWriteTokens}`
+            `cacheR=${oursRes.outcome.cacheReadTokens} cacheW=${oursRes.outcome.cacheWriteTokens} ` +
+            `events=${oursRes.outcome.sessionEvents ?? 0} jsonl=${oursRes.outcome.sessionEntries ?? 0}`
         );
-        if (oursRes.outcome.errored) console.log(`  err: ${oursRes.outcome.errorMessage}`);
+        if (oursRes.outcome.errored) {
+          console.log(`  err: ${oursRes.outcome.errorMessage}`);
+        }
       } else {
         console.log(`[ours]${tag} (skipped)`);
       }
@@ -798,22 +1195,40 @@ async function main(): Promise<void> {
     cols.push([r.task.name, 'verify', fmtVerify(r.pi), fmtVerify(r.ours)]);
     cols.push(['', 'wall', fmtSideMs(r.pi), fmtSideMs(r.ours)]);
     cols.push(['', 'tool calls', fmtSideCalls(r.pi), fmtSideCalls(r.ours)]);
-    cols.push(['', 'input new', fmtSide(r.pi, 'inputTokens'), fmtSide(r.ours, 'inputTokens')]);
-    cols.push(['', 'cache read', fmtSide(r.pi, 'cacheReadTokens'), fmtSide(r.ours, 'cacheReadTokens')]);
-    cols.push(['', 'cache write', fmtSide(r.pi, 'cacheWriteTokens'), fmtSide(r.ours, 'cacheWriteTokens')]);
-    cols.push(['', 'output tok', fmtSide(r.pi, 'outputTokens'), fmtSide(r.ours, 'outputTokens')]);
+    cols.push([
+      '',
+      'input new',
+      fmtSide(r.pi, 'inputTokens'),
+      fmtSide(r.ours, 'inputTokens'),
+    ]);
+    cols.push([
+      '',
+      'cache read',
+      fmtSide(r.pi, 'cacheReadTokens'),
+      fmtSide(r.ours, 'cacheReadTokens'),
+    ]);
+    cols.push([
+      '',
+      'cache write',
+      fmtSide(r.pi, 'cacheWriteTokens'),
+      fmtSide(r.ours, 'cacheWriteTokens'),
+    ]);
+    cols.push([
+      '',
+      'output tok',
+      fmtSide(r.pi, 'outputTokens'),
+      fmtSide(r.ours, 'outputTokens'),
+    ]);
     cols.push(['', 'cost', fmtCost(r.pi), fmtCost(r.ours)]);
+    cols.push(['', 'session events', 'N/A', fmtSide(r.ours, 'sessionEvents')]);
+    cols.push(['', 'jsonl entries', 'N/A', fmtSide(r.ours, 'sessionEntries')]);
   }
 
   const widths = [0, 0, 0, 0].map((_, i) =>
     Math.max(...cols.map((row) => row[i].length))
   );
   for (const row of cols) {
-    console.log(
-      row
-        .map((cell, i) => cell.padEnd(widths[i]))
-        .join('  ')
-    );
+    console.log(row.map((cell, i) => cell.padEnd(widths[i])).join('  '));
   }
 
   // Aggregate verify counts across all iters of all non-skipped tasks.
@@ -824,7 +1239,11 @@ async function main(): Promise<void> {
   console.log(
     `\nOverall: pi ${piPassed}/${piVerifies.length}, ours ${oursPassed}/${oursVerifies.length}.`
   );
-  if (piPassed < piVerifies.length || oursPassed < oursVerifies.length) {
+  if (
+    !dxPassed ||
+    piPassed < piVerifies.length ||
+    oursPassed < oursVerifies.length
+  ) {
     process.exitCode = 1;
   }
 }