From 501c1edab51d527bcd875ddd4f393eccc922a2f9 Mon Sep 17 00:00:00 2001 From: Matt Toohey Date: Sat, 2 May 2026 03:16:01 +1000 Subject: [PATCH] fix: preserve replay message timestamps (#8942) Signed-off-by: Matt Toohey --- crates/goose/src/acp/server.rs | 115 ++++++++++++++- .../__tests__/acpNotificationHandler.test.ts | 135 ++++++++++++++++++ .../api/__tests__/acpReplayMetadata.test.ts | 117 +++++++++++++++ .../src/shared/api/acpNotificationHandler.ts | 46 ++++-- .../src/shared/api/acpReplayAssistant.ts | 9 +- ui/goose2/src/shared/api/acpReplayMetadata.ts | 57 ++++++++ .../src/shared/api/acpSkillReplayChips.ts | 6 +- 7 files changed, 466 insertions(+), 19 deletions(-) create mode 100644 ui/goose2/src/shared/api/__tests__/acpReplayMetadata.test.ts create mode 100644 ui/goose2/src/shared/api/acpReplayMetadata.ts diff --git a/crates/goose/src/acp/server.rs b/crates/goose/src/acp/server.rs index 64af66d7..198bdb4b 100644 --- a/crates/goose/src/acp/server.rs +++ b/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 { + 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, 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) -> Vec { 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( - ContentBlock::Text(TextContent::new(thinking.thinking.clone())), - )), + SessionUpdate::AgentThoughtChunk( + 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, input_tokens: Option, diff --git a/ui/goose2/src/shared/api/__tests__/acpNotificationHandler.test.ts b/ui/goose2/src/shared/api/__tests__/acpNotificationHandler.test.ts index 81341cf4..ec3fcaf6 100644 --- a/ui/goose2/src/shared/api/__tests__/acpNotificationHandler.test.ts +++ b/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([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([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([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, + }); + }); }); diff --git a/ui/goose2/src/shared/api/__tests__/acpReplayMetadata.test.ts b/ui/goose2/src/shared/api/__tests__/acpReplayMetadata.test.ts new file mode 100644 index 00000000..24478537 --- /dev/null +++ b/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 }), + ).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(); + }); +}); diff --git a/ui/goose2/src/shared/api/acpNotificationHandler.ts b/ui/goose2/src/shared/api/acpNotificationHandler.ts index 61435bf8..02857104 100644 --- a/ui/goose2/src/shared/api/acpNotificationHandler.ts +++ b/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(); - handleReplayUserMessageChunk(sessionId, messageId, update.content); + const messageId = getReplayMessageId(update) ?? crypto.randomUUID(); + 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 msg = - findReplayMessageWithToolCall(sessionId, update.toolCallId) ?? - (replayMessageId - ? getBufferedMessage(sessionId, replayMessageId) - : undefined); + const created = getReplayCreated(update); + const replayMessageId = getReplayMessageId(update); + const trackedMessageId = getTrackedReplayAssistantMessageId(sessionId); + const replayMsg = replayMessageId + ? getBufferedMessage(sessionId, replayMessageId) + : 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, diff --git a/ui/goose2/src/shared/api/acpReplayAssistant.ts b/ui/goose2/src/shared/api/acpReplayAssistant.ts index 811a8bbc..2e7b3158 100644 --- a/ui/goose2/src/shared/api/acpReplayAssistant.ts +++ b/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, diff --git a/ui/goose2/src/shared/api/acpReplayMetadata.ts b/ui/goose2/src/shared/api/acpReplayMetadata.ts new file mode 100644 index 00000000..f7d81ae0 --- /dev/null +++ b/ui/goose2/src/shared/api/acpReplayMetadata.ts @@ -0,0 +1,57 @@ +type ReplayMetadataSource = { + _meta?: Record | 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 | 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 { + return typeof value === "object" && value !== null && !Array.isArray(value); +} diff --git a/ui/goose2/src/shared/api/acpSkillReplayChips.ts b/ui/goose2/src/shared/api/acpSkillReplayChips.ts index 872f8d8c..f3d21408 100644 --- a/ui/goose2/src/shared/api/acpSkillReplayChips.ts +++ b/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); }