feature(context): extend plugin system to support custom context management (#22201)

* feat(context-engine): add ContextEngine interface and registry Introduce the pluggable ContextEngine abstraction that allows external plugins to register custom context management strategies. - ContextEngine interface with lifecycle methods: bootstrap, ingest, ingestBatch, afterTurn, assemble, compact, prepareSubagentSpawn, onSubagentEnded, dispose - Module-level singleton registry with registerContextEngine() and resolveContextEngine() (config-driven slot selection) - LegacyContextEngine: pass-through implementation wrapping existing compaction behavior for 100% backward compatibility - ensureContextEnginesInitialized() guard for safe one-time registration - 19 tests covering contract, registry, resolution, and legacy parity * feat(plugins): add context-engine slot and registerContextEngine API Wire the ContextEngine abstraction into the plugin system so external plugins can register context engines via the standard plugin API. - Add 'context-engine' to PluginKind union type - Add 'contextEngine' slot to PluginSlotsConfig (default: 'legacy') - Wire registerContextEngine() through OpenClawPluginApi - Export ContextEngine types from plugin-sdk for external consumers - Restore proper slot-based resolution in registry * feat(context-engine): wire ContextEngine into agent run lifecycle Integrate the ContextEngine abstraction into the core agent run path: - Resolve context engine once per run (reused across retries) - Bootstrap: hydrate canonical store from session file on first run - Assemble: route context assembly through pluggable engine - Auto-compaction guard: disable built-in auto-compaction when the engine declares ownsCompaction (prevents double-compaction) - AfterTurn: post-turn lifecycle hook for ingest + background compaction decisions - Overflow compaction: route through contextEngine.compact() - Dispose: clean up engine resources in finally block - Notify context engine on subagent lifecycle events Legacy engine: all lifecycle methods are pass-through/no-op, preserving 100% backward compatibility for users without a context engine plugin. * feat(plugins): add scoped subagent methods and gateway request scope Expose runtime.subagent.{run, waitForRun, getSession, deleteSession} so external plugins can spawn sub-agent sessions without raw gateway dispatch access. Uses AsyncLocalStorage request-scope bridge to dispatch internally via handleGatewayRequest with a synthetic operator client. Methods are only available during gateway request handling. - Symbol.for-backed global singleton for cross-module-reload safety - Fallback gateway context for non-WS dispatch paths (Telegram/WhatsApp) - Set gateway request scope for all handlers, not just plugin handlers - 3 staleness tests for fallback context hardening * feat(context-engine): route /compact and sessions.get through context engine Wire the /compact command and sessions.get handler through the pluggable ContextEngine interface. - Thread tokenBudget and force parameters to context engine compact - Route /compact through contextEngine.compact() when registered - Wire sessions.get as runtime alias for plugin subagent dispatch - Add .pebbles/ to .gitignore * style: format with oxfmt 0.33.0 Fix duplicate import (ControlUiRootState in server.impl.ts) and import ordering across all changed files. * fix: update extension test mocks for context-engine types Add missing subagent property to bluebubbles PluginRuntime mock. Add missing registerContextEngine to lobster OpenClawPluginApi mock. * fix(subagents): keep deferred delete cleanup retryable * style: format run attempt for CI * fix(rebase): remove duplicate embedded-run imports * test: add missing gateway context mock export * fix: pass resolved auth profile into afterTurn compaction Ensure the embedded runner forwards resolved auth profile context into legacy context-engine compaction params on the normal afterTurn path, matching overflow compaction behavior. This allows downstream LCM summarization to use the intended provider auth/profile consistently. Also fix strict TS typing in external-link token dedupe and align an attempt unit test reasoningLevel value with the current ReasoningLevel enum. Regeneration-Prompt: | We were debugging context-engine compaction where downstream summary calls were missing the right auth/profile context in normal afterTurn flow, while overflow compaction already propagated it. Preserve current behavior and keep changes additive: thread the resolved authProfileId through run -> attempt -> legacy compaction param builder without broad refactors. Add tests that prove the auth profile is included in afterTurn legacy params and that overflow compaction still passes it through run attempts. Keep existing APIs stable, and only adjust small type issues needed for strict compilation. * fix: remove duplicate imports from rebase * feat: add context-engine system prompt additions * fix(rebase): dedupe attempt import declarations * test: fix fetch mock typing in ollama autodiscovery * fix(test): add registerContextEngine to diffs extension mock APIs * test(windows): use path.delimiter in ios-team-id fixture PATH * test(cron): add model formatting and precedence edge case tests Covers: - Provider/model string splitting (whitespace, nested paths, empty segments) - Provider normalization (casing, aliases like bedrock→amazon-bedrock) - Anthropic model alias normalization (opus-4.5→claude-opus-4-5) - Precedence: job payload > session override > config default - Sequential runs with different providers (CI flake regression pattern) - forceNew session preserving stored model overrides - Whitespace/empty model string edge cases - Config model as string vs object format * test(cron): fix model formatting test config types * test(phone-control): add registerContextEngine to mock API * fix: re-export ChannelKind from config-reload-plan * fix: add subagent mock to plugin-runtime-mock test util * docs: add changelog fragment for context engine PR #22201
2026-04-20 21:23:23 +00:00 · 2026-03-06 05:31:59 -08:00
parent fa6c0e1b40
commit fee91fefce
44 changed files with 2308 additions and 103 deletions
--- a/src/context-engine/context-engine.test.ts
+++ b/src/context-engine/context-engine.test.ts
@@ -0,0 +1,337 @@
+import type { AgentMessage } from "@mariozechner/pi-agent-core";
+import { describe, expect, it, beforeEach } from "vitest";
+// ---------------------------------------------------------------------------
+// We dynamically import the registry so we can get a fresh module per test
+// group when needed.  For most groups we use the shared singleton directly.
+// ---------------------------------------------------------------------------
+import { LegacyContextEngine, registerLegacyContextEngine } from "./legacy.js";
+import {
+  registerContextEngine,
+  getContextEngineFactory,
+  listContextEngineIds,
+  resolveContextEngine,
+} from "./registry.js";
+import type {
+  ContextEngine,
+  ContextEngineInfo,
+  AssembleResult,
+  CompactResult,
+  IngestResult,
+} from "./types.js";
+
+// ---------------------------------------------------------------------------
+// Helpers
+// ---------------------------------------------------------------------------
+
+/** Build a config object with a contextEngine slot for testing. */
+// eslint-disable-next-line @typescript-eslint/no-explicit-any
+function configWithSlot(engineId: string): any {
+  return { plugins: { slots: { contextEngine: engineId } } };
+}
+
+function makeMockMessage(role: "user" | "assistant" = "user", text = "hello"): AgentMessage {
+  return { role, content: text, timestamp: Date.now() } as AgentMessage;
+}
+
+/** A minimal mock engine that satisfies the ContextEngine interface. */
+class MockContextEngine implements ContextEngine {
+  readonly info: ContextEngineInfo = {
+    id: "mock",
+    name: "Mock Engine",
+    version: "0.0.1",
+  };
+
+  async ingest(_params: {
+    sessionId: string;
+    message: AgentMessage;
+    isHeartbeat?: boolean;
+  }): Promise<IngestResult> {
+    return { ingested: true };
+  }
+
+  async assemble(params: {
+    sessionId: string;
+    messages: AgentMessage[];
+    tokenBudget?: number;
+  }): Promise<AssembleResult> {
+    return {
+      messages: params.messages,
+      estimatedTokens: 42,
+      systemPromptAddition: "mock system addition",
+    };
+  }
+
+  async compact(_params: {
+    sessionId: string;
+    sessionFile: string;
+    tokenBudget?: number;
+    compactionTarget?: "budget" | "threshold";
+    customInstructions?: string;
+    legacyParams?: Record<string, unknown>;
+  }): Promise<CompactResult> {
+    return {
+      ok: true,
+      compacted: true,
+      reason: "mock compaction",
+      result: {
+        summary: "mock summary",
+        tokensBefore: 100,
+        tokensAfter: 50,
+      },
+    };
+  }
+
+  async dispose(): Promise<void> {
+    // no-op
+  }
+}
+
+// ═══════════════════════════════════════════════════════════════════════════
+// 1. Engine contract tests
+// ═══════════════════════════════════════════════════════════════════════════
+
+describe("Engine contract tests", () => {
+  it("a mock engine implementing ContextEngine can be registered and resolved", async () => {
+    const factory = () => new MockContextEngine();
+    registerContextEngine("mock", factory);
+
+    const resolved = getContextEngineFactory("mock");
+    expect(resolved).toBe(factory);
+
+    const engine = await resolved!();
+    expect(engine).toBeInstanceOf(MockContextEngine);
+    expect(engine.info.id).toBe("mock");
+  });
+
+  it("ingest() returns IngestResult with ingested boolean", async () => {
+    const engine = new MockContextEngine();
+    const result = await engine.ingest({
+      sessionId: "s1",
+      message: makeMockMessage(),
+    });
+
+    expect(result).toHaveProperty("ingested");
+    expect(typeof result.ingested).toBe("boolean");
+    expect(result.ingested).toBe(true);
+  });
+
+  it("assemble() returns AssembleResult with messages array and estimatedTokens", async () => {
+    const engine = new MockContextEngine();
+    const msgs = [makeMockMessage(), makeMockMessage("assistant", "world")];
+    const result = await engine.assemble({
+      sessionId: "s1",
+      messages: msgs,
+    });
+
+    expect(Array.isArray(result.messages)).toBe(true);
+    expect(result.messages).toHaveLength(2);
+    expect(typeof result.estimatedTokens).toBe("number");
+    expect(result.estimatedTokens).toBe(42);
+    expect(result.systemPromptAddition).toBe("mock system addition");
+  });
+
+  it("compact() returns CompactResult with ok, compacted, reason, result fields", async () => {
+    const engine = new MockContextEngine();
+    const result = await engine.compact({
+      sessionId: "s1",
+      sessionFile: "/tmp/session.json",
+    });
+
+    expect(typeof result.ok).toBe("boolean");
+    expect(typeof result.compacted).toBe("boolean");
+    expect(result.ok).toBe(true);
+    expect(result.compacted).toBe(true);
+    expect(result.reason).toBe("mock compaction");
+    expect(result.result).toBeDefined();
+    expect(result.result!.summary).toBe("mock summary");
+    expect(result.result!.tokensBefore).toBe(100);
+    expect(result.result!.tokensAfter).toBe(50);
+  });
+
+  it("dispose() is callable (optional method)", async () => {
+    const engine = new MockContextEngine();
+    // Should complete without error
+    await expect(engine.dispose()).resolves.toBeUndefined();
+  });
+});
+
+// ═══════════════════════════════════════════════════════════════════════════
+// 2. Registry tests
+// ═══════════════════════════════════════════════════════════════════════════
+
+describe("Registry tests", () => {
+  it("registerContextEngine() stores a factory", () => {
+    const factory = () => new MockContextEngine();
+    registerContextEngine("reg-test-1", factory);
+
+    expect(getContextEngineFactory("reg-test-1")).toBe(factory);
+  });
+
+  it("getContextEngineFactory() returns the factory", () => {
+    const factory = () => new MockContextEngine();
+    registerContextEngine("reg-test-2", factory);
+
+    const retrieved = getContextEngineFactory("reg-test-2");
+    expect(retrieved).toBe(factory);
+    expect(typeof retrieved).toBe("function");
+  });
+
+  it("listContextEngineIds() returns all registered ids", () => {
+    // Ensure at least our test entries exist
+    registerContextEngine("reg-test-a", () => new MockContextEngine());
+    registerContextEngine("reg-test-b", () => new MockContextEngine());
+
+    const ids = listContextEngineIds();
+    expect(ids).toContain("reg-test-a");
+    expect(ids).toContain("reg-test-b");
+    expect(Array.isArray(ids)).toBe(true);
+  });
+
+  it("registering the same id overwrites the previous factory", () => {
+    const factory1 = () => new MockContextEngine();
+    const factory2 = () => new MockContextEngine();
+
+    registerContextEngine("reg-overwrite", factory1);
+    expect(getContextEngineFactory("reg-overwrite")).toBe(factory1);
+
+    registerContextEngine("reg-overwrite", factory2);
+    expect(getContextEngineFactory("reg-overwrite")).toBe(factory2);
+    expect(getContextEngineFactory("reg-overwrite")).not.toBe(factory1);
+  });
+});
+
+// ═══════════════════════════════════════════════════════════════════════════
+// 3. Default engine selection
+// ═══════════════════════════════════════════════════════════════════════════
+
+describe("Default engine selection", () => {
+  // Ensure both legacy and a custom test engine are registered before these tests.
+  beforeEach(() => {
+    // Registration is idempotent (Map.set), so calling again is safe.
+    registerLegacyContextEngine();
+    // Register a lightweight custom stub so we don't need external resources.
+    registerContextEngine("test-engine", () => {
+      const engine: ContextEngine = {
+        info: { id: "test-engine", name: "Custom Test Engine", version: "0.0.0" },
+        async ingest() {
+          return { ingested: true };
+        },
+        async assemble({ messages }) {
+          return { messages, estimatedTokens: 0 };
+        },
+        async compact() {
+          return { ok: true, compacted: false };
+        },
+      };
+      return engine;
+    });
+  });
+
+  it("resolveContextEngine() with no config returns the default ('legacy') engine", async () => {
+    const engine = await resolveContextEngine();
+    expect(engine.info.id).toBe("legacy");
+  });
+
+  it("resolveContextEngine() with config contextEngine='legacy' returns legacy engine", async () => {
+    const engine = await resolveContextEngine(configWithSlot("legacy"));
+    expect(engine.info.id).toBe("legacy");
+  });
+
+  it("resolveContextEngine() with config contextEngine='test-engine' returns the custom engine", async () => {
+    const engine = await resolveContextEngine(configWithSlot("test-engine"));
+    expect(engine.info.id).toBe("test-engine");
+  });
+});
+
+// ═══════════════════════════════════════════════════════════════════════════
+// 4. Invalid engine fallback
+// ═══════════════════════════════════════════════════════════════════════════
+
+describe("Invalid engine fallback", () => {
+  it("resolveContextEngine() with config pointing to unregistered engine throws with helpful error", async () => {
+    await expect(resolveContextEngine(configWithSlot("nonexistent-engine"))).rejects.toThrow(
+      /nonexistent-engine/,
+    );
+  });
+
+  it("error message includes the requested id and available ids", async () => {
+    // Ensure at least legacy is registered so we see it in the available list
+    registerLegacyContextEngine();
+
+    try {
+      await resolveContextEngine(configWithSlot("does-not-exist"));
+      // Should not reach here
+      expect.unreachable("Expected resolveContextEngine to throw");
+    } catch (err: unknown) {
+      const message = err instanceof Error ? err.message : String(err);
+      expect(message).toContain("does-not-exist");
+      expect(message).toContain("not registered");
+      // Should mention available engines
+      expect(message).toMatch(/Available engines:/);
+      // At least "legacy" should be listed as available
+      expect(message).toContain("legacy");
+    }
+  });
+});
+
+// ═══════════════════════════════════════════════════════════════════════════
+// 5. LegacyContextEngine parity
+// ═══════════════════════════════════════════════════════════════════════════
+
+describe("LegacyContextEngine parity", () => {
+  it("ingest() returns { ingested: false } (no-op)", async () => {
+    const engine = new LegacyContextEngine();
+    const result = await engine.ingest({
+      sessionId: "s1",
+      message: makeMockMessage(),
+    });
+
+    expect(result).toEqual({ ingested: false });
+  });
+
+  it("assemble() returns messages as-is (pass-through)", async () => {
+    const engine = new LegacyContextEngine();
+    const messages = [
+      makeMockMessage("user", "first"),
+      makeMockMessage("assistant", "second"),
+      makeMockMessage("user", "third"),
+    ];
+
+    const result = await engine.assemble({
+      sessionId: "s1",
+      messages,
+    });
+
+    // Messages should be the exact same array reference (pass-through)
+    expect(result.messages).toBe(messages);
+    expect(result.messages).toHaveLength(3);
+    expect(result.estimatedTokens).toBe(0);
+    expect(result.systemPromptAddition).toBeUndefined();
+  });
+
+  it("dispose() completes without error", async () => {
+    const engine = new LegacyContextEngine();
+    await expect(engine.dispose()).resolves.toBeUndefined();
+  });
+});
+
+// ═══════════════════════════════════════════════════════════════════════════
+// 6. Initialization guard
+// ═══════════════════════════════════════════════════════════════════════════
+
+describe("Initialization guard", () => {
+  it("ensureContextEnginesInitialized() is idempotent (calling twice does not throw)", async () => {
+    const { ensureContextEnginesInitialized } = await import("./init.js");
+
+    expect(() => ensureContextEnginesInitialized()).not.toThrow();
+    expect(() => ensureContextEnginesInitialized()).not.toThrow();
+  });
+
+  it("after init, 'legacy' engine is registered", async () => {
+    const { ensureContextEnginesInitialized } = await import("./init.js");
+    ensureContextEnginesInitialized();
+
+    const ids = listContextEngineIds();
+    expect(ids).toContain("legacy");
+  });
+});
--- a/src/context-engine/index.ts
+++ b/src/context-engine/index.ts
@@ -0,0 +1,19 @@
+export type {
+  ContextEngine,
+  ContextEngineInfo,
+  AssembleResult,
+  CompactResult,
+  IngestResult,
+} from "./types.js";
+
+export {
+  registerContextEngine,
+  getContextEngineFactory,
+  listContextEngineIds,
+  resolveContextEngine,
+} from "./registry.js";
+export type { ContextEngineFactory } from "./registry.js";
+
+export { LegacyContextEngine, registerLegacyContextEngine } from "./legacy.js";
+
+export { ensureContextEnginesInitialized } from "./init.js";
--- a/src/context-engine/init.ts
+++ b/src/context-engine/init.ts
@@ -0,0 +1,23 @@
+import { registerLegacyContextEngine } from "./legacy.js";
+
+/**
+ * Ensures all built-in context engines are registered exactly once.
+ *
+ * The legacy engine is always registered as a safe fallback so that
+ * `resolveContextEngine()` can resolve the default "legacy" slot without
+ * callers needing to remember manual registration.
+ *
+ * Additional engines are registered by their own plugins via
+ * `api.registerContextEngine()` during plugin load.
+ */
+let initialized = false;
+
+export function ensureContextEnginesInitialized(): void {
+  if (initialized) {
+    return;
+  }
+  initialized = true;
+
+  // Always available – safe fallback for the "legacy" slot default.
+  registerLegacyContextEngine();
+}
--- a/src/context-engine/legacy.ts
+++ b/src/context-engine/legacy.ts
@@ -0,0 +1,115 @@
+import type { AgentMessage } from "@mariozechner/pi-agent-core";
+import { registerContextEngine } from "./registry.js";
+import type {
+  ContextEngine,
+  ContextEngineInfo,
+  AssembleResult,
+  CompactResult,
+  IngestResult,
+} from "./types.js";
+
+/**
+ * LegacyContextEngine wraps the existing compaction behavior behind the
+ * ContextEngine interface, preserving 100% backward compatibility.
+ *
+ * - ingest: no-op (SessionManager handles message persistence)
+ * - assemble: pass-through (existing sanitize/validate/limit pipeline in attempt.ts handles this)
+ * - compact: delegates to compactEmbeddedPiSessionDirect
+ */
+export class LegacyContextEngine implements ContextEngine {
+  readonly info: ContextEngineInfo = {
+    id: "legacy",
+    name: "Legacy Context Engine",
+    version: "1.0.0",
+  };
+
+  async ingest(_params: {
+    sessionId: string;
+    message: AgentMessage;
+    isHeartbeat?: boolean;
+  }): Promise<IngestResult> {
+    // No-op: SessionManager handles message persistence in the legacy flow
+    return { ingested: false };
+  }
+
+  async assemble(params: {
+    sessionId: string;
+    messages: AgentMessage[];
+    tokenBudget?: number;
+  }): Promise<AssembleResult> {
+    // Pass-through: the existing sanitize -> validate -> limit -> repair pipeline
+    // in attempt.ts handles context assembly for the legacy engine.
+    // We just return the messages as-is with a rough token estimate.
+    return {
+      messages: params.messages,
+      estimatedTokens: 0, // Caller handles estimation
+    };
+  }
+
+  async afterTurn(_params: {
+    sessionId: string;
+    sessionFile: string;
+    messages: AgentMessage[];
+    prePromptMessageCount: number;
+    autoCompactionSummary?: string;
+    isHeartbeat?: boolean;
+    tokenBudget?: number;
+    legacyCompactionParams?: Record<string, unknown>;
+  }): Promise<void> {
+    // No-op: legacy flow persists context directly in SessionManager.
+  }
+
+  async compact(params: {
+    sessionId: string;
+    sessionFile: string;
+    tokenBudget?: number;
+    force?: boolean;
+    currentTokenCount?: number;
+    compactionTarget?: "budget" | "threshold";
+    customInstructions?: string;
+    legacyParams?: Record<string, unknown>;
+  }): Promise<CompactResult> {
+    // Import dynamically to avoid circular dependencies
+    const { compactEmbeddedPiSessionDirect } =
+      await import("../agents/pi-embedded-runner/compact.js");
+
+    // legacyParams carries the full CompactEmbeddedPiSessionParams fields
+    // set by the caller in run.ts. We spread them and override the fields
+    // that come from the ContextEngine compact() signature directly.
+    const lp = params.legacyParams ?? {};
+
+    // eslint-disable-next-line @typescript-eslint/no-explicit-any -- legacy bridge: legacyParams is an opaque bag matching CompactEmbeddedPiSessionParams
+    const result = await compactEmbeddedPiSessionDirect({
+      ...lp,
+      sessionId: params.sessionId,
+      sessionFile: params.sessionFile,
+      tokenBudget: params.tokenBudget,
+      force: params.force,
+      customInstructions: params.customInstructions,
+      workspaceDir: (lp.workspaceDir as string) ?? process.cwd(),
+    } as Parameters<typeof compactEmbeddedPiSessionDirect>[0]);
+
+    return {
+      ok: result.ok,
+      compacted: result.compacted,
+      reason: result.reason,
+      result: result.result
+        ? {
+            summary: result.result.summary,
+            firstKeptEntryId: result.result.firstKeptEntryId,
+            tokensBefore: result.result.tokensBefore,
+            tokensAfter: result.result.tokensAfter,
+            details: result.result.details,
+          }
+        : undefined,
+    };
+  }
+
+  async dispose(): Promise<void> {
+    // Nothing to clean up for legacy engine
+  }
+}
+
+export function registerLegacyContextEngine(): void {
+  registerContextEngine("legacy", () => new LegacyContextEngine());
+}
--- a/src/context-engine/registry.ts
+++ b/src/context-engine/registry.ts
@@ -0,0 +1,67 @@
+import type { OpenClawConfig } from "../config/config.js";
+import { defaultSlotIdForKey } from "../plugins/slots.js";
+import type { ContextEngine } from "./types.js";
+
+/**
+ * A factory that creates a ContextEngine instance.
+ * Supports async creation for engines that need DB connections etc.
+ */
+export type ContextEngineFactory = () => ContextEngine | Promise<ContextEngine>;
+
+// ---------------------------------------------------------------------------
+// Registry (module-level singleton)
+// ---------------------------------------------------------------------------
+
+const _engines = new Map<string, ContextEngineFactory>();
+
+/**
+ * Register a context engine implementation under the given id.
+ */
+export function registerContextEngine(id: string, factory: ContextEngineFactory): void {
+  _engines.set(id, factory);
+}
+
+/**
+ * Return the factory for a registered engine, or undefined.
+ */
+export function getContextEngineFactory(id: string): ContextEngineFactory | undefined {
+  return _engines.get(id);
+}
+
+/**
+ * List all registered engine ids.
+ */
+export function listContextEngineIds(): string[] {
+  return [..._engines.keys()];
+}
+
+// ---------------------------------------------------------------------------
+// Resolution
+// ---------------------------------------------------------------------------
+
+/**
+ * Resolve which ContextEngine to use based on plugin slot configuration.
+ *
+ * Resolution order:
+ *   1. `config.plugins.slots.contextEngine` (explicit slot override)
+ *   2. Default slot value ("legacy")
+ *
+ * Throws if the resolved engine id has no registered factory.
+ */
+export async function resolveContextEngine(config?: OpenClawConfig): Promise<ContextEngine> {
+  const slotValue = config?.plugins?.slots?.contextEngine;
+  const engineId =
+    typeof slotValue === "string" && slotValue.trim()
+      ? slotValue.trim()
+      : defaultSlotIdForKey("contextEngine");
+
+  const factory = _engines.get(engineId);
+  if (!factory) {
+    throw new Error(
+      `Context engine "${engineId}" is not registered. ` +
+        `Available engines: ${listContextEngineIds().join(", ") || "(none)"}`,
+    );
+  }
+
+  return factory();
+}
--- a/src/context-engine/types.ts
+++ b/src/context-engine/types.ts
@@ -0,0 +1,167 @@
+import type { AgentMessage } from "@mariozechner/pi-agent-core";
+
+// Result types
+
+export type AssembleResult = {
+  /** Ordered messages to use as model context */
+  messages: AgentMessage[];
+  /** Estimated total tokens in assembled context */
+  estimatedTokens: number;
+  /** Optional context-engine-provided instructions prepended to the runtime system prompt */
+  systemPromptAddition?: string;
+};
+
+export type CompactResult = {
+  ok: boolean;
+  compacted: boolean;
+  reason?: string;
+  result?: {
+    summary?: string;
+    firstKeptEntryId?: string;
+    tokensBefore: number;
+    tokensAfter?: number;
+    details?: unknown;
+  };
+};
+
+export type IngestResult = {
+  /** Whether the message was ingested (false if duplicate or no-op) */
+  ingested: boolean;
+};
+
+export type IngestBatchResult = {
+  /** Number of messages ingested from the supplied batch */
+  ingestedCount: number;
+};
+
+export type BootstrapResult = {
+  /** Whether bootstrap ran and initialized the engine's store */
+  bootstrapped: boolean;
+  /** Number of historical messages imported (if applicable) */
+  importedMessages?: number;
+  /** Optional reason when bootstrap was skipped */
+  reason?: string;
+};
+
+export type ContextEngineInfo = {
+  id: string;
+  name: string;
+  version?: string;
+  /** True when the engine manages its own compaction lifecycle. */
+  ownsCompaction?: boolean;
+};
+
+export type SubagentSpawnPreparation = {
+  /** Roll back pre-spawn setup when subagent launch fails. */
+  rollback: () => void | Promise<void>;
+};
+
+export type SubagentEndReason = "deleted" | "completed" | "swept" | "released";
+
+/**
+ * ContextEngine defines the pluggable contract for context management.
+ *
+ * Required methods define a generic lifecycle; optional methods allow engines
+ * to provide additional capabilities (retrieval, lineage, etc.).
+ */
+export interface ContextEngine {
+  /** Engine identifier and metadata */
+  readonly info: ContextEngineInfo;
+
+  /**
+   * Initialize engine state for a session, optionally importing historical context.
+   */
+  bootstrap?(params: { sessionId: string; sessionFile: string }): Promise<BootstrapResult>;
+
+  /**
+   * Ingest a single message into the engine's store.
+   */
+  ingest(params: {
+    sessionId: string;
+    message: AgentMessage;
+    /** True when the message belongs to a heartbeat run. */
+    isHeartbeat?: boolean;
+  }): Promise<IngestResult>;
+
+  /**
+   * Ingest a completed turn batch as a single unit.
+   */
+  ingestBatch?(params: {
+    sessionId: string;
+    messages: AgentMessage[];
+    /** True when the batch belongs to a heartbeat run. */
+    isHeartbeat?: boolean;
+  }): Promise<IngestBatchResult>;
+
+  /**
+   * Execute optional post-turn lifecycle work after a run attempt completes.
+   * Engines can use this to persist canonical context and trigger background
+   * compaction decisions.
+   */
+  afterTurn?(params: {
+    sessionId: string;
+    sessionFile: string;
+    messages: AgentMessage[];
+    /** Number of messages that existed before the prompt was sent. */
+    prePromptMessageCount: number;
+    /** Optional auto-compaction summary emitted by the runtime. */
+    autoCompactionSummary?: string;
+    /** True when this turn belongs to a heartbeat run. */
+    isHeartbeat?: boolean;
+    /** Optional model context token budget for proactive compaction. */
+    tokenBudget?: number;
+    /** Backward-compat only: legacy compaction bridge runtime params. */
+    legacyCompactionParams?: Record<string, unknown>;
+  }): Promise<void>;
+
+  /**
+   * Assemble model context under a token budget.
+   * Returns an ordered set of messages ready for the model.
+   */
+  assemble(params: {
+    sessionId: string;
+    messages: AgentMessage[];
+    tokenBudget?: number;
+  }): Promise<AssembleResult>;
+
+  /**
+   * Compact context to reduce token usage.
+   * May create summaries, prune old turns, etc.
+   */
+  compact(params: {
+    sessionId: string;
+    sessionFile: string;
+    tokenBudget?: number;
+    /** Backward-compat only: force legacy compaction behavior even below threshold. */
+    force?: boolean;
+    /** Optional live token estimate from the caller's active context. */
+    currentTokenCount?: number;
+    /** Controls convergence target; defaults to budget for compatibility. */
+    compactionTarget?: "budget" | "threshold";
+    customInstructions?: string;
+    /** Backward-compat only: full params bag for legacy compaction bridge. */
+    legacyParams?: Record<string, unknown>;
+  }): Promise<CompactResult>;
+
+  /**
+   * Prepare context-engine-managed subagent state before the child run starts.
+   *
+   * Implementations can return a rollback handle that is invoked when spawn
+   * fails after preparation succeeds.
+   */
+  prepareSubagentSpawn?(params: {
+    parentSessionKey: string;
+    childSessionKey: string;
+    ttlMs?: number;
+  }): Promise<SubagentSpawnPreparation | undefined>;
+
+  /**
+   * Notify the context engine that a subagent lifecycle ended.
+   */
+  onSubagentEnded?(params: { childSessionKey: string; reason: SubagentEndReason }): Promise<void>;
+
+  /**
+   * Dispose of any resources held by the engine.
+   */
+  dispose?(): Promise<void>;
+}