@miaws/miaw 1.18.4 → 1.18.5

package/src/cli/cmd/run/session-replay.ts

@@ -0,0 +1,374 @@
+import type { Event, PermissionRequest, QuestionRequest } from "@miaw/sdk/v2"
+import { bootstrapSessionData, createSessionData, reduceSessionData, type SessionData } from "./session-data"
+import { messagePrompt, type SessionMessages } from "./session.shared"
+import { messageTurnSummaryCommit } from "./turn-summary"
+import type { FooterPatch, LocalReplayRow, RunProvider, StreamCommit } from "./types"
+
+type ReplayInput = {
+  messages: SessionMessages
+  permissions: PermissionRequest[]
+  questions: QuestionRequest[]
+  thinking: boolean
+  limits: Record<string, number>
+  providers?: RunProvider[]
+}
+
+type ReplayConfig = {
+  limits: Record<string, number>
+  providers?: RunProvider[]
+  summaries: ReadonlySet<string>
+}
+
+export type SessionReplay = {
+  data: SessionData
+  commits: StreamCommit[]
+  patch?: FooterPatch
+}
+
+type ReplayMessage = {
+  commits: StreamCommit[]
+  patch?: FooterPatch
+}
+
+const SHELL_SYNTHETIC_USER_TEXT = "The following tool was executed by the user"
+
+function apply(data: SessionData, event: Event, sessionID: string, thinking: boolean, limits: Record<string, number>) {
+  return reduceSessionData({
+    data,
+    event,
+    sessionID,
+    thinking,
+    limits,
+  })
+}
+
+function mergePatch(left: FooterPatch | undefined, right: FooterPatch | undefined) {
+  if (!left) {
+    return right
+  }
+
+  if (!right) {
+    return left
+  }
+
+  return {
+    ...left,
+    ...right,
+  }
+}
+
+function active(data: SessionData) {
+  return data.part.size > 0 || data.tools.size > 0
+}
+
+function replayPatch(data: SessionData, patch: FooterPatch | undefined) {
+  if (active(data)) {
+    if (!patch) {
+      return {
+        phase: "running",
+      } satisfies FooterPatch
+    }
+
+    return {
+      ...patch,
+      phase: "running",
+    } satisfies FooterPatch
+  }
+
+  if (data.permissions.length > 0 || data.questions.length > 0) {
+    if (!patch) {
+      return {
+        phase: "idle",
+      } satisfies FooterPatch
+    }
+
+    return {
+      ...patch,
+      phase: "idle",
+    } satisfies FooterPatch
+  }
+
+  if (!patch) {
+    return undefined
+  }
+
+  return {
+    ...patch,
+    phase: "idle",
+    status: "",
+  } satisfies FooterPatch
+}
+
+function isShellSyntheticUser(message: SessionMessages[number]) {
+  if (message.info.role !== "user") {
+    return false
+  }
+
+  const prompt = messagePrompt(message)
+  return (
+    !prompt.text.trim() &&
+    prompt.parts.length === 0 &&
+    message.parts.some((part) => part.type === "text" && part.synthetic && part.text === SHELL_SYNTHETIC_USER_TEXT)
+  )
+}
+
+function isShellSyntheticAssistant(message: SessionMessages[number], shellParents: ReadonlySet<string>) {
+  return (
+    message.info.role === "assistant" &&
+    shellParents.has(message.info.parentID) &&
+    message.parts.some((part) => part.type === "tool" && part.tool === "bash")
+  )
+}
+
+function summaryMessageIDs(messages: SessionMessages): ReadonlySet<string> {
+  const shellParents = new Set(messages.filter(isShellSyntheticUser).map((message) => message.info.id))
+  const parents = new Set<string>()
+  const summaries = new Set<string>()
+
+  for (let idx = messages.length - 1; idx >= 0; idx -= 1) {
+    const message = messages[idx]
+    if (!message || message.info.role !== "assistant") {
+      continue
+    }
+
+    if (isShellSyntheticAssistant(message, shellParents)) {
+      continue
+    }
+
+    if (parents.has(message.info.parentID)) {
+      continue
+    }
+
+    parents.add(message.info.parentID)
+
+    const completed = message.info.time.completed
+    if (typeof completed === "number" && completed > message.info.time.created) {
+      summaries.add(message.info.id)
+    }
+  }
+
+  return summaries
+}
+
+function replayMessage(
+  data: SessionData,
+  message: SessionMessages[number],
+  thinking: boolean,
+  config: ReplayConfig,
+): ReplayMessage {
+  if (message.info.role === "user") {
+    const prompt = messagePrompt(message)
+    if (!prompt.text.trim()) {
+      return {
+        commits: [],
+      }
+    }
+
+    return {
+      commits: [
+        {
+          kind: "user",
+          text: prompt.text,
+          phase: "start",
+          source: "system",
+          messageID: message.info.id,
+        },
+      ],
+    }
+  }
+
+  const commits: StreamCommit[] = []
+  let patch: FooterPatch | undefined
+
+  const info = apply(
+    data,
+    {
+      id: `bootstrap:message:${message.info.id}`,
+      type: "message.updated",
+      properties: {
+        sessionID: message.info.sessionID,
+        info: message.info,
+      },
+    },
+    message.info.sessionID,
+    thinking,
+    config.limits,
+  )
+  commits.push(...info.commits)
+  patch = mergePatch(patch, info.footer?.patch)
+
+  for (const part of message.parts) {
+    const next = apply(
+      data,
+      {
+        id: `bootstrap:part:${part.id}`,
+        type: "message.part.updated",
+        properties: {
+          sessionID: part.sessionID,
+          part,
+          time: 0,
+        },
+      },
+      message.info.sessionID,
+      thinking,
+      config.limits,
+    )
+    patch = mergePatch(patch, next.footer?.patch)
+    commits.push(...next.commits)
+  }
+
+  const summary = config.summaries.has(message.info.id)
+    ? messageTurnSummaryCommit(message, config.providers)
+    : undefined
+  if (summary) {
+    commits.push(summary)
+  }
+
+  return {
+    commits,
+    patch,
+  }
+}
+
+export function replaySession(input: ReplayInput): SessionReplay {
+  const data = createSessionData()
+  const commits: StreamCommit[] = []
+  let patch: FooterPatch | undefined
+  const summaries = summaryMessageIDs(input.messages)
+
+  bootstrapSessionData({
+    data,
+    messages: input.messages,
+    permissions: input.permissions,
+    questions: input.questions,
+  })
+
+  for (const message of input.messages) {
+    const next = replayMessage(data, message, input.thinking, {
+      limits: input.limits,
+      providers: input.providers,
+      summaries,
+    })
+    commits.push(...next.commits)
+    patch = mergePatch(patch, next.patch)
+  }
+
+  return {
+    data,
+    commits,
+    patch: replayPatch(data, patch),
+  }
+}
+
+export function replayLocalRows(
+  messages: SessionMessages,
+  commits: StreamCommit[],
+  rows: LocalReplayRow[],
+): StreamCommit[] {
+  const persisted = new Set(messages.map((message) => message.info.id))
+  return rows.reduce((out, local) => {
+    const row = local.commit
+    if (row.kind === "user" && row.messageID && persisted.has(row.messageID)) {
+      return out
+    }
+
+    if (!row.messageID) {
+      return [...out, row]
+    }
+
+    const exact = local.after
+      ? out.findIndex(
+          (commit) =>
+            commit.kind === local.after?.kind &&
+            commit.text === local.after.text &&
+            commit.phase === local.after.phase &&
+            commit.toolState === local.after.toolState &&
+            (local.after.partID ? commit.partID === local.after.partID : commit.messageID === local.after.messageID),
+        )
+      : -1
+    const anchored =
+      exact !== -1
+        ? exact
+        : local.after
+          ? out.findLastIndex((commit) =>
+              local.after?.partID
+                ? commit.partID === local.after.partID
+                : commit.kind === local.after?.kind && commit.messageID === local.after.messageID,
+            )
+          : -1
+    if (anchored !== -1) {
+      const commit = out[anchored]
+      const visible = local.after?.visible
+      if (commit && visible && commit.text.startsWith(visible) && commit.text.length > visible.length) {
+        return [
+          ...out.slice(0, anchored),
+          { ...commit, text: visible },
+          row,
+          { ...commit, text: commit.text.slice(visible.length) },
+          ...out.slice(anchored + 1),
+        ]
+      }
+
+      return [...out.slice(0, anchored + 1), row, ...out.slice(anchored + 1)]
+    }
+
+    const after = out.findIndex((commit) => commit.kind === "user" && commit.messageID === row.messageID)
+    if (after !== -1) {
+      return [...out.slice(0, after + 1), row, ...out.slice(after + 1)]
+    }
+
+    const before = out.findIndex((commit) => commit.messageID && row.messageID! < commit.messageID)
+    if (before === -1) {
+      return [...out, row]
+    }
+
+    return [...out.slice(0, before), row, ...out.slice(before)]
+  }, commits)
+}
+
+export function replayActiveText(data: SessionData, current: SessionData): StreamCommit[] {
+  return [...current.part.entries()].flatMap(([partID, kind]) => {
+    if (kind === "user" || current.end.has(partID) || data.ids.has(partID)) {
+      return []
+    }
+
+    const text = current.text.get(partID) ?? ""
+    const existing = data.text.get(partID) ?? ""
+    const sent = current.sent.get(partID) ?? 0
+    const existingSent = data.sent.get(partID) ?? 0
+    const visible = current.visible.get(partID) ?? ""
+    const existingVisible = data.visible.get(partID) ?? ""
+    if (!text.startsWith(existing) || existingSent > sent || !visible.startsWith(existingVisible)) {
+      return []
+    }
+
+    data.part.set(partID, kind)
+    data.text.set(partID, text)
+    data.sent.set(partID, sent)
+    data.visible.set(partID, visible)
+    const messageID = current.msg.get(partID)
+    if (messageID) {
+      data.msg.set(partID, messageID)
+      const role = current.role.get(messageID)
+      if (role) {
+        data.role.set(messageID, role)
+      }
+    }
+
+    const chunk = visible.slice(existingVisible.length)
+    if (!chunk) {
+      return []
+    }
+
+    return [
+      {
+        kind,
+        text: chunk,
+        phase: "progress",
+        source: kind,
+        ...(messageID ? { messageID } : {}),
+        partID,
+      },
+    ] satisfies StreamCommit[]
+  })
+}

