Preserve Pi message clocks and unified tool correlation

actions/setup/js/pi_session.cjs

@@ -106,6 +106,11 @@ function messageIdentity(message) {
   return undefined;
 }
 
+/** @param {SessionEvent} event @param {any} message */
+function retainMessageTimestamp(event, message) {
+  if (event.timestamp === undefined && (typeof message?.timestamp === "string" || typeof message?.timestamp === "number")) event.timestamp = message.timestamp;
+}
+
 /** @returns {MessageState} */
 function newMessageState() {
   return { text: new Map(), thinking: new Map(), calls: new Map(), blocks: new Map(), errors: [], finalized: false };
@@ -298,21 +303,24 @@ function transformPiV3Entries(records, mapLegacyEvent = event => event) {
     const id = message.toolCallId;
     if (id !== undefined && (executionEnds.has(id) || completions.has(id))) return;
     if (id !== undefined) completions.add(id);
-    emit(raw, "tool.execution_complete", {
-      ...message,
-      toolCallId: id,
-      toolName: message.toolName,
-      output: message.content,
-      success: message.isError === true || isPiResultError(message) ? false : message.isError === false ? true : undefined,
-    });
+    retainMessageTimestamp(
+      emit(raw, "tool.execution_complete", {
+        ...message,
+        toolCallId: id,
+        toolName: message.toolName,
+        output: message.content,
+        success: message.isError === true || isPiResultError(message) ? false : message.isError === false ? true : undefined,
+      }),
+      message
+    );
   };
 
   const emitMessage = (raw, message, fallback) => {
     if (!message || typeof message.role !== "string") return;
     if (!["assistant", "user", "toolResult"].includes(message.role)) {
       if (!snapshots.some(snapshot => isDeepStrictEqual(snapshot.message, message))) {
         snapshots.push({ message: structuredClone(message), state: fallback });
-        emit(raw, "session.info", { sourceType: "message_snapshot", role: message.role, content: message });
+        retainMessageTimestamp(emit(raw, "session.info", { sourceType: "message_snapshot", role: message.role, content: message }), message);
       }
       return;
     }
@@ -363,15 +371,20 @@ function transformPiV3Entries(records, mapLegacyEvent = event => event) {
       ...((message.id ?? message.responseId) !== undefined ? { messageId: message.id ?? message.responseId } : {}),
     };
     for (const segments of [...state.text.values(), ...state.thinking.values()]) {
-      for (const event of segments.segments) Object.assign(event.data, structuredClone(metadata));
+      for (const event of segments.segments) {
+        Object.assign(event.data, structuredClone(metadata));
+        retainMessageTimestamp(event, message);
+      }
+    }
+    for (const event of [...state.calls.values(), ...state.blocks.values()]) {
+      Object.assign(event.data, structuredClone(metadata));
+      retainMessageTimestamp(event, message);
     }
-    for (const event of state.calls.values()) Object.assign(event.data, structuredClone(metadata));
-    for (const event of state.blocks.values()) Object.assign(event.data, structuredClone(metadata));
     if (!snapshots.some(snapshot => snapshot.state === state && isDeepStrictEqual(snapshot.message, message))) snapshots.push({ message: structuredClone(message), state });
     if (raw.type === "message_end" || raw.type === "turn_end") state.finalized = true;
     if (message.errorMessage !== undefined && !state.errors.some(error => isDeepStrictEqual(error, message.errorMessage))) {
       state.errors.push(structuredClone(message.errorMessage));
-      emit(raw, "session.error", { error: message.errorMessage, model: message.responseModel ?? message.model });
+      retainMessageTimestamp(emit(raw, "session.error", { error: message.errorMessage, model: message.responseModel ?? message.model }), message);
     }
   };
 

actions/setup/js/pi_session.test.cjs

@@ -11,6 +11,79 @@ const parse = records => parsePiLog(records.map(record => JSON.stringify(record)
 const byType = (events, type) => events.filter(event => event.type === type);
 
 describe("Pi CI stream regressions", () => {
+  // Synthetic content in the shape of Smoke Pi 42539928028 and
+  // Smoke Pi Sub Agents 39516073462; no captured payloads are published.
+  it.each([0, 1794184043729, "2026-11-09T00:27:23.729Z"])("retains the observed message clock %s through streaming reconciliation and publication", timestamp => {
+    const message = {
+      role: "assistant",
+      responseId: "sanitized-response",
+      model: "requested-model",
+      responseModel: "served-model",
+      timestamp,
+      content: [
+        { type: "thinking", thinking: "Sanitized reasoning." },
+        { type: "toolCall", id: "sanitized-tool", name: "lookup", arguments: { nested: [false, 0, null, "line\n\u00e9"] } },
+        { type: "text", text: "Sanitized answer.\n" },
+      ],
+      usage: { input: 10, output: 3 },
+    };
+    const records = [
+      { type: "turn_start" },
+      { type: "message_update", assistantMessageEvent: { type: "thinking_delta", contentIndex: 0, delta: "Sanitized reasoning." } },
+      { type: "message_update", assistantMessageEvent: { type: "toolcall_start", contentIndex: 1, id: "sanitized-tool", toolName: "lookup" } },
+      { type: "message_update", assistantMessageEvent: { type: "text_delta", contentIndex: 2, delta: "Sanitized answer.\n" } },
+      { type: "message_end", message },
+      { type: "turn_end", message },
+      { type: "agent_end", messages: [message] },
+    ];
+    const original = structuredClone(records);
+    const events = parse(records).logEntries;
+    const messages = events.filter(event => event.data.messageId === message.responseId);
+    expect(messages.map(event => event.type)).toEqual(["assistant.reasoning", "tool.execution_start", "assistant.message"]);
+    expect(messages.every(event => event.timestamp === timestamp)).toBe(true);
+    const unified = mergeSessionSources([{ component: "agent", phase: "agent", path: "agent-session.jsonl", events }]);
+    for (const event of unified.filter(event => event.type !== "session.result")) {
+      expect(event.timestamp).toBe(timestamp);
+      expect(event.provenance.timestampMs).toBe(typeof timestamp === "number" ? timestamp : Date.parse(timestamp));
+      expect(event.data).toMatchObject({ messageId: message.responseId, model: "served-model" });
+      expect(event.data).not.toHaveProperty("messageMetadata");
+      expect(event).not.toHaveProperty("assistantMessageEvent");
+    }
+    expect(unified.find(event => event.type === "tool.execution_start").data.input).toEqual(message.content[1].arguments);
+    expect(selectSessionResult(events).usage).toMatchObject({ input_tokens: 10, output_tokens: 3 });
+    expect(records).toEqual(original);
+  });
+
+  it("preserves source clocks over embedded message clocks and never invents a missing clock", () => {
+    const message = { role: "assistant", timestamp: 10, content: [{ type: "text", text: "Sanitized answer" }] };
+    const events = transformPiV3Entries([
+      { type: "message_update", timestamp: 5, assistantMessageEvent: { type: "text_delta", delta: "Sanitized answer" } },
+      { type: "message_end", timestamp: 20, message },
+    ]);
+    expect(events).toHaveLength(1);
+    expect(events[0].timestamp).toBe(5);
+    expect(events[0].data.messageMetadata.timestamp).toBe(10);
+    const [untimed] = transformPiV3Entries([{ type: "message_end", message: { ...message, timestamp: undefined } }]);
+    expect(untimed).not.toHaveProperty("timestamp");
+  });
+
+  it("retains child, tool-result and provider-error clocks without duplicating snapshots or parent usage", () => {
+    const child = { role: "assistant", responseId: "child-response", model: "fixture-model", timestamp: 3, content: [{ type: "text", text: "Child answer" }], usage: { input: 5, output: 2 } };
+    const records = [
+      { type: "gh_aw_subagent_dispatch", agent: "reader", invocation_id: "child", timestamp: 1 },
+      { type: "gh_aw_subagent_event", agent: "reader", invocation_id: "child", timestamp: 4, event: { type: "message_end", message: child } },
+      { type: "gh_aw_subagent_event", agent: "reader", invocation_id: "child", timestamp: 5, event: { type: "turn_end", message: child } },
+      { type: "message_end", message: { role: "toolResult", toolCallId: "tool", timestamp: 6, content: [{ type: "text", text: "result" }], isError: false } },
+      { type: "message_end", message: { role: "assistant", timestamp: 7, content: [], errorMessage: "Sanitized provider failure", usage: { input: 0, output: 0 } } },
+    ];
+    const events = parse(records).logEntries;
+    expect(byType(events, "assistant.message")).toMatchObject([{ agentId: "child", timestamp: 3 }]);
+    expect(byType(events, "tool.execution_complete")).toMatchObject([{ timestamp: 6 }]);
+    expect(byType(events, "session.error")).toMatchObject([{ timestamp: 7 }]);
+    expect(byType(events, "subagent.request")).toHaveLength(1);
+    expect(selectSessionResult(events).usage).toMatchObject({ input_tokens: 0, output_tokens: 0 });
+  });
+
   it("retains sampled child text and reasoning in their evidenced dispatch scope through unified projection", () => {
     const original = structuredClone(childObservations);
     const records = [...success, ...childObservations];

actions/setup/js/pi_session_pipeline.test.cjs

@@ -7,7 +7,7 @@ import { parsePiLog } from "./parse_pi_log.cjs";
 import { runLogParser } from "./log_parser_bootstrap.cjs";
 import { writeUnifiedSession } from "./unified_session.cjs";
 import { normalizeUnifiedSessionEvent } from "./unified_session_payload.cjs";
-import { createSessionValidator } from "./scripts/validate_session.cjs";
+import { createSessionValidator, validateSession } from "./scripts/validate_session.cjs";
 
 describe("Pi native-to-unified artifact pipeline", () => {
   let root;
@@ -40,6 +40,35 @@ describe("Pi native-to-unified artifact pipeline", () => {
     fs.rmSync(root, { recursive: true, force: true });
   });
 
+  it.each([false, true])("publishes message clocks and tool correlation from saved canonical events (metadata-only clocks=%s)", legacy => {
+    const records = [...success, ...childObservations];
+    const events = JSON.parse(JSON.stringify(parsePiLog(records.map(JSON.stringify).join("\n")).logEntries));
+    if (legacy) {
+      for (const event of events) if (event.data.messageMetadata?.timestamp !== undefined) delete event.timestamp;
+    }
+    const original = structuredClone(events);
+    fs.writeFileSync(path.join(root, "agent-session.jsonl"), events.map(JSON.stringify).join("\n") + "\n");
+    const unified = writeUnifiedSession({ rootDir: root, engine: "pi", warn: vi.fn() });
+    const agent = unified.filter(event => event.provenance.path === "agent-session.jsonl");
+    const messages = agent.filter(event => event.type === "assistant.message" || event.type === "assistant.reasoning" || event.type === "tool.execution_start");
+    expect(messages).toHaveLength(4);
+    for (const event of messages) {
+      const native = events[event.provenance.index];
+      expect(event.timestamp).toBe(native.data.messageMetadata.timestamp);
+      expect(event.provenance.timestampMs).toBe(native.data.messageMetadata.timestamp);
+      expect(event.data.messageId).toBe(native.data.messageId);
+      expect(event.data.model).toBe(native.data.model);
+      expect(event.data).not.toHaveProperty("messageMetadata");
+    }
+    expect(messages.find(event => event.type === "tool.execution_start").data.input).toEqual({ command: "printf sanitized" });
+    expect(messages.find(event => event.agentId)?.agentId).toBe("legacy:pi:sanitized-reader:1");
+    const bytes = fs.readFileSync(path.join(root, "usage/aw_session.jsonl"), "utf8");
+    expect(validateSession(bytes)).toBe(unified.length);
+    expect(readEvents(path.join(root, "agent-session.jsonl"))).toEqual(original);
+    writeUnifiedSession({ rootDir: root, engine: "pi", warn: vi.fn() });
+    expect(fs.readFileSync(path.join(root, "usage/aw_session.jsonl"), "utf8")).toBe(bytes);
+  });
+
   // Sanitized excerpts from separate runs, combined locally to exercise the
   // collector. This is not a new live execution or an exact single-run fixture.
   it.each([

actions/setup/js/types/agent_session.type-test.ts

@@ -117,3 +117,15 @@ const invalidUnifiedInit: UnifiedSessionEventDataMap["session.init"] = {
   tools: [],
 };
 void invalidUnifiedInit;
+
+const correlatedUnifiedTools: UnifiedSessionEventDataMap["tool.execution_start" | "tool.execution_update" | "tool.execution_complete"][] = [
+  { toolCallId: "call", messageId: "response", model: "served-model", input: false },
+  { toolCallId: "call", messageId: "response", model: "served-model", output: null },
+  { toolCallId: "call", messageId: "response", model: "served-model", success: false },
+];
+void correlatedUnifiedTools;
+const invalidUnifiedToolEnvelope: UnifiedSessionEventDataMap["tool.execution_start"] = {
+  // @ts-expect-error The native envelope remains archival, not part of the unified payload.
+  messageMetadata: { timestamp: 0 },
+};
+void invalidUnifiedToolEnvelope;

actions/setup/js/types/unified_session.d.ts

@@ -834,22 +834,23 @@ export interface UnifiedSessionEventDataMap {
     /** Observed native usage, not an additive accounting contribution. */
     usageSnapshot?: JsonValue;
   };
-  "tool.execution_start": Pick<ToolExecutionStartData, "toolCallId" | "stepIndex" | "toolName" | "input" | "inputTruncated" | "command" | "mcpServerName"> & UnifiedSessionContextData;
-  "tool.execution_update": Pick<ToolExecutionStartData, "toolCallId" | "stepIndex" | "toolName" | "input" | "mcpServerName"> & {
-    output?: JsonValue;
-    partial?: JsonValue;
-    delta?: JsonValue;
-    contentIndex?: JsonValue;
-    sourceType?: JsonValue;
-    /** Repeated observed metadata, never additive inference usage. */
-    usageSnapshot?: JsonValue;
-    durationMs?: JsonValue;
-    status?: JsonValue;
-    command?: JsonValue;
-    taskId?: string;
-    taskType?: string;
-    error?: JsonValue;
-  } & UnifiedSessionContextData;
+  "tool.execution_start": Pick<ToolExecutionStartData, "toolCallId" | "stepIndex" | "toolName" | "input" | "inputTruncated" | "command" | "mcpServerName"> & Pick<UnifiedMessageData, "messageId" | "model"> & UnifiedSessionContextData;
+  "tool.execution_update": Pick<ToolExecutionStartData, "toolCallId" | "stepIndex" | "toolName" | "input" | "mcpServerName"> &
+    Pick<UnifiedMessageData, "messageId" | "model"> & {
+      output?: JsonValue;
+      partial?: JsonValue;
+      delta?: JsonValue;
+      contentIndex?: JsonValue;
+      sourceType?: JsonValue;
+      /** Repeated observed metadata, never additive inference usage. */
+      usageSnapshot?: JsonValue;
+      durationMs?: JsonValue;
+      status?: JsonValue;
+      command?: JsonValue;
+      taskId?: string;
+      taskType?: string;
+      error?: JsonValue;
+    } & UnifiedSessionContextData;
   "tool.execution_complete": Pick<
     ToolExecutionCompleteData,
     | "toolCallId"
@@ -871,6 +872,7 @@ export interface UnifiedSessionEventDataMap {
     | "workflowName"
     | "workflowRunId"
   > &
+    Pick<UnifiedMessageData, "messageId" | "model"> &
     UnifiedSessionContextData;
   "dynamicWorkflows.task_started": DynamicWorkflowObservationData;
   "dynamicWorkflows.task_progress": DynamicWorkflowObservationData;

actions/setup/js/unified_session_order.cjs

@@ -10,7 +10,7 @@
  * @returns {bigint | undefined}
  */
 function sessionTimestampNs(record, unit = "milliseconds") {
-  const value = record.timestamp ?? record.ts ?? record.time ?? record.created_at ?? record.message?.timestamp;
+  const value = record.timestamp ?? record.ts ?? record.time ?? record.created_at ?? record.message?.timestamp ?? record.data?.messageMetadata?.timestamp;
   if (unit === "nanoseconds") {
     if (typeof value !== "string" || !/^\d{1,22}$/.test(value)) return undefined;
     const ns = BigInt(value);

actions/setup/js/unified_session_order.test.cjs

@@ -9,6 +9,31 @@ const header = { type: "session.format", data: { version: 1 }, provenance: { com
 const jsonl = events => [header, ...events].map(JSON.stringify).join("\n") + "\n";
 
 describe("Unified trace ordering", () => {
+  it.each([0, 2000, "2026-11-09T00:27:23.729Z"])("orders saved Pi message metadata clocks without publishing the native envelope (%s)", timestamp => {
+    const clock = typeof timestamp === "number" ? timestamp : Date.parse(timestamp);
+    const saved = {
+      type: "tool.execution_start",
+      data: { toolCallId: "call", messageId: "message", model: "served-model", input: { nested: [false, 0] }, messageMetadata: { timestamp, provider: "fixture", usage: { input: 10 } } },
+    };
+    const original = structuredClone(saved);
+    expect(sessionTimestamp(saved)).toBe(clock);
+    const merged = mergeSessionSources([source("agent", [saved]), source("mcp", [event("before", clock - 1), event("after", clock + 1)])]);
+    expect(merged.map(record => record.id ?? record.data.toolCallId)).toEqual(["before", "call", "after"]);
+    expect(merged[1]).toMatchObject({ timestamp, data: { messageId: "message", model: "served-model" }, provenance: { timestampMs: clock, index: 0 } });
+    expect(merged[1].data).not.toHaveProperty("messageMetadata");
+    expect(validateSession(jsonl(merged))).toBe(4);
+    expect(saved).toEqual(original);
+    expect(mergeSessionSources([source("agent", [merged[1]])])[0].timestamp).toBe(timestamp);
+  });
+
+  it("does not overwrite an explicit source clock or replace invalid clocks with message metadata", () => {
+    const saved = { type: "assistant.message", timestamp: 10, data: { content: "exact", messageMetadata: { timestamp: 20 } } };
+    expect(sessionTimestamp(saved)).toBe(10);
+    expect(mergeSessionSources([source("agent", [saved])])[0].timestamp).toBe(10);
+    expect(sessionTimestamp({ ...saved, timestamp: "invalid" })).toBeUndefined();
+    expect(sessionTimestamp({ type: "assistant.message", data: { content: "untimed" } })).toBeUndefined();
+  });
+
   it.each([[], [source("agent", [])], [source("agent", []), source("otel", [], "nanoseconds")]].map(sources => [sources]))("accepts empty source collections without inventing observations: %j", sources => {
     expect(mergeSessionSources(sources)).toEqual([]);
   });

actions/setup/js/unified_session_payload.cjs

@@ -55,6 +55,8 @@ const AIC_RESOLVABLE_PHASES = new Set(["detection", "evals"]);
 const TOOL_FIELDS = {
   toolCallId: ["toolCallId", "tool_call_id", "tool_use_id"],
   toolName: ["toolName", "tool_name"],
+  messageId: ["messageId", "message_id"],
+  model: ["model"],
   mcpServerName: ["mcpServerName"],
   parentToolCallId: ["parentToolCallId"],
   stepIndex: ["stepIndex", "step_index"],
@@ -863,9 +865,9 @@ function normalizeUnifiedSessionEvent(event, phase) {
     if (Array.isArray(errors)) data.errors = errors.filter(error => error && typeof error === "object" && !Array.isArray(error)).map(projectSafeOutputData);
   }
   const metadata = selectFields(event, { id: ["id"], parentId: ["parentId"], timestamp: ["timestamp", "ts", "time", "created_at"], agentId: ["agentId"], ephemeral: ["ephemeral"] });
-  /** @type {any} */
   const nativeMessage = event.message;
-  if (!Object.hasOwn(metadata, "timestamp") && nativeMessage?.timestamp !== undefined) metadata.timestamp = structuredClone(nativeMessage.timestamp);
+  const nativeTimestamp = (nativeMessage && typeof nativeMessage === "object" && "timestamp" in nativeMessage ? nativeMessage.timestamp : undefined) ?? source.messageMetadata?.timestamp;
+  if (!Object.hasOwn(metadata, "timestamp") && (typeof nativeTimestamp === "string" || typeof nativeTimestamp === "number")) metadata.timestamp = nativeTimestamp;
   return {
     type: event.type,
     ...metadata,