Studio: preserve durations across reasoning blocks (#7520)

* Studio: preserve durations across reasoning blocks

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

* Studio: keep a reasoning group's timer running when it reopens

A rendered reasoning group can be closed and then reopened: parseAssistantContent
coalesces adjacent reasoning parts, so a provider that emits each block as a
complete <think>...</think> chunk lands several blocks in one group. The tracker
wrote a group's duration once and never revisited it, so such a group froze at
its first close and displayed 0 seconds.

Measure from the first time an index becomes visible rather than from the last
startGroup, and reopen a closed group while its reasoning text is still growing.
Gating on growth is what stops the timer running on into the answer. A duration
supplied by the server is now recorded as authoritative so local timing cannot
overwrite it.

Also fill indices that a single delta skips. startGroup(n) could jump past
earlier indices and leave array holes, which JSON.stringify persists as null; a
skipped group became visible and closed inside the same chunk, so it gets a
measured zero instead.

Test discovery now globs tests/, so a second test file cannot be silently
skipped by CI, and tsconfig.test.json puts tests/ under typecheck for the first
time.

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: danielhanchen <danielhanchen@gmail.com>
This commit is contained in:
oobabooga 2026-07-28 08:53:26 -03:00 committed by GitHub
commit fc861cc870
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
11 changed files with 682 additions and 121 deletions

View file

@ -11392,6 +11392,7 @@ class LlamaCppBackend:
# Time each reasoning pass so final answers can replace tool timing.
_reasoning_started_at = None
_reasoning_summary_emitted = False
_deferred_reasoning_summary = None
cumulative_display = "" # Cumulative yielded text (with <think>)
in_thinking = False
has_content_tokens = False
@ -11643,7 +11644,11 @@ class LlamaCppBackend:
and not _reasoning_summary_emitted
):
_reasoning_summary_emitted = True
yield _reasoning_summary_event(_reasoning_started_at)
_summary = _reasoning_summary_event(_reasoning_started_at)
if _suppress_visible_output:
_deferred_reasoning_summary = _summary
else:
yield _summary
has_content_tokens = True
content_accum += token
@ -11927,7 +11932,11 @@ class LlamaCppBackend:
# route's extractor closes the streamed <think>).
if _reasoning_started_at is not None and not _reasoning_summary_emitted:
_reasoning_summary_emitted = True
yield _reasoning_summary_event(_reasoning_started_at)
_summary = _reasoning_summary_event(_reasoning_started_at)
if _suppress_visible_output:
_deferred_reasoning_summary = _summary
else:
yield _summary
cumulative_display = _finalize_reasoning_only_cumulative(
cumulative_display,
reasoning_accum,
@ -12041,6 +12050,8 @@ class LlamaCppBackend:
"type": "content",
"text": forced_visible_text,
}
if _deferred_reasoning_summary is not None:
yield _deferred_reasoning_summary
elif not _suppress_visible_output:
# Turn ended as a plain answer (no [ARGS] followed): the held
# rehearsal tail is real prose, release it.

View file

