Fix Pi extension command turns (#1290)

This commit is contained in:
Mohamed Boudra
2026-06-02 18:09:50 +08:00
committed by GitHub
parent 07ace19b69
commit 5666b3014a
3 changed files with 52 additions and 0 deletions

View File

@@ -97,6 +97,18 @@ class SessionEvents {
.map((event) => event.item);
}
timelineAndCompletionEvents() {
return this.events.flatMap((event) => {
if (event.type === "timeline") {
return [{ type: "timeline" as const, item: event.item }];
}
if (event.type === "turn_completed") {
return [{ type: "turn_completed" as const }];
}
return [];
});
}
nextTurnCompletion(): Promise<Extract<AgentStreamEvent, { type: "turn_completed" }>> {
return this.nextEvent(
(event): event is Extract<AgentStreamEvent, { type: "turn_completed" }> =>
@@ -459,6 +471,28 @@ describe("PiRpcAgentSession", () => {
]);
});
test("surfaces Pi extension command messages and completes when no agent turn starts", async () => {
const { pi, session, events } = await createSession();
const fakeSession = pi.latestSession();
await session.startTurn("/show-status");
fakeSession.emit({
type: "message_end",
message: {
role: "custom",
content: [{ type: "text", text: "Extension command output" }],
},
});
expect(events.timelineAndCompletionEvents()).toEqual([
{
type: "timeline",
item: { type: "assistant_message", text: "Extension command output" },
},
{ type: "turn_completed" },
]);
});
test("adds Pi assistant context to generic provider finish errors", async () => {
const { pi, session, events } = await createSession();

View File

@@ -1499,6 +1499,20 @@ export class PiRpcAgentSession implements AgentSession {
event: Extract<PiAgentSessionEvent, { type: "message_end" }>,
turnId: string | undefined,
): void {
if (event.message.role === "custom") {
const text = getUserMessageText(event.message.content);
if (text) {
this.emit({
type: "timeline",
provider: PI_PROVIDER,
turnId,
item: { type: "assistant_message", text },
});
}
this.completeTurn(turnId, []);
return;
}
if (event.message.role !== "user") {
return;
}

View File

@@ -30,6 +30,10 @@ export type PiAgentMessage =
role: "user";
content: string | Array<PiTextContent | PiImageContent>;
}
| {
role: "custom";
content: string | Array<PiTextContent | PiImageContent>;
}
| {
role: "assistant";
content: PiAssistantContent[];