Merge pull request #1085 from alvinunreal/omos/issue-1080-model-identity

src/tools/task-message.test.ts

@@ -23,6 +23,15 @@ function makePrompt(): ReturnType<typeof mock> {
   return mock(async () => ({}));
 }
 
+function makeSession(prompt: ReturnType<typeof mock>) {
+  return {
+    get: mock(async () => ({
+      data: { model: { providerID: 'openai', id: 'gpt-5.6' } },
+    })),
+    prompt,
+  };
+}
+
 function createTool(board: BackgroundJobBoard) {
   return createTaskMessageTool({
     input: { directory: '/test' } as any,
@@ -43,7 +52,7 @@ describe('task_message', () => {
     const board = new BackgroundJobBoard();
     registerRunningChild(board);
     const prompt = makePrompt();
-    client = { session: { prompt } };
+    client = { session: makeSession(prompt) };
 
     await expect(
       createTool(board).execute(
@@ -56,6 +65,8 @@ describe('task_message', () => {
       path: { id: 'ses_child1' },
       body: {
         agent: 'fixer',
+        model: { providerID: 'openai', modelID: 'gpt-5.6' },
+        variant: 'default',
         noReply: true,
         parts: [{ type: 'text', text: 'Please continue.' }],
       },
@@ -67,7 +78,7 @@ describe('task_message', () => {
     const board = new BackgroundJobBoard();
     registerRunningChild(board);
     const prompt = makePrompt();
-    client = { session: { prompt } };
+    client = { session: makeSession(prompt) };
     const task_message = createTool(board);
 
     await task_message.execute({ task_id: 'ses_child1', message: 'First' }, {
@@ -83,6 +94,316 @@ describe('task_message', () => {
     expect(prompt.mock.calls[1]?.[0].body.noReply).toBe(true);
   });
 
+  test('transports the authoritative current model and variant', async () => {
+    const board = new BackgroundJobBoard();
+    registerRunningChild(board);
+    const transportOrder: string[] = [];
+    const prompt = mock(async () => {
+      transportOrder.push('prompt');
+      return {};
+    });
+    const get = mock(async () => {
+      transportOrder.push('get');
+      return {
+        data: {
+          model: {
+            providerID: 'openai',
+            id: 'gpt-5.6',
+            variant: 'high',
+          },
+        },
+      };
+    });
+    client = { session: { get, prompt } };
+
+    await createTool(board).execute(
+      { task_id: 'ses_child1', message: 'Continue with the fix.' },
+      { sessionID: 'parent-1' } as any,
+    );
+
+    expect(get).toHaveBeenCalledWith({
+      path: { id: 'ses_child1' },
+      query: { directory: '/test' },
+      signal: expect.any(AbortSignal),
+    });
+    expect(transportOrder).toEqual(['get', 'prompt']);
+    expect(prompt).toHaveBeenCalledWith({
+      path: { id: 'ses_child1' },
+      body: {
+        agent: 'fixer',
+        model: { providerID: 'openai', modelID: 'gpt-5.6' },
+        variant: 'high',
+        noReply: true,
+        parts: [{ type: 'text', text: 'Continue with the fix.' }],
+      },
+      throwOnError: true,
+    });
+  });
+
+  test('transports a separate current session variant', async () => {
+    const board = new BackgroundJobBoard();
+    registerRunningChild(board);
+    const prompt = makePrompt();
+    const get = mock(async () => ({
+      data: {
+        model: { providerID: 'openai', id: 'gpt-5.6' },
+        variant: 'medium',
+      },
+    }));
+    client = { session: { get, prompt } };
+
+    await createTool(board).execute(
+      { task_id: 'ses_child1', message: 'Continue with the fix.' },
+      { sessionID: 'parent-1' } as any,
+    );
+
+    expect(prompt.mock.calls[0]?.[0].body).toEqual({
+      agent: 'fixer',
+      model: { providerID: 'openai', modelID: 'gpt-5.6' },
+      variant: 'medium',
+      noReply: true,
+      parts: [{ type: 'text', text: 'Continue with the fix.' }],
+    });
+  });
+
+  test('transports a valid current model without a variant', async () => {
+    const board = new BackgroundJobBoard();
+    registerRunningChild(board);
+    const prompt = makePrompt();
+    client = {
+      session: {
+        get: mock(async () => ({
+          data: { model: { providerID: 'openai', id: 'gpt-5.6' } },
+        })),
+        prompt,
+      },
+    };
+
+    await createTool(board).execute(
+      { task_id: 'ses_child1', message: 'Continue with the fix.' },
+      { sessionID: 'parent-1' } as any,
+    );
+
+    expect(prompt.mock.calls[0]?.[0].body).toEqual({
+      agent: 'fixer',
+      model: { providerID: 'openai', modelID: 'gpt-5.6' },
+      variant: 'default',
+      noReply: true,
+      parts: [{ type: 'text', text: 'Continue with the fix.' }],
+    });
+  });
+
+  test('falls back to the latest user message when get is malformed', async () => {
+    const board = new BackgroundJobBoard();
+    registerRunningChild(board);
+    const prompt = makePrompt();
+    client = {
+      session: {
+        get: mock(async () => ({ data: { model: { providerID: 'openai' } } })),
+        messages: mock(async () => ({
+          data: [
+            {
+              info: {
+                role: 'user',
+                model: {
+                  providerID: 'anthropic',
+                  modelID: 'claude-sonnet',
+                  variant: 'high',
+                },
+              },
+            },
+            { info: { role: 'assistant' } },
+          ],
+        })),
+        prompt,
+      },
+    };
+
+    await createTool(board).execute(
+      { task_id: 'ses_child1', message: 'Continue with the fix.' },
+      { sessionID: 'parent-1' } as any,
+    );
+
+    expect(prompt.mock.calls[0]?.[0].body).toEqual({
+      agent: 'fixer',
+      model: { providerID: 'anthropic', modelID: 'claude-sonnet' },
+      variant: 'high',
+      noReply: true,
+      parts: [{ type: 'text', text: 'Continue with the fix.' }],
+    });
+  });
+
+  test('falls back to the latest user message when get throws', async () => {
+    const board = new BackgroundJobBoard();
+    registerRunningChild(board);
+    const prompt = makePrompt();
+    const messages = mock(async () => ({
+      data: [
+        {
+          info: {
+            role: 'user',
+            model: { providerID: 'anthropic', id: 'claude-sonnet' },
+            variant: 'high',
+          },
+        },
+      ],
+    }));
+    client = {
+      session: {
+        get: mock(async () => {
+          throw new Error('session unavailable');
+        }),
+        messages,
+        prompt,
+      },
+    };
+
+    await createTool(board).execute(
+      { task_id: 'ses_child1', message: 'Continue with the fix.' },
+      { sessionID: 'parent-1' } as any,
+    );
+
+    expect(messages).toHaveBeenCalledWith({
+      path: { id: 'ses_child1' },
+      query: { directory: '/test', limit: 20 },
+      signal: expect.any(AbortSignal),
+    });
+    expect(prompt.mock.calls[0]?.[0].body).toEqual({
+      agent: 'fixer',
+      model: { providerID: 'anthropic', modelID: 'claude-sonnet' },
+      variant: 'high',
+      noReply: true,
+      parts: [{ type: 'text', text: 'Continue with the fix.' }],
+    });
+  });
+
+  test('falls back to the latest user message when get is unavailable', async () => {
+    const board = new BackgroundJobBoard();
+    registerRunningChild(board);
+    const prompt = makePrompt();
+    client = {
+      session: {
+        messages: mock(async () => ({
+          data: [
+            {
+              info: {
+                role: 'user',
+                model: { providerID: 'google', id: 'gemini-pro' },
+              },
+            },
+          ],
+        })),
+        prompt,
+      },
+    };
+
+    await createTool(board).execute(
+      { task_id: 'ses_child1', message: 'Continue with the fix.' },
+      { sessionID: 'parent-1' } as any,
+    );
+
+    expect(prompt.mock.calls[0]?.[0].body).toEqual({
+      agent: 'fixer',
+      model: { providerID: 'google', modelID: 'gemini-pro' },
+      variant: 'default',
+      noReply: true,
+      parts: [{ type: 'text', text: 'Continue with the fix.' }],
+    });
+  });
+
+  test('rejects without prompting when both identity sources are unavailable', async () => {
+    const board = new BackgroundJobBoard();
+    registerRunningChild(board);
+    const prompt = makePrompt();
+    client = {
+      session: {
+        get: mock(async () => ({
+          data: { model: { id: 'missing-provider' } },
+        })),
+        messages: mock(async () => ({ data: [] })),
+        prompt,
+      },
+    };
+
+    await expect(
+      createTool(board).execute(
+        { task_id: 'ses_child1', message: 'Continue with the fix.' },
+        { sessionID: 'parent-1' } as any,
+      ),
+    ).rejects.toThrow('no authoritative model identity');
+    expect(prompt).not.toHaveBeenCalled();
+
+    const job = board.get('ses_child1');
+    expect(job).toBeDefined();
+    if (!job) throw new Error('missing running job');
+    expect(
+      board.acquireCancellationLease(job.taskID, job.generation),
+    ).toBeDefined();
+  });
+
+  test('bounds a hanging identity lookup and releases the lease', async () => {
+    const board = new BackgroundJobBoard();
+    registerRunningChild(board);
+    const prompt = makePrompt();
+    let lookupSignal: AbortSignal | undefined;
+    const get = mock((input: { signal?: AbortSignal }) => {
+      lookupSignal = input.signal;
+      return new Promise<unknown>(() => {});
+    });
+    client = { session: { get, prompt } };
+
+    await expect(
+      createToolWithTimeout(board, 5).execute(
+        { task_id: 'ses_child1', message: 'Continue with the fix.' },
+        { sessionID: 'parent-1' } as any,
+      ),
+    ).rejects.toThrow('model lookup timed out');
+    expect(lookupSignal).toBeDefined();
+    expect(lookupSignal?.aborted).toBe(true);
+    expect(prompt).not.toHaveBeenCalled();
+
+    const job = board.get('ses_child1');
+    expect(job).toBeDefined();
+    if (!job) throw new Error('missing running job');
+    expect(
+      board.acquireCancellationLease(job.taskID, job.generation),
+    ).toBeDefined();
+  });
+
+  test('rechecks the job after lookup before prompting', async () => {
+    const board = new BackgroundJobBoard();
+    registerRunningChild(board);
+    const prompt = makePrompt();
+    const get = mock(async () => {
+      board.updateStatus({ taskID: 'ses_child1', state: 'completed' });
+      return {
+        data: {
+          model: { providerID: 'openai', id: 'gpt-5.6' },
+          variant: 'high',
+        },
+      };
+    });
+    client = { session: { get, prompt } };
+
+    await expect(
+      createTool(board).execute(
+        { task_id: 'ses_child1', message: 'Do not send.' },
+        { sessionID: 'parent-1' } as any,
+      ),
+    ).rejects.toThrow('not running');
+    expect(prompt).not.toHaveBeenCalled();
+
+    const job = board.get('ses_child1');
+    expect(job).toBeDefined();
+    if (!job) throw new Error('missing completed job');
+    const terminalLease = board.acquireTerminalNotificationLease(
+      job.taskID,
+      job.generation,
+    );
+    expect(terminalLease).toBeDefined();
+    if (terminalLease) board.releaseLease(terminalLease);
+  });
+
   test('serializes message transport against cancellation and relaunch', async () => {
     const board = new BackgroundJobBoard();
     registerRunningChild(board);
@@ -93,7 +414,7 @@ describe('task_message', () => {
           releasePrompt = () => resolve({});
         }),
     );
-    client = { session: { prompt } };
+    client = { session: makeSession(prompt) };
 
     const pending = createTool(board).execute(
       { task_id: 'ses_child1', message: 'Hold the lane.' },
@@ -129,7 +450,7 @@ describe('task_message', () => {
     const board = new BackgroundJobBoard();
     registerRunningChild(board);
     const prompt = mock(async () => ({ error: { message: 'HTTP 409' } }));
-    client = { session: { prompt } };
+    client = { session: makeSession(prompt) };
 
     await expect(
       createTool(board).execute(
@@ -156,7 +477,7 @@ describe('task_message', () => {
           settlePrompt = () => resolve({});
         }),
     );
-    client = { session: { prompt } };
+    client = { session: makeSession(prompt) };
 
     await expect(
       createToolWithTimeout(board, 5).execute(
@@ -184,6 +505,9 @@ describe('task_message', () => {
     registerRunningChild(board);
     const prompt = makePrompt();
     const session = {
+      get: async () => ({
+        data: { model: { providerID: 'openai', id: 'gpt-5.6' } },
+      }),
       get prompt() {
         board.drop('ses_child1');
         return prompt;
@@ -274,7 +598,7 @@ describe('task_message', () => {
   test('uses explicit queue wording without legacy delivery terms', async () => {
     const board = new BackgroundJobBoard();
     registerRunningChild(board);
-    client = { session: { prompt: makePrompt() } };
+    client = { session: makeSession(makePrompt()) };
 
     const result = await createTool(board).execute(
       { task_id: 'ses_child1', message: 'Status update' },

src/tools/task-message.ts

@@ -3,6 +3,10 @@ import {
   type ToolDefinition,
   tool,
 } from '@opencode-ai/plugin';
+import {
+  type ContinuationModelSelection,
+  parseContinuationModelSelection,
+} from '../hooks/task-session-manager/continuation-model-selection';
 import type { BackgroundJobStore } from '../utils/background-job-store';
 import { getClient } from '../utils/opencode-client';
 import { OperationTimeoutError, withTimeout } from '../utils/session';
@@ -10,6 +14,7 @@ import { OperationTimeoutError, withTimeout } from '../utils/session';
 const z = tool.schema;
 const MAX_MESSAGE_LENGTH = 500;
 const DEFAULT_MESSAGE_TIMEOUT_MS = 10_000;
+const MODEL_LOOKUP_HISTORY_LIMIT = 20;
 
 class MessageLeaseOperationTimeoutError extends Error {
   constructor(
@@ -82,27 +87,66 @@ export function createTaskMessageTool(options: {
 
         const session = getClient(options.input).session;
         const prompt = session.prompt.bind(session);
-        const transportJob = getCurrentTaskMessageJob(
-          options.backgroundJobBoard,
-          parentSessionID,
-          requested,
-          lease.taskID,
-          lease.generation,
+        const messageTimeoutMs = Math.max(
+          1,
+          options.messageTimeoutMs ?? DEFAULT_MESSAGE_TIMEOUT_MS,
         );
+        const deadline = Date.now() + messageTimeoutMs;
+        const lookupController = new AbortController();
+        let modelSelection: ContinuationModelSelection | undefined;
+        try {
+          modelSelection = await withTimeout(
+            readCurrentChildModel(
+              session,
+              lease.taskID,
+              options.input.directory,
+              lookupController.signal,
+            ),
+            messageTimeoutMs,
+            `Task message model lookup timed out after ${messageTimeoutMs}ms`,
+          );
+        } finally {
+          lookupController.abort();
+        }
+        if (!modelSelection) {
+          throw new Error(
+            `Task ${requested} has no authoritative model identity; refusing message`,
+          );
+        }
+
+        const remainingTimeoutMs = deadline - Date.now();
+        if (remainingTimeoutMs <= 0) {
+          throw new OperationTimeoutError(
+            `Task message transport timed out after ${messageTimeoutMs}ms`,
+          );
+        }
+
         const response = await awaitMessageTransport(
           options.backgroundJobBoard,
           lease,
-          () =>
-            prompt({
+          () => {
+            assertMessageLease(options.backgroundJobBoard, lease, requested);
+            const currentJob = getCurrentTaskMessageJob(
+              options.backgroundJobBoard,
+              parentSessionID,
+              requested,
+              lease.taskID,
+              lease.generation,
+            );
+            const body = {
+              agent: currentJob.agent,
+              model: modelSelection.model,
+              variant: modelSelection.variant ?? 'default',
+              noReply: true,
+              parts: [{ type: 'text', text: args.message.trim() }],
+            } as Parameters<typeof prompt>[0]['body'];
+            return prompt({
               path: { id: lease.taskID },
-              body: {
-                agent: transportJob.agent,
-                noReply: true,
-                parts: [{ type: 'text', text: args.message.trim() }],
-              },
+              body,
               throwOnError: true,
-            }),
-          options.messageTimeoutMs ?? DEFAULT_MESSAGE_TIMEOUT_MS,
+            });
+          },
+          remainingTimeoutMs,
         );
         assertMessageLease(options.backgroundJobBoard, lease, requested);
         assertSuccessfulMessageResponse(response);
@@ -201,6 +245,78 @@ function errorText(error: unknown): string {
   }
 }
 
+async function readCurrentChildModel(
+  session: unknown,
+  taskID: string,
+  directory: string | undefined,
+  signal: AbortSignal,
+): Promise<ContinuationModelSelection | undefined> {
+  if (!isRecord(session)) return undefined;
+
+  const query = directory ? { directory } : undefined;
+  let get: unknown;
+  try {
+    get = session.get;
+  } catch {
+    get = undefined;
+  }
+  if (typeof get === 'function') {
+    try {
+      const response = await get.call(session, {
+        path: { id: taskID },
+        query,
+        signal,
+      });
+      if (isRecord(response) && isRecord(response.data)) {
+        const selection = parseContinuationModelSelection(
+          response.data.model,
+          response.data.variant,
+        );
+        if (selection) return selection;
+      }
+    } catch {
+      // Fall through to the authoritative latest user message.
+      if (signal.aborted) return undefined;
+    }
+  }
+  if (signal.aborted) return undefined;
+
+  let messages: unknown;
+  try {
+    messages = session.messages;
+  } catch {
+    messages = undefined;
+  }
+  if (typeof messages !== 'function') return undefined;
+
+  try {
+    const messageQuery = {
+      ...(directory ? { directory } : {}),
+      limit: MODEL_LOOKUP_HISTORY_LIMIT,
+    };
+    const response = await messages.call(session, {
+      path: { id: taskID },
+      query: messageQuery,
+      signal,
+    });
+    if (!isRecord(response) || !Array.isArray(response.data)) return undefined;
+
+    for (let index = response.data.length - 1; index >= 0; index -= 1) {
+      const message = response.data[index];
+      if (!isRecord(message) || !isRecord(message.info)) continue;
+      if (message.info.role !== 'user') continue;
+      return parseContinuationModelSelection(
+        message.info.model,
+        message.info.variant,
+      );
+    }
+  } catch {
+    return undefined;
+  }
+
+  return undefined;
+}
+
 function getCurrentTaskMessageJob(
   backgroundJobBoard: BackgroundJobStore,
   parentSessionID: string,