Fix ZCode 3.3.6 message contexts

sdk/index.ts

@@ -440,9 +440,14 @@ export type ZCodeSessionSummary = Record<string, unknown> & {
 export type ZCodeMessage = Record<string, unknown> & {
   messageId?: string;
   id?: string;
+  sessionId?: string;
   role?: string;
   turnId?: string;
   createdAt?: string | number;
+  completedAt?: string | number;
+  parentMessageId?: string;
+  info?: Record<string, unknown>;
+  parts?: unknown[];
 };
 
 export type ZCodeTaskSummary = Record<string, unknown> & {

src/protocol/extension-service.ts

@@ -9,6 +9,7 @@ import type {
   ModelRequestHistory,
   ModelRequestRecord,
   ModelTokenUsage,
+  ZCodeMessage,
   ZCodeSessionEvent,
   ZCodeStreamEnvelope,
   ZCodeSubscriptionTarget,
@@ -171,7 +172,7 @@ export class ExtensionZCodeService {
       }
       case "sessions.list": return this.#serviceCall(manifest, "zcode.sessions.read", "zcode-session", "listSessions", payload);
       case "sessions.read": return this.#serviceCall(manifest, "zcode.sessions.read", "zcode-session", "readSession", payload);
-      case "sessions.readMessages": return this.#serviceCall(manifest, "zcode.sessions.read", "zcode-session", "readSessionMessages", payload);
+      case "sessions.readMessages": return this.#readSessionMessages(manifest, payload) as Promise<T>;
       case "sessions.readEvents": return this.#serviceCall(manifest, "zcode.sessions.read", "zcode-session", "readSessionEvents", payload);
       case "sessions.create": return this.#serviceCall(manifest, "zcode.sessions.write", "zcode-session", "createSession", payload);
       case "sessions.resume": return this.#serviceCall(manifest, "zcode.sessions.write", "zcode-session", "resumeSession", payload);
@@ -329,6 +330,30 @@ export class ExtensionZCodeService {
     };
   }
 
+  async #readSessionMessages(manifest: ExtensionManifest, payload: unknown): Promise<ZCodeMessage[]> {
+    const value = requiredRecord(payload);
+    const workspacePath = requiredString(value, "workspacePath");
+    const sessionId = requiredString(value, "sessionId");
+    const workspaceIdentity = string(value.workspaceIdentity);
+    const limit = optionalPositiveInteger(value.limit);
+    const snapshot = requiredRecord(await this.#serviceCall(manifest, "zcode.sessions.read", "zcode-session", "readSession", {
+      workspacePath,
+      ...(workspaceIdentity ? {workspaceIdentity} : {}),
+      sessionId,
+      ...(limit ? {messageLimit: limit} : {}),
+    }));
+    const messages = (Array.isArray(snapshot.messages) ? snapshot.messages : [])
+      .map(normalizeZCodeMessage)
+      .filter((message): message is ZCodeMessage => Boolean(message));
+    const afterMessageId = string(value.afterMessageId);
+    if (afterMessageId) {
+      const index = messages.findIndex((message) => message.messageId === afterMessageId || message.id === afterMessageId);
+      const remaining = index >= 0 ? messages.slice(index + 1) : messages;
+      return limit ? remaining.slice(0, limit) : remaining;
+    }
+    return limit ? messages.slice(-limit) : messages;
+  }
+
   #resolveSessionTarget(sessionId: string): {workspacePath: string; workspaceIdentity?: string; sessionId: string} | undefined {
     const dataRoot = process.env.ZCODE_DATA_BASE_DIR?.trim() || process.env.HOME?.trim() || homedir();
     const databasePath = path.join(dataRoot, ".zcode", "v2", "tasks-index.sqlite");
@@ -462,6 +487,32 @@ function sanitizeModelRequest(sessionId: string, value: unknown): ModelRequestRe
   };
 }
 
