@mastra/longmemeval 0.0.0-error-handler-fix-20251020202607 → 0.0.0-execa-dynamic-import-20260304221256

package/src/cli.ts

@@ -1,16 +1,35 @@
 #!/usr/bin/env node
 
+import 'dotenv/config';
 import { Command } from 'commander';
 import chalk from 'chalk';
-import { readFile, readdir } from 'fs/promises';
+import { readFile, readdir, writeFile } from 'fs/promises';
 import { join } from 'path';
 import { existsSync, statSync } from 'fs';
 import { execSync } from 'child_process';
 
 import { DatasetLoader } from './data/loader';
-import type { EvaluationResult, BenchmarkMetrics, QuestionType } from './data/types';
+import type { EvaluationResult, BenchmarkMetrics, QuestionType, DatasetType, MemoryConfigType } from './data/types';
 import { PrepareCommand } from './commands/prepare';
 import { RunCommand } from './commands/run';
+import { SyncCommand } from './commands/sync';
+import { CleanCommand } from './commands/clean';
+import { ObscureThreadIdsCommand } from './commands/obscure-thread-ids';
+import { SessionsCommand } from './commands/sessions';
+import { DeterministicIdsCommand } from './commands/deterministic-ids';
+import { ListPartialCommand } from './commands/list-partial';
+import { TokensCommand } from './commands/tokens';
+import { PrecomputeEmbeddingsCommand } from './commands/precompute-embeddings';
+import { FindProhibitedCommand } from './commands/find-prohibited';
+import { PartialResultsCommand } from './commands/partial-results';
+import {
+  getRunVariant,
+  resolveConfigAlias,
+  getConfigAliases,
+  getMemoryConfig,
+  RUN_VARIANTS,
+  CONFIG_ALIASES,
+} from './config';
 
 const program = new Command();
 
@@ -86,52 +105,326 @@ function calculateMetrics(results: EvaluationResult[]): BenchmarkMetrics {
   return metrics;
 }
 
+// Helper to load preparation token usage from om-debug.jsonl files
+interface PreparationTokenUsage {
+  observerInputTokens: number;
+  observerOutputTokens: number;
+  reflectorInputTokens: number;
+  reflectorOutputTokens: number;
+  totalInputTokens: number;
+  totalOutputTokens: number;
+  questionCount: number;
+}
+
+// Cache for per-question preparation token usage
+interface QuestionTokenCache {
+  [questionId: string]: {
+    mtime: number; // mtime of om-debug.jsonl when cached
+    usage: {
+      observerInputTokens: number;
+      observerOutputTokens: number;
+      reflectorInputTokens: number;
+      reflectorOutputTokens: number;
+    };
+  };
+}
+
+interface PrepTokenCache {
+  [memoryConfig: string]: QuestionTokenCache;
+}
+
+const PREP_TOKEN_CACHE_FILE = '.prep-token-cache.json';
+
+async function loadPrepTokenCache(): Promise<PrepTokenCache> {
+  try {
+    const content = await readFile(PREP_TOKEN_CACHE_FILE, 'utf-8');
+    return JSON.parse(content);
+  } catch {
+    return {};
+  }
+}
+
+async function savePrepTokenCache(cache: PrepTokenCache): Promise<void> {
+  await writeFile(PREP_TOKEN_CACHE_FILE, JSON.stringify(cache, null, 2));
+}
+
+async function parseDebugFileForTokens(debugFile: string): Promise<{
+  observerInputTokens: number;
+  observerOutputTokens: number;
+  reflectorInputTokens: number;
+  reflectorOutputTokens: number;
+} | null> {
+  try {
+    const content = await readFile(debugFile, 'utf-8');
+
+    const result = {
+      observerInputTokens: 0,
+      observerOutputTokens: 0,
+      reflectorInputTokens: 0,
+      reflectorOutputTokens: 0,
+    };
+
+    // Parse pretty-printed JSON objects
+    const events: any[] = [];
+    let currentJson = '';
+    let braceCount = 0;
+
+    for (const line of content.split('\n')) {
+      const trimmed = line.trim();
+
+      if (braceCount === 0 && trimmed.startsWith('{')) {
+        currentJson = line;
+        braceCount = (trimmed.match(/{/g) || []).length - (trimmed.match(/}/g) || []).length;
+      } else if (braceCount > 0) {
+        currentJson += '\n' + line;
+        braceCount += (trimmed.match(/{/g) || []).length;
+        braceCount -= (trimmed.match(/}/g) || []).length;
+      }
+
+      if (braceCount === 0 && currentJson) {
+        try {
+          events.push(JSON.parse(currentJson));
+        } catch {
+          // Skip malformed JSON
+        }
+        currentJson = '';
+      }
+    }
+
+    for (const event of events) {
+      if (event.usage && event.usage.inputTokens) {
+        if (event.type === 'observation_complete') {
+          result.observerInputTokens += event.usage.inputTokens || 0;
+          result.observerOutputTokens += event.usage.outputTokens || 0;
+        } else if (event.type === 'reflection_complete') {
+          result.reflectorInputTokens += event.usage.inputTokens || 0;
+          result.reflectorOutputTokens += event.usage.outputTokens || 0;
+        }
+      }
+    }
+
+    return result;
+  } catch {
+    return null;
+  }
+}
+
+async function loadPreparationTokenUsage(
+  preparedDataDir: string,
+  memoryConfig: string,
+  questionIds?: string[],
+): Promise<PreparationTokenUsage | null> {
+  const usage: PreparationTokenUsage = {
+    observerInputTokens: 0,
+    observerOutputTokens: 0,
+    reflectorInputTokens: 0,
+    reflectorOutputTokens: 0,
+    totalInputTokens: 0,
+    totalOutputTokens: 0,
+    questionCount: 0,
+  };
+
+  try {
+    const configDir = join(preparedDataDir, 'longmemeval_s', memoryConfig);
+    if (!existsSync(configDir)) {
+      return null;
+    }
+
+    // Load cache
+    const cache = await loadPrepTokenCache();
+    if (!cache[memoryConfig]) {
+      cache[memoryConfig] = {};
+    }
+    const configCache = cache[memoryConfig];
+    let cacheUpdated = false;
+
+    const questionDirs = await readdir(configDir);
+    const dirsToProcess = questionIds ? questionDirs.filter(d => questionIds.includes(d)) : questionDirs;
+
+    for (const questionDir of dirsToProcess) {
+      const debugFile = join(configDir, questionDir, 'om-debug.jsonl');
+      if (!existsSync(debugFile)) continue;
+
+      try {
+        const fileStat = statSync(debugFile);
+        const fileMtime = fileStat.mtimeMs;
+
+        // Check if we have a valid cache entry
+        const cached = configCache[questionDir];
+        if (cached && cached.mtime === fileMtime) {
+          // Use cached values
+          usage.observerInputTokens += cached.usage.observerInputTokens;
+          usage.observerOutputTokens += cached.usage.observerOutputTokens;
+          usage.reflectorInputTokens += cached.usage.reflectorInputTokens;
+          usage.reflectorOutputTokens += cached.usage.reflectorOutputTokens;
+          usage.questionCount++;
+        } else {
+          // Parse the file and update cache
+          const parsed = await parseDebugFileForTokens(debugFile);
+          if (parsed) {
+            usage.observerInputTokens += parsed.observerInputTokens;
+            usage.observerOutputTokens += parsed.observerOutputTokens;
+            usage.reflectorInputTokens += parsed.reflectorInputTokens;
+            usage.reflectorOutputTokens += parsed.reflectorOutputTokens;
+            usage.questionCount++;
+
+            // Update cache
+            configCache[questionDir] = {
+              mtime: fileMtime,
+              usage: parsed,
+            };
+            cacheUpdated = true;
+          }
+        }
+      } catch {
+        // Skip files that can't be read
+      }
+    }
+
+    // Save cache if updated
+    if (cacheUpdated) {
+      await savePrepTokenCache(cache);
+    }
+
+    usage.totalInputTokens = usage.observerInputTokens + usage.reflectorInputTokens;
+    usage.totalOutputTokens = usage.observerOutputTokens + usage.reflectorOutputTokens;
+
+    return usage.questionCount > 0 ? usage : null;
+  } catch {
+    return null;
+  }
+}
+
 program.name('longmemeval').description('LongMemEval benchmark for Mastra Memory').version('0.1.0');
 
