feat(context-engine): pass incoming prompt to assemble (#50848)

Merged via squash. Prepared head SHA: 282dc9264d Co-authored-by: danhdoan <12591333+danhdoan@users.noreply.github.com> Co-authored-by: jalehman <550978+jalehman@users.noreply.github.com> Reviewed-by: @jalehman
2026-03-21 23:11:01 +00:00 · 2026-03-21 07:03:21 +07:00
parent 6a6f1b5351
commit e78129a4d9
5 changed files with 296 additions and 34 deletions
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -190,6 +190,7 @@ Docs: https://docs.openclaw.ai
 - Telegram/routing: fail loud when `message send` targets an unknown non-default Telegram `accountId`, instead of silently falling back to the channel-level bot token and sending through the wrong bot. (#50853) Thanks @hclsys.
 - Web search: align onboarding, configure, and finalize with plugin-owned provider contracts, including disabled-provider recovery, config-aware credential hooks, and runtime-visible summaries. (#50935) Thanks @gumadeiras.
 - Agents/replay: sanitize malformed assistant tool-call replay blocks before provider replay so follow-up Anthropic requests do not inherit the downstream `replace` crash. (#50005) Thanks @jalehman.
+- Plugins/context engines: retry strict legacy `assemble()` calls without the new `prompt` field when older engines reject it, preserving prompt-aware retrieval compatibility for pre-prompt plugins. (#50848) thanks @danhdoan.

 ### Breaking

--- a/src/agents/pi-embedded-runner/run/attempt.ts
+++ b/src/agents/pi-embedded-runner/run/attempt.ts
@@ -2426,6 +2426,7 @@ export async function runEmbeddedAttempt(
              messages: activeSession.messages,
              tokenBudget: params.contextTokenBudget,
              model: params.modelId,
+              ...(params.prompt !== undefined ? { prompt: params.prompt } : {}),
            });
            if (assembled.messages !== activeSession.messages) {
              activeSession.agent.replaceMessages(assembled.messages);
--- a/src/context-engine/context-engine.test.ts
+++ b/src/context-engine/context-engine.test.ts
@@ -145,6 +145,7 @@ class LegacySessionKeyStrictEngine implements ContextEngine {
    sessionKey?: string;
    messages: AgentMessage[];
    tokenBudget?: number;
+    prompt?: string;
  }): Promise<AssembleResult> {
    this.assembleCalls.push({ ...params });
    this.rejectSessionKey(params);
@@ -234,6 +235,58 @@ class SessionKeyRuntimeErrorEngine implements ContextEngine {
  }
 }

+class LegacyAssembleStrictEngine implements ContextEngine {
+  readonly info: ContextEngineInfo = {
+    id: "legacy-assemble-strict",
+    name: "Legacy Assemble Strict Engine",
+  };
+  readonly assembleCalls: Array<Record<string, unknown>> = [];
+
+  async ingest(_params: {
+    sessionId: string;
+    sessionKey?: string;
+    message: AgentMessage;
+    isHeartbeat?: boolean;
+  }): Promise<IngestResult> {
+    return { ingested: true };
+  }
+
+  async assemble(params: {
+    sessionId: string;
+    sessionKey?: string;
+    messages: AgentMessage[];
+    tokenBudget?: number;
+    prompt?: string;
+  }): Promise<AssembleResult> {
+    this.assembleCalls.push({ ...params });
+    if (Object.prototype.hasOwnProperty.call(params, "sessionKey")) {
+      throw new Error("Unrecognized key(s) in object: 'sessionKey'");
+    }
+    if (Object.prototype.hasOwnProperty.call(params, "prompt")) {
+      throw new Error("Unrecognized key(s) in object: 'prompt'");
+    }
+    return {
+      messages: params.messages,
+      estimatedTokens: 3,
+    };
+  }
+
+  async compact(_params: {
+    sessionId: string;
+    sessionKey?: string;
+    sessionFile: string;
+    tokenBudget?: number;
+    compactionTarget?: "budget" | "threshold";
+    customInstructions?: string;
+    runtimeContext?: Record<string, unknown>;
+  }): Promise<CompactResult> {
+    return {
+      ok: true,
+      compacted: false,
+    };
+  }
+}
+
 // ═══════════════════════════════════════════════════════════════════════════
 // 1. Engine contract tests
 // ═══════════════════════════════════════════════════════════════════════════
@@ -640,6 +693,124 @@ describe("LegacyContextEngine parity", () => {
  });
 });

+// ═══════════════════════════════════════════════════════════════════════════
+// 5b. assemble() prompt forwarding
+// ═══════════════════════════════════════════════════════════════════════════
+
+describe("assemble() prompt forwarding", () => {
+  it("forwards prompt to the underlying engine", async () => {
+    const engineId = `prompt-fwd-${Date.now().toString(36)}`;
+    const calls: Array<Record<string, unknown>> = [];
+    registerContextEngine(engineId, () => ({
+      info: { id: engineId, name: "Prompt Tracker", version: "0.0.0" },
+      async ingest() {
+        return { ingested: false };
+      },
+      async assemble(params) {
+        calls.push({ ...params });
+        return { messages: params.messages, estimatedTokens: 0 };
+      },
+      async compact() {
+        return { ok: true, compacted: false };
+      },
+    }));
+
+    const engine = await resolveContextEngine(configWithSlot(engineId));
+    await engine.assemble({
+      sessionId: "s1",
+      messages: [makeMockMessage("user", "hello")],
+      prompt: "hello",
+    });
+
+    expect(calls).toHaveLength(1);
+    expect(calls[0]).toHaveProperty("prompt", "hello");
+  });
+
+  it("omits prompt when not provided", async () => {
+    const engineId = `prompt-omit-${Date.now().toString(36)}`;
+    const calls: Array<Record<string, unknown>> = [];
+    registerContextEngine(engineId, () => ({
+      info: { id: engineId, name: "Prompt Tracker", version: "0.0.0" },
+      async ingest() {
+        return { ingested: false };
+      },
+      async assemble(params) {
+        calls.push({ ...params });
+        return { messages: params.messages, estimatedTokens: 0 };
+      },
+      async compact() {
+        return { ok: true, compacted: false };
+      },
+    }));
+
+    const engine = await resolveContextEngine(configWithSlot(engineId));
+    await engine.assemble({
+      sessionId: "s1",
+      messages: [makeMockMessage("user", "hello")],
+    });
+
+    expect(calls).toHaveLength(1);
+    expect(calls[0]).not.toHaveProperty("prompt");
+  });
+
+  it("does not leak prompt key when caller spreads undefined", async () => {
+    // Guards against the pattern `{ prompt: params.prompt }` when params.prompt
+    // is undefined — JavaScript keeps the key present with value undefined,
+    // which breaks engines that guard with `'prompt' in params`.
+    const engineId = `prompt-undef-${Date.now().toString(36)}`;
+    const calls: Array<Record<string, unknown>> = [];
+    registerContextEngine(engineId, () => ({
+      info: { id: engineId, name: "Prompt Tracker", version: "0.0.0" },
+      async ingest() {
+        return { ingested: false };
+      },
+      async assemble(params) {
+        calls.push({ ...params });
+        return { messages: params.messages, estimatedTokens: 0 };
+      },
+      async compact() {
+        return { ok: true, compacted: false };
+      },
+    }));
+
+    const engine = await resolveContextEngine(configWithSlot(engineId));
+    // Simulate the attempt.ts call-site pattern: conditional spread
+    const callerPrompt: string | undefined = undefined;
+    await engine.assemble({
+      sessionId: "s1",
+      messages: [makeMockMessage("user", "hello")],
+      ...(callerPrompt !== undefined ? { prompt: callerPrompt } : {}),
+    });
+
+    expect(calls).toHaveLength(1);
+    expect(calls[0]).not.toHaveProperty("prompt");
+    expect(Object.keys(calls[0] as object)).not.toContain("prompt");
+  });
+
+  it("retries strict legacy assemble without sessionKey and prompt", async () => {
+    const engineId = `prompt-legacy-${Date.now().toString(36)}`;
+    const strictEngine = new LegacyAssembleStrictEngine();
+    registerContextEngine(engineId, () => strictEngine);
+
+    const engine = await resolveContextEngine(configWithSlot(engineId));
+    const result = await engine.assemble({
+      sessionId: "s1",
+      sessionKey: "agent:main:test",
+      messages: [makeMockMessage("user", "hello")],
+      prompt: "hello",
+    });
+
+    expect(result.estimatedTokens).toBe(3);
+    expect(strictEngine.assembleCalls).toHaveLength(3);
+    expect(strictEngine.assembleCalls[0]).toHaveProperty("sessionKey", "agent:main:test");
+    expect(strictEngine.assembleCalls[0]).toHaveProperty("prompt", "hello");
+    expect(strictEngine.assembleCalls[1]).not.toHaveProperty("sessionKey");
+    expect(strictEngine.assembleCalls[1]).toHaveProperty("prompt", "hello");
+    expect(strictEngine.assembleCalls[2]).not.toHaveProperty("sessionKey");
+    expect(strictEngine.assembleCalls[2]).not.toHaveProperty("prompt");
+  });
+});
+
 // ═══════════════════════════════════════════════════════════════════════════
 // 6. Initialization guard
 // ═══════════════════════════════════════════════════════════════════════════
--- a/src/context-engine/registry.ts
+++ b/src/context-engine/registry.ts
@@ -23,11 +23,24 @@ const SESSION_KEY_COMPAT_METHODS = [
  "assemble",
  "compact",
 ] as const;
+const LEGACY_COMPAT_PARAMS = ["sessionKey", "prompt"] as const;
+const LEGACY_COMPAT_METHOD_KEYS = {
+  bootstrap: ["sessionKey"],
+  maintain: ["sessionKey"],
+  ingest: ["sessionKey"],
+  ingestBatch: ["sessionKey"],
+  afterTurn: ["sessionKey"],
+  assemble: ["sessionKey", "prompt"],
+  compact: ["sessionKey"],
+} as const;

 type SessionKeyCompatMethodName = (typeof SESSION_KEY_COMPAT_METHODS)[number];
 type SessionKeyCompatParams = {
  sessionKey?: string;
+  prompt?: string;
 };
+type LegacyCompatKey = (typeof LEGACY_COMPAT_PARAMS)[number];
+type LegacyCompatParamMap = Partial<Record<LegacyCompatKey, unknown>>;

 function isSessionKeyCompatMethodName(value: PropertyKey): value is SessionKeyCompatMethodName {
  return (
@@ -35,21 +48,29 @@ function isSessionKeyCompatMethodName(value: PropertyKey): value is SessionKeyCo
  );
 }

-function hasOwnSessionKey(params: unknown): params is SessionKeyCompatParams {
+function hasOwnLegacyCompatKey<K extends LegacyCompatKey>(
+  params: unknown,
+  key: K,
+): params is SessionKeyCompatParams & Required<Pick<LegacyCompatParamMap, K>> {
  return (
    params !== null &&
    typeof params === "object" &&
-    Object.prototype.hasOwnProperty.call(params, "sessionKey")
+    Object.prototype.hasOwnProperty.call(params, key)
  );
 }

-function withoutSessionKey<T extends SessionKeyCompatParams>(params: T): T {
+function withoutLegacyCompatKeys<T extends SessionKeyCompatParams>(
+  params: T,
+  keys: Iterable<LegacyCompatKey>,
+): T {
  const legacyParams = { ...params };
-  delete legacyParams.sessionKey;
+  for (const key of keys) {
+    delete legacyParams[key];
+  }
  return legacyParams;
 }

-function issueRejectsSessionKeyStrictly(issue: unknown): boolean {
+function issueRejectsLegacyCompatKeyStrictly(issue: unknown, key: LegacyCompatKey): boolean {
  if (!issue || typeof issue !== "object") {
    return false;
  }
@@ -62,12 +83,12 @@ function issueRejectsSessionKeyStrictly(issue: unknown): boolean {
  if (
    issueRecord.code === "unrecognized_keys" &&
    Array.isArray(issueRecord.keys) &&
-    issueRecord.keys.some((key) => key === "sessionKey")
+    issueRecord.keys.some((issueKey) => issueKey === key)
  ) {
    return true;
  }

-  return isSessionKeyCompatibilityError(issueRecord.message);
+  return isLegacyCompatErrorForKey(issueRecord.message, key);
 }

 function* iterateErrorChain(error: unknown) {
@@ -83,31 +104,45 @@ function* iterateErrorChain(error: unknown) {
  }
 }

-const SESSION_KEY_UNKNOWN_FIELD_PATTERNS = [
-  /\bunrecognized key(?:\(s\)|s)? in object:.*['"`]sessionKey['"`]/i,
-  /\badditional propert(?:y|ies)\b.*['"`]sessionKey['"`]/i,
-  /\bmust not have additional propert(?:y|ies)\b.*['"`]sessionKey['"`]/i,
-  /\b(?:unexpected|extraneous)\s+(?:property|properties|field|fields|key|keys)\b.*['"`]sessionKey['"`]/i,
-  /\b(?:unknown|invalid)\s+(?:property|properties|field|fields|key|keys)\b.*['"`]sessionKey['"`]/i,
-  /['"`]sessionKey['"`].*\b(?:was|is)\s+not allowed\b/i,
-  /"code"\s*:\s*"unrecognized_keys"[^]*"sessionKey"/i,
-] as const;
+const LEGACY_UNKNOWN_FIELD_PATTERNS: Record<LegacyCompatKey, readonly RegExp[]> = {
+  sessionKey: [
+    /\bunrecognized key(?:\(s\)|s)? in object:.*['"`]sessionKey['"`]/i,
+    /\badditional propert(?:y|ies)\b.*['"`]sessionKey['"`]/i,
+    /\bmust not have additional propert(?:y|ies)\b.*['"`]sessionKey['"`]/i,
+    /\b(?:unexpected|extraneous)\s+(?:property|properties|field|fields|key|keys)\b.*['"`]sessionKey['"`]/i,
+    /\b(?:unknown|invalid)\s+(?:property|properties|field|fields|key|keys)\b.*['"`]sessionKey['"`]/i,
+    /['"`]sessionKey['"`].*\b(?:was|is)\s+not allowed\b/i,
+    /"code"\s*:\s*"unrecognized_keys"[^]*"sessionKey"/i,
+  ],
+  prompt: [
+    /\bunrecognized key(?:\(s\)|s)? in object:.*['"`]prompt['"`]/i,
+    /\badditional propert(?:y|ies)\b.*['"`]prompt['"`]/i,
+    /\bmust not have additional propert(?:y|ies)\b.*['"`]prompt['"`]/i,
+    /\b(?:unexpected|extraneous)\s+(?:property|properties|field|fields|key|keys)\b.*['"`]prompt['"`]/i,
+    /\b(?:unknown|invalid)\s+(?:property|properties|field|fields|key|keys)\b.*['"`]prompt['"`]/i,
+    /['"`]prompt['"`].*\b(?:was|is)\s+not allowed\b/i,
+    /"code"\s*:\s*"unrecognized_keys"[^]*"prompt"/i,
+  ],
+} as const;

-function isSessionKeyUnknownFieldValidationMessage(message: string): boolean {
-  return SESSION_KEY_UNKNOWN_FIELD_PATTERNS.some((pattern) => pattern.test(message));
+function isLegacyCompatUnknownFieldValidationMessage(
+  message: string,
+  key: LegacyCompatKey,
+): boolean {
+  return LEGACY_UNKNOWN_FIELD_PATTERNS[key].some((pattern) => pattern.test(message));
 }

-function isSessionKeyCompatibilityError(error: unknown): boolean {
+function isLegacyCompatErrorForKey(error: unknown, key: LegacyCompatKey): boolean {
  for (const candidate of iterateErrorChain(error)) {
    if (Array.isArray(candidate)) {
-      if (candidate.some((entry) => issueRejectsSessionKeyStrictly(entry))) {
+      if (candidate.some((entry) => issueRejectsLegacyCompatKeyStrictly(entry, key))) {
        return true;
      }
      continue;
    }

    if (typeof candidate === "string") {
-      if (isSessionKeyUnknownFieldValidationMessage(candidate)) {
+      if (isLegacyCompatUnknownFieldValidationMessage(candidate, key)) {
        return true;
      }
      continue;
@@ -125,21 +160,21 @@ function isSessionKeyCompatibilityError(error: unknown): boolean {

    if (
      Array.isArray(issueContainer.issues) &&
-      issueContainer.issues.some((issue) => issueRejectsSessionKeyStrictly(issue))
+      issueContainer.issues.some((issue) => issueRejectsLegacyCompatKeyStrictly(issue, key))
    ) {
      return true;
    }

    if (
      Array.isArray(issueContainer.errors) &&
-      issueContainer.errors.some((issue) => issueRejectsSessionKeyStrictly(issue))
+      issueContainer.errors.some((issue) => issueRejectsLegacyCompatKeyStrictly(issue, key))
    ) {
      return true;
    }

    if (
      typeof issueContainer.message === "string" &&
-      isSessionKeyUnknownFieldValidationMessage(issueContainer.message)
+      isLegacyCompatUnknownFieldValidationMessage(issueContainer.message, key)
    ) {
      return true;
    }
@@ -148,25 +183,66 @@ function isSessionKeyCompatibilityError(error: unknown): boolean {
  return false;
 }

-async function invokeWithLegacySessionKeyCompat<TResult, TParams extends SessionKeyCompatParams>(
+function detectRejectedLegacyCompatKeys(
+  error: unknown,
+  allowedKeys: readonly LegacyCompatKey[],
+): Set<LegacyCompatKey> {
+  const rejectedKeys = new Set<LegacyCompatKey>();
+  for (const key of allowedKeys) {
+    if (isLegacyCompatErrorForKey(error, key)) {
+      rejectedKeys.add(key);
+    }
+  }
+  return rejectedKeys;
+}
+
+async function invokeWithLegacyCompat<TResult, TParams extends SessionKeyCompatParams>(
  method: (params: TParams) => Promise<TResult> | TResult,
  params: TParams,
+  allowedKeys: readonly LegacyCompatKey[],
  opts?: {
    onLegacyModeDetected?: () => void;
+    onLegacyKeysDetected?: (keys: Set<LegacyCompatKey>) => void;
+    rejectedKeys?: ReadonlySet<LegacyCompatKey>;
  },
 ): Promise<TResult> {
-  if (!hasOwnSessionKey(params)) {
+  const activeRejectedKeys = new Set(opts?.rejectedKeys ?? []);
+  const availableKeys = allowedKeys.filter((key) => hasOwnLegacyCompatKey(params, key));
+  if (availableKeys.length === 0) {
    return await method(params);
  }

+  let currentParams =
+    activeRejectedKeys.size > 0 ? withoutLegacyCompatKeys(params, activeRejectedKeys) : params;
+
  try {
-    return await method(params);
+    return await method(currentParams);
  } catch (error) {
-    if (!isSessionKeyCompatibilityError(error)) {
-      throw error;
+    let currentError = error;
+    while (true) {
+      const rejectedKeys = detectRejectedLegacyCompatKeys(currentError, availableKeys);
+      let learnedNewKey = false;
+      for (const key of rejectedKeys) {
+        if (!activeRejectedKeys.has(key)) {
+          activeRejectedKeys.add(key);
+          learnedNewKey = true;
+        }
+      }
+
+      if (!learnedNewKey) {
+        throw currentError;
+      }
+
+      opts?.onLegacyModeDetected?.();
+      opts?.onLegacyKeysDetected?.(rejectedKeys);
+      currentParams = withoutLegacyCompatKeys(params, activeRejectedKeys);
+
+      try {
+        return await method(currentParams);
+      } catch (retryError) {
+        currentError = retryError;
+      }
    }
-    opts?.onLegacyModeDetected?.();
-    return await method(withoutSessionKey(params));
  }
 }

@@ -179,6 +255,7 @@ function wrapContextEngineWithSessionKeyCompat(engine: ContextEngine): ContextEn
  }

  let isLegacy = false;
+  const rejectedKeys = new Set<LegacyCompatKey>();
  const proxy: ContextEngine = new Proxy(engine, {
    get(target, property, receiver) {
      if (property === LEGACY_SESSION_KEY_COMPAT) {
@@ -196,13 +273,23 @@ function wrapContextEngineWithSessionKeyCompat(engine: ContextEngine): ContextEn

      return (params: SessionKeyCompatParams) => {
        const method = value.bind(target) as (params: SessionKeyCompatParams) => unknown;
-        if (isLegacy && hasOwnSessionKey(params)) {
-          return method(withoutSessionKey(params));
+        const allowedKeys = LEGACY_COMPAT_METHOD_KEYS[property];
+        if (
+          isLegacy &&
+          allowedKeys.some((key) => rejectedKeys.has(key) && hasOwnLegacyCompatKey(params, key))
+        ) {
+          return method(withoutLegacyCompatKeys(params, rejectedKeys));
        }
-        return invokeWithLegacySessionKeyCompat(method, params, {
+        return invokeWithLegacyCompat(method, params, allowedKeys, {
          onLegacyModeDetected: () => {
            isLegacy = true;
          },
+          onLegacyKeysDetected: (keys) => {
+            for (const key of keys) {
+              rejectedKeys.add(key);
+            }
+          },
+          rejectedKeys,
        });
      };
    },
--- a/src/context-engine/types.ts
+++ b/src/context-engine/types.ts
@@ -183,6 +183,8 @@ export interface ContextEngine {
    /** Current model identifier (e.g. "claude-opus-4", "gpt-4o", "qwen2.5-7b").
     *  Allows context engine plugins to adapt formatting per model. */
    model?: string;
+    /** The incoming user prompt for this turn (useful for retrieval-oriented engines). */
+    prompt?: string;
  }): Promise<AssembleResult>;

  /**