diff --git a/packages/core/src/session/tokens.ts b/packages/core/src/session/tokens.ts new file mode 100644 index 000000000000..dc9d9deb9032 --- /dev/null +++ b/packages/core/src/session/tokens.ts @@ -0,0 +1,92 @@ +export const DEFAULT_MIN_TPS_ELAPSED_MS = 250 +export const DEFAULT_INCLUDE_REASONING = true + +export interface TokenMetrics { + output: number + reasoning: number +} + +export interface TimestampMetrics { + created: number + firstToken?: number + completed?: number +} + +export interface TPSResult { + rate: number + totalTokens: number + elapsedMs: number + isValid: boolean +} + +export function stampFirstToken(time: TimestampMetrics, now: number): number { + if (time.firstToken === undefined) time.firstToken = now + return time.firstToken +} + +export function totalGeneratedTokens(tokens: TokenMetrics, includeReasoning = DEFAULT_INCLUDE_REASONING): number { + return tokens.output + (includeReasoning ? tokens.reasoning : 0) +} + +type TPSMessage = { + summary?: boolean + finish?: string | null + tokens: TokenMetrics + time: TimestampMetrics +} + +function tpsInputs(msg: TPSMessage): { totalTokens: number; elapsedMs: number } | undefined { + if (msg.summary) return undefined + if (!msg.finish) return undefined + if (["tool-calls", "unknown", "error"].includes(msg.finish)) return undefined + + const totalTokens = totalGeneratedTokens(msg.tokens) + if (totalTokens <= 0) return undefined + const { firstToken, completed } = msg.time + if (firstToken === undefined || completed === undefined) return undefined + + return { totalTokens, elapsedMs: completed - firstToken } +} + +export function isValidForTPS(msg: TPSMessage & { + minElapsedMs?: number +}): boolean { + const inputs = tpsInputs(msg) + if (!inputs) return false + const minElapsedMs = msg.minElapsedMs ?? DEFAULT_MIN_TPS_ELAPSED_MS + return inputs.elapsedMs >= minElapsedMs +} + +export function calculateTPS( + totalTokens: number, + elapsedMs: number, + minElapsedMs = DEFAULT_MIN_TPS_ELAPSED_MS, +): TPSResult | undefined { + if (totalTokens <= 0) return undefined + if (elapsedMs < minElapsedMs) return undefined + + const rate = totalTokens / (elapsedMs / 1000) + if (!Number.isFinite(rate) || rate < 0) return undefined + + return { + rate: Math.round(rate), + totalTokens, + elapsedMs, + isValid: true, + } +} + +export function formatTPS(result: TPSResult): string { + return `${result.rate.toLocaleString()} tok/s` +} + +export function getMessageTPS(msg: { + summary?: boolean + finish?: string | null + tokens: TokenMetrics + time: TimestampMetrics +}): TPSResult | undefined { + const inputs = tpsInputs(msg) + if (!inputs) return undefined + return calculateTPS(inputs.totalTokens, inputs.elapsedMs) +} diff --git a/packages/core/test/session-tps.test.ts b/packages/core/test/session-tps.test.ts new file mode 100644 index 000000000000..be562c6a13e3 --- /dev/null +++ b/packages/core/test/session-tps.test.ts @@ -0,0 +1,139 @@ +import { describe, expect, test } from "bun:test" +import { + calculateTPS, + DEFAULT_MIN_TPS_ELAPSED_MS, + getMessageTPS, + isValidForTPS, + stampFirstToken, + type TimestampMetrics, +} from "../src/session/tokens" + +const validMessage = { + finish: "stop", + tokens: { output: 100, reasoning: 50 }, + time: { created: 1000, firstToken: 1100, completed: 2100 }, +} + +describe("getMessageTPS", () => { + test("calculates rounded output and reasoning tokens per second", () => { + expect(getMessageTPS(validMessage)?.rate).toBe(150) + }) + + test("returns no value for summary messages", () => { + expect(getMessageTPS({ ...validMessage, summary: true })).toBeUndefined() + }) + + test("returns no value when finish is missing", () => { + expect(getMessageTPS({ ...validMessage, finish: undefined })).toBeUndefined() + }) + + test("returns no value for tool-call finishes", () => { + expect(getMessageTPS({ ...validMessage, finish: "tool-calls" })).toBeUndefined() + }) + + test("returns no value for unknown finishes", () => { + expect(getMessageTPS({ ...validMessage, finish: "unknown" })).toBeUndefined() + }) + + test("returns no value for error finishes", () => { + expect(getMessageTPS({ ...validMessage, finish: "error" })).toBeUndefined() + }) + + test("returns no value when token total is zero", () => { + expect(getMessageTPS({ ...validMessage, tokens: { output: 0, reasoning: 0 } })).toBeUndefined() + }) + + test("returns no value when first token timestamp is missing", () => { + expect(getMessageTPS({ ...validMessage, time: { created: 1000, completed: 2100 } })).toBeUndefined() + }) + + test("returns no value when completion timestamp is missing", () => { + expect(getMessageTPS({ ...validMessage, time: { created: 1000, firstToken: 1100 } })).toBeUndefined() + }) + + test("returns no value when elapsed time is below 250 milliseconds", () => { + expect( + getMessageTPS({ + ...validMessage, + time: { + ...validMessage.time, + completed: validMessage.time.firstToken! + DEFAULT_MIN_TPS_ELAPSED_MS - 1, + }, + }), + ).toBeUndefined() + }) + + test("accepts elapsed time at the 250 millisecond threshold", () => { + expect( + getMessageTPS({ + ...validMessage, + time: { + ...validMessage.time, + completed: validMessage.time.firstToken! + DEFAULT_MIN_TPS_ELAPSED_MS, + }, + })?.rate, + ).toBe(600) + }) +}) + +describe("isValidForTPS", () => { + test("rejects negative token totals", () => { + expect(isValidForTPS({ ...validMessage, tokens: { output: -1, reasoning: 0 } })).toBe(false) + }) + + test("rejects zero token totals", () => { + expect(isValidForTPS({ ...validMessage, tokens: { output: 0, reasoning: 0 } })).toBe(false) + }) + + test("rejects missing finish values", () => { + expect(isValidForTPS({ ...validMessage, finish: null })).toBe(false) + }) + + test("rejects invalid finish values", () => { + for (const finish of ["tool-calls", "unknown", "error"]) { + expect(isValidForTPS({ ...validMessage, finish })).toBe(false) + } + }) + + test("rejects a missing first token timestamp", () => { + expect(isValidForTPS({ ...validMessage, time: { created: 1000, completed: 2100 } })).toBe(false) + }) + + test("rejects a missing completion timestamp", () => { + expect(isValidForTPS({ ...validMessage, time: { created: 1000, firstToken: 1100 } })).toBe(false) + }) + + test("rejects elapsed time below the configured threshold", () => { + expect( + isValidForTPS({ + ...validMessage, + time: { ...validMessage.time, completed: validMessage.time.firstToken! + 249 }, + }), + ).toBe(false) + }) +}) + +describe("calculateTPS", () => { + test("rounds fractional rates to the nearest integer", () => { + expect(calculateTPS(100, 327)?.rate).toBe(306) + }) + + test("rejects zero and negative token totals", () => { + expect(calculateTPS(0, 1000)).toBeUndefined() + expect(calculateTPS(-1, 1000)).toBeUndefined() + }) + + test("rejects elapsed time below the configured threshold", () => { + expect(calculateTPS(100, 249)).toBeUndefined() + }) +}) + +describe("stampFirstToken", () => { + test("stamps the first event and preserves it for later deltas", () => { + const time: TimestampMetrics = { created: 1000 } + + expect(stampFirstToken(time, 1100)).toBe(1100) + expect(stampFirstToken(time, 1200)).toBe(1100) + expect(time.firstToken).toBe(1100) + }) +}) diff --git a/packages/opencode/src/session/processor.ts b/packages/opencode/src/session/processor.ts index 20aa8a8404d8..f58081a1ed7d 100644 --- a/packages/opencode/src/session/processor.ts +++ b/packages/opencode/src/session/processor.ts @@ -1,4 +1,5 @@ import { LayerNode } from "@opencode-ai/core/effect/layer-node" +import { stampFirstToken } from "@opencode-ai/core/session/tokens" import { PermissionV1 } from "@opencode-ai/core/v1/permission" import { Image } from "@/image/image" import { SessionV1 } from "@opencode-ai/core/v1/session" @@ -120,6 +121,10 @@ const layer = Layer.effect( aborted, }) + const markFirstToken = () => { + stampFirstToken(ctx.assistantMessage.time, Date.now()) + } + const settleToolCall = Effect.fn("SessionProcessor.settleToolCall")(function* (toolCallID: string) { const done = ctx.toolcalls[toolCallID]?.done delete ctx.toolcalls[toolCallID] @@ -279,6 +284,7 @@ const layer = Layer.effect( switch (value.type) { case "reasoning-start": if (value.id in ctx.reasoningMap) return + markFirstToken() ctx.reasoningMap[value.id] = { id: PartID.ascending(), messageID: ctx.assistantMessage.id, @@ -294,6 +300,7 @@ const layer = Layer.effect( case "reasoning-delta": // Match dev: silently drop orphan deltas (no preceding reasoning-start). if (!(value.id in ctx.reasoningMap)) return + markFirstToken() ctx.reasoningMap[value.id].text += value.text if (value.providerMetadata) ctx.reasoningMap[value.id].metadata = value.providerMetadata yield* session.updatePartDelta({ @@ -316,15 +323,18 @@ const layer = Layer.effect( if (ctx.assistantMessage.summary) { throw new Error(`Tool call not allowed while generating summary: ${value.name}`) } + markFirstToken() yield* ensureToolCall(value) return case "tool-input-delta": yield* ensureToolCall(value) + markFirstToken() return case "tool-input-end": { yield* ensureToolCall(value) + markFirstToken() return } @@ -332,6 +342,7 @@ const layer = Layer.effect( if (ctx.assistantMessage.summary) { throw new Error(`Tool call not allowed while generating summary: ${value.name}`) } + markFirstToken() yield* ensureToolCall(value) const input = isRecord(value.input) ? value.input : { value: value.input } yield* updateToolCall(value.id, (match) => ({ @@ -433,6 +444,7 @@ const layer = Layer.effect( return case "step-finish": { + markFirstToken() const completedSnapshot = yield* snapshot.track() yield* Effect.forEach(Object.keys(ctx.reasoningMap), finishReasoning) const usage = Session.getUsage({ @@ -484,6 +496,7 @@ const layer = Layer.effect( } case "text-start": + markFirstToken() ctx.currentText = { id: PartID.ascending(), messageID: ctx.assistantMessage.id, @@ -498,6 +511,7 @@ const layer = Layer.effect( case "text-delta": if (!ctx.currentText) return + markFirstToken() ctx.currentText.text += value.text if (value.providerMetadata) ctx.currentText.metadata = value.providerMetadata yield* session.updatePartDelta({ @@ -532,6 +546,7 @@ const layer = Layer.effect( return case "finish": + markFirstToken() return } }) diff --git a/packages/opencode/test/session/processor-effect.test.ts b/packages/opencode/test/session/processor-effect.test.ts index c67f82d9c71b..09390ea064c1 100644 --- a/packages/opencode/test/session/processor-effect.test.ts +++ b/packages/opencode/test/session/processor-effect.test.ts @@ -209,6 +209,27 @@ const providerErrorLLM = Layer.succeed( const providerErrorEnv = LayerNode.compile(root, [...replacements, [LLM.node, providerErrorLLM]]) const itProviderError = testEffect(providerErrorEnv) +const toolOnlyLLM = Layer.succeed( + LLM.Service, + LLM.Service.of({ + stream: () => + Stream.make( + LLMEvent.stepStart({ index: 0 }), + LLMEvent.toolCall({ id: "call-only", name: "lookup", input: {}, providerExecuted: true }), + LLMEvent.toolResult({ + id: "call-only", + name: "lookup", + result: { type: "json", value: { output: "tool result" } }, + providerExecuted: true, + }), + LLMEvent.stepFinish({ index: 0, reason: "stop" }), + LLMEvent.finish({ reason: "stop" }), + ), + }), +) +const toolOnlyEnv = LayerNode.compile(root, [...replacements, [LLM.node, toolOnlyLLM]]) +const itToolOnly = testEffect(toolOnlyEnv) + const fragmentFailureLLM = Layer.succeed( LLM.Service, LLM.Service.of({ @@ -1117,6 +1138,40 @@ itProviderError.live("session.processor effect tests fail provider-executed erro ), ) +itToolOnly.live("session.processor effect tests stamp first token for tool-only turns", () => + provideTmpdirInstance( + (dir) => + Effect.gen(function* () { + const { processors, session, provider } = yield* boot() + const chat = yield* session.create({}) + const parent = yield* user(chat.id, "tool-only") + const msg = yield* assistant(chat.id, parent.id, path.resolve(dir)) + const mdl = yield* provider.getModel(ref.providerID, ref.modelID) + const handle = yield* processors.create({ assistantMessage: msg, sessionID: chat.id, model: mdl }) + + yield* handle.process({ + user: { + id: parent.id, + sessionID: chat.id, + role: "user", + time: parent.time, + agent: parent.agent, + model: { providerID: ref.providerID, modelID: ref.modelID }, + } satisfies SessionV1.User, + sessionID: chat.id, + model: mdl, + agent: agent(), + system: [], + messages: [{ role: "user", content: "tool-only" }], + tools: {}, + }) + + expect(handle.message.time.firstToken).toBeDefined() + }), + { config: cfg }, + ), +) + itFragmentFailure.live("session.processor effect tests retain partial legacy parts without v2 events", () => provideTmpdirInstance( (dir) => diff --git a/packages/schema/src/v1/session.ts b/packages/schema/src/v1/session.ts index 75e9282f117c..bf9e1bc7d4b5 100644 --- a/packages/schema/src/v1/session.ts +++ b/packages/schema/src/v1/session.ts @@ -456,6 +456,7 @@ export const Assistant = Schema.Struct({ time: Schema.Struct({ created: NonNegativeInt, completed: Schema.optional(NonNegativeInt), + firstToken: Schema.optional(NonNegativeInt), }), error: Schema.optional(AssistantErrorSchema), parentID: MessageID, diff --git a/packages/sdk/js/src/v2/gen/types.gen.ts b/packages/sdk/js/src/v2/gen/types.gen.ts index 72b5e6f30ace..6b90898f951e 100644 --- a/packages/sdk/js/src/v2/gen/types.gen.ts +++ b/packages/sdk/js/src/v2/gen/types.gen.ts @@ -337,6 +337,7 @@ export type AssistantMessage = { time: { created: number completed?: number + firstToken?: number } error?: | ProviderAuthError diff --git a/packages/sdk/openapi.json b/packages/sdk/openapi.json index 5e372b6fb6b8..1dc845130cc6 100644 --- a/packages/sdk/openapi.json +++ b/packages/sdk/openapi.json @@ -16256,6 +16256,10 @@ "completed": { "type": "integer", "minimum": 0 + }, + "firstToken": { + "type": "integer", + "minimum": 0 } }, "required": ["created"], diff --git a/packages/tui/src/routes/session/index.tsx b/packages/tui/src/routes/session/index.tsx index cbdaf0cfa0c7..4397385bb5cb 100644 --- a/packages/tui/src/routes/session/index.tsx +++ b/packages/tui/src/routes/session/index.tsx @@ -76,6 +76,7 @@ import { useClipboard } from "../../context/clipboard" import { nextThinkingMode, reasoningSummary, useThinkingMode, type ThinkingMode } from "../../context/thinking" import { getScrollAcceleration } from "../../util/scroll" import { collapseToolOutput } from "../../util/collapse-tool-output" +import { formatTPS, getMessageTPS } from "@opencode-ai/core/session/tokens" import { usePluginRuntime } from "../../plugin/runtime" import { DialogRetryAction } from "../../component/dialog-retry-action" import { getRevertDiffFiles } from "../../util/revert-diff" @@ -1487,6 +1488,8 @@ function AssistantMessage(props: { message: AssistantMessage; parts: Part[]; las return props.message.time.completed - user.time.created }) + const tps = createMemo(() => getMessageTPS(props.message)) + const childShortcut = useCommandShortcut("session.child.first") const backgroundShortcut = useCommandShortcut("session.background") @@ -1565,6 +1568,9 @@ function AssistantMessage(props: { message: AssistantMessage; parts: Part[]; las · {Locale.duration(duration())} + + · {formatTPS(tps()!)} + · interrupted