package/src/cli/cmd/run/session.shared.ts

@@ -0,0 +1,196 @@
+// Session message extraction and prompt history.
+//
+// Fetches session messages from the SDK and extracts user turn text for
+// the prompt history ring. Also finds the most recently used variant for
+// the current model so the footer can pre-select it.
+import { promptCopy, promptSame } from "./prompt.shared"
+import type { RunInput, RunPrompt } from "./types"
+
+const LIMIT = 200
+
+export type SessionMessages = NonNullable<Awaited<ReturnType<RunInput["sdk"]["session"]["messages"]>>["data"]>
+
+type Turn = {
+  prompt: RunPrompt
+  provider: string | undefined
+  model: string | undefined
+  variant: string | undefined
+}
+
+export type RunSession = {
+  first: boolean
+  turns: Turn[]
+}
+
+function fileName(url: string, filename?: string) {
+  if (filename) {
+    return filename
+  }
+
+  try {
+    const next = new URL(url)
+    if (next.protocol !== "file:") {
+      return url
+    }
+
+    const name = next.pathname.split("/").at(-1)
+    if (name) {
+      return decodeURIComponent(name)
+    }
+  } catch {}
+
+  return url
+}
+
+function fileSource(
+  part: Extract<SessionMessages[number]["parts"][number], { type: "file" }>,
+  text: { start: number; end: number; value: string },
+) {
+  if (part.source) {
+    return {
+      ...structuredClone(part.source),
+      text,
+    }
+  }
+
+  return {
+    type: "file" as const,
+    path: part.filename ?? part.url,
+    text,
+  }
+}
+
+export function messagePrompt(msg: SessionMessages[number]): RunPrompt {
+  const parts: RunPrompt["parts"] = []
+  let text = msg.parts
+    .filter((part): part is Extract<SessionMessages[number]["parts"][number], { type: "text" }> => {
+      return part.type === "text" && !part.synthetic
+    })
+    .map((part) => part.text)
+    .join("")
+  let cursor = Bun.stringWidth(text)
+  const used: Array<{ start: number; end: number }> = []
+
+  const take = (value: string): { start: number; end: number; value: string } | undefined => {
+    let from = 0
+    while (true) {
+      const idx = text.indexOf(value, from)
+      if (idx === -1) {
+        return undefined
+      }
+
+      const start = Bun.stringWidth(text.slice(0, idx))
+      const end = start + Bun.stringWidth(value)
+      if (!used.some((item) => item.start < end && start < item.end)) {
+        return { start, end, value }
+      }
+
+      from = idx + value.length
+    }
+  }
+
+  const add = (value: string) => {
+    const gap = text ? " " : ""
+    const start = cursor + Bun.stringWidth(gap)
+    text += gap + value
+    const end = start + Bun.stringWidth(value)
+    cursor = end
+    return { start, end, value }
+  }
+
+  for (const part of msg.parts) {
+    if (part.type === "file") {
+      const next = part.source?.text ? structuredClone(part.source.text) : take("@" + fileName(part.url, part.filename))
+      const span = next ?? add("@" + fileName(part.url, part.filename))
+      used.push({ start: span.start, end: span.end })
+      parts.push({
+        type: "file",
+        mime: part.mime,
+        filename: part.filename,
+        url: part.url,
+        source: fileSource(part, span),
+      })
+      continue
+    }
+
+    if (part.type !== "agent") {
+      continue
+    }
+
+    const span = part.source ? structuredClone(part.source) : (take("@" + part.name) ?? add("@" + part.name))
+    used.push({ start: span.start, end: span.end })
+    parts.push({
+      type: "agent",
+      name: part.name,
+      source: span,
+    })
+  }
+
+  return { text, parts }
+}
+
+function turn(msg: SessionMessages[number]): Turn | undefined {
+  if (msg.info.role !== "user") {
+    return undefined
+  }
+
+  return {
+    prompt: messagePrompt(msg),
+    provider: msg.info.model.providerID,
+    model: msg.info.model.modelID,
+    variant: msg.info.model.variant,
+  }
+}
+
+export function createSession(messages: SessionMessages): RunSession {
+  return {
+    first: messages.length === 0,
+    turns: messages.flatMap((msg) => {
+      const item = turn(msg)
+      return item ? [item] : []
+    }),
+  }
+}
+
+export async function resolveSession(sdk: RunInput["sdk"], sessionID: string, limit = LIMIT): Promise<RunSession> {
+  const response = await sdk.session.messages({
+    sessionID,
+    limit,
+  })
+  return createSession(response.data ?? [])
+}
+
+export function sessionHistory(session: RunSession, limit = LIMIT): RunPrompt[] {
+  const out: RunPrompt[] = []
+
+  for (const turn of session.turns) {
+    if (!turn.prompt.text.trim()) {
+      continue
+    }
+
+    if (out[out.length - 1] && promptSame(out[out.length - 1], turn.prompt)) {
+      continue
+    }
+
+    out.push(promptCopy(turn.prompt))
+  }
+
+  return out.slice(-limit)
+}
+
+export function sessionVariant(session: RunSession, model: RunInput["model"]): string | undefined {
+  if (!model) {
+    return undefined
+  }
+
+  for (let idx = session.turns.length - 1; idx >= 0; idx -= 1) {
+    const turn = session.turns[idx]
+    if (turn.provider !== model.providerID || turn.model !== model.modelID) {
+      continue
+    }
+
+    return turn.variant
+  }
+
+  return undefined
+}