fix: preserve replay message timestamps (#8942)

Signed-off-by: Matt Toohey <contact@matttoohey.com>

crates/goose/src/acp/server.rs

@@ -1670,6 +1670,42 @@ fn extract_tool_call_update_meta(
     Some(meta_map)
 }
 
+fn replay_message_meta(message: &Message) -> Meta {
+    let mut meta = serde_json::Map::new();
+    meta.insert(
+        "goose".to_string(),
+        serde_json::Value::Object(replay_message_goose_meta(message)),
+    );
+    meta
+}
+
+fn replay_message_goose_meta(message: &Message) -> serde_json::Map<String, serde_json::Value> {
+    let mut goose = serde_json::Map::new();
+    goose.insert("created".to_string(), serde_json::json!(message.created));
+    if let Some(id) = &message.id {
+        goose.insert("messageId".to_string(), serde_json::json!(id));
+    }
+    goose
+}
+
+fn merge_replay_message_meta(meta: Option<Meta>, message: &Message) -> Meta {
+    let replay_goose = replay_message_goose_meta(message);
+    let mut meta = meta.unwrap_or_default();
+    let goose_value = meta
+        .entry("goose".to_string())
+        .or_insert_with(|| serde_json::Value::Object(serde_json::Map::new()));
+
+    if let serde_json::Value::Object(goose) = goose_value {
+        for (key, value) in replay_goose {
+            goose.insert(key, value);
+        }
+    } else {
+        *goose_value = serde_json::Value::Object(replay_goose);
+    }
+
+    meta
+}
+
 fn build_tool_call_content(tool_result: &ToolResult<CallToolResult>) -> Vec<ToolCallContent> {
     match tool_result {
         Ok(result) => result
@@ -2119,7 +2155,8 @@ impl GooseAcpAgent {
                                 ),
                             );
                         }
-                        let chunk = ContentChunk::new(ContentBlock::Text(tc));
+                        let chunk = ContentChunk::new(ContentBlock::Text(tc))
+                            .meta(replay_message_meta(message));
                         let update = match message.role {
                             Role::User => SessionUpdate::UserMessageChunk(chunk),
                             Role::Assistant => SessionUpdate::AgentMessageChunk(chunk),
@@ -2147,7 +2184,8 @@ impl GooseAcpAgent {
                                     ToolCallId::new(tool_request.id.clone()),
                                     format_tool_name(&tool_name),
                                 )
-                                .status(ToolCallStatus::Pending),
+                                .status(ToolCallStatus::Pending)
+                                .meta(replay_message_meta(message)),
                             ),
                         ))?;
                     }
