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

@@ -104,13 +104,7 @@ export async function runAgentTurnWithFallback(params: {
while (true) {
try {
const allowPartialStream = !(
params.followupRun.run.reasoningLevel === "stream" && params.opts?.onReasoningStream
);
const normalizeStreamingText = (payload: ReplyPayload): { text?: string; skip: boolean } => {
if (!allowPartialStream) {
return { skip: true };
}
let text = payload.text;
if (!params.isHeartbeat && text?.includes("HEARTBEAT_OK")) {
const stripped = stripHeartbeatToken(text, {
@@ -290,18 +284,16 @@ export async function runAgentTurnWithFallback(params: {
abortSignal: params.opts?.abortSignal,
blockReplyBreak: params.resolvedBlockStreamingBreak,
blockReplyChunking: params.blockReplyChunking,
onPartialReply: allowPartialStream
? async (payload) => {
const textForTyping = await handlePartialForTyping(payload);
if (!params.opts?.onPartialReply || textForTyping === undefined) {
return;
}
await params.opts.onPartialReply({
text: textForTyping,
mediaUrls: payload.mediaUrls,
});
}
: undefined,
onPartialReply: async (payload) => {
const textForTyping = await handlePartialForTyping(payload);
if (!params.opts?.onPartialReply || textForTyping === undefined) {
return;
}
await params.opts.onPartialReply({
text: textForTyping,
mediaUrls: payload.mediaUrls,
});
},
onAssistantMessageStart: async () => {
await params.typingSignals.signalMessageStart();
await params.opts?.onAssistantMessageStart?.();

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" });