Browse Source

Merge PR #930: preserve continuation model selection

Alvin Unreal 1 week ago
parent
commit
bbe982fecb

+ 1 - 0
src/hooks/task-session-manager/codemap.md

@@ -11,6 +11,7 @@ The directory follows a **Facade + Strategy** pattern where `index.ts` acts as t
 - **index.ts**: Main facade that wires hooks into OpenCode's lifecycle and coordinates between the job board, pending calls, task context tracking, and explicit user waits. Implements the plugin hook interface (`tool.execute.before`, `tool.execute.after`, `experimental.chat.messages.transform`, `event`) and exposes `beginUserWait()` to the `wait_for_user` tool.
 - **input-wait-tracker.ts**: Provides the single `hasInputWait()` seam used by idle reconciliation and continuation evaluation. It combines local question/permission waits with the process-global explicit user-wait latch.
 - **continuation-attempt-gate.ts**: Owns process-global continuation epochs, reservations, and explicit user waits across hook recreation. The wait is encoded as an `attempts` sentinel so pre-upgrade #856 hooks sharing the store also fail closed. Distinct external user-message identity rearms both states.
+- **continuation-model-selection.ts**: Normalizes current-session and chat-hook model shapes before forwarding runtime model and variant choices to idle continuation prompts.
 - **pending-call-tracker.ts**: Tracks in-flight task calls using a capped ordered map (`MAX_PENDING_TASK_CALLS`) to correlate launch output safely. Provides call ID generation, storage, retrieval, and cleanup for pending task invocations.
 - **task-context-tracker.ts**: Manages read context from child sessions with line-count and file caps. Stores context per task ID and provides pruning to prevent unbounded growth.
 

+ 29 - 0
src/hooks/task-session-manager/continuation-evaluator.ts

@@ -11,6 +11,10 @@ import { createInternalAgentTextPart } from '../../utils';
 import type { BackgroundJobStore } from '../../utils/background-job-store';
 import { isRecord as isObjectRecord } from '../../utils/guards';
 import { log } from '../../utils/logger';
+import {
+  type ContinuationModelSelection,
+  parseContinuationModelSelection,
+} from './continuation-model-selection';
 import { isActiveStatus } from './status-utils';
 
 const CONTINUATION_NUDGE =
@@ -95,10 +99,14 @@ export async function evaluateContinuation(
     options: {
       isFallbackInProgress?: (sessionID: string) => boolean;
     };
+    getObservedModelSelection: (
+      sessionID: string,
+    ) => ContinuationModelSelection | undefined;
     sessionSdk?: {
       todo?: (input: unknown) => Promise<{ data?: unknown }>;
       children?: (input: unknown) => Promise<{ data?: unknown }>;
       status?: (input: unknown) => Promise<{ data?: unknown }>;
+      get?: (input: unknown) => Promise<{ data?: unknown }>;
       promptAsync?: (input: unknown) => Promise<unknown>;
     };
   },
@@ -230,6 +238,25 @@ export async function evaluateContinuation(
       return;
     }
 
+    let currentModelSelection: ContinuationModelSelection | undefined;
+    if (deps.sessionSdk.get) {
+      try {
+        const sessionResponse = await deps.sessionSdk.get({
+          path: { id: parentSessionID },
+          throwOnError: true,
+        });
+        const session = isObjectRecord(sessionResponse?.data)
+          ? sessionResponse.data
+          : undefined;
+        currentModelSelection = parseContinuationModelSelection(session?.model);
+      } catch {
+        // Model enrichment is fail-soft. Older OpenCode session payloads do
+        // not expose Session.model, so fall back to the filtered chat hook.
+      }
+    }
+    const modelSelection =
+      currentModelSelection ?? deps.getObservedModelSelection(parentSessionID);
+
     if (
       isEvaluationAborted(parentSessionID, sessionToken, evaluationToken, deps)
     ) {
@@ -248,6 +275,8 @@ export async function evaluateContinuation(
       path: { id: parentSessionID },
       body: {
         agent: 'orchestrator',
+        ...(modelSelection ? { model: modelSelection.model } : {}),
+        ...(modelSelection?.variant ? { variant: modelSelection.variant } : {}),
         parts: [createInternalAgentTextPart(CONTINUATION_NUDGE)],
       },
       throwOnError: true,

+ 46 - 0
src/hooks/task-session-manager/continuation-model-selection.ts

@@ -0,0 +1,46 @@
+import { isRecord as isObjectRecord } from '../../utils/guards';
+
+export type ContinuationModelSelection = {
+  model: {
+    providerID: string;
+    modelID: string;
+  };
+  variant?: string;
+};
+
+/**
+ * Normalize the two runtime model shapes used across supported OpenCode
+ * versions:
+ * - chat.message / promptAsync: { providerID, modelID }
+ * - current Session.model:      { providerID, id }
+ */
+export function parseContinuationModelSelection(
+  value: unknown,
+  variantOverride?: unknown,
+): ContinuationModelSelection | undefined {
+  if (!isObjectRecord(value)) return undefined;
+
+  const providerID =
+    typeof value.providerID === 'string' && value.providerID.length > 0
+      ? value.providerID
+      : undefined;
+  const modelID =
+    typeof value.modelID === 'string' && value.modelID.length > 0
+      ? value.modelID
+      : typeof value.id === 'string' && value.id.length > 0
+        ? value.id
+        : undefined;
+  if (!providerID || !modelID) return undefined;
+
+  const variant =
+    typeof variantOverride === 'string' && variantOverride.length > 0
+      ? variantOverride
+      : typeof value.variant === 'string' && value.variant.length > 0
+        ? value.variant
+        : undefined;
+
+  return {
+    model: { providerID, modelID },
+    ...(variant ? { variant } : {}),
+  };
+}

+ 279 - 0
src/hooks/task-session-manager/index.test.ts

@@ -90,6 +90,55 @@ function createContinuationHook(options?: HookOptions) {
   });
 }
 