@@ -2187,7 +2225,10 @@ impl GooseAcpAgent {
 
                         let update =
                             ToolCallUpdate::new(ToolCallId::new(tool_response.id.clone()), fields)
-                                .meta(extract_tool_call_update_meta(tool_response));
+                                .meta(merge_replay_message_meta(
+                                    extract_tool_call_update_meta(tool_response),
+                                    message,
+                                ));
                         cx.send_notification(SessionNotification::new(
                             args.session_id.clone(),
                             SessionUpdate::ToolCallUpdate(update),
@@ -2196,9 +2237,12 @@ impl GooseAcpAgent {
                     MessageContent::Thinking(thinking) => {
                         cx.send_notification(SessionNotification::new(
                             args.session_id.clone(),
-                            SessionUpdate::AgentThoughtChunk(ContentChunk::new(
+                            SessionUpdate::AgentThoughtChunk(
-                                ContentBlock::Text(TextContent::new(thinking.thinking.clone())),
+                                ContentChunk::new(ContentBlock::Text(TextContent::new(
-                            )),
+                                    thinking.thinking.clone(),
+                                )))
+                                .meta(replay_message_meta(message)),
+                            ),
                         ))?;
                     }
                     _ => {}
@@ -3252,6 +3296,65 @@ print(\"hello, world\")
         );
     }
 
+    #[test]
+    fn test_merge_replay_message_meta_preserves_existing_goose_meta() {
+        let message = Message::new(Role::Assistant, 1_700_000_000, vec![]).with_id("msg_1");
+        let existing = serde_json::from_value(serde_json::json!({
+            "goose": {
+                "mcpApp": {
+                    "resourceUri": "ui://trusted/app",
+                    "extensionName": "weather",
+                    "toolName": "weather__render",
+                },
+            },
+        }))
+        .unwrap();
+
+        let merged = merge_replay_message_meta(Some(existing), &message);
+
+        assert_eq!(
+            merged.get("goose"),
+            Some(&serde_json::json!({
+                "created": 1_700_000_000,
+                "messageId": "msg_1",
+                "mcpApp": {
+                    "resourceUri": "ui://trusted/app",
+                    "extensionName": "weather",
+                    "toolName": "weather__render",
+                },
+            })),
+        );
+    }
+
+    #[test]
+    fn test_merge_replay_message_meta_creates_fresh_when_none() {
+        let message = Message::new(Role::Assistant, 1_700_000_000, vec![]).with_id("msg_2");
+
+        let merged = merge_replay_message_meta(None, &message);
+
+        assert_eq!(
+            merged.get("goose"),
+            Some(&serde_json::json!({
+                "created": 1_700_000_000,
+                "messageId": "msg_2",
+            })),
+        );
+    }
+
+    #[test]
+    fn test_merge_replay_message_meta_omits_message_id_when_none() {
+        let message = Message::new(Role::Assistant, 1_700_000_000, vec![]);
+
+        let merged = merge_replay_message_meta(None, &message);
+
+        assert_eq!(
+            merged.get("goose"),
+            Some(&serde_json::json!({
+                "created": 1_700_000_000,
+            })),
+        );
+    }
+
     fn make_session_with_usage(
         total_tokens: Option<i32>,
         input_tokens: Option<i32>,

ui/goose2/src/shared/api/__tests__/acpNotificationHandler.test.ts

@@ -288,8 +288,64 @@ describe("acpNotificationHandler", () => {
     });
   });
 
+  it("replay preserves timestamps from goose metadata on user and assistant chunks", async () => {
+    const replaySessionId = "replay-timestamp-session";
+    const userCreated = 1_700_000_000;
+    const assistantCreated = 1_700_000_120;
+    useChatStore.setState({
+      loadingSessionIds: new Set<string>([replaySessionId]),
+    });
+
+    await handleSessionNotification({
+      sessionId: replaySessionId,
+      update: {
+        sessionUpdate: "user_message_chunk",
+        content: {
+          type: "text",
+          text: "what time was this sent?",
+        },
+        _meta: {
+          goose: {
+            messageId: "user-from-meta",
+            created: userCreated,
+          },
+        },
+      },
+    } as never);
+
+    await handleSessionNotification({
+      sessionId: replaySessionId,
+      update: {
+        sessionUpdate: "agent_message_chunk",
+        content: {
+          type: "text",
+          text: "At the original replay time.",
+        },
+        _meta: {
+          goose: {
+            messageId: "assistant-from-meta",
+            created: assistantCreated,
+          },
+        },
+      },
+    } as never);
+
+    const buffer = getReplayBuffer(replaySessionId);
+    expect(buffer?.[0]).toMatchObject({
+      id: "user-from-meta",
+      role: "user",
+      created: userCreated * 1000,
+    });
+    expect(buffer?.[1]).toMatchObject({
+      id: "assistant-from-meta",
+      role: "assistant",
+      created: assistantCreated * 1000,
+    });
+  });
+
   it("replay preserves gooseSessionId in MCP app payloads before tracker registration", async () => {
     const replaySessionId = "replay-goose-session-2";
+    const replayCreated = 1_700_000_240;
     useChatStore.setState({
       loadingSessionIds: new Set<string>([replaySessionId]),
     });
@@ -300,6 +356,12 @@ describe("acpNotificationHandler", () => {
         sessionUpdate: "tool_call",
         toolCallId: "tool-1",
         title: "mcp_app_bench__inspect_host_info",
+        _meta: {
+          goose: {
+            messageId: "assistant-tool-only",
+            created: replayCreated,
+          },
+        },
       },
     } as never);
 
@@ -316,6 +378,8 @@ describe("acpNotificationHandler", () => {
               extensionName: "mcp_app_bench",
               resourceUri: "ui://inspect-host-info",
             },
+            messageId: "assistant-tool-only",
+            created: replayCreated,
           },
         },
       },
@@ -323,6 +387,10 @@ describe("acpNotificationHandler", () => {
 
     const buffer = getReplayBuffer(replaySessionId);
     const assistant = buffer?.[0];
+    expect(assistant).toMatchObject({
+      id: "assistant-tool-only",
+      created: replayCreated * 1000,
+    });
     const mcpAppBlock = assistant?.content.find(
       (block) => block.type === "mcpApp",
     );
@@ -333,4 +401,71 @@ describe("acpNotificationHandler", () => {
       }),
     });
   });
