From 9abc0217e9a8d648d503aa6e1e4a69af9301ea9f Mon Sep 17 00:00:00 2001 From: Mohamed Boudra Date: Thu, 23 Apr 2026 15:28:25 +0700 Subject: [PATCH] refactor: extract types and simplify activity grouping --- packages/app/src/types/agent-activity.test.ts | 116 +++++++++ packages/app/src/types/agent-activity.ts | 245 +++++++----------- 2 files changed, 216 insertions(+), 145 deletions(-) create mode 100644 packages/app/src/types/agent-activity.test.ts diff --git a/packages/app/src/types/agent-activity.test.ts b/packages/app/src/types/agent-activity.test.ts new file mode 100644 index 000000000..55ad48bac --- /dev/null +++ b/packages/app/src/types/agent-activity.test.ts @@ -0,0 +1,116 @@ +import { describe, expect, it } from "vitest"; +import { groupActivities, type AgentActivity } from "./agent-activity"; + +function timestamp(seconds: number): Date { + return new Date(`2026-01-01T00:00:${seconds.toString().padStart(2, "0")}.000Z`); +} + +describe("groupActivities", () => { + it("groups text chunks and merges tool call updates at the first tool position", () => { + const activities: AgentActivity[] = [ + { + timestamp: timestamp(1), + update: { + kind: "agent_message_chunk", + content: { type: "text", text: "Hel" }, + }, + }, + { + timestamp: timestamp(2), + update: { + kind: "agent_message_chunk", + content: { type: "text", text: "lo" }, + }, + }, + { + timestamp: timestamp(3), + update: { + kind: "tool_call", + toolCallId: "tool_1", + title: "Read file", + status: "in_progress", + toolKind: "read", + input: { path: "README.md" }, + }, + }, + { + timestamp: timestamp(4), + update: { + kind: "agent_thought_chunk", + content: { type: "text", text: "Checking context" }, + }, + }, + { + timestamp: timestamp(5), + update: { + kind: "tool_call_update", + toolCallId: "tool_1", + status: "completed", + output: { ok: true }, + }, + }, + ]; + + expect(groupActivities(activities)).toEqual([ + { + kind: "grouped_text", + messageType: "agent", + text: "Hello", + startTimestamp: timestamp(1), + endTimestamp: timestamp(2), + }, + { + kind: "merged_tool_call", + toolCallId: "tool_1", + title: "Read file", + status: "completed", + toolKind: "read", + input: { path: "README.md" }, + output: { ok: true }, + content: undefined, + locations: undefined, + startTimestamp: timestamp(3), + endTimestamp: timestamp(5), + }, + { + kind: "grouped_text", + messageType: "thought", + text: "Checking context", + startTimestamp: timestamp(4), + endTimestamp: timestamp(4), + }, + ]); + }); + + it("creates a merged tool call when an update arrives before the initial call", () => { + const activities: AgentActivity[] = [ + { + timestamp: timestamp(1), + update: { + kind: "tool_call_update", + toolCallId: "tool_2", + title: "Shell", + status: "completed", + toolKind: "execute", + output: { exitCode: 0 }, + }, + }, + ]; + + expect(groupActivities(activities)).toEqual([ + { + kind: "merged_tool_call", + toolCallId: "tool_2", + title: "Shell", + status: "completed", + toolKind: "execute", + input: undefined, + output: { exitCode: 0 }, + content: undefined, + locations: undefined, + startTimestamp: timestamp(1), + endTimestamp: timestamp(1), + }, + ]); + }); +}); diff --git a/packages/app/src/types/agent-activity.ts b/packages/app/src/types/agent-activity.ts index 414301ed5..7e100798d 100644 --- a/packages/app/src/types/agent-activity.ts +++ b/packages/app/src/types/agent-activity.ts @@ -1,6 +1,3 @@ -/** - * Discriminated union types for session updates - */ export type SessionUpdate = | UserMessageChunk | AgentMessageChunk @@ -11,6 +8,22 @@ export type SessionUpdate = | AvailableCommandsUpdate | CurrentModeUpdate; +export type ToolCallStatus = "pending" | "in_progress" | "completed" | "failed"; + +export type ToolKind = + | "read" + | "edit" + | "delete" + | "move" + | "search" + | "execute" + | "think" + | "fetch" + | "switch_mode" + | "other"; + +type TextMessageType = "user" | "agent" | "thought"; + export interface UserMessageChunk { kind: "user_message_chunk"; content: { @@ -39,18 +52,8 @@ export interface ToolCall { kind: "tool_call"; toolCallId: string; title: string; - status?: "pending" | "in_progress" | "completed" | "failed"; - toolKind?: - | "read" - | "edit" - | "delete" - | "move" - | "search" - | "execute" - | "think" - | "fetch" - | "switch_mode" - | "other"; + status?: ToolCallStatus; + toolKind?: ToolKind; input?: Record; output?: Record; content?: unknown[]; @@ -61,19 +64,8 @@ export interface ToolCallUpdate { kind: "tool_call_update"; toolCallId: string; title?: string | null; - status?: "pending" | "in_progress" | "completed" | "failed" | null; - toolKind?: - | "read" - | "edit" - | "delete" - | "move" - | "search" - | "execute" - | "think" - | "fetch" - | "switch_mode" - | "other" - | null; + status?: ToolCallStatus | null; + toolKind?: ToolKind | null; input?: Record; output?: Record; content?: unknown[] | null; @@ -102,44 +94,25 @@ export interface CurrentModeUpdate { currentModeId: string; } -/** - * Activity item with timestamp - */ export interface AgentActivity { timestamp: Date; update: SessionUpdate; } -/** - * Grouped text message (consecutive chunks combined) - */ export interface GroupedTextMessage { kind: "grouped_text"; - messageType: "user" | "agent" | "thought"; + messageType: TextMessageType; text: string; startTimestamp: Date; endTimestamp: Date; } -/** - * Merged tool call (initial call + all updates combined) - */ export interface MergedToolCall { kind: "merged_tool_call"; toolCallId: string; title: string; - status: "pending" | "in_progress" | "completed" | "failed"; - toolKind?: - | "read" - | "edit" - | "delete" - | "move" - | "search" - | "execute" - | "think" - | "fetch" - | "switch_mode" - | "other"; + status: ToolCallStatus; + toolKind?: ToolKind; input?: Record; output?: Record; content?: unknown[]; @@ -148,89 +121,95 @@ export interface MergedToolCall { endTimestamp: Date; } -/** - * Group consecutive text chunks into messages and merge tool calls by ID - */ -export function groupActivities( - activities: AgentActivity[], -): Array { - const result: Array = []; +export type GroupedActivity = GroupedTextMessage | MergedToolCall | AgentActivity; - // Track current text group - let currentTextGroup: { - messageType: "user" | "agent" | "thought"; - chunks: string[]; - startTimestamp: Date; - endTimestamp: Date; - } | null = null; +type TextGroup = { + messageType: TextMessageType; + chunks: string[]; + startTimestamp: Date; + endTimestamp: Date; +}; - // Track tool calls by ID - const toolCallsById = new Map< - string, - { - toolCallId: string; - title: string; - status: "pending" | "in_progress" | "completed" | "failed"; - toolKind?: - | "read" - | "edit" - | "delete" - | "move" - | "search" - | "execute" - | "think" - | "fetch" - | "switch_mode" - | "other"; - input?: Record; - output?: Record; - content?: unknown[]; - locations?: unknown[]; - startTimestamp: Date; - endTimestamp: Date; - insertIndex: number; // Track where to insert in result - } - >(); +type ToolCallAccumulator = Omit & { + insertIndex: number; +}; + +function getMessageType( + update: UserMessageChunk | AgentMessageChunk | AgentThoughtChunk, +): TextMessageType { + switch (update.kind) { + case "user_message_chunk": + return "user"; + case "agent_message_chunk": + return "agent"; + case "agent_thought_chunk": + return "thought"; + } +} + +function isTextChunk( + update: SessionUpdate, +): update is UserMessageChunk | AgentMessageChunk | AgentThoughtChunk { + return ( + update.kind === "user_message_chunk" || + update.kind === "agent_message_chunk" || + update.kind === "agent_thought_chunk" + ); +} + +function isGroupedActivity(item: GroupedActivity | null): item is GroupedActivity { + return item !== null; +} + +function toMergedToolCall(toolCall: ToolCallAccumulator): MergedToolCall { + return { + kind: "merged_tool_call", + toolCallId: toolCall.toolCallId, + title: toolCall.title, + status: toolCall.status, + toolKind: toolCall.toolKind, + input: toolCall.input, + output: toolCall.output, + content: toolCall.content, + locations: toolCall.locations, + startTimestamp: toolCall.startTimestamp, + endTimestamp: toolCall.endTimestamp, + }; +} + +export function groupActivities(activities: AgentActivity[]): GroupedActivity[] { + const result: Array = []; + const toolCallsById = new Map(); + let currentTextGroup: TextGroup | null = null; function flushTextGroup() { - if (currentTextGroup) { - result.push({ - kind: "grouped_text", - messageType: currentTextGroup.messageType, - text: currentTextGroup.chunks.join(""), - startTimestamp: currentTextGroup.startTimestamp, - endTimestamp: currentTextGroup.endTimestamp, - }); - currentTextGroup = null; + if (!currentTextGroup) { + return; } + + result.push({ + kind: "grouped_text", + messageType: currentTextGroup.messageType, + text: currentTextGroup.chunks.join(""), + startTimestamp: currentTextGroup.startTimestamp, + endTimestamp: currentTextGroup.endTimestamp, + }); + currentTextGroup = null; } for (const activity of activities) { const update = activity.update; - // Handle text chunks - if ( - update.kind === "user_message_chunk" || - update.kind === "agent_message_chunk" || - update.kind === "agent_thought_chunk" - ) { - const messageType = - update.kind === "user_message_chunk" - ? "user" - : update.kind === "agent_message_chunk" - ? "agent" - : "thought"; - + if (isTextChunk(update)) { + const messageType = getMessageType(update); const text = update.content.text; - // If we have a current group of the same type, add to it if (currentTextGroup && currentTextGroup.messageType === messageType) { currentTextGroup.chunks.push(text); currentTextGroup.endTimestamp = activity.timestamp; } else { flushTextGroup(); - // Start new group currentTextGroup = { messageType, chunks: [text], @@ -238,20 +217,16 @@ export function groupActivities( endTimestamp: activity.timestamp, }; } - } - // Handle tool calls and updates - else if (update.kind === "tool_call" || update.kind === "tool_call_update") { + } else if (update.kind === "tool_call" || update.kind === "tool_call_update") { flushTextGroup(); const toolCallId = update.toolCallId; const existing = toolCallsById.get(toolCallId); if (update.kind === "tool_call") { - // Initial tool call if (!existing) { - // Create new entry and placeholder in result const insertIndex = result.length; - result.push(null as any); // Placeholder + result.push(null); toolCallsById.set(toolCallId, { toolCallId, @@ -267,7 +242,6 @@ export function groupActivities( insertIndex, }); } else { - // Update existing existing.title = update.title; if (update.status) existing.status = update.status; if (update.toolKind) existing.toolKind = update.toolKind; @@ -278,9 +252,7 @@ export function groupActivities( existing.endTimestamp = activity.timestamp; } } else { - // Tool call update if (existing) { - // Merge update into existing if (update.title) existing.title = update.title; if (update.status) existing.status = update.status; if (update.toolKind) existing.toolKind = update.toolKind; @@ -290,9 +262,8 @@ export function groupActivities( if (update.locations) existing.locations = update.locations; existing.endTimestamp = activity.timestamp; } else { - // Update without initial call - create entry const insertIndex = result.length; - result.push(null as any); // Placeholder + result.push(null); toolCallsById.set(toolCallId, { toolCallId, @@ -309,33 +280,17 @@ export function groupActivities( }); } } - } - // Handle other activities - else { + } else { flushTextGroup(); result.push(activity); } } - // Flush final text group flushTextGroup(); - // Replace placeholders with merged tool calls for (const toolCall of toolCallsById.values()) { - result[toolCall.insertIndex] = { - kind: "merged_tool_call", - toolCallId: toolCall.toolCallId, - title: toolCall.title, - status: toolCall.status, - toolKind: toolCall.toolKind, - input: toolCall.input, - output: toolCall.output, - content: toolCall.content, - locations: toolCall.locations, - startTimestamp: toolCall.startTimestamp, - endTimestamp: toolCall.endTimestamp, - }; + result[toolCall.insertIndex] = toMergedToolCall(toolCall); } - return result; + return result.filter(isGroupedActivity); }