diff --git a/packages/browseros-agent/apps/server/src/agent/ai-sdk-agent.ts b/packages/browseros-agent/apps/server/src/agent/ai-sdk-agent.ts index 1565960c..8d3c078f 100644 --- a/packages/browseros-agent/apps/server/src/agent/ai-sdk-agent.ts +++ b/packages/browseros-agent/apps/server/src/agent/ai-sdk-agent.ts @@ -54,8 +54,14 @@ export class AiSdkAgent { private _messages: UIMessage[], private _mcpClients: Array<{ close(): Promise }>, private conversationId: string, + private _toolNames: Set, ) {} + /** Tool names registered on this agent — used to sanitize messages during session rebuilds. */ + get toolNames(): Set { + return this._toolNames + } + static async create(config: AiSdkAgentConfig): Promise { 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)), ) } diff --git a/packages/browseros-agent/apps/server/src/agent/message-validation.ts b/packages/browseros-agent/apps/server/src/agent/message-validation.ts index 13a0e84b..d16c4fef 100644 --- a/packages/browseros-agent/apps/server/src/agent/message-validation.ts +++ b/packages/browseros-agent/apps/server/src/agent/message-validation.ts @@ -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, +): 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) +} diff --git a/packages/browseros-agent/apps/server/src/agent/session-store.ts b/packages/browseros-agent/apps/server/src/agent/session-store.ts index 53714fb5..ee395068 100644 --- a/packages/browseros-agent/apps/server/src/agent/session-store.ts +++ b/packages/browseros-agent/apps/server/src/agent/session-store.ts @@ -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 { diff --git a/packages/browseros-agent/apps/server/src/agent/tool-adapter.ts b/packages/browseros-agent/apps/server/src/agent/tool-adapter.ts index 25434ecd..f36dbfb1 100644 --- a/packages/browseros-agent/apps/server/src/agent/tool-adapter.ts +++ b/packages/browseros-agent/apps/server/src/agent/tool-adapter.ts @@ -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 = {} diff --git a/packages/browseros-agent/apps/server/src/agent/types.ts b/packages/browseros-agent/apps/server/src/agent/types.ts index aa6300ca..659a26a9 100644 --- a/packages/browseros-agent/apps/server/src/agent/types.ts +++ b/packages/browseros-agent/apps/server/src/agent/types.ts @@ -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. */ diff --git a/packages/browseros-agent/apps/server/src/api/services/chat-service.ts b/packages/browseros-agent/apps/server/src/api/services/chat-service.ts index 36de5aa7..910de504 100644 --- a/packages/browseros-agent/apps/server/src/api/services/chat-service.ts +++ b/packages/browseros-agent/apps/server/src/api/services/chat-service.ts @@ -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 { + 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 { - 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 - } } diff --git a/packages/browseros-agent/apps/server/src/tools/framework.ts b/packages/browseros-agent/apps/server/src/tools/framework.ts index a60760b6..45b5d747 100644 --- a/packages/browseros-agent/apps/server/src/tools/framework.ts +++ b/packages/browseros-agent/apps/server/src/tools/framework.ts @@ -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 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< diff --git a/packages/browseros-agent/apps/server/src/tools/page-actions.ts b/packages/browseros-agent/apps/server/src/tools/page-actions.ts index 6db1fd6c..6ce3a609 100644 --- a/packages/browseros-agent/apps/server/src/tools/page-actions.ts +++ b/packages/browseros-agent/apps/server/src/tools/page-actions.ts @@ -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 } = diff --git a/packages/browseros-agent/apps/server/tests/agent/message-validation.test.ts b/packages/browseros-agent/apps/server/tests/agent/message-validation.test.ts new file mode 100644 index 00000000..125e2aa0 --- /dev/null +++ b/packages/browseros-agent/apps/server/tests/agent/message-validation.test.ts @@ -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) + }) +})