+
+  it("replay falls back to tracked assistant when a tool update ID is not buffered", async () => {
+    const replaySessionId = "replay-tool-response-id-session";
+    const assistantCreated = 1_700_000_120;
+    const toolResponseCreated = 1_700_000_240;
+    useChatStore.setState({
+      loadingSessionIds: new Set<string>([replaySessionId]),
+    });
+
+    await handleSessionNotification({
+      sessionId: replaySessionId,
+      update: {
+        sessionUpdate: "agent_message_chunk",
+        content: {
+          type: "text",
+          text: "I'll check that.",
+        },
+        _meta: {
+          goose: {
+            messageId: "assistant-1",
+            created: assistantCreated,
+          },
+        },
+      },
+    } as never);
+
+    await handleSessionNotification({
+      sessionId: replaySessionId,
+      update: {
+        sessionUpdate: "tool_call_update",
+        toolCallId: "tool-1",
+        status: "completed",
+        content: [
+          {
+            type: "content",
+            content: {
+              type: "text",
+              text: "Tool completed.",
+            },
+          },
+        ],
+        _meta: {
+          goose: {
+            messageId: "tool-response-user-message",
+            created: toolResponseCreated,
+          },
+        },
+      },
+    } as never);
+
+    const buffer = getReplayBuffer(replaySessionId);
+    const assistant = buffer?.[0];
+    expect(assistant).toMatchObject({
+      id: "assistant-1",
+      created: assistantCreated * 1000,
+    });
+    expect(assistant?.content.map((block) => block.type)).toEqual([
+      "text",
+      "toolResponse",
+    ]);
+    expect(assistant?.content[1]).toMatchObject({
+      type: "toolResponse",
+      id: "tool-1",
+      result: "Tool completed.",
+      isError: false,
+    });
+  });
 });

ui/goose2/src/shared/api/__tests__/acpReplayMetadata.test.ts

@@ -0,0 +1,117 @@
+import { describe, expect, it } from "vitest";
+import { getReplayCreated, getReplayMessageId } from "../acpReplayMetadata";
+
+describe("getReplayMessageId", () => {
+  it("returns messageId from top-level field", () => {
+    expect(getReplayMessageId({ messageId: "msg_1" })).toBe("msg_1");
+  });
+
+  it("returns messageId from _meta.goose", () => {
+    const source = {
+      _meta: { goose: { messageId: "msg_2" } },
+    };
+    expect(getReplayMessageId(source)).toBe("msg_2");
+  });
+
+  it("prefers top-level messageId over _meta.goose.messageId", () => {
+    const source = {
+      messageId: "top",
+      _meta: { goose: { messageId: "nested" } },
+    };
+    expect(getReplayMessageId(source)).toBe("top");
+  });
+
+  it("returns null when no messageId is present", () => {
+    expect(getReplayMessageId({})).toBeNull();
+  });
+
+  it("returns null for empty string messageId", () => {
+    expect(getReplayMessageId({ messageId: "" })).toBeNull();
+  });
+
+  it("returns null when _meta is null", () => {
+    expect(getReplayMessageId({ _meta: null })).toBeNull();
+  });
+
+  it("returns null when _meta.goose is not an object", () => {
+    expect(getReplayMessageId({ _meta: { goose: "not-object" } })).toBeNull();
+  });
+
+  it("returns null when _meta is an array", () => {
+    expect(
+      getReplayMessageId({ _meta: [] as unknown as Record<string, unknown> }),
+    ).toBeNull();
+  });
+});
+
+describe("getReplayCreated", () => {
+  it("returns milliseconds from a seconds-epoch timestamp", () => {
+    const source = { _meta: { goose: { created: 1_700_000_000 } } };
+    expect(getReplayCreated(source)).toBe(1_700_000_000_000);
+  });
+
+  it("returns milliseconds directly when already in milliseconds", () => {
+    const source = { _meta: { goose: { created: 1_700_000_000_000 } } };
+    expect(getReplayCreated(source)).toBe(1_700_000_000_000);
+  });
+
+  it("falls back to createdAt if created is missing", () => {
+    const source = { _meta: { goose: { createdAt: 1_700_000_000 } } };
+    expect(getReplayCreated(source)).toBe(1_700_000_000_000);
+  });
+
+  it("returns undefined when no timestamp is present", () => {
+    expect(getReplayCreated({})).toBeUndefined();
+  });
+
+  it("returns undefined for non-numeric values", () => {
+    const source = { _meta: { goose: { created: "2024-01-01T00:00:00Z" } } };
+    expect(getReplayCreated(source)).toBeUndefined();
+  });
+
+  it("returns undefined for NaN", () => {
+    const source = { _meta: { goose: { created: NaN } } };
+    expect(getReplayCreated(source)).toBeUndefined();
+  });
+
+  it("returns undefined for Infinity", () => {
+    const source = { _meta: { goose: { created: Infinity } } };
+    expect(getReplayCreated(source)).toBeUndefined();
+  });
+
+  it("returns undefined for negative timestamps", () => {
+    const source = { _meta: { goose: { created: -1 } } };
+    expect(getReplayCreated(source)).toBeUndefined();
+  });
+
+  it("returns undefined for negative epoch values", () => {
+    const source = { _meta: { goose: { created: -1_000_000_000 } } };
+    expect(getReplayCreated(source)).toBeUndefined();
+  });
+
+  it("handles the boundary between seconds and milliseconds", () => {
+    // Just below the threshold: treated as seconds
+    const belowSource = {
+      _meta: { goose: { created: 999_999_999_999 } },
+    };
+    expect(getReplayCreated(belowSource)).toBe(999_999_999_999_000);
+
+    // At the threshold: treated as milliseconds
+    const atSource = {
+      _meta: { goose: { created: 1_000_000_000_000 } },
+    };
+    expect(getReplayCreated(atSource)).toBe(1_000_000_000_000);
+  });
+
+  it("returns zero as a valid timestamp", () => {
+    const source = { _meta: { goose: { created: 0 } } };
+    expect(getReplayCreated(source)).toBe(0);
+  });
+
+  it("returns undefined when _meta.goose is an array", () => {
+    const source = {
+      _meta: { goose: [{ created: 1_700_000_000 }] },
+    };
+    expect(getReplayCreated(source)).toBeUndefined();
+  });
+});

