enhance pi configuration and skills

pi/extensions/openai-server-compaction/index.ts

@@ -0,0 +1,348 @@
+/**
+ * Main extension entrypoint.
+ *
+ * Wires together request patching, remote compaction, runtime state
+ * reconstruction, session lifecycle cleanup, and Pi-native HTTP request patching.
+ */
+import type { ExtensionAPI } from "@earendil-works/pi-coding-agent";
+import type { AgentMessage } from "@earendil-works/pi-agent-core";
+import { isRecord, loadConfig } from "./config.ts";
+import {
+  applyRemoteHistoryPayloadPatch,
+  extractResponsesReasoningConfig,
+  extractResponsesTextConfig,
+  isOpenAICodexResponsesModel,
+  looksLikeResponsesPayload,
+  messageMatchesModel,
+  modelKey,
+  supportsRemoteCompactionModel,
+  thinkingLevelToResponsesReasoning,
+} from "./openai.ts";
+import {
+  buildCompactionSummaryText,
+  buildRemoteCompactionDetails,
+  buildToolsPayload,
+  callRemoteCompactionEndpoint,
+  generateBestEffortLocalSummary,
+  messageToResponseItems,
+  messagesToResponseItems,
+  normalizeResponseItemsForPrompt,
+  reconstructRemoteCompactionStateFromBranch,
+} from "./remote-compaction.ts";
+import {
+  clearAllRuntimeState,
+  clearRemoteCompactionState,
+  clearResponsesRequestShapeState,
+  getRemoteCompactionState,
+  getResponsesRequestShapeState,
+  setRemoteCompactionState,
+  setResponsesRequestShapeState,
+} from "./state.ts";
+
+type TargetModel = Parameters<typeof modelKey>[0];
+
+type BranchEntry = {
+  type: string;
+  id: string;
+  details?: unknown;
+  message?: unknown;
+  thinkingLevel?: unknown;
+};
+
+type SessionContextLike = {
+  sessionManager: {
+    getSessionId(): string;
+    getBranch(): BranchEntry[];
+  };
+};
+
+function getSessionId(ctx: SessionContextLike): string {
+  return ctx.sessionManager.getSessionId();
+}
+
+function getBranchMessages(branchEntries: BranchEntry[]): AgentMessage[] {
+  return branchEntries.flatMap((entry) =>
+    entry.type === "message" && entry.message ? [entry.message as AgentMessage] : [],
+  );
+}
+
+function getBranchThinkingLevel(branchEntries: BranchEntry[]): string | undefined {
+  for (let index = branchEntries.length - 1; index >= 0; index--) {
+    const entry = branchEntries[index];
+    if (entry?.type !== "thinking_level_change") continue;
+    return typeof entry.thinkingLevel === "string" ? entry.thinkingLevel : undefined;
+  }
+  return undefined;
+}
+
+function clearSessionRuntimeState(sessionId: string | undefined): void {
+  clearRemoteCompactionState(sessionId);
+  clearResponsesRequestShapeState(sessionId);
+}
+
+function syncRemoteState(ctx: SessionContextLike): void {
+  const sessionId = getSessionId(ctx);
+  const branchEntries = ctx.sessionManager.getBranch() as Array<{
+    type: string;
+    id: string;
+    details?: unknown;
+    message?: AgentMessage;
+  }>;
+  const state = reconstructRemoteCompactionStateFromBranch({ branchEntries });
+  if (state) {
+    setRemoteCompactionState(sessionId, state);
+  } else {
+    clearRemoteCompactionState(sessionId);
+  }
+}
+
+function getMatchingRemoteState(
+  sessionId: string,
+  model: TargetModel | undefined,
+): ReturnType<typeof getRemoteCompactionState> {
+  if (!model) return undefined;
+  const remoteState = getRemoteCompactionState(sessionId);
+  return remoteState && remoteState.modelKey === modelKey(model) ? remoteState : undefined;
+}
+
+function extendRemoteHistoryIfCompatible(params: {
+  sessionId: string;
+  model: TargetModel | undefined;
+  message: AgentMessage;
+}): void {
+  const remoteState = getMatchingRemoteState(params.sessionId, params.model);
+  if (!remoteState || !params.model) return;
+  if (params.message.role === "assistant" && !messageMatchesModel(params.message, params.model)) {
+    return;
+  }
+
+  const items = messageToResponseItems(params.message);
+  if (items.length === 0) return;
+
+  setRemoteCompactionState(params.sessionId, {
+    ...remoteState,
+    explicitHistory: [...remoteState.explicitHistory, ...items],
+  });
+}
+
+function maybeNotifyRequestFeatures(params: {
+  notifiedModels: Set<string>;
+  hasUI: boolean;
+  notify: boolean;
+  ui: { notify(message: string, level: "info" | "warning"): void };
+  model: TargetModel;
+  features: string[];
+}): void {
+  if (!params.notify || !params.hasUI || params.features.length === 0) return;
+
+  const key = `${String(params.model.provider)}/${String(params.model.id)}`;
+  const noticeKey = `${key}:${params.features.join(",")}`;
+  if (params.notifiedModels.has(noticeKey)) return;
+
+  params.notifiedModels.add(noticeKey);
+  params.ui.notify(`OpenAI compaction active for ${key} (${params.features.join(", ")})`, "info");
+}
+
+export default function openaiServerCompactionExtension(pi: ExtensionAPI) {
+  const notifiedModels = new Set<string>();
+
+  pi.on("session_start", (_event, ctx) => {
+    const sessionId = getSessionId(ctx);
+    clearResponsesRequestShapeState(sessionId);
+    syncRemoteState(ctx);
+  });
+
+  const clearBeforeSessionChange = (_event: unknown, ctx: SessionContextLike): void => {
+    clearSessionRuntimeState(getSessionId(ctx));
+  };
+  pi.on("session_before_switch", clearBeforeSessionChange);
+  pi.on("session_before_fork", clearBeforeSessionChange);
+  pi.on("session_before_tree", clearBeforeSessionChange);
+
+  const syncAfterSessionChange = (_event: unknown, ctx: SessionContextLike): void => {
+    syncRemoteState(ctx);
+  };
+  pi.on("session_tree", syncAfterSessionChange);
+  pi.on("session_compact", syncAfterSessionChange);
+
+  pi.on("model_select", (_event, ctx) => {
+    clearResponsesRequestShapeState(getSessionId(ctx));
+  });
+
+  pi.on("session_shutdown", () => {
+    clearAllRuntimeState();
+  });
+
+  pi.on("session_before_compact", async (event, ctx) => {
+    const cfg = loadConfig(ctx.cwd);
+    const model = ctx.model;
+    if (!cfg.enabled || !model || !supportsRemoteCompactionModel(model)) return undefined;
+
+    const auth = await ctx.modelRegistry.getApiKeyAndHeaders(model);
+    if (!auth.ok || !auth.apiKey) return undefined;
+
+    const tools = buildToolsPayload(pi.getAllTools(), pi.getActiveTools());
+    const sessionId = getSessionId(ctx);
+    const branchEntries = event.branchEntries as BranchEntry[];
+    const remoteState = getMatchingRemoteState(sessionId, model);
+    const observedRequestShape = getResponsesRequestShapeState(sessionId);
+    const fullBranchMessages = getBranchMessages(branchEntries);
+    const responseItems = remoteState
+      ? remoteState.explicitHistory
+      : messagesToResponseItems(fullBranchMessages);
+    const promptResponseItems = normalizeResponseItemsForPrompt(responseItems, model);
+    const thinkingLevel = pi.getThinkingLevel();
+    const fallbackReasoning = model.reasoning
+      ? thinkingLevelToResponsesReasoning(thinkingLevel ?? getBranchThinkingLevel(branchEntries))
+      : undefined;
+    const reasoning = observedRequestShape?.reasoning ?? fallbackReasoning;
+    const text = observedRequestShape?.text;
+
+    const [localResult, remoteResult] = await Promise.allSettled([
+      generateBestEffortLocalSummary({
+        preparation: event.preparation,
+        messages: fullBranchMessages,
+        model,
+        apiKey: auth.apiKey,
+        headers: auth.headers,
+        customInstructions: event.customInstructions,
+        signal: event.signal,
+        thinkingLevel,
+        firstKeptEntryId: event.preparation.firstKeptEntryId,
+        tokensBefore: event.preparation.tokensBefore,
+      }),
+      callRemoteCompactionEndpoint({
+        model,
+        apiKey: auth.apiKey,
+        headers: auth.headers,
+        sessionId,
+        input: promptResponseItems,
+        instructions: ctx.getSystemPrompt(),
+        tools,
+        parallelToolCalls: true,
+        reasoning,
+        text,
+        signal: event.signal,
+      }),
+    ]);
+
+    if (remoteResult.status !== "fulfilled") {
+      if (localResult.status === "fulfilled") {
+        return { compaction: localResult.value };
+      }
+      if (!event.signal.aborted && ctx.hasUI) {
+        const message =
+          remoteResult.reason instanceof Error
+            ? remoteResult.reason.message
+            : String(remoteResult.reason);
+        ctx.ui.notify(
+          `OpenAI remote compaction failed; falling back to default compaction. ${message}`,
+          "warning",
+        );
+      }
+      return undefined;
+    }
+
+    const remoteDetails = buildRemoteCompactionDetails(
+      model,
+      remoteResult.value.output,
+      remoteResult.value.usage,
+    );
+    const localSummary =
+      localResult.status === "fulfilled"
+        ? localResult.value
+        : {
+            summary: buildCompactionSummaryText(model),
+            firstKeptEntryId: event.preparation.firstKeptEntryId,
+            tokensBefore: event.preparation.tokensBefore,
+          };
+
+    return {
+      compaction: {
+        summary: localSummary.summary,
+        firstKeptEntryId: localSummary.firstKeptEntryId,
+        tokensBefore: localSummary.tokensBefore,
+        details: {
+          ...(localSummary.details !== undefined
+            ? { localSummaryDetails: localSummary.details }
+            : {}),
+          remoteCompaction: remoteDetails,
+        },
+      },
+    };
+  });
+
+  pi.on("message_end", (event, ctx) => {
+    const sessionId = getSessionId(ctx);
+    const model = ctx.model;
+
+    extendRemoteHistoryIfCompatible({
+      sessionId,
+      model,
+      message: event.message,
+    });
+  });
+
+  pi.on("before_provider_request", (event, ctx) => {
+    const cfg = loadConfig(ctx.cwd);
+    if (!cfg.enabled) return undefined;
+
+    const model = ctx.model;
+    if (
+      !model ||
+      !supportsRemoteCompactionModel(model) ||
+      !isRecord(event.payload) ||
+      !looksLikeResponsesPayload(event.payload)
+    ) {
+      return undefined;
+    }
+
+    const sessionId = getSessionId(ctx);
+    setResponsesRequestShapeState(sessionId, {
+      updatedAt: Date.now(),
+      reasoning: extractResponsesReasoningConfig(event.payload),
+      text: extractResponsesTextConfig(event.payload),
+    });
+    const remoteState = getMatchingRemoteState(sessionId, model);
+
+    if (isOpenAICodexResponsesModel(model)) {
+      if (!remoteState) return undefined;
+      const payload = applyRemoteHistoryPayloadPatch({
+        payload: event.payload,
+        explicitHistory: normalizeResponseItemsForPrompt(
+          remoteState.explicitHistory,
+          model,
+        ) as unknown[],
+      });
+      maybeNotifyRequestFeatures({
+        notifiedModels,
+        hasUI: ctx.hasUI,
+        notify: cfg.notify,
+        ui: ctx.ui,
+        model,
+        features: ["remote_compaction_history"],
+      });
+      return payload;
+    }
+
+    if (!remoteState) return undefined;
+
+    const payload = applyRemoteHistoryPayloadPatch({
+      payload: event.payload,
+      explicitHistory: normalizeResponseItemsForPrompt(
+        remoteState.explicitHistory,
+        model,
+      ) as unknown[],
+    });
+    maybeNotifyRequestFeatures({
+      notifiedModels,
+      hasUI: ctx.hasUI,
+      notify: cfg.notify,
+      ui: ctx.ui,
+      model,
+      features: ["remote_compaction_history", "native_http"],
+    });
+
+    return payload;
+  });
+}