nikcli 0.0.6

package/src/session/retry.ts

@@ -0,0 +1,90 @@
+import type { NamedError } from "@nikcli-ai/util/error"
+import { MessageV2 } from "./message-v2"
+
+export namespace SessionRetry {
+  export const RETRY_INITIAL_DELAY = 2000
+  export const RETRY_BACKOFF_FACTOR = 2
+  export const RETRY_MAX_DELAY_NO_HEADERS = 30_000
+  export const RETRY_MAX_DELAY = 2_147_483_647
+
+  export async function sleep(ms: number, signal: AbortSignal): Promise<void> {
+    return new Promise((resolve, reject) => {
+      const abortHandler = () => {
+        clearTimeout(timeout)
+        reject(new DOMException("Aborted", "AbortError"))
+      }
+      const timeout = setTimeout(
+        () => {
+          signal.removeEventListener("abort", abortHandler)
+          resolve()
+        },
+        Math.min(ms, RETRY_MAX_DELAY),
+      )
+      signal.addEventListener("abort", abortHandler, { once: true })
+    })
+  }
+
+  export function delay(attempt: number, error?: MessageV2.APIError) {
+    if (error) {
+      const headers = error.data.responseHeaders
+      if (headers) {
+        const retryAfterMs = headers["retry-after-ms"]
+        if (retryAfterMs) {
+          const parsedMs = Number.parseFloat(retryAfterMs)
+          if (!Number.isNaN(parsedMs)) {
+            return parsedMs
+          }
+        }
+
+        const retryAfter = headers["retry-after"]
+        if (retryAfter) {
+          const parsedSeconds = Number.parseFloat(retryAfter)
+          if (!Number.isNaN(parsedSeconds)) {
+            return Math.ceil(parsedSeconds * 1000)
+          }
+          const parsed = Date.parse(retryAfter) - Date.now()
+          if (!Number.isNaN(parsed) && parsed > 0) {
+            return Math.ceil(parsed)
+          }
+        }
+
+        return RETRY_INITIAL_DELAY * Math.pow(RETRY_BACKOFF_FACTOR, attempt - 1)
+      }
+    }
+
+    return Math.min(RETRY_INITIAL_DELAY * Math.pow(RETRY_BACKOFF_FACTOR, attempt - 1), RETRY_MAX_DELAY_NO_HEADERS)
+  }
+
+  export function retryable(error: ReturnType<NamedError["toObject"]>) {
+    if (MessageV2.APIError.isInstance(error)) {
+      if (!error.data.isRetryable) return undefined
+      if (error.data.responseBody?.includes("FreeUsageLimitError"))
+        return `Free usage exceeded, add credits https://opencode.ai/zen`
+      return error.data.message.includes("Overloaded") ? "Provider is overloaded" : error.data.message
+    }
+
+    if (typeof error.data?.message === "string") {
+      try {
+        const json = JSON.parse(error.data.message)
+        if (json.type === "error" && json.error?.type === "too_many_requests") {
+          return "Too Many Requests"
+        }
+        if (json.code.includes("exhausted") || json.code.includes("unavailable")) {
+          return "Provider is overloaded"
+        }
+        if (json.type === "error" && json.error?.code?.includes("rate_limit")) {
+          return "Rate Limited"
+        }
+        if (
+          json.error?.message?.includes("no_kv_space") ||
+          (json.type === "error" && json.error?.type === "server_error") ||
+          !!json.error
+        ) {
+          return "Provider Server Error"
+        }
+      } catch {}
+    }
+
+    return undefined
+  }
+}

package/src/session/revert.ts

