omp-cache-optimizer 1.0.1 → 1.0.2

package/index.ts

@@ -312,6 +312,15 @@ type CacheUsageSample = {
   missingUsageFields: boolean;
 };
 
+type PromptRewriteContext = {
+  options?: BuildSystemPromptOptions;
+  routeSnapshot?: PiRouteSnapshot;
+  routedModel?: PiModel;
+  timestamp: number;
+};
+
+const PROMPT_REWRITE_CONTEXT_TTL_MS = 10_000;
+
 /** Maximum number of recent samples kept per model key (in-memory only, not persisted). */
 const MAX_RECENT_SAMPLES = 50;
 
@@ -968,16 +977,42 @@ function getNonNegativeNumber(record: UnknownRecord, key: string): number | unde
  */
 function getCompat(model: PiModel | undefined): CacheCompat {
   if (!model) return {} as CacheCompat;
-  
-  // The host runtime merges provider.compat with model.compat (model wins on conflicts).
-  // We approximate this by reading from ctx.model which should already have merged compat.
-  // However, for safety, we check both levels if available.
-  const modelCompat = (model.compat ?? {}) as CacheCompat;
 
-  // Note: ctx.model from the host runtime should already contain merged compat,
-  // but we document the two-level structure for clarity.
-  // but we document the two-level structure for clarity
-  return modelCompat;
+  const record = model as PiModel & { compatConfig?: Record<string, unknown> };
+  return {
+    ...((record.compatConfig ?? {}) as CacheCompat),
+    ...((record.compat ?? {}) as CacheCompat),
+  };
+}
+
+function makePromptRewriteContextKey(sessionHash: string | undefined, model: PiModel | undefined): string | undefined {
+  if (!sessionHash || !model) return undefined;
+  return `${sessionHash}:${modelKey(model)}`;
+}
+
+function rememberPromptRewriteContext(
+  contexts: Map<string, PromptRewriteContext>,
+  key: string | undefined,
+  context: PromptRewriteContext,
+): void {
+  if (!key) return;
+  contexts.set(key, context);
+}
+
+function getPromptRewriteContext(
+  contexts: Map<string, PromptRewriteContext>,
+  key: string | undefined,
+  now = Date.now(),
+  ttlMs = PROMPT_REWRITE_CONTEXT_TTL_MS,
+): PromptRewriteContext | undefined {
+  if (!key) return undefined;
+  const context = contexts.get(key);
+  if (!context) return undefined;
+  if (now - context.timestamp > ttlMs) {
+    contexts.delete(key);
+    return undefined;
+  }
+  return context;
 }
 
 /**
@@ -1972,13 +2007,6 @@ function setSystemPrompt(payload: unknown, text: string): boolean {
     return true;
   }
   if (Array.isArray(record.system) && record.system.length > 0) {
-    // Replace first text block, keep structure
-    const first = asRecord(record.system[0]);
-    if (first && typeof first.text === "string") {
-      first.text = text;
-      return true;
-    }
-    // Fallback: convert to single-block string form
     record.system = [{ type: "text", text }];
     return true;
   }
@@ -1986,11 +2014,8 @@ function setSystemPrompt(payload: unknown, text: string): boolean {
   // google-generative-ai: payload.systemInstruction
   const systemInstruction = asRecord(record.systemInstruction);
   if (systemInstruction && Array.isArray(systemInstruction.parts) && systemInstruction.parts.length > 0) {
-    const firstPart = asRecord(systemInstruction.parts[0]);
-    if (firstPart && typeof firstPart.text === "string") {
-      firstPart.text = text;
-      return true;
-    }
+    systemInstruction.parts = [{ text }];
+    return true;
   }
 
   // openai-completions / openai-responses: payload.messages[] first system/developer message
@@ -2005,11 +2030,8 @@ function setSystemPrompt(payload: unknown, text: string): boolean {
           return true;
         }
         if (Array.isArray(r.content) && r.content.length > 0) {
-          const first = asRecord(r.content[0]);
-          if (first && typeof first.text === "string") {
-            first.text = text;
-            return true;
-          }
+          r.content = text;
+          return true;
         }
       }
     }
@@ -5422,6 +5444,10 @@ export const __internals_for_tests = {
   hashSessionId,
   makeSessionModelKey,
   modelKeyFromSessionKey,
+  makePromptRewriteContextKey,
+  rememberPromptRewriteContext,
+  getPromptRewriteContext,
+  PROMPT_REWRITE_CONTEXT_TTL_MS,
   filterRestorableStatsForSession,
   parsePersistedRoutedModelRef,
   routedModelRefToPiModel,
@@ -5495,11 +5521,10 @@ export default function (pi: ExtensionAPI) {
   let latestCacheHint: PiCacheHintSnapshot | undefined;
   // OMP divergence: prompt rewriting moved from before_agent_start to
   // before_provider_request (OMP's before_agent_start can only inject messages,
-  // not mutate systemPrompt). We cache systemPromptOptions + route snapshot here
-  // so before_provider_request can apply the 3-step pipeline to the payload.
-  let pendingPromptOptions: BuildSystemPromptOptions | undefined;
-  let pendingRouteSnapshot: PiRouteSnapshot | undefined;
-  let pendingRoutedModel: PiModel | undefined;
+  // not mutate systemPrompt). Store prompt options per session/model so an
+  // overlapping turn or sub-agent cannot overwrite another request's rewrite
+  // context before before_provider_request fires.
+  const promptRewriteContexts = new Map<string, PromptRewriteContext>();
   const PERSIST_DEBOUNCE_MS = 2000;
   /** In-memory recent usage samples per model key (not persisted, cleared on reload). */
   const recentSamplesByModelKey = new Map<string, CacheUsageSample[]>();
