pi-extensions-tool-display 0.1.1 → 0.2.1

package/src/thinking-label.ts

@@ -0,0 +1,343 @@
+import type { ExtensionAPI, ExtensionContext } from "@earendil-works/pi-coding-agent";
+import { isRecord } from "./tool-metadata.js";
+import { onReloadShutdown } from "./extension-lifecycle.js";
+
+interface ThemeLike {
+  fg(color: string, text: string): string;
+}
+
+interface AssistantMessageLike {
+  role?: unknown;
+  api?: unknown;
+  content?: unknown;
+}
+
+const THINKING_CHAT_PREFIX = "Thinking: ";
+const THINKING_LABEL_PREFIX_PATTERN = /^(?:thinking:\s*)+/i;
+const LEADING_ANSI_FRAGMENT_PATTERN = /^(?:\s*;?\d{1,3}(?:;\d{1,3})*m)+\s*/;
+const MAX_THINKING_CONTENT_DEPTH = 16;
+
+const registeredThinkingApis = new WeakSet<ExtensionAPI>();
+
+const OPENAI_REASONING_APIS = new Set([
+  "openai-completions",
+  "openai-responses",
+  "openai-codex-responses",
+]);
+
+const ANTHROPIC_REASONING_APIS = new Set([
+  "anthropic-messages",
+  "anthropic-responses",
+  "anthropic-completions",
+]);
+
+function normalizeApiName(api: unknown): string | undefined {
+  if (typeof api !== "string") {
+    return undefined;
+  }
+
+  const normalized = api.trim().toLowerCase();
+  return normalized.length > 0 ? normalized : undefined;
+}
+
+function shouldPrefixThinkingForApi(api: unknown): boolean {
+  const normalizedApi = normalizeApiName(api);
+  if (!normalizedApi) {
+    return true;
+  }
+
+  if (OPENAI_REASONING_APIS.has(normalizedApi)) {
+    return true;
+  }
+
+  if (ANTHROPIC_REASONING_APIS.has(normalizedApi)) {
+    return true;
+  }
+
+  if (normalizedApi.startsWith("anthropic-")) {
+    return true;
+  }
+
+  // Keep OpenAI handling explicit to avoid applying this formatter to
+  // unrelated OpenAI transport APIs that may not emit thinking blocks.
+  if (normalizedApi.startsWith("openai-")) {
+    return false;
+  }
+
+  // For non-OpenAI providers, apply the prefix when thinking blocks exist.
+  return true;
+}
+
+function stripAnsi(text: string): string {
+  return text.replace(/\x1b\[[0-9;]*m/g, "");
+}
+
+function stripLeadingAnsiFragments(text: string): string {
+  let current = text;
+  while (true) {
+    const next = current.replace(LEADING_ANSI_FRAGMENT_PATTERN, "");
+    if (next === current) {
+      return current;
+    }
+    current = next;
+  }
+}
+
+function stripThinkingPresentationArtifacts(text: string): string {
+  let current = stripAnsi(text);
+  let removedThinkingLabel = false;
+
+  while (true) {
+    const withoutLabel = current
+      .replace(THINKING_LABEL_PREFIX_PATTERN, "")
+      .trimStart();
+    if (withoutLabel !== current) {
+      current = withoutLabel;
+      removedThinkingLabel = true;
+      continue;
+    }
+
+    const withoutAnsiFragments = stripLeadingAnsiFragments(current).trimStart();
+    if (withoutAnsiFragments !== current) {
+      const fragmentExposedAnotherLabel =
+        withoutAnsiFragments.replace(THINKING_LABEL_PREFIX_PATTERN, "").trimStart() !==
+        withoutAnsiFragments;
+
+      if (removedThinkingLabel || fragmentExposedAnotherLabel) {
+        current = withoutAnsiFragments;
+        continue;
+      }
+    }
+
+    return current;
+  }
+}
+
+function formatThinkingLabel(theme: ThemeLike | undefined, thinkingText: string): string {
+  if (!theme) {
+    return `${THINKING_CHAT_PREFIX}${thinkingText}`;
+  }
+
+  const label = theme.fg("accent", THINKING_CHAT_PREFIX.trimEnd());
+  const body = theme.fg("thinkingText", thinkingText);
+  return `${label} ${body}`;
+}
+
+function prefixThinkingLine(text: string, theme: ThemeLike | undefined): string {
+  const normalizedThinking = stripThinkingPresentationArtifacts(text).trim();
+  if (!normalizedThinking) {
+    return text;
+  }
+
+  return formatThinkingLabel(theme, normalizedThinking);
+}
+
+function normalizeThinkingLineForContext(text: string): string {
+  return stripThinkingPresentationArtifacts(text);
+}
+
+function isThinkingBlock(value: unknown): value is Record<string, unknown> & {
+  type: "thinking";
+  thinking: string;
+} {
+  if (!isRecord(value)) {
+    return false;
+  }
+
+  return value.type === "thinking" && typeof value.thinking === "string";
+}
+
+function mapThinkingContentArray(
+  content: unknown[],
+  mapThinkingText: (text: string) => string,
+  depth = 0,
+  seen: WeakSet<object> = new WeakSet<object>(),
+): { content: unknown[]; changed: boolean } {
+  if (depth > MAX_THINKING_CONTENT_DEPTH || seen.has(content)) {
+    return { content, changed: false };
+  }
+
+  seen.add(content);
+  let changed = false;
+  const nextContent = content.map((block) => {
+    if (Array.isArray(block)) {
+      const nested = mapThinkingContentArray(
+        block,
+        mapThinkingText,
+        depth + 1,
+        seen,
+      );
+      if (nested.changed) {
+        changed = true;
+        return nested.content;
+      }
+      return block as unknown[];
+    }
+
+    if (!isThinkingBlock(block)) {
+      return block;
+    }
+
+    const nextThinking = mapThinkingText(block.thinking);
+    if (nextThinking === block.thinking) {
+      return block;
+    }
+
+    changed = true;
+    return { ...block, thinking: nextThinking };
+  });
+
+  return { content: changed ? nextContent : content, changed };
+}
+
+function withThinkingLabelsForDisplay(
+  content: unknown,
+  theme: ThemeLike | undefined,
+): unknown {
+  if (!Array.isArray(content)) {
+    return content;
+  }
+
+  const mapped = mapThinkingContentArray(content, (thinking) =>
+    prefixThinkingLine(thinking, theme),
+  );
+
+  return mapped.changed ? mapped.content : content;
+}
+
+function sanitizeThinkingBlocksForContext(message: AssistantMessageLike): AssistantMessageLike {
+  if (!Array.isArray(message.content)) {
+    return message;
+  }
+
+  const mapped = mapThinkingContentArray(
+    message.content,
+    normalizeThinkingLineForContext,
+  );
+
+  return mapped.changed ? { ...message, content: mapped.content } : message;
+}
+
+function sanitizeContextMessages(messages: unknown): unknown {
+  if (!Array.isArray(messages)) {
+    return messages;
+  }
+
+  const messageList = messages as unknown[];
+  let changed = false;
+  const nextMessages = messageList.map((message) => {
+    if (!isRecord(message) || message.role !== "assistant") {
+      return message;
+    }
+
+    const sanitized = sanitizeThinkingBlocksForContext(message as AssistantMessageLike);
+    if (sanitized !== message) {
+      changed = true;
+      return sanitized;
+    }
+
+    return message;
+  });
+
+  return changed ? nextMessages : messageList;
+}
+
+function prefixThinkingBlocksForDisplay(
+  message: AssistantMessageLike,
+  theme: ThemeLike | undefined,
+): void {
+  if (!shouldPrefixThinkingForApi(message.api)) {
+    return;
+  }
+
+  const displayContent = withThinkingLabelsForDisplay(message.content, theme);
+  if (displayContent !== message.content) {
+    message.content = displayContent;
+  }
+}
+
+function extractAssistantMessage(event: unknown): AssistantMessageLike | undefined {
+  if (!isRecord(event)) {
+    return undefined;
+  }
+
+  const maybeMessage = event.message;
+  if (!isRecord(maybeMessage)) {
+    return undefined;
+  }
+
+  if (maybeMessage.role !== "assistant") {
+    return undefined;
+  }
+
+  return maybeMessage as AssistantMessageLike;
+}
+
+function processThinkingEvent(
+  event: unknown,
+  ctx: ExtensionContext | undefined,
+  notifyPrefix: string,
+): void {
+  try {
+    const message = extractAssistantMessage(event);
+    if (!message) {
+      return;
+    }
+    prefixThinkingBlocksForDisplay(message, ctx?.ui?.theme);
+  } catch (error) {
+    const errorMessage = error instanceof Error ? error.message : "Unknown error";
+    ctx?.ui?.notify(`${notifyPrefix}: ${errorMessage}`, "warning");
+  }
+}
+
+function handleThinkingMessageUpdateEvent(event: unknown, ctx: ExtensionContext | undefined): void {
+  // Render-only labeling: update the transient message_update payload while
+  // leaving canonical session/LLM context content untouched.
+  processThinkingEvent(event, ctx, "Thinking label formatting failed");
+}
+
+function handleThinkingMessageEndEvent(event: unknown, ctx: ExtensionContext | undefined): void {
+  // Persist themed labels on final assistant messages so the label remains
+  // visible after streaming ends and across session reloads.
+  // Context sanitization strips these presentation artifacts before each LLM call.
+  processThinkingEvent(event, ctx, "Thinking label finalization failed");
+}
+
+function handleThinkingContextEvent(event: unknown, ctx: ExtensionContext | undefined): void {
+  try {
+    if (!isRecord(event) || !Array.isArray(event.messages)) {
+      return;
+    }
+
+    const sanitizedMessages = sanitizeContextMessages(event.messages);
+    if (sanitizedMessages !== event.messages && Array.isArray(sanitizedMessages)) {
+      event.messages.splice(0, event.messages.length, ...(sanitizedMessages as unknown[]));
+    }
+  } catch (error) {
+    const message = error instanceof Error ? error.message : "Unknown error";
+    ctx?.ui?.notify(`Thinking context sanitization failed: ${message}`, "warning");
+  }
+}
+
+export function registerThinkingLabeling(pi: ExtensionAPI): void {
+  if (registeredThinkingApis.has(pi)) {
+    return;
+  }
+  registeredThinkingApis.add(pi);
+
+  onReloadShutdown(pi, () => {
+    registeredThinkingApis.delete(pi);
+  });
+
+  pi.on("message_update", async (event, ctx) => {
+    handleThinkingMessageUpdateEvent(event, ctx);
+  });
+
+  pi.on("message_end", async (event, ctx) => {
+    handleThinkingMessageEndEvent(event, ctx);
+  });
+
+  pi.on("context", async (event, ctx) => {
+    handleThinkingContextEvent(event, ctx);
+  });
+}

package/src/tool-metadata.ts

@@ -0,0 +1,139 @@
+export interface PromptMetadata {
+	promptSnippet?: string;
+	promptGuidelines?: string[];
+}
+
+const MCP_DESCRIPTION_PATTERN = /\bmcp\b/i;
+const MCP_ADAPTER_SOURCE_PATTERN = /(?:^|[/\\@_-])(?:pi-)?mcp(?:[/\\@_-]|$)|pi-mcp-adapter|mcp-adapter/i;
+const MAX_PROMPT_SNIPPET_LENGTH = 120;
+
+export const MCP_PROXY_PROMPT_SNIPPET = "Discover, inspect, and call MCP tools across configured servers";
+export const MCP_PROXY_PROMPT_GUIDELINES = [
+	"Use mcp for MCP discovery first: search by capability, describe one exact tool, then call it.",
+] as const;
+
+export function isRecord(value: unknown): value is Record<string, unknown> {
+	return typeof value === "object" && value !== null && !Array.isArray(value);
+}
+
+export function toRecord(value: unknown): Record<string, unknown> {
+	if (!isRecord(value)) {
+		return {};
+	}
+
+	return value;
+}
+
+export function getTextField(value: unknown, field: string): string | undefined {
+	const record = toRecord(value);
+	const raw = record[field];
+	return typeof raw === "string" && raw.trim().length > 0 ? raw.trim() : undefined;
+}
+
+function normalizeInlineText(value: string): string {
+	return value.trim().replace(/\s+/g, " ");
+}
+
+function trimPromptSnippet(value: string): string {
+	if (value.length <= MAX_PROMPT_SNIPPET_LENGTH) {
+		return value;
+	}
+
+	const truncated = value.slice(0, MAX_PROMPT_SNIPPET_LENGTH).trimEnd();
+	return `${truncated.replace(/[\s.,;:!?-]+$/u, "")}…`;
+}
+
+export function buildPromptSnippetFromDescription(description: string | undefined, fallback: string): string {
+	const normalizedDescription = normalizeInlineText(description || "");
+	const normalizedFallback = normalizeInlineText(fallback);
+	const base = normalizedDescription || normalizedFallback;
+	const firstSentence = base.split(/(?<=[.!?])\s+/u, 1)[0] ?? base;
+	const withoutSentencePunctuation = firstSentence.replace(/[.!?]+$/u, "").trim();
+	return trimPromptSnippet(withoutSentencePunctuation || base);
+}
+
+export function extractPromptMetadata(tool: unknown): PromptMetadata {
+	const source = toRecord(tool);
+	const promptSnippet =
+		typeof source.promptSnippet === "string" && source.promptSnippet.trim().length > 0
+			? source.promptSnippet
+			: undefined;
+	const promptGuidelines = Array.isArray(source.promptGuidelines)
+		? source.promptGuidelines.filter(
+				(guideline): guideline is string =>
+					typeof guideline === "string" && guideline.trim().length > 0,
+			)
+		: undefined;
+
+	return {
+		promptSnippet,
+		promptGuidelines:
+			promptGuidelines && promptGuidelines.length > 0
+				? [...promptGuidelines]
+				: undefined,
+	};
+}
+
+function hasMcpSourceInfo(value: unknown): boolean {
+	const sourceInfo = toRecord(value);
+	for (const [key, raw] of Object.entries(sourceInfo)) {
+		if (typeof raw !== "string" || raw.trim().length === 0) {
+			continue;
+		}
+
+		const normalizedKey = key.toLowerCase();
+		const normalizedValue = raw.trim();
+		if (["source", "type", "kind", "origin"].includes(normalizedKey) && normalizedValue.toLowerCase() === "mcp") {
+			return true;
+		}
+		if (MCP_ADAPTER_SOURCE_PATTERN.test(normalizedValue)) {
+			return true;
+		}
+	}
+
+	return false;
+}
+
+export function isMcpToolCandidate(tool: unknown): boolean {
+	if (!tool || typeof tool !== "object") {
+		return false;
+	}
+
+	const record = tool as Record<string, unknown>;
+	const name = typeof record.name === "string" ? record.name : "";
+	const description = typeof record.description === "string" ? record.description : "";
+	const label = typeof record.label === "string" ? record.label : "";
+
+	if (name === "mcp") {
+		return true;
+	}
+	if (MCP_DESCRIPTION_PATTERN.test(description) || MCP_DESCRIPTION_PATTERN.test(label)) {
+		return true;
+	}
+	if (hasMcpSourceInfo(record.sourceInfo)) {
+		return true;
+	}
+	if (/^mcp[_-]/i.test(name) || /_mcp$/i.test(name)) {
+		return true;
+	}
+	if (name.includes(":")) {
+		return true;
+	}
+	if (/^ctx_/i.test(name)) {
+		return true;
+	}
+
+	const params = record.parameters;
+	if (params && typeof params === "object") {
+		const parameterRecord = params as Record<string, unknown>;
+		if (
+			"mcpServer" in parameterRecord ||
+			"serverUrl" in parameterRecord ||
+			"server_name" in parameterRecord
+		) {
+			return true;
+		}
+	}
+
+	return false;
+}