Просмотр исходного кода

fix(tokens): dedupe orchestrator system injection, scope to main-chat requests, stop duplicate revive notifications

- system.transform dedups by the effective orchestrator prompt instead of
  default-prompt markers, so a custom replacement without them is no
  longer appended twice (P + host + P).
- Auxiliary LLM requests (title generation, compaction) running in an
  orchestrator session no longer receive orchestration instructions:
  the v2 context bridge forwards event.agent (request-scoped), and v1
  hosts gate on the core's environment block, which auxiliaries never
  carry.
- A late-transport success after the local 10s timeout now marks the
  revived-run notification sent (classified via responseError, so an
  {error} envelope is not delivery) and cancels the pending retry; a
  retry parked on selection resolution revalidates sent before
  acquiring the lease. The same terminal result is never delivered to
  the parent twice.
dhaern 1 день назад
Родитель
Сommit
3bd5be801d

+ 217 - 0
src/hooks/task-session-manager/revived-run-tracker.test.ts

@@ -674,4 +674,221 @@ describe('revived run tracker', () => {
     harness.tracker.onTerminal(cancelled);
     expect(harness.tracker.isTracked(next.taskID, next.generation)).toBe(true);
   });
+
+  // Controlled clock for the transport-timeout scenarios below: capture
+  // timer registrations so the 10s transport timeout can be fired without
+  // waiting, and track clearTimeout so a cancelled retry is provable.
+  function installCapturedTimers() {
+    const timers = new Map<number, { delay: number; callback: () => void }>();
+    const cleared = new Set<number>();
+    let nextId = 0;
+    globalThis.setTimeout = ((callback: () => void, delay = 0) => {
+      const id = ++nextId;
+      timers.set(id, { delay, callback });
+      return id;
+    }) as typeof setTimeout;
+    globalThis.clearTimeout = ((id: number) => {
+      cleared.add(id);
+      timers.delete(id);
+    }) as typeof clearTimeout;
+    const settle = async () => {
+      for (let i = 0; i < 15; i += 1) await Promise.resolve();
+    };
+    const fire = (delay: number) => {
+      for (const [id, timer] of [...timers.entries()]) {
+        if (timer.delay !== delay) continue;
+        timers.delete(id);
+        timer.callback();
+        return id;
+      }
+      return undefined;
+    };
+    const soleSurviving = (delay: number) =>
+      [...timers.values()].find((timer) => timer.delay === delay);
+    return { timers, cleared, settle, fire, soleSurviving };
+  }
+
+  test('late transport success after timeout marks sent and cancels the retry', async () => {
+    const clock = installCapturedTimers();
+    let resolvePrompt: ((value: unknown) => void) | undefined;
+    const prompt = mock(
+      () =>
+        new Promise((resolve) => {
+          resolvePrompt = resolve;
+        }),
+    );
+    const harness = createHarness(() => ({ data: [] }), prompt);
+    harness.tracker.register({
+      taskID: harness.run.taskID,
+      generation: harness.run.generation,
+      parentSessionID: 'parent',
+      description: 'inspect the change',
+    });
+    const terminal = harness.board.updateStatus({
+      taskID: harness.run.taskID,
+      expectedGeneration: harness.run.generation,
+      state: 'completed',
+      resultSummary: 'done',
+    });
+    if (!terminal) throw new Error('missing terminal record');
+    harness.tracker.onTerminal(terminal);
+    await clock.settle();
+    expect(harness.prompt).toHaveBeenCalledTimes(1);
+
+    // Local 10s transport timeout fires while promptAsync is still pending.
+    clock.fire(10_000);
+    await clock.settle();
+    const retryTimer = clock.soleSurviving(0);
+    expect(retryTimer).toBeDefined();
+
+    // The original transport settles successfully AFTER the timeout.
+    resolvePrompt?.({});
+    await clock.settle();
+
+    expect(harness.prompt).toHaveBeenCalledTimes(1);
+    expect([...clock.timers.values()]).not.toContain(retryTimer);
+  });
+
+  test('late transport failure after timeout keeps the retry path', async () => {
+    const clock = installCapturedTimers();
+    let rejectPrompt: ((reason: unknown) => void) | undefined;
+    const prompt = mock(
+      () =>
+        new Promise((_resolve, reject) => {
+          rejectPrompt = reject;
+        }),
+    );
+    const harness = createHarness(() => ({ data: [] }), prompt);
+    harness.tracker.register({
+      taskID: harness.run.taskID,
+      generation: harness.run.generation,
+      parentSessionID: 'parent',
+      description: 'inspect the change',
+    });
+    const terminal = harness.board.updateStatus({
+      taskID: harness.run.taskID,
+      expectedGeneration: harness.run.generation,
+      state: 'completed',
+      resultSummary: 'done',
+    });
+    if (!terminal) throw new Error('missing terminal record');
+    harness.tracker.onTerminal(terminal);
+    await clock.settle();
+
+    clock.fire(10_000);
+    await clock.settle();
+    expect(clock.soleSurviving(0)).toBeDefined();
+
+    // The original transport fails after the timeout: the retry must stay
+    // armed and deliver the notification on the next attempt.
+    rejectPrompt?.(new Error('host unavailable'));
+    await clock.settle();
+    expect(clock.soleSurviving(0)).toBeDefined();
+
+    clock.fire(0);
+    await clock.settle();
+    expect(harness.prompt).toHaveBeenCalledTimes(2);
+  });
+
+  test('late transport error envelope after timeout is not delivery', async () => {
+    const clock = installCapturedTimers();
+    let resolvePrompt: ((value: unknown) => void) | undefined;
+    const prompt = mock(
+      () =>
+        new Promise((resolve) => {
+          resolvePrompt = resolve;
+        }),
+    );
+    const harness = createHarness(() => ({ data: [] }), prompt);
+    harness.tracker.register({
+      taskID: harness.run.taskID,
+      generation: harness.run.generation,
+      parentSessionID: 'parent',
+      description: 'inspect the change',
+    });
+    const terminal = harness.board.updateStatus({
+      taskID: harness.run.taskID,
+      expectedGeneration: harness.run.generation,
+      state: 'completed',
+      resultSummary: 'done',
+    });
+    if (!terminal) throw new Error('missing terminal record');
+    harness.tracker.onTerminal(terminal);
+    await clock.settle();
+    expect(harness.prompt).toHaveBeenCalledTimes(1);
+
+    clock.fire(10_000);
+    await clock.settle();
+
+    // The transport PROMISE resolves, but with a host error envelope —
+    // a resolved SDK call without throwOnError is not a delivered
+    // notification. The retry must stay armed.
+    resolvePrompt?.({ error: { message: 'host rejected' } });
+    await clock.settle();
+    expect(clock.soleSurviving(0)).toBeDefined();
+
+    clock.fire(0);
+    await clock.settle();
+    expect(harness.prompt).toHaveBeenCalledTimes(2);
+  });
+
+  test('late success while a retry waits on selection prevents a second send', async () => {
+    const clock = installCapturedTimers();
+    let resolvePrompt: ((value: unknown) => void) | undefined;
+    const prompt = mock(
+      () =>
+        new Promise((resolve) => {
+          resolvePrompt = resolve;
+        }),
+    );
+    // First selection resolves immediately (attempt 1 sends); the second
+    // call (retry) blocks until released, modeling a slow host read.
+    let selectionCalls = 0;
+    let releaseSecondSelection: (() => void) | undefined;
+    const secondGate = new Promise<void>((resolve) => {
+      releaseSecondSelection = resolve;
+    });
+    const harness = createHarness(() => ({ data: [] }), prompt, false, {
+      resolveSelection: async () => {
+        selectionCalls += 1;
+        if (selectionCalls >= 2) await secondGate;
+        return { agent: 'plan', provenance: 'host-persisted' };
+      },
+    });
+    harness.tracker.register({
+      taskID: harness.run.taskID,
+      generation: harness.run.generation,
+      parentSessionID: 'parent',
+      description: 'inspect the change',
+    });
+    const terminal = harness.board.updateStatus({
+      taskID: harness.run.taskID,
+      expectedGeneration: harness.run.generation,
+      state: 'completed',
+      resultSummary: 'done',
+    });
+    if (!terminal) throw new Error('missing terminal record');
+    harness.tracker.onTerminal(terminal);
+    await clock.settle();
+    expect(harness.prompt).toHaveBeenCalledTimes(1);
+
+    // Local timeout: retry armed and fired; the retry passes its entry
+    // guard (sent is still false) and parks on the selection await.
+    clock.fire(10_000);
+    await clock.settle();
+    clock.fire(0);
+    await clock.settle();
+    expect(selectionCalls).toBeGreaterThanOrEqual(2);
+
+    // The ORIGINAL transport settles successfully after everything: the
+    // notification is delivered, sent is marked, and the parked retry
+    // must not acquire the lease or send again.
+    resolvePrompt?.({});
+    await clock.settle();
+    releaseSecondSelection?.();
+    await clock.settle();
+    await clock.settle();
+
+    expect(harness.prompt).toHaveBeenCalledTimes(1);
+  });
 });