+export function normalizeZCodeMessage(value: unknown): ZCodeMessage | undefined {
+  const source = record(value);
+  const info = record(source.info);
+  const time = record(info.time ?? source.time);
+  const messageId = string(source.messageId ?? source.id ?? info.messageId ?? info.id);
+  if (!messageId) return undefined;
+  const sessionId = string(source.sessionId ?? info.sessionId ?? info.sessionID);
+  const role = string(source.role ?? info.role);
+  const turnId = string(source.turnId ?? info.turnId);
+  const parentMessageId = string(source.parentMessageId ?? info.parentMessageId ?? info.parentID);
+  const createdAt = string(source.createdAt ?? time.created) ?? number(source.createdAt ?? time.created);
+  const completedAt = string(source.completedAt ?? time.completed) ?? number(source.completedAt ?? time.completed);
+  return {
+    ...source,
+    info,
+    messageId,
+    id: messageId,
+    sessionId,
+    role,
+    turnId,
+    parentMessageId,
+    createdAt,
+    completedAt,
+  };
+}
+
 function normalizeCallSource(value: unknown, fallbackQuerySource?: string): ModelCallSource {
   const source = record(value);
   const querySource = string(source.querySource) ?? fallbackQuerySource;

src/renderer/extension-runtime.ts

@@ -559,9 +559,16 @@ function insertContribution(target: Element, host: HTMLElement, placement: Contr
 }
 
 function contextFrom(target: Element): ActiveUiContext {
-  const closest = (attribute: string) => target.closest<HTMLElement>(`[${attribute}]`)?.getAttribute(attribute)
+  const nearest = (attribute: string) => target.closest<HTMLElement>(`[${attribute}]`)?.getAttribute(attribute) ?? undefined;
+  const closest = (attribute: string) => nearest(attribute)
     ?? document.querySelector<HTMLElement>(`[${attribute}]`)?.getAttribute(attribute)
     ?? undefined;
+  const messageElement = target.matches("[data-message-id][data-role]")
+    ? target as HTMLElement
+    : target.matches("[data-chat-turn-group-shell]")
+      ? target.querySelector<HTMLElement>('[data-message-id][data-role="assistant"]')
+        ?? target.querySelector<HTMLElement>("[data-message-id][data-role]")
+      : undefined;
   const taskItemKey = closest("data-task-item-key");
   const taskItemId = taskItemKey && taskItemKey.includes(":")
     ? taskItemKey.slice(taskItemKey.lastIndexOf(":") + 1)
@@ -571,9 +578,13 @@ function contextFrom(target: Element): ActiveUiContext {
     workspaceIdentity: closest("data-workspace-identity"),
     taskId: closest("data-task-id") ?? taskItemId,
     sessionId: closest("data-session-id"),
-    turnId: closest("data-anchor-turn-id") ?? closest("data-turn-id"),
-    messageId: closest("data-message-id") ?? closest("data-anchor-assistant-message-id"),
-    role: closest("data-role"),
+    turnId: nearest("data-anchor-turn-id") ?? nearest("data-turn-id")
+      ?? messageElement?.getAttribute("data-anchor-turn-id") ?? messageElement?.getAttribute("data-turn-id")
+      ?? closest("data-anchor-turn-id") ?? closest("data-turn-id"),
+    messageId: nearest("data-anchor-assistant-message-id")
+      ?? messageElement?.getAttribute("data-message-id")
+      ?? closest("data-message-id") ?? closest("data-anchor-assistant-message-id"),
+    role: messageElement?.getAttribute("data-role") ?? closest("data-role"),
     runtimeStatus: closest("data-runtime-status"),
     toolCallId: closest("data-tool-call-id"),
   };

tests/extension-service.test.ts

@@ -3,6 +3,7 @@ import type {ExtensionManifest} from "../sdk/index.ts";
 import {
   ExtensionZCodeService,
   normalizeModelRequestEvent,
+  normalizeZCodeMessage,
   normalizeSessionEvent,
 } from "../src/protocol/extension-service.ts";
 
@@ -146,6 +147,40 @@ describe("extension ZCode service", () => {
     ]);
   });
 
