Skip to content
Open
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
12 changes: 9 additions & 3 deletions packages/shared/src/agent/pi-rpc-client.ts
Original file line number Diff line number Diff line change
Expand Up @@ -576,18 +576,23 @@ export class PiRpcClient {
return;
}

const pending = this.findPending(event);
const control = this.findControl(event);
if (control) {
clearTimeout(control.timeout);
this.pendingControls.delete(control.id);
if (event.type === 'error') {
control.reject(createCodeError('PI_RUNTIME_ERROR', String(event.message ?? 'Pi runtime error.')));
if (event.type === 'error' || event.success === false) {
control.reject(createCodeError(
'PI_RUNTIME_ERROR',
String(event.error ?? event.message ?? 'Pi runtime error.')
));
} else {
control.resolve(event.data);
}
return;
}

const pending = this.findPending(event);

if (!pending) return;

pending.trace.push({
Expand Down Expand Up @@ -684,6 +689,7 @@ export class PiRpcClient {
private findPending(event: Record<string, unknown>) {
const id = typeof event.id === 'string' ? event.id : typeof event.requestId === 'string' ? event.requestId : undefined;
if (id && this.pendingPrompts.has(id)) return this.pendingPrompts.get(id);
if (event.type === 'response') return undefined;
if (this.pendingPrompts.size === 1) return Array.from(this.pendingPrompts.values())[0];
return undefined;
}
Expand Down
56 changes: 55 additions & 1 deletion packages/shared/src/agent/pi-runtime-agent-backend.test.ts
Original file line number Diff line number Diff line change
Expand Up @@ -249,11 +249,65 @@ describe('PiRpcClient', () => {
await stream.abort();
await consume;

expect(events).toContain('message_update');
expect(events).toEqual(['message_update']);
expect(events).not.toContain('response');
expect(events).not.toContain('abort');
expect(result).toMatchObject({ aborted: true });
expect((result as { answer: string }).answer).toBe('Partial');
});

it('ignores stale prompt responses while a prompt is active', async () => {
const client = new PiRpcClient({
spawnProcess: createSpawn(() =>
new FakePiProcess((line, proc) => {
if (line.type !== 'prompt') return;

proc.writeEvent({
id: 'stale-prompt-id',
type: 'response',
command: 'prompt',
success: false,
error: 'stale prompt response',
});
proc.writeEvent({
id: line.id,
type: 'response',
command: 'prompt',
success: true,
});
proc.writeEvent({
id: line.id,
type: 'agent_end',
messages: [{ role: 'assistant', content: [{ type: 'text', text: 'completed' }] }],
});
})
),
});

await expect(client.prompt('hello')).resolves.toMatchObject({ answer: 'completed' });
});

it('falls back to the only active prompt for id-less agent events', async () => {
const client = new PiRpcClient({
spawnProcess: createSpawn(() =>
new FakePiProcess((line, proc) => {
if (line.type !== 'prompt') return;

proc.writeEvent({
type: 'message_update',
assistantMessageEvent: { type: 'text_delta', delta: 'id-less ' },
});
proc.writeEvent({
type: 'agent_end',
messages: [{ role: 'assistant', content: [{ type: 'text', text: 'event' }] }],
});
})
),
});

await expect(client.prompt('hello')).resolves.toMatchObject({ answer: 'event' });
});

it('rejects when Pi refuses a prompt before execution', async () => {
const client = new PiRpcClient({
spawnProcess: createSpawn(() =>
Expand Down