diff --git a/ts/packages/memory/service/src/agentEditionValidation.ts b/ts/packages/memory/service/src/agentEditionValidation.ts index d1ebd3fff3..3e33a1bf88 100644 --- a/ts/packages/memory/service/src/agentEditionValidation.ts +++ b/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 = ""; -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) - .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( diff --git a/ts/packages/memory/service/src/agentSessionParsing.ts b/ts/packages/memory/service/src/agentSessionParsing.ts index 066ae6f1f0..fa37f1330c 100644 --- a/ts/packages/memory/service/src/agentSessionParsing.ts +++ b/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(); 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 { + 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[], diff --git a/ts/packages/memory/service/src/agentSessionStore.ts b/ts/packages/memory/service/src/agentSessionStore.ts index ed35616c95..566fb26081 100644 --- a/ts/packages/memory/service/src/agentSessionStore.ts +++ b/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 = {}; + let legacyDigests: Map | 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; diff --git a/ts/packages/memory/service/src/canonicalJson.ts b/ts/packages/memory/service/src/canonicalJson.ts new file mode 100644 index 0000000000..92df5eae1b --- /dev/null +++ b/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, 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(); + 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)); +} diff --git a/ts/packages/memory/service/src/procedureMarkdown.ts b/ts/packages/memory/service/src/procedureMarkdown.ts index 7de91fcc48..8b7fcf59d5 100644 --- a/ts/packages/memory/service/src/procedureMarkdown.ts +++ b/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) - .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 ?? []) { diff --git a/ts/packages/memory/service/test/agentSessionParsing.spec.ts b/ts/packages/memory/service/test/agentSessionParsing.spec.ts index 398175e9b6..bb07b06fdc 100644 --- a/ts/packages/memory/service/test/agentSessionParsing.spec.ts +++ b/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); + }); +}); diff --git a/ts/packages/memory/service/test/canonicalJson.spec.ts b/ts/packages/memory/service/test/canonicalJson.spec.ts new file mode 100644 index 0000000000..82a41fe79e --- /dev/null +++ b/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; + expect(Object.getPrototypeOf(result)).toBe(Object.prototype); + expect(canonicalJson(value)).toBe('{"__proto__":{"polluted":true}}'); + }); +});