oc-ls-stats/tui.tsx

/** @jsxImportSource @opentui/solid */
import type { TextRenderable } from "@opentui/core"
import type { TuiPlugin, TuiPluginModule } from "@opencode-ai/plugin/tui"
import { onCleanup, createSignal } from "solid-js"
import {
  backoffDelayMs,
  processPoll,
  resolveServerUrls,
  selectSessionUrls,
  serverTypeOverride,
  isStrataHealth,
  mapStrataLive,
  type ServerType,
} from "./src/stats.ts"

type StreamSample = {
  at: number
  tokens: number
}

const STREAM_WINDOW_MS = 5_000
const SLOT_POLL_MS = 500
const BACKOFF_BASE_MS = 1_000
const BACKOFF_MAX_MS = 10_000
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, { prevNd: number; idTask: number; classification: "TG" | "PP" | "neither" }>
  prevPromptTokensBySlot: Record<number, { tokens: number; idTask: number }>
  failure: string | null
}

function estimateStreamTokens(delta: string) {
  return Math.max(1, Math.ceil(Buffer.byteLength(delta, "utf8") / 5))
}

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
  }
}

async function fetchHealth(baseUrl: string): Promise<unknown | null> {
  try {
    const resp = await fetch(`${baseUrl}/health`)
    if (!resp.ok) return null
    return await resp.json()
  } catch {
    return null
  }
}

// Strata reports the live rates under /metrics.live (prefill_tok_s_mean while
// reading, tok_s while generating) rather than in the minimal /slots stub.
async function fetchStrataLive(baseUrl: string): Promise<unknown | null> {
  try {
    const resp = await fetch(`${baseUrl}/metrics`)
    if (!resp.ok) return null
    const body = await resp.json()
    if (body && typeof body === "object" && "live" in body) return (body as { live: unknown }).live
    return null
  } 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
  subscribe: (listener: () => void) => () => void
}) {
  let textRef: TextRenderable | undefined

  const sync = () => {
    if (!textRef) return
    textRef.content = getStatusText()
    props.api.renderer.requestRender()
  }

  const unsubscribe = props.subscribe(sync)
  onCleanup(unsubscribe)

  return (
    <text
      ref={(ref: TextRenderable) => {
        textRef = ref
        sync()
      }}
      fg={props.api.theme.current.textMuted}
    >
      {getStatusText()}
    </text>
  )

  function getStatusText() {
    const s = props.tracker
    if (s.failure) {
      return s.failure
    }
    const pp = s.isPrefilling && s.lastPrefillRate > 0 ? formatPps(s.lastPrefillRate) : undefined
    const tps = s.isGenerating && s.lastGeneratedTps > 0 ? formatTps(s.lastGeneratedTps) : undefined
    let result = "- tps (TG)"
    if (s.isPrefilling && pp) {
      result = `${pp} tps (PP)`
    } else if (s.isGenerating && tps) {
      result = `${tps} tps (TG)`
    }
    return result
  }
}

const tui: TuiPlugin = async (api, options) => {
  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: {},
    prevPromptTokensBySlot: {},
    failure: null,
  }
  const listeners = new Set<() => void>()

  const bump = () => {
    for (const listener of listeners) listener()
  }

  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 backoffByUrl = new Map<string, { failures: number; nextRetryAt: number }>()
  const serverKindByUrl = new Map<string, ServerType>()
  const forcedServerType = serverTypeOverride((options as Record<string, unknown> | undefined)?.serverType)

  const resolveServerKind = async (baseUrl: string): Promise<ServerType> => {
    if (forcedServerType) return forcedServerType
    const cached = serverKindByUrl.get(baseUrl)
    if (cached) return cached
    const health = await fetchHealth(baseUrl)
    if (health != null) {
      const kind = isStrataHealth(health) ? "strata" : "llama"
      serverKindByUrl.set(baseUrl, kind)
      return kind
    }
    return "llama"
  }

  const noteServerFailure = (baseUrl: string) => {
    const failures = (backoffByUrl.get(baseUrl)?.failures ?? 0) + 1
    const delay = backoffDelayMs(failures, BACKOFF_BASE_MS, BACKOFF_MAX_MS)
    const jitter = Math.floor(Math.random() * (delay / 4))
    backoffByUrl.set(baseUrl, {
      failures,
      nextRetryAt: Date.now() + Math.min(delay + jitter, BACKOFF_MAX_MS),
    })
    tracker.failure = "n/a"
    bump()
  }

  const pollMetrics = async () => {
    if (!llamaServerUrls.length) return

    let model = llamaServerModel
    let sessionProviderId: string | undefined
    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
        }
        sessionProviderId = session?.model?.providerID
      }
    }

    const activeUrls = selectSessionUrls(llamaServerUrls, api.state.provider, sessionProviderId)
    if (!activeUrls.length) {
      tracker.failure = "n/a"
      bump()
      return
    }
    for (const baseUrl of activeUrls) {
      const now = Date.now()
      const backoff = backoffByUrl.get(baseUrl)
      if (backoff && now < backoff.nextRetryAt) {
        tracker.failure = "n/a"
        bump()
        continue
      }

      if ((await resolveServerKind(baseUrl)) === "strata") {
        const live = await fetchStrataLive(baseUrl)
        if (live == null) {
          noteServerFailure(baseUrl)
          continue
        }
        backoffByUrl.delete(baseUrl)
        tracker.failure = null
        mapStrataLive(live, tracker)
        bump()
        continue
      }

      const slots = await fetchSlots(baseUrl, model)
      if (slots == null) {
        noteServerFailure(baseUrl)
        continue
      }

      backoffByUrl.delete(baseUrl)
      const slotList = Array.isArray(slots) ? slots : Object.values(slots)
      if (slotList.length === 0) {
        tracker.failure = "n/a"
        bump()
        continue
      }
      tracker.failure = null
      bump()

      processPoll(slotList, tracker)
      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(() => {
    pruneSamples()
  }, 1000)

  const slotTimer = setInterval(async () => {
    await pollMetrics()
  }, SLOT_POLL_MS)

  api.lifecycle.onDispose(() => {
    onDelta()
    onMessage()
    onPart()
    clearInterval(timer)
    clearInterval(slotTimer)
  })

  try {
    llamaServerUrls = resolveServerUrls(options, api.state.provider, api.state.config)
  } catch (e) {
    console.error("[oc-ls-stats] resolveServerUrls error:", e)
    llamaServerUrls = ["http://localhost:8080"]
  }

  api.slots.register({
    slots: {
      session_prompt_right(_ctx, value) {
        return <SessionPromptRight api={api} sessionID={value.session_id} tracker={tracker} subscribe={(listener) => {
          listeners.add(listener)
          return () => {
            listeners.delete(listener)
          }
        }} />
      },
    },
  })
}