@@ -0,0 +1,120 @@
+import z from "zod"
+import { Identifier } from "../id/id"
+import { Snapshot } from "../snapshot"
+import { MessageV2 } from "./message-v2"
+import { Session } from "."
+import { Log } from "../util/log"
+import { splitWhen } from "remeda"
+import { Storage } from "../storage/storage"
+import { Bus } from "../bus"
+import { SessionPrompt } from "./prompt"
+import { SessionSummary } from "./summary"
+
+export namespace SessionRevert {
+  const log = Log.create({ service: "session.revert" })
+
+  export const RevertInput = z.object({
+    sessionID: Identifier.schema("session"),
+    messageID: Identifier.schema("message"),
+    partID: Identifier.schema("part").optional(),
+  })
+  export type RevertInput = z.infer<typeof RevertInput>
+
+  export async function revert(input: RevertInput) {
+    SessionPrompt.assertNotBusy(input.sessionID)
+    const all = await Session.messages({ sessionID: input.sessionID })
+    let lastUser: MessageV2.User | undefined
+    const session = await Session.get(input.sessionID)
+
+    let revert: Session.Info["revert"]
+    const patches: Snapshot.Patch[] = []
+    for (const msg of all) {
+      if (msg.info.role === "user") lastUser = msg.info
+      const remaining = []
+      for (const part of msg.parts) {
+        if (revert) {
+          if (part.type === "patch") {
+            patches.push(part)
+          }
+          continue
+        }
+
+        if (!revert) {
+          if ((msg.info.id === input.messageID && !input.partID) || part.id === input.partID) {
+            const partID = remaining.some((item) => ["text", "tool"].includes(item.type)) ? input.partID : undefined
+            revert = {
+              messageID: !partID && lastUser ? lastUser.id : msg.info.id,
+              partID,
+            }
+          }
+          remaining.push(part)
+        }
+      }
+    }
+
+    if (revert) {
+      const session = await Session.get(input.sessionID)
+      revert.snapshot = session.revert?.snapshot ?? (await Snapshot.track())
+      await Snapshot.revert(patches)
+      if (revert.snapshot) revert.diff = await Snapshot.diff(revert.snapshot)
+      const rangeMessages = all.filter((msg) => msg.info.id >= revert!.messageID)
+      const diffs = await SessionSummary.computeDiff({ messages: rangeMessages })
+      await Storage.write(["session_diff", input.sessionID], diffs)
+      Bus.publish(Session.Event.Diff, {
+        sessionID: input.sessionID,
+        diff: diffs,
+      })
+      return Session.update(input.sessionID, (draft) => {
+        draft.revert = revert
+        draft.summary = {
+          additions: diffs.reduce((sum, x) => sum + x.additions, 0),
+          deletions: diffs.reduce((sum, x) => sum + x.deletions, 0),
+          files: diffs.length,
+        }
+      })
+    }
+    return session
+  }
+
+  export async function unrevert(input: { sessionID: string }) {
+    log.info("unreverting", input)
+    SessionPrompt.assertNotBusy(input.sessionID)
+    const session = await Session.get(input.sessionID)
+    if (!session.revert) return session
+    if (session.revert.snapshot) await Snapshot.restore(session.revert.snapshot)
+    const next = await Session.update(input.sessionID, (draft) => {
+      draft.revert = undefined
+    })
+    return next
+  }
+
+  export async function cleanup(session: Session.Info) {
+    if (!session.revert) return
+    const sessionID = session.id
+    let msgs = await Session.messages({ sessionID })
+    const messageID = session.revert.messageID
+    const [preserve, remove] = splitWhen(msgs, (x) => x.info.id === messageID)
+    msgs = preserve
+    for (const msg of remove) {
+      await Storage.remove(["message", sessionID, msg.info.id])
+      await Bus.publish(MessageV2.Event.Removed, { sessionID: sessionID, messageID: msg.info.id })
+    }
+    const last = preserve.at(-1)
+    if (session.revert.partID && last) {
+      const partID = session.revert.partID
+      const [preserveParts, removeParts] = splitWhen(last.parts, (x) => x.id === partID)
+      last.parts = preserveParts
+      for (const part of removeParts) {
+        await Storage.remove(["part", last.info.id, part.id])
+        await Bus.publish(MessageV2.Event.PartRemoved, {
+          sessionID: sessionID,
+          messageID: last.info.id,
+          partID: part.id,
+        })
+      }
+    }
+    await Session.update(sessionID, (draft) => {
+      draft.revert = undefined
+    })
+  }
+}

package/src/session/stats.ts

