Preserve Pi import model and thinking (#1441)

* Preserve Pi import model and thinking * Format Pi import config helper * Preserve Pi thinking in large imports --------- Co-authored-by: Mohamed Boudra <boudra.moha@gmail.com>
2026-06-10 16:52:15 +08:00
parent 8e6455f1d6
commit 5ea586a926
4 changed files with 366 additions and 21 deletions
@@ -779,6 +779,85 @@ describe("PiRpcAgentClient", () => {
    ]);
  });

+  test("imports JSONL sessions with the recorded model and thinking level", async () => {
+    const root = mkdtempSync(path.join(tmpdir(), "paseo-pi-import-config-"));
+    const cwd = path.join(root, "workspace");
+    const sessionsDir = path.join(root, "sessions");
+    mkdirSync(sessionsDir, { recursive: true });
+    const sessionFile = path.join(sessionsDir, "20260103_session.jsonl");
+    writeFileSync(
+      sessionFile,
+      [
+        JSON.stringify({
+          type: "session",
+          version: 3,
+          id: "pi-import-session",
+          timestamp: "2026-01-03T00:00:00.000Z",
+          cwd,
+        }),
+        JSON.stringify({
+          type: "message",
+          id: "entry-1",
+          timestamp: "2026-01-03T00:00:01.000Z",
+          message: { role: "user", content: "first prompt" },
+        }),
+        JSON.stringify({
+          type: "model_change",
+          id: "model-1",
+          timestamp: "2026-01-03T00:00:02.000Z",
+          provider: "openrouter",
+          modelId: "anthropic/claude-sonnet-4.5",
+        }),
+        JSON.stringify({
+          type: "thinking_level_change",
+          id: "thinking-1",
+          timestamp: "2026-01-03T00:00:03.000Z",
+          thinkingLevel: "high",
+        }),
+      ].join("\n") + "\n",
+      "utf8",
+    );
+    const pi = new FakePi();
+    const client = new PiRpcAgentClient({
+      logger: pino({ level: "silent" }),
+      runtime: pi,
+      providerParams: { sessionDir: sessionsDir },
+    });
+
+    const imported = await client.importSession(
+      { providerHandleId: sessionFile, cwd },
+      { config: createConfig({ cwd }), storedConfig: createConfig({ cwd }) },
+    );
+
+    const actualLaunch = pi.recordedLaunches[0]!;
+    expect(actualLaunch.extensionPaths).toHaveLength(1);
+    expect(actualLaunch.argv).toEqual([
+      "pi",
+      "--mode",
+      "rpc",
+      "--model",
+      "openrouter/anthropic/claude-sonnet-4.5",
+      "--thinking",
+      "high",
+      "--session",
+      sessionFile,
+      "--extension",
+      actualLaunch.extensionPaths[0],
+    ]);
+    expect(imported.config).toMatchObject({
+      provider: "pi",
+      cwd,
+      model: "openrouter/anthropic/claude-sonnet-4.5",
+      thinkingOptionId: "high",
+    });
+    expect(imported.persistence.metadata).toMatchObject({
+      provider: "pi",
+      cwd,
+      model: "openrouter/anthropic/claude-sonnet-4.5",
+      thinkingOptionId: "high",
+    });
+  });
+
  test("lists models from a short-lived Pi session in the requested cwd", async () => {
    const pi = new FakePi();
    const client = createClient(pi);
@@ -57,7 +57,7 @@ import {
 } from "./history-mapper.js";
 import { PiCliRuntime } from "./cli-runtime.js";
 import { revertPiConversation } from "./rewind.js";
-import { listPiImportableSessions } from "./session-descriptor.js";
+import { listPiImportableSessions, readPiImportSessionConfig } from "./session-descriptor.js";
 import type { PiRuntime, PiRuntimeSession } from "./runtime.js";
 import type {
  PiAgentSessionEvent,
@@ -1984,11 +1984,13 @@ export class PiRpcAgentClient implements AgentClient {
  }

  async importSession(input: ImportProviderSessionInput, context: ImportProviderSessionContext) {
+    const importConfig = await readPiImportSessionConfig(input.providerHandleId);
    return importSessionFromPersistence({
      provider: PI_PROVIDER,
      request: input,
      context,
      resumeSession: this.resumeSession.bind(this),
+      config: importConfig,
    });
  }

@@ -0,0 +1,172 @@
+import { mkdtemp, mkdir, writeFile } from "node:fs/promises";
+import { tmpdir } from "node:os";
+import path from "node:path";
+import { expect, test } from "vitest";
+
+import { listPiImportableSessions, readPiImportSessionConfig } from "./session-descriptor.js";
+
+async function writeSession(root: string, lines: unknown[]): Promise<string> {
+  const sessionsDir = path.join(root, "sessions", "project");
+  await mkdir(sessionsDir, { recursive: true });
+  const filePath = path.join(sessionsDir, "2026-06-09T00-00-00-000Z_session.jsonl");
+  await writeFile(filePath, `${lines.map((line) => JSON.stringify(line)).join("\n")}\n`, "utf8");
+  return filePath;
+}
+
+test("Pi import config preserves the latest recorded model and thinking level", async () => {
+  const root = await mkdtemp(path.join(tmpdir(), "paseo-pi-session-model-"));
+  const cwd = path.join(root, "repo");
+  const sessionFile = await writeSession(root, [
+    {
+      type: "session",
+      version: 3,
+      id: "session-1",
+      timestamp: "2026-06-09T00:00:00.000Z",
+      cwd,
+    },
+    {
+      type: "model_change",
+      id: "model-1",
+      timestamp: "2026-06-09T00:00:01.000Z",
+      provider: "openai-codex",
+      modelId: "gpt-5.1",
+    },
+    {
+      type: "thinking_level_change",
+      id: "thinking-1",
+      timestamp: "2026-06-09T00:00:01.500Z",
+      thinkingLevel: "low",
+    },
+    {
+      type: "message",
+      id: "user-1",
+      timestamp: "2026-06-09T00:00:02.000Z",
+      message: {
+        role: "user",
+        content: [{ type: "text", text: "hello" }],
+      },
+    },
+    {
+      type: "model_change",
+      id: "model-2",
+      timestamp: "2026-06-09T00:00:03.000Z",
+      provider: "openrouter",
+      modelId: "anthropic/claude-sonnet-4.5",
+    },
+    {
+      type: "thinking_level_change",
+      id: "thinking-2",
+      timestamp: "2026-06-09T00:00:04.000Z",
+      thinkingLevel: "high",
+    },
+  ]);
+
+  const [descriptor] = await listPiImportableSessions({ sessionDir: path.join(root, "sessions") });
+  const importConfig = await readPiImportSessionConfig(sessionFile);
+
+  expect(descriptor).toMatchObject({
+    providerHandleId: sessionFile,
+    cwd,
+    firstPromptPreview: "hello",
+  });
+  expect(importConfig).toEqual({
+    model: "openrouter/anthropic/claude-sonnet-4.5",
+    thinkingOptionId: "high",
+  });
+});
+
+test("Pi import config can infer model from assistant messages", async () => {
+  const root = await mkdtemp(path.join(tmpdir(), "paseo-pi-session-message-model-"));
+  const cwd = path.join(root, "repo");
+  const sessionFile = await writeSession(root, [
+    {
+      type: "session",
+      version: 3,
+      id: "session-2",
+      timestamp: "2026-06-09T00:00:00.000Z",
+      cwd,
+    },
+    {
+      type: "message",
+      id: "user-1",
+      timestamp: "2026-06-09T00:00:01.000Z",
+      message: {
+        role: "user",
+        content: [{ type: "text", text: "hello" }],
+      },
+    },
+    {
+      type: "message",
+      id: "assistant-1",
+      timestamp: "2026-06-09T00:00:02.000Z",
+      message: {
+        role: "assistant",
+        content: [{ type: "text", text: "hi" }],
+        provider: "google",
+        model: "gemini-2.5-pro",
+      },
+    },
+  ]);
+
+  const importConfig = await readPiImportSessionConfig(sessionFile);
+
+  expect(importConfig).toEqual({
+    model: "google/gemini-2.5-pro",
+  });
+});
+
+test("Pi import config preserves thinking before a later model in large sessions", async () => {
+  const root = await mkdtemp(path.join(tmpdir(), "paseo-pi-session-large-thinking-"));
+  const cwd = path.join(root, "repo");
+  const fillerMessages = Array.from({ length: 2_100 }, (_, index) => ({
+    type: "message",
+    id: `filler-${index}`,
+    timestamp: `2026-06-09T00:${String(Math.floor(index / 60)).padStart(2, "0")}:${String(index % 60).padStart(2, "0")}.000Z`,
+    message: {
+      role: "assistant",
+      content: [{ type: "text", text: `filler ${index}` }],
+    },
+  }));
+  const sessionFile = await writeSession(root, [
+    {
+      type: "session",
+      version: 3,
+      id: "session-3",
+      timestamp: "2026-06-09T00:00:00.000Z",
+      cwd,
+    },
+    {
+      type: "message",
+      id: "user-1",
+      timestamp: "2026-06-09T00:00:01.000Z",
+      message: { role: "user", content: "hello" },
+    },
+    ...fillerMessages,
+    {
+      type: "thinking_level_change",
+      id: "thinking-1",
+      timestamp: "2026-06-09T01:00:00.000Z",
+      thinkingLevel: "low",
+    },
+    {
+      type: "model_change",
+      id: "model-1",
+      timestamp: "2026-06-09T01:00:01.000Z",
+      provider: "openrouter",
+      modelId: "google/gemini-2.5-pro",
+    },
+    {
+      type: "session_info",
+      id: "info-1",
+      timestamp: "2026-06-09T01:00:02.000Z",
+      name: "large session",
+    },
+  ]);
+
+  const importConfig = await readPiImportSessionConfig(sessionFile);
+
+  expect(importConfig).toEqual({
+    model: "openrouter/google/gemini-2.5-pro",
+    thinkingOptionId: "low",
+  });
+});
@@ -33,6 +33,30 @@ interface PiSessionTail {
  title: string | null;
  lastActivityAt: Date | null;
  lastUserMessage: string | null;
+  model: string | null;
+  thinkingOptionId: string | null;
+}
+
+interface PiSessionHead {
+  title: string | null;
+  firstUserMessage: string | null;
+  model: string | null;
+  thinkingOptionId: string | null;
+}
+
+interface PiSessionDescriptor {
+  cwd: string;
+  title: string | null;
+  firstUserMessage: string | null;
+  lastUserMessage: string | null;
+  lastActivityAt: Date;
+  model: string | null;
+  thinkingOptionId: string | null;
+}
+
+export interface PiImportSessionConfig {
+  model?: string;
+  thinkingOptionId?: string;
 }

 export async function listPiImportableSessions(
@@ -56,6 +80,12 @@ export async function listPiImportableSessions(
    .slice(0, limit);
 }

+export async function readPiImportSessionConfig(filePath: string): Promise<PiImportSessionConfig> {
+  const descriptor = await readPiSessionDescriptor(filePath);
+  if (!descriptor) return {};
+  return toPiImportSessionConfig(descriptor);
+}
+
 async function resolvePiSessionsDir(options: PiSessionDescriptorOptions): Promise<string> {
  const env = options.env ?? process.env;
  const homeDir = options.homeDir ?? homedir();
@@ -155,6 +185,22 @@ async function walkJsonlFiles(root: string): Promise<string[]> {
 async function readPiImportableSession(
  filePath: string,
 ): Promise<ImportableProviderSession | null> {
+  const descriptor = await readPiSessionDescriptor(filePath);
+  if (!descriptor) return null;
+
+  return {
+    providerHandleId: filePath,
+    cwd: descriptor.cwd,
+    title: descriptor.title,
+    firstPromptPreview: normalizePromptPreview(descriptor.firstUserMessage),
+    lastPromptPreview: normalizePromptPreview(
+      descriptor.lastUserMessage ?? descriptor.firstUserMessage,
+    ),
+    lastActivityAt: descriptor.lastActivityAt,
+  };
+}
+
+async function readPiSessionDescriptor(filePath: string): Promise<PiSessionDescriptor | null> {
  const firstLine = await readFirstLine(filePath);
  if (!firstLine) return null;
  const header = parseSessionHeader(firstLine);
@@ -164,17 +210,26 @@ async function readPiImportableSession(
  const tailInfo = parseSessionTail(tail);
  const headInfo = await scanSessionHead(filePath);
  const title = tailInfo.title ?? headInfo.title ?? headInfo.firstUserMessage;
+  const model = tailInfo.model ?? headInfo.model;
+  const thinkingOptionId = tailInfo.thinkingOptionId ?? headInfo.thinkingOptionId;
  const lastActivityAt =
    tailInfo.lastActivityAt ?? (await readFileMtime(filePath)) ?? header.createdAt ?? new Date(0);
+
  return {
-    providerHandleId: filePath,
    cwd: header.cwd,
    title,
-    firstPromptPreview: normalizePromptPreview(headInfo.firstUserMessage),
-    lastPromptPreview: normalizePromptPreview(
-      tailInfo.lastUserMessage ?? headInfo.firstUserMessage,
-    ),
+    firstUserMessage: headInfo.firstUserMessage,
+    lastUserMessage: tailInfo.lastUserMessage,
    lastActivityAt,
+    model,
+    thinkingOptionId,
+  };
+}
+
+function toPiImportSessionConfig(descriptor: PiSessionDescriptor): PiImportSessionConfig {
+  return {
+    ...(descriptor.model ? { model: descriptor.model } : {}),
+    ...(descriptor.thinkingOptionId ? { thinkingOptionId: descriptor.thinkingOptionId } : {}),
  };
 }

@@ -231,6 +286,8 @@ function parseSessionTail(tail: string): PiSessionTail {
  let lastActivityAt: Date | null = null;
  let fallbackTimestamp: Date | null = null;
  let lastUserMessage: string | null = null;
+  let model: string | null = null;
+  let thinkingOptionId: string | null = null;

  for (let index = lines.length - 1; index >= 0; index -= 1) {
    const entry = parseJsonRecord(lines[index].trim());
@@ -240,14 +297,20 @@ function parseSessionTail(tail: string): PiSessionTail {
      title = readNonEmptyString(entry.name);
    }

+    if (!model) {
+      model = extractModel(entry);
+    }
+
+    if (!thinkingOptionId) {
+      thinkingOptionId = extractThinkingOptionId(entry);
+    }
+
    const entryTimestamp = parseDate(entry.timestamp);
    if (!fallbackTimestamp && entryTimestamp) {
      fallbackTimestamp = entryTimestamp;
    }

-    if (entry.type !== "message") {
-      continue;
-    }
+    if (entry.type !== "message") continue;

    if (!lastActivityAt && entryTimestamp) {
      lastActivityAt = entryTimestamp;
@@ -256,28 +319,29 @@ function parseSessionTail(tail: string): PiSessionTail {
    if (!lastUserMessage && isRecord(entry.message) && entry.message.role === "user") {
      lastUserMessage = extractMessageText(entry.message.content);
    }
-
-    if (title && lastActivityAt && lastUserMessage) {
-      break;
-    }
  }

-  return { title, lastActivityAt: lastActivityAt ?? fallbackTimestamp, lastUserMessage };
+  return {
+    title,
+    lastActivityAt: lastActivityAt ?? fallbackTimestamp,
+    lastUserMessage,
+    model,
+    thinkingOptionId,
+  };
 }

-async function scanSessionHead(filePath: string): Promise<{
-  title: string | null;
-  firstUserMessage: string | null;
-}> {
+async function scanSessionHead(filePath: string): Promise<PiSessionHead> {
  let content: string;
  try {
    content = await readFile(filePath, "utf8");
  } catch {
-    return { title: null, firstUserMessage: null };
+    return { title: null, firstUserMessage: null, model: null, thinkingOptionId: null };
  }

  let title: string | null = null;
  let firstUserMessage: string | null = null;
+  let model: string | null = null;
+  let thinkingOptionId: string | null = null;
  let lineCount = 0;

  for (const rawLine of content.split(/\r?\n/u)) {
@@ -289,13 +353,16 @@ async function scanSessionHead(filePath: string): Promise<{
      title = readNonEmptyString(entry.name) ?? title;
    }

+    model = extractModel(entry) ?? model;
+    thinkingOptionId = extractThinkingOptionId(entry) ?? thinkingOptionId;
+
    if (!firstUserMessage && entry.type === "message" && isRecord(entry.message)) {
      if (entry.message.role === "user") {
        firstUserMessage = extractMessageText(entry.message.content);
      }
    }

-    if (title && firstUserMessage) {
+    if (title && firstUserMessage && model && thinkingOptionId) {
      break;
    }
    if (lineCount >= FULL_SCAN_LINE_LIMIT && firstUserMessage) {
@@ -303,7 +370,32 @@ async function scanSessionHead(filePath: string): Promise<{
    }
  }

-  return { title, firstUserMessage };
+  return { title, firstUserMessage, model, thinkingOptionId };
+}
+
+function extractModel(entry: Record<string, unknown>): string | null {
+  if (entry.type === "model_change") {
+    return buildModelId(entry.provider, entry.modelId);
+  }
+
+  if (entry.type === "message" && isRecord(entry.message)) {
+    return buildModelId(entry.message.provider, entry.message.model);
+  }
+
+  return null;
+}
+
+function extractThinkingOptionId(entry: Record<string, unknown>): string | null {
+  return entry.type === "thinking_level_change" ? readNonEmptyString(entry.thinkingLevel) : null;
+}
+
+function buildModelId(provider: unknown, modelId: unknown): string | null {
+  const providerName = readNonEmptyString(provider);
+  const modelName = readNonEmptyString(modelId);
+  if (!providerName || !modelName) {
+    return null;
+  }
+  return `${providerName}/${modelName}`;
 }

 function parseJsonRecord(line: string): Record<string, unknown> | null {