Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
91 changes: 3 additions & 88 deletions src/server.ts
Original file line number Diff line number Diff line change
Expand Up @@ -158,7 +158,7 @@ export const toolDefinitions: ToolDef[] = [
{
name: 'comet_ask',
description:
'Send a prompt to Perplexity Comet and poll until the agent responds or times out. Supports newChat to start fresh.',
'Send a prompt to Perplexity Comet and return immediately. Supports newChat to start fresh. Use comet_poll or comet_wait to get the response.',
inputSchema: buildInputSchema(askShape),
},
{
Expand Down Expand Up @@ -355,7 +355,7 @@ export async function startServer(): Promise<void> {
// 2. comet_ask
server.tool(
'comet_ask',
'Send a prompt to Perplexity Comet and poll until the agent responds or times out. Supports newChat to start fresh.',
'Send a prompt to Perplexity Comet and return immediately. Supports newChat to start fresh. Use comet_poll or comet_wait to get the response.',
askShape,
async ({ prompt, newChat, timeout }) => {
try {
Expand Down Expand Up @@ -398,92 +398,7 @@ export async function startServer(): Promise<void> {
const submitResult = await client.safeEvaluate(buildSubmitPromptScript())
logger.debug('Submit result:', extractValue(submitResult))

// POLLING LOOP
const startTime = Date.now()
let sawNewResponse = false
let timedOut = false
const collectedSteps: string[] = []
let lastResponse = ''
let stallCount = 0
const MAX_STALL_POLLS = 10

while (!timedOut && Date.now() - startTime < effectiveTimeout) {
await sleep(config.pollInterval)

const statusRaw = await client.safeEvaluate(buildGetAgentStatusScript(activeSelectors))
const status = parseAgentStatus(extractValue(statusRaw))

// Collect new steps
for (const step of status.steps) {
if (!collectedSteps.includes(step)) {
collectedSteps.push(step)
}
}

// Check for new response — proseCount is the primary signal
const proseIncreased = (status.proseCount ?? 0) > preSendState.proseCount
// Only consider response "changed" if:
// 1. proseCount increased (new prose element added), OR
// 2. Fresh page had no prose before, and now there's a substantial response
const responseChanged =
proseIncreased || (!preSendState.lastProseText && hasSubstantialResponse(status))

if (responseChanged && status.response) {
// Track response growth for auto-extend
if (status.response.length > lastResponse.length) {
stallCount = 0
} else if (sawNewResponse) {
stallCount++
}
sawNewResponse = true
lastResponse = status.response
}

// Stall detection — if response stopped growing, give up after MAX_STALL_POLLS
if (stallCount >= MAX_STALL_POLLS) break

if ((status.status === 'completed' || status.status === 'idle') && sawNewResponse) {
// Wait for response to stabilize — poll until length stops growing
let settledResponse = lastResponse
for (let settle = 0; settle < 5; settle++) {
await sleep(1000)
const settledRaw = await client.safeEvaluate(
buildGetAgentStatusScript(activeSelectors),
)
const settledStatus = parseAgentStatus(extractValue(settledRaw))
const candidate = settledStatus.response || settledResponse
if (candidate.length <= settledResponse.length) break
settledResponse = candidate
}

const parts: string[] = []
if (settledResponse) parts.push(settledResponse)
if (collectedSteps.length > 0) {
parts.push(`\n\nSteps:\n${collectedSteps.map((s) => ` - ${s}`).join('\n')}`)
}
return textResult(parts.join('') || 'Agent completed with no visible response.')
}

if (
status.status === 'idle' &&
!sawNewResponse &&
status.response &&
!preSendState.lastProseText
) {
return textResult(status.response)
}
}

// Mark timed out to prevent any further polling
timedOut = true

// Timeout
const timeoutParts: string[] = ['Agent is still working. Use comet_poll to check status.']
if (collectedSteps.length > 0) {
timeoutParts.push(`\nSteps so far:\n${collectedSteps.map((s) => ` - ${s}`).join('\n')}`)
}
if (lastResponse) timeoutParts.push(`\nPartial response:\n${lastResponse}`)
return textResult(timeoutParts.join('\n'))
return textResult('Prompt submitted successfully. Use comet_poll to track status or comet_wait to block until completion.')
} catch (err) {
return toMcpError(err)
}
Expand Down
224 changes: 5 additions & 219 deletions tests/integration/tools/core-tools.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -138,241 +138,27 @@ describe('Core tool handlers', () => {
// ---------------------------------------------------------------------------

describe('comet_ask', () => {
it('quick response — returns agent response', async () => {
it('returns immediate submission message without polling', async () => {
let callCount = 0
const responseText =
'The answer is 42, which is the meaning of life according to Douglas Adams'

mocks.safeEvaluate.mockImplementation(async () => {
callCount++
// First call: pre-send state
if (callCount === 1) {
return { result: { value: '{"proseCount":0,"lastProseText":""}' } }
}
// Second call: type prompt
if (callCount === 2) {
return { result: { value: 'typed' } }
}
// Third call: submit
if (callCount === 3) {
return { result: { value: 'submitted' } }
}
// Fourth+ calls: status polling (completed)
return {
result: {
value: JSON.stringify({
status: 'completed',
steps: ['Searching web'],
currentStep: 'Searching web',
response: responseText,
hasStopButton: false,
}),
},
}
})

const handler = getHandler('comet_ask')
const result = await handler({ prompt: 'What is 42?' })

expect(result.content[0].text).toContain(responseText)
})

it('timeout — returns still working message', async () => {
mocks.safeEvaluate.mockResolvedValue({
result: {
value: JSON.stringify({
status: 'working',
steps: [],
currentStep: '',
response: '',
hasStopButton: true,
}),
},
return { result: { value: '{"proseCount":0,"lastProseText":""}' } }
})

const handler = getHandler('comet_ask')
const result = await handler({ prompt: 'test', timeout: 300 })
const result = await handler({ prompt: 'test' })

expect(result.content[0].text).toContain('Agent is still working')
expect(result.content[0].text).toContain('Prompt submitted successfully')
expect(result.content[0].text).toContain('comet_poll')
})

it('error handling — returns MCP error when safeEvaluate throws', async () => {
mocks.safeEvaluate.mockRejectedValue(new Error('Script error'))

const handler = getHandler('comet_ask')
const result = await handler({ prompt: 'test' })

expect(result.isError).toBe(true)
expect(result.content[0].text).toContain('Error')
})

it('sequential queries — returns new response when proseCount increases', async () => {
// Simulates BUG-2: second query should detect new response via proseCount
// even when old response text is still on the page
let callCount = 0
const oldResponse = 'This is the old response from the first query that is still on the page.'
const newResponse = 'This is the new response from the second query with different content.'

mocks.safeEvaluate.mockImplementation(async () => {
callCount++
// First call: pre-send state (old response still on page, proseCount=1)
if (callCount === 1) {
return {
result: { value: JSON.stringify({ proseCount: 1, lastProseText: oldResponse }) },
}
}
// Second call: type prompt
if (callCount === 2) return { result: { value: 'typed' } }
// Third call: submit
if (callCount === 3) return { result: { value: 'submitted' } }
// Fourth+ calls: status polling — proseCount now 2 (new prose added)
return {
result: {
value: JSON.stringify({
status: 'completed',
steps: ['Searching web'],
currentStep: 'Searching web',
response: newResponse,
hasStopButton: false,
proseCount: 2,
}),
},
}
})

const handler = getHandler('comet_ask')
const result = await handler({ prompt: 'What is 3+3?' })

expect(result.content[0].text).toContain(newResponse)
expect(result.content[0].text).not.toContain(oldResponse)
})

it('comet_ask stops polling after timeout — no runaway polling', async () => {
let evalCalls = 0
mocks.safeEvaluate.mockImplementation(async () => {
evalCalls++
if (evalCalls === 1) return { result: { value: '{"proseCount":0,"lastProseText":""}' } }
if (evalCalls === 2) return { result: { value: 'typed' } }
if (evalCalls === 3) return { result: { value: 'submitted' } }
return {
result: {
value: JSON.stringify({
status: 'working',
steps: [],
currentStep: '',
response: '',
hasStopButton: true,
}),
},
}
})

const handler = getHandler('comet_ask')
const result = await handler({ prompt: 'test', timeout: 300 })
expect(result.content[0].text).toContain('Agent is still working')

// Verify no runaway polling after timeout
const callsAfterTimeout = evalCalls
await new Promise((r) => setTimeout(r, 500))
expect(evalCalls).toBe(callsAfterTimeout)
})

it('smart polling — auto-extends when response is growing', async () => {
let callCount = 0
const growingResponses = ['A'.repeat(60), 'A'.repeat(120), 'A'.repeat(200)]
mocks.safeEvaluate.mockImplementation(async () => {
callCount++
if (callCount === 1) return { result: { value: '{"proseCount":0,"lastProseText":""}' } }
if (callCount === 2) return { result: { value: 'typed' } }
if (callCount === 3) return { result: { value: 'submitted' } }
const responseIdx = Math.min(callCount - 4, growingResponses.length - 1)
return {
result: {
value: JSON.stringify({
status: callCount > 6 ? 'completed' : 'working',
steps: [],
currentStep: '',
response: growingResponses[responseIdx],
hasStopButton: callCount <= 6,
proseCount: 1,
}),
},
}
})

const handler = getHandler('comet_ask')
// 300ms timeout would normally be too short, but growing response should keep it alive
const result = await handler({ prompt: 'test', timeout: 300 })
// Should have gotten the full response since it was growing
expect(result.content[0].text).toContain('A'.repeat(200))
})

it('smart polling — gives up after stall', async () => {
let callCount = 0
const stalledResponse = 'B'.repeat(60)
mocks.safeEvaluate.mockImplementation(async () => {
callCount++
if (callCount === 1) return { result: { value: '{"proseCount":0,"lastProseText":""}' } }
if (callCount === 2) return { result: { value: 'typed' } }
if (callCount === 3) return { result: { value: 'submitted' } }
return {
result: {
value: JSON.stringify({
status: 'working',
steps: [],
currentStep: '',
response: stalledResponse,
hasStopButton: true,
proseCount: 1,
}),
},
}
})

const handler = getHandler('comet_ask')
const result = await handler({ prompt: 'test', timeout: 30000 })
// Should time out because response stopped growing (stall detection)
expect(result.content[0].text).toContain('still working')
})

it('ignores old substantial response — no false positive from hasSubstantialResponse', async () => {
// Regression: hasSubstantialResponse was OR'd into responseChanged,
// causing old responses to be treated as new when proseCount didn't increase.
const oldResponse =
'This is an old response from a previous query that is still on the page and is quite long.'
let callCount = 0
mocks.safeEvaluate.mockImplementation(async () => {
callCount++
// pre-send state: old response still on page
if (callCount === 1) {
return {
result: { value: JSON.stringify({ proseCount: 1, lastProseText: oldResponse }) },
}
}
if (callCount === 2) return { result: { value: 'typed' } }
if (callCount === 3) return { result: { value: 'submitted' } }
// Polling: agent hasn't started yet, old response still visible
return {
result: {
value: JSON.stringify({
status: 'working',
steps: [],
currentStep: '',
response: oldResponse,
hasStopButton: true,
proseCount: 1,
}),
},
}
})

const handler = getHandler('comet_ask')
const result = await handler({ prompt: 'New question?', timeout: 1500 })

// Should NOT return the old response as if it were the new answer
// Instead should timeout since no new response detected
expect(result.content[0].text).toContain('still working')
})
})

// ---------------------------------------------------------------------------
Expand Down
Loading