1  import type { ExtensionAPI } from "@mariozechner/pi-coding-agent";
  2  import { execSync } from "child_process";
  3  import * as fs from "fs";
  4  import * as os from "os";
  5  import * as path from "path";
  6  import * as readline from "readline";
  7  
  8  /**
  9   * Processes Pi agent JSONL session logs into readable text format
 10   * Output directory: ~/.pi/agent/pi-sessions-extracted/
 11   */
 12  
 13  interface SessionMeta {
 14    sessionId: string;
 15    startedAt: Date | null;
 16    cwd: string;
 17  }
 18  
 19  interface ToolCallInfo {
 20    name: string;
 21    cmd: string;
 22    include: boolean;
 23  }
 24  
 25  interface ToolFilter {
 26    mode: "all" | "includeOnly";
 27    includeNames: Set<string>;
 28    excludeNames: Set<string>;
 29  }
 30  
 31  interface ExportOptions {
 32    includeThinking: boolean;
 33    toolFilter: ToolFilter | null;
 34  }
 35  
 36  interface ConversationState {
 37    meta: SessionMeta;
 38    conversation: string[];
 39    pendingToolCalls: Map<string, ToolCallInfo>;
 40  }
 41  
 42  /**
 43   * Extract text from various content formats (string, array of content blocks)
 44   */
 45  function extractTextFromContent(content: unknown): string {
 46    if (content === null || content === undefined) {
 47      return "";
 48    }
 49    if (typeof content === "string") {
 50      return content.trim();
 51    }
 52    if (Array.isArray(content)) {
 53      const parts: string[] = [];
 54      for (const item of content) {
 55        if (typeof item !== "object" || item === null) {
 56          continue;
 57        }
 58        const itemObj = item as Record<string, unknown>;
 59        if (itemObj.type !== "text") {
 60          continue;
 61        }
 62        const text = itemObj.text;
 63        if (typeof text === "string" && text.trim()) {
 64          parts.push(text.trim());
 65        }
 66      }
 67      return parts.join("\n").trim();
 68    }
 69    return "";
 70  }
 71  
 72  /**
 73   * Parse ISO-ish timestamp string to Date
 74   */
 75  function parseIsoTimestamp(ts: unknown): Date | null {
 76    if (!ts || typeof ts !== "string") {
 77      return null;
 78    }
 79    try {
 80      let normalized = ts;
 81      if (ts.endsWith("Z")) {
 82        normalized = ts.slice(0, -1) + "+00:00";
 83      }
 84      const date = new Date(normalized);
 85      return isNaN(date.getTime()) ? null : date;
 86    } catch {
 87      return null;
 88    }
 89  }
 90  
 91  /**
 92   * Create a filesystem-safe slug from a string
 93   */
 94  function slug(s: string): string {
 95    const out: string[] = [];
 96    let prevUnderscore = false;
 97    for (const ch of s) {
 98      const isAlphaNum = /[a-zA-Z0-9]/.test(ch);
 99      const isAllowed = isAlphaNum || ch === "-" || ch === "_";
100      if (isAllowed) {
101        out.push(ch);
102        prevUnderscore = false;
103      } else if (!prevUnderscore) {
104        out.push("_");
105        prevUnderscore = true;
106      }
107    }
108    const result = out.join("").replace(/^_+|_+$/g, "");
109    return result || "unknown";
110  }
111  
112  /**
113   * Iterate through JSONL file line by line, yielding parsed objects
114   */
115  async function* iterJsonl(
116    filePath: string
117  ): AsyncGenerator<Record<string, unknown>> {
118    const fileStream = fs.createReadStream(filePath, { encoding: "utf-8" });
119    const rl = readline.createInterface({
120      input: fileStream,
121      crlfDelay: Infinity,
122    });
123  
124    for await (const line of rl) {
125      const trimmed = line.trim();
126      if (!trimmed) {
127        continue;
128      }
129      try {
130        const obj = JSON.parse(trimmed);
131        if (typeof obj === "object" && obj !== null && !Array.isArray(obj)) {
132          yield obj as Record<string, unknown>;
133        }
134      } catch {
135        // Skip malformed JSON lines
136        continue;
137      }
138    }
139  }
140  
141  function normalizeToolName(name: string): string {
142    return name.trim().toLowerCase();
143  }
144  
145  function createConversationState(): ConversationState {
146    return {
147      meta: {
148        sessionId: "",
149        startedAt: null,
150        cwd: "",
151      },
152      conversation: [],
153      pendingToolCalls: new Map(),
154    };
155  }
156  
157  function formatThinkingBlock(thinking: string): string {
158    return `[thinking]\n${thinking.trim()}\n[/thinking]`;
159  }
160  
161  function extractToolCommand(args: unknown): string {
162    if (typeof args === "object" && args !== null && !Array.isArray(args)) {
163      const argsObj = args as Record<string, unknown>;
164      const cmdVal = argsObj.command || argsObj.cmd;
165      if (typeof cmdVal === "string") {
166        return cmdVal;
167      }
168      try {
169        return JSON.stringify(args, Object.keys(args).sort());
170      } catch {
171        return String(args);
172      }
173    }
174    return String(args);
175  }
176  
177  function shouldIncludeTool(toolName: string, toolFilter: ToolFilter | null): boolean {
178    if (!toolFilter) {
179      return false;
180    }
181  
182    const normalized = normalizeToolName(toolName);
183    if (toolFilter.excludeNames.has(normalized)) {
184      return false;
185    }
186    if (toolFilter.mode === "includeOnly") {
187      return toolFilter.includeNames.has(normalized);
188    }
189    return true;
190  }
191  
192  function applyMessageRecord(
193    rec: Record<string, unknown>,
194    state: ConversationState,
195    options: ExportOptions
196  ): void {
197    const rtype = rec.type;
198  
199    if (rtype === "session" && !state.meta.sessionId) {
200      state.meta.sessionId = String(rec.id || "");
201      state.meta.startedAt = parseIsoTimestamp(rec.timestamp);
202      state.meta.cwd = String(rec.cwd || "");
203      return;
204    }
205  
206    if (rtype !== "message") {
207      return;
208    }
209  
210    const msg = rec.message;
211    if (typeof msg !== "object" || msg === null || Array.isArray(msg)) {
212      return;
213    }
214  
215    const msgObj = msg as Record<string, unknown>;
216    const role = msgObj.role;
217  
218    if (role === "user") {
219      const text = extractTextFromContent(msgObj.content);
220      if (text) {
221        state.conversation.push(`USER: ${text}`);
222      }
223      return;
224    }
225  
226    if (role === "assistant") {
227      appendAssistantMessage(msgObj, state, options);
228      return;
229    }
230  
231    if (role === "toolResult") {
232      appendToolResultMessage(msgObj, state, options.toolFilter);
233    }
234  }
235  
236  function appendAssistantMessage(
237    msgObj: Record<string, unknown>,
238    state: ConversationState,
239    options: ExportOptions
240  ): void {
241    const content = msgObj.content;
242    const textParts: string[] = [];
243    const toolParts: string[] = [];
244  
245    if (typeof content === "string") {
246      if (content.trim()) {
247        textParts.push(content.trim());
248      }
249    } else if (Array.isArray(content)) {
250      for (const item of content) {
251        if (typeof item !== "object" || item === null) {
252          continue;
253        }
254        const itemObj = item as Record<string, unknown>;
255        const itemType = itemObj.type;
256  
257        if (itemType === "text") {
258          const text = itemObj.text;
259          if (typeof text === "string" && text.trim()) {
260            textParts.push(text.trim());
261          }
262          continue;
263        }
264  
265        if (itemType === "thinking") {
266          const thinking = itemObj.thinking;
267          if (
268            options.includeThinking &&
269            typeof thinking === "string" &&
270            thinking.trim()
271          ) {
272            textParts.push(formatThinkingBlock(thinking));
273          }
274          continue;
275        }
276  
277        if (itemType !== "toolCall") {
278          continue;
279        }
280  
281        const toolName = String(itemObj.name || "tool");
282        const toolId = String(itemObj.id || "");
283        const cmd = extractToolCommand(itemObj.arguments);
284        const include = shouldIncludeTool(toolName, options.toolFilter);
285  
286        if (toolId) {
287          state.pendingToolCalls.set(toolId, { name: toolName, cmd, include });
288        }
289        if (include) {
290          toolParts.push(`[tool:${toolName}] ${cmd}`.trim());
291        }
292      }
293    }
294  
295    let fullText = textParts.filter(Boolean).join("\n").trim();
296    if (toolParts.length > 0) {
297      fullText = (fullText ? `${fullText}\n` : "") + toolParts.join("\n");
298    }
299  
300    if (fullText) {
301      state.conversation.push(`ASSISTANT: ${fullText}`);
302    }
303  }
304  
305  function appendToolResultMessage(
306    msgObj: Record<string, unknown>,
307    state: ConversationState,
308    toolFilter: ToolFilter | null
309  ): void {
310    const toolName = String(msgObj.toolName || msgObj.tool_name || "tool");
311    const toolCallId = String(msgObj.toolCallId || msgObj.tool_call_id || "");
312  
313    let label = toolName;
314    let include = shouldIncludeTool(toolName, toolFilter);
315  
316    if (toolCallId && state.pendingToolCalls.has(toolCallId)) {
317      const pending = state.pendingToolCalls.get(toolCallId)!;
318      state.pendingToolCalls.delete(toolCallId);
319      label = pending.name || toolName;
320      include = pending.include;
321      if (pending.cmd) {
322        label = `${label}: ${pending.cmd}`;
323      }
324    }
325  
326    if (!include) {
327      return;
328    }
329  
330    const isError = Boolean(msgObj.isError || msgObj.is_error);
331    const out = extractTextFromContent(msgObj.content);
332    if (out && out !== "(no content)") {
333      state.conversation.push(
334        `SYSTEM [${label} output${isError ? " ERROR" : ""}]:\n${out}`
335      );
336    }
337  }
338  
339  /**
340   * Extract conversation messages from a Pi session JSONL file
341   */
342  async function extractConversation(
343    jsonlFile: string,
344    options: ExportOptions
345  ): Promise<{ meta: SessionMeta; conversation: string[] }> {
346    const state = createConversationState();
347  
348    for await (const rec of iterJsonl(jsonlFile)) {
349      applyMessageRecord(rec, state, options);
350    }
351  
352    return { meta: state.meta, conversation: state.conversation };
353  }
354  
355  /**
356   * Extract conversation messages from the *currently selected branch* (root → current leaf)
357   * using the in-memory SessionManager tree state.
358   *
359   * This intentionally includes *all* turns on the branch, including turns that may no longer
360   * be shown in the TUI after compaction.
361   */
362  function extractConversationFromBranch(
363    sessionManager: any,
364    options: ExportOptions
365  ): { meta: SessionMeta; conversation: string[]; leafId: string | null } {
366    const state = createConversationState();
367    const header =
368      typeof sessionManager.getHeader === "function"
369        ? sessionManager.getHeader()
370        : null;
371  
372    if (header && typeof header === "object") {
373      const headerObj = header as Record<string, unknown>;
374      state.meta.sessionId = typeof headerObj.id === "string" ? headerObj.id : "";
375      state.meta.startedAt = parseIsoTimestamp(headerObj.timestamp);
376      if (typeof headerObj.cwd === "string") {
377        state.meta.cwd = headerObj.cwd;
378      }
379    }
380  
381    if (!state.meta.cwd && typeof sessionManager.getCwd === "function") {
382      state.meta.cwd = String(sessionManager.getCwd() || "");
383    }
384  
385    const leafId =
386      typeof sessionManager.getLeafId === "function"
387        ? (sessionManager.getLeafId() as string | null)
388        : null;
389  
390    const branchEntries: unknown[] =
391      leafId && typeof sessionManager.getBranch === "function"
392        ? (sessionManager.getBranch(leafId) as unknown[])
393        : [];
394  
395    for (const entry of branchEntries) {
396      if (typeof entry !== "object" || entry === null) {
397        continue;
398      }
399      applyMessageRecord(entry as Record<string, unknown>, state, options);
400    }
401  
402    return { meta: state.meta, conversation: state.conversation, leafId };
403  }
404  
405  function buildMarkdownContent(
406    meta: SessionMeta,
407    sessionFile: string,
408    conversation: string[],
409    lastTurns: number | null,
410    branchInfo?: { leafId: string | null }
411  ): { content: string; filename: string } | null {
412    const finalConversation =
413      lastTurns && lastTurns > 0 ? sliceLastNTurns(conversation, lastTurns) : conversation;
414  
415    if (finalConversation.length === 0) {
416      return null;
417    }
418  
419    const project = meta.cwd ? path.basename(meta.cwd) : "unknown";
420    const projectSlug = slug(project);
421  
422    const started = meta.startedAt || new Date();
423    const stamp = formatTimestamp(started);
424    const sid = (meta.sessionId || "unknown").slice(0, 8);
425    const filename = `${projectSlug}_pi_${stamp}_${sid}.md`;
426  
427    const lines: string[] = [];
428    lines.push("PI SESSION (processed)");
429    if (branchInfo) {
430      lines.push("mode: branch");
431      lines.push(`leaf: ${branchInfo.leafId === null ? "null" : branchInfo.leafId}`);
432    }
433    if (meta.sessionId) {
434      lines.push(`id: ${meta.sessionId}`);
435    }
436    if (meta.startedAt) {
437      lines.push(`started: ${meta.startedAt.toISOString()}`);
438    }
439    if (meta.cwd) {
440      lines.push(`cwd: ${meta.cwd}`);
441    }
442    lines.push(`source: ${sessionFile}`);
443    if (lastTurns && lastTurns > 0) {
444      lines.push(`turns: last ${lastTurns}`);
445    }
446    lines.push("");
447  
448    for (const msg of finalConversation) {
449      lines.push(msg.trimEnd());
450      lines.push("");
451    }
452  
453    return { content: lines.join("\n"), filename };
454  }
455  
456  /**
457   * Generate markdown content from the current branch without writing to file
458   */
459  function generateMarkdownFromBranch(
460    sessionManager: any,
461    sessionFile: string,
462    options: ExportOptions,
463    lastTurns: number | null = null
464  ): { content: string; filename: string } | null {
465    const { meta, conversation, leafId } = extractConversationFromBranch(
466      sessionManager,
467      options
468    );
469    return buildMarkdownContent(meta, sessionFile, conversation, lastTurns, {
470      leafId,
471    });
472  }
473  
474  /**
475   * Generate markdown content from full session file without writing to file
476   */
477  async function generateMarkdownFromSession(
478    jsonlFile: string,
479    options: ExportOptions,
480    lastTurns: number | null = null
481  ): Promise<{ content: string; filename: string } | null> {
482    const { meta, conversation } = await extractConversation(jsonlFile, options);
483    return buildMarkdownContent(meta, jsonlFile, conversation, lastTurns);
484  }
485  
486  /**
487   * Copy text to clipboard (macOS)
488   */
489  function copyToClipboard(text: string): void {
490    execSync("pbcopy", { input: text });
491  }
492  
493  /**
494   * Format date as YYYYMMDD_HHMMSS
495   */
496  function formatTimestamp(date: Date): string {
497    const pad = (n: number) => String(n).padStart(2, "0");
498    return (
499      `${date.getFullYear()}${pad(date.getMonth() + 1)}${pad(date.getDate())}_` +
500      `${pad(date.getHours())}${pad(date.getMinutes())}${pad(date.getSeconds())}`
501    );
502  }
503  
504  /**
505   * Slice a flattened conversation array to the last N turns
506   *
507   * A turn is defined as a unit of [USER message -> ASSISTANT message], including any
508   * SYSTEM tool outputs that occur in between, until the next USER message begins.
509   */
510  function sliceLastNTurns(conversation: string[], turns: number): string[] {
511    if (!Number.isFinite(turns) || turns <= 0) {
512      return conversation;
513    }
514  
515    const userMessageIndices: number[] = [];
516    for (let i = 0; i < conversation.length; i++) {
517      if (conversation[i].startsWith("USER:")) {
518        userMessageIndices.push(i);
519      }
520    }
521  
522    if (userMessageIndices.length === 0) {
523      return conversation;
524    }
525  
526    const startIndex =
527      userMessageIndices.length <= turns
528        ? 0
529        : userMessageIndices[userMessageIndices.length - turns];
530  
531    return conversation.slice(startIndex);
532  }
533  
534  function parseToolFilter(tokensLower: string[]): ToolFilter {
535    const includeNames = new Set<string>();
536    const excludeNames = new Set<string>();
537  
538    for (const token of tokensLower) {
539      if (!["+", "-"].includes(token[0]) || token.length < 2) {
540        continue;
541      }
542  
543      const name = normalizeToolName(token.slice(1));
544      if (!name) {
545        continue;
546      }
547  
548      if (token.startsWith("+")) {
549        includeNames.add(name);
550      } else {
551        excludeNames.add(name);
552      }
553    }
554  
555    return {
556      mode: includeNames.size > 0 ? "includeOnly" : "all",
557      includeNames,
558      excludeNames,
559    };
560  }
561  
562  function hasToolFilterTokens(tokensLower: string[]): boolean {
563    return tokensLower.some((token) => /^[+-].+/.test(token));
564  }
565  
566  function describeToolFilter(toolFilter: ToolFilter | null): string | null {
567    if (!toolFilter) {
568      return null;
569    }
570  
571    const includes = [...toolFilter.includeNames].sort().map((name) => `+${name}`);
572    const excludes = [...toolFilter.excludeNames].sort().map((name) => `-${name}`);
573  
574    if (toolFilter.mode === "includeOnly") {
575      return ["tool calls", ...includes, ...excludes].join(" ");
576    }
577    if (excludes.length > 0) {
578      return ["tool calls", ...excludes].join(" ");
579    }
580    return "tool calls";
581  }
582  
583  const OUTPUT_DIR = path.join(os.homedir(), ".pi", "agent", "pi-sessions-extracted");
584  
585  export default function (pi: ExtensionAPI) {
586    pi.registerCommand("md", {
587      description:
588        "Export current session as markdown (current /tree branch) on clipboard or to file. Tool calls and thinking blocks are excluded by default. Use '/md tc' to include tool calls, optionally filtered with exact tool names like '/md tc -bash -read' or '/md tc +ask'. Use '/md t' to include thinking blocks. Use '/md all' for full file. Pass a number (e.g. '/md 2' or '/md tc t 2') to export only the last N turns.",
589      handler: async (args, ctx) => {
590        const sessionFile = ctx.sessionManager.getSessionFile();
591  
592        if (!sessionFile) {
593          ctx.ui.notify("No session file (ephemeral session)", "error");
594          return;
595        }
596  
597        const argsTrimmed = args.trim();
598        const tokens = argsTrimmed ? argsTrimmed.split(/\s+/).filter(Boolean) : [];
599        const tokensLower = tokens.map((t) => t.toLowerCase());
600  
601        const includeThinking =
602          tokensLower.includes("t") ||
603          tokensLower.includes("think") ||
604          tokensLower.includes("thinking");
605  
606        const includeToolCalls = tokensLower.includes("tc");
607        const exportAll = tokensLower.includes("all") || tokensLower.includes("file");
608        const toolFilterTokensPresent = hasToolFilterTokens(tokensLower);
609  
610        if (toolFilterTokensPresent && !includeToolCalls) {
611          ctx.ui.notify("Tool filters require 'tc' (e.g. /md tc -bash or /md tc +ask)", "error");
612          return;
613        }
614  
615        if (tokensLower.includes("+all") || tokensLower.includes("-all")) {
616          ctx.ui.notify("'+all' and '-all' are not supported; use '/md tc' or '/md tc +tool'", "error");
617          return;
618        }
619  
620        const toolFilter = includeToolCalls ? parseToolFilter(tokensLower) : null;
621        const options: ExportOptions = {
622          includeThinking,
623          toolFilter,
624        };
625  
626        const turnsToken = tokens.find((t) => /^\d+$/.test(t)) || null;
627        const lastTurns = turnsToken ? parseInt(turnsToken, 10) : null;
628        if (turnsToken && (!Number.isFinite(lastTurns) || lastTurns < 1)) {
629          ctx.ui.notify("Turn limit must be >= 1 (e.g. /md 2)", "error");
630          return;
631        }
632  
633        const turnSuffix = lastTurns ? ` (last ${lastTurns} turns)` : "";
634        const extras: string[] = [];
635        if (includeThinking) {
636          extras.push("thinking");
637        }
638        const toolFilterDescription = describeToolFilter(toolFilter);
639        if (toolFilterDescription) {
640          extras.push(toolFilterDescription);
641        }
642        const extrasSuffix = extras.length > 0 ? ` (with ${extras.join(" + ")})` : "";
643        const title = `Export session${turnSuffix}${extrasSuffix} as Markdown`;
644        const choice = await ctx.ui.select(`${title}\n\nSelect export method:`, [
645          "Copy to clipboard",
646          `Save to .md file in ${OUTPUT_DIR}/`,
647        ]);
648  
649        if (!choice) {
650          return;
651        }
652  
653        try {
654          const result = exportAll
655            ? await generateMarkdownFromSession(sessionFile, options, lastTurns)
656            : generateMarkdownFromBranch(ctx.sessionManager, sessionFile, options, lastTurns);
657  
658          if (!result) {
659            const mode = exportAll ? "session file" : "current branch";
660            ctx.ui.notify(`No meaningful conversation found in ${mode}`, "error");
661            return;
662          }
663  
664          const suffix = extrasSuffix;
665          const mode = exportAll ? " (full file)" : " (branch)";
666  
667          if (choice === "Copy to clipboard") {
668            copyToClipboard(result.content);
669            ctx.ui.notify(`Copied to clipboard${suffix}${mode}`, "success");
670            return;
671          }
672  
673          const outputFile = path.join(OUTPUT_DIR, result.filename);
674          fs.mkdirSync(OUTPUT_DIR, { recursive: true });
675          fs.writeFileSync(outputFile, result.content, "utf-8");
676          ctx.ui.notify(`Saved${suffix}${mode}: ${outputFile}`, "success");
677        } catch (err) {
678          const errMsg = err instanceof Error ? err.message : String(err);
679          ctx.ui.notify(`Export failed: ${errMsg}`, "error");
680        }
681      },
682    });
683  }