feat: Add tests for proactive mode

This commit is contained in:
Vanalite 2025-10-28 12:19:00 +07:00
parent e9f469b623
commit a14872666a
4 changed files with 669 additions and 2 deletions

View File

@ -0,0 +1,124 @@
import { describe, it, expect, vi } from 'vitest'
import { render, screen } from '@testing-library/react'
import Capabilities from '../Capabilities'
// Mock Tooltip components
vi.mock('@/components/ui/tooltip', () => ({
Tooltip: ({ children }: { children: React.ReactNode }) => <div>{children}</div>,
TooltipContent: ({ children }: { children: React.ReactNode }) => <div>{children}</div>,
TooltipProvider: ({ children }: { children: React.ReactNode }) => <div>{children}</div>,
TooltipTrigger: ({ children }: { children: React.ReactNode }) => <div>{children}</div>,
}))
// Mock Tabler icons
vi.mock('@tabler/icons-react', () => ({
IconEye: () => <div data-testid="icon-eye">Eye Icon</div>,
IconTool: () => <div data-testid="icon-tool">Tool Icon</div>,
IconSparkles: () => <div data-testid="icon-sparkles">Sparkles Icon</div>,
IconAtom: () => <div data-testid="icon-atom">Atom Icon</div>,
IconWorld: () => <div data-testid="icon-world">World Icon</div>,
IconCodeCircle2: () => <div data-testid="icon-code">Code Icon</div>,
}))
describe('Capabilities', () => {
it('should render vision capability with eye icon', () => {
render(<Capabilities capabilities={['vision']} />)
const eyeIcon = screen.getByTestId('icon-eye')
expect(eyeIcon).toBeInTheDocument()
})
it('should render tools capability with tool icon', () => {
render(<Capabilities capabilities={['tools']} />)
const toolIcon = screen.getByTestId('icon-tool')
expect(toolIcon).toBeInTheDocument()
})
it('should render proactive capability with sparkles icon', () => {
render(<Capabilities capabilities={['proactive']} />)
const sparklesIcon = screen.getByTestId('icon-sparkles')
expect(sparklesIcon).toBeInTheDocument()
})
it('should render reasoning capability with atom icon', () => {
render(<Capabilities capabilities={['reasoning']} />)
const atomIcon = screen.getByTestId('icon-atom')
expect(atomIcon).toBeInTheDocument()
})
it('should render web_search capability with world icon', () => {
render(<Capabilities capabilities={['web_search']} />)
const worldIcon = screen.getByTestId('icon-world')
expect(worldIcon).toBeInTheDocument()
})
it('should render embeddings capability with code icon', () => {
render(<Capabilities capabilities={['embeddings']} />)
const codeIcon = screen.getByTestId('icon-code')
expect(codeIcon).toBeInTheDocument()
})
it('should render multiple capabilities', () => {
render(<Capabilities capabilities={['tools', 'vision', 'proactive']} />)
expect(screen.getByTestId('icon-tool')).toBeInTheDocument()
expect(screen.getByTestId('icon-eye')).toBeInTheDocument()
expect(screen.getByTestId('icon-sparkles')).toBeInTheDocument()
})
it('should render all capabilities in correct order', () => {
render(<Capabilities capabilities={['tools', 'vision', 'proactive', 'reasoning', 'web_search', 'embeddings']} />)
expect(screen.getByTestId('icon-tool')).toBeInTheDocument()
expect(screen.getByTestId('icon-eye')).toBeInTheDocument()
expect(screen.getByTestId('icon-sparkles')).toBeInTheDocument()
expect(screen.getByTestId('icon-atom')).toBeInTheDocument()
expect(screen.getByTestId('icon-world')).toBeInTheDocument()
expect(screen.getByTestId('icon-code')).toBeInTheDocument()
})
it('should handle empty capabilities array', () => {
const { container } = render(<Capabilities capabilities={[]} />)
expect(container.querySelector('[data-testid^="icon-"]')).not.toBeInTheDocument()
})
it('should handle unknown capabilities gracefully', () => {
const { container } = render(<Capabilities capabilities={['unknown_capability']} />)
expect(container).toBeInTheDocument()
})
it('should display proactive tooltip with correct text', () => {
render(<Capabilities capabilities={['proactive']} />)
// The tooltip content should be 'Proactive'
expect(screen.getByTestId('icon-sparkles')).toBeInTheDocument()
})
it('should render proactive icon between tools/vision and reasoning', () => {
const { container } = render(<Capabilities capabilities={['tools', 'vision', 'proactive', 'reasoning']} />)
// All icons should be rendered
expect(screen.getByTestId('icon-tool')).toBeInTheDocument()
expect(screen.getByTestId('icon-eye')).toBeInTheDocument()
expect(screen.getByTestId('icon-sparkles')).toBeInTheDocument()
expect(screen.getByTestId('icon-atom')).toBeInTheDocument()
expect(container.querySelector('[data-testid="icon-sparkles"]')).toBeInTheDocument()
})
it('should apply correct CSS classes to proactive icon', () => {
render(<Capabilities capabilities={['proactive']} />)
const sparklesIcon = screen.getByTestId('icon-sparkles')
expect(sparklesIcon).toBeInTheDocument()
// Icon should have size-3.5 class (same as tools, reasoning, etc.)
expect(sparklesIcon.parentElement).toBeInTheDocument()
})
})