@@ -0,0 +1,404 @@
+import type { Argv } from "yargs"
+
+import { Session } from "../session"
+
+import { Storage } from "../storage/storage"
+import { Project } from "../project/project"
+import { Instance } from "../project/instance"
+import { bootstrap } from "@/cli/bootstrap"
+import { cmd } from "@/cli/cmd/cmd"
+
+interface SessionStats {
+  totalSessions: number
+  totalMessages: number
+  totalCost: number
+  totalTokens: {
+    input: number
+    output: number
+    reasoning: number
+    cache: {
+      read: number
+      write: number
+    }
+  }
+  toolUsage: Record<string, number>
+  modelUsage: Record<
+    string,
+    {
+      messages: number
+      tokens: {
+        input: number
+        output: number
+      }
+      cost: number
+    }
+  >
+  dateRange: {
+    earliest: number
+    latest: number
+  }
+  days: number
+  costPerDay: number
+  tokensPerSession: number
+  medianTokensPerSession: number
+}
+
+export const StatsCommand = cmd({
+  command: "stats",
+  describe: "show token usage and cost statistics",
+  builder: (yargs: Argv) => {
+    return yargs
+      .option("days", {
+        describe: "show stats for the last N days (default: all time)",
+        type: "number",
+      })
+      .option("tools", {
+        describe: "number of tools to show (default: all)",
+        type: "number",
+      })
+      .option("models", {
+        describe: "show model statistics (default: hidden). Pass a number to show top N, otherwise shows all",
+      })
+      .option("project", {
+        describe: "filter by project (default: all projects, empty string: current project)",
+        type: "string",
+      })
+  },
+  handler: async (args) => {
+    await bootstrap(process.cwd(), async () => {
+      const stats = await aggregateSessionStats(args.days, args.project)
+
+      let modelLimit: number | undefined
+      if (args.models === true) {
+        modelLimit = Infinity
+      } else if (typeof args.models === "number") {
+        modelLimit = args.models
+      }
+
+      displayStats(stats, args.tools, modelLimit)
+    })
+  },
+})
+
+async function getCurrentProject(): Promise<Project.Info> {
+  return Instance.project
+}
+
+async function getAllSessions(): Promise<Session.Info[]> {
+  const sessions: Session.Info[] = []
+
+  const projectKeys = await Storage.list(["project"])
+  const projects = await Promise.all(projectKeys.map((key) => Storage.read<Project.Info>(key)))
+
+  for (const project of projects) {
+    if (!project) continue
+
+    const sessionKeys = await Storage.list(["session", project.id])
+    const projectSessions = await Promise.all(sessionKeys.map((key) => Storage.read<Session.Info>(key)))
+
+    for (const session of projectSessions) {
+      if (session) {
+        sessions.push(session)
+      }
+    }
+  }
+
+  return sessions
+}
+
+export async function aggregateSessionStats(days?: number, projectFilter?: string): Promise<SessionStats> {
+  const sessions = await getAllSessions()
+  const MS_IN_DAY = 24 * 60 * 60 * 1000
+
+  const cutoffTime = (() => {
+    if (days === undefined) return 0
+    if (days === 0) {
+      const now = new Date()
+      now.setHours(0, 0, 0, 0)
+      return now.getTime()
+    }
+    return Date.now() - days * MS_IN_DAY
+  })()
+
+  const windowDays = (() => {
+    if (days === undefined) return
+    if (days === 0) return 1
+    return days
+  })()
+
+  let filteredSessions = cutoffTime > 0 ? sessions.filter((session) => session.time.updated >= cutoffTime) : sessions
+
+  if (projectFilter !== undefined) {
+    if (projectFilter === "") {
+      const currentProject = await getCurrentProject()
+      filteredSessions = filteredSessions.filter((session) => session.projectID === currentProject.id)
+    } else {
+      filteredSessions = filteredSessions.filter((session) => session.projectID === projectFilter)
+    }
+  }
+
+  const stats: SessionStats = {
+    totalSessions: filteredSessions.length,
+    totalMessages: 0,
+    totalCost: 0,
+    totalTokens: {
+      input: 0,
+      output: 0,
+      reasoning: 0,
+      cache: {
+        read: 0,
+        write: 0,
+      },
+    },
+    toolUsage: {},
+    modelUsage: {},
+    dateRange: {
+      earliest: Date.now(),
+      latest: Date.now(),
+    },
+    days: 0,
+    costPerDay: 0,
+    tokensPerSession: 0,
+    medianTokensPerSession: 0,
+  }
+
+  if (filteredSessions.length > 1000) {
+    console.log(`Large dataset detected (${filteredSessions.length} sessions). This may take a while...`)
+  }
+
+  if (filteredSessions.length === 0) {
+    stats.days = windowDays ?? 0
+    return stats
+  }
+
+  let earliestTime = Date.now()
+  let latestTime = 0
+
+  const sessionTotalTokens: number[] = []
+
+  const BATCH_SIZE = 20
+  for (let i = 0; i < filteredSessions.length; i += BATCH_SIZE) {
+    const batch = filteredSessions.slice(i, i + BATCH_SIZE)
+
+    const batchPromises = batch.map(async (session) => {
+      const messages = await Session.messages({ sessionID: session.id })
+
+      let sessionCost = 0
+      let sessionTokens = { input: 0, output: 0, reasoning: 0, cache: { read: 0, write: 0 } }
+      let sessionToolUsage: Record<string, number> = {}
+      let sessionModelUsage: Record<
+        string,
+        {
+          messages: number
+          tokens: {
+            input: number
+            output: number
+          }
+          cost: number
+        }
+      > = {}
+
+      for (const message of messages) {
+        if (message.info.role === "assistant") {
+          sessionCost += message.info.cost || 0
+
+          const modelKey = `${message.info.providerID}/${message.info.modelID}`
+          if (!sessionModelUsage[modelKey]) {
+            sessionModelUsage[modelKey] = {
+              messages: 0,
+              tokens: { input: 0, output: 0 },
+              cost: 0,
+            }
+          }
+          sessionModelUsage[modelKey].messages++
+          sessionModelUsage[modelKey].cost += message.info.cost || 0
+
+          if (message.info.tokens) {
+            sessionTokens.input += message.info.tokens.input || 0
+            sessionTokens.output += message.info.tokens.output || 0
+            sessionTokens.reasoning += message.info.tokens.reasoning || 0
+            sessionTokens.cache.read += message.info.tokens.cache?.read || 0
+            sessionTokens.cache.write += message.info.tokens.cache?.write || 0
+
+            sessionModelUsage[modelKey].tokens.input += message.info.tokens.input || 0
+            sessionModelUsage[modelKey].tokens.output +=
+              (message.info.tokens.output || 0) + (message.info.tokens.reasoning || 0)
+          }
+        }
+
+        for (const part of message.parts) {
+          if (part.type === "tool" && part.tool) {
+            sessionToolUsage[part.tool] = (sessionToolUsage[part.tool] || 0) + 1
+          }
+        }
+      }
+
+      return {
+        messageCount: messages.length,
+        sessionCost,
+        sessionTokens,
+        sessionTotalTokens: sessionTokens.input + sessionTokens.output + sessionTokens.reasoning,
+        sessionToolUsage,
+        sessionModelUsage,
+        earliestTime: cutoffTime > 0 ? session.time.updated : session.time.created,
+        latestTime: session.time.updated,
+      }
+    })
+
+    const batchResults = await Promise.all(batchPromises)
+
+    for (const result of batchResults) {
+      earliestTime = Math.min(earliestTime, result.earliestTime)
+      latestTime = Math.max(latestTime, result.latestTime)
+      sessionTotalTokens.push(result.sessionTotalTokens)
+
+      stats.totalMessages += result.messageCount
+      stats.totalCost += result.sessionCost
+      stats.totalTokens.input += result.sessionTokens.input
+      stats.totalTokens.output += result.sessionTokens.output
+      stats.totalTokens.reasoning += result.sessionTokens.reasoning
+      stats.totalTokens.cache.read += result.sessionTokens.cache.read
+      stats.totalTokens.cache.write += result.sessionTokens.cache.write
+
+      for (const [tool, count] of Object.entries(result.sessionToolUsage)) {
+        stats.toolUsage[tool] = (stats.toolUsage[tool] || 0) + count
+      }
+
+      for (const [model, usage] of Object.entries(result.sessionModelUsage)) {
+        if (!stats.modelUsage[model]) {
+          stats.modelUsage[model] = {
+            messages: 0,
+            tokens: { input: 0, output: 0 },
+            cost: 0,
+          }
+        }
+        stats.modelUsage[model].messages += usage.messages
+        stats.modelUsage[model].tokens.input += usage.tokens.input
+        stats.modelUsage[model].tokens.output += usage.tokens.output
+        stats.modelUsage[model].cost += usage.cost
+      }
+    }
+  }
+
+  const rangeDays = Math.max(1, Math.ceil((latestTime - earliestTime) / MS_IN_DAY))
+  const effectiveDays = windowDays ?? rangeDays
+  stats.dateRange = {
+    earliest: earliestTime,
+    latest: latestTime,
+  }
+  stats.days = effectiveDays
+  stats.costPerDay = stats.totalCost / effectiveDays
+  const totalTokens = stats.totalTokens.input + stats.totalTokens.output + stats.totalTokens.reasoning
+  stats.tokensPerSession = filteredSessions.length > 0 ? totalTokens / filteredSessions.length : 0
+  sessionTotalTokens.sort((a, b) => a - b)
+  const mid = Math.floor(sessionTotalTokens.length / 2)
+  stats.medianTokensPerSession =
+    sessionTotalTokens.length === 0
+      ? 0
+      : sessionTotalTokens.length % 2 === 0
+        ? (sessionTotalTokens[mid - 1] + sessionTotalTokens[mid]) / 2
+        : sessionTotalTokens[mid]
+
+  return stats
+}
+
+export function displayStats(stats: SessionStats, toolLimit?: number, modelLimit?: number) {
+  const width = 56
+
+  function renderRow(label: string, value: string): string {
+    const availableWidth = width - 1
+    const paddingNeeded = availableWidth - label.length - value.length
+    const padding = Math.max(0, paddingNeeded)
+    return `│${label}${" ".repeat(padding)}${value} │`
+  }
+
+  // Overview section
+  console.log("┌────────────────────────────────────────────────────────┐")
+  console.log("│                       OVERVIEW                         │")
+  console.log("├────────────────────────────────────────────────────────┤")
+  console.log(renderRow("Sessions", stats.totalSessions.toLocaleString()))
+  console.log(renderRow("Messages", stats.totalMessages.toLocaleString()))
+  console.log(renderRow("Days", stats.days.toString()))
+  console.log("└────────────────────────────────────────────────────────┘")
+  console.log()
+
+  // Cost & Tokens section
+  console.log("┌────────────────────────────────────────────────────────┐")
+  console.log("│                    COST & TOKENS                       │")
+  console.log("├────────────────────────────────────────────────────────┤")
+  const cost = isNaN(stats.totalCost) ? 0 : stats.totalCost
+  const costPerDay = isNaN(stats.costPerDay) ? 0 : stats.costPerDay
+  const tokensPerSession = isNaN(stats.tokensPerSession) ? 0 : stats.tokensPerSession
+  console.log(renderRow("Total Cost", `$${cost.toFixed(2)}`))
+  console.log(renderRow("Avg Cost/Day", `$${costPerDay.toFixed(2)}`))
+  console.log(renderRow("Avg Tokens/Session", formatNumber(Math.round(tokensPerSession))))
+  const medianTokensPerSession = isNaN(stats.medianTokensPerSession) ? 0 : stats.medianTokensPerSession
+  console.log(renderRow("Median Tokens/Session", formatNumber(Math.round(medianTokensPerSession))))
+  console.log(renderRow("Input", formatNumber(stats.totalTokens.input)))
+  console.log(renderRow("Output", formatNumber(stats.totalTokens.output)))
+  console.log(renderRow("Cache Read", formatNumber(stats.totalTokens.cache.read)))
+  console.log(renderRow("Cache Write", formatNumber(stats.totalTokens.cache.write)))
+  console.log("└────────────────────────────────────────────────────────┘")
+  console.log()
+
+  // Model Usage section
+  if (modelLimit !== undefined && Object.keys(stats.modelUsage).length > 0) {
+    const sortedModels = Object.entries(stats.modelUsage).sort(([, a], [, b]) => b.messages - a.messages)
+    const modelsToDisplay = modelLimit === Infinity ? sortedModels : sortedModels.slice(0, modelLimit)
+
+    console.log("┌────────────────────────────────────────────────────────┐")
+    console.log("│                      MODEL USAGE                       │")
+    console.log("├────────────────────────────────────────────────────────┤")
+
+    for (const [model, usage] of modelsToDisplay) {
+      console.log(`│ ${model.padEnd(54)} │`)
+      console.log(renderRow("  Messages", usage.messages.toLocaleString()))
+      console.log(renderRow("  Input Tokens", formatNumber(usage.tokens.input)))
+      console.log(renderRow("  Output Tokens", formatNumber(usage.tokens.output)))
+      console.log(renderRow("  Cost", `$${usage.cost.toFixed(4)}`))
+      console.log("├────────────────────────────────────────────────────────┤")
+    }
+    // Remove last separator and add bottom border
+    process.stdout.write("\x1B[1A") // Move up one line
+    console.log("└────────────────────────────────────────────────────────┘")
+  }
+  console.log()
+
+  // Tool Usage section
+  if (Object.keys(stats.toolUsage).length > 0) {
+    const sortedTools = Object.entries(stats.toolUsage).sort(([, a], [, b]) => b - a)
+    const toolsToDisplay = toolLimit ? sortedTools.slice(0, toolLimit) : sortedTools
+
+    console.log("┌────────────────────────────────────────────────────────┐")
+    console.log("│                      TOOL USAGE                        │")
+    console.log("├────────────────────────────────────────────────────────┤")
+
+    const maxCount = Math.max(...toolsToDisplay.map(([, count]) => count))
+    const totalToolUsage = Object.values(stats.toolUsage).reduce((a, b) => a + b, 0)
+
+    for (const [tool, count] of toolsToDisplay) {
+      const barLength = Math.max(1, Math.floor((count / maxCount) * 20))
+      const bar = "█".repeat(barLength)
+      const percentage = ((count / totalToolUsage) * 100).toFixed(1)
+
+      const maxToolLength = 18
+      const truncatedTool = tool.length > maxToolLength ? tool.substring(0, maxToolLength - 2) + ".." : tool
+      const toolName = truncatedTool.padEnd(maxToolLength)
+
+      const content = ` ${toolName} ${bar.padEnd(20)} ${count.toString().padStart(3)} (${percentage.padStart(4)}%)`
+      const padding = Math.max(0, width - content.length - 1)
+      console.log(`│${content}${" ".repeat(padding)} │`)
+    }
+    console.log("└────────────────────────────────────────────────────────┘")
+  }
+  console.log()
+}
+
+function formatNumber(num: number): string {
+  if (num >= 1000000) {
+    return (num / 1000000).toFixed(1) + "M"
+  } else if (num >= 1000) {
+    return (num / 1000).toFixed(1) + "K"
+  }
+  return num.toString()
+}

