fix: preserve assistant partial stream during reasoning

This commit is contained in:
Ayaan Zaidi
2026-02-19 10:45:06 +05:30
committed by Vignesh
parent 2cbf15eb66
commit 221d50bc18
4 changed files with 105 additions and 18 deletions

View File

@@ -91,6 +91,7 @@ function createMinimalRun(params?: {
storePath?: string;
typingMode?: TypingMode;
blockStreamingEnabled?: boolean;
runOverrides?: Partial<FollowupRun["run"]>;
}) {
const typing = createMockTypingController();
const opts = params?.opts;
@@ -124,6 +125,7 @@ function createMinimalRun(params?: {
},
timeoutMs: 1_000,
blockReplyBreak: "message_end",
...params?.runOverrides,
},
} as unknown as FollowupRun;
@@ -411,6 +413,25 @@ describe("runReplyAgent typing (heartbeat)", () => {
expect(typing.startTypingOnText).not.toHaveBeenCalled();
});
it("keeps assistant partial streaming enabled when reasoning mode is stream", async () => {
const onPartialReply = vi.fn();
const onReasoningStream = vi.fn();
state.runEmbeddedPiAgentMock.mockImplementationOnce(async (params: AgentRunParams) => {
await params.onReasoningStream?.({ text: "Reasoning:\n_step_" });
await params.onPartialReply?.({ text: "answer chunk" });
return { payloads: [{ text: "final" }], meta: {} };
});
const { run } = createMinimalRun({
opts: { onPartialReply, onReasoningStream },
runOverrides: { reasoningLevel: "stream" },
});
await run();
expect(onReasoningStream).toHaveBeenCalled();
expect(onPartialReply).toHaveBeenCalledWith({ text: "answer chunk", mediaUrls: undefined });
});
it("suppresses typing in never mode", async () => {
state.runEmbeddedPiAgentMock.mockImplementationOnce(async (params: AgentRunParams) => {
await params.onPartialReply?.({ text: "hi" });