@@ -5913,17 +5938,18 @@ export default function (pi: ExtensionAPI) {
       ? findModelInRegistry(_ctx.modelRegistry, routeSnapshot.provider, routeSnapshot.modelId) ?? routeSnapshotToPiModel(routeSnapshot, _ctx.model)
       : undefined;
 
-    // OMP divergence: before_agent_start in OMP can only inject messages (return
-    // { message }), NOT mutate systemPrompt. We cache the prompt options + route
-    // snapshot here so before_provider_request can apply the 3-step pipeline to
-    // the provider payload. If OMP does not supply systemPromptOptions, skill
-    // compression and stable-prefix reorder are skipped (only churn strip runs).
     const eventRecord = asRecord(event);
-    pendingPromptOptions = (eventRecord?.systemPromptOptions as BuildSystemPromptOptions | undefined) ?? undefined;
-    pendingRouteSnapshot = routeSnapshot;
-    pendingRoutedModel = routedModel ?? _ctx.model;
-
+    const options = (eventRecord?.systemPromptOptions as BuildSystemPromptOptions | undefined) ?? undefined;
     const model = routedModel ?? _ctx.model;
+    const contextKey = makePromptRewriteContextKey(sessionHashFromContext(_ctx), model);
+    rememberPromptRewriteContext(promptRewriteContexts, contextKey, {
+      options,
+      routeSnapshot,
+      routedModel: model,
+      timestamp: Date.now(),
+    });
+
+    const modelForHint = model;
     const promptCacheKey = getSessionPromptCacheKey(_ctx);
     const cacheRetention = process.env[PI_CACHE_RETENTION_ENV] === LONG_CACHE_RETENTION_VALUE ? LONG_CACHE_RETENTION_VALUE : undefined;
     const rawSystemPrompt = typeof eventRecord?.systemPrompt === "string" ? eventRecord.systemPrompt : "";
@@ -5931,9 +5957,9 @@ export default function (pi: ExtensionAPI) {
       sessionIdHash: currentSessionHashSet ? currentSessionHash : sessionHashFromContext(_ctx),
       virtualProvider: routeSnapshot?.virtualProvider ?? _ctx.model?.provider,
       virtualModelId: routeSnapshot?.virtualModelId ?? _ctx.model?.id,
-      upstreamProvider: routeSnapshot?.provider ?? model?.provider,
-      upstreamModelId: routeSnapshot?.modelId ?? model?.id,
-      api: model?.api,
+      upstreamProvider: routeSnapshot?.provider ?? modelForHint?.provider,
+      upstreamModelId: routeSnapshot?.modelId ?? modelForHint?.id,
+      api: modelForHint?.api,
       systemPrompt: rawSystemPrompt,
       promptCacheKey,
       cacheRetention,
@@ -5961,21 +5987,24 @@ export default function (pi: ExtensionAPI) {
       requestModel &&
       !isResponsesPromptRewriteBypassApi(requestModel.api)
     ) {
+      const contextKey = makePromptRewriteContextKey(sessionHashFromContext(ctx), requestModel);
+      const rewriteContext = getPromptRewriteContext(promptRewriteContexts, contextKey);
+      const promptOptions = rewriteContext?.options;
       const original = extractSystemPrompt(resultPayload);
       if (original && original.trim().length > 0) {
         // Step 1: strip per-turn churn from <session-overview>.
         const stripped = stripSessionOverviewChurn(original);
 
         // Step 2: compress skills XML → one-line index (requires cached options).
-        const compressed = pendingPromptOptions
-          ? compressSkillsInSystemPrompt(stripped, pendingPromptOptions)
+        const compressed = promptOptions
+          ? compressSkillsInSystemPrompt(stripped, promptOptions)
           : stripped;
 
         // Step 3: lift stable content above dynamic content (requires cached options).
         let finalPrompt = compressed;
         let changed = false;
-        if (pendingPromptOptions) {
-          const optimized = optimizeSystemPrompt(compressed, pendingPromptOptions);
+        if (promptOptions) {
+          const optimized = optimizeSystemPrompt(compressed, promptOptions);
           if (optimized.changed && optimized.systemPrompt.trim().length > 0) {
             finalPrompt = optimized.systemPrompt;
             changed = true;

package/package.json

@@ -1,6 +1,6 @@
 {
   "name": "omp-cache-optimizer",
-  "version": "1.0.1",
+  "version": "1.0.2",
   "description": "Improve OMP prompt/KV cache hit rates with stable prompts, OpenAI-compatible cache keys, proxy compat warnings, and footer cache stats.",
   "keywords": [
     "omp-package",