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:
Dani Akash
2026-03-27 18:30:25 +05:30
committed by GitHub
parent aacb47f7ee
commit febaf58f91
9 changed files with 471 additions and 49 deletions

View File

@@ -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)),
)
}

View File

@@ -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)
}

View File

@@ -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 {

View File

@@ -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 = {}

View File

@@ -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. */

View File

@@ -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
}
}

View File

@@ -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<

View File

@@ -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 } =

View File

@@ -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)
})
})