Skip to content

Commit 00563a1

Browse files
jahoomajahooma
committed
Refactor cancel response logic, include user message if main prompt threw an error
1 parent f68502a commit 00563a1

File tree

3 files changed

+404
-50
lines changed

3 files changed

+404
-50
lines changed

‎packages/agent-runtime/src/util/messages.ts‎

Lines changed: 0 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -93,29 +93,6 @@ export function buildUserMessageContent(
9393
]
9494
}
9595

96-
export function getCancelledAdditionalMessages(args: {
97-
prompt: string | undefined
98-
params: Record<string, any> | undefined
99-
content?: Array<TextPart | ImagePart>
100-
pendingAgentResponse: string
101-
systemMessage: string
102-
}): Message[] {
103-
const { prompt, params, content, pendingAgentResponse, systemMessage } = args
104-
105-
const messages: Message[] = [
106-
{
107-
role: 'user',
108-
content: buildUserMessageContent(prompt, params, content),
109-
tags: ['USER_PROMPT'],
110-
},
111-
userMessage(
112-
`<previous_assistant_message>${pendingAgentResponse}</previous_assistant_message>\n\n${withSystemTags(systemMessage)}`,
113-
),
114-
]
115-
116-
return messages
117-
}
118-
11996
export function parseUserMessage(str: string): string | undefined {
12097
const match = str.match(/<user_message>(.*?)<\/user_message>/s)
12198
return match ? match[1] : undefined

‎sdk/src/__tests__/run-cancellation.test.ts‎

Lines changed: 359 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -22,6 +22,365 @@ describe('Run Cancellation Handling', () => {
2222
mock.restore()
2323
})
2424

25+
it('does not duplicate user message when server responds with session state', async () => {
26+
spyOn(databaseModule, 'getUserInfoFromApiKey').mockResolvedValue({
27+
id: 'user-123',
28+
email: 'test@example.com',
29+
discord_id: null,
30+
referral_code: null,
31+
stripe_customer_id: null,
32+
banned: false,
33+
})
34+
spyOn(databaseModule, 'fetchAgentFromDatabase').mockResolvedValue(null)
35+
spyOn(databaseModule, 'startAgentRun').mockResolvedValue('run-1')
36+
spyOn(databaseModule, 'finishAgentRun').mockResolvedValue(undefined)
37+
spyOn(databaseModule, 'addAgentStep').mockResolvedValue('step-1')
38+
39+
// Server session state already includes the user's message (as the server would normally do)
40+
const serverSessionState = getInitialSessionState(getStubProjectFileContext())
41+
serverSessionState.mainAgentState.messageHistory.push(
42+
userMessage('Please fix the bug'), // Server added this
43+
assistantMessage('I will help you with that.'),
44+
)
45+
46+
spyOn(mainPromptModule, 'callMainPrompt').mockImplementation(
47+
async (params: Parameters<typeof mainPromptModule.callMainPrompt>[0]) => {
48+
const { sendAction, promptId } = params
49+
50+
await sendAction({
51+
action: {
52+
type: 'prompt-response',
53+
promptId,
54+
sessionState: serverSessionState,
55+
output: {
56+
type: 'lastMessage',
57+
value: [],
58+
},
59+
},
60+
})
61+
62+
return {
63+
sessionState: serverSessionState,
64+
output: {
65+
type: 'lastMessage' as const,
66+
value: [],
67+
},
68+
}
69+
},
70+
)
71+
72+
const client = new CodebuffClient({
73+
apiKey: 'test-key',
74+
})
75+
76+
const result = await client.run({
77+
agent: 'base2',
78+
prompt: 'Please fix the bug',
79+
})
80+
81+
// The user's message should NOT be duplicated
82+
const messageHistory = result.sessionState!.mainAgentState.messageHistory
83+
84+
const userMessages = messageHistory.filter((m) => m.role === 'user')
85+
86+
// Should have exactly 1 user message, not 2
87+
expect(userMessages.length).toBe(1)
88+
89+
// Total messages should be 2 (user + assistant), not 3
90+
expect(messageHistory.length).toBe(2)
91+
})
92+
93+
it('does not duplicate user message when cancelled and server already processed the prompt', async () => {
94+
spyOn(databaseModule, 'getUserInfoFromApiKey').mockResolvedValue({
95+
id: 'user-123',
96+
email: 'test@example.com',
97+
discord_id: null,
98+
referral_code: null,
99+
stripe_customer_id: null,
100+
banned: false,
101+
})
102+
spyOn(databaseModule, 'fetchAgentFromDatabase').mockResolvedValue(null)
103+
spyOn(databaseModule, 'startAgentRun').mockResolvedValue('run-1')
104+
spyOn(databaseModule, 'finishAgentRun').mockResolvedValue(undefined)
105+
spyOn(databaseModule, 'addAgentStep').mockResolvedValue('step-1')
106+
107+
const abortController = new AbortController()
108+
109+
// Server session state already includes the user's message (server processed it)
110+
const serverSessionState = getInitialSessionState(getStubProjectFileContext())
111+
serverSessionState.mainAgentState.messageHistory.push(
112+
userMessage('Please fix the bug'), // Server added the user's message
113+
assistantMessage('I will help you with that.'),
114+
)
115+
116+
spyOn(mainPromptModule, 'callMainPrompt').mockImplementation(
117+
async (params: Parameters<typeof mainPromptModule.callMainPrompt>[0]) => {
118+
const { sendAction, promptId } = params
119+
120+
// Stream some content
121+
await sendAction({
122+
action: {
123+
type: 'response-chunk',
124+
userInputId: promptId,
125+
chunk: 'Working on it...',
126+
},
127+
})
128+
129+
// User cancels
130+
abortController.abort()
131+
132+
// Server still responds with its session state
133+
await sendAction({
134+
action: {
135+
type: 'prompt-response',
136+
promptId,
137+
sessionState: serverSessionState,
138+
output: {
139+
type: 'lastMessage',
140+
value: [],
141+
},
142+
},
143+
})
144+
145+
return {
146+
sessionState: serverSessionState,
147+
output: {
148+
type: 'lastMessage' as const,
149+
value: [],
150+
},
151+
}
152+
},
153+
)
154+
155+
const client = new CodebuffClient({
156+
apiKey: 'test-key',
157+
})
158+
159+
const result = await client.run({
160+
agent: 'base2',
161+
prompt: 'Please fix the bug',
162+
signal: abortController.signal,
163+
})
164+
165+
// The user's message should NOT be duplicated
166+
const messageHistory = result.sessionState!.mainAgentState.messageHistory
167+
168+
// Count user messages (excluding system interruption messages)
169+
const userPromptMessages = messageHistory.filter(
170+
(m) => m.role === 'user' &&
171+
m.content.some((c: any) => c.type === 'text' && c.text.includes('fix the bug'))
172+
)
173+
174+
// Should have exactly 1 user message with the prompt, not 2
175+
expect(userPromptMessages.length).toBe(1)
176+
177+
// Total messages should be: 1 user + 1 assistant (original) + 1 partial assistant (streamed) + 1 interruption = 4
178+
// NOT: 2 users + 1 assistant + 1 partial assistant + 1 interruption = 5
179+
expect(messageHistory.length).toBe(4)
180+
})
181+
182+
it('preserves user message when callMainPrompt throws an error', async () => {
183+
spyOn(databaseModule, 'getUserInfoFromApiKey').mockResolvedValue({
184+
id: 'user-123',
185+
email: 'test@example.com',
186+
discord_id: null,
187+
referral_code: null,
188+
stripe_customer_id: null,
189+
banned: false,
190+
})
191+
spyOn(databaseModule, 'fetchAgentFromDatabase').mockResolvedValue(null)
192+
spyOn(databaseModule, 'startAgentRun').mockResolvedValue('run-1')
193+
spyOn(databaseModule, 'finishAgentRun').mockResolvedValue(undefined)
194+
spyOn(databaseModule, 'addAgentStep').mockResolvedValue('step-1')
195+
196+
// Simulate callMainPrompt throwing an error (network failure, server error, etc.)
197+
spyOn(mainPromptModule, 'callMainPrompt').mockRejectedValue(
198+
new Error('Network connection failed'),
199+
)
200+
201+
const client = new CodebuffClient({
202+
apiKey: 'test-key',
203+
})
204+
205+
const result = await client.run({
206+
agent: 'base2',
207+
prompt: 'Please fix the bug in my code',
208+
})
209+
210+
// Should return an error output
211+
expect(result.output.type).toBe('error')
212+
expect((result.output as { type: 'error'; message: string }).message).toBe('Network connection failed')
213+
214+
// The user's message should be preserved in the session state
215+
expect(result.sessionState).toBeDefined()
216+
const messageHistory = result.sessionState!.mainAgentState.messageHistory
217+
218+
// Should have: user message + interruption message
219+
expect(messageHistory.length).toBeGreaterThanOrEqual(2)
220+
221+
// Find the user's original prompt message (should have USER_PROMPT tag)
222+
const userPromptMessage = messageHistory.find(
223+
(m) => m.role === 'user' && m.tags?.includes('USER_PROMPT'),
224+
)
225+
expect(userPromptMessage).toBeDefined()
226+
227+
// Verify the message content contains the original prompt
228+
const textContent = userPromptMessage!.content.find((c: any) => c.type === 'text') as { type: 'text'; text: string } | undefined
229+
expect(textContent).toBeDefined()
230+
expect(textContent!.text).toContain('Please fix the bug in my code')
231+
})
232+
233+
it('does not add empty assistant message when no streaming content', async () => {
234+
spyOn(databaseModule, 'getUserInfoFromApiKey').mockResolvedValue({
235+
id: 'user-123',
236+
email: 'test@example.com',
237+
discord_id: null,
238+
referral_code: null,
239+
stripe_customer_id: null,
240+
banned: false,
241+
})
242+
spyOn(databaseModule, 'fetchAgentFromDatabase').mockResolvedValue(null)
243+
spyOn(databaseModule, 'startAgentRun').mockResolvedValue('run-1')
244+
spyOn(databaseModule, 'finishAgentRun').mockResolvedValue(undefined)
245+
spyOn(databaseModule, 'addAgentStep').mockResolvedValue('step-1')
246+
247+
const abortController = new AbortController()
248+
const serverSessionState = getInitialSessionState(getStubProjectFileContext())
249+
serverSessionState.mainAgentState.messageHistory.push(
250+
userMessage('User prompt'),
251+
)
252+
const originalHistoryLength = serverSessionState.mainAgentState.messageHistory.length
253+
254+
spyOn(mainPromptModule, 'callMainPrompt').mockImplementation(
255+
async (params: Parameters<typeof mainPromptModule.callMainPrompt>[0]) => {
256+
const { sendAction, promptId } = params
257+
258+
// Abort immediately WITHOUT any streaming chunks
259+
abortController.abort()
260+
261+
await sendAction({
262+
action: {
263+
type: 'prompt-response',
264+
promptId,
265+
sessionState: serverSessionState,
266+
output: {
267+
type: 'lastMessage',
268+
value: [],
269+
},
270+
},
271+
})
272+
273+
return {
274+
sessionState: serverSessionState,
275+
output: {
276+
type: 'lastMessage' as const,
277+
value: [],
278+
},
279+
}
280+
},
281+
)
282+
283+
const client = new CodebuffClient({
284+
apiKey: 'test-key',
285+
})
286+
287+
const result = await client.run({
288+
agent: 'base2',
289+
prompt: 'test prompt',
290+
signal: abortController.signal,
291+
})
292+
293+
const messageHistory = result.sessionState!.mainAgentState.messageHistory
294+
295+
// Should only have: original history + 1 interruption message (NO empty assistant message)
296+
expect(messageHistory.length).toBe(originalHistoryLength + 1)
297+
298+
// The last message should be the interruption (user role), not an empty assistant message
299+
const lastMessage = messageHistory[messageHistory.length - 1]
300+
expect(lastMessage.role).toBe('user')
301+
expect((lastMessage.content[0] as { type: 'text'; text: string }).text).toContain('User interrupted')
302+
303+
// Verify there's no empty assistant message before the interruption
304+
const secondToLastMessage = messageHistory[messageHistory.length - 2]
305+
// This should be the original 'User prompt' message, not an empty assistant
306+
expect(secondToLastMessage.role).toBe('user')
307+
})
308+
309+
it('preserves user message with USER_PROMPT tag when error thrown during callMainPrompt', async () => {
310+
spyOn(databaseModule, 'getUserInfoFromApiKey').mockResolvedValue({
311+
id: 'user-123',
312+
email: 'test@example.com',
313+
discord_id: null,
314+
referral_code: null,
315+
stripe_customer_id: null,
316+
banned: false,
317+
})
318+
spyOn(databaseModule, 'fetchAgentFromDatabase').mockResolvedValue(null)
319+
spyOn(databaseModule, 'startAgentRun').mockResolvedValue('run-1')
320+
spyOn(databaseModule, 'finishAgentRun').mockResolvedValue(undefined)
321+
spyOn(databaseModule, 'addAgentStep').mockResolvedValue('step-1')
322+
323+
let streamedContent = ''
324+
spyOn(mainPromptModule, 'callMainPrompt').mockImplementation(
325+
async (params: Parameters<typeof mainPromptModule.callMainPrompt>[0]) => {
326+
const { sendAction, promptId } = params
327+
328+
// Simulate some partial streaming before error
329+
await sendAction({
330+
action: {
331+
type: 'response-chunk',
332+
userInputId: promptId,
333+
chunk: 'Starting to analyze...',
334+
},
335+
})
336+
337+
// Then throw an error (simulating connection drop)
338+
throw new Error('Connection reset by peer')
339+
},
340+
)
341+
342+
const client = new CodebuffClient({
343+
apiKey: 'test-key',
344+
})
345+
346+
const result = await client.run({
347+
agent: 'base2',
348+
prompt: 'Implement the feature',
349+
handleStreamChunk: (chunk) => {
350+
if (typeof chunk === 'string') {
351+
streamedContent += chunk
352+
}
353+
},
354+
})
355+
356+
// Verify we received some streamed content before the error
357+
expect(streamedContent).toBe('Starting to analyze...')
358+
359+
// Should have error output
360+
expect(result.output.type).toBe('error')
361+
362+
// Session state should be preserved
363+
expect(result.sessionState).toBeDefined()
364+
const messageHistory = result.sessionState!.mainAgentState.messageHistory
365+
366+
// Should have: user message (with USER_PROMPT tag) + partial assistant + interruption
367+
expect(messageHistory.length).toBe(3)
368+
369+
// First message should be the user's prompt with the tag
370+
const firstMessage = messageHistory[0]
371+
expect(firstMessage.role).toBe('user')
372+
expect(firstMessage.tags).toContain('USER_PROMPT')
373+
374+
// Second message should be the partial assistant response
375+
const secondMessage = messageHistory[1]
376+
expect(secondMessage.role).toBe('assistant')
377+
expect((secondMessage.content[0] as { type: 'text'; text: string }).text).toBe('Starting to analyze...')
378+
379+
// Third message should be the interruption/error message
380+
const thirdMessage = messageHistory[2]
381+
expect(thirdMessage.role).toBe('user')
382+
})
383+
25384
it('preserves session state from server when aborted and appends interruption message', async () => {
26385
spyOn(databaseModule, 'getUserInfoFromApiKey').mockResolvedValue({
27386
id: 'user-123',

0 commit comments

Comments
 (0)