feat: track context window utilization (#2064)

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

Co-authored-by: Anshi Agarwal <anshiagarwal@microsoft.com>
Co-authored-by: Brady Gaster <41929050+bradygaster@users.noreply.github.com>
Copilot-Session: c8ae7013-1eb9-4f99-8762-ea407e3633f3

Author: 3 people

.changeset/1718-context-utilization.md

@@ -0,0 +1,6 @@
+---
+"@bradygaster/squad-cli": patch
+"@bradygaster/squad-sdk": patch
+---
+
+Track context-window utilization from runtime and estimated token data, publish typed events and OpenTelemetry metrics, warn interactive-shell users at a configurable threshold, and report estimated utilization from `squad watch`.

docs/src/content/docs/reference/cli.md

@@ -354,6 +354,7 @@ First match wins.
 |----------|---------|--------|
 | `SQUAD_CLIENT` | Detected client platform | `cli`, `vscode` |
 | `COPILOT_TOKEN` | Copilot auth token (SDK usage) | Token string |
+| `SQUAD_CONTEXT_WARNING_THRESHOLD` | Interactive-shell context utilization warning threshold | Fraction greater than `0` and at most `1` (default: `0.8`) |
 
 ---
 

docs/src/content/docs/reference/config.md

@@ -195,6 +195,7 @@ First match wins.
 |----------|---------|
 | `SQUAD_CLIENT` | Detected client (`cli` or `vscode`) |
 | `COPILOT_TOKEN` | Auth token for SDK usage |
+| `SQUAD_CONTEXT_WARNING_THRESHOLD` | Interactive-shell context warning threshold as a fraction greater than `0` and at most `1` (default: `0.8`) |
 
 ---
 

docs/src/content/docs/reference/sdk.md

@@ -317,7 +317,28 @@ await telemetry.shutdown();
 await shutdownOTel();
 ```
 
-The root package also exports streaming helpers, cost tracking, offline mode, runtime events, OpenTelemetry metrics, and internationalization helpers.
+`StreamingPipeline` can also track the latest context-window occupancy for each attached session. Runtime `context_usage` events are exact; regular usage events provide an estimate when a model context-window resolver is configured.
+
+```typescript
+import { RuntimeEventBus, StreamingPipeline } from '@bradygaster/squad-sdk';
+
+const eventBus = new RuntimeEventBus();
+const pipeline = new StreamingPipeline({
+  eventBus,
+  contextWarningThreshold: 0.8,
+  resolveContextWindow: model => model === 'gpt-5.4' ? 200_000 : undefined,
+});
+
+pipeline.onContextUtilization(snapshot => {
+  if (snapshot.thresholdCrossed) {
+    console.warn(`${snapshot.sessionId} context is ${(snapshot.utilization * 100).toFixed(1)}% full`);
+  }
+});
+```
+
+Every accepted sample is also emitted as a typed `context:utilization` event and recorded in the `squad.context.utilization` OpenTelemetry gauge. Exact runtime samples take precedence over later estimates for the same session.
+
+The root package also exports cost tracking, offline mode, runtime events, OpenTelemetry metrics, and internationalization helpers.
 
 ## Other Public Modules
 

packages/squad-cli/src/cli/commands/watch/agent-spawn.ts

@@ -10,12 +10,85 @@
  */
 
 import { execFile, execFileSync } from 'node:child_process';
+import { existsSync, mkdirSync, readFileSync, unlinkSync } from 'node:fs';
+import { randomUUID } from 'node:crypto';
+import path from 'node:path';
 import type { WatchContext } from './types.js';
 import { withAdditionalMcpConfig } from '../../core/copilot-invocation.js';
 
 /** True when running on Windows — used to gate `shell: true`. */
 export const IS_WINDOWS = process.platform === 'win32';
 
+export interface CopilotUsageOutput {
+  currentModel: string;
+  lastCallInputTokens: number;
+  lastCallOutputTokens: number;
+}
+
+export interface AgentSpawnResult {
+  success: boolean;
+  error?: string;
+  usage?: CopilotUsageOutput;
+  usageError?: string;
+}
+
+export interface CopilotUsageCapture {
+  args: string[];
+  filePath: string;
+}
+
+function isNonNegativeInteger(value: unknown): value is number {
+  return typeof value === 'number' && Number.isInteger(value) && value >= 0;
+}
+
+/** Add Copilot's per-invocation usage output flag under the effective state directory. */
+export function withCopilotUsageOutput(
+  args: string[],
+  stateRoot: string,
+): CopilotUsageCapture {
+  const logDir = path.join(stateRoot, 'log');
+  mkdirSync(logDir, { recursive: true });
+  const filePath = path.join(logDir, `context-usage-${randomUUID()}.json`);
+  return {
+    args: [...args, '--usage-output-file', filePath],
+    filePath,
+  };
+}
+
+/** Strictly parse the subset of Copilot usage output needed for context estimation. */
+export function parseCopilotUsageOutput(value: unknown): CopilotUsageOutput | undefined {
+  if (!value || typeof value !== 'object') return undefined;
+  const record = value as Record<string, unknown>;
+  if (
+    typeof record['currentModel'] !== 'string'
+    || record['currentModel'].trim() === ''
+    || !isNonNegativeInteger(record['lastCallInputTokens'])
+    || !isNonNegativeInteger(record['lastCallOutputTokens'])
+  ) {
+    return undefined;
+  }
+  return {
+    currentModel: record['currentModel'],
+    lastCallInputTokens: record['lastCallInputTokens'],
+    lastCallOutputTokens: record['lastCallOutputTokens'],
+  };
+}
+
+function readCopilotUsageOutput(filePath: string): {
+  usage?: CopilotUsageOutput;
+  error?: string;
+} {
+  if (!existsSync(filePath)) {
+    return { error: 'Copilot usage output was not created' };
+  }
+  try {
+    const usage = parseCopilotUsageOutput(JSON.parse(readFileSync(filePath, 'utf8')));
+    return usage ? { usage } : { error: 'Copilot usage output has an invalid shape' };
+  } catch (error) {
+    return { error: `Could not read Copilot usage output: ${(error as Error).message}` };
+  }
+}
+
 /**
  * Escape an argument for safe use with cmd.exe when `shell: true`.
  *
@@ -206,9 +279,10 @@ export function spawnAgent(
   cwd: string,
   timeoutMs: number,
   pidTracking?: { tracker: NonNullable<WatchContext['pidTracker']>; label: string },
-): Promise<{ success: boolean; error?: string }> {
+  usageOutputFile?: string,
+): Promise<AgentSpawnResult> {
   const safeArgs = escapeArgs(args);
-  return new Promise<{ success: boolean; error?: string }>((resolve) => {
+  return new Promise<AgentSpawnResult>((resolve) => {
     const cp = execFile(
       cmd,
       safeArgs,
@@ -219,12 +293,31 @@ export function spawnAgent(
         shell: IS_WINDOWS,
       },
       (err) => {
+        const usageResult = usageOutputFile
+          ? readCopilotUsageOutput(usageOutputFile)
+          : {};
+        if (usageOutputFile && existsSync(usageOutputFile)) {
+          try {
+            unlinkSync(usageOutputFile);
+          } catch (cleanupError) {
+            usageResult.error ??= `Could not remove Copilot usage output: ${(cleanupError as Error).message}`;
+          }
+        }
         if (err) {
           const execErr = err as Error & { killed?: boolean };
           const msg = execErr.killed ? 'Timed out' : execErr.message;
-          resolve({ success: false, error: msg });
+          resolve({
+            success: false,
+            error: msg,
+            usage: usageResult.usage,
+            usageError: usageResult.error,
+          });
         } else {
-          resolve({ success: true });
+          resolve({
+            success: true,
+            usage: usageResult.usage,
+            usageError: usageResult.error,
+          });
         }
       },
     );

packages/squad-cli/src/cli/commands/watch/capabilities/execute.ts

@@ -7,9 +7,15 @@ import { existsSync } from 'node:fs';
 import path from 'node:path';
 import type { WatchCapability, WatchContext, PreflightResult, CapabilityResult } from '../types.js';
 import type { MachineCapabilities } from '@bradygaster/squad-sdk/ralph/capabilities';
+import { SquadClient } from '@bradygaster/squad-sdk/client';
 import { createVerboseLogger } from '../verbose.js';
 import { loadAgentCharter } from '../../../shell/spawn.js';
-import { buildCopilotCommand, spawnAgent } from '../agent-spawn.js';
+import {
+  buildCopilotCommand,
+  spawnAgent,
+  withCopilotUsageOutput,
+} from '../agent-spawn.js';
+import type { AgentSpawnResult, CopilotUsageOutput } from '../agent-spawn.js';
 
 /** Normalized work item for execution. */
 export interface ExecutableWorkItem {
@@ -138,7 +144,7 @@ async function executeAll(
   issues: ExecutableWorkItem[],
   context: WatchContext,
   timeoutMs: number,
-): Promise<{ success: boolean; error?: string }> {
+): Promise<AgentSpawnResult> {
   const prompt = buildAgentPrompt(issues, context.teamRoot, context.stateRoot);
 
   // Load Ralph's charter to give the spawned session full specialist context.
@@ -151,15 +157,67 @@ async function executeAll(
   }
 
   const fullPrompt = charterPrefix + prompt;
-  const { cmd, args } = buildCopilotCommand(fullPrompt, context);
+  const { cmd, args: baseArgs } = buildCopilotCommand(fullPrompt, context);
+  let usageCapture: ReturnType<typeof withCopilotUsageOutput> | undefined;
+  let usageSetupError: string | undefined;
+  if (!context.agentCmd) {
+    try {
+      usageCapture = withCopilotUsageOutput(baseArgs, context.stateRoot);
+    } catch (error) {
+      usageSetupError = `Could not initialize Copilot usage capture: ${(error as Error).message}`;
+    }
+  }
+  const args = usageCapture?.args ?? baseArgs;
 
   // Track child PID for cleanup on exit/crash
   const issueNums = issues.map(i => `#${i.number}`).join(',');
   const pidTracking = context.pidTracker
     ? { tracker: context.pidTracker, label: `copilot-session-${issueNums}` }
     : undefined;
 
-  return spawnAgent(cmd, args, context.teamRoot, timeoutMs, pidTracking);
+  const result = await spawnAgent(
+    cmd,
+    args,
+    context.teamRoot,
+    timeoutMs,
+    pidTracking,
+    usageCapture?.filePath,
+  );
+  return usageSetupError ? { ...result, usageError: usageSetupError } : result;
+}
+
+interface EstimatedContextUtilization {
+  model: string;
+  occupiedTokens: number;
+  contextWindowTokens: number;
+  utilization: number;
+  source: 'estimated';
+}
+
+async function estimateContextUtilization(
+  usage: CopilotUsageOutput,
+  teamRoot: string,
+): Promise<EstimatedContextUtilization | undefined> {
+  const client = new SquadClient({ cwd: teamRoot });
+  try {
+    const models = await client.listModels();
+    const model = models.find(candidate =>
+      candidate.id === usage.currentModel || candidate.name === usage.currentModel);
+    const contextWindowTokens = model?.capabilities.limits.max_context_window_tokens;
+    if (!Number.isInteger(contextWindowTokens) || !contextWindowTokens || contextWindowTokens <= 0) {
+      return undefined;
+    }
+    const occupiedTokens = usage.lastCallInputTokens + usage.lastCallOutputTokens;
+    return {
+      model: usage.currentModel,
+      occupiedTokens,
+      contextWindowTokens,
+      utilization: occupiedTokens / contextWindowTokens,
+      source: 'estimated',
+    };
+  } finally {
+    await client.disconnect();
+  }
 }
 
 export class ExecuteCapability implements WatchCapability {
@@ -209,13 +267,44 @@ export class ExecuteCapability implements WatchCapability {
 
       // Single agent invocation with all issues — agent reads ralph-instructions.md
       const result = await executeAll(eligible, context, timeout);
+      let contextUtilization: EstimatedContextUtilization | undefined;
+      let contextTelemetryError = result.usageError;
+      if (result.usage) {
+        try {
+          contextUtilization = await estimateContextUtilization(result.usage, context.teamRoot);
+          if (!contextUtilization) {
+            contextTelemetryError = `No context-window limit found for model ${result.usage.currentModel}`;
+          }
+        } catch (error) {
+          contextTelemetryError = `Could not resolve model context window: ${(error as Error).message}`;
+        }
+      }
+
+      const contextSummary = contextUtilization
+        ? `; context ${(contextUtilization.utilization * 100).toFixed(1)}% (${contextUtilization.source})`
+        : contextTelemetryError
+          ? `; context telemetry unavailable (${contextTelemetryError})`
+          : '';
+      if (contextUtilization) {
+        vlog.log(
+          `Execute: context ${(contextUtilization.utilization * 100).toFixed(1)}% `
+          + `(${contextUtilization.occupiedTokens}/${contextUtilization.contextWindowTokens} tokens, estimated)`,
+        );
+      } else if (contextTelemetryError) {
+        vlog.log(`Execute: ${contextTelemetryError}`);
+      }
 
       return {
         success: result.success,
         summary: result.success
-          ? `agent dispatched with ${eligible.length} issues`
-          : `agent failed: ${result.error}`,
-        data: { dispatched: eligible.length, success: result.success },
+          ? `agent dispatched with ${eligible.length} issues${contextSummary}`
+          : `agent failed: ${result.error}${contextSummary}`,
+        data: {
+          dispatched: eligible.length,
+          success: result.success,
+          ...(contextUtilization ? { contextUtilization } : {}),
+          ...(contextTelemetryError ? { contextTelemetryError } : {}),
+        },
       };
     } catch (e) {
       return { success: false, summary: `execute error: ${(e as Error).message}` };