CoolFace
Datasetpublic

basant307/AI_Governance_Project

sourceHugging Faceapache-2.0updated 2mo agoView on Hugging Face
0likes45downloads
client.ts3039 linesDownload Raw Back to core
1/**2 * @license3 * Copyright 2026 Qwen Team4 * SPDX-License-Identifier: Apache-2.05 */6 7// External dependencies8import { createUserContent } from '@google/genai';9import type {10  Content,11  GenerateContentConfig,12  GenerateContentResponse,13  PartListUnion,14  Tool,15} from '@google/genai';16import process from 'node:process';17 18// Config19import { ApprovalMode, type Config } from '../config/config.js';20import { createDebugLogger } from '../utils/debugLogger.js';21import { cleanupOldToolResults } from '../utils/toolResultCleanup.js';22import { Storage } from '../config/storage.js';23import { recordStartupEvent } from '../utils/startupEventSink.js';24import {25  microcompactHistory,26  type MicrocompactMeta,27  type MicrocompactOptions,28} from '../services/microcompaction/microcompact.js';29import {30  activeGoalEquals,31  getActiveGoal,32  type ActiveGoal,33} from '../goals/activeGoalStore.js';34import { abortGoalForStopHookCap } from '../goals/goalHook.js';35import { formatStopHookBlockingCapWarning } from '../hooks/stopHookCap.js';36import { buildContextUsage } from '../hooks/context-usage.js';37import {38  DEFAULT_TOKEN_LIMIT,39  escalatedOutputTokenLimit,40  parsePositiveIntegerEnvValue,41} from './tokenLimits.js';42import { createSessionStartProfiler } from './session-start-profiler.js';43 44const debugLogger = createDebugLogger('CLIENT');45 46// Core modules47import { GeminiChat } from './geminiChat.js';48import { getRecentGitStatus } from '../utils/gitUtils.js';49import {50  getArenaSystemReminder,51  getCoreSystemPrompt,52  getCustomSystemPrompt,53  getPlanModeSystemReminder,54} from './prompts.js';55import {56  CompressionStatus,57  GeminiEventType,58  Turn,59  type ChatCompressionInfo,60  type ServerGeminiStreamEvent,61} from './turn.js';62 63// Services64import { LoopDetectionService } from '../services/loopDetectionService.js';65import { CommitAttributionService } from '../services/commitAttribution.js';66 67// Tools68import type { RelevantAutoMemoryPromptResult } from '../memory/manager.js';69import { AUTO_SKILL_THRESHOLD } from '../memory/manager.js';70import {71  DEFAULT_AUTO_SKILL_MAX_TURNS,72  DEFAULT_AUTO_SKILL_TIMEOUT_MS,73} from '../memory/skillReviewAgentPlanner.js';74import { isProjectSkillPath } from '../skills/skill-paths.js';75import { ToolNames } from '../tools/tool-names.js';76 77// Telemetry78import {79  NextSpeakerCheckEvent,80  logNextSpeakerCheck,81  startInteractionSpan,82  endInteractionSpan,83  getActiveInteractionSpan,84  addUserPromptAttributes,85} from '../telemetry/index.js';86import { uiTelemetryService } from '../telemetry/uiTelemetry.js';87 88// Forked agent cache89import {90  saveCacheSafeParams,91  clearCacheSafeParams,92} from '../utils/forkedAgent.js';93 94// Utilities95import {96  formatDateForContext,97  buildChangedAgentsReminder,98  buildChangedMcpToolsReminder,99  buildChangedSkillsReminder,100  getDirectoryContextString,101  getInitialChatHistory,102  getStartupContextLength,103  type AgentAvailabilityEntry,104} from '../utils/environmentContext.js';105import {106  collectAvailableSkillEntries,107  type AvailableSkillEntry,108} from '../tools/skill-utils.js';109import type { DeferredToolSummary } from '../tools/tool-registry.js';110import {111  buildApiHistoryFromConversation,112  replayUiTelemetryFromConversation,113} from '../services/sessionService.js';114import { reportError } from '../utils/errorReporting.js';115import { getErrorMessage } from '../utils/errors.js';116import { checkNextSpeaker } from '../utils/nextSpeakerChecker.js';117import {118  flatMapTextParts,119  prependToFirstTextPart,120} from '../utils/partUtils.js';121import { promptIdContext } from '../utils/promptIdContext.js';122import { retryWithBackoff, isUnattendedMode } from '../utils/retry.js';123import { subagentNameContext } from '../utils/subagentNameContext.js';124import { escapeSystemReminderTags } from '../utils/xml.js';125import { ApiRetryEvent } from '../telemetry/types.js';126import { logApiRetry } from '../telemetry/loggers.js';127import { shouldUsePlanOnlyReminderInSubagentContext } from '../agents/runtime/subagent-plan-tool-policy.js';128 129// Hook types and utilities130import {131  MessageBusType,132  type HookExecutionRequest,133  type HookExecutionResponse,134} from '../confirmation-bus/types.js';135import { partToString } from '../utils/partUtils.js';136import { createHookOutput, SessionStartSource } from '../hooks/types.js';137import fsPromises from 'node:fs/promises';138 139// IDE integration140import { ideContextStore } from '../ide/ideContext.js';141import { type File, type IdeContext } from '../ide/types.js';142import { PermissionMode, type StopHookOutput } from '../hooks/types.js';143 144const MAX_TURNS = 100;145const MAX_RECENT_TOOL_NAMES_FOR_MEMORY = 20;146 147export enum SendMessageType {148  UserQuery = 'userQuery',149  ToolResult = 'toolResult',150  Retry = 'retry',151  Hook = 'hook',152  /** Cron-fired prompt. Behaves like UserQuery but skips UserPromptSubmit hook. */153  Cron = 'cron',154  /** Background agent notification. Display item is added by the drain loop. */155  Notification = 'notification',156  /**157   * A message delivered to the leader from a teammate. Behaves like a158   * fresh top-level interaction (loop-detector reset + interaction span)159   * but is not a user prompt — it does not bump commit attribution or get160   * recorded as a user message.161   */162  Teammate = 'teammate',163}164 165export interface SendMessageOptions {166  type: SendMessageType;167  /** Track stop hook iterations to prevent infinite loops and display loop info */168  stopHookState?: {169    iterationCount: number;170    reasons: string[];171  };172  /** Display text for notification messages (persisted for session resume). */173  notificationDisplayText?: string;174  /** Model override from skill execution. When present, overrides the session model for this turn. */175  modelOverride?: string;176}177 178const EMPTY_RELEVANT_AUTO_MEMORY_RESULT: RelevantAutoMemoryPromptResult = {179  prompt: '',180  selectedDocs: [],181  strategy: 'none',182};183 184function wrapIdeContext(contextText: string): string {185  const safeContextText = escapeSystemReminderTags(contextText);186  return `<system-reminder>\n${safeContextText}\n</system-reminder>`;187}188 189/**190 * Handle for a non-blocking auto-memory recall prefetch.191 *192 * Lifecycle:193 *  1. Created on UserQuery/Cron — the recall promise fires immediately,194 *     `pendingMemoryPrefetch` is set to this handle.195 *  2. Consumed at either of two opportunistic points: a zero-wait196 *     `settledAt !== null` poll just before the UserQuery main request,197 *     or — if recall hadn't settled yet — on the first ToolResult turn.198 *  3. Aborted-and-discarded by every cleanup path (resetChat,199 *     MaxSessionTurns, etc.) or replaced when a new UserQuery arrives.200 */201type MemoryPrefetchHandle = {202  promise: Promise<RelevantAutoMemoryPromptResult>;203  /** Set by promise.finally(). null until the promise settles. */204  settledAt: number | null;205  /** True after memory has been injected — prevents double-inject. */206  consumed: boolean;207  controller: AbortController;208};209 210/** Tools that can write to the skills directory, used to detect skillsModifiedInSession. */211const SKILL_WRITE_TOOL_NAMES: ReadonlySet<string> = new Set([212  ToolNames.WRITE_FILE,213  ToolNames.EDIT,214]);215 216export class GeminiClient {217  private chat?: GeminiChat;218  private initializedSessionId: string | undefined;219  private sessionTurnCount = 0;220  private toolCallCount = 0;221  private skillsModifiedInSession = false;222  private cachedGitStatus: string | null | undefined;223  private readonly surfacedRelevantAutoMemoryPaths = new Set<string>();224  private shutdownRequested = false;225 226  private readonly loopDetector: LoopDetectionService;227  private lastPromptId: string | undefined = undefined;228  private lastSentIdeContext: IdeContext | undefined;229  private forceFullIdeContext = true;230  private recentCompletedToolNames: string[] = [];231  private pendingMemoryPrefetch: MemoryPrefetchHandle | undefined;232  private lastSessionStartContext: string | undefined;233  private lastSessionStartSource: SessionStartSource | undefined;234  private announcedDeferredToolNames = new Set<string>();235  // MCP-only subset the model has actually seen via startup or delta reminders.236  // `announcedDeferredToolNames` is broader and exists for deferred tool-search237  // dedup; MCP add/remove deltas need this narrower model-visible set.238  private announcedMcpToolNames = new Set<string>();239  private pendingAddedMcpTools = new Map<string, DeferredToolSummary>();240  private pendingRemovedMcpToolNames = new Set<string>();241  // Dedup state for the per-turn skill/command "now available" delta reminders242  // (drainSkillAndCommandReminders). Keys are "skill:<name>" / "cmd:<name>". The243  // set is seeded on the first drain from the current skills (the startup244  // snapshot already listed them) and reset whenever the startup prelude is245  // rebuilt (startChat), so a resumed/compacted session re-seeds from its fresh246  // snapshot instead of re-announcing — mirrors Claude Code's247  // suppressNextSkillListing / "don't re-inject on compact".248  private announcedSkillReminderKeys = new Set<string>();249  private skillRemindersInitialized = false;250  private announcedAgentReminderNames = new Set<string>();251  private agentRemindersInitialized = false;252 253  private static skillEntryKey(e: AvailableSkillEntry): string {254    return e.level !== undefined ? `skill:${e.name}` : `cmd:${e.name}`;255  }256 257  /**258   * Seeds skill-reminder dedup from the entries actually rendered into the259   * startup snapshot. Mirrors `rememberAnnouncedDeferredTools`: the dedup is260   * seeded from what the model actually SAW, not from whatever happens to be261   * current at the first drain (which may include late-registered MCP262   * prompts/commands the snapshot never listed).263   */264  private seedSkillReminderDedupFromSnapshot(265    snapshotEntries: AvailableSkillEntry[],266  ): void {267    this.announcedSkillReminderKeys = new Set(268      snapshotEntries.map(GeminiClient.skillEntryKey),269    );270    this.skillRemindersInitialized = true;271  }272 273  private async seedAgentReminderDedupFromCurrent(): Promise<void> {274    try {275      const agents = await this.config.getSubagentManager().listSubagents();276      this.announcedAgentReminderNames = new Set(277        agents.map((agent) => agent.name),278      );279      this.agentRemindersInitialized = true;280    } catch (error) {281      debugLogger.warn('seedAgentReminderDedupFromCurrent failed', error);282      this.announcedAgentReminderNames.clear();283      this.agentRemindersInitialized = false;284    }285  }286 287  /**288   * Tracks the most recently injected date string to prevent injecting289   * duplicate or conflicting dates when a session spans midnight.290   * Only UserQuery turns inject dates; Cron/ToolResult turns reuse the291   * startup-context date which is still current within the same session.292   */293  private lastInjectedDate: string | undefined;294 295  /**296   * Promises for pending background memory tasks (dream / extract).297   * Each promise resolves with a count of memory files touched (0 = nothing written).298   * Consumed by the CLI via `consumePendingMemoryTaskPromises()`.299   */300  private pendingMemoryTaskPromises: Array<Promise<number>> = [];301 302  /**303   * Timestamp (epoch ms) of the last completed API call.304   * Used to detect idle periods for thinking block cleanup.305   * Starts as null — on the first query there is no prior thinking to clean,306   * so the idle check is skipped until the first API call completes.307   */308  private lastApiCompletionTimestamp: number | null = null;309  /** Cleanup checkpoint for long-running Hook continuations such as /goal. */310  private lastHookMicrocompactionTimestamp: number | null = null;311 312  constructor(private readonly config: Config) {313    this.loopDetector = new LoopDetectionService(config);314  }315 316  async initialize(sessionStartSource?: SessionStartSource) {317    const sessionId = this.config.getSessionId();318    this.lastPromptId = sessionId;319 320    if (this.isInitialized() && this.initializedSessionId === sessionId) {321      return;322    }323 324    // Check if we're resuming from a previous session325    const resumedSessionData = this.config.getResumedSessionData();326    if (resumedSessionData) {327      const resumeTokenCounts = replayUiTelemetryFromConversation(328        resumedSessionData.conversation,329        this.config.getSessionId(),330      );331      // Convert resumed session to API history format332      // Each ChatRecord's message field is already a Content object333      const resumedHistory = buildApiHistoryFromConversation(334        resumedSessionData.conversation,335      );336      this.seedRecentCompletedToolNamesFromHistory(resumedHistory);337      await this.startChat(338        resumedHistory,339        sessionStartSource ?? SessionStartSource.Resume,340      );341      const chat = this.getChat();342      if (resumeTokenCounts) {343        chat.seedResumeTokenCounts(344          resumeTokenCounts.promptTokenCount,345          resumeTokenCounts.outputTokenCount,346        );347      } else {348        chat.setLastPromptTokenCount(349          uiTelemetryService.getLastPromptTokenCount(),350        );351      }352 353      // Restore attribution state from the last snapshot in the session354      this.restoreAttributionFromSession(resumedSessionData.conversation);355    } else {356      if (sessionStartSource !== undefined) {357        await this.startChat(undefined, sessionStartSource);358      } else {359        await this.startChat();360      }361    }362 363    this.initializedSessionId = sessionId;364 365    // Clean up stale tool result files from previous sessions (fire-and-forget)366    void cleanupOldToolResults(Storage.getGlobalTempDir(), 24 * 60 * 60 * 1000);367  }368 369  /**370   * Restore attribution state from the last snapshot in a resumed session.371   */372  private restoreAttributionFromSession(conversation: {373    messages: Array<{ subtype?: string; systemPayload?: unknown }>;374  }): void {375    // Find the last attribution snapshot in the session376    let lastSnapshot: unknown = null;377    for (const msg of conversation.messages) {378      if (379        msg.subtype === 'attribution_snapshot' &&380        msg.systemPayload &&381        typeof msg.systemPayload === 'object' &&382        'snapshot' in msg.systemPayload383      ) {384        lastSnapshot = (msg.systemPayload as { snapshot: unknown }).snapshot;385      }386    }387    if (lastSnapshot && typeof lastSnapshot === 'object') {388      try {389        CommitAttributionService.getInstance().restoreFromSnapshot(390          lastSnapshot as import('../services/commitAttribution.js').AttributionSnapshot,391        );392        debugLogger.debug('Restored attribution state from session snapshot');393      } catch {394        debugLogger.warn('Failed to restore attribution snapshot');395      }396    }397  }398 399  async addHistory(content: Content) {400    this.getChat().addHistory(content);401  }402 403  getChat(): GeminiChat {404    if (!this.chat) {405      throw new Error('Chat not initialized');406    }407    return this.chat;408  }409 410  isInitialized(): boolean {411    return this.chat !== undefined;412  }413 414  getHistory(curated: boolean = false): Content[] {415    return this.getChat().getHistory(curated);416  }417 418  getHistoryShallow(curated: boolean = false): Content[] {419    const chat = this.getChat();420    return chat.getHistoryShallow?.(curated) ?? chat.getHistory(curated);421  }422 423  getHistoryTail(count: number, curated: boolean = false): Content[] {424    return this.getChat().getHistoryTail(count, curated);425  }426 427  private getHistoryTailShallow(428    count: number,429    curated: boolean = false,430  ): Content[] {431    const chat = this.getChat();432    return (433      chat.getHistoryTailShallow?.(count, curated) ??434      chat.getHistoryTail?.(count, curated) ??435      chat.getHistory(curated).slice(-count)436    );437  }438 439  private peekLastHistoryEntry(): Content | undefined {440    const chat = this.getChat();441    return chat.peekLastHistoryEntry?.() ?? chat.getHistory().at(-1);442  }443 444  private getHistoryLength(): number {445    const chat = this.getChat();446    return chat.getHistoryLength?.() ?? chat.getHistory().length;447  }448 449  private getLastModelMessageText(): string | undefined {450    const chat = this.getChat();451    if (chat.getLastModelMessageText) {452      return chat.getLastModelMessageText();453    }454    const history = chat.getHistoryShallow?.() ?? chat.getHistory();455    for (let i = history.length - 1; i >= 0; i--) {456      const message = history[i];457      if (message?.role !== 'model') continue;458      const text =459        message.parts460          ?.filter(461            (part): part is { text: string } =>462              typeof part.text === 'string' && !part.thought,463          )464          .map((part) => part.text)465          .join('') ?? '';466      return text || undefined;467    }468    return undefined;469  }470 471  /**472   * Walk-only accessor for the set of `functionResponse.id` strings in473   * raw history. Callers that only need the dedup id set (notably474   * `useGeminiStream.handleCompletedTools`) MUST prefer this over475   * {@link getHistory}, which deep-clones the entire conversation via476   * `structuredClone` on every call. On long sessions with sizable477   * tool outputs the clone is a multi-millisecond hit on the React UI478   * thread; running it on every tool-completion batch caused visible479   * frame drops during streaming. See480   * `GeminiChat.getHistoryFunctionResponseIds` for the implementation.481   */482  getHistoryFunctionResponseIds(): Set<string> {483    return this.getChat().getHistoryFunctionResponseIds();484  }485 486  /**487   * Pop orphaned trailing user entries from the in-memory chat history.488   * Used by:489   *   - The Retry submit path (sendMessageStream below), which drops a490   *     prior failed attempt before re-sending.491   *   - The auto-restore-on-cancel flow in AppContainer, which rewinds492   *     a user prompt out of the UI transcript and the disk-backed493   *     ↑-history; this is the third place the cancelled prompt lives.494   *     Without calling this from auto-restore, the next request's wire495   *     payload would carry two consecutive user turns — the cancelled496   *     one and the new one — and the model would see context the user497   *     thought had been undone.498   */499  stripOrphanedUserEntriesFromHistory(): Content[] {500    const chat = this.getChat();501    const before = chat.getHistoryLength();502    const strippedEntries = chat.stripOrphanedUserEntriesFromHistory();503    const after = chat.getHistoryLength();504    if (after >= before) {505      // Nothing to strip — leave caches and IDE context alone.506      return strippedEntries;507    }508    // Stripped trailing user entries can include read_file509    // functionResponses from a failed-then-retried request. The510    // FileReadCache would still record those reads, so the retry's511    // re-issued Read could hit the file_unchanged placeholder while512    // the model has nothing to fall back on. Clear to be safe.513    debugLogger.debug(514      `[FILE_READ_CACHE] clear after stripOrphanedUserEntriesFromHistory(prev=${before}, new=${after})`,515    );516    this.config.getFileReadCache().clear();517    // The stripped user turn may have carried the IDE context (open files,518    // workspace state) that `lastSentIdeContext` advanced past. Without519    // forcing a resend, the next request would either skip IDE context520    // entirely or send only a diff against a now-removed baseline. Match521    // the invalidation `setHistory()` / `truncateHistory()` already do.522    this.forceFullIdeContext = true;523    return strippedEntries;524  }525 526  /**527   * Synthesize a `functionResponse` for every dangling `model[functionCall]`528   * in chat history whose corresponding tool_result never landed. Inverse of529   * {@link stripOrphanedUserEntriesFromHistory}, which only handles trailing530   * `user` entries.531   *532   * This `GeminiClient` method is the resume-path entry point — called once533   * from {@link startChat} after the transcript loads, covering `--resume`534   * of a session that crashed between a partial-tool_use push and the535   * tool's eventual completion.536   *537   * The other two coverage points (Retry submit path after538   * `stripOrphanedUserEntriesFromHistory`, and the defensive pass at the539   * start of every UserQuery / Cron send) live one layer down inside540   * `GeminiChat.sendMessageStream` and call the standalone541   * `repairOrphanedToolUseTurns(history)` function directly — they don't542   * route through this wrapper. Anyone tracing the repair-pass coupling543   * between the client and chat layers should follow that path544   * separately rather than expect everything to funnel through here.545   *546   * Synthesizes an `error` `functionResponse`. The React tool scheduler547   * (`useGeminiStream.handleCompletedTools`) MUST dedupe by `callId` against548   * the live history before submitting its own `tool_result` — otherwise a549   * late real result lands as a second `user[tool_result]` block (orphan550   * because the synthetic already consumed the matching `tool_use`).551   */552  repairOrphanedToolUseTurnsInHistory(reason?: string): {553    injected: Array<{ callId: string; name: string }>;554    droppedDuplicates: Array<{ callId: string; name: string }>;555  } {556    const result = this.getChat().repairOrphanedToolUseTurns(reason);557    if (result.injected.length > 0) {558      debugLogger.warn(559        `[REPAIR] Synthesized ${result.injected.length} functionResponse(s) ` +560          `for dangling tool_use(s): ${result.injected561            .map((e) => `${e.name}(${e.callId})`)562            .join(', ')}`,563      );564    }565    if (result.droppedDuplicates.length > 0) {566      // Surface the duplicate-cleanup pass so investigators tracing567      // a dedup-drop log have a breadcrumb pointing back to the568      // repair function. Without this a duplicate-only repair (no569      // synthesis, no hoist) leaves zero diagnostic trail and a570      // future callId-collision bug would silently delete the571      // wrong fr.572      debugLogger.warn(573        `[REPAIR] Dropped ${result.droppedDuplicates.length} duplicate ` +574          `functionResponse(s) for callId(s): ${result.droppedDuplicates575            .map((e) => `${e.name}(${e.callId})`)576            .join(', ')}`,577      );578    }579    return result;580  }581 582  setHistory(history: Content[]) {583    this.getChat().setHistory(history);584    // Replacing history wholesale drops any prior read_file tool585    // results the FileReadCache still believes the model has seen.586    // Without clearing, a follow-up Read of an unchanged file would587    // return the file_unchanged placeholder for bytes that no longer588    // exist in the new history.589    debugLogger.debug('[FILE_READ_CACHE] clear after setHistory');590    this.config.getFileReadCache().clear();591    this.forceFullIdeContext = true;592  }593 594  truncateHistory(keepCount: number) {595    // Use the O(1) length getter rather than getHistory() — the latter596    // structuredClone's the entire history just to read .length, which597    // gets expensive in long-running sessions.598    const prevLen = this.getChat().getHistoryLength();599    this.getChat().truncateHistory(keepCount);600    // Decide whether to invalidate based on the *actual* post-truncate601    // length, not on the keepCount argument. Comparing keepCount alone602    // misses pathological inputs (e.g. NaN: slice(0, NaN) returns [],603    // emptying history, but `NaN < prevLen` is false and would skip604    // the clear, reintroducing the file_unchanged placeholder bug).605    const newLen = this.getChat().getHistoryLength();606    if (newLen < prevLen) {607      debugLogger.debug(608        `[FILE_READ_CACHE] clear after truncateHistory(keep=${keepCount}, prev=${prevLen}, new=${newLen})`,609      );610      this.config.getFileReadCache().clear();611    }612    this.forceFullIdeContext = true;613  }614 615  async setTools(): Promise<void> {616    if (!this.isInitialized()) {617      return;618    }619 620    const toolRegistry = this.config.getToolRegistry();621    await toolRegistry.warmAll();622    const deferredTools = this.resolveDeferredToolsForReminder();623    const toolDeclarations = toolRegistry.getFunctionDeclarations();624    const tools: Tool[] = [{ functionDeclarations: toolDeclarations }];625    this.getChat().setTools(tools);626    this.queueAddedMcpToolsReminder(deferredTools ?? []);627    recordStartupEvent('gemini_tools_updated', {628      toolCount: toolDeclarations.length,629      deferredCount: deferredTools?.length ?? 0,630    });631  }632 633  /**634   * Signal that shutdown is imminent. Subsequent calls to background memory635   * tasks (extract, dream, skill review) will be skipped so the process can636   * exit cleanly without spawning new work.637   */638  requestShutdown(): void {639    this.shutdownRequested = true;640  }641 642  /**643   * Abort and release the pending auto-memory prefetch in one step.644   * Safe to call when no prefetch is pending — does nothing. Centralises645   * the abort-then-clear idiom so every cleanup path (resetChat, early646   * returns, finally) cannot half-fix one without the other.647   *648   * If the handle has already settled (recall completed but consume point649   * hadn't run yet), the settled result is discarded — logged at debug so650   * operators can diagnose missing-memory scenarios.651   */652  private cancelPendingMemoryPrefetch(): void {653    const handle = this.pendingMemoryPrefetch;654    if (!handle) return;655    if (handle.settledAt !== null && !handle.consumed) {656      debugLogger.debug('Discarding settled but unconsumed memory prefetch.');657    }658    handle.controller.abort();659    this.pendingMemoryPrefetch = undefined;660  }661 662  /**663   * Atomically consume the pending prefetch if it has already settled.664   * Returns the recall result (caller decides where to inject it in665   * `requestToSend`), or `null` if there's nothing to consume yet.666   *667   * Centralises the consume-and-mark dance so the UserQuery and ToolResult668   * inject sites can't drift on the guard logic.669   */670  private async tryConsumeMemoryPrefetch(): Promise<RelevantAutoMemoryPromptResult | null> {671    const handle = this.pendingMemoryPrefetch;672    if (!handle || handle.settledAt === null || handle.consumed) {673      return null;674    }675    handle.consumed = true;676    this.pendingMemoryPrefetch = undefined;677    const result = await handle.promise; // already settled, returns immediately678    if (result.prompt) {679      for (const doc of result.selectedDocs) {680        this.surfacedRelevantAutoMemoryPaths.add(doc.filePath);681      }682    }683    return result;684  }685 686  async resetChat(): Promise<void> {687    const memBefore = process.memoryUsage();688    const historyLength = this.chat?.getHistoryLength() ?? 0;689    if (debugLogger.isEnabled()) {690      debugLogger.debug(691        `[RESET_CHAT_START] Starting resetChat, ` +692          `historyLength=${historyLength}, ` +693          `heapUsed=${(memBefore.heapUsed / 1024 / 1024).toFixed(1)}MB, ` +694          `rss=${(memBefore.rss / 1024 / 1024).toFixed(1)}MB`,695      );696    }697 698    this.initializedSessionId = undefined;699    this.surfacedRelevantAutoMemoryPaths.clear();700    this.cachedGitStatus = undefined;701    this.lastApiCompletionTimestamp = null;702    this.lastHookMicrocompactionTimestamp = null;703    this.recentCompletedToolNames = [];704    // startChat() rewrites the chat to its initial state. Any prior705    // read_file tool results the FileReadCache still tracks are no706    // longer in history, so a follow-up Read would serve a placeholder707    // pointing at content the model can no longer retrieve.708    debugLogger.debug('[FILE_READ_CACHE] clear after resetChat');709    this.config.getFileReadCache().clear();710    // Clean up old tool result overflow files on /clear711    void cleanupOldToolResults(Storage.getGlobalTempDir(), 24 * 60 * 60 * 1000);712    this.config.getBaseLlmClient().clearPerModelGeneratorCache();713    // Abort any in-flight auto-memory recall so the stale controller714    // does not leak into the next session.715    this.cancelPendingMemoryPrefetch();716    // Drop any deferred tools revealed this session so /clear really gives717    // a clean slate. We don't clear inside startChat itself because that path718    // is also taken by compression (which preserves the session), and719    // compression should keep previously-revealed tools so the model can720    // continue using them without re-running ToolSearch.721    this.config.getToolRegistry().clearRevealedDeferredTools();722    await this.startChat(undefined, SessionStartSource.Clear);723    this.initializedSessionId = this.config.getSessionId();724 725    const memAfter = process.memoryUsage();726    const newHistoryLength = this.chat?.getHistoryLength() ?? 0;727    if (debugLogger.isEnabled()) {728      debugLogger.debug(729        `[RESET_CHAT_END] resetChat completed, ` +730          `oldHistoryLength=${historyLength}, ` +731          `newHistoryLength=${newHistoryLength}, ` +732          `heapUsed=${(memAfter.heapUsed / 1024 / 1024).toFixed(1)}MB, ` +733          `rss=${(memAfter.rss / 1024 / 1024).toFixed(1)}MB, ` +734          `heapDiff=${((memAfter.heapUsed - memBefore.heapUsed) / 1024 / 1024).toFixed(1)}MB`,735      );736    }737  }738 739  getLoopDetectionService(): LoopDetectionService {740    return this.loopDetector;741  }742 743  async addDirectoryContext(): Promise<void> {744    if (!this.chat) {745      return;746    }747 748    this.getChat().addHistory({749      role: 'user',750      parts: [{ text: await getDirectoryContextString(this.config) }],751    });752  }753 754  async addWorkingDirectoryChangedContext(755    oldDir: string,756    newDir: string,757  ): Promise<void> {758    if (!this.chat) {759      return;760    }761 762    this.cachedGitStatus = undefined;763    await this.refreshSystemInstruction();764    this.getChat().addHistory({765      role: 'user',766      parts: [767        {768          text:769            `The session's working directory has changed from ${oldDir} to ${newDir} via /cd. ` +770            `The startup directory context above is stale. All tool calls and relative paths now resolve from ${newDir}.`,771        },772      ],773    });774    await this.addDirectoryContext();775  }776 777  private getCachedGitStatus(): string | null {778    if (this.cachedGitStatus === undefined) {779      // Mirror claude-code: append git status (branch + recent commits) to the780      // system prompt so the main agent treats version history as authoritative781      // context, not background noise. Only injected when cwd is a git repo.782      this.cachedGitStatus = getRecentGitStatus(this.config.getCwd());783    }784    return this.cachedGitStatus;785  }786 787  private getMainSessionSystemInstruction(): string {788    const userMemory = this.config.getUserMemory();789    const overrideSystemPrompt = this.config.getSystemPrompt();790    const appendSystemPrompt = this.config.getAppendSystemPrompt();791    const gitStatus = this.getCachedGitStatus();792 793    if (overrideSystemPrompt) {794      const base = getCustomSystemPrompt(795        overrideSystemPrompt,796        userMemory,797        appendSystemPrompt,798      );799      return gitStatus ? base + '\n\n' + gitStatus : base;800    }801 802    const base = getCoreSystemPrompt(803      userMemory,804      this.config.getModel(),805      appendSystemPrompt,806    );807    return gitStatus ? base + '\n\n' + gitStatus : base;808  }809 810  async refreshStartupContextReminder(): Promise<void> {811    if (!this.chat) {812      return;813    }814 815    const currentHistory = this.getChat().getHistory();816    const startupLength = getStartupContextLength(currentHistory);817    if (startupLength === 0) {818      return;819    }820 821    // Slice by the detected prelude length, not a hardcoded 1: a restored822    // legacy session stores startup context as a [user(env), model("Got823    // it…")] pair (getStartupContextLength === 2), so slice(1) would leave824    // the orphaned model-ack entry behind when re-prepending the prelude.825    const remaining = currentHistory.slice(startupLength);826    const [[startupContext], snapshotEntries] = await getInitialChatHistory(827      this.config,828    );829    this.seedSkillReminderDedupFromSnapshot(snapshotEntries);830    await this.seedAgentReminderDedupFromCurrent();831    this.getChat().setHistory(832      startupContext ? [startupContext, ...remaining] : remaining,833    );834  }835 836  /**837   * Re-prepend a fresh startup-context prelude after auto-compaction.838   *839   * Auto-compaction runs in-place inside `GeminiChat.sendMessageStream`840   * (`setHistory([summary, ack, ...kept])`) and does NOT route through841   * `tryCompressChat` → `startChat`, so — unlike manual `/compress` — the842   * startup prelude at history[0] is consumed into the summary and never843   * rebuilt. Without this, workspace/env context, deferred-tool metadata,844   * and MCP server instructions are lost for the rest of the session (before845   * this PR they lived in the system instruction and survived compaction).846   *847   * Unlike `refreshStartupContextReminder` (which replaces an existing848   * prelude and no-ops when absent), this prepends when absent. No-ops if a849   * prelude is already present so it can't double-prepend.850   */851  async restoreStartupContextAfterCompaction(): Promise<void> {852    if (!this.chat) {853      return;854    }855 856    const currentHistory = this.getChat().getHistory();857    if (getStartupContextLength(currentHistory) !== 0) {858      return;859    }860 861    const [[startupContext], snapshotEntries] = await getInitialChatHistory(862      this.config,863    );864    this.seedSkillReminderDedupFromSnapshot(snapshotEntries);865    await this.seedAgentReminderDedupFromCurrent();866    if (startupContext) {867      this.getChat().setHistory([startupContext, ...currentHistory]);868    }869  }870 871  /**872   * Rebuilds the main-session system instruction from the current873   * `userMemory` / model / prompt overrides and re-binds it to the live chat.874   *875   * Use this after mutating inputs that feed into the system instruction876   * (e.g. user memory refreshed from `output-language.md`) so the change877   * takes effect on the next turn without restarting the session. No-op if878   * no chat has been started yet.879   */880  async refreshSystemInstruction(): Promise<void> {881    if (!this.chat) {882      return;883    }884    await this.config.getToolRegistry().warmAll();885    this.chat.setSystemInstruction(this.getMainSessionSystemInstruction());886    if (this.lastSessionStartContext && this.lastSessionStartSource) {887      this.chat.applySessionStartContext(888        this.lastSessionStartContext,889        this.lastSessionStartSource,890      );891    }892  }893 894  /**895   * Computes the deferred-tools list that should be announced through896   * user-role system reminders.897   *898   * Caller MUST `await toolRegistry.warmAll()` first — this method only899   * inspects the registry's eager state and would otherwise miss factory-900   * backed deferred tools.901   *902   * Side effect: when ToolSearch is not registered (e.g. `--exclude-tools903   * tool_search` or a deny rule), every deferred tool is eagerly revealed904   * here so it lands in the declaration list. Skipping this would leave the905   * tool both off the declarations AND off the deferred-summary list (since906   * `undefined` is returned in that branch) — a silent disappearance that's907   * harder to diagnose than seeing the tool name absent from `/mcp` output.908   *909   * Returns `undefined` when ToolSearch is unavailable: reminders must not910   * advertise tools the model has no way to load on demand.911   */912  private resolveDeferredToolsForReminder(): DeferredToolSummary[] | undefined {913    const toolRegistry = this.config.getToolRegistry();914    const deferredSummary = toolRegistry.getDeferredToolSummary();915    const toolSearchAvailable = !!toolRegistry.getTool(ToolNames.TOOL_SEARCH);916    if (!toolSearchAvailable) {917      if (deferredSummary.length > 0) {918        for (const t of deferredSummary) {919          toolRegistry.revealDeferredTool(t.name);920        }921      }922      return undefined;923    }924    return deferredSummary.filter(925      (t) => !toolRegistry.isDeferredToolRevealed(t.name),926    );927  }928 929  private rememberAnnouncedDeferredTools(930    deferredTools: readonly DeferredToolSummary[] | undefined,931  ): void {932    this.announcedDeferredToolNames = new Set(933      (deferredTools ?? []).map((tool) => tool.name),934    );935    this.announcedMcpToolNames = new Set(936      (deferredTools ?? [])937        .filter((tool) => tool.serverName)938        .map((tool) => tool.name),939    );940    this.pendingAddedMcpTools.clear();941    this.pendingRemovedMcpToolNames.clear();942  }943 944  private queueAddedMcpToolsReminder(945    deferredTools: readonly DeferredToolSummary[],946  ): void {947    const currentDeferredNames = new Set(948      deferredTools.map((tool) => tool.name),949    );950    const currentMcpToolNames = new Set(951      deferredTools.filter((tool) => tool.serverName).map((tool) => tool.name),952    );953    for (const name of this.pendingAddedMcpTools.keys()) {954      if (!currentDeferredNames.has(name)) {955        this.pendingAddedMcpTools.delete(name);956      }957    }958    for (const name of this.pendingRemovedMcpToolNames) {959      if (currentMcpToolNames.has(name)) {960        this.pendingRemovedMcpToolNames.delete(name);961      }962    }963 964    // Drop announced names that are no longer deferred (e.g. an MCP server965    // disconnected and removeMcpToolsByServer() pruned its tools). Without966    // this, a tool that reconnects later is still in announcedDeferredToolNames967    // and gets silently skipped below, so the user never sees the "new tools968    // available" reminder even though setTools() re-declared the tool.969    for (const name of this.announcedDeferredToolNames) {970      if (!currentDeferredNames.has(name)) {971        this.announcedDeferredToolNames.delete(name);972      }973    }974    for (const name of this.announcedMcpToolNames) {975      if (!currentMcpToolNames.has(name)) {976        this.pendingRemovedMcpToolNames.add(name);977      }978    }979 980    for (const tool of deferredTools) {981      if (tool.serverName) {982        if (!this.announcedMcpToolNames.has(tool.name)) {983          this.pendingAddedMcpTools.set(tool.name, tool);984        }985      }986      this.announcedDeferredToolNames.add(tool.name);987    }988  }989 990  private drainPendingAddedMcpToolsReminder(): void {991    if (992      this.pendingAddedMcpTools.size === 0 &&993      this.pendingRemovedMcpToolNames.size === 0994    ) {995      return;996    }997 998    const addedMcpTools = Array.from(this.pendingAddedMcpTools.values());999    const removedMcpToolNames = Array.from(this.pendingRemovedMcpToolNames);1000    const reminder = buildChangedMcpToolsReminder(1001      addedMcpTools,1002      removedMcpToolNames,1003    );1004 1005    if (!reminder) {1006      return;1007    }1008 1009    this.getChat().addHistory({1010      role: 'user',1011      parts: [{ text: reminder }],1012    });1013 1014    for (const name of removedMcpToolNames) {1015      this.announcedMcpToolNames.delete(name);1016    }1017    for (const tool of addedMcpTools) {1018      this.announcedMcpToolNames.add(tool.name);1019    }1020    this.pendingAddedMcpTools.clear();1021    this.pendingRemovedMcpToolNames.clear();1022  }1023 1024  /**1025   * Per-turn delta for skills/commands that became invocable after session start1026   * — skills enabled mid-session (e.g. via `/skills`) and MCP prompts added after1027   * startup. Emitted as a tail `<system-reminder>` only, so it never mutates the1028   * cached tools/system/messages prefix. Deduped via `announcedSkillReminderKeys`.1029   *1030   * The first call after a (re)built startup prelude seeds the announced set from1031   * the current skills and emits nothing — the startup snapshot already listed1032   * them (mirrors Claude Code's `suppressNextSkillListing` and its decision not1033   * to re-inject the listing after compaction). Conditional path-activations are1034   * announced inline on the tool result by `coreToolScheduler`, so they are1035   * recorded here as announced (not re-queued) to avoid a double announcement.1036   */1037  private async drainSkillAndCommandReminders(): Promise<void> {1038    const toolRegistry = this.config.getToolRegistry();1039    // Only relevant when the model can actually invoke skills (subagents often1040    // run without the Skill tool).1041    if (!toolRegistry?.getTool(ToolNames.SKILL)) {1042      return;1043    }1044    const skillManager = this.config.getSkillManager();1045    if (!skillManager) {1046      return;1047    }1048 1049    let entries: AvailableSkillEntry[];1050    try {1051      ({ entries } = await collectAvailableSkillEntries(1052        skillManager,1053        this.config,1054      ));1055    } catch (error) {1056      debugLogger.warn(1057        'drainSkillAndCommandReminders: collectAvailableSkillEntries failed',1058        error,1059      );1060      return;1061    }1062 1063    const currentKeys = new Set(entries.map(GeminiClient.skillEntryKey));1064    const wasInitialized = this.skillRemindersInitialized;1065    const removedNames: string[] = [];1066 1067    // Prune announced keys no longer present so a later re-enable / reconnect1068    // re-announces (mirrors the MCP added-tools prune above).1069    for (const key of this.announcedSkillReminderKeys) {1070      if (!currentKeys.has(key)) {1071        if (wasInitialized) {1072          removedNames.push(key.slice(key.indexOf(':') + 1));1073        }1074        this.announcedSkillReminderKeys.delete(key);1075      }1076    }1077 1078    // Safety net: if seedSkillReminderDedupFromSnapshot was never called (e.g.1079    // edge-case construction path), mark initialized but do NOT seed from1080    // current entries — no startup snapshot was shown to the model, so all1081    // entries are genuinely new and should be announced by the code below.1082    // Seeding here used to silently swallow late registrations (cmd:* keys1083    // and MCP prompts discovered after startChat) by marking them as1084    // "already announced" when the model had never seen them.1085    if (!this.skillRemindersInitialized) {1086      this.skillRemindersInitialized = true;1087    }1088 1089    // Consume skill keys that coreToolScheduler announced inline on a tool1090    // result this turn (e.g. path-activated conditional skills). Mark them as1091    // announced so the drain below does not re-announce them. This fixes the1092    // subagent shared-SkillManager case: the inline reminder lands in the1093    // subagent's discarded transcript, but the parent's drain now skips those1094    // keys because the scheduler recorded them on the shared Config.1095    const inlineKeys = this.config.consumeInlineAnnouncedSkillKeys();1096    for (const key of inlineKeys) {1097      this.announcedSkillReminderKeys.add(key);1098    }1099 1100    // Announce every genuinely new skill/command that was not already1101    // announced — either in the startup snapshot, a prior drain, or inline1102    // by coreToolScheduler above.1103    const newEntries: AvailableSkillEntry[] = [];1104    for (const entry of entries) {1105      const key = GeminiClient.skillEntryKey(entry);1106      if (this.announcedSkillReminderKeys.has(key)) {1107        continue;1108      }1109      this.announcedSkillReminderKeys.add(key);1110      newEntries.push(entry);1111    }1112 1113    if (newEntries.length === 0 && removedNames.length === 0) {1114      return;1115    }1116    const reminder = buildChangedSkillsReminder(newEntries, removedNames);1117    if (!reminder) {1118      return;1119    }1120    this.getChat().addHistory({1121      role: 'user',1122      parts: [{ text: reminder }],1123    });1124  }1125 1126  private async drainAgentReminders(): Promise<void> {1127    const toolRegistry = this.config.getToolRegistry();1128    if (!toolRegistry?.getTool(ToolNames.AGENT)) {1129      return;1130    }1131 1132    if (!this.agentRemindersInitialized) {1133      await this.seedAgentReminderDedupFromCurrent();1134      return;1135    }1136 1137    let agents: AgentAvailabilityEntry[];1138    try {1139      agents = await this.config.getSubagentManager().listSubagents();1140    } catch (error) {1141      debugLogger.warn('drainAgentReminders: listSubagents failed', error);1142      return;1143    }1144 1145    const currentByName = new Map(agents.map((agent) => [agent.name, agent]));1146    const addedAgents: AgentAvailabilityEntry[] = [];1147    const removedAgentNames: string[] = [];1148 1149    for (const name of this.announcedAgentReminderNames) {1150      if (!currentByName.has(name)) {1151        removedAgentNames.push(name);1152      }1153    }1154 1155    for (const agent of currentByName.values()) {1156      if (this.announcedAgentReminderNames.has(agent.name)) {1157        continue;1158      }1159      addedAgents.push({1160        name: agent.name,1161        description: agent.description,1162      });1163    }1164 1165    const reminder = buildChangedAgentsReminder(addedAgents, removedAgentNames);1166    if (!reminder) {1167      return;1168    }1169    this.getChat().addHistory({1170      role: 'user',1171      parts: [{ text: reminder }],1172    });1173 1174    for (const name of removedAgentNames) {1175      this.announcedAgentReminderNames.delete(name);1176    }1177    for (const agent of addedAgents) {1178      this.announcedAgentReminderNames.add(agent.name);1179    }1180  }1181 1182  private toPermissionMode(approvalMode: ApprovalMode): PermissionMode {1183    switch (approvalMode) {1184      case ApprovalMode.DEFAULT:1185        return PermissionMode.Default;1186      case ApprovalMode.PLAN:1187        return PermissionMode.Plan;1188      case ApprovalMode.AUTO_EDIT:1189        return PermissionMode.AutoEdit;1190      case ApprovalMode.AUTO:1191        return PermissionMode.Auto;1192      case ApprovalMode.YOLO:1193        return PermissionMode.Yolo;1194      default:1195        return PermissionMode.Default;1196    }1197  }1198 1199  private async fireSessionStartHook(1200    source: SessionStartSource,

Showing the first 1,200 of 3039 lines. Download the file for the rest.

basant307/AI_Governance_Project · CoolFace