From f93abb7564aacbc6073a4067fc186dca6afdf81a Mon Sep 17 00:00:00 2001 From: Federico Jaramillo Martinez Date: Thu, 7 May 2026 22:33:16 +0200 Subject: [PATCH] Unify streamed tool message normalization --- src/client/src/chatMessages.ts | 2 +- src/client/src/chatTranscript.ts | 18 +++++++++++------- src/client/src/components/PromptEditor.ts | 15 ++++++++------- src/client/src/sessionSocket.ts | 4 ++-- src/server/sessions/piSessionService.ts | 20 ++++++++++++++++++-- 5 files changed, 40 insertions(+), 19 deletions(-) diff --git a/src/client/src/chatMessages.ts b/src/client/src/chatMessages.ts index 3bf27ed..990a500 100644 --- a/src/client/src/chatMessages.ts +++ b/src/client/src/chatMessages.ts @@ -20,7 +20,7 @@ export function appendText(messages: ChatLine[], role: ChatLine["role"], text: s return [...messages, textMessage(role, text)]; } -function normalizeMessage(message: unknown): ChatLine[] { +export function normalizeMessage(message: unknown): ChatLine[] { if (getString(message, "role") === "bashExecution") return [normalizeBashExecution(message)]; const role = normalizeRole(getString(message, "role")); const parts = normalizeContent(getProperty(message, "content"), message); diff --git a/src/client/src/chatTranscript.ts b/src/client/src/chatTranscript.ts index ff92abd..04f6342 100644 --- a/src/client/src/chatTranscript.ts +++ b/src/client/src/chatTranscript.ts @@ -1,12 +1,12 @@ -import { appendText, textMessage } from "./chatMessages"; -import type { ChatLine, ChatPart } from "./components/shared"; +import { appendText, normalizeMessage, textMessage } from "./chatMessages"; +import type { ChatLine } from "./components/shared"; import { appendShellChunk, finalizeShellMessage, shellStartMessage } from "./shellMessages"; import type { SessionUiEvent } from "./sessionSocket"; export function applyTranscriptEvent(messages: ChatLine[], event: SessionUiEvent): ChatLine[] | undefined { if (event.type === "assistant.delta") return appendText(messages, "assistant", event.text); - if (event.type === "tool.start") return appendPart(messages, "assistant", { type: "toolCall", toolName: event.toolName, summary: event.summary }); - if (event.type === "tool.end") return [...messages, { role: "tool", parts: [{ type: "toolResult", toolName: event.toolName, text: event.text, isError: event.isError }] }]; + if (event.type === "tool.start") return appendNormalized(messages, { role: "assistant", content: [{ type: "toolCall", name: event.toolName, arguments: event.args }] }); + if (event.type === "tool.end") return appendNormalized(messages, { role: "toolResult", toolName: event.toolName, content: event.content ?? [{ type: "text", text: event.text }], isError: event.isError }); if (event.type === "shell.start") return [...messages, shellStartMessage(event.command, event.excludeFromContext)]; if (event.type === "shell.chunk") return appendShellChunk(messages, event.chunk); if (event.type === "shell.end") return finalizeShellMessage(messages, event); @@ -15,8 +15,12 @@ export function applyTranscriptEvent(messages: ChatLine[], event: SessionUiEvent return undefined; } -function appendPart(messages: ChatLine[], role: ChatLine["role"], part: ChatPart): ChatLine[] { +function appendNormalized(messages: ChatLine[], rawMessage: unknown): ChatLine[] { + return normalizeMessage(rawMessage).reduce(appendLine, messages); +} + +function appendLine(messages: ChatLine[], line: ChatLine): ChatLine[] { const last = messages.at(-1); - if (last?.role === role) return [...messages.slice(0, -1), { ...last, parts: [...last.parts, part] }]; - return [...messages, { role, parts: [part] }]; + if (last?.role === line.role) return [...messages.slice(0, -1), { ...last, parts: [...last.parts, ...line.parts] }]; + return [...messages, line]; } diff --git a/src/client/src/components/PromptEditor.ts b/src/client/src/components/PromptEditor.ts index a17ee7f..acb3206 100644 --- a/src/client/src/components/PromptEditor.ts +++ b/src/client/src/components/PromptEditor.ts @@ -97,7 +97,7 @@ export class PromptEditor extends LitElement { .map((command) => ({ kind: "command", replaceFrom: trigger.from, - replaceTo: this.draft.length, + replaceTo: trigger.to, insertText: `/${command.name}`, detail: command.source, ...(command.description === undefined ? {} : { description: command.description }), @@ -107,18 +107,19 @@ export class PromptEditor extends LitElement { if (version !== this.requestVersion) return; this.completions = files .slice(0, 12) - .map((file) => ({ kind: "file", replaceFrom: trigger.from, replaceTo: this.draft.length, insertText: `@${file.path}`, detail: file.kind })); + .map((file) => ({ kind: "file", replaceFrom: trigger.from, replaceTo: trigger.to, insertText: `@${file.path}`, detail: file.kind })); } } - private currentTrigger(): { kind: "command" | "file"; query: string; from: number; fileKind?: FileSuggestion["kind"] } | undefined { - const beforeCursor = this.draft; - if (beforeCursor.endsWith("@ ")) return { kind: "file", query: "", from: beforeCursor.length - 2, fileKind: "untracked" }; + private currentTrigger(): { kind: "command" | "file"; query: string; from: number; to: number; fileKind?: FileSuggestion["kind"] } | undefined { + const cursor = this.textarea?.selectionStart ?? this.draft.length; + const beforeCursor = this.draft.slice(0, cursor); + if (beforeCursor.endsWith("@ ")) return { kind: "file", query: "", from: beforeCursor.length - 2, to: cursor, fileKind: "untracked" }; const tokenStart = Math.max(beforeCursor.lastIndexOf(" "), beforeCursor.lastIndexOf("\n")) + 1; const token = beforeCursor.slice(tokenStart); - if (token.startsWith("/")) return { kind: "command", query: token.slice(1), from: tokenStart }; - if (token.startsWith("@")) return { kind: "file", query: token.slice(1), from: tokenStart }; + if (token.startsWith("/") && tokenStart === 0) return { kind: "command", query: token.slice(1), from: tokenStart, to: cursor }; + if (token.startsWith("@")) return { kind: "file", query: token.slice(1), from: tokenStart, to: cursor }; return undefined; } diff --git a/src/client/src/sessionSocket.ts b/src/client/src/sessionSocket.ts index 8114b0a..016667f 100644 --- a/src/client/src/sessionSocket.ts +++ b/src/client/src/sessionSocket.ts @@ -2,8 +2,8 @@ import { globalSessionEvents, sessionEvents, type SessionActivity, type SessionS export type SessionUiEvent = | { type: "assistant.delta"; text: string } - | { type: "tool.start"; toolName: string; summary: string } - | { type: "tool.end"; toolName: string; text: string; isError: boolean } + | { type: "tool.start"; toolName: string; summary: string; args?: unknown } + | { type: "tool.end"; toolName: string; text: string; isError: boolean; content?: unknown } | { type: "shell.start"; command: string; excludeFromContext?: boolean } | { type: "shell.chunk"; chunk: string } | { type: "shell.end"; output?: string; exitCode?: number | null; cancelled?: boolean; truncated?: boolean; fullOutputPath?: string; isError?: boolean } diff --git a/src/server/sessions/piSessionService.ts b/src/server/sessions/piSessionService.ts index a2edc95..749ecbf 100644 --- a/src/server/sessions/piSessionService.ts +++ b/src/server/sessions/piSessionService.ts @@ -355,10 +355,12 @@ function toClientEvent(event: unknown): unknown { return { type: "assistant.delta", text: getString(assistantMessageEvent, "delta") ?? "" }; } if (eventType === "tool_execution_start") { - return { type: "tool.start", toolName: getString(event, "toolName") ?? "", toolCallId: getString(event, "toolCallId") ?? "", summary: summarizeToolArgs(getProperty(event, "args")) }; + const args = getProperty(event, "args"); + return { type: "tool.start", toolName: getString(event, "toolName") ?? "", toolCallId: getString(event, "toolCallId") ?? "", summary: summarizeToolArgs(args), args }; } if (eventType === "tool_execution_end") { - return { type: "tool.end", toolName: getString(event, "toolName") ?? "", toolCallId: getString(event, "toolCallId") ?? "", text: stringifyToolResult(getProperty(event, "result")), isError: getBoolean(event, "isError") === true }; + const result = getProperty(event, "result"); + return { type: "tool.end", toolName: getString(event, "toolName") ?? "", toolCallId: getString(event, "toolCallId") ?? "", text: stringifyToolResult(result), content: toolResultContent(result), isError: getBoolean(event, "isError") === true }; } if (eventType === "agent_start") return { type: "agent.start" }; if (eventType === "agent_end") return { type: "agent.end" }; @@ -387,12 +389,26 @@ function shortToolValue(value: unknown): string { return ""; } +function toolResultContent(result: unknown): unknown { + if (isRecord(result)) { + const content = getProperty(result, "content"); + if (content !== undefined) return content; + const text = getString(result, "text") ?? getString(result, "output"); + if (text !== undefined) return [{ type: "text", text }]; + } + if (typeof result === "string") return [{ type: "text", text: result }]; + return result; +} + function stringifyToolResult(result: unknown): string { if (typeof result === "string") return result; if (Array.isArray(result)) return result.map(stringifyToolResult).filter((text) => text !== "").join("\n"); if (isRecord(result)) { + if (getString(result, "type") === "image") return "[image]"; const text = getString(result, "text") ?? getString(result, "content") ?? getString(result, "output"); if (text !== undefined) return text; + const content = getProperty(result, "content"); + if (Array.isArray(content)) return stringifyToolResult(content); return JSON.stringify(result, null, 2); } return stringifyPrimitive(result);