Harden Assistant output hygiene

This commit is contained in:
rcourtman
2026-06-05 15:15:56 +01:00
parent bbef81274c
commit 75706f8a2c
7 changed files with 276 additions and 51 deletions
@@ -120,7 +120,9 @@ runtime cost control, and shared AI transport surfaces.
`tool_start` / `tool_end`, approval, or question blocks; if a provider emits
`pulse_*` / `patrol_*` calls, DSML, XML/function-call envelopes, or JSON
tool-call shapes as text content, the chat runtime must strip them before
streaming, persistence, and frontend rendering.
streaming, persistence, and frontend rendering. Token accounting and other
provider metadata remain runtime/accounting data, not normal transcript
prose.
4. Add or change Patrol, alert-analysis, or remediation transport through `internal/api/ai_handlers.go`, `internal/api/ai_intelligence_handlers.go`, and `frontend-modern/src/api/patrol.ts`
Provider preflight diagnostics returned from `internal/api/ai_handlers.go`
must reuse the Patrol runtime failure classifier in `internal/ai/` and
@@ -310,14 +310,6 @@ export const MessageItem: Component<MessageItemProps> = (props) => {
</div>
</div>
</Show>
<Show when={props.message.tokens && !props.message.isStreaming}>
<div class="mt-1 flex justify-end">
<span class="text-[9px] text-muted font-mono">
{props.message.tokens!.input} in · {props.message.tokens!.output} out
</span>
</div>
</Show>
</div>
</div>
</div>
@@ -374,7 +374,7 @@ describe('MessageItem', () => {
});
describe('token display', () => {
it('shows token counts when tokens are provided and not streaming', () => {
it('keeps token counts out of the visible transcript', () => {
render(() => (
<MessageItem
message={makeMessage({
@@ -386,21 +386,6 @@ describe('MessageItem', () => {
/>
));
expect(screen.getByText('500 in · 200 out')).toBeInTheDocument();
});
it('does not show token counts when streaming', () => {
render(() => (
<MessageItem
message={makeMessage({
role: 'assistant',
tokens: { input: 500, output: 200 },
isStreaming: true,
})}
{...makeHandlers()}
/>
));
expect(screen.queryByText('500 in · 200 out')).not.toBeInTheDocument();
});
@@ -1,5 +1,10 @@
import { describe, expect, it } from 'vitest';
import { stripAssistantOutputArtifacts } from '../assistantOutputHygiene';
import {
appendVisibleTextBeforeAssistantOutputArtifacts,
createAssistantOutputArtifactStreamState,
flushPendingAssistantOutputText,
stripAssistantOutputArtifacts,
} from '../assistantOutputHygiene';
describe('stripAssistantOutputArtifacts', () => {
it('strips plain Pulse function-call leaks while preserving prose before them', () => {
@@ -22,10 +27,52 @@ describe('stripAssistantOutputArtifacts', () => {
expect(result.stripped).toBe(true);
});
it('strips pulse-like tool calls even when the frontend has not learned a new tool name yet', () => {
const result = stripAssistantOutputArtifacts(
'Checking it now.\npulse_future_tool(target_host="current_resource")',
);
expect(result.text).toBe('Checking it now.');
expect(result.stripped).toBe(true);
});
it('leaves ordinary prose and unrelated function calls alone', () => {
expect(stripAssistantOutputArtifacts('Call helper(target="x") in the example.')).toEqual({
text: 'Call helper(target="x") in the example.',
stripped: false,
});
});
it('holds split tool-name prefixes until the next stream delta proves the shape', () => {
const state = createAssistantOutputArtifactStreamState();
expect(appendVisibleTextBeforeAssistantOutputArtifacts(state, 'I will check pu')).toEqual({
text: 'I will check ',
stripped: false,
});
expect(
appendVisibleTextBeforeAssistantOutputArtifacts(
state,
'lse_read(target_host="current_resource", command="lsblk")',
),
).toEqual({
text: '',
stripped: true,
});
expect(flushPendingAssistantOutputText(state)).toBe('');
});
it('releases a held prefix when the next delta proves it is normal prose', () => {
const state = createAssistantOutputArtifactStreamState();
expect(appendVisibleTextBeforeAssistantOutputArtifacts(state, 'The p')).toEqual({
text: 'The ',
stripped: false,
});
expect(appendVisibleTextBeforeAssistantOutputArtifacts(state, 'latform is healthy.')).toEqual({
text: 'platform is healthy.',
stripped: false,
});
expect(flushPendingAssistantOutputText(state)).toBe('');
});
});
@@ -800,6 +800,31 @@ describe('useChat', () => {
dispose();
});
it('strips serialized Pulse tool calls split across streamed content deltas', async () => {
const { getFireEvent } = setupWithEventCapture();
const { value: chat, dispose } = withRoot(() => useChat({ sessionId: 's' }));
await chat.sendMessage('how many devices in this');
const fire = getFireEvent();
fire({ type: 'content', data: 'I will check pu' });
fire({
type: 'content',
data: 'lse_read(target_host="current_resource", command="ls /dev | wc -l")',
});
fire({ type: 'content', data: 'raw arguments that should stay hidden' });
const assistant = chat.messages().find((m) => m.role === 'assistant')!;
expect(assistant.content).toBe('I will check ');
expect(assistant.content).not.toContain('pulse_read');
expect(assistant.content).not.toContain('target_host');
expect(assistant.content).not.toContain('raw arguments');
expect(assistant.streamEvents?.filter((e) => e.type === 'content')).toEqual([
{ type: 'content', content: 'I will check ' },
]);
dispose();
});
it('resumes visible content after a governed tool boundary clears a raw leak', async () => {
const { getFireEvent } = setupWithEventCapture();
const { value: chat, dispose } = withRoot(() => useChat({ sessionId: 's' }));
@@ -821,7 +846,7 @@ describe('useChat', () => {
const assistant = chat.messages().find((m) => m.role === 'assistant')!;
expect(assistant.content).toBe(
'I will inspect the device nodes.There are 42 device entries.',
'I will inspect the device nodes. There are 42 device entries.',
);
expect(assistant.content).not.toContain('pulse_read');
expect(assistant.content).not.toContain('raw arguments');
@@ -3,10 +3,14 @@ const RAW_TOOL_MARKERS = [
'</DSML',
'<||DSML||',
'</||DSML||',
'</DSML',
'<||/DSML||',
'<|DSML|',
'</|DSML|',
'<||DSML||',
'</||DSML||',
'<|/DSML|',
'<||/DSML||',
'<tool_call',
'</tool_call',
'<tool_calls',
@@ -21,10 +25,22 @@ const RAW_TOOL_MARKERS = [
];
const jsonToolCallLeakRe =
/(?:^|\n)[ \t]*(?:```[ \t]*(?:json|JSON)?[ \t]*\n?[ \t]*)?\{[ \t\n]*"name"[ \t]*:[ \t]*"((?:pulse|patrol)_[a-zA-Z0-9_]*)"/;
const functionToolCallLeakRe = /(?:^|[^a-zA-Z0-9_])((?:pulse|patrol)_[a-zA-Z0-9_]*)[ \t\r\n]*\(/;
/(?:^|\n)[ \t]*(?:```[ \t]*(?:json|JSON)?[ \t]*\n?[ \t]*)?\{[ \t\n]*"name"[ \t]*:[ \t]*"([a-zA-Z_][a-zA-Z0-9_]*)"/g;
const functionToolCallLeakRe = /(?:^|[^a-zA-Z0-9_])([a-zA-Z_][a-zA-Z0-9_]*)[ \t\r\n]*\(/g;
const minimaxToolCallLeakRe = /^minimax:tool_call\b/m;
export interface AssistantOutputArtifactStreamState {
visibleText: string;
pendingText: string;
}
export function createAssistantOutputArtifactStreamState(): AssistantOutputArtifactStreamState {
return {
visibleText: '',
pendingText: '',
};
}
export function stripAssistantOutputArtifacts(content: string): {
text: string;
stripped: boolean;
@@ -36,6 +52,47 @@ export function stripAssistantOutputArtifacts(content: string): {
return { text: content.slice(0, idx).trimEnd(), stripped: true };
}
export function appendVisibleTextBeforeAssistantOutputArtifacts(
state: AssistantOutputArtifactStreamState,
content: string,
): {
text: string;
stripped: boolean;
} {
if (!content && !state.pendingText) {
return { text: '', stripped: false };
}
const text = state.pendingText + content;
state.pendingText = '';
const existing = state.visibleText;
const candidate = existing + text;
const idx = assistantOutputArtifactIndex(candidate);
if (idx < 0) {
const { visible, held } = splitTrailingPotentialToolNamePrefix(text);
state.visibleText += visible;
state.pendingText = held;
return { text: visible, stripped: false };
}
const safeText = candidate.slice(0, idx).trimEnd();
const visibleDelta = safeText.length > existing.length ? safeText.slice(existing.length) : '';
state.visibleText = safeText;
state.pendingText = '';
return { text: visibleDelta, stripped: true };
}
export function flushPendingAssistantOutputText(state: AssistantOutputArtifactStreamState): string {
if (!state.pendingText) {
return '';
}
const text = state.pendingText;
state.pendingText = '';
state.visibleText += text;
return text;
}
function assistantOutputArtifactIndex(content: string): number {
if (!content) return -1;
@@ -51,15 +108,8 @@ function assistantOutputArtifactIndex(content: string): number {
record(content.indexOf(marker));
}
const jsonMatch = jsonToolCallLeakRe.exec(content);
if (jsonMatch) {
record(jsonMatch.index);
}
const functionMatch = functionToolCallLeakRe.exec(content);
if (functionMatch?.[1]) {
record(functionMatch.index + functionMatch[0].lastIndexOf(functionMatch[1]));
}
record(findJSONToolCallLeak(content));
record(findFunctionToolCallLeak(content));
const minimaxMatch = minimaxToolCallLeakRe.exec(content);
if (minimaxMatch) {
@@ -68,3 +118,65 @@ function assistantOutputArtifactIndex(content: string): number {
return first;
}
function findJSONToolCallLeak(content: string): number {
for (const match of content.matchAll(jsonToolCallLeakRe)) {
const name = match[1] || '';
if (isAssistantToolLikeName(name)) {
return match.index ?? -1;
}
}
return -1;
}
function findFunctionToolCallLeak(content: string): number {
for (const match of content.matchAll(functionToolCallLeakRe)) {
const name = match[1] || '';
if (isAssistantToolLikeName(name)) {
return (match.index ?? 0) + match[0].lastIndexOf(name);
}
}
return -1;
}
function splitTrailingPotentialToolNamePrefix(content: string): {
visible: string;
held: string;
} {
if (!content) {
return { visible: '', held: '' };
}
let start = content.length;
while (start > 0 && isToolNameCharacter(content[start - 1])) {
start -= 1;
}
if (start === content.length) {
return { visible: content, held: '' };
}
const token = content.slice(start);
if (isKnownAssistantToolNamePrefix(token)) {
return { visible: content.slice(0, start), held: token };
}
return { visible: content, held: '' };
}
function isToolNameCharacter(char: string): boolean {
return /[a-zA-Z0-9_]/.test(char);
}
function isAssistantToolLikeName(name: string): boolean {
return /^(?:pulse|patrol)_[a-zA-Z0-9_]+$/.test(name);
}
function isKnownAssistantToolNamePrefix(prefix: string): boolean {
if (!prefix) return false;
return (
'pulse_'.startsWith(prefix) ||
'patrol_'.startsWith(prefix) ||
/^pulse_[a-zA-Z0-9_]*$/.test(prefix) ||
/^patrol_[a-zA-Z0-9_]*$/.test(prefix)
);
}
@@ -10,7 +10,13 @@ import {
import { notificationStore } from '@/stores/notifications';
import { logger } from '@/utils/logger';
import { normalizeChatToolName } from '@/utils/chatIdentifiers';
import { stripAssistantOutputArtifacts } from '../assistantOutputHygiene';
import {
appendVisibleTextBeforeAssistantOutputArtifacts,
createAssistantOutputArtifactStreamState,
flushPendingAssistantOutputText,
stripAssistantOutputArtifacts,
type AssistantOutputArtifactStreamState,
} from '../assistantOutputHygiene';
import type {
ChatMessage,
ToolExecution,
@@ -72,6 +78,25 @@ export function useChat(options: UseChatOptions = {}) {
let pendingBackendAbort: Promise<void> | null = null;
let isDrainingQueuedFollowUps = false;
const suppressedRawContentMessageIds = new Set<string>();
const outputArtifactStreamStates = new Map<string, AssistantOutputArtifactStreamState>();
const outputArtifactStateFor = (assistantId: string) => {
let state = outputArtifactStreamStates.get(assistantId);
if (!state) {
state = createAssistantOutputArtifactStreamState();
outputArtifactStreamStates.set(assistantId, state);
}
return state;
};
const clearOutputArtifactState = (assistantId: string) => {
outputArtifactStreamStates.delete(assistantId);
};
const clearSuppressedOutputBoundary = (assistantId: string) => {
suppressedRawContentMessageIds.delete(assistantId);
clearOutputArtifactState(assistantId);
};
const abortBackendSession = (targetSessionId: string): Promise<void> | null => {
const normalizedSessionId = targetSessionId.trim();
@@ -201,6 +226,24 @@ export function useChat(options: UseChatOptions = {}) {
};
};
const appendMessageContent = (msg: ChatMessage, content: string): string => {
const existing = msg.content || '';
if (!existing || !content) {
return existing + content;
}
const events = msg.streamEvents || [];
const lastEvent = events[events.length - 1];
if (!lastEvent || lastEvent.type === 'content') {
return existing + content;
}
if (/\s$/.test(existing) || /^\s|^[,.;:!?)]/.test(content)) {
return existing + content;
}
return `${existing} ${content}`;
};
// Process stream events
const extractText = (value: unknown): string => {
if (typeof value === 'string') return value;
@@ -387,17 +430,19 @@ export function useChat(options: UseChatOptions = {}) {
}
const content = extractText(event.data);
if (!content) return msg;
const visible = stripAssistantOutputArtifacts(content);
const visible = appendVisibleTextBeforeAssistantOutputArtifacts(
outputArtifactStateFor(assistantId),
content,
);
if (visible.stripped) {
suppressedRawContentMessageIds.add(assistantId);
}
if (!visible.text) return msg;
const existing = msg.content || '';
// Add to streamEvents for chronological display
const updated = addStreamEvent(msg, { type: 'content', content: visible.text });
return {
...updated,
content: existing + visible.text,
content: appendMessageContent(msg, visible.text),
};
}
@@ -416,7 +461,7 @@ export function useChat(options: UseChatOptions = {}) {
}
case 'tool_start': {
suppressedRawContentMessageIds.delete(assistantId);
clearSuppressedOutputBoundary(assistantId);
const data = (event.data || {}) as {
id?: string;
name?: string;
@@ -453,7 +498,7 @@ export function useChat(options: UseChatOptions = {}) {
}
case 'tool_end': {
suppressedRawContentMessageIds.delete(assistantId);
clearSuppressedOutputBoundary(assistantId);
const data = event.data as {
id?: string;
name: string;
@@ -556,7 +601,7 @@ export function useChat(options: UseChatOptions = {}) {
}
case 'approval_needed': {
suppressedRawContentMessageIds.delete(assistantId);
clearSuppressedOutputBoundary(assistantId);
const data = event.data as {
command: string;
tool_id: string;
@@ -637,7 +682,7 @@ export function useChat(options: UseChatOptions = {}) {
}
case 'question': {
suppressedRawContentMessageIds.delete(assistantId);
clearSuppressedOutputBoundary(assistantId);
const data = event.data as { question_id: string; questions: Array<any> };
const pendingQuestion: PendingQuestion = {
@@ -672,22 +717,37 @@ export function useChat(options: UseChatOptions = {}) {
}
case 'done': {
const pendingText = suppressedRawContentMessageIds.has(assistantId)
? ''
: flushPendingAssistantOutputText(outputArtifactStateFor(assistantId));
suppressedRawContentMessageIds.delete(assistantId);
clearOutputArtifactState(assistantId);
const flushedMsg = pendingText
? {
...addStreamEvent(msg, { type: 'content', content: pendingText }),
content: appendMessageContent(msg, pendingText),
}
: msg;
const tokens = extractTokens(event.data);
if (tokens && (tokens.input > 0 || tokens.output > 0)) {
return {
...msg,
...flushedMsg,
isStreaming: false,
pendingTools: [],
tokens,
workflowStatus: undefined,
};
}
return { ...msg, isStreaming: false, pendingTools: [], workflowStatus: undefined };
return {
...flushedMsg,
isStreaming: false,
pendingTools: [],
workflowStatus: undefined,
};
}
case 'error': {
suppressedRawContentMessageIds.delete(assistantId);
clearSuppressedOutputBoundary(assistantId);
const errorMsg = extractErrorMessage(event.data);
// Keep any content streamed before the failure; surface the error
// as a distinct, recoverable block rather than overwriting the answer.
@@ -922,12 +982,14 @@ export function useChat(options: UseChatOptions = {}) {
setMessages(
msgs.map((m) => {
const toolCalls = m.tool_calls || [];
const content =
m.role === 'assistant' ? stripAssistantOutputArtifacts(m.content).text : m.content;
const streamEvents =
m.role === 'assistant' ? buildPersistedStreamEvents(m.content, toolCalls) : undefined;
m.role === 'assistant' ? buildPersistedStreamEvents(content, toolCalls) : undefined;
return {
id: m.id,
role: m.role,
content: m.content,
content,
timestamp: new Date(m.timestamp),
model: m.model,
toolCalls,