Make memory service canonical JSON locale-independent and lossless

ts/packages/memory/service/src/agentEditionValidation.ts

@@ -1,6 +1,7 @@
 // Copyright (c) Microsoft Corporation.
 // Licensed under the MIT License.
 
+import { canonicalJsonValue } from "./canonicalJson.js";
 import {
     normalizeRunbookBindingArguments,
     validateRunbookArgumentReadiness,
@@ -139,20 +140,8 @@ export interface AgentEdition {
 
 export const agentEditionMarkdownMarker = "<!-- typeagent-agent-edition:1 -->";
 
-function sortedJson(value: unknown): unknown {
-    if (Array.isArray(value)) return value.map(sortedJson);
-    if (value !== null && typeof value === "object") {
-        return Object.fromEntries(
-            Object.entries(value as Record<string, unknown>)
-                .sort(([left], [right]) => left.localeCompare(right))
-                .map(([key, child]) => [key, sortedJson(child)]),
-        );
-    }
-    return value;
-}
-
 export function canonicalizeProcedure(value: unknown): string {
-    return `${JSON.stringify(sortedJson(value), undefined, 2)}\n`;
+    return `${JSON.stringify(canonicalJsonValue(value), undefined, 2)}\n`;
 }
 
 export function getProcedureEvidenceReferences(

ts/packages/memory/service/src/agentSessionParsing.ts

@@ -9,6 +9,11 @@ import type {
     AgentSessionImportRequest,
     AgentSessionParseResult,
 } from "./agentSessionTypes.js";
+import {
+    canonicalJsonValue,
+    codePointOrder,
+    uniqueKey,
+} from "./canonicalJson.js";
 import { redactRunbookText, redactRunbookValue } from "./runbookRedaction.js";
 import { validateIsoTimestamp } from "./timestampValidation.js";
 
@@ -139,19 +144,18 @@ function timestamp(value: unknown): string {
     return text;
 }
 
+// Set only while legacyAgentSessionEventDigests re-parses a chunk. Parsing is
+// synchronous, so no other parse can observe it.
+let keyCollision: "suffix" | "overwrite" = "suffix";
+
+// Event text feeds agentSessionEventDigest, which sessions store to detect
+// conflicting re-sent events, so this keeps the original code-point key order.
 function sortedValue(value: unknown): unknown {
-    if (Array.isArray(value)) return value.map(sortedValue);
-    if (value !== null && typeof value === "object") {
-        return Object.fromEntries(
-            Object.entries(value)
-                .sort(([a], [b]) => (a < b ? -1 : a > b ? 1 : 0))
-                .map(([key, child]) => [
-                    redactRunbookText(key),
-                    sortedValue(child),
-                ]),
-        );
-    }
-    return value;
+    return canonicalJsonValue(value, {
+        keyOrder: codePointOrder,
+        mapKey: redactRunbookText,
+        keyCollision,
+    });
 }
 
 function redactedStrings(value: unknown, depth: number): unknown {
@@ -161,11 +165,18 @@ function redactedStrings(value: unknown, depth: number): unknown {
         return value.map((child) => redactedStrings(child, depth + 1));
     }
     if (value !== null && typeof value === "object") {
+        // Keep key order: redactedText compares this with the parsed input.
+        const used = new Set<string>();
         return Object.fromEntries(
-            Object.entries(value).map(([key, child]) => [
-                redactRunbookText(key),
-                redactedStrings(child, depth + 1),
-            ]),
+            Object.entries(value).map(([key, child]) => {
+                const redacted = redactRunbookText(key);
+                return [
+                    keyCollision === "overwrite"
+                        ? redacted
+                        : uniqueKey(used, redacted),
+                    redactedStrings(child, depth + 1),
+                ];
+            }),
         );
     }
     return value;
@@ -563,6 +574,30 @@ export function parseAgentSessionChunk(
     return result;
 }
 
+/**
+ * Event digests as computed before colliding redacted keys were kept, when a
+ * later key silently replaced an earlier one. Used only to recognize events
+ * stored by that version when they are re-sent.
+ * TODO: remove once sessions recorded with lossy key redaction are gone.
+ */
+export function legacyAgentSessionEventDigests(
+    request: AgentSessionImportRequest,
+): Map<string, string> {
+    keyCollision = "overwrite";
+    try {
+        return new Map(
+            parseAgentSessionChunk(request).events.map((event) => [
+                event.id,
+                agentSessionEventDigest(event),
+            ]),
+        );
+    } catch {
+        return new Map();
+    } finally {
+        keyCollision = "suffix";
+    }
+}
+
 /** Group explicit tasks/agents without inferring causality from adjacent events. */
 export function buildAgentSessionEpisodes(
     events: AgentSessionEvent[],

ts/packages/memory/service/src/agentSessionStore.ts

@@ -4,11 +4,13 @@
 import { createHash } from "node:crypto";
 import { lstat, readFile, readdir, rm } from "node:fs/promises";
 import path from "node:path";
+import { canonicalJson } from "./canonicalJson.js";
 import { writeRunbookJson } from "./durableRunbookJson.js";
 import { redactRunbookText } from "./runbookRedaction.js";
 import {
     agentSessionEventDigest,
     buildAgentSessionEpisodes,
+    legacyAgentSessionEventDigests,
     parseAgentSessionChunk,
 } from "./agentSessionParsing.js";
 import {
@@ -84,7 +86,13 @@ interface SessionHost {
 }
 
 function digest(value: unknown): string {
-    return createHash("sha256").update(JSON.stringify(value)).digest("hex");
+    return createHash("sha256").update(canonicalJson(value)).digest("hex");
+}
+
+// Request hashes stored before canonical hashing depend on key order.
+// TODO: remove once sessions recorded with key-order-dependent hashes are gone.
+function legacyRequestHash(request: AgentSessionImportRequest): string {
+    return createHash("sha256").update(JSON.stringify(request)).digest("hex");
 }
 
 function validateIdentifier(value: string): void {
@@ -338,7 +346,11 @@ export class AgentSessionStore {
         const requestKey = digest(request.idempotencyKey);
         const requestHash = digest(request);
         const prior = stored.requests.find((entry) => entry.key === requestKey);
-        if (prior && prior.hash !== requestHash)
+        if (
+            prior &&
+            prior.hash !== requestHash &&
+            prior.hash !== legacyRequestHash(request)
+        )
             throw new Error(
                 "Agent session idempotency key has different content",
             );
@@ -389,14 +401,18 @@ export class AgentSessionStore {
         if (request.expectedCursor !== stored.receipt.cursor)
             throw new Error("Agent session continuation cursor changed");
         const seen: Record<string, string> = {};
+        let legacyDigests: Map<string, string> | undefined;
         const events = parsed.events.filter((event) => {
             const eventKey = digest(event.id);
             const value = agentSessionEventDigest(event);
             const previous = stored.seen[eventKey];
-            if (previous && previous !== value)
-                throw new Error(
-                    "Agent session event identity has conflicting content",
-                );
+            if (previous && previous !== value) {
+                legacyDigests ??= legacyAgentSessionEventDigests(request);
+                if (legacyDigests.get(event.id) !== previous)
+                    throw new Error(
+                        "Agent session event identity has conflicting content",
+                    );
+            }
             if (previous) {
                 parsed.duplicates++;
                 return false;

ts/packages/memory/service/src/canonicalJson.ts

@@ -0,0 +1,68 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+export type KeyOrder = (left: string, right: string) => number;
+
+export function codePointOrder(left: string, right: string): number {
+    return left < right ? -1 : left > right ? 1 : 0;
+}
+
+// A fixed "en" collator instead of the process default locale, so canonical
+// text and the hashes stored from it do not depend on the host's LANG.
+// Existing stored hashes were produced with en-locale localeCompare, which this
+// matches. Code-point order breaks ties between keys that collate as equal
+// (for example, the same text in different Unicode normalization forms).
+const enCollator = new Intl.Collator("en");
+
+export function collatedOrder(left: string, right: string): number {
+    return enCollator.compare(left, right) || codePointOrder(left, right);
+}
+
+export interface CanonicalJsonOptions {
+    keyOrder?: KeyOrder;
+    // Rewrites each key after sorting by the original key.
+    mapKey?: (key: string) => string;
+    // "overwrite" reproduces Object.fromEntries: a later key that maps to the
+    // same text replaces the earlier value. Only for recomputing old digests.
+    keyCollision?: "suffix" | "overwrite";
+}
+
+/**
+ * Distinct keys can map to the same text (two redacted secrets). Suffix the
+ * later one instead of silently dropping its value.
+ */
+export function uniqueKey(used: Set<string>, key: string): string {
+    let unique = key;
+    for (let n = 2; used.has(unique); n++) unique = `${key} (${n})`;
+    used.add(unique);
+    return unique;
+}
+
+/** Recursively sorts object keys so equal values serialize identically. */
+export function canonicalJsonValue(
+    value: unknown,
+    options: CanonicalJsonOptions = {},
+): unknown {
+    if (Array.isArray(value))
+        return value.map((item) => canonicalJsonValue(item, options));
+    if (value === null || typeof value !== "object") return value;
+    const { keyOrder = collatedOrder, mapKey, keyCollision } = options;
+    const used = new Set<string>();
+    const entries = Object.entries(value)
+        .sort(([left], [right]) => keyOrder(left, right))
+        .map(([key, child]): [string, unknown] => {
+            const mapped = mapKey ? mapKey(key) : key;
+            return [
+                keyCollision === "overwrite" ? mapped : uniqueKey(used, mapped),
+                canonicalJsonValue(child, options),
+            ];
+        });
+    return Object.fromEntries(entries);
+}
+
+export function canonicalJson(
+    value: unknown,
+    options?: CanonicalJsonOptions,
+): string {
+    return JSON.stringify(canonicalJsonValue(value, options));
+}

ts/packages/memory/service/src/procedureMarkdown.ts

@@ -8,18 +8,7 @@ import {
     agentEditionToMarkdown,
     validateAgentEdition,
 } from "./agentEditionValidation.js";
-
-function sortJson(value: unknown): unknown {
-    if (Array.isArray(value)) return value.map(sortJson);
-    if (value !== null && typeof value === "object") {
-        return Object.fromEntries(
-            Object.entries(value as Record<string, unknown>)
-                .sort(([left], [right]) => left.localeCompare(right))
-                .map(([key, child]) => [key, sortJson(child)]),
-        );
-    }
-    return value;
-}
+import { canonicalJson } from "./canonicalJson.js";
 
 function validateIdentifier(kind: string, value: string): void {
     if (!/^[A-Za-z0-9][A-Za-z0-9._:-]{0,199}$/.test(value)) {
@@ -72,7 +61,7 @@ export function procedureToMarkdown(document: ProcedureDocument): string {
     );
     lines.push("", "## Sources", "");
     for (const citation of document.citations) {
-        lines.push(`- ${JSON.stringify(sortJson(citation))}`);
+        lines.push(`- ${canonicalJson(citation)}`);
     }
     if (document.citations.length === 0) lines.push("_None_");
     for (const section of document.additionalSections ?? []) {

ts/packages/memory/service/test/agentSessionParsing.spec.ts

@@ -4,6 +4,7 @@
 import {
     agentSessionEventDigest,
     buildAgentSessionEpisodes,
+    legacyAgentSessionEventDigests,
     parseAgentSessionChunk,
 } from "../src/agentSessionParsing.js";
 import type {
@@ -529,3 +530,44 @@ describe("bounded causal agent session episodes", () => {
         expect(buildAgentSessionEpisodes([])).toEqual([]);
     });
 });
+
+describe("colliding redacted keys", () => {
+    const chunk = () =>
+        request(
+            JSON.stringify(
+                sdk(
+                    "tool.execution_start",
+                    {
+                        toolCallId: "call-1",
+                        toolName: "read",
+                        arguments: {
+                            "sk-aaaaaaaaaaaaaaaaaaaa": "first",
+                            "sk-bbbbbbbbbbbbbbbbbbbb": "second",
+                        },
+                    },
+                    { id: "start" },
+                ),
+            ),
+            { importType: "copilot-session-jsonl" },
+        );
+
+    it("keeps every value when redacted keys collide", () => {
+        const [parsed] = parseAgentSessionChunk(chunk()).events;
+        expect(JSON.parse(parsed.text).data.arguments).toEqual({
+            "[REDACTED]": "first",
+            "[REDACTED] (2)": "second",
+        });
+    });
+
+    it("recomputes digests stored when colliding keys were dropped", () => {
+        const [parsed] = parseAgentSessionChunk(chunk()).events;
+        const current = agentSessionEventDigest(parsed);
+        const legacy = legacyAgentSessionEventDigests(chunk()).get(parsed.id);
+        expect(legacy).toBeDefined();
+        expect(legacy).not.toBe(current);
+        // The legacy re-parse must not leak into later parsing.
+        expect(
+            agentSessionEventDigest(parseAgentSessionChunk(chunk()).events[0]),
+        ).toBe(current);
+    });
+});

ts/packages/memory/service/test/canonicalJson.spec.ts

@@ -0,0 +1,50 @@
+// Copyright (c) Microsoft Corporation.
+// Licensed under the MIT License.
+
+import {
+    canonicalJson,
+    canonicalJsonValue,
+    codePointOrder,
+} from "../src/canonicalJson.js";
+
+describe("canonicalJson", () => {
+    it("serializes equal values identically regardless of key order", () => {
+        expect(
+            canonicalJson({ b: 1, a: { d: [2, { y: 1, x: 2 }], c: 3 } }),
+        ).toBe(canonicalJson({ a: { c: 3, d: [2, { x: 2, y: 1 }] }, b: 1 }));
+    });
+
+    it("sorts keys with en collation by default", () => {
+        // Code-point order would put "aZ" first.
+        expect(
+            Object.keys(canonicalJsonValue({ aZ: 1, ab: 2 }) as object),
+        ).toEqual(["ab", "aZ"]);
+    });
+
+    it("sorts keys in code-point order when requested", () => {
+        expect(
+            Object.keys(
+                canonicalJsonValue(
+                    { ab: 1, aZ: 2 },
+                    { keyOrder: codePointOrder },
+                ) as object,
+            ),
+        ).toEqual(["aZ", "ab"]);
+    });
+
+    it("keeps every value when mapped keys collide", () => {
+        expect(
+            canonicalJsonValue(
+                { secretA: 1, secretB: 2 },
+                { mapKey: () => "[REDACTED]" },
+            ),
+        ).toEqual({ "[REDACTED]": 1, "[REDACTED] (2)": 2 });
+    });
+
+    it("keeps __proto__ as an own key", () => {
+        const value = JSON.parse('{"__proto__": {"polluted": true}}');
+        const result = canonicalJsonValue(value) as Record<string, unknown>;
+        expect(Object.getPrototypeOf(result)).toBe(Object.prototype);
+        expect(canonicalJson(value)).toBe('{"__proto__":{"polluted":true}}');
+    });
+});