+function createContinuationSessionClient(
+  promptAsync: unknown,
+  overrides?: Record<string, unknown>,
+): Record<string, unknown> {
+  return {
+    todo: mock(async () => ({ data: [{ status: 'in_progress' }] })),
+    children: mock(async () => ({ data: [] })),
+    status: mock(async () => ({ data: {} })),
+    promptAsync,
+    ...overrides,
+  };
+}
+
+function createRuntimeUserTurn(options: {
+  sessionID?: string;
+  messageID: string;
+  providerID: string;
+  modelID: string;
+  variant?: string;
+}) {
+  const sessionID = options.sessionID ?? 'parent-1';
+  const model = {
+    providerID: options.providerID,
+    modelID: options.modelID,
+  };
+  const parts = [{ type: 'text', text: 'continue with this model' }];
+  return {
+    input: {
+      sessionID,
+      messageID: options.messageID,
+      model,
+      ...(options.variant ? { variant: options.variant } : {}),
+      parts,
+    },
+    output: {
+      message: {
+        id: options.messageID,
+        sessionID,
+        role: 'user' as const,
+        model: {
+          ...model,
+          ...(options.variant ? { variant: options.variant } : {}),
+        },
+      },
+      parts,
+    },
+  };
+}
+
 function createMessages(sessionID: string, text = 'user message') {
   return {
     messages: [
@@ -4035,6 +4084,236 @@ describe('task-session-manager hook', () => {
     );
   });
 
+  test('preserves the current session model and variant on continuation nudges', async () => {
+    const promptAsync = mock(async () => ({}));
+    const get = mock(async () => ({
+      data: {
+        model: {
+          providerID: 'runtime-provider',
+          id: 'selected-model',
+          variant: 'selected-variant',
+        },
+      },
+    }));
+    const { hook } = createContinuationHook({
+      idleReconcileDelayMs: 0,
+      sessionClient: createContinuationSessionClient(promptAsync, {
+        get,
+      }),
+    });
+
+    await hook.event({
+      event: { type: 'session.idle', properties: { sessionID: 'parent-1' } },
+    });
+    await flushContinuation();
+
+    expect(get).toHaveBeenCalledWith({
+      path: { id: 'parent-1' },
+      throwOnError: true,
+    });
+    expect(promptAsync).toHaveBeenCalledTimes(1);
+    expect(promptAsync).toHaveBeenCalledWith(
+      expect.objectContaining({
+        body: expect.objectContaining({
+          model: {
+            providerID: 'runtime-provider',
+            modelID: 'selected-model',
+          },
+          variant: 'selected-variant',
+        }),
+      }),
+    );
+  });
+
+  test('falls back to the latest external user model when session lookup fails', async () => {
+    const promptAsync = mock(async () => ({}));
+    const userTurn = createRuntimeUserTurn({
+      messageID: 'user-1',
+      providerID: 'runtime-provider',
+      modelID: 'selected-model',
+      variant: 'selected-variant',
+    });
+    const { hook } = createContinuationHook({
+      idleReconcileDelayMs: 0,
+      sessionClient: createContinuationSessionClient(promptAsync, {
+        get: mock(async () => {
+          throw new Error('session lookup unavailable');
+        }),
+      }),
+    });
+
+    hook.observeChatMessage(
+      {
+        sessionID: userTurn.input.sessionID,
+        messageID: userTurn.input.messageID,
+        parts: userTurn.input.parts,
+      },
+      userTurn.output,
+    );
+    await hook.event({
+      event: { type: 'session.idle', properties: { sessionID: 'parent-1' } },
+    });
+    await flushContinuation();
+
+    expect(promptAsync).toHaveBeenCalledTimes(1);
+    expect(promptAsync).toHaveBeenCalledWith(
+      expect.objectContaining({
+        body: expect.objectContaining({
+          model: {
+            providerID: 'runtime-provider',
+            modelID: 'selected-model',
+          },
+          variant: 'selected-variant',
+        }),
+      }),
+    );
+  });
+
+  test('treats a current session model without variant as authoritative', async () => {
+    const promptAsync = mock(async (_input: unknown) => ({}));
+    const { hook } = createContinuationHook({
+      idleReconcileDelayMs: 0,
+      sessionClient: createContinuationSessionClient(promptAsync, {
+        get: mock(async () => ({
+          data: {
+            model: {
+              providerID: 'current-provider',
+              id: 'current-model',
+            },
+          },
+        })),
+      }),
+    });
+    const previousTurn = createRuntimeUserTurn({
+      messageID: 'user-1',
+      providerID: 'previous-provider',
+      modelID: 'previous-model',
+      variant: 'previous-variant',
+    });
+    hook.observeChatMessage(previousTurn.input, previousTurn.output);
+
+    await hook.event({
+      event: { type: 'session.idle', properties: { sessionID: 'parent-1' } },
+    });
+    await flushContinuation();
+
+    const request = promptAsync.mock.calls[0]?.[0] as {
+      body: Record<string, unknown>;
+    };
+    expect(request.body.model).toEqual({
+      providerID: 'current-provider',
+      modelID: 'current-model',
+    });
+    expect(request.body).not.toHaveProperty('variant');
+  });
+
+  test('only external messages replace the model fallback and clear its variant', async () => {
+    const promptAsync = mock(async (_input: unknown) => ({}));
+    const { hook } = createContinuationHook({
+      idleReconcileDelayMs: 0,
+      sessionClient: createContinuationSessionClient(promptAsync),
+    });
+    const selectedTurn = createRuntimeUserTurn({
+      messageID: 'user-1',
+      providerID: 'selected-provider',
+      modelID: 'selected-model',
+      variant: 'selected-variant',
+    });
+    hook.observeChatMessage(selectedTurn.input, selectedTurn.output);
+    const newTurn = createRuntimeUserTurn({
+      messageID: 'user-2',
+      providerID: 'new-provider',
+      modelID: 'new-model',
+    });
+    hook.observeChatMessage(newTurn.input, newTurn.output);
+    hook.observeChatMessage(
+      {
+        sessionID: 'parent-1',
+        messageID: 'synthetic-1',
+        model: { providerID: 'static-provider', modelID: 'static-model' },
+        variant: 'static-variant',
+      },
+      {
+        message: {
+          id: 'synthetic-1',
+          sessionID: 'parent-1',
+          role: 'user',
+        },
+        parts: [
+          {
+            type: 'text',
+            text: 'synthetic continuation',
+            synthetic: true,
+          },
+        ],
+      },
+    );
+
+    await hook.event({
+      event: { type: 'session.idle', properties: { sessionID: 'parent-1' } },
+    });
+    await flushContinuation();
+
+    const request = promptAsync.mock.calls[0]?.[0] as {
+      body: Record<string, unknown>;
+    };
+    expect(request.body.model).toEqual({
+      providerID: 'new-provider',
+      modelID: 'new-model',
+    });
+    expect(request.body).not.toHaveProperty('variant');
+  });
+
+  test('a user message invalidates continuation while current model lookup is pending', async () => {
+    let resolveGet!: (value: {
+      data: {
+        model: { providerID: string; id: string; variant: string };
+      };
+    }) => void;
+    const get = mock(
+      () =>
+        new Promise<{
+          data: {
+            model: { providerID: string; id: string; variant: string };
+          };
+        }>((resolve) => {
+          resolveGet = resolve;
+        }),
+    );
+    const promptAsync = mock(async () => ({}));
+    const { hook } = createContinuationHook({
+      idleReconcileDelayMs: 0,
+      sessionClient: createContinuationSessionClient(promptAsync, {
+        get,
+      }),
+    });
+
+    await hook.event({
+      event: { type: 'session.idle', properties: { sessionID: 'parent-1' } },
+    });
+    await flushContinuation();
+    expect(get).toHaveBeenCalledTimes(1);
+
+    const newTurn = createRuntimeUserTurn({
+      messageID: 'user-2',
+      providerID: 'new-provider',
+      modelID: 'new-model',
+    });
+    hook.observeChatMessage(newTurn.input, newTurn.output);
+    resolveGet({
+      data: {
+        model: {
+          providerID: 'stale-provider',
+          id: 'stale-model',
+          variant: 'stale-variant',
+        },
+      },
+    });
+    await flushContinuation();
+
+    expect(promptAsync).not.toHaveBeenCalled();
+  });
+
   test('paired idle events submit at most one continuation', async () => {
     const promptAsync = mock(async () => ({}));
     const { hook } = createContinuationHook({

+ 39 - 3
src/hooks/task-session-manager/index.ts

@@ -18,6 +18,10 @@ import {
   updateFromInjectedCompletion,
 } from './board-injection';
 import { evaluateContinuation as evaluateContinuationFn } from './continuation-evaluator';
+import {
+  type ContinuationModelSelection,
+  parseContinuationModelSelection,
+} from './continuation-model-selection';
 import { createContinuationTokenManager } from './continuation-token-manager';
 import { handleEvent } from './event-router';
 import { createIdleReconciler } from './idle-reconciliation';
@@ -85,6 +89,10 @@ export function createTaskSessionManagerHook(
   const processedInjectedCompletions = new Set<string>();
   const processedInjectedCompletionOrder: string[] = [];
   const terminalJobsInjectedByParent = new Map<string, InjectedTerminalJobs>();
+  const observedContinuationModels = new Map<
+    string,
+    ContinuationModelSelection
+  >();
 
   // Forward refs for circular deps — set after corresponding managers exist.
   // These are captured by closure in createIdleReconciler and only called
@@ -137,6 +145,7 @@ export function createTaskSessionManagerHook(
     todo?: (input: unknown) => Promise<SdkResponse>;
     children?: (input: unknown) => Promise<SdkResponse>;
     status?: (input: unknown) => Promise<SdkResponse>;
+    get?: (input: unknown) => Promise<SdkResponse>;
     promptAsync?: (input: unknown) => Promise<unknown>;
   };
   const sessionSdk = (_ctx.client as unknown as { session?: SessionSdk })
@@ -150,6 +159,8 @@ export function createTaskSessionManagerHook(
       inputWaits,
       options,
       sessionSdk,
+      getObservedModelSelection: (sessionID) =>
+        observedContinuationModels.get(sessionID),
     });
 
   if (options.coordinator) {
@@ -161,6 +172,7 @@ export function createTaskSessionManagerHook(
         continuationTokens.clearContinuation(sessionId);
       }
       inputWaits.clearInputWaits(sessionId);
+      observedContinuationModels.delete(sessionId);
       idleReconciler.clearIdleTimers(sessionId);
       // During a foreground fallback abort/re-prompt cycle, the session
       // is being torn down and immediately recreated with a fallback model.
@@ -242,6 +254,21 @@ export function createTaskSessionManagerHook(
       ) {
         return;
       }
+      const outputModel = isObjectRecord(outputMessage?.model)
+        ? outputMessage.model
+        : undefined;
+      const variant =
+        typeof inputMessage?.variant === 'string'
+          ? inputMessage.variant
+          : outputModel?.variant;
+      const modelSelection =
+        parseContinuationModelSelection(inputMessage?.model, variant) ??
+        parseContinuationModelSelection(outputModel, variant);
+      if (modelSelection) {
+        observedContinuationModels.set(sessionID, modelSelection);
+      } else {
+        observedContinuationModels.delete(sessionID);
+      }
       continuationTokens.rearmForUserMessage(sessionID, messageIdentity);
     },
 
@@ -325,8 +352,16 @@ export function createTaskSessionManagerHook(
           error?: { name?: string };
         };
       };
-    }): Promise<void> =>
-      handleEvent(input, {
+    }): Promise<void> => {
+      if (input.event.type === 'server.instance.disposed') {
+        observedContinuationModels.clear();
+      } else if (input.event.type === 'session.deleted') {
+        const sessionID =
+          input.event.properties?.info?.id ?? input.event.properties?.sessionID;
+        if (sessionID) observedContinuationModels.delete(sessionID);
+      }
+
+      return handleEvent(input, {
         inputWaits,
         continuationTokens,
         options,
@@ -336,6 +371,7 @@ export function createTaskSessionManagerHook(
         taskContextTracker,
         terminalJobsInjectedByParent,
         retainedBoardSnapshots: injectionState.retainedBoardSnapshots,
-      }),
+      });
+    },
   };
 }

+ 10 - 0
src/index.ts

@@ -1137,6 +1137,11 @@ const OhMyOpenCodeLite: Plugin = async (ctx) => {
       input: {
         sessionID: string;
         agent?: string;
+        model?: {
+          providerID: string;
+          modelID: string;
+        };
+        variant?: string;
         parts?: unknown[];
         /** OpenCode chat.message message identity when present. */
         messageID?: string;
@@ -1147,6 +1152,11 @@ const OhMyOpenCodeLite: Plugin = async (ctx) => {
           agent?: string;
           role?: string;
           sessionID?: string;
+          model?: {
+            providerID: string;
+            modelID: string;
+            variant?: string;
+          };
         };
         parts?: unknown[];
       },