Files
oc-ls-stats/tui.tsx
T

548 lines
16 KiB
TypeScript

/** @jsxImportSource @opentui/solid */
import type { TuiPlugin, TuiPluginModule } from "@opencode-ai/plugin/tui"
import { createMemo, createSignal } from "solid-js"
import { appendFileSync } from "fs"
const DEBUG_ENABLED = true
const DEBUG_LOG = "/tmp/oc-ls-stats-debug.log"
function debug(...args: unknown[]) {
if (!DEBUG_ENABLED) return
try {
const line = `${args.map(a => typeof a === "object" ? JSON.stringify(a) : String(a)).join(" ")}\n`
appendFileSync(DEBUG_LOG, line)
} catch {}
}
type StreamSample = {
at: number
tokens: number
}
const STREAM_WINDOW_MS = 5_000
const SLOT_POLL_MS = 500
type MessageTiming = {
sessionID: string
requestStartAt: number
firstResponseAt?: number
firstTokenAt?: number
lastTokenAt?: number
}
type TrackerState = {
streamSamplesBySession: Record<string, StreamSample[]>
messageTimingByID: Record<string, MessageTiming>
lastPrefillRate: number
lastGeneratedTps: number
isPrefilling: boolean
prefillSlotId: number | null
prefillCapturedTokens: number | null
prefillStartAt: number
isGenerating: boolean
generateSlotId: number | null
generatePrevNd: number
generateStartAt: number
prevNdBySlot: Record<number, { baseline: number; hasIncreased: boolean }>
failure: string | null
}
function estimateStreamTokens(delta: string) {
return Math.max(1, Math.ceil(Buffer.byteLength(delta, "utf8") / 5))
}
function parseJSONC(text: string): unknown {
let result = ""
let inString = false
let escape = false
let i = 0
while (i < text.length) {
const ch = text[i]
if (escape) {
result += ch
escape = false
i++
continue
}
if (ch === "\\") {
result += ch
escape = true
i++
continue
}
if (ch === '"') {
result += ch
inString = !inString
i++
continue
}
if (inString) {
result += ch
i++
continue
}
if (ch === "/" && text[i + 1] === "/") {
while (i < text.length && text[i] !== "\n") i++
continue
}
if (ch === "/" && text[i + 1] === "*") {
i += 2
while (i < text.length - 1 && !(text[i] === "*" && text[i + 1] === "/")) i++
i += 2
result += " "
continue
}
if (ch === ",") {
const next = text.indexOf(/[\s}\]]/, i + 1)
if (next !== -1 && /[\s}\]]/.test(text[next])) {
i++
continue
}
}
result += ch
i++
}
return JSON.parse(result)
}
function stripBaseUrlPath(baseURL: string): string {
try {
const url = new URL(baseURL)
return `${url.origin}`
} catch {
return baseURL.replace(/\/v1(\/.*)?$/, "").replace(/\/$/, "") || baseURL
}
}
function extractProviderUrls(config: unknown): string[] {
if (!config || typeof config !== "object") { debug("extractProviderUrls: config is null/undefined or not object"); return [] }
const obj = config as Record<string, unknown>
const provider = obj.provider
if (!provider || typeof provider !== "object") { debug("extractProviderUrls: no provider in config"); return [] }
const urls: string[] = []
let total = 0
let matched = 0
for (const [key, val] of Object.entries(provider as Record<string, unknown>)) {
total++
if (val && typeof val === "object") {
const opts = (val as Record<string, unknown>).options
if (opts && typeof opts === "object") {
const baseURL = (opts as Record<string, unknown>).baseURL || (opts as Record<string, unknown>).base_url
debug(`extractProviderUrls: provider="${key}" baseURL=${baseURL}`)
if (typeof baseURL === "string" && key.toLowerCase().includes("llama") && !key.toLowerCase().includes("ollama")) {
urls.push(stripBaseUrlPath(baseURL))
matched++
debug(`extractProviderUrls: MATCH provider="${key}" -> ${urls[urls.length-1]}`)
} else {
debug(`extractProviderUrls: SKIP provider="${key}" (llama=${key.toLowerCase().includes("llama")} ollama=${key.toLowerCase().includes("ollama")})`)
}
}
}
}
debug(`extractProviderUrls: total=${total} matched=${matched} urls=${JSON.stringify(urls)}`)
return urls
}
async function loadConfigUrls(api: Parameters<TuiPlugin>[0]): Promise<string[]> {
try {
const configPath = api.state.path.config
const resp = await fetch(`file://${configPath}`)
if (!resp.ok) return []
const text = await resp.text()
const config = parseJSONC(text)
return extractProviderUrls(config)
} catch {
return []
}
}
async function fetchSlots(baseUrl: string, model?: string): Promise<unknown | null> {
try {
const url = model ? `${baseUrl}/slots?model=${encodeURIComponent(model)}` : `${baseUrl}/slots`
const resp = await fetch(url)
if (!resp.ok) return null
return await resp.json()
} catch {
return null
}
}
function formatPps(value: number) {
if (!Number.isFinite(value) || value <= 0) return undefined
return `${Math.round(value)}`
}
function formatTps(value: number) {
if (!Number.isFinite(value) || value <= 0) return undefined
return `${Math.round(value)}`
}
function SessionPromptRight(props: {
api: Parameters<TuiPlugin>[0]
sessionID: string
tracker: TrackerState
version: () => number
clock: () => number
}) {
const liveTps = createMemo(() => {
props.version()
props.clock()
const rate = props.tracker.lastGeneratedTps
if (!Number.isFinite(rate) || rate <= 0) return undefined
return rate
})
const prefillRate = createMemo(() => {
props.version()
if (!props.tracker.isPrefilling) return undefined
const rate = props.tracker.lastPrefillRate
if (!Number.isFinite(rate) || rate <= 0) return undefined
return rate
})
const text = createMemo(() => {
if (props.tracker.failure) return props.tracker.failure
const pp = prefillRate()
const tps = liveTps()
if (props.tracker.isPrefilling && pp) {
return `${formatPps(pp)} tps (PP)`
}
if (props.tracker.isGenerating && tps) {
return `${formatTps(tps)} tps (TG)`
}
return `- tps (TG)`
})
return <>{text() ? <text fg={props.api.theme.current.textMuted}>{text()}</text> : null}</>
}
const tui: TuiPlugin = async (api) => {
console.log("[oc-ls-stats] TUI plugin loaded!")
const tracker: TrackerState = {
streamSamplesBySession: {},
messageTimingByID: {},
lastPrefillRate: 0,
lastGeneratedTps: 0,
isPrefilling: false,
prefillSlotId: null,
prefillCapturedTokens: null,
prefillStartAt: 0,
isGenerating: false,
generateSlotId: null,
generatePrevNd: 0,
generateStartAt: 0,
prevNdBySlot: {},
failure: null,
}
const [version, setVersion] = createSignal(0)
const [clock, setClock] = createSignal(Date.now())
const bump = () => setVersion((value) => value + 1)
const pruneSamples = (now = Date.now()) => {
let changed = false
for (const [sessionID, samples] of Object.entries(tracker.streamSamplesBySession)) {
const next = samples.filter((sample) => now - sample.at <= STREAM_WINDOW_MS)
if (next.length !== samples.length) {
changed = true
if (next.length > 0) tracker.streamSamplesBySession[sessionID] = next
else delete tracker.streamSamplesBySession[sessionID]
}
}
if (changed) bump()
}
const clearLiveSamples = (sessionID: string) => {
if (!tracker.streamSamplesBySession[sessionID]?.length) return
delete tracker.streamSamplesBySession[sessionID]
bump()
}
const appendSample = (sessionID: string, messageID: string, sample: StreamSample) => {
const now = sample.at
tracker.streamSamplesBySession[sessionID] = [
...(tracker.streamSamplesBySession[sessionID] ?? []).filter((item) => now - item.at <= STREAM_WINDOW_MS),
sample,
]
const timing = tracker.messageTimingByID[messageID]
if (timing) {
tracker.messageTimingByID[messageID] = timing.firstTokenAt
? { ...timing, lastTokenAt: now }
: {
...timing,
firstResponseAt: timing.firstResponseAt ?? now,
firstTokenAt: now,
lastTokenAt: now,
}
}
bump()
}
let llamaServerUrls: string[] = []
let llamaServerModel: string | undefined
const pollMetrics = async () => {
if (!llamaServerUrls.length) return
let model = llamaServerModel
if (!model) {
const routeSessionID = api.route.current.params?.sessionID
if (routeSessionID) {
const session = api.state.session.get(routeSessionID)
if (session?.model?.id) {
model = session.model.id
}
}
}
let anyPrefilling = false
let anyGenerating = false
let genEnded = false
const prevProcessingSlotIds = new Set(Object.keys(tracker.prevNdBySlot).map(Number))
const currProcessingSlotIds = new Set<number>()
debug(`pollMetrics: urls=${JSON.stringify(llamaServerUrls)} model=${model}`)
for (const baseUrl of llamaServerUrls) {
debug(`pollMetrics: polling ${baseUrl}`)
const slots = await fetchSlots(baseUrl, model)
debug(`pollMetrics: ${baseUrl} response=${JSON.stringify(slots)?.slice(0, 500)}`)
const slotList = Array.isArray(slots) ? slots : slots ? Object.values(slots) : []
debug(`pollMetrics: ${baseUrl} slotList.length=${slotList.length}`)
if (slotList.length === 0) {
tracker.failure = "n/a"
bump()
continue
}
tracker.failure = null
bump()
debug(`poll baseUrl=${baseUrl} slots=${JSON.stringify(slotList)}`)
for (const slot of slotList) {
if (slot?.is_processing) {
currProcessingSlotIds.add(slot.id)
}
}
for (const slotId of prevProcessingSlotIds) {
if (!currProcessingSlotIds.has(slotId)) {
const prev = tracker.prevNdBySlot[slotId]
if (prev) {
prev.hasIncreased = false
prev.baseline = 0
} else {
delete tracker.prevNdBySlot[slotId]
}
}
}
let prefillSlot: any = null
let generateSlot: any = null
for (const slot of slotList) {
if (!slot?.is_processing) continue
const slotId = slot.id
const nd = slot.next_token?.[0]?.n_decoded ?? 0
const entry = tracker.prevNdBySlot[slotId]
const baseline = entry?.baseline ?? nd
const hasIncreased = entry?.hasIncreased ?? false
if (!hasIncreased && nd <= baseline) {
prefillSlot = slot
break
} else if (nd > baseline) {
tracker.prevNdBySlot[slotId] = { baseline, hasIncreased: true }
generateSlot = slot
}
}
for (const slot of slotList) {
if (slot?.is_processing) {
const slotId = slot.id
const nd = slot.next_token?.[0]?.n_decoded ?? 0
const entry = tracker.prevNdBySlot[slotId]
if (entry === undefined) {
tracker.prevNdBySlot[slotId] = { baseline: nd, hasIncreased: false }
}
}
}
if (prefillSlot) {
anyPrefilling = true
const slotId = prefillSlot.id
const now = Date.now()
const nt = prefillSlot.n_prompt_tokens ?? 0
if (tracker.prefillSlotId === slotId) {
const dt = (now - tracker.prefillStartAt) / 1000
const delta = nt - tracker.prefillCapturedTokens!
if (dt > 0 && delta > 0) {
tracker.lastPrefillRate = delta / dt
}
} else {
tracker.prefillSlotId = slotId
tracker.prefillCapturedTokens = nt
tracker.prefillStartAt = now
tracker.lastPrefillRate = 0
debug(`prefill new slot id=${slotId} nt=${nt}`)
}
} else {
if (tracker.isPrefilling) {
tracker.prefillSlotId = null
tracker.prefillCapturedTokens = null
tracker.prefillStartAt = 0
}
}
if (generateSlot) {
anyGenerating = true
const slotId = generateSlot.id
const now = Date.now()
const nd = generateSlot.next_token?.[0]?.n_decoded ?? 0
if (tracker.generateSlotId === slotId) {
const dt = (now - tracker.generateStartAt) / 1000
const delta = nd - tracker.generatePrevNd
if (dt > 0 && delta > 0) {
tracker.lastGeneratedTps = delta / dt
debug(`gen slot=${slotId} nd=${nd} prev=${tracker.generatePrevNd} delta=${delta} dt=${dt.toFixed(2)} tps=${tracker.lastGeneratedTps.toFixed(1)}`)
}
tracker.generatePrevNd = nd
tracker.generateStartAt = now
} else {
tracker.generateSlotId = slotId
tracker.generatePrevNd = nd
tracker.generateStartAt = now
tracker.lastGeneratedTps = 0
debug(`gen new slot id=${slotId} nd=${nd}`)
}
} else {
if (tracker.isGenerating) {
tracker.generateSlotId = null
tracker.generatePrevNd = 0
tracker.generateStartAt = 0
tracker.lastGeneratedTps = 0
genEnded = true
debug(`gen no slot found ending`)
}
}
tracker.isPrefilling = anyPrefilling
tracker.isGenerating = anyGenerating
}
if (tracker.lastPrefillRate > 0 || tracker.lastGeneratedTps > 0 || genEnded) {
bump()
}
}
const onDelta = api.event.on("message.part.delta", (evt) => {
if (evt.properties.field !== "text") return
const parts = api.state.part(evt.properties.messageID)
const part = parts.find((item) => item.id === evt.properties.partID)
if (!part) return
if (part.type !== "text" && part.type !== "reasoning") return
appendSample(evt.properties.sessionID, evt.properties.messageID, {
at: Date.now(),
tokens: estimateStreamTokens(evt.properties.delta),
})
})
const onMessage = api.event.on("message.updated", (evt) => {
if (evt.properties.info.role !== "assistant") return
if (!evt.properties.info.time.completed) {
const existing = tracker.messageTimingByID[evt.properties.info.id]
tracker.messageTimingByID[evt.properties.info.id] = {
sessionID: evt.properties.sessionID,
requestStartAt: evt.properties.info.time.created,
firstResponseAt: existing?.firstResponseAt,
firstTokenAt: existing?.firstTokenAt,
lastTokenAt: existing?.lastTokenAt,
}
bump()
return
}
delete tracker.messageTimingByID[evt.properties.info.id]
pruneSamples(evt.properties.info.time.completed)
bump()
})
const onPart = api.event.on("message.part.updated", (evt) => {
if (evt.properties.part.type !== "tool") return
if (
evt.properties.part.state.status === "running" ||
evt.properties.part.state.status === "completed" ||
evt.properties.part.state.status === "error"
) {
clearLiveSamples(evt.properties.sessionID)
}
const timing = tracker.messageTimingByID[evt.properties.part.messageID]
if (!timing) return
if (evt.properties.part.state.status === "pending") {
tracker.messageTimingByID[evt.properties.part.messageID] = {
...timing,
firstResponseAt: timing.firstResponseAt ?? evt.properties.time,
}
bump()
return
}
if (evt.properties.part.state.status !== "running") return
bump()
})
const timer = setInterval(() => {
setClock(Date.now())
pruneSamples()
}, 1000)
const slotTimer = setInterval(async () => {
await pollMetrics()
}, SLOT_POLL_MS)
api.lifecycle.onDispose(() => {
onDelta()
onMessage()
onPart()
clearInterval(timer)
clearInterval(slotTimer)
})
try {
const configUrls = await loadConfigUrls(api)
debug(`init: configUrls=${JSON.stringify(configUrls)} fallback=${JSON.stringify(["http://localhost:8080"])}`)
llamaServerUrls = configUrls.length > 0 ? configUrls : ["http://localhost:8080"]
} catch (e) {
console.error("[oc-ls-stats] loadConfigUrls error:", e)
llamaServerUrls = ["http://localhost:8080"]
}
debug(`init: final llamaServerUrls=${JSON.stringify(llamaServerUrls)}`)
api.slots.register({
slots: {
session_prompt_right(_ctx, value) {
return <SessionPromptRight api={api} sessionID={value.session_id} tracker={tracker} version={version} clock={clock} />
},
},
})
}
const plugin: TuiPluginModule & { id: string } = {
id: "@troed/oc-ls-stats",
tui,
}
export default plugin