+// Helper to show available variants and configs
+function showAvailableOptions() {
+  console.log(chalk.blue('\n📋 Available Run Variants:\n'));
+  for (const [name, variant] of Object.entries(RUN_VARIANTS)) {
+    const subsetInfo = variant.subset ? `${variant.subset} questions` : 'all questions';
+    console.log(
+      chalk.bold(`  ${name.padEnd(10)}`),
+      chalk.gray(`- ${variant.description}`),
+      chalk.dim(`(${subsetInfo}, prepare: ${variant.prepareConcurrency}x, bench: ${variant.benchConcurrency}x)`),
+    );
+  }
+
+  console.log(chalk.blue('\n📋 Available Memory Configs:\n'));
+  const configNames = getConfigAliases();
+  // Build reverse map: full name -> short aliases
+  const shortAliases: Record<string, string[]> = {};
+  for (const [alias, fullName] of Object.entries(CONFIG_ALIASES)) {
+    if (alias !== fullName) {
+      if (!shortAliases[fullName]) shortAliases[fullName] = [];
+      shortAliases[fullName].push(alias);
+    }
+  }
+  for (const configName of configNames) {
+    const configAliases = shortAliases[configName];
+    if (configAliases && configAliases.length > 0) {
+      console.log(chalk.bold(`  ${configName.padEnd(35)}`), chalk.gray(`(alias: ${configAliases.join(', ')})`));
+    } else {
+      console.log(chalk.bold(`  ${configName}`));
+    }
+  }
+
+  console.log(chalk.gray('\n\nUsage:'));
+  console.log(chalk.gray('  pnpm run prepare <config> [-v quick|full]'));
+  console.log(chalk.gray('  pnpm run bench <config> [-v quick|full]'));
+  console.log(chalk.gray('\nExamples:'));
+  console.log(chalk.gray('  pnpm run prepare om              # uses quick variant by default'));
+  console.log(chalk.gray('  pnpm run prepare om -v full      # full benchmark'));
+  console.log(chalk.gray('  pnpm run bench om'));
+  console.log(chalk.gray('  pnpm run bench om-glm -v full'));
+  console.log(chalk.gray('\nAdditional flags:'));
+  console.log(chalk.gray('  -v, --variant      Run variant (quick, full, rip) - default: quick'));
+  console.log(chalk.gray('  --offset <n>       Skip first n questions'));
+  console.log(chalk.gray('  --question-id <id> Process specific question'));
+  console.log(chalk.gray('  -y, --yes          Skip confirmation prompt'));
+}
+
 // Prepare command
 program
-  .command('prepare')
+  .command('prepare [config]')
   .description('Prepare LongMemEval data by processing through mock agents')
-  .option('-d, --dataset <dataset>', 'Dataset to use', 'longmemeval_s')
-  .option(
-    '-c, --memory-config <config>',
-    'Memory configuration (last-k, semantic-recall, semantic-recall-reranked, working-memory, working-memory-tailored, combined, combined-tailored)',
-    'semantic-recall',
-  )
+  .option('-v, --variant <variant>', 'Run variant (quick, full, rip)', 'quick')
   .option('-o, --output <dir>', 'Output directory for prepared data', './prepared-data')
-  .option('--subset <n>', 'Prepare only a subset of n questions', parseInt)
-  .option('--concurrency <n>', 'Number of questions to process in parallel', parseInt)
+  .option('--subset <n>', 'Override subset size', parseInt)
+  .option('--offset <n>', 'Skip first n questions', parseInt)
+  .option('--concurrency <n>', 'Override concurrency', parseInt)
   .option('--question-id <id>', 'Prepare a specific question by ID')
   .option('--resume-from-message-id <id>', 'Resume processing from a specific message ID')
   .option('--session-limit <n>', 'Limit processing to n sessions after resume point', parseInt)
   .option('--session-offset <n>', 'Start processing from the nth session (1-based)', parseInt)