@ -602,7 +602,7 @@ def test_consumed_tool_final_pass_emits_latest_reasoning_summary(monkeypatch):
]
payloads: list[dict] = []
backend = _make_backend(monkeypatch, [tool_stream, final_stream], payloads)
_patch_monotonic(monkeypatch, [200.0, 201.0, 203.0, 300.0, 400.0, 405.0, 405.0])
_patch_monotonic(monkeypatch, [200.0, 201.0, 203.0, 300.0, 400.0, 405.0, 410.0])
def fake_execute_tool(name, arguments, **_kwargs):
return "Rendered HTML canvas: Done."
@ -1495,6 +1495,7 @@ def test_forced_reprompt_plain_final_answer_is_visible(monkeypatch):
streams = [
[_sse({"content": "I will use render_html now."}), _done()],
[
_sse({"reasoning_content": "I reconsidered the request."}),
_sse({"content": "No tool is needed. Final answer: use a red square."}),
_done(),
],
@ -1531,8 +1532,19 @@ def test_forced_reprompt_plain_final_answer_is_visible(monkeypatch):
content_texts = [event.get("text", "") for event in events if event.get("type") == "content"]
assert content_texts == [
"I will use render_html now.",
"No tool is needed. Final answer: use a red square.",
(
"<think>I reconsidered the request.</think>"
"No tool is needed. Final answer: use a red square."
),
]
summaries = [event for event in events if event.get("type") == "reasoning_summary"]
assert len(summaries) == 1
visible_answer_index = next(
index
for index, event in enumerate(events)
if event.get("type") == "content" and "No tool is needed" in event.get("text", "")
)
assert visible_answer_index < events.index(summaries[0])
assert len(payloads) == 2
@ -1774,24 +1786,14 @@ def test_reprompted_tool_call_still_streams_final_answer(monkeypatch):
streams = [
[_sse({"content": "I will use render_html now."}), _done()],
[
_sse({"reasoning_content": "I should render the requested HTML."}),
_sse(
{
"tool_calls": [
{
"index": 0,
"id": "call_forced",
"type": "function",
"function": {
"name": "render_html",
"arguments": json.dumps(
{
"code": "<html><body>forced</body></html>",
"title": "Forced",
}
),
},
}
]
"content": (
'<tool_call>{"name":"render_html","arguments":'
'{"code":"<html><body>forced</body></html>",'
'"title":"Forced"}}</tool_call>'
)
}
),
_done(),
@ -1835,6 +1837,7 @@ def test_reprompted_tool_call_still_streams_final_answer(monkeypatch):
assert len(calls) == 1
content_texts = [event.get("text", "") for event in events if event.get("type") == "content"]
assert content_texts == ["I will use render_html now.", "Final note after tool."]
assert not any(event.get("type") == "reasoning_summary" for event in events)
assert len(payloads) == 3

View file

@ -11,7 +11,8 @@
"build": "tsc -b && vite build",
"lint": "eslint .",
"preview": "vite preview",
"typecheck": "tsc -b --pretty false",
"test": "node --experimental-strip-types --test \"tests/**/*.test.ts\"",
"typecheck": "tsc -b --pretty false && tsc -p tsconfig.test.json --pretty false",
"i18n:check": "node --experimental-strip-types --no-warnings src/i18n/check-parity.ts",
"biome:check": "biome check",
"biome:fix": "biome check --write"

View file

@ -11,6 +11,7 @@ import {
CollapsibleContent,
CollapsibleTrigger,
} from "@/components/ui/collapsible";
import { resolveReasoningGroupDuration } from "@/features/chat";
import { useCollapseScrollLock } from "@/hooks/use-collapse-scroll-lock";
import { cn } from "@/lib/utils";
import {
@ -339,9 +340,11 @@ const ReasoningGroupImpl: ReasoningGroupComponent = ({
});
const persistedDuration = useAuiState(({ message }) => {
const d = (message.metadata?.custom as Record<string, unknown>)
?.reasoningDuration;
return typeof d === "number" ? d : 0;
return resolveReasoningGroupDuration(
message.parts,
startIndex,
message.metadata?.custom as Record<string, unknown> | undefined,
);
});
const [manualOpen, setManualOpen] = useState(false);
@ -412,7 +415,7 @@ const ReasoningGroupImpl: ReasoningGroupComponent = ({
className="min-w-0 flex-1"
active={isReasoningStreaming}
// Prefer server timing when available.
duration={persistedDuration || duration}
duration={persistedDuration ?? duration}
/>
<div className="flex w-16 shrink-0 justify-end">
{isOpen && !isReasoningStreaming && (

View file

@ -86,9 +86,15 @@ import {
} from "../utils/last-local-model-load";
import { getImageInputUnavailableReason } from "../utils/image-input-support";
import {
hasClosedThinkTag,
extractDeltaText,
hasUnclosedThinkTag,
parseAssistantContent,
} from "../utils/parse-assistant-content";
import {
countReasoningGroups,
createReasoningDurationTracker,
lastReasoningGroupTextLength,
} from "../utils/reasoning-duration";
import { resolveLoadMaxSeqLength } from "../presets/preset-policy";
import {
generateAudio,
@ -617,67 +623,6 @@ function estimateTokenCount(text: string): number | undefined {
return Math.max(1, Math.round(trimmed.length / 4));
}
/**
* Normalize a streamed `delta.content` to a plain text string.
*
* OpenAI Chat Completions originally typed `delta.content` as a string, but
* some providers now emit an array of structured content parts; concatenating
* those directly would stringify each as `[object Object]`. This guards that.
*
* Handled part shapes:
* { type: "text" | "output_text", text | content: "..." } text body
* { type: "thinking" | "reasoning", thinking | text: "..." } wrapped as
* inline `<think>...</think>` so `parseAssistantContent` lifts it into
* a reasoning part (else Mistral magistral and similar reasoning-part
* providers lose their thinking panel).
*
* Unknown part types are skipped better to drop a stray field than
* stringify an object into the rendered chat.
*/
function extractDeltaText(delta: unknown): string {
const extractReasoningText = (payload: unknown): string => {
if (typeof payload === "string") return payload;
if (Array.isArray(payload)) {
return payload.map((item) => extractReasoningText(item)).join("");
}
if (!payload || typeof payload !== "object") return "";
const obj = payload as Record<string, unknown>;
for (const key of ["thinking", "text", "content", "reasoning", "summary"]) {
if (key in obj) {
const text = extractReasoningText(obj[key]);
if (text) return text;
}
}
return "";
};
if (typeof delta === "string") return delta;
if (!Array.isArray(delta)) return "";
let out = "";
for (const part of delta) {
if (typeof part === "string") {
out += part;
continue;
}
if (!part || typeof part !== "object") continue;
const obj = part as {
type?: string;
text?: string;
content?: string;
thinking?: string;
};
if (obj.type === "text" || obj.type === "output_text") {
if (typeof obj.text === "string") out += obj.text;
else if (typeof obj.content === "string") out += obj.content;
} else if (obj.type === "thinking" || obj.type === "reasoning") {
const thinking = extractReasoningText(obj);
if (thinking) out += `<think>${thinking}</think>`;
}
}
return out;
}
function buildTiming(
streamStartTime: number,
totalChunks: number,
@ -2932,8 +2877,7 @@ export function createOpenAIStreamAdapter(
owner: serverCancel,
});
let cumulativeText = "";
let reasoningStartAt: number | null = null;
let reasoningDuration = 0;
const reasoningDurationTracker = createReasoningDurationTracker();
// True while wrapping a `delta.reasoning_content` stream in
// <think>...</think> for parseAssistantContent. Lives outside the
// SSE loop because the close tag fires when content arrives.
@ -3073,9 +3017,11 @@ export function createOpenAIStreamAdapter(
return merged;
};
const closeReasoningContent = () => {
if (!reasoningContentOpen) return;
cumulativeText += "</think>";
reasoningContentOpen = false;
if (reasoningContentOpen) {
cumulativeText += "</think>";
reasoningContentOpen = false;
}
reasoningDurationTracker.finishGroup();
};
// Anthropic document_citations payload, converted to Sources-panel
// parts at end-of-stream so inline [N] markers have matching entries.
@ -3631,8 +3577,9 @@ export function createOpenAIStreamAdapter(
const reasoningMs = (
chunk as { _reasoningDurationMs?: number } | null | undefined
)?._reasoningDurationMs;
if (typeof reasoningMs === "number" && Number.isFinite(reasoningMs)) {
reasoningDuration = Math.max(0, Math.round(reasoningMs / 1000));
if (
reasoningDurationTracker.recordServerDuration(reasoningMs)
) {
continue;
}
@ -3776,7 +3723,7 @@ export function createOpenAIStreamAdapter(
totalChunks,
firstTokenTime,
),
custom: { reasoningDuration },
custom: reasoningDurationTracker.metadata(),
},
};
}
@ -4071,7 +4018,7 @@ export function createOpenAIStreamAdapter(
totalChunks,
firstTokenTime,
),
custom: { reasoningDuration },
custom: reasoningDurationTracker.metadata(),
},
};
continue;
@ -4110,7 +4057,10 @@ export function createOpenAIStreamAdapter(
}
const rawDelta = chunk.choices?.[0]?.delta?.content;
// Normalize structured delta.content (mistral magistral).
const delta = extractDeltaText(rawDelta);
const {
text: delta,
structuredReasoningContinues,
} = extractDeltaText(rawDelta);
// Latest Gemini text-part thoughtSignature for next-turn replay.
const deltaExtraContent = (
chunk.choices?.[0]?.delta as
@ -4264,7 +4214,7 @@ export function createOpenAIStreamAdapter(
totalChunks,
firstTokenTime,
),
custom: { reasoningDuration },
custom: reasoningDurationTracker.metadata(),
},
};
continue;
@ -4281,6 +4231,7 @@ export function createOpenAIStreamAdapter(
if (reasoning) {
if (!reasoningContentOpen) {
reasoningDurationTracker.startGroup();
cumulativeText += `<think>${reasoning}`;
reasoningContentOpen = true;
} else {
@ -4288,7 +4239,9 @@ export function createOpenAIStreamAdapter(
}
}
if (delta) {
closeReasoningContent();
if (reasoningContentOpen) {
closeReasoningContent();
}
cumulativeText += delta;
}
// Strip a trailing ${...} template-literal fragment from
@ -4299,35 +4252,48 @@ export function createOpenAIStreamAdapter(
"",
);
}
const textParts = parseAssistantContent(cumulativeText);
const assistantContent = buildAssistantContent(cumulativeText);
// Fallback when no server-side reasoning_summary arrives.
const parsedReasoningGroupCount =
countReasoningGroups(assistantContent);
if (
textParts.some((part) => part.type === "reasoning") &&
!reasoningStartAt
parsedReasoningGroupCount >
reasoningDurationTracker.groupCount
) {
reasoningStartAt = Date.now();
}
if (
hasClosedThinkTag(cumulativeText) &&
reasoningStartAt &&
!reasoningDuration
) {
reasoningDuration = Math.round(
(Date.now() - reasoningStartAt) / 1000,
reasoningDurationTracker.startGroup(
parsedReasoningGroupCount - 1,
);
}
if (parsedReasoningGroupCount > 0) {
// Providers that close every reasoning block atomically
// (structured parts wrapped as <think>..</think>) end the group
// on each chunk. Reopen while the reasoning text is still
// growing so the timer spans the whole pass.
reasoningDurationTracker.resumeGroup(
parsedReasoningGroupCount - 1,
lastReasoningGroupTextLength(assistantContent),
);
}
if (
reasoningDurationTracker.hasActiveGroup &&
!reasoningContentOpen &&
!structuredReasoningContinues &&
!hasUnclosedThinkTag(cumulativeText)
) {
reasoningDurationTracker.finishGroup();
}
if (textParts.length > 0 || toolCallParts.length > 0) {
if (assistantContent.length > 0) {
yield {
content: buildAssistantContent(cumulativeText),
content: assistantContent,
metadata: {
timing: buildTiming(
streamStartTime,
totalChunks,
firstTokenTime,
),
custom: { reasoningDuration },
custom: reasoningDurationTracker.metadata(),
},
};
}
@ -4430,12 +4396,7 @@ export function createOpenAIStreamAdapter(
);
// Finalize reasoning-only streams.
if (reasoningStartAt && !reasoningDuration) {
reasoningDuration = Math.max(
0,
Math.round((Date.now() - reasoningStartAt) / 1000),
);
}
reasoningDurationTracker.finishGroup();
yield {
content: [
...buildAssistantContent(cumulativeText),
@ -4445,7 +4406,7 @@ export function createOpenAIStreamAdapter(
metadata: {
timing: finalTiming,
custom: {
reasoningDuration,
...reasoningDurationTracker.metadata(),
// Persisted refusal flag driving the two-pass prune.
anthropicRefusal: anthropicRefusalSeen || undefined,
serverTimings: meta?.timings ?? undefined,
@ -4504,6 +4465,30 @@ export function createOpenAIStreamAdapter(
});
}
}
if (!abortSignal.aborted) {
closeReasoningContent();
const partialContent = buildAssistantContent(cumulativeText);
if (partialContent.length > 0) {
const partialTiming = buildTiming(
streamStartTime,
totalChunks,
firstTokenTime,
Date.now() - streamStartTime,
estimateTokenCount(cumulativeText),
toolCallParts.length,
);
yield {
content: partialContent,
metadata: {
timing: partialTiming,
custom: {
...reasoningDurationTracker.metadata(),
timing: partialTiming,
},
},
};
}
}
throw err;
} finally {
runSignal.removeEventListener("abort", onAbortCancel);

View file

@ -91,6 +91,7 @@ export { clearAllChats, countAllChats } from "./utils/clear-all-chats";
export { pasteClipboardFiles } from "./utils/clipboard-files";
export { listStoredChatThreads } from "./utils/chat-history-storage";
export { emitChatAttachmentDeleted } from "./utils/chat-attachment-events";
export { resolveReasoningGroupDuration } from "./utils/reasoning-duration";
export { ArtifactCard } from "./artifacts/artifact-card";
export { ResearchMessage } from "./components/research-message";
export {

View file

@ -8,6 +8,78 @@ type ContentPart = NonNullable<ChatModelRunResult["content"]>[number];
const THINK_OPEN_TAG = "<think>";
const THINK_CLOSE_TAG = "</think>";
/**
* Normalize streamed string or structured delta content to inline text.
* Structured reasoning-only chunks remain distinguishable so their fallback
* timer can span consecutive chunks even though each chunk carries closed tags.
*/
export function extractDeltaText(delta: unknown): {
text: string;
structuredReasoningContinues: boolean;
} {
const extractReasoningText = (payload: unknown): string => {
if (typeof payload === "string") return payload;
if (Array.isArray(payload)) {
return payload.map((item) => extractReasoningText(item)).join("");
}
if (!payload || typeof payload !== "object") return "";
const obj = payload as Record<string, unknown>;
for (const key of ["thinking", "text", "content", "reasoning", "summary"]) {
if (key in obj) {
const text = extractReasoningText(obj[key]);
if (text) return text;
}
}
return "";
};
if (typeof delta === "string") {
return { text: delta, structuredReasoningContinues: false };
}
if (!Array.isArray(delta)) {
return { text: "", structuredReasoningContinues: false };
}
let text = "";
let structuredReasoningContinues = false;
for (const part of delta) {
if (typeof part === "string") {
text += part;
if (part) {
structuredReasoningContinues = false;
}
continue;
}
if (!part || typeof part !== "object") continue;
const obj = part as {
type?: string;
text?: string;
content?: string;
thinking?: string;
};
if (obj.type === "text" || obj.type === "output_text") {
const visibleText =
typeof obj.text === "string"
? obj.text
: typeof obj.content === "string"
? obj.content
: "";
text += visibleText;
if (visibleText) {
structuredReasoningContinues = false;
}
} else if (obj.type === "thinking" || obj.type === "reasoning") {
const thinking = extractReasoningText(obj);
if (thinking) {
text += `${THINK_OPEN_TAG}${thinking}${THINK_CLOSE_TAG}`;
structuredReasoningContinues = true;
}
}
}
return { text, structuredReasoningContinues };
}
// ContentPart from @assistant-ui/react has readonly fields, so coalescing via
// `last.text += text` fails (TS2540). Instead replace the last element with a
// fresh merged object: same allocation cost as mutation but type-safe.
@ -64,6 +136,6 @@ export function parseAssistantContent(
return parts;
}
export function hasClosedThinkTag(raw: string): boolean {
return raw.includes(THINK_CLOSE_TAG);
export function hasUnclosedThinkTag(raw: string): boolean {
return raw.lastIndexOf(THINK_OPEN_TAG) > raw.lastIndexOf(THINK_CLOSE_TAG);
}

View file

@ -0,0 +1,218 @@
// SPDX-License-Identifier: AGPL-3.0-only
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
type MessagePartLike = {
type?: unknown;
text?: unknown;
};
type ReasoningMetadata = {
reasoningDuration?: unknown;
reasoningDurations?: unknown;
};
function asDuration(value: unknown): number | undefined {
return typeof value === "number" && Number.isFinite(value) && value >= 0
? value
: undefined;
}
function getReasoningGroupIndex(
parts: readonly MessagePartLike[],
endIndex: number,
): number {
let index = -1;
let previousWasReasoning = false;
const limit = Math.min(endIndex, parts.length - 1);
for (let partIndex = 0; partIndex <= limit; partIndex += 1) {
const isReasoning = parts[partIndex]?.type === "reasoning";
if (isReasoning && !previousWasReasoning) {
index += 1;
}
previousWasReasoning = isReasoning;
}
return index;
}
export function countReasoningGroups(
parts: readonly MessagePartLike[],
): number {
return getReasoningGroupIndex(parts, parts.length - 1) + 1;
}
/**
* Total reasoning text in the LAST reasoning group (the group any new
* reasoning would join). The adapter compares this across chunks to tell "the
* model is still thinking" from "the model has moved on to the answer": a
* provider that closes every reasoning block atomically would otherwise freeze
* the group's timer at its first close.
*/
export function lastReasoningGroupTextLength(
parts: readonly MessagePartLike[],
): number {
let total = 0;
let inGroup = false;
for (let index = parts.length - 1; index >= 0; index -= 1) {
if (parts[index]?.type !== "reasoning") {
if (inGroup) break;
continue;
}
inGroup = true;
const text = parts[index]?.text;
total += typeof text === "string" ? text.length : 0;
}
return total;
}
export function resolveReasoningGroupDuration(
parts: readonly MessagePartLike[],
startIndex: number,
custom: ReasoningMetadata | null | undefined,
): number | undefined {
const index = getReasoningGroupIndex(parts, startIndex);
if (index < 0) {
return undefined;
}
if (Array.isArray(custom?.reasoningDurations)) {
return asDuration(custom.reasoningDurations[index]);
}
if (index !== getReasoningGroupIndex(parts, parts.length - 1)) {
return undefined;
}
return asDuration(custom?.reasoningDuration);
}
export function createReasoningDurationTracker(
now: () => number = Date.now,
) {
let durations: number[] = [];
// First time each group index became visible. A group can be closed and
// reopened -- a provider that emits several complete <think>...</think>
// blocks in a row has them coalesced into one rendered group -- so the
// duration is always measured from the first sighting, not the last.
const startedAt: number[] = [];
let activeIndex: number | null = null;
let groupCount = 0;
// Reasoning text seen so far per group, used to decide whether a closed
// group is still growing and should reopen.
const reasoningLength: number[] = [];
// The group a server summary would land on. The backend emits one summary at
// the end of each visible reasoning pass, before the next pass can begin, so
// "the group that started most recently" is the correct target. (A FIFO queue
// is tempting but wrong: it mis-assigns as soon as one group has no summary.)
let serverSummaryTargetIndex: number | null = null;
// Indices whose duration came from the server; local timing must not
// overwrite an authoritative value.
const serverClaimed = new Set<number>();
const setDuration = (index: number, duration: number) => {
if (durations[index] === duration) {
return;
}
const next = [...durations];
next[index] = duration;
durations = next;
};
const measure = (index: number, finishedAt: number) => {
if (serverClaimed.has(index)) {
return;
}
const from = startedAt[index];
if (from === undefined) {
return;
}
setDuration(index, Math.max(0, Math.round((finishedAt - from) / 1000)));
};
const finishGroupAt = (finishedAt: number) => {
if (activeIndex === null) {
return;
}
const index = activeIndex;
activeIndex = null;
measure(index, finishedAt);
};
return {
get groupCount() {
return groupCount;
},
get hasActiveGroup() {
return activeIndex !== null;
},
startGroup(index = groupCount) {
if (activeIndex === index) {
return;
}
const at = now();
finishGroupAt(at);
// A single delta can reveal more than one group at once. Any index we
// skipped became visible and closed within this same chunk, so give it a
// measured zero rather than leaving a hole in the persisted array.
for (let skipped = groupCount; skipped < index; skipped += 1) {
if (startedAt[skipped] === undefined) {
startedAt[skipped] = at;
}
measure(skipped, at);
}
if (startedAt[index] === undefined) {
startedAt[index] = at;
}
activeIndex = index;
groupCount = Math.max(groupCount, index + 1);
serverSummaryTargetIndex = index;
},
/**
* Reopen a group that already closed, but only while its reasoning text is
* still growing. Providers that emit each reasoning block as a complete
* <think>...</think> chunk close the group on every chunk; without this the
* group would freeze at the first close. Gating on growth is what keeps the
* timer from running on into the answer.
*/
resumeGroup(index: number, currentReasoningLength: number) {
const seen = reasoningLength[index] ?? 0;
if (currentReasoningLength <= seen) {
return;
}
reasoningLength[index] = currentReasoningLength;
if (activeIndex === index || startedAt[index] === undefined) {
return;
}
finishGroupAt(now());
activeIndex = index;
},
finishGroup() {
finishGroupAt(now());
},
recordServerDuration(reasoningMs: unknown): boolean {
if (
typeof reasoningMs !== "number" ||
!Number.isFinite(reasoningMs) ||
reasoningMs < 0
) {
return false;
}
if (serverSummaryTargetIndex !== null) {
serverClaimed.add(serverSummaryTargetIndex);
setDuration(
serverSummaryTargetIndex,
Math.max(0, Math.round(reasoningMs / 1000)),
);
serverSummaryTargetIndex = null;
}
return true;
},
metadata() {
if (durations.length === 0) {
return {};
}
return {
reasoningDuration: durations.at(-1) ?? 0,
reasoningDurations: durations,
};
},
};
}

View file

@ -0,0 +1,236 @@
// SPDX-License-Identifier: AGPL-3.0-only
// Copyright 2026-present the Unsloth AI Inc. team. All rights reserved. See /studio/LICENSE.AGPL-3.0
import assert from "node:assert/strict";
import test from "node:test";
import {
countReasoningGroups,
createReasoningDurationTracker,
lastReasoningGroupTextLength,
resolveReasoningGroupDuration,
} from "../src/features/chat/utils/reasoning-duration.ts";
import { extractDeltaText } from "../src/features/chat/utils/parse-assistant-content.ts";
const separatedReasoning = [
{ type: "reasoning" },
{ type: "tool-call" },
{ type: "reasoning" },
{ type: "text" },
];
test("selects per-group durations while preserving legacy messages", () => {
const current = {
reasoningDuration: 5,
reasoningDurations: [2, 5],
};
assert.equal(resolveReasoningGroupDuration(separatedReasoning, 0, current), 2);
assert.equal(resolveReasoningGroupDuration(separatedReasoning, 2, current), 5);
assert.equal(countReasoningGroups(separatedReasoning), 2);
const legacy = { reasoningDuration: 5 };
assert.equal(
resolveReasoningGroupDuration(separatedReasoning, 0, legacy),
undefined,
);
assert.equal(resolveReasoningGroupDuration(separatedReasoning, 2, legacy), 5);
const contiguous = [
{ type: "reasoning" },
{ type: "reasoning" },
{ type: "text" },
];
assert.equal(countReasoningGroups(contiguous), 1);
assert.equal(
resolveReasoningGroupDuration(contiguous, 0, {
reasoningDurations: [3],
}),
3,
);
});
test("tracks the exact reasoning, tool, reasoning sequence", () => {
let now = 0;
const tracker = createReasoningDurationTracker(() => now);
tracker.startGroup();
now = 1_200;
tracker.recordServerDuration(2_000);
tracker.finishGroup();
tracker.startGroup();
now = 5_600;
tracker.recordServerDuration(5_000);
tracker.finishGroup();
assert.deepEqual(tracker.metadata(), {
reasoningDuration: 5,
reasoningDurations: [2, 5],
});
});
test("keeps groups aligned when summaries are missing or orphaned", () => {
let now = 0;
const tracker = createReasoningDurationTracker(() => now);
tracker.startGroup();
now = 2_000;
tracker.finishGroup();
tracker.startGroup();
now = 7_000;
tracker.recordServerDuration(5_000);
tracker.finishGroup();
tracker.recordServerDuration(9_000);
assert.deepEqual(tracker.metadata(), {
reasoningDuration: 5,
reasoningDurations: [2, 5],
});
});
test("accepts zero after closure and rejects malformed server timing", () => {
let now = 0;
const tracker = createReasoningDurationTracker(() => now);
tracker.startGroup();
now = 1_000;
tracker.finishGroup();
assert.equal(tracker.recordServerDuration(0), true);
assert.equal(tracker.recordServerDuration(-1), false);
assert.deepEqual(tracker.metadata(), {
reasoningDuration: 0,
reasoningDurations: [0],
});
});
test("omits unknown timing and falls back to elapsed time", () => {
let now = 0;
const tracker = createReasoningDurationTracker(() => now);
tracker.startGroup();
assert.deepEqual(tracker.metadata(), {});
now = 3_200;
tracker.finishGroup();
assert.deepEqual(tracker.metadata(), {
reasoningDuration: 3,
reasoningDurations: [3],
});
});
test("keeps structured reasoning active only when it is the final content", () => {
assert.deepEqual(
extractDeltaText([{ type: "reasoning", text: "First" }]),
{
text: "<think>First</think>",
structuredReasoningContinues: true,
},
);
assert.deepEqual(
extractDeltaText([
{ type: "reasoning", text: "Last thought" },
{ type: "text", text: "Answer" },
]),
{
text: "<think>Last thought</think>Answer",
structuredReasoningContinues: false,
},
);
assert.deepEqual(
extractDeltaText([
{ type: "text", text: "Preface" },
{ type: "reasoning", text: "First thought" },
]),
{
text: "Preface<think>First thought</think>",
structuredReasoningContinues: true,
},
);
});
test("keeps a coalesced reasoning group growing across atomic blocks", () => {
let now = 1_770_000_000_000;
const tracker = createReasoningDurationTracker(() => now);
// A provider that closes every reasoning block in its own chunk still
// belongs to ONE rendered group, so the timer must span all of them.
tracker.startGroup();
tracker.resumeGroup(0, "first block".length);
tracker.finishGroup();
now += 3_000;
tracker.resumeGroup(0, "first blocksecond block".length);
tracker.finishGroup();
// The answer that follows adds no reasoning text, so the timer stops here.
now += 3_000;
tracker.resumeGroup(0, "first blocksecond block".length);
tracker.finishGroup();
assert.deepEqual(tracker.metadata(), {
reasoningDuration: 3,
reasoningDurations: [3],
});
});
test("never persists a hole when one delta reveals several groups", () => {
let now = 1_770_000_000_000;
const tracker = createReasoningDurationTracker(() => now);
// Index 0 was never started explicitly: it became visible and closed inside
// the same chunk that revealed index 1.
tracker.startGroup(1);
now += 4_000;
tracker.finishGroup();
const metadata = tracker.metadata();
const durations = metadata.reasoningDurations as number[];
assert.equal(durations.length, 2);
assert.ok(durations.every((value) => typeof value === "number"));
assert.deepEqual(JSON.parse(JSON.stringify(durations)), [0, 4]);
});
test("a server duration is never overwritten by local timing", () => {
let now = 1_770_000_000_000;
const tracker = createReasoningDurationTracker(() => now);
tracker.startGroup();
tracker.recordServerDuration(2_000);
now += 30_000;
tracker.resumeGroup(0, 99);
tracker.finishGroup();
assert.deepEqual(tracker.metadata(), {
reasoningDuration: 2,
reasoningDurations: [2],
});
});
test("lastReasoningGroupTextLength measures only the last reasoning group", () => {
assert.equal(
lastReasoningGroupTextLength([
{ type: "reasoning", text: "aaaa" },
{ type: "tool-call" },
{ type: "reasoning", text: "bb" },
{ type: "reasoning", text: "c" },
]),
3,
);
// The answer that follows is not reasoning, so it does not count -- but the
// group itself is still measured, which is what lets resumeGroup see that the
// reasoning has stopped growing.
assert.equal(
lastReasoningGroupTextLength([
{ type: "reasoning", text: "aaaa" },
{ type: "text", text: "answer" },
]),
4,
);
assert.equal(
lastReasoningGroupTextLength([{ type: "text", text: "answer only" }]),
0,
);
assert.equal(lastReasoningGroupTextLength([]), 0);
});

View file

@ -0,0 +1,28 @@
{
"compilerOptions": {
"target": "ES2022",
"lib": ["ES2022", "DOM", "DOM.Iterable"],
"module": "ESNext",
"types": ["node"],
"skipLibCheck": true,
"moduleResolution": "bundler",
"allowImportingTsExtensions": true,
"verbatimModuleSyntax": true,
"moduleDetection": "force",
"noEmit": true,
"jsx": "react-jsx",
"strict": true,
"noUnusedLocals": false,
"noUnusedParameters": true,
"erasableSyntaxOnly": true,
"noFallthroughCasesInSwitch": true,
"noUncheckedSideEffectImports": true,
"baseUrl": ".",
"paths": {
"@/*": ["./src/*"]
}
},
"include": ["tests"]
}