Fix context prompt snapshot ordering

Author: Wistomize

apps/api/src/evidence-reference-context.ts

@@ -4,10 +4,9 @@ import type {
   EvidenceResolutionIssue
 } from "@datafoundry/contracts";
 import {
-  createContextItem,
-  createContextSourceMetadata,
-  type ContextItem,
-  type RuntimeContextSource
+  createAgentContextItem,
+  createAgentContextSourceMetadata,
+  type AgentContextItem
 } from "@datafoundry/agent-runtime";
 import type {
   MetadataStore,
@@ -27,14 +26,14 @@ type ResolveEvidenceReferenceContextInput = {
 
 export type ResolvedEvidenceReferenceContext = {
   diagnostics: EvidenceResolutionDiagnostics;
-  items: ContextItem[];
+  items: AgentContextItem[];
 };
 
 /** Resolves client EvidenceRef handles into server-authoritative ContextItems for prompt assembly. */
 export function resolveEvidenceReferenceContext(
   input: ResolveEvidenceReferenceContextInput
 ): ResolvedEvidenceReferenceContext {
-  const items: ContextItem[] = [];
+  const items: AgentContextItem[] = [];
   const accepted: string[] = [];
   const dropped: EvidenceResolutionIssue[] = [];
   const seen = new Set<string>();
@@ -58,21 +57,10 @@ export function resolveEvidenceReferenceContext(
   return { diagnostics: { accepted, dropped }, items };
 }
 
-/** Creates a runtime context source from already-resolved evidence ContextItems. */
-export function createEvidenceRuntimeSource(items: ContextItem[]): RuntimeContextSource | undefined {
-  if (items.length === 0) {
-    return undefined;
-  }
-  return {
-    sourceType: "evidence-focus",
-    collect: () => items
-  };
-}
-
 const resolveEvidenceRef = (
   input: ResolveEvidenceReferenceContextInput,
   ref: EvidenceRef
-): { items: ContextItem[]; issue?: never } | { items?: never; issue: EvidenceResolutionIssue } => {
+): { items: AgentContextItem[]; issue?: never } | { items?: never; issue: EvidenceResolutionIssue } => {
   try {
     if (ref.source.artifactId) {
       return { items: artifactEvidenceItems(input, ref, ref.source.artifactId) };
@@ -99,7 +87,7 @@ const artifactEvidenceItems = (
   input: ResolveEvidenceReferenceContextInput,
   ref: EvidenceRef,
   artifactId: string
-): ContextItem[] => {
+): AgentContextItem[] => {
   const artifact = input.metadataStore.artifacts.get({ user_id: input.userId, artifact_id: artifactId });
   assertSession(ref, artifact.session_id);
   const preview = parseJsonValue(artifact.preview_json);
@@ -125,7 +113,7 @@ const sqlAuditEvidenceItems = (
   input: ResolveEvidenceReferenceContextInput,
   ref: EvidenceRef,
   auditLogId: string
-): ContextItem[] => {
+): AgentContextItem[] => {
   const audit = input.metadataStore.sqlAuditLogs.get({ user_id: input.userId, audit_log_id: auditLogId });
   if (audit.run_id) {
     const run = input.metadataStore.runs.get({ user_id: input.userId, run_id: audit.run_id });
@@ -154,7 +142,7 @@ const sqlAuditEvidenceItems = (
 const toolEventEvidenceItems = (
   input: ResolveEvidenceReferenceContextInput,
   ref: EvidenceRef
-): ContextItem[] => {
+): AgentContextItem[] => {
   const run = requireEvidenceRun(input, ref);
   const events = input.metadataStore.runEvents.listByRun({ user_id: input.userId, run_id: run.id });
   const event = findToolEvent(events, ref);
@@ -180,11 +168,11 @@ const toolEventEvidenceItems = (
 const createEvidenceItem = (
   input: ResolveEvidenceReferenceContextInput,
   ref: EvidenceRef,
-  visibility: ContextItem["visibility"],
+  visibility: AgentContextItem["visibility"],
   content: unknown,
-  trust: ContextItem["trust"],
+  trust: AgentContextItem["trust"],
   datasourceId?: string
-): ContextItem => createContextItem({
+): AgentContextItem => createAgentContextItem({
   id: `evidence:${ref.id}:${visibility}`,
   sourceType: "evidence-focus",
   sourceId: ref.id,
@@ -194,7 +182,7 @@ const createEvidenceItem = (
   retention: visibility === "model" ? "active" : "reference",
   priority: visibility === "model" ? 70 : 65,
   content,
-  metadata: createContextSourceMetadata({
+  metadata: createAgentContextSourceMetadata({
     dedupeKeys: evidenceDedupeKeys(ref),
     exclusivityKey: `evidence:${ref.id}`,
     overlapKeys: evidenceOverlapKeys(ref),

apps/api/src/run-agent-assembly.ts

@@ -5,9 +5,9 @@ import {
   createDataFoundry,
   createDataFoundryRunContext,
   type AgentRunContext,
+  type AgentContextItem,
   type AgUiEventEmitter,
   type GoalRuntimeAdapter,
-  type RuntimeContextSource,
   type TaskStateRuntime,
   type WorkspaceAttachment
 } from "@datafoundry/agent-runtime";
@@ -54,7 +54,7 @@ type CreateRunAgentAssemblyInput = {
   dataGateway: DataGateway;
   effectiveRunConfig: EffectiveRunConfig;
   emitter: AgUiEventEmitter;
-  evidenceRuntimeSources?: RuntimeContextSource[] | undefined;
+  evidenceContextItems?: AgentContextItem[] | undefined;
   fileAssetService: FileAssetService;
   goal?: EffectiveRunConfig["goal"] | undefined;
   interactionResume?: InteractionResume | undefined;
@@ -143,9 +143,7 @@ export const createRunAgentAssembly = async (
     ...(input.modelContextProfile ? { modelContextProfile: input.modelContextProfile } : {}),
     modelProvider: input.modelProvider,
     ...(input.modelSettings ? { modelSettings: input.modelSettings } : {}),
-    ...(input.evidenceRuntimeSources?.length
-      ? { additionalRuntimeSources: input.evidenceRuntimeSources }
-      : {}),
+    ...(input.evidenceContextItems?.length ? { evidenceContextItems: input.evidenceContextItems } : {}),
     ...(input.longTermMemories.length > 0 ? { longTermMemory: { records: input.longTermMemories } } : {}),
     runContext: input.runContext,
     selectedSkills: input.selectedSkills,

apps/api/src/run-config-resolver.ts

@@ -2,7 +2,7 @@ import type { RunAgentInput } from "@ag-ui/client";
 import {
   createModelProviderFromEnv,
   createModelProviderFromProfile,
-  type ModelContextProfile
+  type AgentModelContextProfile
 } from "@datafoundry/agent-runtime";
 import type { MetadataStore } from "@datafoundry/metadata";
 import {
@@ -26,7 +26,7 @@ export type ResolvedRunConfig = {
     temperature?: number;
     topP?: number;
   };
-  modelContextProfile?: ModelContextProfile;
+  modelContextProfile?: AgentModelContextProfile;
   reasoningModel?: boolean;
   runTimeoutMs?: number;
   skillSelection: SkillSelectionResult;
@@ -377,7 +377,7 @@ const resolveModelContextProfile = (
   metadataStore: MetadataStore,
   userId: string,
   workspaceId: string
-): ModelContextProfile | undefined => {
+): AgentModelContextProfile | undefined => {
   if (!profileId || profileId === "server-default") {
     return undefined;
   }

apps/api/src/server.ts

@@ -48,10 +48,7 @@ import {
   sendAuthError
 } from "./auth/routes.js";
 import { persistCurrentUserMessage } from "./conversation-memory.js";
-import {
-  createEvidenceRuntimeSource,
-  resolveEvidenceReferenceContext
-} from "./evidence-reference-context.js";
+import { resolveEvidenceReferenceContext } from "./evidence-reference-context.js";
 import { createRunAgentAssembly, createRunAgentContext } from "./run-agent-assembly.js";
 import { resolveRunConfig } from "./run-config-resolver.js";
 import { resolveRunIdentity } from "./run-identity-orchestrator.js";
@@ -524,7 +521,6 @@ class DataFoundryAgUiAgent extends AbstractAgent {
           userId: this.input.user.id,
           workspaceId: this.input.workspaceId
         });
-        const evidenceRuntimeSource = createEvidenceRuntimeSource(evidenceContext.items);
         const taskPlanProjector = new TaskPlanProjector(runContext);
         const toolCallResultBridge = new ToolCallResultBridge();
         const runAbortController = new AbortController();
@@ -552,7 +548,7 @@ class DataFoundryAgUiAgent extends AbstractAgent {
           dataGateway: this.input.dataGateway,
           artifactService: this.input.artifactService,
           effectiveRunConfig,
-          ...(evidenceRuntimeSource ? { evidenceRuntimeSources: [evidenceRuntimeSource] } : {}),
+          ...(evidenceContext.items.length ? { evidenceContextItems: evidenceContext.items } : {}),
           fileAssetService: this.input.fileAssetService,
           emitter: { emit },
           ...(effectiveRunConfig.goal ? { goal: effectiveRunConfig.goal } : {}),

…_tests__/dashscope-prompt-compat.test.ts …non-empty-message-content-compat.test.tsapps/web/src/app/data-tasks/__tests__/dashscope-prompt-compat.test.ts renamed to apps/web/src/app/data-tasks/__tests__/non-empty-message-content-compat.test.ts apps/web/src/app/data-tasks/__tests__/dashscope-prompt-compat.test.ts renamed to apps/web/src/app/data-tasks/__tests__/non-empty-message-content-compat.test.ts

@@ -1,18 +1,20 @@
 import { describe, expect, it } from "vitest";
 
 import {
-  ensureDashScopeCompatiblePrompt,
-  shouldApplyDashScopePromptCompat,
-} from "../../../../../../packages/agent-runtime/src/context/dashscope-prompt-compat";
-
-describe("dashscope prompt compat", () => {
-  it("enables compat for openai-compatible providers", () => {
-    expect(shouldApplyDashScopePromptCompat("openai-compatible")).toBe(true);
-    expect(shouldApplyDashScopePromptCompat("unsupported-provider")).toBe(false);
+  ensureNonEmptyMessageContentPrompt,
+  shouldApplyNonEmptyMessageContentCompat,
+} from "../../../../../../packages/agent-runtime/src/provider-compat/non-empty-message-content-compat";
+
+describe("non-empty message content compat", () => {
+  it("enables compat only for providers that require non-empty message content", () => {
+    expect(shouldApplyNonEmptyMessageContentCompat({
+      prompt_compat: { requires_non_empty_message_content: true },
+    })).toBe(true);
+    expect(shouldApplyNonEmptyMessageContentCompat({})).toBe(false);
   });
 
   it("adds placeholder text to assistant messages that only contain tool calls", () => {
-    const prompt = ensureDashScopeCompatiblePrompt([
+    const prompt = ensureNonEmptyMessageContentPrompt([
       {
         role: "assistant",
         content: [
@@ -30,7 +32,7 @@ describe("dashscope prompt compat", () => {
   });
 
   it("adds placeholder text to assistant messages that only contain reasoning", () => {
-    const prompt = ensureDashScopeCompatiblePrompt([
+    const prompt = ensureNonEmptyMessageContentPrompt([
       {
         role: "assistant",
         content: [{ type: "reasoning", text: "thinking..." }],
@@ -41,7 +43,7 @@ describe("dashscope prompt compat", () => {
   });
 
   it("fills empty tool result outputs", () => {
-    const prompt = ensureDashScopeCompatiblePrompt([
+    const prompt = ensureNonEmptyMessageContentPrompt([
       {
         role: "tool",
         content: [

packages/agent-runtime/src/context/inventory/context-run-state.ts

@@ -42,6 +42,19 @@ export class ContextRunState {
     return this.currentPackage;
   }
 
+  replaceMatchingItems(shouldReplace: (item: ContextItem) => boolean, contextPackage: ContextPackage): ContextPackage {
+    const retainedItems = this.currentPackage.items.filter((item) => !shouldReplace(item));
+    const incomingItems = contextPackage.items.filter(shouldReplace);
+    this.currentPackage = this.builder.build([...retainedItems, ...incomingItems], {
+      packageId: this.currentPackage.packageId,
+      revision: this.currentPackage.revision + 1,
+      resourceId: this.identity.resourceId,
+      sessionId: this.identity.sessionId,
+      runId: this.identity.runId
+    });
+    return this.currentPackage;
+  }
+
   replaceSourceItems(sourceTypes: string[], contextPackage: ContextPackage): ContextPackage {
     const sourceTypeSet = new Set(sourceTypes);
     const retainedItems = this.currentPackage.items.filter((item) => !sourceTypeSet.has(item.sourceType));

packages/agent-runtime/src/context/protocol/mastra/mastra-context-budget-processor.ts

@@ -1,6 +1,7 @@
 import type { ProcessInputStepArgs, ProcessInputStepResult, Processor } from "@mastra/core/processors";
 
 import { ContextPackageBuilder } from "../../inventory/context-package-builder.js";
+import type { ContextItem } from "../../inventory/context-item.js";
 import type { ContextRunState } from "../../inventory/context-run-state.js";
 import { ContextSourcePolicy } from "../../policy/context-source-policy.js";
 import { isContextSourceOmissionDecision } from "../../policy/context-source-policy.js";
@@ -19,6 +20,7 @@ import {
   type MastraContextProtocolOutput
 } from "./mastra-context-protocol-adapter.js";
 import { createMastraConversationContextItems } from "./mastra-conversation-context-adapter.js";
+import { applyMastraPromptSnapshot } from "./mastra-prompt-snapshot-applicator.js";
 import type { MastraToolObservationRouter } from "./mastra-tool-observation-router.js";
 
 export type MastraContextBudgetProcessorOptions = {
@@ -58,7 +60,7 @@ export class MastraContextBudgetProcessor implements Processor<"context-budget">
         runId: this.options.runState.identity.runId
       }
     );
-    const contextPackage = this.options.runState.merge(livePackage);
+    const contextPackage = this.options.runState.replaceMatchingItems(isProtocolMessageItem, livePackage);
     const sourcePolicyResult = this.sourcePolicy.applyPackage(contextPackage);
     const groupPlan = this.materializer.createGroupPlan({
       contextPackage,
@@ -93,10 +95,27 @@ export class MastraContextBudgetProcessor implements Processor<"context-budget">
       createMastraContextCompiledEventPayload(contextPackage, result.plan, this.options.modelName)
     );
     const protocolOutput = this.protocolAdapter.toProtocol(promptView);
+    const systemMessages = protocolOutput.systemMessages as NonNullable<ProcessInputStepResult["systemMessages"]>;
+
+    if (!args.messageList) {
+      return {
+        messages: protocolOutput.messages,
+        systemMessages
+      };
+    }
+
+    applyMastraPromptSnapshot(args.messageList, { messages: protocolOutput.messages });
 
     return {
-      messages: protocolOutput.messages,
-      systemMessages: protocolOutput.systemMessages as NonNullable<ProcessInputStepResult["systemMessages"]>
+      messageList: args.messageList,
+      systemMessages
     };
   }
 }
+
+const PROTOCOL_MESSAGE_KINDS = new Set(["message", "source-message", "system"]);
+
+const isProtocolMessageItem = (item: ContextItem): boolean => {
+  const messageKind = item.metadata.messageKind;
+  return typeof messageKind === "string" && PROTOCOL_MESSAGE_KINDS.has(messageKind);
+};

packages/agent-runtime/src/context/protocol/mastra/mastra-context-processor-boundary.ts

@@ -7,7 +7,7 @@ import {
 } from "../../policy/context-reduction-strategy.js";
 import { createDefaultContextSourcePolicy } from "../../policy/context-source-authority-profile.js";
 import { ContextStepPlanner } from "../../policy/context-step-planner.js";
-import { ModelContextProfileRegistry } from "../../policy/model-context-profile.js";
+import { ModelContextProfileRegistry, type ModelContextProfile } from "../../policy/model-context-profile.js";
 import { PromptTokenCounter } from "../../policy/prompt-token-counter.js";
 import { ContextPromptMaterializer } from "../../projection/context-prompt-materializer.js";
 import type { ContextSourcePromptMaterializer } from "../../projection/context-source-prompt-materializer.js";
@@ -42,6 +42,7 @@ export type CreateMastraContextProcessorBoundaryInput = {
   dispatcher: ToolObservationDispatcher;
   eventSink: ContextProtocolEventSink;
   longTermMemory?: CreateDefaultRuntimeContextSourceRegistryInput["longTermMemory"];
+  modelContextProfile?: ModelContextProfile;
   modelName: string | undefined;
   runScope: RuntimeContextRunScope;
   runState: ContextRunState;
@@ -56,7 +57,9 @@ export type MastraContextProcessorBoundary = {
 export const createMastraContextProcessorBoundary = (
   input: CreateMastraContextProcessorBoundaryInput
 ): MastraContextProcessorBoundary => {
-  const profileRegistry = input.contextCompilation?.profileRegistry ?? new ModelContextProfileRegistry();
+  const profileRegistry = input.contextCompilation?.profileRegistry ?? new ModelContextProfileRegistry({
+    ...(input.modelContextProfile ? { defaultProfile: input.modelContextProfile } : {})
+  });
   const tokenCounter = input.contextCompilation?.tokenCounter ?? new PromptTokenCounter();
   const planner = new ContextStepPlanner({
     profileRegistry,

packages/agent-runtime/src/context/protocol/mastra/mastra-prompt-snapshot-applicator.ts

@@ -0,0 +1,28 @@
+import type { MastraDBMessage } from "@mastra/core/agent";
+import type { ProcessInputStepArgs } from "@mastra/core/processors";
+
+export type MastraPromptSnapshot = {
+  messages: MastraDBMessage[];
+};
+
+type MastraMessageList = NonNullable<ProcessInputStepArgs["messageList"]>;
+
+// Mastra applies returned message arrays with per-message remove/add. Context compilation
+// produces a complete prompt snapshot, so apply it atomically to avoid assistant-message merging.
+export const applyMastraPromptSnapshot = (
+  messageList: MastraMessageList,
+  snapshot: MastraPromptSnapshot
+): void => {
+  const sourceCheck = messageList.makeMessageSourceChecker();
+  const currentIds = messageList.get.all.db().map((message) => message.id).filter(isString);
+
+  if (currentIds.length > 0) {
+    messageList.removeByIds(currentIds);
+  }
+
+  for (const message of snapshot.messages) {
+    messageList.add(message, sourceCheck.getSource(message) ?? "input");
+  }
+};
+
+const isString = (value: unknown): value is string => typeof value === "string";