1100 lines · 33.7 KB · TypeScript
RawDownload
/** * Unified Context Tracker * * Provides comprehensive context tracking for all sources of context injection: * - CLAUDE.md files (enterprise, user, project, directory) * - Mentioned files (@mentions) * - Tool outputs * * This builds on claudeMdTracker.ts and extends it to track all context sources. */ import { estimateTokens } from '@shared/utils/tokenFormatting'; import { MAX_MENTIONED_FILE_TOKENS } from '../types/contextInjection'; import { buildDisplayItems, findLastOutput, linkToolCallsToResults } from './aiGroupEnhancer';import {  createGlobalInjections,  detectClaudeMdFromFilePath,  extractFileRefsFromResponses,  extractReadToolPaths,  extractUserMentionPaths,  generateInjectionId,  getDisplayName,} from './claudeMdTracker'; import type { ClaudeMdInjection, ClaudeMdSource } from '../types/claudeMd';import type {  ClaudeMdContextInjection,  CompactionTokenDelta,  ContextInjection,  ContextPhase,  ContextPhaseInfo,  ContextStats,  MentionedFileInfo,  MentionedFileInjection,  NewCountsByCategory,  TaskCoordinationBreakdown,  TaskCoordinationInjection,  ThinkingTextBreakdown,  ThinkingTextInjection,  TokensByCategory,  ToolOutputInjection,  ToolTokenBreakdown,  UserMessageInjection,} from '../types/contextInjection';import type { ClaudeMdFileInfo } from '../types/data';import type {  AIGroup,  AIGroupDisplayItem,  ChatItem,  LinkedToolItem,  UserGroup,} from '../types/groups'; // =============================================================================// Constants// ============================================================================= /** Category identifier for mentioned file injections */const CATEGORY_MENTIONED_FILE = 'mentioned-file' as const; /** Tool names that constitute task coordination overhead */const TASK_COORDINATION_TOOL_NAMES = new Set([  'SendMessage',  'TeamCreate',  'TeamDelete',  'TaskCreate',  'TaskUpdate',  'TaskList',  'TaskGet',]); // =============================================================================// ID Generation Functions// ============================================================================= /** * Generate a unique ID for a mentioned file injection. * Uses a similar approach to generateInjectionId but with 'mf-' prefix. */function generateMentionedFileId(path: string): string {  let hash = 0;  for (let i = 0; i < path.length; i++) {    const char = path.charCodeAt(i);    hash = (hash << 5) - hash + char;    hash = hash & hash; // Convert to 32bit integer  }  const positiveHash = Math.abs(hash).toString(16);  return `mf-${positiveHash}`;} /** * Generate a unique ID for a tool output injection. */function generateToolOutputId(turnIndex: number): string {  return `tool-output-ai-${turnIndex}`;} /** * Generate unique ID for thinking-text injection. */function generateThinkingTextId(turnIndex: number): string {  return `thinking-text-ai-${turnIndex}`;} /** * Generate unique ID for task coordination injection. */function generateTaskCoordinationId(turnIndex: number): string {  return `task-coord-ai-${turnIndex}`;} /** * Generate unique ID for user message injection. */function generateUserMessageId(turnIndex: number): string {  return `user-msg-ai-${turnIndex}`;} // =============================================================================// Injection Wrapping Functions// ============================================================================= /** * Wrap a ClaudeMdInjection with the 'claude-md' category for union compatibility. */function wrapClaudeMdInjection(injection: ClaudeMdInjection): ClaudeMdContextInjection {  return {    ...injection,    category: 'claude-md' as const,  };} // =============================================================================// Mentioned File Injection Creation// ============================================================================= /** * Parameters for creating a mentioned file injection. */interface CreateMentionedFileInjectionParams {  /** Absolute file path */  path: string;  /** Display name (relative path or filename) */  displayName: string;  /** Estimated token count for this file */  estimatedTokens: number;  /** Turn index where this file was first mentioned */  turnIndex: number;  /** AI group ID for navigation */  aiGroupId: string;  /** Whether the file exists on disk */  exists?: boolean;} /** * Create a MentionedFileInjection object. */function createMentionedFileInjection(  params: CreateMentionedFileInjectionParams): MentionedFileInjection {  return {    id: generateMentionedFileId(params.path),    category: CATEGORY_MENTIONED_FILE,    path: params.path,    displayName: params.displayName,    estimatedTokens: params.estimatedTokens,    firstSeenTurnIndex: params.turnIndex,    firstSeenInGroup: params.aiGroupId,    exists: params.exists ?? true,  };} // =============================================================================// Tool Output Aggregation// ============================================================================= /** * Aggregate tool outputs from all linked tools in a turn. * Also includes tokens from user-invoked skills (via /skill-name commands). * Returns a ToolOutputInjection if there are any tool outputs with tokens. */function aggregateToolOutputs(  linkedTools: Map<string, LinkedToolItem>,  turnIndex: number,  aiGroupId: string,  displayItems?: AIGroupDisplayItem[]): ToolOutputInjection | null {  const toolBreakdown: ToolTokenBreakdown[] = [];  let totalTokens = 0;   for (const linkedTool of linkedTools.values()) {    // Skip task coordination tools - they are tracked separately    if (TASK_COORDINATION_TOOL_NAMES.has(linkedTool.name)) {      continue;    }     // Calculate total context tokens for the tool operation    // This matches getToolContextTokens in LinkedToolItem.tsx and ErrorDetector    //    // callTokens: What Claude generated (Write file content, Edit old/new strings)    // resultTokens: What Claude reads back (success message, Read file content)    // skillTokens: Additional context for Skill tools    const callTokens = linkedTool.callTokens ?? 0;    const resultTokens = linkedTool.result?.tokenCount ?? 0;    const skillTokens = linkedTool.skillInstructionsTokenCount ?? 0;    const toolTokenCount = callTokens + resultTokens + skillTokens;     if (toolTokenCount > 0) {      // Rename "Task" to "Task (Subagent)" for clarity in the UI      const displayName = linkedTool.name === 'Task' ? 'Task (Subagent)' : linkedTool.name;      toolBreakdown.push({        toolName: displayName,        tokenCount: toolTokenCount,        isError: linkedTool.result?.isError ?? false,        toolUseId: linkedTool.id,      });      totalTokens += toolTokenCount;    }  }   // Include user-invoked slash tokens from display items  // These are slashes invoked via /xxx commands  if (displayItems) {    for (const item of displayItems) {      if (item.type === 'slash' && item.slash.instructionsTokenCount) {        toolBreakdown.push({          toolName: `/${item.slash.name}`,          tokenCount: item.slash.instructionsTokenCount,          isError: false,        });        totalTokens += item.slash.instructionsTokenCount;      }    }  }   // Return null if no tokens from tools  if (totalTokens === 0) {    return null;  }   return {    id: generateToolOutputId(turnIndex),    category: 'tool-output',    turnIndex,    aiGroupId,    estimatedTokens: totalTokens,    toolCount: toolBreakdown.length,    toolBreakdown,  };} // =============================================================================// Task Coordination Aggregation// ============================================================================= /** * Aggregate task coordination tokens from linked tools and display items. * Tracks SendMessage, TeamCreate, TaskCreate, and other task tools, * plus teammate_message items injected into the session. */function aggregateTaskCoordination(  linkedTools: Map<string, LinkedToolItem>,  turnIndex: number,  aiGroupId: string,  displayItems?: AIGroupDisplayItem[]): TaskCoordinationInjection | null {  const breakdown: TaskCoordinationBreakdown[] = [];  let totalTokens = 0;   // Scan linked tools for task coordination tools  for (const linkedTool of linkedTools.values()) {    if (!TASK_COORDINATION_TOOL_NAMES.has(linkedTool.name)) {      continue;    }     const callTokens = linkedTool.callTokens ?? 0;    const resultTokens = linkedTool.result?.tokenCount ?? 0;    const skillTokens = linkedTool.skillInstructionsTokenCount ?? 0;    const toolTokenCount = callTokens + resultTokens + skillTokens;     if (toolTokenCount > 0) {      // Extract a label from tool input for SendMessage (recipient name)      let label = linkedTool.name;      if (linkedTool.name === 'SendMessage' && linkedTool.input) {        const recipient = linkedTool.input.recipient as string | undefined;        if (recipient) {          label = `SendMessage → ${recipient}`;        }      }       breakdown.push({        type: linkedTool.name === 'SendMessage' ? 'send-message' : 'task-tool',        toolName: linkedTool.name,        tokenCount: toolTokenCount,        label,      });      totalTokens += toolTokenCount;    }  }   // Scan display items for teammate messages  if (displayItems) {    for (const item of displayItems) {      if (item.type === 'teammate_message' && item.teammateMessage.tokenCount) {        breakdown.push({          type: 'teammate-message',          tokenCount: item.teammateMessage.tokenCount,          label: item.teammateMessage.teammateId,        });        totalTokens += item.teammateMessage.tokenCount;      }    }  }   if (totalTokens === 0) {    return null;  }   return {    id: generateTaskCoordinationId(turnIndex),    category: 'task-coordination',    turnIndex,    aiGroupId,    estimatedTokens: totalTokens,    breakdown,  };} // =============================================================================// User Message Injection Creation// ============================================================================= /** * Create a UserMessageInjection from a user group. * Uses rawText (includes commands and @mentions) for token estimation * since that's what's actually sent to the API. * * @returns UserMessageInjection or null if empty text or 0 tokens */function createUserMessageInjection(  userGroup: UserGroup,  turnIndex: number,  aiGroupId: string): UserMessageInjection | null {  const text = userGroup.content.rawText ?? userGroup.content.text ?? '';  if (!text) return null;   const tokens = estimateTokens(text);  if (tokens === 0) return null;   const textPreview = text.length > 80 ? text.slice(0, 80) + '…' : text;   return {    id: generateUserMessageId(turnIndex),    category: 'user-message',    turnIndex,    aiGroupId,    estimatedTokens: tokens,    textPreview,  };} // =============================================================================// Thinking/Text Output Aggregation// ============================================================================= /** * Aggregates thinking and text output tokens for a single turn. * Creates a ThinkingTextInjection that tracks all thinking blocks and text outputs. * * @param displayItems - Display items from the AI group * @param turnIndex - The turn index (0-based) * @param aiGroupId - The AI group ID for navigation * @returns ThinkingTextInjection or null if no tokens */function aggregateThinkingText(  displayItems: AIGroupDisplayItem[],  turnIndex: number,  aiGroupId: string): ThinkingTextInjection | null {  const breakdown: ThinkingTextBreakdown[] = [];  let totalTokens = 0;  let thinkingTokens = 0;  let textTokens = 0;   for (const item of displayItems) {    if (item.type === 'thinking' && item.tokenCount && item.tokenCount > 0) {      thinkingTokens += item.tokenCount;      totalTokens += item.tokenCount;    } else if (item.type === 'output' && item.tokenCount && item.tokenCount > 0) {      textTokens += item.tokenCount;      totalTokens += item.tokenCount;    }  }   if (thinkingTokens > 0) {    breakdown.push({ type: 'thinking', tokenCount: thinkingTokens });  }  if (textTokens > 0) {    breakdown.push({ type: 'text', tokenCount: textTokens });  }   if (totalTokens === 0) {    return null;  }   return {    id: generateThinkingTextId(turnIndex),    category: 'thinking-text',    turnIndex,    aiGroupId,    estimatedTokens: totalTokens,    breakdown,  };} // =============================================================================// Stats Computation// ============================================================================= /** * Parameters for computing context stats for an AI group. */interface ComputeContextStatsParams {  /** The AI group being processed */  aiGroup: AIGroup;  /** The preceding user group (if any) */  userGroup: UserGroup | null;  /** Linked tools map from the enhanced AI group */  linkedTools: Map<string, LinkedToolItem>;  /** Display items from enhanced AI group (includes user skills) */  displayItems?: AIGroupDisplayItem[];  /** Whether this is the first AI group in the session */  isFirstGroup: boolean;  /** Accumulated injections from previous groups */  previousInjections: ContextInjection[];  /** Project root path for resolving relative paths */  projectRoot: string;  /** Token data for CLAUDE.md files (global sources) */  claudeMdTokenData?: Record<string, ClaudeMdFileInfo>;  /** Token data for mentioned files */  mentionedFileTokenData?: Map<string, MentionedFileInfo>;  /** Token data for validated directory CLAUDE.md files (keyed by full path) */  directoryTokenData?: Record<string, ClaudeMdFileInfo>;} /** * Helper to check if a path is absolute. */function isAbsolutePath(path: string): boolean {  return (    path.startsWith('/') ||    path.startsWith('~/') ||    path.startsWith('~\\') ||    path === '~' ||    path.startsWith('\\\\') ||    /^[a-zA-Z]:[\\/]/.test(path)  );} /** * Helper to join paths, handling various path formats properly. * Handles: * - Absolute paths: /full/path/file.tsx (returned as-is) * - Relative paths with ./: ./apps/foo/bar.tsx (strips ./) * - Parent paths with ../: ../other/file.tsx (walks up directories) * - Plain paths: apps/foo/bar.tsx (joins with base) * - Paths with @ prefix: @apps/foo/bar.tsx (strips @ then joins) */function joinPaths(base: string, relative: string): string {  if (isAbsolutePath(relative)) {    return relative;  }   const cleanBase = trimTrailingSeparator(base);   // Handle @ prefix (file mention marker) - strip it if present  let cleanRelative = relative;  if (cleanRelative.startsWith('@')) {    cleanRelative = cleanRelative.slice(1);  }   // Handle ./ prefix (current directory)  if (cleanRelative.startsWith('./')) {    cleanRelative = cleanRelative.slice(2);  }   // Handle ../ prefixes (parent directory)  const separator = cleanBase.includes('\\') ? '\\' : '/';  const hasUnixRoot = cleanBase.startsWith('/');  const hasUncRoot = cleanBase.startsWith('\\\\');  const normalizedRelative = normalizeSeparators(cleanRelative, separator);  const baseParts = splitPath(cleanBase);  let remainingRelative = normalizedRelative;  while (remainingRelative.startsWith(`..${separator}`)) {    remainingRelative = remainingRelative.slice(3);    if (baseParts.length > 1) {      baseParts.pop();

Lines 1–500 of 1100.