mirror of
https://github.com/browseros-ai/BrowserOS.git
synced 2026-05-13 15:46:22 +00:00
fix: guard filesystem tools behind workspace selection and handle mid-conversation changes (#595)
* fix: remove filesystem tools when no workspace is selected - Make workingDir optional on ResolvedAgentConfig - Remove resolveSessionDir() fallback that always created a session dir, masking the no-workspace state and keeping filesystem tools available - Gate buildFilesystemToolSet() on workingDir being defined - Add workspace change detection mid-conversation — rebuilds the agent session when workspace is added, removed, or switched (same pattern as existing MCP server change detection) - download_file falls back to tmpdir() when no workspace is set - Memory/soul tools are unaffected — they use ~/BrowserOS/ paths * fix: sanitize message history when session rebuilds with different tools When a session is rebuilt due to workspace or MCP changes, the carried-over message history may contain tool parts for tools that no longer exist in the new session. The AI SDK validates messages against the current toolset and rejects parts with no matching schema. - Add toolNames getter to AiSdkAgent exposing registered tool names - Add sanitizeMessagesForToolset() to strip tool parts referencing removed tools from carried-over messages - Apply sanitization in both MCP and workspace session rebuilds * fix: prepend tool-change context to user message on session rebuild When workspace or MCP integrations change mid-conversation, prepend a [Context: ...] block to the user's message explaining what changed. This prevents the LLM from hallucinating tool usage based on patterns in the carried-over conversation history. Context messages vary by change type: - Workspace removed: lists unavailable filesystem tools, suggests selecting a working directory - Workspace added: confirms filesystem tools are available with path - Workspace switched: notes the new working directory - MCP changed: notes that some integration tools may have changed Only fires on the first message after a rebuild. Invisible in the UI. * fix: make MCP change context specific about which apps were added/removed Diff the old and new MCP server keys to produce specific context like: - "The following app integrations were disconnected: Gmail, Slack." - "The following app integrations were connected: Linear." instead of a generic "some tools may no longer be available" message. * refactor: extract shared rebuildSession helper in ChatService Eliminates the duplicated 20-line dispose→create→sanitize→store flow that existed separately in both the MCP and workspace change-detection blocks. Co-authored-by: Dani Akash <DaniAkash@users.noreply.github.com> * test: add sanitizeMessagesForToolset test suite Tests for the message sanitization that runs when a session rebuilds with a different toolset (workspace or MCP change mid-conversation): - Preserves messages with no tool parts - Preserves tool parts when tool is in the toolset - Strips tool parts when tool is NOT in the toolset - Strips multiple removed tool parts from same message - Keeps browser tools while removing filesystem tools - Removes messages that become empty after stripping - Preserves non-tool parts (reasoning, step-start, file) - Returns same references when no filtering needed - Handles empty message array and empty toolset * style: fix biome formatting in chat-service.ts --------- Co-authored-by: claude[bot] <41898282+claude[bot]@users.noreply.github.com>
This commit is contained in:
@@ -54,8 +54,14 @@ export class AiSdkAgent {
|
||||
private _messages: UIMessage[],
|
||||
private _mcpClients: Array<{ close(): Promise<void> }>,
|
||||
private conversationId: string,
|
||||
private _toolNames: Set<string>,
|
||||
) {}
|
||||
|
||||
/** Tool names registered on this agent — used to sanitize messages during session rebuilds. */
|
||||
get toolNames(): Set<string> {
|
||||
return this._toolNames
|
||||
}
|
||||
|
||||
static async create(config: AiSdkAgentConfig): Promise<AiSdkAgent> {
|
||||
const contextWindow =
|
||||
config.resolvedConfig.contextWindowSize ??
|
||||
@@ -160,10 +166,11 @@ export class AiSdkAgent {
|
||||
}
|
||||
}
|
||||
|
||||
// Add filesystem tools (Pi coding agent) — skip in chat mode (read-only)
|
||||
const filesystemTools = config.resolvedConfig.chatMode
|
||||
? {}
|
||||
: buildFilesystemToolSet(config.resolvedConfig.workingDir)
|
||||
// Add filesystem tools — skip in chat mode (read-only) and when no workspace is selected
|
||||
const filesystemTools =
|
||||
!config.resolvedConfig.chatMode && config.resolvedConfig.workingDir
|
||||
? buildFilesystemToolSet(config.resolvedConfig.workingDir)
|
||||
: {}
|
||||
const memoryTools = config.resolvedConfig.chatMode
|
||||
? {}
|
||||
: buildMemoryToolSet()
|
||||
@@ -269,6 +276,7 @@ export class AiSdkAgent {
|
||||
[],
|
||||
clients,
|
||||
config.resolvedConfig.conversationId,
|
||||
new Set(Object.keys(tools)),
|
||||
)
|
||||
}
|
||||
|
||||
|
||||
@@ -44,3 +44,37 @@ export function hasMessageContent(message: UIMessage): boolean {
|
||||
export function filterValidMessages(messages: UIMessage[]): UIMessage[] {
|
||||
return messages.filter(hasMessageContent)
|
||||
}
|
||||
|
||||
/**
|
||||
* Remove tool parts that reference tools not present in the given toolset.
|
||||
*
|
||||
* When a session is rebuilt with a different set of tools (e.g., workspace
|
||||
* removed mid-conversation or MCP server disconnected), the carried-over
|
||||
* message history may contain tool parts for tools that no longer exist.
|
||||
* The AI SDK validates messages against the current toolset and rejects
|
||||
* parts with no matching schema.
|
||||
*
|
||||
* Tool parts use the type format `tool-${toolName}` (static tools) or
|
||||
* `dynamic-tool` (dynamic tools). This function filters out static tool
|
||||
* parts whose tool name is not in the provided set.
|
||||
*/
|
||||
export function sanitizeMessagesForToolset(
|
||||
messages: UIMessage[],
|
||||
toolNames: Set<string>,
|
||||
): UIMessage[] {
|
||||
return messages
|
||||
.map((msg) => {
|
||||
const filteredParts = msg.parts.filter((part) => {
|
||||
// Static tool parts have type `tool-${toolName}`
|
||||
if (typeof part.type === 'string' && part.type.startsWith('tool-')) {
|
||||
const toolName = part.type.slice(5)
|
||||
if (!toolNames.has(toolName)) return false
|
||||
}
|
||||
return true
|
||||
})
|
||||
|
||||
if (filteredParts.length === msg.parts.length) return msg
|
||||
return { ...msg, parts: filteredParts }
|
||||
})
|
||||
.filter(hasMessageContent)
|
||||
}
|
||||
|
||||
@@ -9,6 +9,8 @@ export interface AgentSession {
|
||||
browserContext?: BrowserContext
|
||||
/** MCP server names used when the session was created, for change detection. */
|
||||
mcpServerKey?: string
|
||||
/** Workspace directory when the session was created, for change detection. */
|
||||
workingDir?: string
|
||||
}
|
||||
|
||||
export class SessionStore {
|
||||
|
||||
@@ -38,7 +38,7 @@ function contentToModelOutput(
|
||||
export function buildBrowserToolSet(
|
||||
registry: ToolRegistry,
|
||||
browser: Browser,
|
||||
workingDir: string,
|
||||
workingDir: string | undefined,
|
||||
session?: { origin?: 'sidepanel' | 'newtab'; originPageId?: number },
|
||||
): ToolSet {
|
||||
const toolSet: ToolSet = {}
|
||||
|
||||
@@ -35,7 +35,7 @@ export interface ResolvedAgentConfig {
|
||||
reasoningSummary?: string
|
||||
contextWindowSize?: number
|
||||
userSystemPrompt?: string
|
||||
workingDir: string
|
||||
workingDir?: string
|
||||
/** Whether the model supports image inputs (vision). Defaults to true. */
|
||||
supportsImages?: boolean
|
||||
/** Eval mode - enables window management tools. Defaults to false. */
|
||||
|
||||
@@ -4,16 +4,16 @@
|
||||
* SPDX-License-Identifier: AGPL-3.0-or-later
|
||||
*/
|
||||
|
||||
import { mkdir, utimes } from 'node:fs/promises'
|
||||
import path from 'node:path'
|
||||
import { createAgentUIStreamResponse, type UIMessage } from 'ai'
|
||||
import { AiSdkAgent } from '../../agent/ai-sdk-agent'
|
||||
import { formatUserMessage } from '../../agent/format-message'
|
||||
import { filterValidMessages } from '../../agent/message-validation'
|
||||
import type { SessionStore } from '../../agent/session-store'
|
||||
import {
|
||||
filterValidMessages,
|
||||
sanitizeMessagesForToolset,
|
||||
} from '../../agent/message-validation'
|
||||
import type { AgentSession, SessionStore } from '../../agent/session-store'
|
||||
import type { ResolvedAgentConfig } from '../../agent/types'
|
||||
import type { Browser } from '../../browser/browser'
|
||||
import { getSessionsDir } from '../../lib/browseros-dir'
|
||||
import type { KlavisClient } from '../../lib/clients/klavis/klavis-client'
|
||||
import { resolveLLMConfig } from '../../lib/clients/llm/config'
|
||||
import { logger } from '../../lib/logger'
|
||||
@@ -40,8 +40,6 @@ export class ChatService {
|
||||
|
||||
const llmConfig = await resolveLLMConfig(request, this.deps.browserosId)
|
||||
|
||||
const workingDir = await this.resolveSessionDir(request)
|
||||
|
||||
const agentConfig: ResolvedAgentConfig = {
|
||||
conversationId: request.conversationId,
|
||||
provider: llmConfig.provider,
|
||||
@@ -59,7 +57,7 @@ export class ChatService {
|
||||
reasoningSummary: request.reasoningSummary,
|
||||
contextWindowSize: request.contextWindowSize,
|
||||
userSystemPrompt: request.userSystemPrompt,
|
||||
workingDir,
|
||||
workingDir: request.userWorkingDir,
|
||||
supportsImages: request.supportsImages,
|
||||
chatMode: request.mode === 'chat',
|
||||
isScheduledTask: request.isScheduledTask,
|
||||
@@ -70,6 +68,7 @@ export class ChatService {
|
||||
|
||||
let session = sessionStore.get(request.conversationId)
|
||||
let isNewSession = false
|
||||
const contextChanges: string[] = []
|
||||
|
||||
// Build a stable key from enabled MCP servers for change detection
|
||||
const mcpServerKey = this.buildMcpServerKey(request.browserContext)
|
||||
@@ -81,23 +80,68 @@ export class ChatService {
|
||||
previous: session.mcpServerKey,
|
||||
current: mcpServerKey,
|
||||
})
|
||||
const previousMessages = session.agent.messages
|
||||
await session.agent.dispose()
|
||||
sessionStore.remove(request.conversationId)
|
||||
const previousMcpKey = session.mcpServerKey
|
||||
session = await this.rebuildSession(
|
||||
session,
|
||||
request,
|
||||
agentConfig,
|
||||
mcpServerKey,
|
||||
)
|
||||
|
||||
const browserContext = await this.resolvePageIds(request.browserContext)
|
||||
const agent = await AiSdkAgent.create({
|
||||
resolvedConfig: agentConfig,
|
||||
browser: this.deps.browser,
|
||||
registry: this.deps.registry,
|
||||
browserContext,
|
||||
klavisClient: this.deps.klavisClient,
|
||||
browserosId: this.deps.browserosId,
|
||||
aiSdkDevtoolsEnabled: this.deps.aiSdkDevtoolsEnabled,
|
||||
const oldServers = new Set(
|
||||
(previousMcpKey ?? '').split(',').filter(Boolean),
|
||||
)
|
||||
const newServers = new Set(mcpServerKey.split(',').filter(Boolean))
|
||||
const added = [...newServers].filter((s) => !oldServers.has(s))
|
||||
const removed = [...oldServers].filter((s) => !newServers.has(s))
|
||||
|
||||
const parts: string[] = []
|
||||
if (removed.length > 0) {
|
||||
parts.push(
|
||||
`The following app integrations were disconnected: ${removed.join(', ')}. Their tools are no longer available.`,
|
||||
)
|
||||
}
|
||||
if (added.length > 0) {
|
||||
parts.push(
|
||||
`The following app integrations were connected: ${added.join(', ')}. Their tools are now available.`,
|
||||
)
|
||||
}
|
||||
if (parts.length === 0) {
|
||||
parts.push(
|
||||
'Connected app integrations changed during this conversation. Use only tools that are currently registered.',
|
||||
)
|
||||
}
|
||||
contextChanges.push(parts.join(' '))
|
||||
}
|
||||
|
||||
// Detect workspace change mid-conversation → rebuild session
|
||||
if (session && session.workingDir !== request.userWorkingDir) {
|
||||
logger.info('Workspace changed mid-conversation, rebuilding session', {
|
||||
conversationId: request.conversationId,
|
||||
previous: session.workingDir ?? '(none)',
|
||||
current: request.userWorkingDir ?? '(none)',
|
||||
})
|
||||
session = { agent, browserContext, mcpServerKey }
|
||||
session.agent.messages = previousMessages
|
||||
sessionStore.set(request.conversationId, session)
|
||||
const previousWorkingDir = session.workingDir
|
||||
session = await this.rebuildSession(
|
||||
session,
|
||||
request,
|
||||
agentConfig,
|
||||
mcpServerKey,
|
||||
)
|
||||
|
||||
if (!request.userWorkingDir) {
|
||||
contextChanges.push(
|
||||
'The user disconnected the workspace during this conversation. Filesystem tools (filesystem_read, filesystem_write, filesystem_edit, filesystem_bash, filesystem_grep, filesystem_find, filesystem_ls) are no longer available. Return all output directly in chat. If the user asks for file operations, suggest they select a working directory from the chat toolbar.',
|
||||
)
|
||||
} else if (!previousWorkingDir) {
|
||||
contextChanges.push(
|
||||
`The user connected a workspace during this conversation. Filesystem tools are now available. Working directory: ${request.userWorkingDir}`,
|
||||
)
|
||||
} else {
|
||||
contextChanges.push(
|
||||
`The user switched workspace during this conversation. Filesystem tools now use the new working directory: ${request.userWorkingDir}`,
|
||||
)
|
||||
}
|
||||
}
|
||||
|
||||
if (!session) {
|
||||
@@ -142,7 +186,13 @@ export class ChatService {
|
||||
browserosId: this.deps.browserosId,
|
||||
aiSdkDevtoolsEnabled: this.deps.aiSdkDevtoolsEnabled,
|
||||
})
|
||||
session = { agent, hiddenWindowId, browserContext, mcpServerKey }
|
||||
session = {
|
||||
agent,
|
||||
hiddenWindowId,
|
||||
browserContext,
|
||||
mcpServerKey,
|
||||
workingDir: request.userWorkingDir,
|
||||
}
|
||||
sessionStore.set(request.conversationId, session)
|
||||
}
|
||||
|
||||
@@ -176,7 +226,13 @@ export class ChatService {
|
||||
request.selectedText,
|
||||
request.selectedTextSource,
|
||||
)
|
||||
session.agent.appendUserMessage(userContent)
|
||||
|
||||
// Prepend tool-change context when session was rebuilt mid-conversation
|
||||
const contextPrefix =
|
||||
contextChanges.length > 0
|
||||
? `${contextChanges.map((c) => `[Context: ${c}]`).join('\n')}\n\n`
|
||||
: ''
|
||||
session.agent.appendUserMessage(contextPrefix + userContent)
|
||||
|
||||
return createAgentUIStreamResponse({
|
||||
agent: session.agent.toolLoopAgent,
|
||||
@@ -263,22 +319,44 @@ export class ChatService {
|
||||
})
|
||||
}
|
||||
|
||||
private async rebuildSession(
|
||||
session: AgentSession,
|
||||
request: ChatRequest,
|
||||
agentConfig: ResolvedAgentConfig,
|
||||
mcpServerKey: string,
|
||||
): Promise<AgentSession> {
|
||||
const previousMessages = session.agent.messages
|
||||
await session.agent.dispose()
|
||||
this.deps.sessionStore.remove(request.conversationId)
|
||||
|
||||
const browserContext = await this.resolvePageIds(request.browserContext)
|
||||
const agent = await AiSdkAgent.create({
|
||||
resolvedConfig: agentConfig,
|
||||
browser: this.deps.browser,
|
||||
registry: this.deps.registry,
|
||||
browserContext,
|
||||
klavisClient: this.deps.klavisClient,
|
||||
browserosId: this.deps.browserosId,
|
||||
aiSdkDevtoolsEnabled: this.deps.aiSdkDevtoolsEnabled,
|
||||
})
|
||||
const newSession: AgentSession = {
|
||||
agent,
|
||||
browserContext,
|
||||
mcpServerKey,
|
||||
workingDir: request.userWorkingDir,
|
||||
}
|
||||
newSession.agent.messages = sanitizeMessagesForToolset(
|
||||
previousMessages,
|
||||
agent.toolNames,
|
||||
)
|
||||
this.deps.sessionStore.set(request.conversationId, newSession)
|
||||
return newSession
|
||||
}
|
||||
|
||||
private buildMcpServerKey(browserContext?: BrowserContext): string {
|
||||
const managed = browserContext?.enabledMcpServers?.slice().sort() ?? []
|
||||
const custom =
|
||||
browserContext?.customMcpServers?.map((s) => s.url).sort() ?? []
|
||||
return [...managed, ...custom].join(',')
|
||||
}
|
||||
|
||||
private async resolveSessionDir(request: ChatRequest): Promise<string> {
|
||||
const dir = request.userWorkingDir
|
||||
? request.userWorkingDir
|
||||
: path.join(getSessionsDir(), request.conversationId)
|
||||
await mkdir(dir, { recursive: true })
|
||||
if (!request.userWorkingDir) {
|
||||
const now = new Date()
|
||||
await utimes(dir, now, now).catch(() => {})
|
||||
}
|
||||
return dir
|
||||
}
|
||||
}
|
||||
|
||||
@@ -1,3 +1,4 @@
|
||||
import { tmpdir } from 'node:os'
|
||||
import { resolve } from 'node:path'
|
||||
import type { z } from 'zod'
|
||||
import type { Browser } from '../browser/browser'
|
||||
@@ -18,7 +19,7 @@ export type ToolHandler = (
|
||||
) => Promise<void>
|
||||
|
||||
export interface ToolDirectories {
|
||||
workingDir: string
|
||||
workingDir?: string
|
||||
resourcesDir?: string
|
||||
}
|
||||
|
||||
@@ -38,7 +39,7 @@ export function resolveWorkingPath(
|
||||
targetPath: string,
|
||||
cwd?: string,
|
||||
): string {
|
||||
return resolve(cwd ?? ctx.directories.workingDir, targetPath)
|
||||
return resolve(cwd ?? ctx.directories.workingDir ?? tmpdir(), targetPath)
|
||||
}
|
||||
|
||||
export function defineTool<
|
||||
|
||||
@@ -1,4 +1,5 @@
|
||||
import { mkdir, mkdtemp, rename, rm } from 'node:fs/promises'
|
||||
import { tmpdir } from 'node:os'
|
||||
import { join } from 'node:path'
|
||||
import { z } from 'zod'
|
||||
import { defineTool, resolveWorkingPath } from './framework'
|
||||
@@ -121,10 +122,9 @@ export const download_file = defineTool({
|
||||
}),
|
||||
handler: async (args, ctx, response) => {
|
||||
const resolvedDir = resolveWorkingPath(ctx, args.path, args.cwd)
|
||||
await mkdir(ctx.directories.workingDir, { recursive: true })
|
||||
const tempDir = await mkdtemp(
|
||||
join(ctx.directories.workingDir, 'browseros-dl-'),
|
||||
)
|
||||
const baseDir = ctx.directories.workingDir ?? tmpdir()
|
||||
await mkdir(baseDir, { recursive: true })
|
||||
const tempDir = await mkdtemp(join(baseDir, 'browseros-dl-'))
|
||||
|
||||
try {
|
||||
const { filePath, suggestedFilename } =
|
||||
|
||||
@@ -0,0 +1,299 @@
|
||||
/**
|
||||
* @license
|
||||
* Copyright 2025 BrowserOS
|
||||
*
|
||||
* Message Validation — Test Suite
|
||||
*
|
||||
* Tests for sanitizeMessagesForToolset, which strips tool parts from
|
||||
* carried-over messages when a session is rebuilt with a different toolset
|
||||
* (e.g., workspace removed or MCP server disconnected mid-conversation).
|
||||
*
|
||||
* Without this sanitization, the AI SDK throws a validation error because
|
||||
* it finds tool parts in the message history that have no matching schema.
|
||||
*/
|
||||
|
||||
import { describe, expect, it } from 'bun:test'
|
||||
import type { UIMessage } from 'ai'
|
||||
import {
|
||||
hasMessageContent,
|
||||
sanitizeMessagesForToolset,
|
||||
} from '../../src/agent/message-validation'
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Helpers
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
function makeUserMessage(text: string, id?: string): UIMessage {
|
||||
return {
|
||||
id: id ?? crypto.randomUUID(),
|
||||
role: 'user',
|
||||
parts: [{ type: 'text', text }],
|
||||
}
|
||||
}
|
||||
|
||||
function makeAssistantMessage(
|
||||
parts: UIMessage['parts'],
|
||||
id?: string,
|
||||
): UIMessage {
|
||||
return {
|
||||
id: id ?? crypto.randomUUID(),
|
||||
role: 'assistant',
|
||||
parts,
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// sanitizeMessagesForToolset
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
describe('sanitizeMessagesForToolset', () => {
|
||||
const allTools = new Set([
|
||||
'navigate_page',
|
||||
'click',
|
||||
'take_snapshot',
|
||||
'filesystem_read',
|
||||
'filesystem_write',
|
||||
'memory_search',
|
||||
])
|
||||
|
||||
const noFilesystemTools = new Set([
|
||||
'navigate_page',
|
||||
'click',
|
||||
'take_snapshot',
|
||||
'memory_search',
|
||||
])
|
||||
|
||||
it('preserves messages with no tool parts', () => {
|
||||
const messages: UIMessage[] = [
|
||||
makeUserMessage('Hello'),
|
||||
makeAssistantMessage([{ type: 'text', text: 'Hi there!' }]),
|
||||
]
|
||||
|
||||
const result = sanitizeMessagesForToolset(messages, noFilesystemTools)
|
||||
expect(result).toHaveLength(2)
|
||||
expect(result[0].parts).toHaveLength(1)
|
||||
expect(result[1].parts).toHaveLength(1)
|
||||
})
|
||||
|
||||
it('preserves tool parts when tool is in the toolset', () => {
|
||||
const messages: UIMessage[] = [
|
||||
makeAssistantMessage([
|
||||
{ type: 'text', text: 'Taking a snapshot...' },
|
||||
{
|
||||
type: 'tool-take_snapshot',
|
||||
toolCallId: 'call-1',
|
||||
toolName: 'take_snapshot',
|
||||
state: 'result',
|
||||
input: { page: 1 },
|
||||
output: { content: 'snapshot data' },
|
||||
} as unknown as UIMessage['parts'][number],
|
||||
]),
|
||||
]
|
||||
|
||||
const result = sanitizeMessagesForToolset(messages, allTools)
|
||||
expect(result).toHaveLength(1)
|
||||
expect(result[0].parts).toHaveLength(2)
|
||||
})
|
||||
|
||||
it('strips tool parts when tool is NOT in the toolset', () => {
|
||||
const messages: UIMessage[] = [
|
||||
makeAssistantMessage([
|
||||
{ type: 'text', text: 'Reading file...' },
|
||||
{
|
||||
type: 'tool-filesystem_read',
|
||||
toolCallId: 'call-1',
|
||||
toolName: 'filesystem_read',
|
||||
state: 'result',
|
||||
input: { path: '/tmp/test.txt' },
|
||||
output: { content: 'file data' },
|
||||
} as unknown as UIMessage['parts'][number],
|
||||
]),
|
||||
]
|
||||
|
||||
const result = sanitizeMessagesForToolset(messages, noFilesystemTools)
|
||||
expect(result).toHaveLength(1)
|
||||
// Only the text part should remain
|
||||
expect(result[0].parts).toHaveLength(1)
|
||||
expect(result[0].parts[0].type).toBe('text')
|
||||
})
|
||||
|
||||
it('strips multiple removed tool parts from same message', () => {
|
||||
const messages: UIMessage[] = [
|
||||
makeAssistantMessage([
|
||||
{ type: 'text', text: 'Working on files...' },
|
||||
{
|
||||
type: 'tool-filesystem_read',
|
||||
toolCallId: 'call-1',
|
||||
toolName: 'filesystem_read',
|
||||
state: 'result',
|
||||
input: { path: '/tmp/a.txt' },
|
||||
output: {},
|
||||
} as unknown as UIMessage['parts'][number],
|
||||
{
|
||||
type: 'tool-filesystem_write',
|
||||
toolCallId: 'call-2',
|
||||
toolName: 'filesystem_write',
|
||||
state: 'result',
|
||||
input: { path: '/tmp/b.txt', content: 'data' },
|
||||
output: {},
|
||||
} as unknown as UIMessage['parts'][number],
|
||||
]),
|
||||
]
|
||||
|
||||
const result = sanitizeMessagesForToolset(messages, noFilesystemTools)
|
||||
expect(result).toHaveLength(1)
|
||||
expect(result[0].parts).toHaveLength(1)
|
||||
expect(result[0].parts[0].type).toBe('text')
|
||||
})
|
||||
|
||||
it('keeps browser tool parts while removing filesystem tool parts', () => {
|
||||
const messages: UIMessage[] = [
|
||||
makeAssistantMessage([
|
||||
{
|
||||
type: 'tool-take_snapshot',
|
||||
toolCallId: 'call-1',
|
||||
toolName: 'take_snapshot',
|
||||
state: 'result',
|
||||
input: { page: 1 },
|
||||
output: {},
|
||||
} as unknown as UIMessage['parts'][number],
|
||||
{
|
||||
type: 'tool-filesystem_read',
|
||||
toolCallId: 'call-2',
|
||||
toolName: 'filesystem_read',
|
||||
state: 'result',
|
||||
input: { path: '/tmp/test.txt' },
|
||||
output: {},
|
||||
} as unknown as UIMessage['parts'][number],
|
||||
]),
|
||||
]
|
||||
|
||||
const result = sanitizeMessagesForToolset(messages, noFilesystemTools)
|
||||
expect(result).toHaveLength(1)
|
||||
expect(result[0].parts).toHaveLength(1)
|
||||
expect((result[0].parts[0] as { type: string }).type).toBe(
|
||||
'tool-take_snapshot',
|
||||
)
|
||||
})
|
||||
|
||||
it('removes messages that become empty after stripping', () => {
|
||||
const messages: UIMessage[] = [
|
||||
makeUserMessage('Read this file'),
|
||||
makeAssistantMessage([
|
||||
{
|
||||
type: 'tool-filesystem_read',
|
||||
toolCallId: 'call-1',
|
||||
toolName: 'filesystem_read',
|
||||
state: 'result',
|
||||
input: { path: '/tmp/test.txt' },
|
||||
output: {},
|
||||
} as unknown as UIMessage['parts'][number],
|
||||
]),
|
||||
]
|
||||
|
||||
const result = sanitizeMessagesForToolset(messages, noFilesystemTools)
|
||||
// The assistant message had only a tool part — after stripping, it's empty
|
||||
// and should be filtered out by hasMessageContent
|
||||
expect(result).toHaveLength(1)
|
||||
expect(result[0].role).toBe('user')
|
||||
})
|
||||
|
||||
it('preserves non-tool part types (reasoning, step-start, file)', () => {
|
||||
const messages: UIMessage[] = [
|
||||
makeAssistantMessage([
|
||||
{ type: 'text', text: 'Let me think...' },
|
||||
{
|
||||
type: 'reasoning',
|
||||
reasoning: 'Analyzing the request',
|
||||
} as unknown as UIMessage['parts'][number],
|
||||
{
|
||||
type: 'step-start',
|
||||
} as unknown as UIMessage['parts'][number],
|
||||
]),
|
||||
]
|
||||
|
||||
const result = sanitizeMessagesForToolset(messages, noFilesystemTools)
|
||||
expect(result).toHaveLength(1)
|
||||
expect(result[0].parts).toHaveLength(3)
|
||||
})
|
||||
|
||||
it('returns same message references when no filtering needed', () => {
|
||||
const messages: UIMessage[] = [
|
||||
makeUserMessage('Hello'),
|
||||
makeAssistantMessage([{ type: 'text', text: 'Hi!' }]),
|
||||
]
|
||||
|
||||
const result = sanitizeMessagesForToolset(messages, noFilesystemTools)
|
||||
// Messages that don't need filtering should be the same reference
|
||||
expect(result[0]).toBe(messages[0])
|
||||
expect(result[1]).toBe(messages[1])
|
||||
})
|
||||
|
||||
it('handles empty message array', () => {
|
||||
const result = sanitizeMessagesForToolset([], noFilesystemTools)
|
||||
expect(result).toHaveLength(0)
|
||||
})
|
||||
|
||||
it('handles empty toolset (all tools removed)', () => {
|
||||
const messages: UIMessage[] = [
|
||||
makeAssistantMessage([
|
||||
{ type: 'text', text: 'Working...' },
|
||||
{
|
||||
type: 'tool-navigate_page',
|
||||
toolCallId: 'call-1',
|
||||
toolName: 'navigate_page',
|
||||
state: 'result',
|
||||
input: {},
|
||||
output: {},
|
||||
} as unknown as UIMessage['parts'][number],
|
||||
]),
|
||||
]
|
||||
|
||||
const result = sanitizeMessagesForToolset(messages, new Set())
|
||||
expect(result).toHaveLength(1)
|
||||
expect(result[0].parts).toHaveLength(1)
|
||||
expect(result[0].parts[0].type).toBe('text')
|
||||
})
|
||||
})
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// hasMessageContent (existing function, verify edge cases)
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
describe('hasMessageContent', () => {
|
||||
it('rejects messages with empty parts array', () => {
|
||||
const msg: UIMessage = {
|
||||
id: '1',
|
||||
role: 'assistant',
|
||||
parts: [],
|
||||
}
|
||||
expect(hasMessageContent(msg)).toBe(false)
|
||||
})
|
||||
|
||||
it('rejects messages with only whitespace text', () => {
|
||||
const msg: UIMessage = {
|
||||
id: '1',
|
||||
role: 'assistant',
|
||||
parts: [{ type: 'text', text: ' \n ' }],
|
||||
}
|
||||
expect(hasMessageContent(msg)).toBe(false)
|
||||
})
|
||||
|
||||
it('accepts messages with non-text parts', () => {
|
||||
const msg: UIMessage = {
|
||||
id: '1',
|
||||
role: 'assistant',
|
||||
parts: [
|
||||
{
|
||||
type: 'tool-click',
|
||||
toolCallId: 'call-1',
|
||||
toolName: 'click',
|
||||
state: 'result',
|
||||
input: {},
|
||||
output: {},
|
||||
} as unknown as UIMessage['parts'][number],
|
||||
],
|
||||
}
|
||||
expect(hasMessageContent(msg)).toBe(true)
|
||||
})
|
||||
})
|
||||
Reference in New Issue
Block a user