+ 34 - 3
src/hooks/task-session-manager/revived-run-tracker.ts

@@ -291,7 +291,13 @@ export function createRevivedRunTracker(options: {
             .resolveSelection(run.parentSessionID)
             .catch((): undefined => undefined)
         : undefined;
-      if (disposed || runs.get(run.taskID) !== run) return;
+      if (disposed || runs.get(run.taskID) !== run || run.notification.sent) {
+        return;
+      }
+      // Revalidate AFTER the selection await: a late success from a
+      // previous attempt may have marked this notification sent while the
+      // retry was pending here — sending again would duplicate the
+      // terminal result (Oracle r2 P1.2).
       const latestBeforeSend = options.backgroundJobBoard.get(run.taskID);
       if (
         !latestBeforeSend ||
@@ -351,6 +357,19 @@ export function createRevivedRunTracker(options: {
               parts: [createInternalAgentTextPart(text)],
             },
           }),
+        // Late settlement after the local timeout: a SUCCESS means the
+        // host DID accept the notification — mark it delivered and cancel
+        // the pending retry so the same terminal result is never sent to
+        // the parent twice. A late FAILURE keeps the retry scheduled.
+        (outcome) => {
+          if (!outcome.ok) return;
+          if (disposed || runs.get(run.taskID) !== run) return;
+          run.notification.sent = true;
+          if (run.notification.retryTimer) {
+            clearTimeout(run.notification.retryTimer);
+            run.notification.retryTimer = undefined;
+          }
+        },
       );
       const error = responseError(response);
       if (error !== undefined) throw new Error(stringifyError(error));
@@ -440,6 +459,7 @@ async function awaitNotificationTransport<T>(
   backgroundJobBoard: BackgroundJobStore,
   lease: BackgroundJobLease,
   operation: () => Promise<T>,
+  onLateSettlement?: (outcome: { ok: boolean }) => void,
 ): Promise<T> {
   let settled = false;
   let timedOut = false;
@@ -449,12 +469,23 @@ async function awaitNotificationTransport<T>(
     .then(
       (value) => {
         settled = true;
-        if (timedOut) backgroundJobBoard.releaseLease(lease);
+        if (timedOut) {
+          backgroundJobBoard.releaseLease(lease);
+          // A resolved promise is NOT delivery: the SDK can resolve with
+          // an `{ error }` envelope when throwOnError is off. Classify
+          // with the same check the normal path uses (Oracle r2 P1.1).
+          onLateSettlement?.({
+            ok: responseError(value) === undefined,
+          });
+        }
         return value;
       },
       (error: unknown) => {
         settled = true;
-        if (timedOut) backgroundJobBoard.releaseLease(lease);
+        if (timedOut) {
+          backgroundJobBoard.releaseLease(lease);
+          onLateSettlement?.({ ok: false });
+        }
         throw error;
       },
     );

+ 134 - 0
src/index.test.ts

@@ -1289,6 +1289,140 @@ describe('plugin config model inheritance', () => {
   });
 });
 
+describe('system.transform orchestrator injection', () => {
+  let originalEnv: typeof process.env;
+  const configDirs: string[] = [];
+
+  beforeEach(() => {
+    originalEnv = { ...process.env };
+  });
+
+  afterEach(async () => {
+    process.env = originalEnv;
+    while (configDirs.length > 0) {
+      const configDir = configDirs.pop();
+      if (configDir) {
+        await rm(configDir, { recursive: true, force: true });
+      }
+    }
+  });
+
+  async function loadPluginWithOrchestratorSession(
+    config: Record<string, unknown> = {},
+  ) {
+    const configDir = await mkdtemp('/tmp/oh-my-system-transform-');
+    configDirs.push(configDir);
+    await Bun.write(
+      `${configDir}/oh-my-opencode-slim.json`,
+      JSON.stringify(config),
+    );
+    process.env = {
+      ...originalEnv,
+      OPENCODE_CONFIG_DIR: configDir,
+      XDG_DATA_HOME: `${configDir}/data`,
+      XDG_CACHE_HOME: `${configDir}/cache`,
+      OPENCODE_LOG_DIR: `${configDir}/logs`,
+    };
+    const client = createPluginClient(async () => ({}));
+    const hooks = await plugin({
+      client,
+      directory: configDir,
+      worktree: configDir,
+      serverUrl: new URL('http://127.0.0.1:4096'),
+    } as never);
+    // Session tracked as orchestrator (how chat.message records it).
+    await hooks['chat.message']?.(
+      {
+        sessionID: 'ses-orc',
+        agent: 'orchestrator',
+        model: { providerID: 'test', modelID: 'm' },
+      } as never,
+      {} as never,
+    );
+    return hooks;
+  }
+
+  const ENV_BLOCK = [
+    'You are powered by the model named test/m.',
+    '<env>',
+    '  Working directory: /tmp',
+    '</env>',
+  ].join('\n');
+
+  test('does not duplicate a custom orchestrator prompt already present', async () => {
+    // Configure a REAL custom replacement without default-prompt markers:
+    // the effective prompt is this string, and the dedup must key on it.
+    const customPrompt = 'Mi prompt custom sin marcadores.';
+    const hooks = await loadPluginWithOrchestratorSession({
+      agents: { orchestrator: { prompt: customPrompt } },
+    });
+    try {
+      const system = [`${ENV_BLOCK}\n\n${customPrompt}`];
+      await hooks['experimental.chat.system.transform']?.(
+        { sessionID: 'ses-orc' } as never,
+        { system } as never,
+      );
+      // Exactly one copy of the effective prompt (split = parts + 1) and
+      // no default-prompt content appended after it.
+      expect(system[0]?.split(customPrompt).length).toBe(2);
+      expect(system[0]).toBe(`${ENV_BLOCK}\n\n${customPrompt}`);
+    } finally {
+      await hooks.dispose?.();
+    }
+  });
+
+  test('skips auxiliary requests (title/compaction) in an orchestrator session', async () => {
+    const hooks = await loadPluginWithOrchestratorSession();
+    try {
+      // Title/compaction requests carry their own short system and no
+      // environment block.
+      const system = [
+        'You are a title generator. You output ONLY a thread title.',
+      ];
+      await hooks['experimental.chat.system.transform']?.(
+        { sessionID: 'ses-orc' } as never,
+        { system } as never,
+      );
+      expect(system[0]).not.toContain('<Role>');
+      expect(system[0]).toBe(
+        'You are a title generator. You output ONLY a thread title.',
+      );
+    } finally {
+      await hooks.dispose?.();
+    }
+  });
+
+  test('injects on a main chat request in an orchestrator session', async () => {
+    const hooks = await loadPluginWithOrchestratorSession();
+    try {
+      const system = [ENV_BLOCK];
+      await hooks['experimental.chat.system.transform']?.(
+        { sessionID: 'ses-orc' } as never,
+        { system } as never,
+      );
+      expect(system[0]).toContain('<Role>');
+    } finally {
+      await hooks.dispose?.();
+    }
+  });
+
+  test('request-scoped agent overrides session tracking', async () => {
+    const hooks = await loadPluginWithOrchestratorSession();
+    try {
+      // v2 bridge forwards the request agent: an auxiliary request says
+      // its real agent even though the session is tracked as orchestrator.
+      const system = [ENV_BLOCK];
+      await hooks['experimental.chat.system.transform']?.(
+        { sessionID: 'ses-orc', agent: 'title' } as never,
+        { system } as never,
+      );
+      expect(system[0]).not.toContain('<Role>');
+    } finally {
+      await hooks.dispose?.();
+    }
+  });
+});
+
 describe('multiplexer host gating', () => {
   let originalEnv: typeof process.env;
 

+ 43 - 23
src/index.ts

@@ -102,7 +102,10 @@ import {
   createSessionSelectionReader,
   resolveCurrentSelection,
 } from './utils/session-selection';
-import { collapseSystemInPlace } from './utils/system-collapse';
+import {
+  collapseSystemInPlace,
+  looksLikeMainChatRequest,
+} from './utils/system-collapse';
 import { createV2Setup } from './v2';
 import {
   isInternalAdmission,
@@ -1676,39 +1679,56 @@ export const OhMyOpenCodeLite: Plugin = async (ctx) => {
     // agentDefs (which has custom replacement or append prompts applied)
     // instead of rebuilding the default.
     'experimental.chat.system.transform': async (
-      input: { sessionID?: string },
+      input: { sessionID?: string; agent?: unknown },
       output: { system: string[] },
     ): Promise<void> => {
-      const agentName = input.sessionID
+      // Request-scoped agent when the host provides one (the v2 context
+      // bridge forwards `event.agent`). v1 hosts only pass sessionID, so
+      // there we fall back to the session's tracked agent — which is the
+      // SESSION agent, not the request agent: auxiliary LLM requests
+      // (title generation, compaction) run in the same session under
+      // their own agent and must not receive orchestrator instructions.
+      const requestAgent =
+        typeof input.agent === 'string' && input.agent
+          ? input.agent
+          : undefined;
+      const sessionAgent = input.sessionID
         ? sessionMetadata.getAgent(input.sessionID)
         : undefined;
-      if (agentName === 'orchestrator') {
-        const alreadyInjected = output.system.some(
-          (s) =>
-            typeof s === 'string' &&
-            s.includes('<Role>') &&
-            s.includes('orchestrator'),
+      const isOrchestratorRequest =
+        requestAgent !== undefined
+          ? requestAgent === 'orchestrator'
+          : sessionAgent === 'orchestrator' &&
+            looksLikeMainChatRequest(output.system);
+      if (isOrchestratorRequest) {
+        const orchestratorDef = agentDefs.find(
+          (a) => a.name === 'orchestrator',
         );
-        if (!alreadyInjected) {
+        const orchestratorPrompt =
+          typeof orchestratorDef?.config?.prompt === 'string'
+            ? orchestratorDef.config.prompt
+            : buildOrchestratorPrompt(
+                runtime.disabledAgents,
+                undefined,
+                true,
+                true,
+                hostFlavor,
+              );
+        // Dedup by the EFFECTIVE prompt, not by default-prompt markers:
+        // a custom replacement without `<Role>` previously slipped past
+        // the marker check and was appended twice (P + host + P).
+        const alreadyInjected =
+          !!orchestratorPrompt &&
+          output.system.some(
+            (s) => typeof s === 'string' && s.includes(orchestratorPrompt),
+          );
+        if (!alreadyInjected && orchestratorPrompt) {
           // Place the orchestrator prompt after AGENTS.md so the user's
           // behavioral rules (language, code conventions, etc.) retain
           // their intended priority. AGENTS.md is injected by OpenCode
           // core into system[0]; prepending the orchestrator prompt before
           // it buries user-defined rules under thousands of lines of
           // orchestration instructions.
-          const orchestratorDef = agentDefs.find(
-            (a) => a.name === 'orchestrator',
-          );
-          const orchestratorPrompt =
-            typeof orchestratorDef?.config?.prompt === 'string'
-              ? orchestratorDef.config.prompt
-              : buildOrchestratorPrompt(
-                  runtime.disabledAgents,
-                  undefined,
-                  true,
-                  true,
-                  hostFlavor,
-                );
           output.system[0] = `${output.system[0] || ''}\n\n${orchestratorPrompt}`;
         }
       }

+ 22 - 0
src/utils/system-collapse.ts

@@ -22,3 +22,25 @@ export function collapseSystemInPlace(system: string[]): void {
     system.push(joined);
   }
 }
+
+/**
+ * Heuristic for v1 hosts, where the system transform only receives the
+ * sessionID: the session's tracked agent says "orchestrator" but auxiliary
+ * LLM requests (title generation, compaction) run in the SAME session
+ * under their own agent. Those requests are built with `system: []` and
+ * never include the core's environment block, while every main chat
+ * request does (the core prepends it to the request system). Structured
+ * signal, not message text. Degradation if upstream renames both
+ * markers: main-chat requests stop matching and the fallback stops
+ * injecting there (missing serve-mode prompt) — auxiliaries were never
+ * injected, so they cannot regress. A custom auxiliary prompt
+ * CONTAINING these strings would false-positive; accepted as bounded
+ * compatibility with this core.
+ */
+export function looksLikeMainChatRequest(system: string[]): boolean {
+  return system.some(
+    (entry) =>
+      typeof entry === 'string' &&
+      (entry.includes('<env>') || entry.includes('You are powered by')),
+  );
+}

+ 5 - 1
src/v2/setup.ts

@@ -322,7 +322,11 @@ export function createSessionContextHandler(
       try {
         const sysStrings = event.system.map((s) => s.text ?? '');
         await deps.systemTransform(
-          { sessionID: event.sessionID },
+          // Forward the request-scoped agent so the transform can tell a
+          // real orchestrator request from an auxiliary (title/compaction)
+          // request running in the same session — v1 hosts lack this and
+          // fall back to a structural heuristic.
+          { sessionID: event.sessionID, agent: event.agent },
           { system: sysStrings },
         );
         event.system = sysStrings.map((text) => ({