@springbrand/agent-runtime 0.1.0

package/src/pi/assembly/context.ts

@@ -0,0 +1,510 @@
+import {
+  buildSessionContext,
+  compact as compactPiTranscript,
+  DEFAULT_COMPACTION_SETTINGS,
+  estimateContextTokens,
+  prepareCompaction,
+  type AgentMessage,
+  type CompactionEntry,
+  type SessionTreeEntry,
+} from "@earendil-works/pi-agent-core";
+import type {
+  AssistantMessage,
+  Model,
+  Models,
+  RetryPolicy,
+  Usage,
+} from "@earendil-works/pi-ai";
+import { estimateStringTokens } from "agents/experimental/memory/utils";
+import type { RuntimeMemoryPort } from "../../kernel/bindings";
+import type { RuntimeDegradation } from "../../kernel/degradation";
+import type { RuntimeMemoryProfile } from "../../kernel/profile";
+import type { PiExtensionContextContribution } from "./extensions";
+
+/**
+ * 压缩这一次 LLM 调用的重试策略。
+ *
+ * 为什么必须显式传:pi 的 `retryAssistantCall` 是 opt-in 的 —— d.ts 原话
+ * 「When `policy` is undefined or disabled, the first response is returned
+ * unchanged」。升到 0.83 不传这个参数,行为与 0.80.6 完全一样(零重试)。
+ *
+ * 为什么压缩值得重试:压缩失败不是「少一次摘要」,是**下一轮直接撞 provider
+ * 的上下文窗口上限、整个 turn 挂掉**。而它跑在 Workers 里,瞬时网络错误是常态。
+ *
+ * 为什么只有 2 次、退避从 500ms 起:压缩发生在 turn 的关键路径上,DO 的墙钟
+ * 是真金白银(等 LLM 不占 CPU 配额但占用户的等待)。2 次退避总计约 1.5s,
+ * 够穿过瞬时抖动,又不至于把一次失败拖成用户可感知的卡顿。
+ */
+const COMPACTION_RETRY: RetryPolicy = {
+  enabled: true,
+  maxRetries: 2,
+  baseDelayMs: 500,
+};
+
+const COLD_MEMORY_INDEX = "/shared/memories/MEMORY.md";
+const TRUNCATION_MARKER = "\n[truncated to memory budget]";
+
+function sumUsage(responses: readonly AssistantMessage[]): Usage | undefined {
+  const first = responses[0];
+  if (!first) return undefined;
+
+  const total = structuredClone(first.usage);
+  for (const response of responses.slice(1)) {
+    const usage = response.usage;
+    total.input += usage.input;
+    total.output += usage.output;
+    total.cacheRead += usage.cacheRead;
+    total.cacheWrite += usage.cacheWrite;
+    if (usage.cacheWrite1h !== undefined) {
+      total.cacheWrite1h = (total.cacheWrite1h ?? 0) + usage.cacheWrite1h;
+    }
+    if (usage.reasoning !== undefined) {
+      total.reasoning = (total.reasoning ?? 0) + usage.reasoning;
+    }
+    total.totalTokens += usage.totalTokens;
+    total.cost.input += usage.cost.input;
+    total.cost.output += usage.cost.output;
+    total.cost.cacheRead += usage.cost.cacheRead;
+    total.cost.cacheWrite += usage.cost.cacheWrite;
+    total.cost.total += usage.cost.total;
+  }
+  return total;
+}
+
+// #region Transcript compaction
+
+/**
+ * 把当前 Pi 会话分支压缩成下一轮可以直接使用的消息。
+ *
+ * @remarks
+ * Transcript 在每轮开始前传入当前分支、模型和 token 阈值，并在返回压缩条目后自行追加到 Session。
+ *
+ * 这里先用 Pi 的 Session 管线重建上下文，任何失败都退回未压缩消息，以避免压缩失败损坏持久化会话。
+ *
+ * 核心术语见包入口 `index.ts`。
+ */
+export async function compactPiContext(input: Readonly<{
+  branch: readonly SessionTreeEntry[];
+  compactAfterTokens: number;
+  models: Models;
+  model: Model<any>;
+  apiKey?: string;
+  signal?: AbortSignal;
+}>): Promise<{
+  messages: AgentMessage[];
+  compactionEntry?: Omit<
+    CompactionEntry,
+    "id" | "parentId" | "timestamp"
+  >;
+  modelUsage?: {
+    api: string;
+    provider: string;
+    model: string;
+    responseModel?: string;
+    stopReason?: string;
+    usage: Usage;
+  };
+  didCompact: boolean;
+  degraded: boolean;
+}> {
+  const branch = [...input.branch];
+  let currentMessages = branch.flatMap((entry) =>
+    entry.type === "message" ? [entry.message] : []
+  );
+  try {
+    currentMessages = buildSessionContext(branch).messages;
+    const settings = DEFAULT_COMPACTION_SETTINGS;
+    if (
+      !settings.enabled ||
+      estimateContextTokens(currentMessages).tokens <=
+        input.compactAfterTokens
+    ) {
+      return {
+        messages: currentMessages,
+        didCompact: false,
+        degraded: false,
+      };
+    }
+
+    const preparation = prepareCompaction(branch, settings);
+    if (!preparation.ok) {
+      return {
+        messages: currentMessages,
+        didCompact: false,
+        degraded: true,
+      };
+    }
+    if (!preparation.value) {
+      return {
+        messages: currentMessages,
+        didCompact: false,
+        degraded: false,
+      };
+    }
+
+    const responses: AssistantMessage[] = [];
+    const models = {
+      completeSimple: async (
+        model: Model<any>,
+        context: Parameters<Models["completeSimple"]>[1],
+        options?: Parameters<Models["completeSimple"]>[2],
+      ) => {
+        const response = await input.models.completeSimple(model, context, {
+          ...options,
+          ...(input.apiKey ? { apiKey: input.apiKey } : {}),
+        });
+        responses.push(response);
+        return response;
+      },
+    } as Models;
+    const compacted = await compactPiTranscript(
+      preparation.value,
+      models,
+      input.model,
+      undefined,
+      input.signal,
+      undefined,
+      COMPACTION_RETRY,
+    );
+    if (!compacted.ok) {
+      return {
+        messages: currentMessages,
+        didCompact: false,
+        degraded: true,
+      };
+    }
+
+    const compactionEntry = {
+      type: "compaction",
+      ...compacted.value,
+    } satisfies Omit<
+      CompactionEntry,
+      "id" | "parentId" | "timestamp"
+    >;
+    const previewEntry: CompactionEntry = {
+      ...compactionEntry,
+      id: `compaction-preview:${branch[branch.length - 1]?.id ?? "root"}`,
+      parentId: branch[branch.length - 1]?.id ?? null,
+      timestamp: new Date().toISOString(),
+    };
+    const firstResponse = responses[0];
+    const responseModel = firstResponse?.responseModel && responses.every(
+        (response) => response.responseModel === firstResponse.responseModel,
+      )
+      ? firstResponse.responseModel
+      : undefined;
+    const stopReason = firstResponse && responses.every(
+        (response) => response.stopReason === firstResponse.stopReason,
+      )
+      ? firstResponse.stopReason
+      : undefined;
+    const usage = sumUsage(responses);
+    return {
+      messages: buildSessionContext(
+        [...branch, previewEntry],
+      ).messages,
+      compactionEntry,
+      ...(firstResponse && usage
+        ? {
+            modelUsage: {
+              api: firstResponse.api,
+              provider: firstResponse.provider,
+              model: firstResponse.model,
+              ...(responseModel ? { responseModel } : {}),
+              ...(stopReason ? { stopReason } : {}),
+              usage,
+            },
+          }
+        : {}),
+      didCompact: true,
+      degraded: false,
+    };
+  } catch {
+    return {
+      messages: currentMessages,
+      didCompact: false,
+      degraded: true,
+    };
+  }
+}
+
+// 作用：把一段上下文裁到指定 token 预算内。
+// 调用：内存和 Extension 上下文在拼入 system prompt 前调用。
+// 原因：二分搜索每次都用同一个 token 估算器验证，改成字符数截断会突破真实预算。
+function truncateToBudget(content: string, maxTokens: number): string {
+  if (maxTokens <= 0) return "";
+  if (estimateStringTokens(content) <= maxTokens) return content;
+
+  const marker =
+    estimateStringTokens(TRUNCATION_MARKER) <= maxTokens
+      ? TRUNCATION_MARKER
+      : "";
+  const characters = [...content];
+  let low = 0;
+  let high = characters.length;
+  while (low < high) {
+    const middle = Math.ceil((low + high) / 2);
+    const candidate = `${characters.slice(0, middle).join("")}${marker}`;
+    if (estimateStringTokens(candidate) <= maxTokens) {
+      low = middle;
+    } else {
+      high = middle - 1;
+    }
+  }
+
+  return `${characters.slice(0, low).join("")}${marker}`;
+}
+
+// 作用：给一段动态上下文加上人可读的来源标题。
+// 调用：内存和 Extension 读取器在生成 system prompt 块时调用。
+// 原因：固定分隔格式能保留来源边界，不能随意拼成无标记文本。
+function renderBlock(
+  label: string,
+  description: string,
+  content: string,
+): string {
+  const separator = "═".repeat(46);
+  return `${separator}\n${label} (${description})\n${separator}\n${content}`;
+}
+
+// #endregion
+
+// #region System context sources
+
+// 作用：同时收集当前 Session 内存和用户管理的冷内存索引。
+// 调用：系统上下文装配在 Runtime 锁定当前配置时调用。
+// 原因：内存是可选能力，各来源失败只记 degradation，不应阻断整个 Runtime。
+async function readMemory(
+  profile: RuntimeMemoryProfile,
+  port: Pick<RuntimeMemoryPort, "get"> | undefined,
+  workspace: { readFile(path: string): Promise<string | null> } | undefined,
+): Promise<{
+  blocks: string[];
+  degradations: RuntimeDegradation[];
+}> {
+  if (!profile.enabled) return { blocks: [], degradations: [] };
+
+  const blocks: string[] = [];
+  const degradations: RuntimeDegradation[] = [];
+  const hot = port
+    ? await Promise.allSettled([
+        port.get("memory"),
+        port.get("preferences"),
+      ])
+    : [];
+
+  if (!port) {
+    degradations.push({
+      capability: "memory",
+      reason: "unavailable",
+      detail: "hot",
+    });
+  } else {
+    const memory = hot[0];
+    const preferences = hot[1];
+    let hotUnavailable = false;
+    if (memory.status === "fulfilled" && memory.value) {
+      blocks.push(
+        renderBlock(
+          "MEMORY",
+          "Working facts for the current Session only",
+          truncateToBudget(memory.value, profile.memoryTokens),
+        ),
+      );
+    } else if (memory.status === "rejected") {
+      hotUnavailable = true;
+    }
+    if (preferences.status === "fulfilled" && preferences.value) {
+      blocks.push(
+        renderBlock(
+          "PREFERENCES",
+          "Working preferences for the current Session only",
+          truncateToBudget(
+            preferences.value,
+            profile.preferencesTokens,
+          ),
+        ),
+      );
+    } else if (preferences.status === "rejected") {
+      hotUnavailable = true;
+    }
+    if (hotUnavailable) {
+      degradations.push({
+        capability: "memory",
+        reason: "unavailable",
+        detail: "hot",
+      });
+    }
+  }
+
+  if (!workspace) {
+    degradations.push({
+      capability: "memory",
+      reason: "unavailable",
+      detail: "cold",
+    });
+  } else {
+    try {
+      const index = await workspace.readFile(COLD_MEMORY_INDEX);
+      if (index) {
+        blocks.push(
+          renderBlock(
+            "MEMORY_INDEX",
+            "Index of user-managed cold memory under /shared/memories; entries are read on demand",
+            index,
+          ),
+        );
+      }
+    } catch {
+      degradations.push({
+        capability: "memory",
+        reason: "unavailable",
+        detail: "cold",
+      });
+    }
+  }
+
+  return { blocks, degradations };
+}
+
+// 作用：把已授权 Skill 的持久化 catalog 整理成提示词。
+// 调用：系统上下文装配在需要告诉模型当前可用 Skill 时调用。
+// 原因：catalog 已在发布和绑定时固定，装配期不得为此访问远端内容来源。
+function renderSkillCatalog(
+  bindings: readonly {
+    name: string;
+    description: string;
+  }[],
+): string | null {
+  if (bindings.length === 0) return null;
+
+  return [
+    "AUTHORIZED SKILLS",
+    "Use an authorized Skill when its description matches. Read its instructions or files only through the authorized Skill tool.",
+    "",
+    ...bindings.map(
+      ({ name, description }) => `- ${name}: ${description}`,
+    ),
+  ].join("\n");
+}
+
+// 作用：读取 Extension 声明的上下文并按各自预算截断。
+// 调用：系统上下文装配在已加载 Extension 提供了上下文声明时调用。
+// 原因：各块使用 `Promise.allSettled` 隔离失败，不能让一个 Extension 拖垮整份 system prompt。
+async function readExtensionContexts(
+  contributions: readonly PiExtensionContextContribution[],
+  read:
+    | ((
+        contribution: PiExtensionContextContribution,
+      ) => Promise<string | null>)
+    | undefined,
+): Promise<{
+  blocks: string[];
+  degradations: RuntimeDegradation[];
+}> {
+  if (contributions.length === 0) {
+    return { blocks: [], degradations: [] };
+  }
+  if (!read) {
+    return {
+      blocks: [],
+      degradations: [{
+        capability: "extension",
+        reason: "unavailable",
+        detail: "context",
+      }],
+    };
+  }
+
+  const results = await Promise.allSettled(
+    contributions.map((contribution) => read(contribution)),
+  );
+  const blocks: string[] = [];
+  const degradations: RuntimeDegradation[] = [];
+
+  results.forEach((result, index) => {
+    const contribution = contributions[index]!;
+    if (result.status === "rejected") {
+      degradations.push({
+        capability: "extension",
+        reason: "unavailable",
+        detail: contribution.label,
+      });
+      return;
+    }
+    if (!result.value) return;
+    const content = contribution.maxTokens === undefined
+      ? result.value
+      : truncateToBudget(result.value, contribution.maxTokens);
+    if (!content) return;
+    blocks.push(
+      renderBlock(
+        contribution.label,
+        contribution.description
+          ?? `Context contributed by ${contribution.extensionName}`,
+        content,
+      ),
+    );
+  });
+
+  return { blocks, degradations };
+}
+
+// #endregion
+
+// #region Public assembly
+
+/**
+ * 从已授权的 Runtime 能力重建 Pi 本轮使用的 system prompt。
+ *
+ * @remarks
+ * Runtime 在锁定每轮配置时调用，调用方只应传入 Snapshot 已选中的 Port、Skill 和 Extension 上下文。
+ *
+ * 上下文每次重建而不写入 Pi 消息，这使刷新可选能力不会污染持久化 transcript；各可选来源失败只通过 degradation 返回。
+ *
+ * 核心术语见包入口 `index.ts`。
+ */
+export async function assemblePiSystemContext(input: Readonly<{
+  baseSystemPrompt: string;
+  memory: RuntimeMemoryProfile;
+  memoryPort?: Pick<RuntimeMemoryPort, "get">;
+  workspace?: { readFile(path: string): Promise<string | null> };
+  skills: readonly {
+    name: string;
+    description: string;
+  }[];
+  extensionContexts?: readonly PiExtensionContextContribution[];
+  readExtensionContext?: (
+    contribution: PiExtensionContextContribution,
+  ) => Promise<string | null>;
+}>): Promise<Readonly<{
+  systemPrompt: string;
+  compactAfterTokens: number;
+  degradations: readonly RuntimeDegradation[];
+}>> {
+  const [memory, extensions] = await Promise.all([
+    readMemory(input.memory, input.memoryPort, input.workspace),
+    readExtensionContexts(
+      input.extensionContexts ?? [],
+      input.readExtensionContext,
+    ),
+  ]);
+  const skills = renderSkillCatalog(input.skills);
+  const parts = [
+    input.baseSystemPrompt,
+    ...memory.blocks,
+    ...(skills ? [skills] : []),
+    ...extensions.blocks,
+  ].filter(Boolean);
+
+  return Object.freeze({
+    systemPrompt: parts.join("\n\n"),
+    compactAfterTokens: input.memory.compactAfterTokens,
+    degradations: Object.freeze([
+      ...memory.degradations,
+      ...extensions.degradations,
+    ]),
+  });
+}
+
+// #endregion