// ---------------------------------------------------------------------------
// V2 (opencode >= 2.0) implementation
//
// V2 TUI plugin system replaces the V1 tui(api) function with
// setup(context): the readout registers through context.ui.slot on
// prompt.footer.status (the V2 successor of session_prompt_right), events
// arrive through context.data.on, and provider/session state comes from
// context.data.location / context.data.session. The V1 tui() above is kept
// untouched; both implementations are exposed from the same default export
// (V1 calls tui, V2 calls setup) and share the tracker helpers from
// src/stats.ts.
// ---------------------------------------------------------------------------

const v2Setup = async (context: any) => {
  const location = context.location ?? context.data.location.default()
  const themeMuted = context.theme?.text?.muted

  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: {},
    prevPromptTokensBySlot: {},
    failure: null,
  }

  // part id -> part type, learned from message.part.updated; used to skip
  // non-text/reasoning deltas (e.g. tool output) like the V1 part lookup
  const partTypeByID = new Map<string, string>()

  function getStatusText() {
    const s = tracker
    if (s.failure) {
      return s.failure
    }
    const pp = s.isPrefilling && s.lastPrefillRate > 0 ? formatPps(s.lastPrefillRate) : undefined
    const tps = s.isGenerating && s.lastGeneratedTps > 0 ? formatTps(s.lastGeneratedTps) : undefined
    let result = "- tps (TG)"
    if (s.isPrefilling && pp) {
      result = `${pp} tps (PP)`
    } else if (s.isGenerating && tps) {
      result = `${tps} tps (TG)`
    }
    return result
  }

  const [statusText, setStatusText] = createSignal<string>("- tps (TG)")

  const bump = () => {
    setStatusText(getStatusText())
  }

  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 providerList: unknown[] = []
  const backoffByUrl = new Map<string, { failures: number; nextRetryAt: number }>()
  const serverKindByUrl = new Map<string, ServerType>()
  const forcedServerType = serverTypeOverride((context.options as Record<string, unknown> | undefined)?.serverType)

  const resolveServerKind = async (baseUrl: string): Promise<ServerType> => {
    if (forcedServerType) return forcedServerType
    const cached = serverKindByUrl.get(baseUrl)
    if (cached) return cached
    const health = await fetchHealth(baseUrl)
    if (health != null) {
      const kind = isStrataHealth(health) ? "strata" : "llama"
      serverKindByUrl.set(baseUrl, kind)
      return kind
    }
    return "llama"
  }

  const noteServerFailure = (baseUrl: string) => {
    const failures = (backoffByUrl.get(baseUrl)?.failures ?? 0) + 1
    const delay = backoffDelayMs(failures, BACKOFF_BASE_MS, BACKOFF_MAX_MS)
    const jitter = Math.floor(Math.random() * (delay / 4))
    backoffByUrl.set(baseUrl, {
      failures,
      nextRetryAt: Date.now() + Math.min(delay + jitter, BACKOFF_MAX_MS),
    })
    tracker.failure = "n/a"
    bump()
  }

  const currentSessionID = (): string | undefined => {
    const route = context.ui.router.current?.()
    if (route?.sessionID) return route.sessionID
    const loc = context.location ?? context.data.location.default()
    return loc?.sessionID
  }

  const pollMetrics = async () => {
    if (!llamaServerUrls.length) return

    let model: string | undefined
    let sessionProviderId: string | undefined
    const routeSessionID = currentSessionID()
    if (routeSessionID) {
      const session = context.data.session.get(routeSessionID)
      if (session?.model?.id) {
        model = session.model.id
      }
      sessionProviderId = session?.model?.providerID
    }

    const activeUrls = selectSessionUrls(llamaServerUrls, providerList, sessionProviderId)
    if (!activeUrls.length) {
      tracker.failure = "n/a"
      bump()
      return
    }
    for (const baseUrl of activeUrls) {
      const now = Date.now()
      const backoff = backoffByUrl.get(baseUrl)
      if (backoff && now < backoff.nextRetryAt) {
        tracker.failure = "n/a"
        bump()
        continue
      }

      if ((await resolveServerKind(baseUrl)) === "strata") {
        const live = await fetchStrataLive(baseUrl)
        if (live == null) {
          noteServerFailure(baseUrl)
          continue
        }
        backoffByUrl.delete(baseUrl)
        tracker.failure = null
        mapStrataLive(live, tracker)
        bump()
        continue
      }

      const slots = await fetchSlots(baseUrl, model)
      if (slots == null) {
        noteServerFailure(baseUrl)
        continue
      }

      backoffByUrl.delete(baseUrl)
      const slotList = Array.isArray(slots) ? slots : Object.values(slots)
      if (slotList.length === 0) {
        tracker.failure = "n/a"
        bump()
        continue
      }
      tracker.failure = null
      bump()

      processPoll(slotList, tracker)
      bump()
    }
  }

  const propsOf = (evt: any) => evt?.data ?? evt?.properties ?? {}

  const offDelta = context.data.on("message.part.delta", (evt: any) => {
    const p = propsOf(evt)
    if (p.field !== "text") return
    const partType = partTypeByID.get(p.partID)
    if (partType && partType !== "text" && partType !== "reasoning") return
    appendSample(p.sessionID, p.messageID, {
      at: Date.now(),
      tokens: estimateStreamTokens(String(p.delta ?? "")),
    })
  })

  const offMessage = context.data.on("message.updated", (evt: any) => {
    const p = propsOf(evt)
    if (p.info?.role !== "assistant") return

    if (!p.info?.time?.completed) {
      const existing = tracker.messageTimingByID[p.info.id]
      tracker.messageTimingByID[p.info.id] = {
        sessionID: p.sessionID,
        requestStartAt: p.info.time.created,
        firstResponseAt: existing?.firstResponseAt,
        firstTokenAt: existing?.firstTokenAt,
        lastTokenAt: existing?.lastTokenAt,
      }
      bump()
      return
    }

    delete tracker.messageTimingByID[p.info.id]
    pruneSamples(p.info?.time?.completed)
    bump()
  })

  const offPart = context.data.on("message.part.updated", (evt: any) => {
    const p = propsOf(evt)
    const part = p.part
    if (!part?.id) return
    if (typeof part.type === "string") partTypeByID.set(part.id, part.type)
    if (part.type !== "tool") return
    if (
      part.state?.status === "running" ||
      part.state?.status === "completed" ||
      part.state?.status === "error"
    ) {
      clearLiveSamples(p.sessionID)
    }
    const messageID = part.messageID
    if (!messageID) return
    const timing = tracker.messageTimingByID[messageID]
    if (!timing) return
    if (part.state?.status === "pending") {
      tracker.messageTimingByID[messageID] = {
        ...timing,
        firstResponseAt: timing.firstResponseAt ?? p.time,
      }
      bump()
      return
    }
    if (part.state?.status === "running") bump()
  })

  const timer = setInterval(() => {
    pruneSamples()
  }, 1000)

  const slotTimer = setInterval(() => {
    void pollMetrics()
  }, SLOT_POLL_MS)

  try {
    await context.data.location.provider.sync(location)
    providerList = context.data.location.provider.list(location) ?? []
    llamaServerUrls = resolveServerUrls(context.options, providerList, undefined)
  } catch (e) {
    console.error("[oc-ls-stats] resolveServerUrls error:", e)
    llamaServerUrls = ["http://localhost:8080"]
  }

  const unregisterSlot = context.ui.slot({
    append: "prompt.footer.status",
    render: () => <text fg={themeMuted}>{statusText()}</text>,
  })

  return () => {
    offDelta()
    offMessage()
    offPart()
    clearInterval(timer)
    clearInterval(slotTimer)
    if (typeof unregisterSlot === "function") unregisterSlot()
  }
}

const plugin = {
  id: "@troed/oc-ls-stats",
  tui,
  setup: v2Setup,
}

export default plugin