View File

@ -437,4 +437,31 @@ describe('ChatInput', () => {
expect(() => renderWithRouter()).not.toThrow()
})
})
describe('Proactive Mode', () => {
it('should render ChatInput with proactive capable model', async () => {
await act(async () => {
renderWithRouter()
})
expect(screen.getByTestId('chat-input')).toBeInTheDocument()
})
it('should handle proactive capability detection', async () => {
await act(async () => {
renderWithRouter()
})
expect(screen.getByTestId('chat-input')).toBeInTheDocument()
})
it('should work with models that have multiple capabilities', async () => {
await act(async () => {
renderWithRouter()
})
expect(screen.getByTestId('chat-input')).toBeInTheDocument()
})
})
})

View File

@ -170,6 +170,7 @@ vi.mock('@/lib/completion', () => ({
sendCompletion: vi.fn(),
postMessageProcessing: vi.fn(),
isCompletionResponse: vi.fn(),
captureProactiveScreenshots: vi.fn(() => Promise.resolve([])),
}))
vi.mock('@/lib/messages', () => ({
@ -225,4 +226,26 @@ describe('useChat', () => {
expect(result.current).toBeDefined()
})
describe('Proactive Mode', () => {
it('should detect proactive mode when model has proactive capability', () => {
const { result } = renderHook(() => useChat())
expect(result.current).toBeDefined()
expect(typeof result.current).toBe('function')
})
it('should handle model with tools, vision, and proactive capabilities', () => {
const { result } = renderHook(() => useChat())
expect(result.current).toBeDefined()
})
it('should work with models that have proactive capability', () => {
const { result } = renderHook(() => useChat())
expect(result.current).toBeDefined()
expect(typeof result.current).toBe('function')
})
})
})

View File

@ -1,5 +1,5 @@
import { describe, it, expect, vi, beforeEach } from 'vitest'
import {
import {
newUserThreadContent,
newAssistantThreadContent,
emptyThreadContent,
@ -8,7 +8,8 @@ import {
stopModel,
normalizeTools,
extractToolCall,
postMessageProcessing
postMessageProcessing,
captureProactiveScreenshots
} from '../completion'
// Mock dependencies
@ -72,6 +73,54 @@ vi.mock('../extension', () => ({
ExtensionManager: {},
}))
vi.mock('@/hooks/useServiceHub', () => ({
getServiceHub: vi.fn(() => ({
mcp: vi.fn(() => ({
getTools: vi.fn(() => Promise.resolve([])),
callToolWithCancellation: vi.fn(() => ({
promise: Promise.resolve({
content: [{ type: 'text', text: 'mock result' }],
error: '',
}),
cancel: vi.fn(),
})),
})),
rag: vi.fn(() => ({
getToolNames: vi.fn(() => Promise.resolve([])),
callTool: vi.fn(() => Promise.resolve({
content: [{ type: 'text', text: 'mock rag result' }],
error: '',
})),
})),
})),
}))
vi.mock('@/hooks/useAttachments', () => ({
useAttachments: {
getState: vi.fn(() => ({ enabled: true })),
},
}))
vi.mock('@/hooks/useAppState', () => ({
useAppState: {
getState: vi.fn(() => ({
setCancelToolCall: vi.fn(),
})),
},
}))
vi.mock('@/lib/platform/const', () => ({
PlatformFeatures: {
ATTACHMENTS: true,
},
}))
vi.mock('@/lib/platform/types', () => ({
PlatformFeature: {
ATTACHMENTS: 'ATTACHMENTS',
},
}))
describe('completion.ts', () => {
beforeEach(() => {
vi.clearAllMocks()
@ -187,4 +236,448 @@ describe('completion.ts', () => {
expect(result.length).toBe(0)
})
})
describe('Proactive Mode - Browser MCP Tool Detection', () => {
// We need to access the private function, so we'll test it through postMessageProcessing
it('should detect browser tool names with "browser" prefix', async () => {
const { getServiceHub } = await import('@/hooks/useServiceHub')
const mockGetTools = vi.fn(() => Promise.resolve([]))
const mockMcp = {
getTools: mockGetTools,
callToolWithCancellation: vi.fn(() => ({
promise: Promise.resolve({ content: [{ type: 'text', text: 'result' }], error: '' }),
cancel: vi.fn(),
}))
}
vi.mocked(getServiceHub).mockReturnValue({
mcp: () => mockMcp,
rag: () => ({ getToolNames: () => Promise.resolve([]) })
} as any)
const calls = [{
id: 'call_1',
type: 'function' as const,
function: { name: 'browserbase_navigate', arguments: '{"url": "test.com"}' }
}]
const builder = {
addToolMessage: vi.fn(),
getMessages: vi.fn(() => [])
} as any
const message = { thread_id: 'test-thread', metadata: {} } as any
const abortController = new AbortController()
await postMessageProcessing(
calls,
builder,
message,
abortController,
{},
undefined,
false,
true // isProactiveMode = true
)
// Verify tool was executed
expect(mockMcp.callToolWithCancellation).toHaveBeenCalled()
})
it('should detect browserbase tools', async () => {
const { getServiceHub } = await import('@/hooks/useServiceHub')
const mockCallTool = vi.fn(() => ({
promise: Promise.resolve({ content: [{ type: 'text', text: 'result' }], error: '' }),
cancel: vi.fn(),
}))
vi.mocked(getServiceHub).mockReturnValue({
mcp: () => ({
getTools: () => Promise.resolve([]),
callToolWithCancellation: mockCallTool
}),
rag: () => ({ getToolNames: () => Promise.resolve([]) })
} as any)
const calls = [{
id: 'call_1',
type: 'function' as const,
function: { name: 'browserbase_screenshot', arguments: '{}' }
}]
const builder = {
addToolMessage: vi.fn(),
getMessages: vi.fn(() => [])
} as any
const message = { thread_id: 'test-thread', metadata: {} } as any
const abortController = new AbortController()
await postMessageProcessing(calls, builder, message, abortController, {}, undefined, false, true)
expect(mockCallTool).toHaveBeenCalled()
})
it('should detect multi_browserbase tools', async () => {
const { getServiceHub } = await import('@/hooks/useServiceHub')
const mockCallTool = vi.fn(() => ({
promise: Promise.resolve({ content: [{ type: 'text', text: 'result' }], error: '' }),
cancel: vi.fn(),
}))
vi.mocked(getServiceHub).mockReturnValue({
mcp: () => ({
getTools: () => Promise.resolve([]),
callToolWithCancellation: mockCallTool
}),
rag: () => ({ getToolNames: () => Promise.resolve([]) })
} as any)
const calls = [{
id: 'call_1',
type: 'function' as const,
function: { name: 'multi_browserbase_stagehand_navigate', arguments: '{}' }
}]
const builder = {
addToolMessage: vi.fn(),
getMessages: vi.fn(() => [])
} as any
const message = { thread_id: 'test-thread', metadata: {} } as any
const abortController = new AbortController()
await postMessageProcessing(calls, builder, message, abortController, {}, undefined, false, true)
expect(mockCallTool).toHaveBeenCalled()
})
it('should not treat non-browser tools as browser tools', async () => {
const { getServiceHub } = await import('@/hooks/useServiceHub')
const mockGetTools = vi.fn(() => Promise.resolve([]))
vi.mocked(getServiceHub).mockReturnValue({
mcp: () => ({
getTools: mockGetTools,
callToolWithCancellation: vi.fn(() => ({
promise: Promise.resolve({ content: [{ type: 'text', text: 'result' }], error: '' }),
cancel: vi.fn(),
}))
}),
rag: () => ({ getToolNames: () => Promise.resolve([]) })
} as any)
const calls = [{
id: 'call_1',
type: 'function' as const,
function: { name: 'fetch_url', arguments: '{"url": "test.com"}' }
}]
const builder = {
addToolMessage: vi.fn(),
getMessages: vi.fn(() => [])
} as any
const message = { thread_id: 'test-thread', metadata: {} } as any
const abortController = new AbortController()
await postMessageProcessing(calls, builder, message, abortController, {}, undefined, false, true)
// Proactive screenshots should not be called for non-browser tools
expect(mockGetTools).not.toHaveBeenCalled()
})
})
describe('Proactive Mode - Screenshot Capture', () => {
it('should capture screenshot and snapshot when available', async () => {
const { getServiceHub } = await import('@/hooks/useServiceHub')
const mockScreenshotResult = {
content: [{ type: 'image', data: 'base64screenshot', mimeType: 'image/png' }],
error: '',
}
const mockSnapshotResult = {
content: [{ type: 'text', text: 'snapshot html' }],
error: '',
}
const mockGetTools = vi.fn(() => Promise.resolve([
{ name: 'browserbase_screenshot', inputSchema: {} },
{ name: 'browserbase_snapshot', inputSchema: {} }
]))
const mockCallTool = vi.fn()
.mockReturnValueOnce({
promise: Promise.resolve(mockScreenshotResult),
cancel: vi.fn(),
})
.mockReturnValueOnce({
promise: Promise.resolve(mockSnapshotResult),
cancel: vi.fn(),
})
vi.mocked(getServiceHub).mockReturnValue({
mcp: () => ({
getTools: mockGetTools,
callToolWithCancellation: mockCallTool
})
} as any)
const abortController = new AbortController()
const results = await captureProactiveScreenshots(abortController)
expect(results).toHaveLength(2)
expect(results[0]).toEqual(mockScreenshotResult)
expect(results[1]).toEqual(mockSnapshotResult)
expect(mockCallTool).toHaveBeenCalledTimes(2)
})
it('should handle missing screenshot tool gracefully', async () => {
const { getServiceHub } = await import('@/hooks/useServiceHub')
const mockGetTools = vi.fn(() => Promise.resolve([
{ name: 'some_other_tool', inputSchema: {} }
]))
vi.mocked(getServiceHub).mockReturnValue({
mcp: () => ({
getTools: mockGetTools,
callToolWithCancellation: vi.fn()
})
} as any)
const abortController = new AbortController()
const results = await captureProactiveScreenshots(abortController)
expect(results).toHaveLength(0)
})
it('should handle screenshot capture errors gracefully', async () => {
const { getServiceHub } = await import('@/hooks/useServiceHub')
const mockGetTools = vi.fn(() => Promise.resolve([
{ name: 'browserbase_screenshot', inputSchema: {} }
]))
const mockCallTool = vi.fn(() => ({
promise: Promise.reject(new Error('Screenshot failed')),
cancel: vi.fn(),
}))
vi.mocked(getServiceHub).mockReturnValue({
mcp: () => ({
getTools: mockGetTools,
callToolWithCancellation: mockCallTool
})
} as any)
const abortController = new AbortController()
const results = await captureProactiveScreenshots(abortController)
// Should return empty array on error, not throw
expect(results).toHaveLength(0)
})
it('should respect abort controller', async () => {
const { getServiceHub } = await import('@/hooks/useServiceHub')
const mockGetTools = vi.fn(() => Promise.resolve([
{ name: 'browserbase_screenshot', inputSchema: {} }
]))
const mockCallTool = vi.fn(() => ({
promise: new Promise((resolve) => setTimeout(() => resolve({
content: [{ type: 'image', data: 'base64', mimeType: 'image/png' }],
error: '',
}), 100)),
cancel: vi.fn(),
}))
vi.mocked(getServiceHub).mockReturnValue({
mcp: () => ({
getTools: mockGetTools,
callToolWithCancellation: mockCallTool
})
} as any)
const abortController = new AbortController()
abortController.abort()
const results = await captureProactiveScreenshots(abortController)
// Should not attempt to capture if already aborted
expect(results).toHaveLength(0)
})
})
describe('Proactive Mode - Screenshot Filtering', () => {
it('should filter out old image_url content from tool messages', () => {
const builder = {
messages: [
{ role: 'user', content: 'Hello' },
{
role: 'tool',
content: [
{ type: 'text', text: 'Tool result' },
{ type: 'image_url', image_url: { url: 'data:image/png;base64,old' } }
],
tool_call_id: 'old_call'
},
{ role: 'assistant', content: 'Response' },
]
}
expect(builder.messages).toHaveLength(3)
})
})
describe('Proactive Mode - Integration', () => {
it('should trigger proactive screenshots after browser tool execution', async () => {
const { getServiceHub } = await import('@/hooks/useServiceHub')
const mockScreenshotResult = {
content: [{ type: 'image', data: 'proactive_screenshot', mimeType: 'image/png' }],
error: '',
}
const mockGetTools = vi.fn(() => Promise.resolve([
{ name: 'browserbase_screenshot', inputSchema: {} }
]))
let callCount = 0
const mockCallTool = vi.fn(() => {
callCount++
if (callCount === 1) {
// First call: the browser tool itself
return {
promise: Promise.resolve({
content: [{ type: 'text', text: 'navigated to page' }],
error: '',
}),
cancel: vi.fn(),
}
} else {
// Second call: proactive screenshot
return {
promise: Promise.resolve(mockScreenshotResult),
cancel: vi.fn(),
}
}
})
vi.mocked(getServiceHub).mockReturnValue({
mcp: () => ({
getTools: mockGetTools,
callToolWithCancellation: mockCallTool
}),
rag: () => ({ getToolNames: () => Promise.resolve([]) })
} as any)
const calls = [{
id: 'call_1',
type: 'function' as const,
function: { name: 'browserbase_navigate', arguments: '{"url": "test.com"}' }
}]
const builder = {
addToolMessage: vi.fn(),
getMessages: vi.fn(() => [])
} as any
const message = { thread_id: 'test-thread', metadata: {} } as any
const abortController = new AbortController()
await postMessageProcessing(
calls,
builder,
message,
abortController,
{},
undefined,
false,
true
)
// Should have called: 1) browser tool, 2) getTools, 3) proactive screenshot
expect(mockCallTool).toHaveBeenCalledTimes(2)
expect(mockGetTools).toHaveBeenCalled()
expect(builder.addToolMessage).toHaveBeenCalledTimes(2)
})
it('should not trigger proactive screenshots when mode is disabled', async () => {
const { getServiceHub } = await import('@/hooks/useServiceHub')
const mockGetTools = vi.fn(() => Promise.resolve([
{ name: 'browserbase_screenshot', inputSchema: {} }
]))
const mockCallTool = vi.fn(() => ({
promise: Promise.resolve({
content: [{ type: 'text', text: 'navigated' }],
error: '',
}),
cancel: vi.fn(),
}))
vi.mocked(getServiceHub).mockReturnValue({
mcp: () => ({
getTools: mockGetTools,
callToolWithCancellation: mockCallTool
}),
rag: () => ({ getToolNames: () => Promise.resolve([]) })
} as any)
const calls = [{
id: 'call_1',
type: 'function' as const,
function: { name: 'browserbase_navigate', arguments: '{}' }
}]
const builder = {
addToolMessage: vi.fn(),
getMessages: vi.fn(() => [])
} as any
const message = { thread_id: 'test-thread', metadata: {} } as any
const abortController = new AbortController()
await postMessageProcessing(
calls,
builder,
message,
abortController,
{},
undefined,
false,
false
)
expect(mockCallTool).toHaveBeenCalledTimes(1)
expect(mockGetTools).not.toHaveBeenCalled()
})
it('should not trigger proactive screenshots for non-browser tools', async () => {
const { getServiceHub } = await import('@/hooks/useServiceHub')
const mockGetTools = vi.fn(() => Promise.resolve([]))
const mockCallTool = vi.fn(() => ({
promise: Promise.resolve({
content: [{ type: 'text', text: 'fetched data' }],
error: '',
}),
cancel: vi.fn(),
}))
vi.mocked(getServiceHub).mockReturnValue({
mcp: () => ({
getTools: mockGetTools,
callToolWithCancellation: mockCallTool
}),
rag: () => ({ getToolNames: () => Promise.resolve([]) })
} as any)
const calls = [{
id: 'call_1',
type: 'function' as const,
function: { name: 'fetch_url', arguments: '{"url": "test.com"}' }
}]
const builder = {
addToolMessage: vi.fn(),
getMessages: vi.fn(() => [])
} as any
const message = { thread_id: 'test-thread', metadata: {} } as any
const abortController = new AbortController()
await postMessageProcessing(
calls,
builder,
message,
abortController,
{},
undefined,
false,
true
)
expect(mockCallTool).toHaveBeenCalledTimes(1)
expect(mockGetTools).not.toHaveBeenCalled()
})
})
})