ui/goose2/src/shared/api/acpNotificationHandler.ts

@@ -25,6 +25,7 @@ import {
   ensureReplayAssistantMessage,
   getTrackedReplayAssistantMessageId,
 } from "./acpReplayAssistant";
+import { getReplayCreated, getReplayMessageId } from "./acpReplayMetadata";
 import {
   getLocalSessionId,
   subscribeToSessionRegistration,
@@ -150,7 +151,8 @@ function handleReplay(
     case "agent_message_chunk": {
       const msg = ensureReplayAssistantMessage(
         sessionId,
-        update.messageId ?? null,
+        getReplayMessageId(update),
+        getReplayCreated(update),
       );
       if (msg && update.content.type === "text" && "text" in update.content) {
         const last = msg.content[msg.content.length - 1];
@@ -166,32 +168,54 @@ function handleReplay(
     case "user_message_chunk": {
       clearReplayAssistantMessage(sessionId);
       if (update.content.type !== "text" || !("text" in update.content)) break;
-      const messageId = update.messageId ?? crypto.randomUUID();
+      const messageId = getReplayMessageId(update) ?? crypto.randomUUID();
-      handleReplayUserMessageChunk(sessionId, messageId, update.content);
+      handleReplayUserMessageChunk(
+        sessionId,
+        messageId,
+        update.content,
+        getReplayCreated(update),
+      );
       break;
     }
 
     case "tool_call": {
-      const msg = ensureReplayAssistantMessage(sessionId);
+      const created = getReplayCreated(update);
+      const msg = ensureReplayAssistantMessage(
+        sessionId,
+        getReplayMessageId(update),
+        created,
+      );
       msg.content.push({
         type: "toolRequest",
         id: update.toolCallId,
         name: update.title,
         arguments: {},
         status: "executing",
-        startedAt: Date.now(),
+        startedAt: created ?? Date.now(),
       });
       break;
     }
 
     case "tool_call_update": {
-      const replayMessageId = getTrackedReplayAssistantMessageId(sessionId);
+      const created = getReplayCreated(update);
-      const msg =
+      const replayMessageId = getReplayMessageId(update);
-        findReplayMessageWithToolCall(sessionId, update.toolCallId) ??
+      const trackedMessageId = getTrackedReplayAssistantMessageId(sessionId);
-        (replayMessageId
+      const replayMsg = replayMessageId
-          ? getBufferedMessage(sessionId, replayMessageId)
+        ? getBufferedMessage(sessionId, replayMessageId)
-          : undefined);
+        : undefined;
+      const trackedMsg =
+        trackedMessageId && trackedMessageId !== replayMessageId
+          ? getBufferedMessage(sessionId, trackedMessageId)
+          : undefined;
+      const existingMsg = findReplayMessageWithToolCall(
+        sessionId,
+        update.toolCallId,
+      );
+      const msg = existingMsg ?? replayMsg ?? trackedMsg;
       if (msg) {
+        if (created !== undefined && !existingMsg && msg === replayMsg) {
+          msg.created = created;
+        }
         if (update.title) {
           const tc = msg.content.find(
             (c) => c.type === "toolRequest" && c.id === update.toolCallId,

ui/goose2/src/shared/api/acpReplayAssistant.ts

@@ -15,12 +15,16 @@ export function getTrackedReplayAssistantMessageId(
 export function ensureReplayAssistantMessage(
   sessionId: string,
   preferredMessageId?: string | null,
+  created?: number,
 ): Message {
   const trackedMessageId = replayAssistantMessageIds.get(sessionId);
 
   if (preferredMessageId) {
     const preferredMessage = getBufferedMessage(sessionId, preferredMessageId);
     if (preferredMessage?.role === "assistant") {
+      if (created !== undefined) {
+        preferredMessage.created = created;
+      }
       replayAssistantMessageIds.set(sessionId, preferredMessageId);
       return preferredMessage;
     }
@@ -33,6 +37,9 @@ export function ensureReplayAssistantMessage(
         trackedMessage.id = preferredMessageId;
         replayAssistantMessageIds.set(sessionId, preferredMessageId);
       }
+      if (created !== undefined) {
+        trackedMessage.created = created;
+      }
       return trackedMessage;
     }
   }
@@ -42,7 +49,7 @@ export function ensureReplayAssistantMessage(
   const message: Message = {
     id: messageId,
     role: "assistant",
-    created: Date.now(),
+    created: created ?? Date.now(),
     content: [],
     metadata: {
       userVisible: true,

ui/goose2/src/shared/api/acpReplayMetadata.ts

@@ -0,0 +1,57 @@
+type ReplayMetadataSource = {
+  _meta?: Record<string, unknown> | null;
+  messageId?: string | null;
+};
+
+export function getReplayMessageId(
+  source: ReplayMetadataSource,
+): string | null {
+  if (source.messageId) {
+    return source.messageId;
+  }
+
+  const metaMessageId = getGooseReplayMeta(source)?.messageId;
+  if (typeof metaMessageId === "string" && metaMessageId.length > 0) {
+    return metaMessageId;
+  }
+
+  return null;
+}
+
+export function getReplayCreated(
+  source: ReplayMetadataSource,
+): number | undefined {
+  const goose = getGooseReplayMeta(source);
+  return coerceReplayTimestamp(goose?.created ?? goose?.createdAt);
+}
+
+function getGooseReplayMeta(
+  source: ReplayMetadataSource,
+): Record<string, unknown> | null {
+  if (!isRecord(source._meta)) {
+    return null;
+  }
+
+  const goose = source._meta.goose;
+  return isRecord(goose) ? goose : null;
+}
+
+function coerceReplayTimestamp(value: unknown): number | undefined {
+  if (typeof value === "number") {
+    return normalizeEpochMilliseconds(value);
+  }
+
+  return undefined;
+}
+
+function normalizeEpochMilliseconds(value: number): number | undefined {
+  if (!Number.isFinite(value) || value < 0) {
+    return undefined;
+  }
+
+  return value < 1_000_000_000_000 ? value * 1000 : value;
+}
+
+function isRecord(value: unknown): value is Record<string, unknown> {
+  return typeof value === "object" && value !== null && !Array.isArray(value);
+}

ui/goose2/src/shared/api/acpSkillReplayChips.ts

@@ -43,6 +43,7 @@ export function handleReplayUserMessageChunk(
   sessionId: string,
   messageId: string,
   content: { text: string },
+  created?: number,
 ): void {
   const buffer = ensureReplayBuffer(sessionId);
   const existing = getBufferedMessage(sessionId, messageId);
@@ -62,7 +63,7 @@ export function handleReplayUserMessageChunk(
     buffer.push({
       id: messageId,
       role: "user",
-      created: Date.now(),
+      created: created ?? Date.now(),
       content: [textBlock],
       metadata: {
         userVisible: true,
@@ -71,6 +72,9 @@ export function handleReplayUserMessageChunk(
       },
     });
   } else {
+    if (created !== undefined) {
+      existing.created = created;
+    }
     existing.content.push(textBlock);
     attachReplayChips(sessionId, messageId, existing, chips);
   }