fix: add subagent context hygiene warnings

src/hooks/todo-continuation/index.test.ts

@@ -1,6 +1,10 @@
 import { describe, expect, mock, test } from 'bun:test';
 import { SLIM_INTERNAL_INITIATOR_MARKER } from '../../utils';
 import { createTodoContinuationHook } from './index';
+import {
+  SUBAGENT_CONTEXT_HYGIENE_INSTRUCTION_OPEN,
+  SUBAGENT_CONTEXT_HYGIENE_REMINDER,
+} from './subagent-context-hygiene';
 import {
   TODO_FINAL_ACTIVE_REMINDER,
   TODO_HYGIENE_REMINDER,
@@ -22,9 +26,26 @@ describe('createTodoContinuationHook', () => {
         parts?: Array<{ type?: string; text?: string }>;
       }>;
     };
+    providersResult?: {
+      data?: {
+        providers?: Array<{
+          id: string;
+          models?: Record<string, { limit?: { context?: number } }>;
+        }>;
+      };
+    };
   }) {
     return {
+      directory: '/repo',
       client: {
+        config: {
+          providers: mock(
+            async () =>
+              overrides?.providersResult ?? {
+                data: { providers: [] },
+              },
+          ),
+        },
         session: {
           todo: mock(async () => overrides?.todoResult ?? { data: [] }),
           messages: mock(async () => overrides?.messagesResult ?? { data: [] }),
@@ -94,6 +115,49 @@ describe('createTodoContinuationHook', () => {
       .join('\n');
   }
 
+  function providersWithContextLimit(context?: number) {
+    return {
+      data: {
+        providers: [
+          {
+            id: 'openai',
+            models: {
+              'gpt-5.5': context ? { limit: { context } } : {},
+            },
+          },
+        ],
+      },
+    };
+  }
+
+  async function sendContextUsage(
+    hook: ReturnType<typeof createTodoContinuationHook>,
+    input: {
+      sessionID: string;
+      agent: string;
+      input: number;
+      cacheRead?: number;
+    },
+  ) {
+    await hook.handleEvent({
+      event: {
+        type: 'message.updated',
+        properties: {
+          info: {
+            sessionID: input.sessionID,
+            agent: input.agent,
+            providerID: 'openai',
+            modelID: 'gpt-5.5',
+            tokens: {
+              input: input.input,
+              cache: { read: input.cacheRead ?? 0 },
+            },
+          },
+        },
+      },
+    });
+  }
+
   describe('tool toggle', () => {
     test('calling auto_continue execute with { enabled: true } sets state', async () => {
       const ctx = createMockContext();
@@ -651,6 +715,155 @@ describe('createTodoContinuationHook', () => {
     });
   });
 
+  describe('subagent context hygiene', () => {
+    test('injects cache-friendly context warning for subagents over configured context limit after tool use', async () => {
+      const ctx = createMockContext({
+        providersResult: providersWithContextLimit(1000),
+      });
+      const hook = createTodoContinuationHook(ctx);
+      const output = userMessages('review this diff', 'sub1', 'oracle');
+
+      hook.handleChatMessage({ sessionID: 'sub1', agent: 'oracle' });
+      await hook.handleMessagesTransform(output);
+      await sendContextUsage(hook, {
+        sessionID: 'sub1',
+        agent: 'oracle',
+        input: 500,
+        cacheRead: 75,
+      });
+      await hook.handleToolExecuteAfter({ tool: 'read', sessionID: 'sub1' });
+      await hook.handleMessagesTransform(output);
+
+      expect(allMessageText(output)).toContain(
+        SUBAGENT_CONTEXT_HYGIENE_REMINDER,
+      );
+      expect(allMessageText(output)).toContain(
+        SUBAGENT_CONTEXT_HYGIENE_INSTRUCTION_OPEN,
+      );
+    });
+
+    test('does not inject context warning when model has no configured context limit', async () => {
+      const ctx = createMockContext({
+        providersResult: providersWithContextLimit(),
+      });
+      const hook = createTodoContinuationHook(ctx);
+      const output = userMessages('review this diff', 'sub1', 'oracle');
+
+      hook.handleChatMessage({ sessionID: 'sub1', agent: 'oracle' });
+      await hook.handleMessagesTransform(output);
+      await sendContextUsage(hook, {
+        sessionID: 'sub1',
+        agent: 'oracle',
+        input: 900,
+        cacheRead: 100,
+      });
+      await hook.handleToolExecuteAfter({ tool: 'read', sessionID: 'sub1' });
+      await hook.handleMessagesTransform(output);
+
+      expect(allMessageText(output)).not.toContain(
+        SUBAGENT_CONTEXT_HYGIENE_REMINDER,
+      );
+    });
+
+    test('does not warn resumed high-context subagent until it uses another tool', async () => {
+      const ctx = createMockContext({
+        providersResult: providersWithContextLimit(1000),
+      });
+      const hook = createTodoContinuationHook(ctx);
+      const output = userMessages('continue previous review', 'sub1', 'oracle');
+
+      hook.handleChatMessage({ sessionID: 'sub1', agent: 'oracle' });
+      await hook.handleMessagesTransform(output);
+      await sendContextUsage(hook, {
+        sessionID: 'sub1',
+        agent: 'oracle',
+        input: 700,
+      });
+      await hook.handleMessagesTransform(output);
+
+      expect(allMessageText(output)).not.toContain(
+        SUBAGENT_CONTEXT_HYGIENE_REMINDER,
+      );
+
+      await hook.handleToolExecuteAfter({ tool: 'grep', sessionID: 'sub1' });
+      await hook.handleMessagesTransform(output);
+
+      expect(allMessageText(output)).toContain(
+        SUBAGENT_CONTEXT_HYGIENE_REMINDER,
+      );
+    });
+
+    test('does not duplicate context warning in same transformed payload', async () => {
+      const ctx = createMockContext({
+        providersResult: providersWithContextLimit(1000),
+      });
+      const hook = createTodoContinuationHook(ctx);
+      const output = userMessages('review this diff', 'sub1', 'oracle');
+
+      hook.handleChatMessage({ sessionID: 'sub1', agent: 'oracle' });
+      await hook.handleMessagesTransform(output);
+      await sendContextUsage(hook, {
+        sessionID: 'sub1',
+        agent: 'oracle',
+        input: 600,
+      });
+      await hook.handleToolExecuteAfter({ tool: 'read', sessionID: 'sub1' });
+      await hook.handleMessagesTransform(output);
+      await hook.handleToolExecuteAfter({ tool: 'glob', sessionID: 'sub1' });
+      await hook.handleMessagesTransform(output);
+
+      expect(
+        allMessageText(output).split(SUBAGENT_CONTEXT_HYGIENE_REMINDER).length -
+          1,
+      ).toBe(1);
+    });
+
+    test('strips only the exact generated context warning block', async () => {
+      const ctx = createMockContext();
+      const hook = createTodoContinuationHook(ctx);
+      const userText = [
+        'custom prompt',
+        '',
+        '<instruction name="subagent_context_hygiene">',
+        'User-authored instruction that must stay.',
+        '</instruction>',
+      ].join('\n');
+      const output = userMessages(userText, 'sub1', 'oracle');
+
+      hook.handleChatMessage({ sessionID: 'sub1', agent: 'oracle' });
+      await hook.handleMessagesTransform(output);
+
+      expect(allMessageText(output)).toContain(
+        'User-authored instruction that must stay.',
+      );
+      expect(allMessageText(output)).not.toContain(
+        SUBAGENT_CONTEXT_HYGIENE_REMINDER,
+      );
+    });
+
+    test('never injects context warning into orchestrator sessions', async () => {
+      const ctx = createMockContext({
+        providersResult: providersWithContextLimit(1000),
+      });
+      const hook = createTodoContinuationHook(ctx);
+      const output = userMessages('do work', 'main1', 'orchestrator');
+
+      hook.handleChatMessage({ sessionID: 'main1', agent: 'orchestrator' });
+      await hook.handleMessagesTransform(output);
+      await sendContextUsage(hook, {
+        sessionID: 'main1',
+        agent: 'orchestrator',
+        input: 700,
+      });
+      await hook.handleToolExecuteAfter({ tool: 'read', sessionID: 'main1' });
+      await hook.handleMessagesTransform(output);
+
+      expect(allMessageText(output)).not.toContain(
+        SUBAGENT_CONTEXT_HYGIENE_REMINDER,
+      );
+    });
+  });
+
   describe('continuation scheduling', () => {
     test('session idle + enabled + incomplete todos → schedules continuation', async () => {
       const ctx = createMockContext({

src/hooks/todo-continuation/index.ts

@@ -6,6 +6,7 @@ import {
   SLIM_INTERNAL_INITIATOR_MARKER,
   withTimeout,
 } from '../../utils';
+import { createSubagentContextHygiene } from './subagent-context-hygiene';
 import { createTodoHygiene } from './todo-hygiene';
 
 const HOOK_NAME = 'todo-continuation';
@@ -200,6 +201,7 @@ export function createTodoContinuationHook(
   const autoEnable = config?.autoEnable ?? false;
   const autoEnableThreshold = config?.autoEnableThreshold ?? 4;
   const requestSignatureBySession = new Map<string, string>();
+  const contextLimitByModel = new Map<string, number | undefined>();
 
   const state: ContinuationState = {
     enabled: false,
@@ -245,6 +247,49 @@ export function createTodoContinuationHook(
     log: (message, meta) => log(`[${HOOK_NAME}] ${message}`, meta),
   });
 
+  const subagentContextHygiene = createSubagentContextHygiene({
+    getContextLimit: async (providerID, modelID) => {
+      const key = `${providerID}/${modelID}`;
+      if (contextLimitByModel.has(key)) {
+        return contextLimitByModel.get(key);
+      }
+
+      try {
+        const result = await ctx.client.config.providers();
+        const providers =
+          (
+            result.data as
+              | {
+                  providers?: Array<{
+                    id: string;
+                    models?: Record<string, { limit?: { context?: unknown } }>;
+                  }>;
+                }
+              | undefined
+          )?.providers ?? [];
+        const provider = providers.find((entry) => entry.id === providerID);
+        const context = provider?.models?.[modelID]?.limit?.context;
+        const limit =
+          typeof context === 'number' && Number.isFinite(context) && context > 0
+            ? context
+            : undefined;
+        contextLimitByModel.set(key, limit);
+        return limit;
+      } catch (error) {
+        log(
+          `[${HOOK_NAME}] Skipped subagent context hygiene: failed to fetch providers`,
+          {
+            providerID,
+            modelID,
+            error: error instanceof Error ? error.message : String(error),
+          },
+        );
+        return undefined;
+      }
+    },
+    log: (message, meta) => log(`[${HOOK_NAME}] ${message}`, meta),
+  });
+
   function inferSessionID(
     messages: ChatTransformMessage[],
     index: number,
@@ -349,6 +394,8 @@ export function createTodoContinuationHook(
       return;
     }
 
+    await subagentContextHygiene.handleMessagesTransform(output);
+
     if (lastUserMessage.agent && lastUserMessage.agent !== 'orchestrator') {
       return;
     }
@@ -435,6 +482,7 @@ export function createTodoContinuationHook(
       return;
     }
 
+    subagentContextHygiene.handleChatMessage(input);
     state.sawChatMessage = true;
     if (input.agent === 'orchestrator') {
       registerOrchestratorSession(input.sessionID);
@@ -469,6 +517,26 @@ export function createTodoContinuationHook(
     const { event } = input;
     const properties = event.properties ?? {};
 
+    await subagentContextHygiene.handleEvent({
+      type: event.type,
+      properties: {
+        info: properties.info as
+          | {
+              id?: string;
+              sessionID?: string;
+              agent?: string;
+              providerID?: string;
+              modelID?: string;
+              tokens?: {
+                input?: unknown;
+                cache?: { read?: unknown };
+              };
+            }
+          | undefined,
+        sessionID: properties.sessionID as string | undefined,
+      },
+    });
+
     hygiene.handleEvent({
       type: event.type,
       properties: {
@@ -870,7 +938,10 @@ export function createTodoContinuationHook(
 
   return {
     tool: { auto_continue: autoContinue },
-    handleToolExecuteAfter: hygiene.handleToolExecuteAfter,
+    handleToolExecuteAfter: async (input, output) => {
+      await hygiene.handleToolExecuteAfter(input, output);
+      await subagentContextHygiene.handleToolExecuteAfter(input);
+    },
     handleMessagesTransform,
     handleEvent,
     handleChatMessage,

src/hooks/todo-continuation/subagent-context-hygiene.ts

@@ -0,0 +1,299 @@
+export const SUBAGENT_CONTEXT_HYGIENE_REMINDER =
+  'Context warning reached; focus only on request-relevant files and finalize when ready.';
+
+export const SUBAGENT_CONTEXT_HYGIENE_INSTRUCTION_OPEN =
+  '<instruction name="subagent_context_hygiene">';
+
+const INSTRUCTION_CLOSE = '</instruction>';
+const CONTEXT_WARNING_THRESHOLD = 0.45;
+const GENERATED_INSTRUCTION = `${SUBAGENT_CONTEXT_HYGIENE_INSTRUCTION_OPEN}\n${SUBAGENT_CONTEXT_HYGIENE_REMINDER}\n${INSTRUCTION_CLOSE}`;
+
+interface Tokens {
+  input?: unknown;
+  cache?: { read?: unknown };
+}
+
+interface MessagePart {
+  type?: string;
+  text?: string;
+  [key: string]: unknown;
+}
+
+interface ChatTransformMessage {
+  info: {
+    id?: string;
+    role?: string;
+    agent?: string;
+    sessionID?: string;
+  };
+  parts: MessagePart[];
+}
+
+interface LastUserMessage {
+  sessionID?: string;
+  agent?: string;
+  signature: string;
+  message: ChatTransformMessage;
+}
+
+interface EventInput {
+  type: string;
+  properties?: {
+    info?: {
+      id?: string;
+      sessionID?: string;
+      agent?: string;
+      providerID?: string;
+      modelID?: string;
+      tokens?: Tokens;
+    };
+    sessionID?: string;
+  };
+}
+
+interface ToolInput {
+  tool: string;
+  sessionID?: string;
+}
+
+interface Options {
+  getContextLimit: (
+    providerID: string,
+    modelID: string,
+  ) => Promise<number | undefined>;
+  log?: (message: string, meta?: Record<string, unknown>) => void;
+}
+
+interface UsageState {
+  providerID: string;
+  modelID: string;
+  used: number;
+}
+
+function isFinitePositiveNumber(value: unknown): value is number {
+  return typeof value === 'number' && Number.isFinite(value) && value > 0;
+}
+
+function usedContextTokens(tokens: Tokens | undefined): number | undefined {
+  if (!tokens) return undefined;
+
+  const input = tokens.input;
+  const cacheRead = tokens.cache?.read;
+  if (!isFinitePositiveNumber(input) && !isFinitePositiveNumber(cacheRead)) {
+    return undefined;
+  }
+
+  return (
+    (isFinitePositiveNumber(input) ? input : 0) +
+    (isFinitePositiveNumber(cacheRead) ? cacheRead : 0)
+  );
+}
+
+function stripInstruction(text: string): string {
+  const trimmed = text.trimEnd();
+  if (!trimmed.endsWith(GENERATED_INSTRUCTION)) {
+    return trimmed;
+  }
+
+  return trimmed.slice(0, -GENERATED_INSTRUCTION.length).trimEnd();
+}
+
+function appendInstruction(message: ChatTransformMessage): void {
+  const textPart = [...message.parts]
+    .reverse()
+    .find((part) => part.type === 'text' && typeof part.text === 'string');
+  if (!textPart) return;
+
+  const baseText = stripInstruction(textPart.text ?? '');
+  textPart.text = baseText
+    ? `${baseText}\n\n${GENERATED_INSTRUCTION}`
+    : GENERATED_INSTRUCTION;
+}
+
+function stripInstructionFromMessage(message: ChatTransformMessage): void {
+  const textPart = [...message.parts]
+    .reverse()
+    .find((part) => part.type === 'text' && typeof part.text === 'string');
+  if (!textPart) return;
+
+  textPart.text = stripInstruction(textPart.text ?? '');
+}
+
+function isExternalUserMessage(message: ChatTransformMessage): boolean {
+  return message.info.role === 'user';
+}
+
+function inferSessionID(
+  messages: ChatTransformMessage[],
+  index: number,
+): string | undefined {
+  const direct = messages[index]?.info.sessionID;
+  if (direct) return direct;
+
+  for (let i = index - 1; i >= 0; i--) {
+    const sessionID = messages[i]?.info.sessionID;
+    if (sessionID) return sessionID;
+  }
+
+  for (let i = index + 1; i < messages.length; i++) {
+    const sessionID = messages[i]?.info.sessionID;
+    if (sessionID) return sessionID;
+  }
+
+  return undefined;
+}
+
+function getLastExternalUserMessage(
+  messages: ChatTransformMessage[],
+): LastUserMessage | null {
+  for (let i = messages.length - 1; i >= 0; i--) {
+    const message = messages[i];
+    if (!isExternalUserMessage(message)) continue;
+
+    const partSignature = message.parts
+      .map((part) => {
+        if (part.type === 'text' && typeof part.text === 'string') {
+          return `${part.type}:${stripInstruction(part.text).trim()}`;
+        }
+        return part.type ?? 'unknown';
+      })
+      .join('|');
+    const ordinal = messages
+      .slice(0, i + 1)
+      .filter((item) => isExternalUserMessage(item)).length;
+
+    return {
+      sessionID: inferSessionID(messages, i),
+      agent: message.info.agent,
+      message,
+      signature: message.info.id
+        ? `${message.info.id}:${partSignature}`
+        : `${ordinal}:${partSignature}`,
+    };
+  }
+
+  return null;
+}
+
+export function createSubagentContextHygiene(options: Options) {
+  const pending = new Set<string>();
+  const sessionAgents = new Map<string, string>();
+  const usageBySession = new Map<string, UsageState>();
+  const requestSignatureBySession = new Map<string, string>();
+
+  function clear(sessionID: string): void {
+    pending.delete(sessionID);
+    usageBySession.delete(sessionID);
+    requestSignatureBySession.delete(sessionID);
+    sessionAgents.delete(sessionID);
+  }
+
+  function isSubagentSession(sessionID: string, agent?: string): boolean {
+    const resolved = agent ?? sessionAgents.get(sessionID);
+    return Boolean(resolved && resolved !== 'orchestrator');
+  }
+
+  async function isOverThreshold(sessionID: string): Promise<boolean> {
+    const usage = usageBySession.get(sessionID);
+    if (!usage) return false;
+
+    const limit = await options.getContextLimit(
+      usage.providerID,
+      usage.modelID,
+    );
+    if (!isFinitePositiveNumber(limit)) {
+      return false;
+    }
+
+    return usage.used / limit >= CONTEXT_WARNING_THRESHOLD;
+  }
+
+  return {
+    handleChatMessage(input: { sessionID: string; agent?: string }): void {
+      if (!input.agent) return;
+      sessionAgents.set(input.sessionID, input.agent);
+    },
+
+    async handleToolExecuteAfter(input: ToolInput): Promise<void> {
+      if (!input.sessionID || !isSubagentSession(input.sessionID)) {
+        return;
+      }
+
+      if (await isOverThreshold(input.sessionID)) {
+        pending.add(input.sessionID);
+        options.log?.('Armed subagent context hygiene reminder', {
+          sessionID: input.sessionID,
+          tool: input.tool,
+        });
+      }
+    },
+
+    async handleMessagesTransform(output: {
+      messages: ChatTransformMessage[];
+    }): Promise<void> {
+      const lastUserMessage = getLastExternalUserMessage(output.messages);
+      if (!lastUserMessage?.sessionID) {
+        return;
+      }
+
+      const { sessionID } = lastUserMessage;
+      const agent = lastUserMessage.agent ?? sessionAgents.get(sessionID);
+      if (!isSubagentSession(sessionID, agent)) {
+        stripInstructionFromMessage(lastUserMessage.message);
+        return;
+      }
+
+      if (
+        requestSignatureBySession.get(sessionID) !== lastUserMessage.signature
+      ) {
+        requestSignatureBySession.set(sessionID, lastUserMessage.signature);
+        stripInstructionFromMessage(lastUserMessage.message);
+        pending.delete(sessionID);
+        return;
+      }
+
+      if (pending.has(sessionID)) {
+        appendInstruction(lastUserMessage.message);
+        pending.delete(sessionID);
+      } else {
+        stripInstructionFromMessage(lastUserMessage.message);
+      }
+    },
+
+    async handleEvent(event: EventInput): Promise<void> {
+      if (event.type === 'session.deleted') {
+        const sessionID =
+          event.properties?.sessionID ?? event.properties?.info?.id;
+        if (sessionID) clear(sessionID);
+        return;
+      }
+
+      if (event.type !== 'message.updated') {
+        return;
+      }
+
+      const info = event.properties?.info;
+      const sessionID = info?.sessionID;
+      if (!sessionID) return;
+
+      if (info.agent) {
+        sessionAgents.set(sessionID, info.agent);
+      }
+
+      if (!info.providerID || !info.modelID) {
+        return;
+      }
+
+      const used = usedContextTokens(info.tokens);
+      if (used === undefined) {
+        return;
+      }
+
+      usageBySession.set(sessionID, {
+        providerID: info.providerID,
+        modelID: info.modelID,
+        used,
+      });
+    },
+  };
+}