+  test("reads normalized messages through the 3.3.6 session snapshot compatibility path", async () => {
+    const calls: unknown[] = [];
+    const session = {
+      async readSession(payload: unknown) {
+        calls.push(payload);
+        return {messages: [
+          {info: {messageId: "message-1", sessionId: "session-1", role: "user", time: {created: 10}}, parts: []},
+          {info: {id: "message-2", sessionID: "session-1", role: "assistant", parentID: "message-1", time: {created: 20, completed: 30}}, parts: []},
+          {info: {messageId: "message-3", sessionId: "session-1", role: "assistant", time: {created: 40}}, parts: []},
+        ]};
+      },
+      async readSessionMessages() {
+        throw new Error("The incompatible 3.3.6 method must not be called");
+      },
+    };
+    const service = createService({}, {
+      async service(name: string) { return name === "zcode-session" ? session : {}; },
+    });
+    const messages = await service.invoke<Array<Record<string, unknown>>>(
+      manifest(["zcode.sessions.read"]),
+      "sessions.readMessages",
+      {workspacePath: "D:\\project", sessionId: "session-1", afterMessageId: "message-1", limit: 2},
+    );
+    expect(calls).toEqual([{workspacePath: "D:\\project", sessionId: "session-1", messageLimit: 2}]);
+    expect(messages).toEqual([
+      expect.objectContaining({
+        messageId: "message-2", id: "message-2", sessionId: "session-1", role: "assistant",
+        parentMessageId: "message-1", createdAt: 20, completedAt: 30,
+      }),
+      expect.objectContaining({messageId: "message-3", id: "message-3", createdAt: 40}),
+    ]);
+    expect(normalizeZCodeMessage({info: {}, parts: []})).toBeUndefined();
+  });
+
   test("normalizes task, workspace, and broadcast streams without discarding raw envelopes", () => {
     const subscriptions: Array<{channel: string; event: string}> = [];
     const service = createService({}, {

tests/renderer-runtime.test.ts

@@ -1,6 +1,6 @@
 import {afterEach, beforeEach, describe, expect, test} from "bun:test";
 import {Window} from "happy-dom";
-import type {ExtensionHostCapabilities, ExtensionManifest, RendererExtension} from "../sdk/index.ts";
+import type {ActiveUiContext, ExtensionHostCapabilities, ExtensionManifest, RendererExtension} from "../sdk/index.ts";
 import type {PluginStatus} from "../src/shared/schemas.ts";
 import type {ZdpBridge} from "../src/renderer/globals.d.ts";
 import {RendererExtensionRuntime} from "../src/renderer/extension-runtime.ts";
@@ -114,6 +114,37 @@ describe("renderer extension runtime", () => {
     await expect(restricted.mountPage("runtime-test", "dashboard", container as never))
       .rejects.toThrow("ui.pages");
   });
+
+  test("supplies the assistant descendant as context for a real ZCode turn shell without turn ids", async () => {
+    const runtime = new RendererExtensionRuntime(createBridge(["ui.chat"]).value);
+    let mountedContext: ActiveUiContext | undefined;
+    runtime.register({
+      id: "runtime-test",
+      activate(context) {
+        context.ui.contribute("chat.turn.after", (_container, active) => {
+          mountedContext = active;
+        });
+      },
+    });
+    runtime.sync([status(1, true, ["ui.chat"])]);
+    const root = window.document.getElementById("root")!;
+    root.innerHTML = `<main data-testid="chat-view" data-session-id="session-1">
+      <section data-chat-turn-group-shell="true">
+        <article data-message-id="user-message" data-role="user"></article>
+        <article data-message-id="assistant-message" data-role="assistant"></article>
+      </section>
+    </main>`;
+    runtime.refreshUi();
+    await settle();
+
+    expect(mountedContext).toMatchObject({
+      sessionId: "session-1",
+      messageId: "assistant-message",
+      role: "assistant",
+    });
+    expect(mountedContext?.turnId).toBeUndefined();
+    await runtime.dispose();
+  });
 });
 
 function appendChatMessage(messageId: string): HTMLElement {