package/src/session/status.ts

@@ -0,0 +1,84 @@
+import { BusEvent } from "@/bus/bus-event"
+import { Bus } from "@/bus"
+import { Instance } from "@/project/instance"
+import z from "zod"
+
+export namespace SessionStatus {
+  export const Info = z
+    .union([
+      z.object({
+        type: z.literal("idle"),
+      }),
+      z.object({
+        type: z.literal("retry"),
+        attempt: z.number(),
+        message: z.string(),
+        next: z.number(),
+      }),
+      z.object({
+        type: z.literal("busy"),
+      }),
+    ])
+    .meta({
+      ref: "SessionStatus",
+    })
+  export type Info = z.infer<typeof Info>
+
+  export const Event = {
+    Status: BusEvent.define(
+      "session.status",
+      z.object({
+        sessionID: z.string(),
+        status: Info,
+      }),
+    ),
+    // deprecated
+    Idle: BusEvent.define(
+      "session.idle",
+      z.object({
+        sessionID: z.string(),
+      }),
+    ),
+  }
+
+  const state = Instance.state(() => {
+    const data: Record<string, Info> = {}
+    return data
+  })
+
+  export function get(sessionID: string) {
+    return (
+      state()[sessionID] ?? {
+        type: "idle",
+      }
+    )
+  }
+
+  export function list() {
+    return state()
+  }
+
+  export function set(sessionID: string, status: Info) {
+    Bus.publish(Event.Status, {
+      sessionID,
+      status,
+    })
+    if (status.type === "idle") {
+      // deprecated
+      Bus.publish(Event.Idle, {
+        sessionID,
+      })
+      delete state()[sessionID]
+      return
+    }
+    state()[sessionID] = status
+  }
+
+  export function hydrate(sessionID: string, status: Info) {
+    if (status.type === "idle") {
+      delete state()[sessionID]
+      return
+    }
+    state()[sessionID] = status
+  }
+}