-  .action(async options => {
+  .option('--from-failures [path]', 'Re-prepare failed questions (uses latest failures.json if no path given)')
+  .option('--older-than <duration>', 'Only re-prepare questions older than this duration (e.g., "1h", "30m", "2d")')
+  .option('--dry-run', 'Show what would be re-prepared without actually doing it (use with --from-failures)')
+  .option('--force-regenerate', 'Force regeneration by deleting existing prepared data first')
+  .option('-y, --yes', 'Skip confirmation prompt')
+  // Legacy options for backwards compatibility
+  .option('-d, --dataset <dataset>', 'Dataset to use (legacy)')
+  .option('-c, --memory-config <config>', 'Memory configuration (legacy)')
+  .action(async (config, options) => {
     try {
-      console.log(chalk.blue('\n🚀 LongMemEval Data Preparation\n'));
-      console.log(chalk.gray(`Dataset: ${options.dataset}`));
-      console.log(chalk.gray(`Memory Config: ${options.memoryConfig}`));
-      if (options.subset) {
-        console.log(chalk.gray(`Subset: ${options.subset} questions`));
+      // If no config provided, show help
+      if (!config && !options.dataset) {
+        showAvailableOptions();
+        process.exit(0);
       }
-      if (options.questionId) {
-        console.log(chalk.gray(`Question ID: ${options.questionId}`));
+
+      // Resolve variant and config (support legacy options)
+      let resolvedVariant: ReturnType<typeof getRunVariant>;
+      let resolvedConfig: string;
+      let dataset: string;
+
+      if (options.dataset && options.memoryConfig) {
+        // Legacy mode: use -d and -c flags
+        dataset = options.dataset;
+        resolvedConfig = options.memoryConfig;
+        resolvedVariant = {
+          name: 'custom',
+          description: 'Custom run',
+          dataset: dataset as any,
+          subset: options.subset,
+          prepareConcurrency: options.concurrency ?? 4,
+          benchConcurrency: options.concurrency ?? 10,
+        };
+      } else if (config) {
+        // New mode: config as positional, variant as flag (default: quick)
+        resolvedVariant = getRunVariant(options.variant);
+        resolvedConfig = resolveConfigAlias(config);
+        dataset = resolvedVariant.dataset;
+      } else {
+        console.error(chalk.red('Error: Please provide a <config>'));
+        console.error(chalk.gray('Run without arguments to see available options'));
+        process.exit(1);
       }
-      if (options.resumeFromMessageId) {
-        console.log(chalk.gray(`Resume from message ID: ${options.resumeFromMessageId}`));
+
+      // Apply overrides
+      const subset = options.subset ?? resolvedVariant.subset;
+      const perTypeCount = resolvedVariant.perTypeCount;
+      const concurrency = options.concurrency ?? resolvedVariant.prepareConcurrency;
+
+      console.log(chalk.blue('\n🚀 LongMemEval Data Preparation\n'));
+      console.log(chalk.gray(`Variant: ${resolvedVariant.name}`));
+      console.log(chalk.gray(`Dataset: ${dataset}`));
+      console.log(chalk.gray(`Memory Config: ${resolvedConfig}`));
+      console.log(chalk.gray(`Concurrency: ${concurrency}`));
+      if (perTypeCount) {
+        console.log(chalk.gray(`Stratified Sample: ${perTypeCount} per type`));
+      } else if (subset) {
+        console.log(chalk.gray(`Subset: ${subset} questions`));
       }
-      if (options.sessionLimit) {
-        console.log(chalk.gray(`Session limit: ${options.sessionLimit} sessions`));
+      if (options.offset) {
+        console.log(chalk.gray(`Offset: skipping first ${options.offset} questions`));
       }
-      if (options.sessionOffset) {
-        console.log(chalk.gray(`Session offset: Start from session ${options.sessionOffset}`));
+      if (options.questionId) {
+        console.log(chalk.gray(`Question ID: ${options.questionId}`));
       }
       console.log();
 
       // Check for OpenAI API key (needed for embeddings in semantic-recall)
-      if (
-        (options.memoryConfig === 'semantic-recall' || options.memoryConfig === 'combined') &&
-        !process.env.OPENAI_API_KEY
-      ) {
+      if ((resolvedConfig === 'semantic-recall' || resolvedConfig === 'combined') && !process.env.OPENAI_API_KEY) {
         console.error(chalk.red('Error: OPENAI_API_KEY environment variable is required for semantic recall'));
         console.error(chalk.gray('Please set it in your environment or .env file'));
         process.exit(1);
@@ -139,55 +432,74 @@ program
 
       // Validate dataset option
       const validDatasets = ['longmemeval_s', 'longmemeval_m', 'longmemeval_oracle'];
-      if (!validDatasets.includes(options.dataset)) {
-        console.error(chalk.red(`Invalid dataset: ${options.dataset}`));
+      if (!validDatasets.includes(dataset)) {
+        console.error(chalk.red(`Invalid dataset: ${dataset}`));
         console.error(chalk.gray(`Valid options: ${validDatasets.join(', ')}`));
         process.exit(1);
       }
 
       // Check if dataset exists and download if needed
-      await ensureDatasetExists(options.dataset);
-
-      // Show warning and ask for confirmation
-      console.log(chalk.yellow('\n⚠️  WARNING'));
-      console.log(chalk.yellow('━'.repeat(50)));
-      console.log(chalk.bold('\nPreparing this data can be very expensive!\n'));
-      console.log('This process will:');
-      console.log('  • Process many conversations through AI models');
-      console.log('  • Generate embeddings for semantic recall');
-      console.log('  • Potentially use significant API credits\n');
-      console.log(chalk.gray('Memory configs like "working-memory" and "combined" are especially costly.\n'));
-
-      const readline = await import('readline');
-      const rl = readline.createInterface({
-        input: process.stdin,
-        output: process.stdout,
-      });
-
-      const answer = await new Promise<string>(resolve => {
-        rl.question(chalk.bold('Are you sure you want to continue? (y/N): '), resolve);
-      });
-      rl.close();
+      await ensureDatasetExists(dataset);
 
-      if (answer.toLowerCase() !== 'y' && answer.toLowerCase() !== 'yes') {
-        console.log(chalk.gray('\nCancelled by user.'));
-        process.exit(0);
+      // For readOnlyConfig, no preparation is needed
+      const configDef = getMemoryConfig(resolvedConfig as MemoryConfigType);
+      if (configDef.readOnlyConfig && configDef.baseConfig) {
+        console.log(
+          chalk.green(`✓ Config "${resolvedConfig}" is read-only and uses data from "${configDef.baseConfig}"`),
+        );
+        console.log(chalk.gray(`  No preparation needed. Run benchmark directly with: pnpm bench ${resolvedConfig}`));
+        console.log(chalk.gray(`  Make sure "${configDef.baseConfig}" is prepared first.\n`));
+        return;
       }
 
-      console.log(); // Add spacing before continuing
+      // Show warning and ask for confirmation (skip if -y flag is passed)
+      if (!options.yes) {
+        console.log(chalk.yellow('\n⚠️  WARNING'));
+        console.log(chalk.yellow('━'.repeat(50)));
+        console.log(chalk.bold('\nPreparing this data can be very expensive!\n'));
+        console.log('This process will:');
+        console.log('  • Process many conversations through AI models');
+        console.log('  • Generate embeddings for semantic recall');
+        console.log('  • Potentially use significant API credits\n');
+        console.log(chalk.gray('Memory configs like "working-memory" and "combined" are especially costly.\n'));
+
+        const readline = await import('readline');
+        const rl = readline.createInterface({
+          input: process.stdin,
+          output: process.stdout,
+        });
+
+        const answer = await new Promise<string>(resolve => {
+          rl.question(chalk.bold('Are you sure you want to continue? (y/N): '), resolve);
+        });
+        rl.close();
+
+        if (answer.toLowerCase() !== 'y' && answer.toLowerCase() !== 'yes') {
+          console.log(chalk.gray('\nCancelled by user.'));
+          process.exit(0);
+        }
+
+        console.log(); // Add spacing before continuing
+      }
 
       // Run prepare command
       const prepareCommand = new PrepareCommand();
       await prepareCommand.run({
-        dataset: options.dataset,
-        memoryConfig: options.memoryConfig,
+        dataset: dataset as DatasetType,
+        memoryConfig: resolvedConfig as MemoryConfigType,
         outputDir: options.output,
-        subset: options.subset,
-        concurrency: options.concurrency,
+        subset: perTypeCount ? undefined : subset, // Don't use subset if using perTypeCount
+        perTypeCount,
+        offset: options.offset,
+        concurrency,
         questionId: options.questionId,
         resumeFromMessageId: options.resumeFromMessageId,
         sessionLimit: options.sessionLimit,
         sessionOffset: options.sessionOffset,
+        fromFailures: options.fromFailures,
+        olderThan: options.olderThan,
+        dryRun: options.dryRun,
+        forceRegenerate: options.forceRegenerate,
       });
 
       // Force exit after completion
@@ -200,25 +512,99 @@ program
     }
   });
 
-// Run benchmark command
+// Run benchmark command (aliased as 'bench' too)
 program
-  .command('run')
+  .command('run [config]')
+  .alias('bench')
   .description('Run LongMemEval benchmark using prepared data')
-  .requiredOption('-d, --dataset <dataset>', 'Dataset to use (longmemeval_s, longmemeval_m, longmemeval_oracle)')
-  .requiredOption('-m, --model <model>', 'Model to use (e.g., gpt-4o, claude-3-opus)')
-  .option(
-    '-c, --memory-config <config>',
-    'Memory configuration (last-k, semantic-recall, semantic-recall-reranked, working-memory, working-memory-tailored, combined, combined-tailored)',
-    'semantic-recall',
-  )
+  .option('-v, --variant <variant>', 'Run variant (quick, full, rip, sample, sample-comb)', 'quick')
   .option('-o, --output <dir>', 'Output directory for results', './results')
   .option('--prepared-data <dir>', 'Directory containing prepared data', './prepared-data')
-  .option('--subset <n>', 'Run on subset of n questions', parseInt)
-  .option('--concurrency <n>', 'Number of parallel requests (default: 5)', parseInt)
+  .option('--subset <n>', 'Override subset size', parseInt)
+  .option('--offset <n>', 'Skip first n questions', parseInt)
+  .option('--concurrency <n>', 'Override concurrency', parseInt)
   .option('--question-id <id>', 'Focus on a specific question by ID')
-  .action(async options => {
+  .option('-t, --type <type>', 'Filter to a specific question type (e.g., multi-session, knowledge-update)')
+  .option('--comb-offset <n>', 'Comb sampling: stride between questions (for sample-comb variant)', parseInt)
+  .option('--start-offset <n>', 'Comb sampling: starting index (for sample-comb variant)', parseInt)
+  .option('--no-fixed', 'Skip improved/fixed question evaluation')
+  .option(
+    '--from-failures [path]',
+    'Re-run only failed questions from a previous run (path to failures.json or "latest")',
+  )
+  .option(
+    '--older-than <duration>',
+    'With --from-failures: only re-run questions prepared before this duration (e.g., "1h", "30m", "2d")',
+  )
+  .option('--resume [run-id]', 'Resume a partial run (omit run-id to auto-detect most recent)')
+  // Legacy options for backwards compatibility
+  .option('-d, --dataset <dataset>', 'Dataset to use (legacy)')
+  .option('-c, --memory-config <config>', 'Memory configuration (legacy)')
+  .action(async (config, options) => {
     try {
+      // If no config provided, show help
+      if (!config && !options.dataset) {
+        showAvailableOptions();
+        process.exit(0);
+      }
+
+      // Resolve variant and config (support legacy options)
+      let resolvedVariant: ReturnType<typeof getRunVariant>;
+      let resolvedConfig: string;
+      let dataset: string;
+
+      if (options.dataset && options.memoryConfig) {
+        // Legacy mode: use -d and -c flags
+        dataset = options.dataset;
+        resolvedConfig = options.memoryConfig;
+        resolvedVariant = {
+          name: 'custom',
+          description: 'Custom run',
+          dataset: dataset as any,
+          subset: options.subset,
+          prepareConcurrency: options.concurrency ?? 4,
+          benchConcurrency: options.concurrency ?? 10,
+        };
+      } else if (config) {
+        // New mode: config as positional, variant as flag (default: quick)
+        resolvedVariant = getRunVariant(options.variant);
+        resolvedConfig = resolveConfigAlias(config);
+        dataset = resolvedVariant.dataset;
+      } else {
+        console.error(chalk.red('Error: Please provide a <config>'));
+        console.error(chalk.gray('Run without arguments to see available options'));
+        process.exit(1);
+      }
+
+      // Apply overrides
+      const subset = options.subset ?? resolvedVariant.subset;
+      const perTypeCount = resolvedVariant.perTypeCount;
+      const combSampleSize = resolvedVariant.combSampleSize;
+      const combOffset = options.combOffset ?? resolvedVariant.combOffset;
+      const combStartOffset = options.startOffset ?? resolvedVariant.combStartOffset;
+      const concurrency = options.concurrency ?? resolvedVariant.benchConcurrency;
+
       console.log(chalk.blue('\n🚀 LongMemEval Benchmark Runner\n'));
+      console.log(chalk.gray(`Variant: ${resolvedVariant.name}`));
+      console.log(chalk.gray(`Dataset: ${dataset}`));
+      console.log(chalk.gray(`Memory Config: ${resolvedConfig}`));
+      console.log(chalk.gray(`Concurrency: ${concurrency}`));
+      if (combSampleSize) {
+        console.log(
+          chalk.gray(`Comb Sample: ${combSampleSize} per type (offset=${combOffset}, start=${combStartOffset})`),
+        );
+      } else if (perTypeCount) {
+        console.log(chalk.gray(`Stratified Sample: ${perTypeCount} per type`));
+      } else if (subset) {
+        console.log(chalk.gray(`Subset: ${subset} questions`));
+      }
+      if (options.offset) {
+        console.log(chalk.gray(`Offset: skipping first ${options.offset} questions`));
+      }
+      if (options.questionId) {
+        console.log(chalk.gray(`Question ID: ${options.questionId}`));
+      }
+      console.log();
 
       // Check for OpenAI API key
       if (!process.env.OPENAI_API_KEY) {
@@ -229,8 +615,8 @@ program
 
       // Validate dataset option
       const validDatasets = ['longmemeval_s', 'longmemeval_m', 'longmemeval_oracle'];
-      if (!validDatasets.includes(options.dataset)) {
-        console.error(chalk.red(`Invalid dataset: ${options.dataset}`));
+      if (!validDatasets.includes(dataset)) {
+        console.error(chalk.red(`Invalid dataset: ${dataset}`));
         console.error(chalk.gray(`Valid options: ${validDatasets.join(', ')}`));
         process.exit(1);
       }
@@ -238,14 +624,23 @@ program
       // Run benchmark using prepared data
       const runCommand = new RunCommand();
       await runCommand.run({
-        dataset: options.dataset,
-        memoryConfig: options.memoryConfig,
-        model: options.model,
+        dataset: dataset as DatasetType,
+        memoryConfig: resolvedConfig as MemoryConfigType,
         preparedDataDir: options.preparedData,
         outputDir: options.output,
-        subset: options.subset,
-        concurrency: options.concurrency,
+        subset: combSampleSize || perTypeCount ? undefined : subset, // Don't use subset if using sampling
+        perTypeCount,
+        combSampleSize,
+        combOffset,
+        combStartOffset,
+        offset: options.offset,
+        concurrency,
         questionId: options.questionId,
+        questionType: options.type,
+        skipFixed: options.fixed === false, // --no-fixed sets options.fixed to false
+        fromFailures: options.fromFailures,
+        resume: options.resume,
+        olderThan: options.olderThan,
       });
 
       // Force exit after completion
@@ -387,14 +782,104 @@ program
     }
   });
 
+// Sync command - sync improved_question/improved_answer from dataset to prepared meta.json files
+program
+  .command('sync')
+  .description('Sync improved_question and improved_answer from dataset JSON to prepared meta.json files')
+  .requiredOption('-d, --dataset <dataset>', 'Dataset to sync from (longmemeval_s, longmemeval_m, longmemeval_oracle)')
+  .option('-c, --memory-config <config>', 'Memory configuration', 'working-memory')
+  .option('--prepared-data <dir>', 'Directory containing prepared data', './prepared-data')
+  .action(async options => {
+    try {
+      console.log(chalk.blue('\n🔄 Syncing Improved Questions/Answers\n'));
+
+      // Validate dataset option
+      const validDatasets = ['longmemeval_s', 'longmemeval_m', 'longmemeval_oracle'];
+      if (!validDatasets.includes(options.dataset)) {
+        console.error(chalk.red(`Invalid dataset: ${options.dataset}`));
+        console.error(chalk.gray(`Valid options: ${validDatasets.join(', ')}`));
+        process.exit(1);
+      }
+
+      const syncCommand = new SyncCommand();
+      await syncCommand.run({
+        dataset: options.dataset,
+        memoryConfig: options.memoryConfig,
+        preparedDataDir: options.preparedData,
+      });
+    } catch (error) {
+      console.error(chalk.red('\nError:'), error);
+      process.exit(1);
+    }
+  });
+
 // Results command - shows latest results for each memory configuration
+program
+  .command('clean')
+  .description('Delete prepared data by offset/subset')
+  .requiredOption('-d, --dataset <name>', 'Dataset name (longmemeval_s, longmemeval_m, longmemeval_oracle)')
+  .option('-c, --memory-config <type>', 'Memory configuration to clean', 'observational-memory')
+  .option('-o, --offset <n>', 'Skip first N questions (delete from N+1 onwards)', parseInt)
+  .option('-s, --subset <n>', 'Only delete N questions (after offset)', parseInt)
+  .option('-q, --question-id <id>', 'Delete a specific question by ID')
+  .option('-p, --prepared-data <dir>', 'Prepared data directory', './prepared-data')
+  .option('--partial', 'Only delete partially prepared questions (have progress.json but no meta.json)')
+  .option('--dry-run', 'Show what would be deleted without actually deleting')
+  .action(async options => {
+    try {
+      const cleanCommand = new CleanCommand();
+      await cleanCommand.run({
+        dataset: options.dataset,
+        memoryConfig: options.memoryConfig,
+        preparedDataDir: options.preparedData,
+        offset: options.offset,
+        subset: options.subset,
+        questionId: options.questionId,
+        dryRun: options.dryRun,
+        partial: options.partial,
+      });
+    } catch (error) {
+      console.error(chalk.red('Error:'), error instanceof Error ? error.message : error);
+      process.exit(1);
+    }
+  });
+
+// Obscure thread IDs command - replace thread IDs with hashed versions to prevent LLM bias
+program
+  .command('obscure-thread-ids')
+  .description('Replace thread IDs in prepared om.json files with hashed versions to prevent LLM bias')
+  .requiredOption('-d, --dataset <name>', 'Dataset name (longmemeval_s, longmemeval_m, longmemeval_oracle)')
+  .option('-c, --memory-config <type>', 'Memory configuration', 'observational-memory')
+  .option('-p, --prepared-data <dir>', 'Prepared data directory', './prepared-data')
+  .option('--dry-run', 'Show what would be changed without actually modifying files')
+  .action(async options => {
+    try {
+      const obscureCommand = new ObscureThreadIdsCommand();
+      await obscureCommand.run({
+        dataset: options.dataset,
+        memoryConfig: options.memoryConfig,
+        preparedDataDir: options.preparedData,
+        dryRun: options.dryRun,
+      });
+    } catch (error) {
+      console.error(chalk.red('Error:'), error instanceof Error ? error.message : error);
+      process.exit(1);
+    }
+  });
+
 program
   .command('results')
   .description('Show latest benchmark results for each memory configuration')
   .option('-r, --results <dir>', 'Results directory', './results')
   .option('-d, --dataset <dataset>', 'Filter by dataset')
-  .option('-a, --all', 'Show all results, not just latest')
+  .option('-l, --latest', 'Show only the latest result per config')
+  .option('--min-questions <n>', 'Minimum questions to include (default: 20)', parseInt)
+  .option('-s, --sort <by>', 'Sort by: date, accuracy, fixed (default: date)', 'date')
+  .option('--no-fixed', 'Hide fixed accuracy numbers')
   .action(async options => {
+    const minQuestions = options.minQuestions ?? 20;
+    const sortBy = options.sort ?? 'date';
+    const showFixed = options.fixed !== false;
     try {
       console.log(chalk.blue('\n📊 Benchmark Results Summary\n'));
 
@@ -414,6 +899,7 @@ program
         metrics: any;
         config: any;
         timestamp: string;
+        metricsPath: string;
       }> = [];
 
       // First, try new structure
@@ -437,11 +923,17 @@ program
                 continue;
               }
 
+              // Filter by minimum questions
+              if (data.total_questions < minQuestions) {
+                continue;
+              }
+
               allRuns.push({
                 runId: runDir,
                 metrics: data,
                 config: data.config,
                 timestamp: data.timestamp,
+                metricsPath,
               });
             } catch (error) {
               // Skip runs with missing or invalid metrics
@@ -465,11 +957,17 @@ program
             continue;
           }
 
+          // Filter by minimum questions
+          if (data.total_questions < minQuestions) {
+            continue;
+          }
+
           allRuns.push({
             runId: runDir,
             metrics: data,
             config: data.config,
             timestamp: data.timestamp,
+            metricsPath,
           });
         } catch (error) {
           // Skip runs with missing or invalid metrics
@@ -481,106 +979,175 @@ program
         return;
       }
 
-      // Group by memory configuration
-      const byMemoryConfig = new Map<string, typeof allRuns>();
-      for (const run of allRuns) {
-        const key = `${run.config.dataset}_${run.config.memoryConfig}`;
-        if (!byMemoryConfig.has(key)) {
-          byMemoryConfig.set(key, []);
+      // Sort all runs based on sortBy option (best/newest at bottom for terminal viewing)
+      let runsToShow = [...allRuns];
+
+      if (options.latest) {
+        // Group by config and take only the latest from each
+        const byConfig = new Map<string, (typeof allRuns)[0]>();
+        for (const run of allRuns) {
+          const key = `${run.config.dataset}_${run.config.memoryConfig}`;
+          const existing = byConfig.get(key);
+          if (!existing || run.timestamp > existing.timestamp) {
+            byConfig.set(key, run);
+          }
         }
-        byMemoryConfig.get(key)!.push(run);
+        runsToShow = Array.from(byConfig.values());
       }
 
-      // Sort groups by worst performing to best performing (based on latest run)
-      const sortedConfigs = Array.from(byMemoryConfig.entries()).sort(([_aKey, aRuns], [_bKey, bRuns]) => {
-        // Get latest run for each config (already sorted by timestamp)
-        const aLatest = aRuns.sort((a, b) => b.timestamp.localeCompare(a.timestamp))[0];
-        const bLatest = bRuns.sort((a, b) => b.timestamp.localeCompare(a.timestamp))[0];
-
-        // Sort by overall accuracy (worst first)
-        return aLatest.metrics.overall_accuracy - bLatest.metrics.overall_accuracy;
+      // Sort runs (oldest/worst first, newest/best at bottom for terminal viewing)
+      runsToShow.sort((a, b) => {
+        if (sortBy === 'date') {
+          // Sort by date (oldest first, newest at bottom)
+          return a.timestamp.localeCompare(b.timestamp);
+        } else if (sortBy === 'fixed') {
+          // Sort by fixed accuracy (worst first, best at bottom), fall back to vanilla if no fixed
+          const aFixed = a.metrics.fixed_overall_accuracy ?? a.metrics.overall_accuracy;
+          const bFixed = b.metrics.fixed_overall_accuracy ?? b.metrics.overall_accuracy;
+          return aFixed - bFixed;
+        } else {
+          // Sort by vanilla accuracy (worst first, best at bottom)
+          return a.metrics.overall_accuracy - b.metrics.overall_accuracy;
+        }
       });
 
-      for (const [_configKey, runs] of sortedConfigs) {
-        // Sort runs by timestamp (newest first)
-        runs.sort((a, b) => b.timestamp.localeCompare(a.timestamp));
-
-        // Show latest or all
-        const runsToShow = options.all ? runs : [runs[0]];
-
-        for (const run of runsToShow) {
-          // Get terminal width, default to 80 if not available
-          const terminalWidth = process.stdout.columns || 80;
-          const lineWidth = Math.min(terminalWidth - 1, 80); // Cap at 80 for readability
-
-          console.log(chalk.bold('\n' + '═'.repeat(lineWidth) + '\n'));
-
-          // Configuration header
-          console.log(chalk.bold('Configuration:\n'));
-          console.log(chalk.gray('Dataset:'), chalk.cyan(run.config.dataset));
-          console.log(chalk.gray('Model:'), chalk.cyan(run.config.model));
-          console.log(chalk.gray('Memory Config:'), chalk.cyan(run.config.memoryConfig));
-          if (run.config.subset) {
-            console.log(chalk.gray('Subset:'), chalk.cyan(`${run.config.subset} questions`));
-          }
-          console.log(chalk.gray('Run ID:'), chalk.dim(run.runId));
-          console.log(chalk.gray('Timestamp:'), chalk.dim(new Date(run.timestamp).toLocaleString()));
-          console.log(chalk.gray('─'.repeat(Math.min(lineWidth, 60))));
-
-          // Display metrics using same format as regular runs
-          const metrics = run.metrics;
+      for (const run of runsToShow) {
+        // Get terminal width, default to 80 if not available
+        const terminalWidth = process.stdout.columns || 80;
+        const lineWidth = Math.min(terminalWidth - 1, 80); // Cap at 80 for readability
 
-          // Recalculate overall accuracy using the new formula (average of type averages)
-          const typeAccuracies = Object.values(metrics.accuracy_by_type).map((t: any) => t.accuracy);
-          const recalculatedOverall =
-            typeAccuracies.length > 0 ? typeAccuracies.reduce((sum, acc) => sum + acc, 0) / typeAccuracies.length : 0;
-          metrics.overall_accuracy = recalculatedOverall;
+        console.log(chalk.bold('\n' + '═'.repeat(lineWidth) + '\n'));
 
-          // Question type breakdown
+        // Configuration header
+        console.log(chalk.bold('Configuration:\n'));
+        console.log(chalk.gray('Dataset:'), chalk.cyan(run.config.dataset));
+        console.log(chalk.gray('Model:'), chalk.cyan(run.config.model));
+        console.log(chalk.gray('Memory Config:'), chalk.cyan(run.config.memoryConfig));
+        if (run.config.subset) {
+          console.log(chalk.gray('Subset:'), chalk.cyan(`${run.config.subset} questions`));
+        }
+        console.log(chalk.gray('Run ID:'), chalk.dim(run.runId));
+        console.log(chalk.gray('Timestamp:'), chalk.dim(new Date(run.timestamp).toLocaleString()));
+        // Make path relative to cwd
+        const relativePath = require('path').relative(process.cwd(), require('path').resolve(run.metricsPath));
+        console.log(chalk.gray('Metrics:'), chalk.blue(relativePath));
+        console.log(chalk.gray('─'.repeat(Math.min(lineWidth, 60))));
+
+        // Display metrics using same format as regular runs
+        const metrics = run.metrics;
+
+        // Recalculate overall accuracy using the new formula (average of type averages)
+        const typeAccuracies = Object.values(metrics.accuracy_by_type).map((t: any) => t.accuracy);
+        const recalculatedOverall =
+          typeAccuracies.length > 0 ? typeAccuracies.reduce((sum, acc) => sum + acc, 0) / typeAccuracies.length : 0;
+        metrics.overall_accuracy = recalculatedOverall;
+
+        // Check if fixed accuracy data exists
+        const hasFixedAccuracy =
+          showFixed && metrics.fixed_accuracy_by_type && Object.keys(metrics.fixed_accuracy_by_type).length > 0;
+
+        // Question type breakdown
+        if (hasFixedAccuracy) {
+          console.log(chalk.bold('\nAccuracy by Question Type:'), chalk.gray('(vanilla → fixed)'));
+        } else {
           console.log(chalk.bold('\nAccuracy by Question Type:'));
+        }
 
-          // Sort question types alphabetically
-          const sortedTypes = Object.entries(metrics.accuracy_by_type).sort(([a], [b]) => a.localeCompare(b));
+        // Sort question types alphabetically
+        const sortedTypes = Object.entries(metrics.accuracy_by_type).sort(([a], [b]) => a.localeCompare(b));
 
-          for (const [type, typeMetrics] of sortedTypes) {
-            const { correct, total, accuracy } = typeMetrics as any;
-            const typeColor = accuracy >= 0.8 ? 'green' : accuracy >= 0.6 ? 'yellow' : 'red';
+        for (const [type, typeMetrics] of sortedTypes) {
+          const { correct, total, accuracy } = typeMetrics as any;
+          const typeColor = accuracy >= 0.8 ? 'green' : accuracy >= 0.6 ? 'yellow' : 'red';
 
-            // Create a simple progress bar
-            const barLength = 20;
-            const filledLength = Math.round(accuracy * barLength);
-            const bar = '█'.repeat(filledLength) + '░'.repeat(barLength - filledLength);
+          // Create a simple progress bar
+          const barLength = 20;
+          const filledLength = Math.round(accuracy * barLength);
+          const bar = '█'.repeat(filledLength) + '░'.repeat(barLength - filledLength);
 
-            console.log(
-              chalk.gray(`  ${type.padEnd(25)}:`),
-              chalk[typeColor](`${(accuracy * 100).toFixed(1).padStart(5)}%`),
-              chalk.gray(`[${bar}]`),
-              chalk.gray(`(${correct}/${total})`),
-            );
+          let fixedPart = '';
+          if (hasFixedAccuracy && metrics.fixed_accuracy_by_type[type]) {
+            const fixedMetrics = metrics.fixed_accuracy_by_type[type] as any;
+            const fixedColor = fixedMetrics.accuracy >= 0.8 ? 'green' : fixedMetrics.accuracy >= 0.6 ? 'yellow' : 'red';
+            fixedPart = chalk.gray(' → ') + chalk[fixedColor](`${(fixedMetrics.accuracy * 100).toFixed(1)}%`);
           }
 
-          // Abstention is hidden - it tests LLM reasoning ability rather than memory system performance
+          console.log(
+            chalk.gray(`  ${type.padEnd(25)}:`),
+            chalk[typeColor](`${(accuracy * 100).toFixed(1).padStart(5)}%`) + fixedPart,
+            chalk.gray(`[${bar}]`),
+            chalk.gray(`(${correct}/${total})`),
+          );
+        }
 
-          // Overall summary at the bottom
-          console.log();
-          const accuracyColor =
-            metrics.overall_accuracy >= 0.8 ? 'green' : metrics.overall_accuracy >= 0.6 ? 'yellow' : 'red';
+        // Abstention is hidden - it tests LLM reasoning ability rather than memory system performance
+
+        // Overall summary at the bottom
+        console.log();
+        const accuracyColor =
+          metrics.overall_accuracy >= 0.8 ? 'green' : metrics.overall_accuracy >= 0.6 ? 'yellow' : 'red';
+
+        if (hasFixedAccuracy && metrics.fixed_overall_accuracy != null) {
+          const fixedOverallColor =
+            metrics.fixed_overall_accuracy >= 0.8 ? 'green' : metrics.fixed_overall_accuracy >= 0.6 ? 'yellow' : 'red';
+          console.log(
+            chalk.bold('Overall Accuracy:'),
+            chalk[accuracyColor](`${(metrics.overall_accuracy * 100).toFixed(2)}%`),
+            chalk.gray('→'),
+            chalk[fixedOverallColor](`${(metrics.fixed_overall_accuracy * 100).toFixed(2)}%`),
+            chalk.gray('(fixed)'),
+          );
+        } else {
           console.log(
             chalk.bold('Overall Accuracy:'),
             chalk[accuracyColor](`${(metrics.overall_accuracy * 100).toFixed(2)}%`),
             chalk.gray(`(average of ${Object.keys(metrics.accuracy_by_type).length} question types)`),
           );
         }
+
+        // Token usage summary (if available)
+        const formatTokens = (n: number) => n.toLocaleString();
+
+        if (metrics.total_usage) {
+          console.log(
+            chalk.gray('Answering Tokens:'),
+            chalk.cyan(formatTokens(metrics.total_usage.inputTokens)),
+            chalk.gray('input,'),
+            chalk.cyan(formatTokens(metrics.total_usage.outputTokens)),
+            chalk.gray('output'),
+            chalk.gray(
+              `(avg ${formatTokens(Math.round(metrics.total_usage.inputTokens / metrics.total_questions))}/q)`,
+            ),
+          );
+        }
+
+        // Load preparation token usage from om-debug.jsonl files
+        const prepUsage = await loadPreparationTokenUsage('./prepared-data', run.config.memoryConfig);
+        if (prepUsage && prepUsage.totalInputTokens > 0) {
+          console.log(
+            chalk.gray('Preparation Tokens:'),
+            chalk.cyan(formatTokens(prepUsage.totalInputTokens)),
+            chalk.gray('input,'),
+            chalk.cyan(formatTokens(prepUsage.totalOutputTokens)),
+            chalk.gray('output'),
+            chalk.gray(
+              `(Observer: ${formatTokens(prepUsage.observerInputTokens)}, Reflector: ${formatTokens(prepUsage.reflectorInputTokens)})`,
+            ),
+          );
+        }
       }
 
       // Get terminal width for final separator
       const terminalWidth = process.stdout.columns || 80;
       const lineWidth = Math.min(terminalWidth - 1, 80);
 
+      // Count unique configurations
+      const uniqueConfigs = new Set(allRuns.map(r => `${r.config.dataset}_${r.config.memoryConfig}`));
+
       console.log(chalk.bold('\n' + '═'.repeat(lineWidth)));
-      console.log(chalk.gray(`\nFound ${allRuns.length} total runs across ${byMemoryConfig.size} configurations`));
-      if (!options.all && byMemoryConfig.size > 0) {
-        console.log(chalk.gray('Use --all to see all runs, not just the latest'));
+      console.log(chalk.gray(`\nFound ${allRuns.length} total runs across ${uniqueConfigs.size} configurations`));
+      if (!options.latest && uniqueConfigs.size > 0 && allRuns.length > uniqueConfigs.size) {
+        console.log(chalk.gray('Use --latest to see only the latest run per config'));
       }
     } catch (error) {
       console.error(chalk.red('\nError:'), error);
@@ -670,8 +1237,8 @@ async function ensureDatasetExists(dataset: string) {
   console.log(chalk.blue('Downloading dataset...\n'));
 
   try {
-    // Run the download script
-    execSync('pnpm download', { stdio: 'inherit' });
+    // Run the download script with specific dataset
+    execSync(`pnpm download -- --dataset ${dataset}`, { stdio: 'inherit' });
 
     // Verify download succeeded
     if (!existsSync(datasetPath) || statSync(datasetPath).size < 1000000) {
@@ -687,4 +1254,249 @@ async function ensureDatasetExists(dataset: string) {
   }
 }
 
+// Sessions command - browse answer sessions for a question
+program
+  .command('sessions')
+  .description('Browse answer sessions for a specific question ID')
+  .requiredOption('-d, --dataset <dataset>', 'Dataset to use (longmemeval_s, longmemeval_m, longmemeval_oracle)')
+  .requiredOption('-q, --question-id <id>', 'Question ID to browse')
+  .option('-a, --all', 'Show all haystack sessions, not just answer sessions')
+  .action(async options => {
+    try {
+      const sessionsCommand = new SessionsCommand();
+      await sessionsCommand.run({
+        dataset: options.dataset,
+        questionId: options.questionId,
+        showAll: options.all,
+      });
+    } catch (error) {
+      console.error(chalk.red('Error:'), error instanceof Error ? error.message : error);
+      process.exit(1);
+    }
+  });
+
+// Deterministic IDs command - update message IDs to be deterministic
+program
+  .command('deterministic-ids')
+  .description('Update message IDs in prepared data to be deterministic (thread_id_msg_index)')
+  .option('--prepared-data-dir <dir>', 'Directory containing prepared data')
+  .option('-q, --question-id <id>', 'Only update a specific question')
+  .action(async options => {
+    try {
+      const command = new DeterministicIdsCommand({
+        preparedDataDir: options.preparedDataDir,
+        questionId: options.questionId,
+      });
+      await command.run();
+    } catch (error) {
+      console.error(chalk.red('Error:'), error instanceof Error ? error.message : error);
+      process.exit(1);
+    }
+  });
+
+// List partial command - find partially prepared questions
+program
+  .command('list-partial')
+  .description('List partially prepared or failed questions')
+  .option('--prepared-data-dir <dir>', 'Directory containing prepared data')
+  .action(async options => {
+    try {
+      const command = new ListPartialCommand({
+        preparedDataDir: options.preparedDataDir,
+      });
+      await command.run();
+    } catch (error) {
+      console.error(chalk.red('Error:'), error instanceof Error ? error.message : error);
+      process.exit(1);
+    }
+  });
+
+// Tokens command - estimate token counts for questions
+program
+  .command('tokens')
+  .description('Estimate token counts for LongMemEval questions')
+  .requiredOption('-d, --dataset <dataset>', 'Dataset to use (longmemeval_s, longmemeval_m, longmemeval_oracle)')
+  .option('-q, --question-id <id>', 'Analyze a specific question by ID')
+  .option('-o, --offset <n>', 'Skip first n questions', parseInt)
+  .option('-s, --subset <n>', 'Analyze only n questions (after offset)', parseInt)
+  .option('-p, --prepared-data <dir>', 'Directory containing prepared data', './prepared-data')
+  .option('--sessions', 'Show per-session breakdown')
+  .option('--top <n>', 'Show top N largest questions in aggregate view', parseInt)
+  .option('--observations-only', 'Only count observation tokens from prepared data (fast)')
+  .option('-c, --config <config>', 'Config to use for observations-only mode', 'observational-memory')
+  .action(async options => {
+    try {
+      const tokensCommand = new TokensCommand();
+      await tokensCommand.run({
+        dataset: options.dataset,
+        questionId: options.questionId,
+        offset: options.offset,
+        subset: options.subset,
+        preparedDataDir: options.preparedData,
+        showSessions: options.sessions,
+        topN: options.top,
+        observationsOnly: options.observationsOnly,
+        config: options.config,
+      });
+    } catch (error) {
+      console.error(chalk.red('Error:'), error instanceof Error ? error.message : error);
+      process.exit(1);
+    }
+  });
+
+program
+  .command('precompute-embeddings')
+  .description('Precompute embeddings for RAG-based observation filtering')
+  .requiredOption('-d, --dataset <dataset>', 'Dataset to use (longmemeval_s, longmemeval_m, longmemeval_oracle)')
+  .requiredOption('-c, --memory-config <config>', 'Memory configuration to use')
+  .option('-o, --offset <n>', 'Skip first n questions', parseInt)
+  .option('-s, --subset <n>', 'Process only n questions (after offset)', parseInt)
+  .option('-p, --prepared-data <dir>', 'Directory containing prepared data', './prepared-data')
+  .option('-b, --batch-size <n>', 'Batch size for embedding (default: 100)', parseInt)
+  .option('--cooldown <ms>', 'Cooldown in ms between questions (default: 1000)', parseInt)
+  .action(async options => {
+    try {
+      const command = new PrecomputeEmbeddingsCommand();
+      await command.run({
+        dataset: options.dataset,
+        memoryConfig: options.memoryConfig,
+        preparedDataDir: options.preparedData,
+        offset: options.offset,
+        subset: options.subset,
+        batchSize: options.batchSize,
+        cooldown: options.cooldown,
+      });
+    } catch (error) {
+      console.error(chalk.red('Error:'), error instanceof Error ? error.message : error);
+      process.exit(1);
+    }
+  });
+
+program
+  .command('find-prohibited')
+  .description('Find messages that trigger Gemini PROHIBITED_CONTENT filter using binary search')
+  .requiredOption('-d, --dataset <dataset>', 'Dataset to use (longmemeval_s, longmemeval_m, longmemeval_oracle)')
+  .requiredOption('-q, --question-id <id>', 'Question ID to analyze')
+  .action(async options => {
+    try {
+      const command = new FindProhibitedCommand();
+      await command.run({
+        dataset: options.dataset,
+        questionId: options.questionId,
+      });
+    } catch (error) {
+      console.error(chalk.red('Error:'), error instanceof Error ? error.message : error);
+      process.exit(1);
+    }
+  });
+
+program
+  .command('partial [config]')
+  .description('Check partial or complete benchmark results')
+  .option('-o, --output <dir>', 'Results directory', './results')
+  .option('--run-id <id>', 'Specific run ID to check')
+  .action(async (config, options) => {
+    try {
+      const command = new PartialResultsCommand();
+      await command.run({
+        config,
+        runId: options.runId,
+        outputDir: options.output,
+      });
+    } catch (error) {
+      console.error(chalk.red('Error:'), error instanceof Error ? error.message : error);
+      process.exit(1);
+    }
+  });
+
+// Investigate command
+import { InvestigateCommand } from './commands/investigate';
+
+program
+  .command('investigate [run-id]')
+  .description('Setup and manage investigation of failed questions')
+  .option('-l, --list', 'List all runs with failures')
+  .option('-c, --config <config>', 'Filter by config name (used with --list)')
+  .option('--status', 'Show investigation status')
+  .option('--next', 'Open next uninvestigated question')
+  .option('--done <question-id>', 'Mark a question as investigated')
+  .option('--fixed <question-id>', 'Mark a question as fix-implemented')
+  .option('--sync', 'Sync fixes to longmemeval_s.json')
+  .option('-o, --output <dir>', 'Investigation output directory', './investigations')
+  .option('-r, --results <dir>', 'Results directory', './results')
+  .option('-p, --prepared-data <dir>', 'Prepared data directory', './prepared-data')
+  .option('--data <dir>', 'Dataset directory', './data')
+  .option('--editor <cmd>', 'Editor command to open files', process.env.EDITOR || 'code')
+  .option('--inspect <question-id>', "Inspect a question's data")
+  .option('--search <keyword>', 'Search observations for a keyword')
+  .option('--trace <keyword>', 'Trace information flow for a keyword')
+  .option('-q, --question-id <id>', 'Question ID for search/trace/date')
+  .option('--date <date>', 'View observations around a date (e.g., "2023/05/29", "May 29")')
+  .option('--context <days>', 'Days of context around --date (default: 1)', '1')
+  .option('--session <idx>', 'View a specific session from original dataset')
+  .option('--list-sessions', 'List all sessions with dates for a question')
+  .option('--check-stale', 'Check if prepared data is stale (pre-cursor-fix)')
+  .option('--stale-only', 'Only list stale questions (use with --check-stale)')
+  .option('--check-duplicates', 'Check for duplicate thread blocks in observations')
+  .option('--baseline', 'Run comprehensive data quality check for a question')
+  .option('--search-original <keyword>', 'Search original dataset for a keyword (shows full context)')
+  .option('--improve <question-id>', 'Add improvements to a question in the dataset')
+  .option('--improve-question <text>', 'Improved question text (use with --improve)')
+  .option('--improve-answer <text>', 'Improved answer text (use with --improve)')
+  .option('--improve-note <text>', 'Improvement note (use with --improve)')
+  .option(
+    '--category <category>',
+    'Failure category: observer-miss, reflector-loss, agent-reasoning, dataset-error, data-freshness, knowledge-update, rag-miss, other',
+  )
+  .option('--clear-improved [fields]', 'Clear improved fields: all, question, answer, note, category (comma-separated)')
+  .option('--prepare-stale', '(deprecated) Use --print-prepare-command instead')
+  .option('--print-prepare-command', 'Print a prepare command for stale/partial questions')
+  .option('--dry-run', 'Show what would be prepared without actually doing it')
+  .action(async (runId, options) => {
+    try {
+      const command = new InvestigateCommand({
+        outputDir: options.output,
+        resultsDir: options.results,
+        preparedDataDir: options.preparedData,
+        datasetDir: options.data,
+        editor: options.editor,
+      });
+      await command.run({
+        runId,
+        list: options.list,
+        config: options.config,
+        status: options.status,
+        next: options.next,
+        done: options.done,
+        fixed: options.fixed,
+        sync: options.sync,
+        inspect: options.inspect,
+        search: options.search,
+        trace: options.trace,
+        questionId: options.questionId,
+        date: options.date,
+        context: options.context ? parseInt(options.context, 10) : 1,
+        session: options.session !== undefined ? parseInt(options.session, 10) : undefined,
+        listSessions: options.listSessions,
+        checkStale: options.checkStale,
+        staleOnly: options.staleOnly,
+        checkDuplicates: options.checkDuplicates,
+        baseline: options.baseline,
+        searchOriginal: options.searchOriginal,
+        improve: options.improve,
+        improveQuestion: options.improveQuestion,
+        improveAnswer: options.improveAnswer,
+        improveNote: options.improveNote,
+        category: options.category,
+        clearImproved: options.clearImproved,
+        prepareStale: options.prepareStale,
+        printPrepareCommand: options.printPrepareCommand,
+        dryRun: options.dryRun,
+      });
+    } catch (error) {
+      console.error(chalk.red('Error:'), error instanceof Error ? error.message : error);
+      process.exit(1);
+    }
+  });
+
 program.parse();