diff --git a/apps/cli/src/bootstrap/extensions.ts b/apps/cli/src/bootstrap/extensions.ts index 92c33f41..33e12c13 100644 --- a/apps/cli/src/bootstrap/extensions.ts +++ b/apps/cli/src/bootstrap/extensions.ts @@ -28,6 +28,7 @@ export interface StepExtensionFactoryDeps { feedbackIdentity: StepExtensionOptions["feedbackIdentity"]; /** Initial Step permission policy parsed from CLI/runtime options. */ permission: StepExtensionOptions["permission"]; + usageHost?: StepExtensionOptions["usageHost"]; /** Optional StepCode provider extension, present only when configured. */ stepCodeProviderExtension: InlineExtension | undefined; } @@ -45,6 +46,7 @@ export function createStepExtensionFactories(deps: StepExtensionFactoryDeps): In stepSettings: deps.stepSettings, feedbackIdentity: deps.feedbackIdentity, permission: deps.permission, + usageHost: deps.usageHost, traceHeaderPolicy: deps.traceHeaderPolicy, }), createStepCapabilitiesExtensionInline({ telemetry: deps.telemetry }), diff --git a/apps/cli/src/main.ts b/apps/cli/src/main.ts index 27fe6b34..f692a597 100644 --- a/apps/cli/src/main.ts +++ b/apps/cli/src/main.ts @@ -13,6 +13,7 @@ import { ambiguousPluginServerNames, applyStepCodeConfigDefaults, buildStepSystemPromptAppendix, + checkStepPlanBeforeInteractive, configureHttpDispatcher, createStepCode, createStepCodeProviderInlineExtension, @@ -43,6 +44,7 @@ import { migrateLegacyStepCredential, needsStepLoginBeforeInteractive, normalizeStepSessionSelectorArgs, + PlanUsageController, parseStepUpdateCommand, prepareMain, prepareStepInstallation, @@ -50,6 +52,7 @@ import { readGlobalStepConfig, readOrCreateStepDeviceId, readStoredCredential, + resolvePlanUsageAccount, resolveStepAgentDir, resolveStepConfigDir, resolveStepConfigRoot, @@ -67,6 +70,7 @@ import { STEP_PROVIDER_ID, STEPCODE_VERSION, type StepSettingsManager, + scanUsageOverview, setStderrDevLogStorageRootDirectory, stopThemeWatcher, syncStepLoginProfileEndpoint, @@ -81,7 +85,15 @@ import { createStepExtensionFactories } from "#bootstrap/extensions"; import { captureRawStdout, sdkStdioRequested } from "#bootstrap/stdout-capture"; import { InteractiveMode, runPrintMode, runRpcMode } from "#modes/index"; import { createSdkStdioMode } from "#modes/sdk-stdio"; -import { selectConfig, selectSession, showFirstTimeSetup, showStartupInput, showStartupSelector } from "#ui/index"; +import { + selectConfig, + selectSession, + showFirstTimeSetup, + showStartupInput, + showStartupSelector, + showUsagePanel, + usagePresentationRows, +} from "#ui/index"; import { observability } from "./observability.ts"; /** @@ -159,6 +171,14 @@ const telemetryIdentityReady = initializeTelemetryIdentity(); configureHttpDispatcher(); let currentStepSettingsManager: StepSettingsManager | undefined; +let usageRuntime: Parameters[0] | undefined; +const usageAccount = new PlanUsageController(() => + resolvePlanUsageAccount( + getStepAuthPath(), + process.env, + process.argv.some((arg) => arg === "--api-key" || arg.startsWith("--api-key=")), + ), +); const hasExplicitStepCredential = process.argv.slice(2).some((arg) => arg === "--api-key" || arg.startsWith("--api-key=")) || Boolean(process.env.STEP_API_KEY?.trim()); @@ -187,6 +207,34 @@ const stepMainOptions: MainOptions = { extensionFactories: createStepExtensionFactories({ telemetry: telemetryReporter, stepSettings: () => currentStepSettingsManager, + usageHost: { + account: usageAccount, + resources: () => { + const loader = usageRuntime?.session.resourceLoader; + return { + agents: loader?.getAgentsFiles().agentsFiles.map((file) => file.path) ?? [], + skills: + loader?.getSkills().skills.map((skill) => ({ name: skill.name, filePath: skill.filePath })) ?? [], + }; + }, + scan: (cwd, scope, signal) => + scanUsageOverview( + [ + join(resolveStepAgentDir(), "sessions"), + ...(usageRuntime ? [usageRuntime.session.sessionManager.getSessionDir()] : []), + ], + join(resolveStepStorageRoot(), "cache", "usage"), + cwd, + scope, + signal, + ), + present: showUsagePanel, + output: (ctx, model) => { + const text = usagePresentationRows(model).join("\n"); + if (ctx.mode === "rpc") ctx.ui.notify(text, "info"); + else process.stderr.write(`${text}\n`); + }, + }, feedbackIdentity: readCurrentFeedbackIdentity, permission: { approvalMode: stepPermissionArgs.approvalMode, @@ -208,7 +256,10 @@ const stepMainOptions: MainOptions = { defaultTheme: getStepDefaultTheme(), defaultProvider: stepCodeConfig?.defaultProvider ?? stepPersistedDefaults.provider ?? STEP_DEFAULT_PROVIDER, defaultModel: stepCodeConfig?.defaultModel ?? stepPersistedDefaults.model ?? STEP_DEFAULT_MODEL, - runtimeHostFactory: createStepCode, + runtimeHostFactory: (runtime) => { + usageRuntime = runtime; + return createStepCode(runtime); + }, stdioModeFactory: sdkStdio ? createSdkStdioMode({ writeFrame: rawStdoutWrite }) : undefined, interactiveModeOptions: { showChangelog: false, @@ -218,6 +269,7 @@ const stepMainOptions: MainOptions = { stepLogin: (host) => runStepLogin({ authPath: getStepAuthPath(), + planUsage: usageAccount, createHost: () => host, themeName: getStepDefaultTheme(), }), @@ -572,6 +624,7 @@ try { } else { const outcome = await runStepLogin({ authPath: getStepAuthPath(), + planUsage: usageAccount, themeName: getStepDefaultTheme(), noBrowser: loginArgs.includes("--no-browser"), }); @@ -579,8 +632,9 @@ try { syncStepLoginProfileEndpoint(getStepAuthPath()); process.stdout.write(`Signed in${outcome.profile ? ` with ${outcome.profile.title}` : ""}.\n`); if (outcome.credentialsPath) process.stdout.write(`Credential written: ${outcome.credentialsPath}\n`); + if (outcome.warning) process.stderr.write(`${outcome.warning}\n`); } else { - process.stderr.write("Sign-in cancelled. No credential was written.\n"); + process.stderr.write("Sign-in cancelled.\n"); process.exitCode = 1; } } @@ -643,21 +697,44 @@ try { stdoutIsTTY: process.stdout.isTTY, args: parsedInteractiveArgs, }); + const selectedProvider = + parsedInteractiveArgs.provider ?? + (parsedInteractiveArgs.model?.includes("/") + ? parsedInteractiveArgs.model.split("/", 1)[0] + : undefined) ?? + process.env.STEP_PROVIDER ?? + process.env.STEP_MODEL_PROVIDER ?? + stepCodeConfig?.defaultProvider ?? + stepPersistedDefaults.provider ?? + STEP_DEFAULT_PROVIDER; if ( - !hasConfiguredStepCodeCredential(stepCodeConfig) && - needsStepLoginBeforeInteractive({ - authPath: getStepAuthPath(), - interactive: interactiveStartup, - }) + interactiveStartup && + selectedProvider === STEP_PROVIDER_ID && + !hasExplicitStepCredential && + !hasConfiguredStepCodeCredential(stepCodeConfig) ) { - const outcome = await runStepLogin({ + const missingCredential = needsStepLoginBeforeInteractive({ authPath: getStepAuthPath(), - themeName: getStepDefaultTheme(), + interactive: interactiveStartup, }); - if (outcome.kind === "exit") { - shouldLaunchMain = false; - } else { - syncStepLoginProfileEndpoint(getStepAuthPath()); + const check = missingCredential + ? {} + : await checkStepPlanBeforeInteractive({ + authPath: getStepAuthPath(), + interactive: interactiveStartup, + planUsage: usageAccount, + }); + if (check.warning) process.stderr.write(`${check.warning}\n`); + if (missingCredential || check.requiredPlan) { + const outcome = await runStepLogin({ + authPath: getStepAuthPath(), + planUsage: usageAccount, + requiredPlan: check.requiredPlan, + themeName: getStepDefaultTheme(), + }); + if (outcome.warning) process.stderr.write(`${outcome.warning}\n`); + if (outcome.kind === "exit") shouldLaunchMain = false; + else syncStepLoginProfileEndpoint(getStepAuthPath()); } } if (shouldLaunchMain) { diff --git a/apps/cli/src/ui/index.ts b/apps/cli/src/ui/index.ts index 72c01243..1fa252d5 100644 --- a/apps/cli/src/ui/index.ts +++ b/apps/cli/src/ui/index.ts @@ -25,3 +25,5 @@ export { showStartupSelector, startStartupTui, } from "./startup-ui.ts"; + +export { showUsagePanel, usagePresentationRows } from "./view/dialogs/usage-panel.ts"; diff --git a/apps/cli/src/ui/interactive-mode.ts b/apps/cli/src/ui/interactive-mode.ts index 15520bef..9b8fd289 100644 --- a/apps/cli/src/ui/interactive-mode.ts +++ b/apps/cli/src/ui/interactive-mode.ts @@ -3307,11 +3307,12 @@ export class InteractiveMode { return new Promise((resolve, reject) => { let component: Component & { dispose?(): void }; let closed = false; + let ownOverlay: OverlayHandle | undefined; const close = (result: T) => { if (closed) return; closed = true; - if (isOverlay) this.ui.hideOverlay(); + if (isOverlay) ownOverlay?.hide(); else restoreEditor(); // Note: both branches above already call requestRender resolve(result); @@ -3324,7 +3325,14 @@ export class InteractiveMode { Promise.resolve(factory(this.ui, theme, this.keybindings, close)) .then((c) => { - if (closed) return; + if (closed) { + try { + c.dispose?.(); + } catch { + /* ignore dispose errors */ + } + return; + } component = c; if (isOverlay) { // Resolve overlay options - can be static or dynamic function @@ -3341,6 +3349,7 @@ export class InteractiveMode { return w ? { width: w } : undefined; }; const handle = this.ui.showOverlay(component, resolveOptions()); + ownOverlay = handle; // Expose handle to caller for visibility control options?.onHandle?.(handle); } else { @@ -5630,8 +5639,9 @@ export class InteractiveMode { const outcome = await this.options.stepLogin(this.createStepLoginHost()); if (outcome.kind === "completed") { await this.handleReloadCommand(); + if (outcome.warning) this.showWarning(outcome.warning); } else { - this.showStatus("Sign-in cancelled. No credential was written."); + this.showStatus("Sign-in cancelled."); } } catch (error: unknown) { this.showError(`Failed to login to Step: ${error instanceof Error ? error.message : String(error)}`); diff --git a/apps/cli/src/ui/view/dialogs/usage-panel.ts b/apps/cli/src/ui/view/dialogs/usage-panel.ts new file mode 100644 index 00000000..68d39682 --- /dev/null +++ b/apps/cli/src/ui/view/dialogs/usage-panel.ts @@ -0,0 +1,325 @@ +import { + type AccountUsageState, + type ExtensionCommandContext, + stripAnsi, + theme, + type UsagePresentation, +} from "@step-harness/coding-agent"; +import { type Component, Key, matchesKey, truncateToWidth, wrapTextWithAnsi } from "@step-harness/pi-tui"; +import { renderStepDialogFrame } from "./step-dialog.ts"; + +const count = (n: number | bigint): string => new Intl.NumberFormat("en-US").format(n); +const dateTime = (ms: number | null): string => + ms === null + ? "Unknown" + : new Intl.DateTimeFormat("en-US", { + month: "short", + day: "numeric", + hour: "2-digit", + minute: "2-digit", + timeZoneName: "shortOffset", + }).format(ms); +const field = (name: string, value: string): string => `${name.padEnd(22)} ${value}`; +const safe = (text: string): string => stripAnsi(text).replace(/[\u0000-\u001f\u007f-\u009f\u2028\u2029]/gu, ""); + +function remainingQuota(usedPercent: number | null): string { + if (usedPercent === null) return "Remaining percentage unavailable"; + const remaining = 100 - usedPercent, + filled = Math.round((remaining / 100) * 12); + return `[${theme.fg("accent", "▆".repeat(filled))}${theme.fg("dim", "▁".repeat(12 - filled))}] ${remaining.toFixed(1)}% left`; +} + +function accountRows(state: AccountUsageState, cached = false): string[] { + const snapshot = state.snapshot, + now = Date.now(); + const lines = [state.status === "loading" ? "Refreshing account…" : ""]; + if (state.message) lines.push(state.message); + if (state.retryAt && state.retryAt > now) lines.push(`Retry in ${Math.ceil((state.retryAt - now) / 1000)}s`); + if (!snapshot) { + if (!state.message) lines.push(cached ? "Not cached · open /usage account" : "Account quota not cached."); + return lines.filter(Boolean); + } + lines.push( + `Plan · ${snapshot.plan.name ?? (snapshot.plan.state === "unavailable" ? "Name unavailable" : snapshot.plan.family.toUpperCase())}`, + ); + lines.push( + `${state.status === "stale" ? "Stale · " : cached ? "Cached · " : ""}Updated ${dateTime(snapshot.observed_at_ms)}`, + ); + if (snapshot.plan.subscription?.expires_at_ms) + lines.push(field("Plan expires:", dateTime(snapshot.plan.subscription.expires_at_ms))); + if (snapshot.plan.family === "none") lines.push("No personal plan or credit buckets"); + for (const window of snapshot.quota.windows) { + const name = window.id === "five_hour" ? "5-hour" : "Weekly"; + const expired = + window.resets_at_ms !== null && window.resets_at_ms <= now && (state.receivedAt ?? 0) < window.resets_at_ms; + if (expired) { + lines.push(field(`${name} limit:`, "Snapshot expired · open /usage account")); + continue; + } + if (!window.limited) { + lines.push(field(`${name} limit:`, "Unlimited")); + continue; + } + lines.push(field(`${name} limit:`, remainingQuota(window.used_percent))); + lines.push(field("", `Resets ${dateTime(window.resets_at_ms)}`)); + } + if (snapshot.quota.credit_buckets === null) lines.push("Credit bucket details temporarily unavailable"); + for (const [index, bucket] of (snapshot.quota.credit_buckets ?? []).entries()) { + if (bucket.unsupported) { + lines.push(`Unsupported bucket type: ${safe(bucket.type)}`); + continue; + } + if ( + (bucket.expires_at_ms !== null && bucket.expires_at_ms <= now) || + (bucket.resets_at_ms !== null && bucket.resets_at_ms <= now && (state.receivedAt ?? 0) < bucket.resets_at_ms) + ) { + lines.push(`Bucket ${index + 1}: snapshot expired · refresh account`); + continue; + } + lines.push( + field( + bucket.type === "subscription" ? "Subscription Credits:" : `Top-up ${index + 1}:`, + remainingQuota(bucket.used_percent), + ), + ); + if (bucket.resets_at_ms !== null) lines.push(field("", `Resets ${dateTime(bucket.resets_at_ms)}`)); + lines.push(field("", `Expires ${dateTime(bucket.expires_at_ms)}`)); + } + if (!cached) { + lines.push( + snapshot.quota.allowed + ? "Quota check: available" + : `Quota check: restricted${snapshot.quota.reason ? ` (${snapshot.quota.reason})` : ""}`, + ); + lines.push( + "Account usage includes other projects and devices.", + "Quota availability does not guarantee model access.", + ); + } + return lines.filter(Boolean); +} + +function heatmap(model: UsagePresentation, width: number): string[] { + const overview = model.overview; + if (!overview) return [model.overviewError ?? "Scanning local activity…"]; + const days = overview.days, + prefix = new Date(`${days[0].date}T12:00:00`).getDay(); + const weeks: (typeof days)[] = []; + let week: typeof days = []; + for (let i = 0; i < prefix; i++) week.push({ date: "", tokens: 0, chats: 0 }); + for (const day of days) { + week.push(day); + if (week.length === 7) { + weeks.push(week); + week = []; + } + } + if (week.length) weeks.push(week); + const visible = Math.max(1, Math.min(weeks.length, Math.floor((width - 8) / 2))), + shown = weeks.slice(-visible), + max = Math.max(1, ...days.map((day) => day.tokens)); + const rows = [`365 days · ${overview.scope === "all" ? "All projects" : "Current project"}`]; + const monthRow = Array(visible * 2).fill(" "); + let previousMonth = ""; + for (const [index, w] of shown.entries()) { + const d = w.find((day) => day.date); + if (!d) continue; + const month = d.date.slice(0, 7); + if (month !== previousMonth) { + const label = new Intl.DateTimeFormat("en-US", { month: "short" }).format(new Date(`${d.date}T12:00:00`)); + for (let j = 0; j < label.length && index * 2 + j < monthRow.length; j++) monthRow[index * 2 + j] = label[j]; + previousMonth = month; + } + } + rows.push(` ${monthRow.join("")}`); + for (let row = 0; row < 7; row++) + rows.push( + `${["Su", "Mo", "Tu", "We", "Th", "Fr", "Sa"][row]} ${shown + .map((w) => { + const d = w[row]; + if (!d?.date) return " "; + if (d.date === model.selectedDate) return theme.fg("accent", "▣ "); + return d.tokens === 0 + ? theme.fg("dim", "· ") + : theme.fg("accent", `${["░", "▒", "▓", "█"][Math.min(3, Math.ceil((d.tokens / max) * 4) - 1)]} `); + }) + .join("")}`, + ); + rows.push( + ` Less ${theme.fg("dim", "·")} ${theme.fg("accent", "░ ▒ ▓ █")} More · ${theme.fg("accent", "▣")} Selected day`, + ); + if (visible < weeks.length) rows.push(`Showing the latest ${visible} weeks; totals cover the full year.`); + const selected = days.find((d) => d.date === model.selectedDate) ?? days.at(-1)!; + rows.push( + `${selected.date} · ${count(selected.tokens)} tokens · ${selected.chats} chats`, + "", + "Activity insights", + field("Total chats", count(overview.chats)), + field("Active days", count(overview.activeDays)), + field("Total tokens", count(overview.tokens)), + ); + const knownReasoning = overview.reasoning.filter(([level]) => level !== "unknown"), + calls = knownReasoning.reduce((n, [, total]) => n + total, 0); + rows.push( + field( + "Most used reasoning", + knownReasoning.length + ? `${knownReasoning[0][0]} · ${((knownReasoning[0][1] / calls) * 100).toFixed(0)}% of recorded calls` + : "Not recorded", + ), + field( + "Skills discovered", + overview.discoveredSkills === null ? "Not recorded" : count(overview.discoveredSkills), + ), + field("Total skills used", overview.usedSkills === null ? "Not recorded" : count(overview.usedSkills)), + "", + "Most used plugins and skills", + ); + rows.push( + ...(overview.topResources.length + ? overview.topResources.map(([name, runs]) => field(safe(name), `${runs} runs`)) + : ["Not recorded"]), + ); + if (overview.limited.length) rows.push("", `Coverage: ${overview.limited.join(" · ")}`); + rows.push("", "Local activity · fork/clone copies deduplicated", "P project scope · [ / ] select day"); + return rows; +} + +export function usagePresentationRows(model: UsagePresentation, width = 100, details = false): string[] { + const s = model.status; + if (model.view === "status") { + const rows = [ + `Step Code (${s.version})`, + "", + field("Model:", `${s.model} (reasoning ${s.reasoning})`), + field("Model provider:", s.provider), + field("Directory:", s.directory), + field("Permissions:", `${s.permissions} · Project ${s.trusted ? "trusted" : "untrusted"}`), + field("Agents.md:", s.agents.length ? s.agents.join("\n ") : "None"), + field( + "Account:", + `${s.account.accountId ?? (s.account.presence === "present" ? "Identity unknown" : "Not signed in")}${s.account.region ? ` (${s.account.region})` : ""}`, + ), + field("Collaboration mode:", s.collaboration), + field("Session:", s.session), + "", + ...accountRows(s.cached, true), + ]; + if (details) + rows.push( + "", + "Details", + field("State:", s.state), + field("Context:", s.context), + field("Model endpoint:", s.endpoint ?? "Not configured"), + field("Credential:", `${s.account.presence} · ${s.account.source} · validity not checked`), + field("Config source:", s.account.profile ?? "Unknown"), + ); + rows.push("", "Local snapshot · no network request", "D details · Esc close"); + return rows; + } + const rows = ["Usage", model.view === "overview" ? "[1 Overview] 2 Account" : " 1 Overview [2 Account]", ""]; + if (model.view === "overview") rows.push(...heatmap(model, width)); + else { + rows.push( + ...accountRows(model.account), + "", + "Current main session · local", + "Includes restored/inherited history · subagents excluded", + field("Input:", count(model.session.input)), + field("Cache read:", count(model.session.cacheRead)), + field( + "Cache hit:", + model.session.cacheHit === null ? "Unknown" : `${(model.session.cacheHit * 100).toFixed(1)}%`, + ), + field("Output:", count(model.session.output)), + field("Total:", `${count(model.session.total)} tokens`), + ); + if (model.session.invalid || model.session.missing) + rows.push(`Usage coverage: ${model.session.invalid} invalid, ${model.session.missing} missing records`); + if (s.provider !== "step") rows.push(`Step account quota · current model uses ${s.provider}`); + rows.push("", "R refresh account"); + } + rows.push("1/2 or ←/→ switch tab · ↑/↓ scroll · Esc close"); + return rows; +} + +export class UsagePanel implements Component { + private readonly model: UsagePresentation; + private readonly redraw: () => void; + private readonly close: () => void; + private readonly rows: () => number; + private offset = 0; + private details = false; + constructor(model: UsagePresentation, redraw: () => void, close: () => void, rows: () => number) { + this.model = model; + this.redraw = redraw; + this.close = close; + this.rows = rows; + } + invalidate(): void {} + handleInput(data: string): void { + if (matchesKey(data, Key.escape)) { + this.close(); + return; + } + if (matchesKey(data, Key.up) || matchesKey(data, Key.pageUp)) + this.offset = Math.max(0, this.offset - (matchesKey(data, Key.pageUp) ? 10 : 1)); + else if (matchesKey(data, Key.down) || matchesKey(data, Key.pageDown)) + this.offset += matchesKey(data, Key.pageDown) ? 10 : 1; + else if (this.model.view === "status") { + if (data === "d" || data === "D") this.details = !this.details; + } else if (data === "1" || data === "2" || matchesKey(data, Key.left) || matchesKey(data, Key.right)) { + this.offset = 0; + this.model.select( + data === "1" + ? "overview" + : data === "2" + ? "account" + : this.model.view === "overview" + ? "account" + : "overview", + ); + } else if (data === "r" || data === "R") this.model.refresh(); + else if ((data === "p" || data === "P") && this.model.view === "overview") + this.model.selectScope(this.model.scope === "all" ? "project" : "all"); + else if (["[", "]"].includes(data) && this.model.view === "overview" && this.model.overview) { + const days = this.model.overview.days, + index = Math.max( + 0, + days.findIndex((d) => d.date === this.model.selectedDate), + ); + this.model.selectedDate = days[Math.max(0, Math.min(days.length - 1, index + (data === "]" ? 1 : -1)))].date; + } + this.redraw(); + } + render(width: number): string[] { + const inner = Math.max(1, width - 4), + rows = usagePresentationRows(this.model, inner, this.details).flatMap((row) => wrapTextWithAnsi(row, inner)); + const height = Math.max(5, this.rows() - 6); + this.offset = Math.max(0, Math.min(this.offset, rows.length - height)); + return renderStepDialogFrame( + rows.slice(this.offset, this.offset + height).map((row) => truncateToWidth(row, inner)), + width, + ); + } +} + +export async function showUsagePanel(ctx: ExtensionCommandContext, model: UsagePresentation): Promise { + await ctx.ui.custom((tui, _theme, _keys, done) => { + model.onChange = () => tui.requestRender(); + model.onClose = () => done(); + const panel = new UsagePanel( + model, + () => tui.requestRender(), + () => done(), + () => tui.terminal.rows, + ); + return Object.assign(panel, { + dispose: () => { + model.onChange = undefined; + model.onClose = undefined; + }, + }); + }); +} diff --git a/apps/cli/test/interactive-mode-step-login.test.ts b/apps/cli/test/interactive-mode-step-login.test.ts index b2db2e56..b3c44cdd 100644 --- a/apps/cli/test/interactive-mode-step-login.test.ts +++ b/apps/cli/test/interactive-mode-step-login.test.ts @@ -106,6 +106,19 @@ describe("Step login command routing", () => { expect(showWarning).toHaveBeenCalledWith("Wait for the active turn to finish before signing in"); }); + it("shows a quota-check warning after reloading a completed login", async () => { + const warning = "Could not verify Step Plan availability. Retry with /usage account."; + const showWarning = vi.fn(); + const reload = vi.fn(async () => {}); + await handleLoginCommand.call({ + options: { stepLogin: vi.fn(async () => ({ kind: "completed", warning })) }, + createStepLoginHost: () => ({}), handleReloadCommand: reload, + showWarning, showStatus: vi.fn(), showError: vi.fn(), + }); + expect(reload).toHaveBeenCalledOnce(); + expect(showWarning).toHaveBeenCalledWith(warning); + }); + it("does not reopen the Step login panel when already signed in", async () => { const root = await mkdtemp(join(tmpdir(), "interactive-step-login-auth-")); const stepLogin = vi.fn(async () => ({ kind: "exit" as const })); diff --git a/apps/cli/test/usage-panel.test.ts b/apps/cli/test/usage-panel.test.ts new file mode 100644 index 00000000..76df06a1 --- /dev/null +++ b/apps/cli/test/usage-panel.test.ts @@ -0,0 +1,154 @@ +import { afterEach, describe, expect, it, vi } from "vitest"; +import { stripTerminalSequences, visibleWidth } from "@step-harness/pi-tui"; +import { type ExtensionCommandContext, initTheme, type UsagePresentation } from "@step-harness/coding-agent"; +import { Container } from "../../../packages/tui/src/tui.ts"; +import { Text } from "../../../packages/tui/src/components/text.ts"; +import { TuiMainScreen } from "../../../packages/tui/src/tui-main-screen.ts"; +import { TuiAltScreen } from "../../../packages/tui/src/tui-alt-screen.ts"; +import { ScrollView } from "../../../packages/tui/src/components/scroll-view.ts"; +import { VStack } from "../../../packages/tui/src/components/v-stack.ts"; +import { VirtualTerminal } from "../../../packages/tui/test/virtual-terminal.ts"; +import { InteractiveMode } from "../src/ui/interactive-mode.ts"; +import { showUsagePanel, UsagePanel, usagePresentationRows } from "../src/ui/view/dialogs/usage-panel.ts"; + +function model(view: UsagePresentation["view"] = "account"): UsagePresentation { + const account = { status: "ready" as const, account: { presence: "present" as const, eligible: true, source: "plan", accountId: "123" }, receivedAt: Date.now(), snapshot: { schema_version: 1 as const, observed_at_ms: Date.now(), scope: "personal_account" as const, account: { id: "123" }, plan: { state: "known" as const, family: "coding" as const, name: "Pro", subscription: null }, quota: { state: "known" as const, allowed: true, reason: null, windows: [{ id: "weekly" as const, limited: true, used_percent: 13, resets_at_ms: Date.now() + 100000 }, { id: "five_hour" as const, limited: false, used_percent: null, resets_at_ms: null }], credit_buckets: [] } } }; + return { view, account, status: { version: "1.0", model: "step-model", provider: "step", reasoning: "high", directory: "/work", permissions: "Ask", trusted: true, agents: ["/work/AGENTS.md"], collaboration: "Default", session: "session", state: "idle", context: "10/100", cached: account, account: account.account }, session: { input: 20, output: 2, total: 22, cacheRead: 3, cacheWrite: 0, cacheHit: .15, invalid: 0, missing: 0 }, scope: "all", selectedDate: "2026-10-13", select: vi.fn(), refresh: vi.fn(), selectScope: vi.fn() }; +} +afterEach(() => initTheme("dark")); + +describe.each(["regular", "fullscreen"] as const)("Usage panel mounting in %s mode", (mode) => { + it.each(["overview", "account", "status"] as const)("places %s in the editor area and restores the draft on close", async (view) => { + initTheme("step-blue"); + const terminal = new VirtualTerminal(80, 24); + const ui = mode === "regular" ? new TuiMainScreen(terminal) : new TuiAltScreen(terminal); + const draft = "unfinished draft"; + const editor = Object.assign(new Text("EDITOR", 0, 0), { + getText: () => draft, + setText: vi.fn(), + handleInput: vi.fn(), + }); + const editorContainer = new Container(); + editorContainer.addChild(editor); + const transcript = new Text("CONVERSATION", 0, 0); + const footer = new Text("FOOTER", 0, 0); + if (mode === "fullscreen") { + (ui as TuiAltScreen).setLayoutRoot(new VStack([ + { component: new ScrollView(transcript, { follow: "end", primary: true }), basis: 0, grow: 1, minSize: 1 }, + { component: new VStack([ + { component: editorContainer, shrink: 1, minSize: 3 }, + { component: footer, shrink: 1, minSize: 1 }, + ]), basis: "auto", minSize: 1 }, + ])); + } else { + ui.addChild(transcript); + ui.addChild(editorContainer); + ui.addChild(footer); + } + const host = { + editor, editorContainer, ui, keybindings: {}, disposeActiveSelector: vi.fn(), + redraw: { requestRender: () => ui.requestRender() }, + }; + // Exercise the real extension mount; only unrelated host dependencies are stubbed. + const mount = Reflect.get(InteractiveMode.prototype, "showExtensionCustom") as ExtensionCommandContext["ui"]["custom"]; + const ctx = { ui: { custom: mount.bind(host) } } as unknown as ExtensionCommandContext; + const m = model(view); + let shown: Promise | undefined; + const flush = async (): Promise => { + await Promise.resolve(); + ui.requestRender(true); + await terminal.waitForRender(); + }; + ui.setFocus(editor); + ui.start(); + try { + shown = showUsagePanel(ctx, m); + for (const [width, height] of [[80, 24], [40, 18], [120, 32]]) { + terminal.resize(width, height); + await flush(); + const lines = terminal.getViewport(); + expect(ui.hasOverlay()).toBe(false); + expect(editorContainer.children).toHaveLength(1); + expect(editorContainer.children[0]).toBeInstanceOf(UsagePanel); + expect(lines.join("\n")).not.toContain("EDITOR"); + expect(lines.join("\n")).toContain("CONVERSATION"); + const top = lines.find((line) => line.startsWith("╭")); + expect(top?.length).toBe(width); + const bottom = lines.findIndex((line) => line.startsWith("╰")); + expect(bottom).toBeGreaterThanOrEqual(0); + expect(lines.findIndex((line) => line.trimEnd() === "FOOTER"), lines.join("\n")).toBeGreaterThan(bottom); + if (width === 40 && view !== "overview") { + terminal.sendInput("\u001b[B"); + await flush(); + expect(terminal.getViewport()).not.toEqual(lines); + expect(terminal.getViewport().some((line) => line.trimEnd() === "FOOTER")).toBe(true); + terminal.sendInput("\u001b[A"); + await flush(); + } + } + if (view === "status") { + m.status.model = "updated-model"; + m.onChange?.(); + await flush(); + expect(terminal.getViewport().join("\n")).toContain("updated-model"); + } + terminal.sendInput("\u001b"); + await shown; + await flush(); + expect(editorContainer.children).toEqual([editor]); + expect(editor.setText).toHaveBeenCalledWith(draft); + expect(m.onChange).toBeUndefined(); + expect(m.onClose).toBeUndefined(); + terminal.sendInput("x"); + expect(editor.handleInput).toHaveBeenCalledWith("x"); + } finally { + m.onClose?.(); + await shown; + ui.stop(); + } + }); +}); +describe("Usage and Status presentation", () => { + it("shows only Overview/Account tabs and consistent remaining quota wording", () => { + const account = usagePresentationRows(model()).join("\n"), status = usagePresentationRows(model("status")).join("\n"); + expect(account).toContain("1 Overview [2 Account]"); expect(account).not.toMatch(/3 Session|4 History/u); expect(account).toContain("87.0% left"); expect(account).not.toContain("% used"); expect(status).toContain("87.0% left"); expect(status).toContain("Permissions:"); expect(status).toContain("Session:"); expect(status).toContain("no network request"); + }); + it.each([ + { used: 0, total: "100", remaining: "100", expected: `[${"▆".repeat(12)}] 100.0% left` }, + { used: 25, total: "100", remaining: "75", expected: `[${"▆".repeat(9)}${"▁".repeat(3)}] 75.0% left` }, + { used: 100, total: "100", remaining: "0", expected: `[${"▁".repeat(12)}] 0.0% left` }, + { used: null, total: "0", remaining: "0", expected: "Remaining percentage unavailable" }, + ])("shows Credit percentages without confusing full, empty and unknown balances: $used", ({ used, total, remaining, expected }) => { + for (const view of ["account", "status"] as const) { + const m = model(view); + m.account.snapshot!.plan.family = "token"; + m.account.snapshot!.quota.windows = []; + m.account.snapshot!.quota.credit_buckets = [{ type: "subscription", total, remaining, used_percent: used, resets_at_ms: null, expires_at_ms: null }]; + const rows = usagePresentationRows(m).map(stripTerminalSequences); + const index = rows.findIndex((row) => row.startsWith("Subscription Credits:")); + expect(rows[index]).toContain(expected); + expect(rows.join("\n")).not.toContain(`${remaining} / ${total} Credits left`); + if (used === null) expect(rows.join("\n")).not.toMatch(/\d+(?:\.\d+)?% left/u); + } + }); + it("shows percentages instead of raw Credit amounts and hides unknown buckets", () => { + const m = model(); m.account.snapshot!.plan.family = "token"; m.account.snapshot!.quota.windows = []; + m.account.snapshot!.quota.credit_buckets = [{ type: "topup", total: "9007199254740993", remaining: "9007199254740992", used_percent: 0, resets_at_ms: null, expires_at_ms: null }, { type: "future", unsupported: true, total: "0", remaining: "0", used_percent: null, resets_at_ms: null, expires_at_ms: null }]; + const rendered = usagePresentationRows(m).join("\n"); expect(rendered).toContain("100.0% left"); expect(rendered).not.toContain("9,007,199,254,740,993"); expect(rendered).not.toContain("9,007,199,254,740,992"); expect(rendered).toContain("Unsupported bucket type: future"); + }); + it("reuses ANSI stripping for historical plugin and skill names", () => { + const m = model("overview"); + m.overview = { days: [{ date: "2026-10-13", tokens: 1, chats: 1 }], tokens: 1, chats: 1, activeDays: 1, + reasoning: [], discoveredSkills: 1, usedSkills: 1, topResources: [["\x1b[31m$example\x1b[0m", 1]], limited: [], scope: "all" }; + const rows = usagePresentationRows(m).map(stripTerminalSequences).join("\n"); + expect(rows).toContain("$example"); + expect(rows).not.toContain("[31m"); + }); + + it("fits narrow terminals, keeps only the documented tab shortcuts and scrolls", () => { + initTheme("step-blue"); const m = model(), close = vi.fn(), panel = new UsagePanel(m, vi.fn(), close, () => 18); + for (const width of [24, 60, 120]) for (const row of panel.render(width)) expect(visibleWidth(row)).toBeLessThanOrEqual(width); + panel.handleInput("3"); expect(m.select).not.toHaveBeenCalled(); panel.handleInput("1"); expect(m.select).toHaveBeenCalledWith("overview"); panel.handleInput("r"); expect(m.refresh).toHaveBeenCalledOnce(); panel.handleInput("\u001b"); expect(close).toHaveBeenCalledOnce(); + m.view = "status"; expect(panel.render(60).map(stripTerminalSequences).join("\n")).toContain("Model:"); panel.handleInput("\u001b[6~"); expect(panel.render(60).map(stripTerminalSequences).join("\n")).toContain("Weekly"); + }); +}); diff --git a/docs/step-configuration.md b/docs/step-configuration.md index c75e74b7..5291256e 100644 --- a/docs/step-configuration.md +++ b/docs/step-configuration.md @@ -19,6 +19,29 @@ and `auth.json` keep their own formats — they hold model definitions and credentials, not settings — and a stale Step endpoint recorded in `models.json` by an older release is still repaired in place at startup. +## Step Plan login and API-key fallback + +Browser sign-in checks the selected region's `/step_plan/v1/usage` endpoint +after saving the credential. An empty interactive launch also checks a stored +Step Plan login before opening the session. The model list is not proof of an +active subscription: it can contain fallback models for an account without a +plan. + +When the usage response confirms there is no active subscription or usable +top-up credit, the login screen offers an API key, the subscription page, +checking again after purchase, or another account / login method. The API-key +option opens the existing Platform key-entry screen for the same region and +only replaces the credential when the key is submitted. Canceling preserves +the stored Plan credential. Mainland and overseas profiles follow the same +flow and keep their own endpoints. + +An active plan with exhausted quota remains a Plan login. A top-up-only account +with usable credit can also continue. Network failures, missing usage endpoints, +invalid responses, and uncertain account state produce a warning instead of +being classified as a missing subscription. Explicit API keys, Platform logins, +other selected providers, and non-interactive commands skip the startup Plan +check. No additional model request is made for this check. + ## Mandatory command approval Ask and Approve for Me cannot automatically approve commands matched by the diff --git a/docs/tui-rendering-pipeline.md b/docs/tui-rendering-pipeline.md index 060b302e..9ed4a2ab 100644 --- a/docs/tui-rendering-pipeline.md +++ b/docs/tui-rendering-pipeline.md @@ -180,7 +180,9 @@ alt screen 进入时（`tui-alt-screen.ts:316-318`）：`\x1b[?1049h` + 关 auto - fullscreen 模式改用布局根：transcript 滚动视图（grow）+ dock（`interactive-mode.ts:1102-1126`）。 - 挂载后 `setFocus(editor)`，`ui.start()`，再 `applyFromSettings()` 上主题（`interactive-mode.ts:1138-1157`）。 -焦点与 overlay：对话框统一 `ui.showOverlay(component, options)` + `ui.setFocus(component)`，关闭时回焦编辑器（`interactive-mode.ts:3113-3131`、`3350-3360`）。overlay 的焦点恢复（被 overlay 内部组件抢走再还回）由 `TuiBase` 的状态机保证（`tui.ts:525-588`）。 +选择器与自定义面板有两种挂载方式。`/settings` 通过 `showSelector` 替换 `editorContainer` 内的编辑器；`/usage` 与 `/status` 使用 `ctx.ui.custom` 默认的 inline 模式，采用相同的输入区域和终端全宽，底栏保留在面板下方。关闭时恢复编辑器、未提交草稿与键盘焦点；异步更新和关闭回调随面板释放。 + +只有显式选择 overlay 的对话框才通过 `ui.showOverlay(component, options)` 叠加到现有内容上。overlay 的焦点恢复（被 overlay 内部组件抢走再还回）由 `TuiBase` 的状态机保证（`tui.ts:525-588`）；普通命令面板不应为了展示内容而默认居中覆盖聊天记录。 ### 7.2 事件 → 组件树 → 重绘 diff --git a/packages/coding-agent/docs/account-usage.md b/packages/coding-agent/docs/account-usage.md new file mode 100644 index 00000000..ecca18a1 --- /dev/null +++ b/packages/coding-agent/docs/account-usage.md @@ -0,0 +1,47 @@ +# Account usage and local activity + +`/usage` opens two tabs: **Overview** and **Account**. `/usage overview` and `/usage account` open either tab directly. `/usage --refresh` and `/usage account --refresh` refresh Account. Session and History are not separate tabs or commands. + +## Overview + +Overview reads this machine's Step session JSONL files and shows the last 365 local calendar dates. It includes an activity heatmap, daily token/chat totals, reasoning settings recorded at dispatch, discovered/used skills, and plugin/skill runs. `P` switches all projects/current project; `[` and `]` select a date. Narrow terminals show recent weeks while totals still cover the year. + +All active heatmap cells use the theme's accent color. Denser blocks indicate more tokens; `▣` marks the selected date and `·` indicates no token usage. + +Tokens include assistant replies, tool-reported usage, compaction/branch summaries, and incremental AutoCLM records. Input includes ordinary input plus cache reads/writes. All main-session entries are counted, including restored history and branches; cumulative AutoCLM summaries are not counted again. Subagent tokens contribute to Overview, attributed to their parent project; subagent sessions do not increase the chat count. + +Forks/exports carry small origin manifests so copied entries retain their original session/project. Missing legacy parent files and older sessions without activity/provenance records are reported as limited coverage. Skill activation means an actual read of a loaded SKILL.md; plugin ownership comes from discovered manifests/tool registrations. Historical metadata is not inferred from today's configuration. + +Reading runs in a worker and does not migrate or rewrite session files. Only statistics/identity metadata is retained in a rebuildable private cache. Message text, prompts, tool arguments, credentials and errors are not copied to that cache. Append-only prefixes are verified, incomplete tails wait for the next scan, corrupt caches rebuild, and concurrent commits compare source identity/prefix under a short lock. + +Limits are explicit: 10,000 files, 25,000 selected records, 16 MiB per source line, depth 64, 200,000 JSON tokens, 256 KiB retained metadata per line, 4 KiB per retained string, 1 second per parse and 30 seconds per scan. Budget-limited/invalid records produce coverage notices. Node additionally limits worker heap; the same parser budgets apply in Bun. A process crash can leave a cache lock until its 5-second stale threshold; it does not affect source sessions. + +## Account + +Account queries the personal Step Plan usage endpoint with the stored CN/Global Plan credential. Model/gateway keys supplied by `STEP_API_KEY` or `--api-key`, platform profiles and unknown profiles do not query that endpoint. Redirects are rejected. Keys stay in the controller, outside renderer/output data. + +CODING plans show 5-hour and weekly percentages **remaining** and reset times. Nonpositive upstream limits mean Unlimited. TOKEN plans show the remaining percentage for each subscription/top-up Credit bucket and reset/expiry times, without displaying raw Credit balances. Account and Status use compact 12-column bars with reduced-height blocks to represent remaining quota. A missing percentage, including a zero-total bucket, is shown as unavailable rather than 0% or 100%. Credit quantities remain exact decimal strings in the snapshot; percentages use the API's validated `used_percent` without converting Credit quantities to floating point. Unknown future bucket types are marked unsupported; partial bucket errors are not displayed as zero balance. + +The current main-session summary is local and includes restored/inherited history and incremental compression usage. It excludes separate child-session files. Cache hit is cache reads divided by complete input; an empty denominator is Unknown. + +Account responses have a 30-second memory cache and a 5-minute display lifetime. Identity/profile/endpoint/key changes invalidate it. Requests are single-flight, time out after 12 seconds and are cancelled when the panel closes or changes tabs/sessions. Attempt generations reject late responses. TTL/cooldown use a monotonic clock; HTTP Retry-After supports seconds and HTTP-date. Visible temporary failures retain a stale snapshot; 401/403 clear the snapshot without deleting credentials. No background polling runs. + +## Status and controls + +`/status` reads the live local model, provider, reasoning, directory, permissions, loaded AGENTS.md files, account metadata, collaboration mode and session identity. It shows cached quota **remaining**; it never triggers an account request. `D` expands local diagnostic details. Credential presence is not a claim of online validity. + +`1`/`2` or left/right arrows switch Usage tabs, `R` refreshes Account, up/down arrows scroll, and Esc closes. The panel uses the editor area, with the footer underneath, and restores editor input on close. + +RPC command output uses `extension_ui_request` notifications. Other noninteractive CLI modes use stderr, keeping protocol stdout clear. SDK embedders can use `PlanUsageController`, `scanUsageOverview`, `collectSessionUsage` and the exported presentation types; a custom command host supplies its own presenter/output sink. + +## Implementation boundaries + +The existing `core/usage-totals.ts` owns entry selection and token validation for session statistics, Account and Overview. `AgentSession.getSessionStats()`, the cost breakdown and Account also share its accumulator; Overview groups validated increments by date and original project after deduplication. Account's complete input includes cache reads/writes; the existing session-stat input retains its provider convention and reports caches separately. Invalid token increments are skipped consistently, while absent/invalid cost does not discard otherwise valid tokens. No second session ledger is maintained. + +`step/usage/recording.ts` records activity through existing extension hooks and custom session entries. `commands.ts` owns command parsing and presentation lifecycle; `apps/cli` supplies the presenter and output sink. The bounded historical scanner/cache stays in the product layer because the ordinary session loader materializes complete messages and migrates legacy data. Fork/export hooks preserve origins when parent files disappear; removing those hooks would lose attribution. They do not change model context or dispatch behavior. + +## Release verification + +Node packages emit `activity-worker.js` beside the chunk that creates it. Bun standalone builds explicitly embed `activity-worker.ts`; release archives still need the existing theme/assets. Test workers from a directory outside the checkout. + +The backend must provide schema 1 at `GET /step_plan/v1/usage` before Account is available. A 404 explains that the endpoint is not deployed. Overview and Status remain local. Platform bill amounts, organization Credit and cross-device activity history are outside this feature. diff --git a/packages/coding-agent/docs/index.md b/packages/coding-agent/docs/index.md index 5b1afbd8..1ba79715 100644 --- a/packages/coding-agent/docs/index.md +++ b/packages/coding-agent/docs/index.md @@ -33,6 +33,7 @@ For the full first-run flow, see [Quickstart](quickstart.md). ## Start here - [Quickstart](quickstart.md) - install, authenticate, and run a first session. +- [Account usage and local activity](account-usage.md) - Overview, Account and local Status. - [Using Step](usage.md) - interactive mode, slash commands, context files, and CLI reference. - [Providers](providers.md) - subscription and API-key setup for built-in providers. - [llama.cpp](llama-cpp.md) - run a local router and manage models with `/llama`. diff --git a/packages/coding-agent/docs/session-format.md b/packages/coding-agent/docs/session-format.md index 153b9503..8a97a1d9 100644 --- a/packages/coding-agent/docs/session-format.md +++ b/packages/coding-agent/docs/session-format.md @@ -119,6 +119,8 @@ interface Usage { The exported pi-ai `StopReason` type also includes `"pending"`, but that value is reserved for partial messages in streaming events. Terminal `done`/`error` messages replace it with a completion reason before step persists the assistant message, so `"pending"` should never appear in session JSONL. +Session totals and the cost breakdown use the shared [`usage-totals.ts`](../src/core/usage-totals.ts) ledger. It includes assistant, tool, compaction, branch-summary and incremental `step-auto-clm-usage` records; cumulative maintenance summaries are excluded. Token increments must contain nonnegative safe integers and must not overflow the accumulated total. A missing/invalid cost does not hide valid tokens. `SessionStats.tokens.input` excludes caches, which are reported separately; the [Account summary](account-usage.md) combines them into complete input. + ### Extended Message Types (from pi-coding-agent) ```typescript diff --git a/packages/coding-agent/package.json b/packages/coding-agent/package.json index baebe106..ff9262e9 100644 --- a/packages/coding-agent/package.json +++ b/packages/coding-agent/package.json @@ -28,7 +28,7 @@ "clean": "shx rm -rf dist", "build": "npm run build:unbundled && node ../../scripts/build-coding-agent-bundle.mjs", "build:unbundled": "tsgo -p tsconfig.build.json && npm run copy-assets", - "build:binary": "npm --prefix ../tui run build && npm --prefix ../telemetry run build && npm --prefix ../providers run build && npm --prefix ../agent-core run build && npm run build && bun build --compile --no-compile-autoload-bunfig --env=STEPCODE_BUILD_* ../../apps/cli/src/bun/stepcode.ts ./src/utils/image-resize-worker.ts --outfile dist/step-bin && npm run copy-binary-assets", + "build:binary": "npm --prefix ../tui run build && npm --prefix ../telemetry run build && npm --prefix ../providers run build && npm --prefix ../agent-core run build && npm run build && bun build --compile --no-compile-autoload-bunfig --env=STEPCODE_BUILD_* ../../apps/cli/src/bun/stepcode.ts ./src/utils/image-resize-worker.ts ./src/step/usage/activity-worker.ts --outfile dist/step-bin && npm run copy-binary-assets", "copy-assets": "shx mkdir -p dist/theme && shx cp src/theme/*.json dist/theme/ && shx mkdir -p dist/core/export-html/vendor && shx cp src/core/export-html/template.html src/core/export-html/template.css src/core/export-html/template.js dist/core/export-html/ && shx cp src/core/export-html/vendor/*.js dist/core/export-html/vendor/", "copy-binary-assets": "shx cp package.json dist/ && shx cp README.md dist/ && shx mkdir -p dist/theme && shx cp src/theme/*.json dist/theme/ && shx mkdir -p dist/export-html/vendor && shx cp src/core/export-html/template.html dist/export-html/ && shx cp src/core/export-html/vendor/*.js dist/export-html/vendor/ && shx cp -r docs dist/ && shx cp -r examples dist/ && node ../../scripts/copy-photon-wasm.mjs dist", "test": "vitest --run", diff --git a/packages/coding-agent/src/core/agent-session.ts b/packages/coding-agent/src/core/agent-session.ts index 850d41c6..8504afa3 100644 --- a/packages/coding-agent/src/core/agent-session.ts +++ b/packages/coding-agent/src/core/agent-session.ts @@ -114,6 +114,7 @@ import { ModelRegistry } from "./model-registry.ts"; import type { ModelRuntime } from "./model-runtime.ts"; import { expandPromptTemplate, type PromptTemplate } from "./prompt-templates.ts"; import { isRateLimitError, rateLimitDelay } from "./rate-limit.ts"; +import { readRequestUsage } from "./request-provenance.ts"; import type { ResourceExtensionPaths, ResourceLoader } from "./resource-loader.ts"; import { exportSessionToJsonl } from "./session-export.ts"; import type { BranchSummaryEntry, CompactionEntry, SessionEntry, SessionManager } from "./session-manager.ts"; @@ -126,7 +127,7 @@ import { type BashOperations, createLocalBashOperations } from "./tools/bash.ts" import { createAllToolDefinitions } from "./tools/index.ts"; import { createToolDefinitionFromAgentTool } from "./tools/tool-definition-wrapper.ts"; import { boundToolResultContent } from "./tools/tool-output.ts"; -import { addUsageToTotals, createUsageTotals, getAutoClmUsage } from "./usage-totals.ts"; +import { addUsageToTotals, createUsageTotals, entryUsage } from "./usage-totals.ts"; export { type ParsedSkillBlock, parseSkillBlock } from "../utils/skill-block.ts"; @@ -1162,7 +1163,11 @@ export class AgentSession { event.message.role === "toolResult" ) { // Regular LLM message - persist as SessionMessageEntry - this.sessionManager.appendMessage(event.message); + const entryId = this.sessionManager.appendMessage(event.message); + if (event.message.role === "assistant" && this._extensionRunner?.hasHandlers("usage_recorded")) { + const provenance = readRequestUsage(event.message); + if (provenance) await this._extensionRunner.emit({ type: "usage_recorded", entryId, provenance }); + } } // Other message types (bashExecution, compactionSummary, branchSummary) are persisted elsewhere @@ -4000,6 +4005,7 @@ export class AgentSession { * actually billed across the session. */ getSessionStats(): SessionStats { + const entries = this.sessionManager.getEntries(); let userMessages = 0; let assistantMessages = 0; let toolResults = 0; @@ -4007,12 +4013,8 @@ export class AgentSession { let toolCalls = 0; const usageTotals = createUsageTotals(); - for (const entry of this.sessionManager.getEntries()) { - const maintenanceUsage = getAutoClmUsage(entry); - if (maintenanceUsage) addUsageToTotals(usageTotals, maintenanceUsage); - if ((entry.type === "branch_summary" || entry.type === "compaction") && entry.usage) { - addUsageToTotals(usageTotals, entry.usage); - } + for (const entry of entries) { + addUsageToTotals(usageTotals, entryUsage(entry).raw); if (entry.type !== "message") continue; totalMessages++; const message = entry.message; @@ -4020,16 +4022,12 @@ export class AgentSession { userMessages++; } else if (message.role === "toolResult") { toolResults++; - if (message.usage) { - addUsageToTotals(usageTotals, message.usage); - } } else if (message.role === "assistant") { assistantMessages++; const assistantMsg = message as AssistantMessage; if (Array.isArray(assistantMsg.content)) { toolCalls += assistantMsg.content.filter((c) => c.type === "toolCall").length; } - addUsageToTotals(usageTotals, assistantMsg.usage); } } diff --git a/packages/coding-agent/src/core/extensions/types.ts b/packages/coding-agent/src/core/extensions/types.ts index ad76bd3f..9c9e0bb7 100644 --- a/packages/coding-agent/src/core/extensions/types.ts +++ b/packages/coding-agent/src/core/extensions/types.ts @@ -858,6 +858,13 @@ export interface MessageEndEvent { message: AgentMessage; } +/** Local provenance emitted after the finalized assistant message is persisted. */ +export interface UsageRecordedEvent { + type: "usage_recorded"; + entryId: string; + provenance: import("../request-provenance.ts").RequestUsageProvenance; +} + /** Fired when a tool starts executing */ export interface ToolExecutionStartEvent { type: "tool_execution_start"; @@ -1166,6 +1173,7 @@ export type ExtensionEvent = | MessageStartEvent | MessageUpdateEvent | MessageEndEvent + | UsageRecordedEvent | ToolExecutionStartEvent | ToolExecutionUpdateEvent | ToolExecutionEndEvent @@ -1354,6 +1362,7 @@ export interface ExtensionAPI { on(event: "message_start", handler: ExtensionHandler): void; on(event: "message_update", handler: ExtensionHandler): void; on(event: "message_end", handler: ExtensionHandler): void; + on(event: "usage_recorded", handler: ExtensionHandler): void; on(event: "tool_execution_start", handler: ExtensionHandler): void; on(event: "tool_execution_update", handler: ExtensionHandler): void; on(event: "tool_execution_end", handler: ExtensionHandler): void; diff --git a/packages/coding-agent/src/core/model-request-observer.ts b/packages/coding-agent/src/core/model-request-observer.ts index e6351e34..99719fef 100644 --- a/packages/coding-agent/src/core/model-request-observer.ts +++ b/packages/coding-agent/src/core/model-request-observer.ts @@ -66,6 +66,8 @@ export interface ModelRequestObserver { } export interface ObserveModelRequestStreamOptions { + /** Local, best-effort terminal hook; independent of optional remote observers. */ + readonly onTerminal?: (message: AssistantMessage) => void; /** Timestamp captured immediately before the provider stream is requested. */ readonly startedAt?: number; /** Session identity forwarded to the observer, when available. */ @@ -140,11 +142,11 @@ export function observeModelRequestStream( observer: ModelRequestObserver | undefined, options: ObserveModelRequestStreamOptions = {}, ): AssistantMessageEventStream { - if (!observer) return source; + if (!observer && !options.onTerminal) return source; const startedAt = options.startedAt ?? Date.now(); const started = createStartedRequest(model, options, startedAt); - notifyStarted(observer, started); + if (observer) notifyStarted(observer, started); // Always use Pi's public factory instead of cloning `source.constructor`. // ModelRuntime normally returns an AssistantMessageEventStream, but lazy and @@ -159,6 +161,12 @@ export function observeModelRequestStream( const finish = (message: AssistantMessage, unexpectedError?: unknown): void => { if (settled) return; settled = true; + try { + options.onTerminal?.(message); + } catch { + /* Local diagnostics cannot change streaming. */ + } + if (!observer) return; const completedAt = Date.now(); const statusCode = readStatusCode(options.getStatusCode); const effectiveModel = readModel(model, options.getModel); diff --git a/packages/coding-agent/src/core/request-provenance.ts b/packages/coding-agent/src/core/request-provenance.ts new file mode 100644 index 00000000..64dea0fa --- /dev/null +++ b/packages/coding-agent/src/core/request-provenance.ts @@ -0,0 +1,17 @@ +import type { AssistantMessage } from "@step-harness/providers/compat"; + +export interface RequestUsageProvenance { + requestId: string; + sessionId: string; + sourceLeafId: string | null; + reasoningAtDispatch: string; + requestedProvider: string; + requestedModel: string; +} +const provenance = new WeakMap(); +export function rememberRequestUsage(message: AssistantMessage, value: RequestUsageProvenance): void { + provenance.set(message, value); +} +export function readRequestUsage(message: AssistantMessage): RequestUsageProvenance | undefined { + return provenance.get(message); +} diff --git a/packages/coding-agent/src/core/sdk.ts b/packages/coding-agent/src/core/sdk.ts index f07d5a34..af5c39c0 100644 --- a/packages/coding-agent/src/core/sdk.ts +++ b/packages/coding-agent/src/core/sdk.ts @@ -1,3 +1,4 @@ +import { randomUUID } from "node:crypto"; import { join } from "node:path"; import { Agent, type AgentMessage, setDefaultStreamFn, type ThinkingLevel } from "@step-harness/agent-core"; import type { AssistantMessageEventStream, ModelsSimpleStreamOptions } from "@step-harness/providers"; @@ -17,6 +18,7 @@ import { import { findInitialModel } from "./model-resolver.ts"; import { ModelRuntime } from "./model-runtime.ts"; import { mergeProviderAttributionHeaders } from "./provider-attribution.ts"; +import { rememberRequestUsage } from "./request-provenance.ts"; import type { ResourceLoader } from "./resource-loader.ts"; import { DefaultResourceLoader } from "./resource-loader.ts"; import { getDefaultSessionDir, SessionManager } from "./session-manager.ts"; @@ -334,6 +336,14 @@ export async function createAgentSession(options: CreateAgentSessionOptions = {} }, convertToLlm: convertToLlmWithBlockImages, streamFn: async (model, context, options) => { + const usageProvenance = { + requestId: randomUUID(), + sessionId: options?.sessionId ?? sessionManager.getSessionId(), + sourceLeafId: sessionManager.getLeafId(), + reasoningAtDispatch: options?.reasoning ?? "off", + requestedProvider: model.provider, + requestedModel: model.id, + }; const startedAt = Date.now(); let statusCode = 0; // ModelRuntime can replace the model's base URL after resolving auth @@ -384,6 +394,7 @@ export async function createAgentSession(options: CreateAgentSessionOptions = {} throw error; } return observeModelRequestStream(model, stream, modelRequestObserver, { + onTerminal: (message) => rememberRequestUsage(message, usageProvenance), startedAt, sessionId: options?.sessionId, getStatusCode: () => statusCode, diff --git a/packages/coding-agent/src/core/session-export.ts b/packages/coding-agent/src/core/session-export.ts index 3607242c..c5229e0b 100644 --- a/packages/coding-agent/src/core/session-export.ts +++ b/packages/coding-agent/src/core/session-export.ts @@ -2,6 +2,7 @@ import { existsSync, mkdirSync, writeFileSync } from "node:fs"; import { dirname } from "node:path"; import { resolvePath } from "../utils/paths.ts"; import { CURRENT_SESSION_VERSION, type SessionHeader, type SessionManager } from "./session-manager.ts"; +import { createUsageLineage } from "./usage-lineage.ts"; /** Write the current session branch and optional trailing export-only entries as JSONL. */ export function exportSessionToJsonl( @@ -36,6 +37,15 @@ export function exportSessionToJsonl( for (const entry of createTrailingEntries?.(parentId, timestamp) ?? []) { lines.push(JSON.stringify(entry)); } + for (const entry of createUsageLineage( + sessionManager.getHeader(), + sessionManager.getEntries(), + sessionManager.getBranch(), + header.id, + timestamp, + parentId, + )) + lines.push(JSON.stringify(entry)); writeFileSync(filePath, `${lines.join("\n")}\n`); return filePath; diff --git a/packages/coding-agent/src/core/session-manager.ts b/packages/coding-agent/src/core/session-manager.ts index 7a198581..a1f05598 100644 --- a/packages/coding-agent/src/core/session-manager.ts +++ b/packages/coding-agent/src/core/session-manager.ts @@ -26,6 +26,7 @@ import { createCompactionSummaryMessage, createCustomMessage, } from "./messages.ts"; +import { createUsageLineage } from "./usage-lineage.ts"; export const CURRENT_SESSION_VERSION = 3; @@ -1423,6 +1424,8 @@ export class SessionManager { * Returns the new session file path, or undefined if not persisting. */ createBranchedSession(leafId: string): string | undefined { + const usageSourceHeader = this.getHeader(), + usageSourceEntries = this.getEntries(); const previousSessionFile = this.sessionFile; const path = this.getBranch(leafId); if (path.length === 0) { @@ -1482,7 +1485,15 @@ export class SessionManager { parentId = labelEntry.id; } - this.fileEntries = [header, ...pathWithoutLabels, ...labelEntries]; + const lineage = createUsageLineage( + usageSourceHeader, + usageSourceEntries, + pathWithoutLabels, + newSessionId, + timestamp, + parentId, + ); + this.fileEntries = [header, ...pathWithoutLabels, ...labelEntries, ...lineage]; this.sessionId = newSessionId; this.sessionFile = newSessionFile; this._buildIndex(); @@ -1643,6 +1654,16 @@ export class SessionManager { appendFileSync(newSessionFile, `${JSON.stringify(entry)}\n`); } } + const copied = sourceEntries.filter((entry): entry is SessionEntry => entry.type !== "session"); + for (const entry of createUsageLineage( + sourceHeader, + copied, + copied, + newSessionId, + timestamp, + copied.at(-1)?.id ?? null, + )) + appendFileSync(newSessionFile, `${JSON.stringify(entry)}\n`); return new SessionManager(resolvedTargetCwd, dir, newSessionFile, true); } diff --git a/packages/coding-agent/src/core/usage-lineage.ts b/packages/coding-agent/src/core/usage-lineage.ts new file mode 100644 index 00000000..7ebd2b52 --- /dev/null +++ b/packages/coding-agent/src/core/usage-lineage.ts @@ -0,0 +1,87 @@ +import { randomUUID } from "node:crypto"; +import type { CustomEntry, SessionEntry, SessionHeader } from "./session-manager.ts"; + +/** Preserve Step's local statistics origin when the source file later disappears. */ +export function createUsageLineage( + source: SessionHeader | null, + full: readonly SessionEntry[], + copied: readonly SessionEntry[], + owner: string, + timestamp: string, + parentId: string | null, +): CustomEntry[] { + if (!source || !full.some((e) => e.type === "custom" && e.customType === "step.activity_capabilities")) return []; + const inherited = new Map(); + for (const e of full) + if (e.type === "custom" && e.customType === "step.usage_lineage") { + const d = e.data as { + version?: number; + sessionId?: string; + origins?: Record; + }; + if (d?.version === 1 && d.sessionId === source.id && d.origins) + for (const [id, origin] of Object.entries(d.origins)) + if ( + origin && + typeof origin.sessionId === "string" && + typeof origin.cwd === "string" && + ["user", "subagent"].includes(origin.role) + ) + inherited.set(id, origin); + } + const cap = full.find( + (e) => + e.type === "custom" && + e.customType === "step.activity_capabilities" && + (e.data as { sessionId?: string })?.sessionId === source.id, + ) as CustomEntry<{ sessionRole?: string; parentCwd?: string }> | undefined; + const role = + cap?.data?.sessionRole === "subagent" || source.id.startsWith("subagent-") || source.id.startsWith("workflow-") + ? "subagent" + : "user"; + const origin = { sessionId: source.id, cwd: cap?.data?.parentCwd ?? source.cwd, role } as const; + const copiedIds = new Set(copied.map((entry) => entry.id)); + const supplemental: CustomEntry[] = []; + for (const entry of full) { + if (entry.type !== "custom" || copiedIds.has(entry.id)) continue; + const data = entry.data as { usageEntryId?: string; sessionId?: string } | undefined; + if ( + (entry.customType === "step.usage_provenance" && + typeof data?.usageEntryId === "string" && + copiedIds.has(data.usageEntryId)) || + (entry.customType === "step.activity_capabilities" && data?.sessionId === source.id) + ) { + supplemental.push({ ...entry, parentId }); + parentId = entry.id; + } + } + const entries: CustomEntry[] = [...supplemental]; + let origins: Record = {}, + bytes = 0; + const flush = (): void => { + if (!Object.keys(origins).length) return; + const id = randomUUID(); + entries.push({ + type: "custom", + customType: "step.usage_lineage", + id, + parentId, + timestamp, + data: { version: 1, sessionId: owner, origins }, + }); + parentId = id; + origins = {}; + bytes = 0; + }; + for (const entry of [...copied, ...supplemental]) { + if (entry.type === "custom" && entry.customType === "step.usage_lineage") continue; + const value = inherited.get(entry.id) ?? origin; + const size = Buffer.byteLength(JSON.stringify([entry.id, value])); + if (size > 8192) continue; + if (bytes + size > 96 * 1024) flush(); + origins[entry.id] = value; + bytes += size; + } + flush(); + return entries; +} diff --git a/packages/coding-agent/src/core/usage-totals.ts b/packages/coding-agent/src/core/usage-totals.ts index 917d5d35..35c533e1 100644 --- a/packages/coding-agent/src/core/usage-totals.ts +++ b/packages/coding-agent/src/core/usage-totals.ts @@ -1,6 +1,21 @@ import type { Usage } from "@step-harness/providers/compat"; import type { SessionEntry } from "./session-manager.ts"; +export interface SessionUsage extends NormalizedUsage { + cacheHit: number | null; + invalid: number; + missing: number; +} + +export interface NormalizedUsage { + /** Complete input, including cache reads and writes. */ + input: number; + output: number; + cacheRead: number; + cacheWrite: number; + total: number; +} + export interface UsageTotals { input: number; output: number; @@ -19,12 +34,87 @@ export function createUsageTotals(): UsageTotals { }; } -export function addUsageToTotals(totals: UsageTotals, usage: Usage): void { - totals.input += usage.input; +export function normalizeUsage(raw: unknown): NormalizedUsage | undefined { + if (!raw || typeof raw !== "object") return undefined; + const u = raw as Record; + if ( + ![u.input, u.output, u.cacheRead, u.cacheWrite].every( + (n) => typeof n === "number" && Number.isSafeInteger(n) && n >= 0, + ) + ) + return undefined; + const input = (u.input as number) + (u.cacheRead as number) + (u.cacheWrite as number), + output = u.output as number; + if (!Number.isSafeInteger(input + output)) return undefined; + return { + input, + output, + cacheRead: u.cacheRead as number, + cacheWrite: u.cacheWrite as number, + total: input + output, + }; +} + +export function entryUsage(value: unknown): { raw?: unknown; expected: boolean; assistant: boolean } { + if (!value || typeof value !== "object") return { expected: false, assistant: false }; + const entry = value as Record; + const message = entry.message as Record | undefined; + if (entry.type === "message" && message?.role === "assistant") + return { raw: message.usage, expected: true, assistant: true }; + if (entry.type === "message" && message?.role === "toolResult") + return { raw: message.usage, expected: message.usage !== undefined, assistant: false }; + if (entry.type === "compaction" || entry.type === "branch_summary") + return { raw: entry.usage, expected: entry.usage !== undefined, assistant: false }; + if (entry.type === "custom" && entry.customType === "step-auto-clm-usage") { + const data = entry.data as Record | undefined; + return { raw: data?.usage, expected: true, assistant: false }; + } + return { expected: false, assistant: false }; +} + +/** Add one validated increment. Input here excludes cache reads/writes, as in provider Usage. */ +export function addUsageToTotals(totals: UsageTotals, raw: unknown): boolean { + const usage = normalizeUsage(raw); + if ( + !usage || + !Number.isSafeInteger(totals.input + totals.output + totals.cacheRead + totals.cacheWrite + usage.total) + ) + return false; + totals.input += usage.input - usage.cacheRead - usage.cacheWrite; totals.output += usage.output; totals.cacheRead += usage.cacheRead; totals.cacheWrite += usage.cacheWrite; - totals.cost += usage.cost.total; + const cost = (raw as { cost?: { total?: unknown } }).cost?.total; + if (typeof cost === "number" && Number.isFinite(cost) && cost >= 0 && Number.isFinite(totals.cost + cost)) + totals.cost += cost; + return true; +} + +/** Project the shared session ledger into the Account panel's complete-input convention. */ +export function collectSessionUsage( + entries: readonly SessionEntry[] | readonly Record[], +): SessionUsage { + const totals = createUsageTotals(); + let invalid = 0, + missing = 0; + for (const entry of entries) { + const picked = entryUsage(entry); + if (!addUsageToTotals(totals, picked.raw) && picked.expected) { + if (picked.raw === undefined) missing++; + else invalid++; + } + } + const input = totals.input + totals.cacheRead + totals.cacheWrite; + return { + input, + output: totals.output, + cacheRead: totals.cacheRead, + cacheWrite: totals.cacheWrite, + total: input + totals.output, + cacheHit: input > 0 ? totals.cacheRead / input : null, + invalid, + missing, + }; } /** Maintenance requests do not become ordinary task messages; their billing still belongs to the session. */ @@ -32,13 +122,7 @@ export function getAutoClmUsage(entry: SessionEntry): Usage | undefined { if (entry.type !== "custom" || entry.customType !== "step-auto-clm-usage") return undefined; const data = entry.data as { usage?: Usage } | undefined; const usage = data?.usage; - if ( - !usage || - !usage.cost || - ![usage.input, usage.output, usage.cacheRead, usage.cacheWrite, usage.cost.total].every( - (value) => Number.isFinite(value) && value >= 0, - ) - ) + if (!usage?.cost || !normalizeUsage(usage) || !Number.isFinite(usage.cost.total) || usage.cost.total < 0) return undefined; return usage; } @@ -54,29 +138,19 @@ export function getUsageCostBreakdown(entries: SessionEntry[]): UsageCostBreakdo const totalsByKey = new Map(); for (const entry of entries) { - let key: string | undefined; - let usage: Usage | undefined; - if (entry.type === "message" && entry.message.role === "assistant") { - key = `${entry.message.provider}/${entry.message.responseModel ?? entry.message.model}`; - usage = entry.message.usage; - } else if (entry.type === "message" && entry.message.role === "toolResult" && entry.message.usage) { - key = "Tools/summaries"; - usage = entry.message.usage; - } else if ((entry.type === "branch_summary" || entry.type === "compaction") && entry.usage) { - key = "Tools/summaries"; - usage = entry.usage; - } else { - usage = getAutoClmUsage(entry); - if (usage) key = "Tools/summaries"; - } - if (!key || !usage) continue; + const picked = entryUsage(entry); + if (!picked.expected) continue; + const key = + entry.type === "message" && entry.message.role === "assistant" + ? `${entry.message.provider}/${entry.message.responseModel ?? entry.message.model}` + : "Tools/summaries"; let totals = totalsByKey.get(key); if (!totals) { totals = createUsageTotals(); totalsByKey.set(key, totals); } - addUsageToTotals(totals, usage); + addUsageToTotals(totals, picked.raw); } return Array.from(totalsByKey, ([key, totals]) => ({ diff --git a/packages/coding-agent/src/features/step-plan.ts b/packages/coding-agent/src/features/step-plan.ts index cf568888..ebf178ec 100644 --- a/packages/coding-agent/src/features/step-plan.ts +++ b/packages/coding-agent/src/features/step-plan.ts @@ -55,11 +55,18 @@ function unique(values: readonly string[]): string[] { return [...new Set(values)]; } +const publishedPlanModes = new WeakMap(); +export function getStepPlanMode(manager: object): string { + const value = publishedPlanModes.get(manager); + return value === undefined ? "Unknown" : value ? "Plan" : "Default"; +} + /** * Status-bar chip: "Planning (user|agent|unknown)" while plan mode is on, * cleared (normal) otherwise. */ function updateStatus(ctx: ExtensionContext, enabled: boolean, planSource: StepPlanSource | undefined): void { + publishedPlanModes.set(ctx.sessionManager, enabled); if (enabled) { ctx.ui.setStatus("plan-mode", ctx.ui.theme.fg("warning", `Planning (${planSource ?? "unknown"})`)); } else { diff --git a/packages/coding-agent/src/features/step-subagent.ts b/packages/coding-agent/src/features/step-subagent.ts index 0416da58..c32bfe86 100644 --- a/packages/coding-agent/src/features/step-subagent.ts +++ b/packages/coding-agent/src/features/step-subagent.ts @@ -107,6 +107,8 @@ export interface StepSubagentRunResult { export interface StepSubagentRunInput { onModelRequest?: (event: import("../core/rate-limit.ts").ModelRequestProgress) => void; + parentCwd?: string; + parentSessionId?: string; agent: StepAgentConfig; task: string; cwd: string; diff --git a/packages/coding-agent/src/features/step.ts b/packages/coding-agent/src/features/step.ts index b722753c..853d5836 100644 --- a/packages/coding-agent/src/features/step.ts +++ b/packages/coding-agent/src/features/step.ts @@ -43,6 +43,8 @@ import { import { registerStepStreamRecovery } from "./step-stream-recovery.ts"; export interface StepExtensionOptions { + /** Local usage services and presentation supplied by the product host. */ + usageHost?: import("../step/usage/commands.ts").StepUsageHost; /** Initial Step policy, normally populated from CLI/runtime options. */ permission?: StepPermissionControllerOptions; /** Optional Step settings decorator used to restore and persist policy. */ @@ -122,7 +124,14 @@ export function createStepExtension(options: StepExtensionOptions = {}): Extensi // The same plugins also contribute skills and commands, which the resource // loader reads from agent/project directories rather than the plugin root. createStepPluginResourcesExtension()(pi); - registerStepPiCommandAdapters(pi, options.telemetry, options.stepSettings, options.feedbackIdentity); + registerStepPiCommandAdapters( + pi, + options.telemetry, + options.stepSettings, + options.feedbackIdentity, + options.usageHost, + () => getStepPermissionPreset(permissions.getState().preset)?.label ?? "Unknown", + ); const builtInSlashCommands = new Set(BUILTIN_SLASH_COMMANDS.map((command) => command.name)); // Extension commands are wrapped at registration time below. Inputs that // reach Pi's normal prompt path are the remaining (usually unknown) slash diff --git a/packages/coding-agent/src/features/subagent/execute.ts b/packages/coding-agent/src/features/subagent/execute.ts index 015633d8..7c12fd9d 100644 --- a/packages/coding-agent/src/features/subagent/execute.ts +++ b/packages/coding-agent/src/features/subagent/execute.ts @@ -264,6 +264,8 @@ export async function executeSubagent( }, laneRuntime?: StepSubagentLaneRuntime, ): Promise> { + const parentCwd = ctx.cwd, + parentSessionId = ctx.sessionManager?.getSessionId?.(); const scope = params.agentScope ?? "user"; const discovery = await discoverStepAgents(ctx.cwd, { agentDir: options.agentDir, @@ -417,6 +419,8 @@ export async function executeSubagent( emit(mode); }; const child = await options.runner({ + parentCwd, + parentSessionId, agent, task: task.task, cwd: childCwd, diff --git a/packages/coding-agent/src/features/subagent/rpc-adapter.ts b/packages/coding-agent/src/features/subagent/rpc-adapter.ts index 72bb9304..e77e4141 100644 --- a/packages/coding-agent/src/features/subagent/rpc-adapter.ts +++ b/packages/coding-agent/src/features/subagent/rpc-adapter.ts @@ -200,6 +200,8 @@ export function buildSubagentChildEnv(input: StepSubagentRunInput): NodeJS.Proce return { ...process.env, [CHILD_MARKER]: "1", + STEPCODE_USAGE_PARENT_CWD: input.parentCwd ?? input.workflowAcl?.baseCwd, + STEPCODE_USAGE_PARENT_SESSION_ID: input.parentSessionId, STEP_DISABLE_WORKFLOW: "1", STEP_DISABLE_CRON: "1", STEP_DISABLE_GOAL: "1", diff --git a/packages/coding-agent/src/index.ts b/packages/coding-agent/src/index.ts index b41b0226..9451c774 100644 --- a/packages/coding-agent/src/index.ts +++ b/packages/coding-agent/src/index.ts @@ -441,7 +441,12 @@ export { type ProjectTrustStoreEntry, type ProjectTrustUpdate, } from "./core/trust-manager.ts"; -export { addUsageToTotals, createUsageTotals, getUsageCostBreakdown } from "./core/usage-totals.ts"; +export { + addUsageToTotals, + collectSessionUsage, + createUsageTotals, + getUsageCostBreakdown, +} from "./core/usage-totals.ts"; export { createStepExtension, createStepExtensionInline, @@ -596,6 +601,7 @@ export { STEP_INIT_PROMPT } from "./step/init-prompt.ts"; export { prepareStepInstallation } from "./step/installation.ts"; export { maybeUpdateStep, runStepUpdateCommand } from "./step/local-update.ts"; export { + checkStepPlanBeforeInteractive, isStepInteractiveLoginStartup, needsStepLoginBeforeInteractive, readStepLoginCredential, @@ -755,6 +761,16 @@ export { export { buildStepThemeOptions, runStepThemePrompt, type StepThemeOption } from "./step/theme-prompt.ts"; export { createStepToolProfile } from "./step/tool-profile.ts"; export { resolveStepTraceHeaderBaseUrls } from "./step/trace-headers.ts"; +export { decodePlanUsage, PlanUsageController, resolvePlanUsageAccount } from "./step/usage/account.ts"; +export { overviewDates, scanUsageOverview } from "./step/usage/activity.ts"; +export type { StepUsageHost } from "./step/usage/commands.ts"; +export type { + AccountUsageState, + PlanUsageSnapshot, + UsageOverview, + UsagePresentation, + UsageStatus, +} from "./step/usage/types.ts"; export { resolveStepCodeVersion, STEPCODE_BUILD_VERSION_ENV, diff --git a/packages/coding-agent/src/step/login-flow.ts b/packages/coding-agent/src/step/login-flow.ts index 73f8d2c1..bedb09b2 100644 --- a/packages/coding-agent/src/step/login-flow.ts +++ b/packages/coding-agent/src/step/login-flow.ts @@ -20,6 +20,8 @@ import { type StepLoginStep, } from "./onboarding.ts"; import { StepOnboardingView } from "./onboarding-view.ts"; +import { PlanUsageController, resolvePlanUsageAccount } from "./usage/account.ts"; +import type { AccountUsageState } from "./usage/types.ts"; import { resolveStepCodeVersion } from "./version.ts"; export interface StepLoginHost { @@ -36,6 +38,7 @@ export interface StepLoginOutcome { readonly kind: "completed" | "exit"; readonly profile?: StepLoginProfile; readonly credentialsPath?: string; + readonly warning?: string; } export interface RunStepLoginOptions { @@ -49,6 +52,48 @@ export interface RunStepLoginOptions { * containers, where a launcher may "succeed" without a window ever appearing. */ readonly noBrowser?: boolean; + /** Reuse the process account controller, including identity-scoped caching. */ + readonly planUsage?: PlanUsageController; + /** A confirmed missing plan discovered during interactive startup. */ + readonly requiredPlan?: "step_plan" | "step_plan_oversea"; +} + +/** Model discovery can return fallback models even when the account has no plan. */ +export function getMissingStepPlanProfile(state: AccountUsageState): "step_plan" | "step_plan_oversea" | undefined { + const { snapshot, account } = state; + if (state.status !== "ready" || !snapshot || !account.eligible) return undefined; + if (account.profile !== "step_plan" && account.profile !== "step_plan_oversea") return undefined; + return snapshot.plan.state === "known" && + snapshot.plan.family === "none" && + snapshot.plan.subscription === null && + snapshot.quota.state === "known" && + !snapshot.quota.allowed && + snapshot.quota.reason === "no_active_subscription" && + snapshot.quota.credit_buckets?.length === 0 + ? account.profile + : undefined; +} + +export function stepPlanCheckWarning(state: AccountUsageState): string | undefined { + return state.account.eligible && + (state.status === "error" || + state.status === "stale" || + (state.snapshot?.plan.state === "unavailable" && !state.snapshot.quota.allowed)) + ? `Could not verify Step Plan availability. ${state.message ?? "Retry with /usage account when online."}` + : undefined; +} + +export async function checkStepPlanBeforeInteractive(input: { + readonly authPath: string; + readonly interactive: boolean; + readonly planUsage: PlanUsageController; + readonly env?: Record; +}): Promise<{ requiredPlan?: "step_plan" | "step_plan_oversea"; warning?: string }> { + if (!input.interactive || (input.env ?? process.env).STEP_API_KEY?.trim()) return {}; + const profile = readStepLoginProfile(input.authPath); + if (profile !== "step_plan" && profile !== "step_plan_oversea") return {}; + const state = await input.planUsage.query(); + return { requiredPlan: getMissingStepPlanProfile(state), warning: stepPlanCheckWarning(state) }; } export async function writeStepLoginCredential(input: { @@ -174,10 +219,16 @@ export async function runStepLogin(options: RunStepLoginOptions = {}): Promise resolvePlanUsageAccount(authPath, options.env)); let browserAbort: AbortController | undefined; - let savedProfile: StepLoginProfile | undefined; - let savedPath: string | undefined; + let savedProfile = options.requiredPlan + ? profiles.find((profile) => profile.id === options.requiredPlan) + : undefined; + let savedPath = options.requiredPlan ? authPath : undefined; + let warning: string | undefined; let settle: ((outcome: StepLoginOutcome) => void) | undefined; const finished = new Promise((resolve) => { settle = resolve; @@ -190,6 +241,13 @@ export async function runStepLogin(options: RunStepLoginOptions = {}): Promise dispatch({ type: "backspace" }), onBack: () => dispatch({ type: "back" }), onQuit: () => dispatch({ type: "quit" }), + onRetryPlan: () => dispatch({ type: "checkPlan" }), + onOpenPlan: () => { + if (step.kind !== "planRequired" || options.noBrowser) return; + const choice = step.choice; + const profile = profiles.find((candidate) => candidate.id === choice); + if (profile) openBrowser(`${profile.authBaseUrl}/step-plan`); + }, requestRender: () => host.requestRender(), }); @@ -197,6 +255,7 @@ export async function runStepLogin(options: RunStepLoginOptions = {}): Promise { + const saving = step; try { const profile = profiles.find((candidate) => candidate.id === choice); if (!profile) throw new Error(`Unknown Step login profile: ${choice}`); @@ -275,11 +348,14 @@ export async function runStepLogin(options: RunStepLoginOptions = {}): Promise new Date()))().toISOString(), }); + if (step !== saving) return; savedProfile = profile; savedPath = authPath; + warning = undefined; syncStepLoginProfileEndpoint(authPath, options.env ?? process.env); - dispatch({ type: "saved" }); + dispatch({ type: profile.credentialSource === "browser" ? "checkPlan" : "saved" }); } catch (error) { + if (step !== saving) return; dispatch({ type: "fail", message: `Could not save credentials: ${error instanceof Error ? error.message : String(error)}`, @@ -287,6 +363,7 @@ export async function runStepLogin(options: RunStepLoginOptions = {}): Promise {}; + const pageToolsPublished = new Promise((resolve) => { + published = resolve; + }); + harness.runtime.refreshTools = () => { + if (harness.extension.tools.size === 2) published(); + }; await harness.start(); await vi.waitFor(() => expect(readiness).toHaveBeenCalledOnce()); expect(readiness).toHaveBeenCalledWith( @@ -612,7 +619,10 @@ require("node:readline").createInterface({ input: process.stdin }).on("line", li await expect(readFile(started)).rejects.toMatchObject({ code: "ENOENT" }); expect(getStepMcpStatuses()).toContainEqual({ name: "steppage__steppage", status: "connecting", toolCount: 0 }); release({ command, installed: true }); - await vi.waitFor(() => expect(harness.extension.tools.size).toBe(2)); + // Real stdio startup can exceed waitFor's default one-second budget under + // load. Await publication within this test's unchanged overall timeout. + await pageToolsPublished; + expect(harness.extension.tools.size).toBe(2); expect(JSON.parse(await readFile(started, "utf8"))).toEqual({ setupSecret: null, installerUrl: null }); expect(getStepMcpStatuses()).toContainEqual({ name: "steppage__steppage", status: "connected", toolCount: 1 }); expect(harness.notify).not.toHaveBeenCalled(); diff --git a/packages/coding-agent/src/step/mcp.ts b/packages/coding-agent/src/step/mcp.ts index 055d0d3e..e6680b2d 100644 --- a/packages/coding-agent/src/step/mcp.ts +++ b/packages/coding-agent/src/step/mcp.ts @@ -37,6 +37,7 @@ const STEPPAGE_DEPLOY_TOOL_NAME = "page_deploy"; const STEPPAGE_MANAGEMENT_URL = "https://platform.stepfun.com/sites"; interface ConnectedServer { + pluginId?: string; readonly name: string; readonly client: Client; readonly transport: StdioClientTransport | StreamableHTTPClientTransport; @@ -71,6 +72,7 @@ interface ServerDeclaration { } interface DiscoveredServer { + pluginId?: string; name: string; declaration: ServerDeclaration; provision?: StepPluginProvision; @@ -83,6 +85,19 @@ export interface StepMcpStatus { } let currentMcpStatuses: StepMcpStatus[] = []; +const currentMcpToolOwners = new WeakMap>(); +export function getStepMcpToolOwner(toolName: string, session: object): string | undefined { + return currentMcpToolOwners.get(session)?.get(toolName); +} +function publishMcpToolOwners( + owners: Map, + server: ConnectedServer, + tools: McpTool[], + previous: McpTool[] = [], +): void { + for (const tool of previous) owners.delete(remoteToolName(server, tool)); + if (server.pluginId) for (const tool of tools) owners.set(remoteToolName(server, tool), server.pluginId); +} export function getStepMcpStatuses(): StepMcpStatus[] { return currentMcpStatuses.map((status) => ({ ...status })); @@ -105,6 +120,7 @@ export function formatStepMcpStatuses(statuses: readonly StepMcpStatus[] = curre export function createStepMcpExtension(): ExtensionFactory { return (pi: ExtensionAPI): void => { let servers: ConnectedServer[] = []; + let owners = new Map(); let startup: Promise | undefined; let cancellation: AbortController | undefined; let statuses: StepMcpStatus[] = []; @@ -115,6 +131,9 @@ export function createStepMcpExtension(): ExtensionFactory { cancellation = controller; statuses = []; currentMcpStatuses = statuses; + owners = new Map(); + if (ctx.sessionManager) currentMcpToolOwners.set(ctx.sessionManager, owners); + const sessionOwners = owners; startup = (async () => { // Finish mounting the interactive session before doing discovery or // spawning processes. Awaiting Promise.all in session_start kept the @@ -169,6 +188,7 @@ export function createStepMcpExtension(): ExtensionFactory { pi.registerTools(remoteTools, { remove: server.tools.map((tool) => remoteToolName(server, tool)), }); + publishMcpToolOwners(sessionOwners, server, tools, server.tools); server.tools = tools; statuses[index] = { name: item.name, status: "connected", toolCount: tools.length }; } catch (error) { @@ -189,6 +209,7 @@ export function createStepMcpExtension(): ExtensionFactory { }; try { const server = await connectStepMcpServer(item, controller.signal, onToolsChanged); + server.pluginId = item.pluginId; connected = server; const remoteTools = server.tools.map((tool) => createRemoteTool(server, tool)); // Publishing a server's catalog refreshes the registry and the @@ -197,12 +218,14 @@ export function createStepMcpExtension(): ExtensionFactory { await yieldToEventLoop(); server.signal.throwIfAborted(); pi.registerTools(remoteTools); + publishMcpToolOwners(sessionOwners, server, server.tools); published = true; servers.push(server); server.signal.addEventListener( "abort", () => { published = false; + publishMcpToolOwners(sessionOwners, server, [], server.tools); if (controller.signal.aborted) return; pi.registerTools([], { remove: server.tools.map((tool) => remoteToolName(server, tool)) }); statuses[index] = { name: item.name, status: "failed", toolCount: 0 }; @@ -251,6 +274,7 @@ export function createStepMcpExtension(): ExtensionFactory { cancellation?.abort(); await startup; const closing = servers; + for (const server of closing) publishMcpToolOwners(owners, server, [], server.tools); servers = []; startup = undefined; if (currentMcpStatuses === statuses) currentMcpStatuses = []; @@ -370,8 +394,10 @@ export async function discoverStepMcpServers(cwd: string, projectTrusted: boolea global && isManagedStepPageCommand(global.declaration, parsed.manifest.provision) && isManagedStepPageCommand(value, parsed.manifest.provision) - ) + ) { global.provision = parsed.manifest.provision; + global.pluginId = parsed.manifest.id; + } continue; } if (value.enabled === false) { @@ -381,6 +407,7 @@ export async function discoverStepMcpServers(cwd: string, projectTrusted: boolea if (!hasTransport(value)) continue; seen.add(name); const discovered: DiscoveredServer = { + pluginId: parsed.manifest.id, name, declaration: anchorPluginServerCwd( pluginDir, diff --git a/packages/coding-agent/src/step/onboarding-view.ts b/packages/coding-agent/src/step/onboarding-view.ts index 84d10428..6681af20 100644 --- a/packages/coding-agent/src/step/onboarding-view.ts +++ b/packages/coding-agent/src/step/onboarding-view.ts @@ -23,6 +23,8 @@ export interface StepOnboardingViewCallbacks { onBackspace(): void; onBack(): void; onQuit(): void; + onRetryPlan?(): void; + onOpenPlan?(): void; requestRender(): void; } @@ -31,6 +33,7 @@ export class StepOnboardingView extends Container implements Component, Focusabl private readonly profiles: readonly StepLoginProfile[]; private readonly callbacks: StepOnboardingViewCallbacks; private readonly selectList: SelectList; + private readonly planActions: SelectList; private readonly input: Input; private step: StepLoginStep = { kind: "pickMode", error: null }; private focusedState = false; @@ -60,6 +63,31 @@ export class StepOnboardingView extends Container implements Component, Focusabl if (profile) this.callbacks.onChoose(profile.id); }; this.selectList.onCancel = () => this.callbacks.onQuit(); + this.planActions = new SelectList( + [ + { value: "key", label: "1. Use an API key" }, + { value: "subscribe", label: "2. Open Step Plan subscription page" }, + { value: "retry", label: "3. Check again" }, + { value: "account", label: "4. Change account or login method" }, + ], + 4, + { + selectedPrefix: (text) => theme.fg("accent", text), + selectedText: (text) => theme.fg("accent", text), + description: (text) => theme.fg("muted", text), + scrollInfo: (text) => theme.fg("muted", text), + noMatch: (text) => theme.fg("muted", text), + }, + ); + this.planActions.onSelect = (item) => { + if (this.step.kind !== "planRequired") return; + if (item.value === "key") + this.callbacks.onChoose(this.step.choice === "step_plan_oversea" ? "platform_oversea" : "platform_cn"); + else if (item.value === "subscribe") this.callbacks.onOpenPlan?.(); + else if (item.value === "retry") this.callbacks.onRetryPlan?.(); + else this.callbacks.onBack(); + }; + this.planActions.onCancel = () => this.callbacks.onQuit(); this.input = new Input(); this.input.onSubmit = (value) => this.callbacks.onSubmitApiKey(value); @@ -77,6 +105,7 @@ export class StepOnboardingView extends Container implements Component, Focusabl } setStep(step: StepLoginStep): void { + if (step.kind === "planRequired" && this.step.kind !== step.kind) this.planActions.setSelectedIndex(0); this.step = step; this.input.setValue(step.kind === "apiKeyEntry" ? step.value : ""); this.input.focused = this.focusedState && step.kind === "apiKeyEntry"; @@ -109,6 +138,14 @@ export class StepOnboardingView extends Container implements Component, Focusabl this.selectList.handleInput(data); return; } + if (this.step.kind === "planRequired") { + if (data === "q") this.callbacks.onQuit(); + else if (/^[1-4]$/u.test(data)) { + this.planActions.setSelectedIndex(Number(data) - 1); + this.planActions.handleInput("\r"); + } else this.planActions.handleInput(data); + return; + } if (this.step.kind === "apiKeyEntry") { if (keybindings.matches(data, "tui.select.cancel")) { this.callbacks.onBack(); @@ -121,7 +158,10 @@ export class StepOnboardingView extends Container implements Component, Focusabl else if (after.length < before.length) this.callbacks.onBackspace(); return; } - if (this.step.kind === "continueInBrowser" && keybindings.matches(data, "tui.select.cancel")) { + if ( + (this.step.kind === "continueInBrowser" || this.step.kind === "checkingPlan") && + keybindings.matches(data, "tui.select.cancel") + ) { this.callbacks.onBack(); } } @@ -131,7 +171,7 @@ export class StepOnboardingView extends Container implements Component, Focusabl const rows = this.renderRows(safeWidth); // In particular, never truncate the authorization URL or its query string. // Remote users need every character, including on a narrow terminal. - if (this.step.kind === "continueInBrowser") { + if (this.step.kind === "continueInBrowser" || this.step.kind === "planRequired") { return rows.flatMap((row) => wrapTextWithAnsi(row, safeWidth)); } return rows.map((row) => truncateToWidth(row, safeWidth, "", false)); @@ -154,9 +194,11 @@ export class StepOnboardingView extends Container implements Component, Focusabl rows.push("", muted(" ↑/↓ select · Enter continue · q quit")); return rows; } - case "apiKeyEntry": + case "apiKeyEntry": { + const choice = this.step.choice; rows.push( "Use your own API key for usage-based billing", + muted(this.profiles.find((profile) => profile.id === choice)?.title ?? ""), "", "Paste or type your API key below. It is stored locally, outside config.json.", "", @@ -165,6 +207,7 @@ export class StepOnboardingView extends Container implements Component, Focusabl if (this.step.error) rows.push(theme.fg("error", this.step.error)); rows.push("", muted(" Enter submit · Esc back")); return rows; + } case "continueInBrowser": rows.push( "Continue sign-in in your browser:", @@ -178,6 +221,23 @@ export class StepOnboardingView extends Container implements Component, Focusabl case "saving": rows.push(muted("Saving credentials...")); return rows; + case "checkingPlan": + rows.push(muted("Checking Step Plan subscription and available credits..."), "", muted(" Esc back")); + return rows; + case "planRequired": { + const choice = this.step.choice; + const profile = this.profiles.find((profile) => profile.id === choice); + rows.push( + theme.fg("warning", "No active Step Plan subscription or usable top-up credits."), + muted(this.step.choice === "step_plan_oversea" ? "Account region: Global" : "Account region: CN"), + "", + ...this.planActions.render(Math.max(1, width - 2)), + "", + profile ? theme.fg("accent", `${profile.authBaseUrl}/step-plan`) : "", + muted(" ↑/↓ select · Enter continue · Esc/q quit"), + ); + return rows; + } case "done": case "exit": return []; diff --git a/packages/coding-agent/src/step/onboarding.ts b/packages/coding-agent/src/step/onboarding.ts index eb2f6ae1..f3c4b507 100644 --- a/packages/coding-agent/src/step/onboarding.ts +++ b/packages/coding-agent/src/step/onboarding.ts @@ -119,6 +119,8 @@ export type StepLoginStep = readonly authUrl: string; } | { readonly kind: "saving"; readonly choice: StepLoginChoice } + | { readonly kind: "checkingPlan"; readonly choice: StepLoginChoice } + | { readonly kind: "planRequired"; readonly choice: StepLoginChoice } | { readonly kind: "done" } | { readonly kind: "exit" }; @@ -134,6 +136,8 @@ export type StepLoginEvent = | { readonly type: "backspace" } | { readonly type: "back" } | { readonly type: "saved" } + | { readonly type: "checkPlan" } + | { readonly type: "planRequired" } | { readonly type: "fail"; readonly message: string } | { readonly type: "quit" }; @@ -174,9 +178,26 @@ export function reduceStepLogin(step: StepLoginStep, event: StepLoginEvent): Ste if (event.type === "fail") return { kind: "pickMode", error: event.message }; return step; case "saving": + if (event.type === "checkPlan") return { kind: "checkingPlan", choice: step.choice }; if (event.type === "saved") return { kind: "done" }; if (event.type === "fail") return { kind: "pickMode", error: event.message }; return step; + case "checkingPlan": + if (event.type === "planRequired") return { kind: "planRequired", choice: step.choice }; + if (event.type === "saved") return { kind: "done" }; + if (event.type === "back") return { kind: "pickMode", error: null }; + return step; + case "planRequired": + if (event.type === "choose") + return { + kind: "apiKeyEntry", + choice: step.choice === "step_plan_oversea" ? "platform_oversea" : "platform_cn", + value: "", + error: null, + }; + if (event.type === "checkPlan") return { kind: "checkingPlan", choice: step.choice }; + if (event.type === "back") return { kind: "pickMode", error: null }; + return step; case "done": case "exit": return step; diff --git a/packages/coding-agent/src/step/plugins.ts b/packages/coding-agent/src/step/plugins.ts index d36aec2e..c4f5242c 100644 --- a/packages/coding-agent/src/step/plugins.ts +++ b/packages/coding-agent/src/step/plugins.ts @@ -711,11 +711,18 @@ function provisionedServerEnvironments(manifest: StepPluginManifest): Array>(); +export function getStepPluginSkillOwner(filePath: string, session: object): string | undefined { + for (const [root, owner] of pluginSkillOwners.get(session) ?? []) if (isPathContained(root, filePath)) return owner; + return undefined; +} + export async function discoverStepPluginResourcePaths(input: { userDir?: string; projectDir?: string; /** Project plugins are read only when the project is trusted, as MCP discovery does. */ projectTrusted?: boolean; + onSkillOwner?: (root: string, owner: string) => void; }): Promise<{ skillPaths: string[]; promptPaths: string[]; warnings: string[] }> { const warnings: string[] = []; const seenByKind = { skills: new Set(), commands: new Set() }; @@ -757,6 +764,7 @@ export async function discoverStepPluginResourcePaths(input: { if (seenByKind[key].has(canonical)) continue; seenByKind[key].add(canonical); collected[key].push(canonical); + if (key === "skills") input.onSkillOwner?.(await fs.realpath(canonical), manifest.id); } } } @@ -1178,15 +1186,20 @@ interface InteractivePluginOptions extends StepPluginCommandOptions { * thrown: a broken plugin must not cost the session its other resources. */ export function createStepPluginResourcesExtension(): ExtensionFactory { + let generation = 0; return (pi: ExtensionAPI): void => { pi.on("resources_discover", async (_event, ctx) => { const projectDir = defaultStepPluginsDir(process.env, { cwd: ctx.cwd, project: true }); + const attempt = ++generation, + owners = new Map(); try { const discovered = await discoverStepPluginResourcePaths({ projectDir, projectTrusted: ctx.isProjectTrusted(), + onSkillOwner: (root, owner) => owners.set(root, owner), userDir: defaultStepPluginsDir(process.env), }); + if (attempt === generation && ctx.sessionManager) pluginSkillOwners.set(ctx.sessionManager, owners); // Surface a malformed declaration once, on the channel the loader // already uses for resource diagnostics, rather than throwing: the // rest of the plugins' resources must still load. diff --git a/packages/coding-agent/src/step/slash-commands.ts b/packages/coding-agent/src/step/slash-commands.ts index 130f0cfa..5fa22873 100644 --- a/packages/coding-agent/src/step/slash-commands.ts +++ b/packages/coding-agent/src/step/slash-commands.ts @@ -26,6 +26,7 @@ import { formatStepMcpStatuses } from "./mcp.ts"; import { registerStepPluginCommand } from "./plugins.ts"; import { resolveStepStorageRoot } from "./storage-root.ts"; import { type StepTelemetryReporter, trackStepTelemetry } from "./telemetry.ts"; +import { registerStepUsageCommands, type StepUsageHost } from "./usage/commands.ts"; /** * Register Step's product command spellings that have a direct Pi action. @@ -40,7 +41,24 @@ export function registerStepPiCommandAdapters( telemetry?: StepTelemetryReporter, stepSettings?: () => { getStepSettings(): { feedbackEnabled?: boolean } } | undefined, feedbackIdentity?: () => FeedbackIdentity, + usageHost?: StepUsageHost, + permissions: () => string = () => "Unknown", ): void { + const usage = usageHost ? registerStepUsageCommands(pi, usageHost, permissions) : undefined; + if (usage) + registerTrackedCommand( + pi, + "usage", + { + description: "Show local activity or personal account quota", + getArgumentCompletions: (prefix) => + ["overview", "account", "--refresh", "account --refresh"] + .filter((value) => value.startsWith(prefix)) + .map((value) => ({ value, label: value })), + handler: usage.usage, + }, + telemetry, + ); registerStepPluginCommand(pi, { telemetry }); registerTrackedCommand( pi, @@ -92,8 +110,9 @@ export function registerStepPiCommandAdapters( "status", { description: "Show current session and model status", - handler: async (_args, ctx) => { - ctx.ui.notify(formatStepStatus(ctx), "info"); + handler: async (args, ctx) => { + if (usage) await usage.status(args, ctx); + else ctx.ui.notify(formatStepStatus(ctx), "info"); }, }, telemetry, diff --git a/packages/coding-agent/src/step/usage/account.ts b/packages/coding-agent/src/step/usage/account.ts new file mode 100644 index 00000000..9fef6b2f --- /dev/null +++ b/packages/coding-agent/src/step/usage/account.ts @@ -0,0 +1,424 @@ +import { createHash } from "node:crypto"; +import { readStoredCredential } from "../../core/auth-storage.ts"; +import { stripAnsi } from "../../utils/ansi.ts"; +import { resolveStepLoginProfiles } from "../onboarding.ts"; +import type { AccountUsageState, LocalPlanAccount, PlanAccountContext, PlanUsageSnapshot } from "./types.ts"; + +export function resolvePlanUsageAccount( + authPath: string, + env: Record = process.env, + explicitKey = false, +): PlanAccountContext { + const stored = readStoredCredential("step", authPath) as + | { type?: string; access?: unknown; key?: unknown; uid?: unknown; profile?: unknown } + | undefined; + const key = stored?.type === "oauth" && typeof stored.access === "string" ? stored.access.trim() : ""; + const profile = typeof stored?.profile === "string" && stored.profile ? stored.profile : "step_plan"; + const uid = typeof stored?.uid === "string" && /^\d+$/u.test(stored.uid) ? stored.uid : undefined; + let projection: LocalPlanAccount = { + presence: + key || + (stored?.type === "api_key" && typeof stored.key === "string" && stored.key.trim()) || + env.STEP_API_KEY?.trim() || + explicitKey + ? "present" + : "missing", + eligible: false, + source: "stored", + profile, + accountId: uid, + }; + if (explicitKey || env.STEP_API_KEY?.trim()) + projection = { + ...projection, + source: "gateway", + accountId: undefined, + profile: undefined, + reason: "Current API key is a gateway/model credential. Sign in with a personal Plan profile to query usage.", + }; + else if (!key) projection = { ...projection, reason: "Sign in with step login to query account usage." }; + else if (!["step_plan", "step_plan_oversea"].includes(profile)) + projection = { ...projection, reason: "This login profile does not support personal Plan usage." }; + else { + const login = resolveStepLoginProfiles(env).find((p) => p.id === profile); + try { + const base = new URL(login?.baseUrl ?? ""); + if ( + base.username || + base.password || + base.search || + base.hash || + !["https:", "http:"].includes(base.protocol) || + (base.protocol === "http:" && !["localhost", "127.0.0.1", "[::1]"].includes(base.hostname)) + ) + throw new Error("endpoint"); + const path = base.pathname.replace(/\/+$/u, ""); + base.pathname = `${path}${path.endsWith("/v1") ? "" : "/v1"}/usage`; + projection = { + ...projection, + eligible: true, + source: "plan", + region: profile === "step_plan_oversea" ? "Global" : "CN", + endpoint: base.href, + }; + } catch { + projection = { ...projection, reason: "Account endpoint is not valid for usage queries." }; + } + } + const identity = createHash("sha256") + .update( + JSON.stringify([projection.source, profile, uid, projection.endpoint, key, env.STEP_API_KEY, explicitKey]), + ) + .digest("hex"); + return { projection, identity, ...(projection.eligible ? { key } : {}) }; +} + +export function usageText(value: string): string { + return stripAnsi(value) + .replace(/[\u0000-\u001f\u007f-\u009f\u2028\u2029]/gu, "") + .slice(0, 256); +} +function record(value: unknown): Record { + if (!value || typeof value !== "object" || Array.isArray(value)) throw new Error("Invalid usage snapshot"); + return value as Record; +} +function decimal(value: unknown, positive = false): string { + if ( + typeof value !== "string" || + !/^\d{1,19}$/u.test(value) || + BigInt(value) > 9223372036854775807n || + (positive && BigInt(value) === 0n) + ) + throw new Error("Invalid usage snapshot"); + return value; +} +function time(value: unknown): number | null { + if (value === null) return null; + if (typeof value !== "number" || !Number.isSafeInteger(value) || value <= 0 || value > 8640000000000000) + throw new Error("Invalid usage snapshot"); + return value; +} +function percent(value: unknown): number | null { + if (value === null) return null; + if (typeof value !== "number" || !Number.isFinite(value) || value < 0 || value > 100) + throw new Error("Invalid usage snapshot"); + return value; +} +function text(value: unknown): string | null { + if (value === null) return null; + if (typeof value !== "string" || value.length > 1024) throw new Error("Invalid usage snapshot"); + return usageText(value); +} + +export function decodePlanUsage(value: unknown, now = Date.now()): PlanUsageSnapshot { + const raw = record(value), + account = record(raw.account), + plan = record(raw.plan), + quota = record(raw.quota); + if ( + raw.schema_version !== 1 || + raw.scope !== "personal_account" || + typeof raw.observed_at_ms !== "number" || + !Number.isSafeInteger(raw.observed_at_ms) || + raw.observed_at_ms <= 0 || + raw.observed_at_ms > now + 120000 || + !["known", "unavailable"].includes(String(plan.state)) || + !["coding", "token", "none"].includes(String(plan.family)) || + !["known", "partial"].includes(String(quota.state)) || + typeof quota.allowed !== "boolean" || + !Array.isArray(quota.windows) + ) + throw new Error("Invalid usage snapshot"); + const windows = quota.windows.map((v) => { + const w = record(v); + if (!["five_hour", "weekly"].includes(String(w.id)) || typeof w.limited !== "boolean") + throw new Error("Invalid usage snapshot"); + const p = percent(w.used_percent); + if ((w.limited && p === null) || (!w.limited && p !== null)) throw new Error("Invalid usage snapshot"); + return { + id: w.id as "five_hour" | "weekly", + limited: w.limited, + used_percent: p, + resets_at_ms: time(w.resets_at_ms), + }; + }); + if ( + (plan.family === "coding" && (windows.length !== 2 || new Set(windows.map((w) => w.id)).size !== 2)) || + (plan.family !== "coding" && windows.length > 0) || + (quota.state === "partial" && plan.family !== "token") + ) + throw new Error("Invalid usage snapshot"); + let buckets: PlanUsageSnapshot["quota"]["credit_buckets"] = null; + if (quota.credit_buckets !== null) { + if (!Array.isArray(quota.credit_buckets) || quota.credit_buckets.length > 256) + throw new Error("Invalid usage snapshot"); + buckets = quota.credit_buckets.map((v) => { + const b = record(v); + if (typeof b.type !== "string" || b.type.length > 128) throw new Error("Invalid usage snapshot"); + if (!["subscription", "topup"].includes(b.type)) + return { + type: usageText(b.type), + unsupported: true, + total: "0", + remaining: "0", + used_percent: null, + resets_at_ms: null, + expires_at_ms: null, + }; + return { + type: b.type, + total: decimal(b.total), + remaining: decimal(b.remaining), + used_percent: percent(b.used_percent), + resets_at_ms: time(b.resets_at_ms), + expires_at_ms: time(b.expires_at_ms), + }; + }); + } + if ( + (quota.state === "partial" && buckets !== null) || + (quota.state === "known" && buckets === null) || + (plan.family !== "token" && buckets?.length) + ) + throw new Error("Invalid usage snapshot"); + let subscription: PlanUsageSnapshot["plan"]["subscription"] = null; + if (plan.subscription !== null) { + const s = record(plan.subscription); + subscription = { id: decimal(s.id, true), expires_at_ms: time(s.expires_at_ms) }; + } + if ( + (plan.state === "unavailable" && (plan.name !== null || subscription !== null)) || + (plan.family === "none" && (subscription !== null || quota.allowed)) + ) + throw new Error("Invalid usage snapshot"); + return { + schema_version: 1, + observed_at_ms: raw.observed_at_ms, + scope: "personal_account", + account: { id: decimal(account.id, true) }, + plan: { + state: plan.state as "known" | "unavailable", + family: plan.family as "coding" | "token" | "none", + name: text(plan.name), + subscription, + }, + quota: { + state: quota.state as "known" | "partial", + allowed: quota.allowed, + reason: text(quota.reason), + windows, + credit_buckets: buckets, + }, + }; +} + +async function readUsageBody(response: Response): Promise { + if (!response.body) throw new Error("Invalid usage snapshot"); + const reader = response.body.getReader(), + decoder = new TextDecoder("utf-8", { fatal: true }); + let bytes = 0, + body = ""; + try { + for (;;) { + const next = await reader.read(); + if (next.done) break; + bytes += next.value.byteLength; + if (bytes > 256 * 1024) throw new Error("Invalid usage snapshot"); + body += decoder.decode(next.value, { stream: true }); + } + body += decoder.decode(); + return JSON.parse(body); + } catch (error) { + void reader.cancel().catch(() => {}); + throw error; + } finally { + reader.releaseLock(); + } +} + +function crossedSnapshotBoundary(state: AccountUsageState, now: number): boolean { + const snapshot = state.snapshot; + if (!snapshot || state.receivedAt === undefined) return false; + const crossed = (value: number | null): boolean => value !== null && value <= now && state.receivedAt! < value; + return ( + snapshot.quota.windows.some((window) => crossed(window.resets_at_ms)) || + (snapshot.quota.credit_buckets ?? []).some( + (bucket) => crossed(bucket.resets_at_ms) || crossed(bucket.expires_at_ms), + ) || + (snapshot.plan.subscription !== null && crossed(snapshot.plan.subscription.expires_at_ms)) + ); +} + +export class PlanUsageController { + private readonly resolve: () => PlanAccountContext; + private readonly fetchFn: typeof fetch; + private readonly clock: () => number; + private readonly timeoutMS: number; + private readonly monotonic: () => number; + private receivedTick?: number; + private retryTick?: number; + private identity = ""; + private generation = 0; + private abort?: AbortController; + private pending?: Promise; + private current: AccountUsageState = { + status: "idle", + account: { presence: "missing", eligible: false, source: "none" }, + }; + constructor( + resolve: () => PlanAccountContext, + fetchFn: typeof fetch = fetch, + clock: () => number = Date.now, + timeoutMS = 12000, + monotonic: () => number = clock === Date.now ? () => performance.now() : clock, + ) { + this.monotonic = monotonic; + this.resolve = resolve; + this.fetchFn = fetchFn; + this.clock = clock; + this.timeoutMS = timeoutMS; + } + readCached(): AccountUsageState { + const context = this.resolve(); + if (context.identity !== this.identity) { + this.cancel(); + this.receivedTick = undefined; + this.retryTick = undefined; + this.identity = context.identity; + this.current = { + status: context.projection.eligible ? "idle" : "unsupported", + account: context.projection, + message: context.projection.reason, + }; + } + if (this.current.snapshot && this.monotonic() - (this.receivedTick ?? -Infinity) > 300000) + this.current = { + ...this.current, + status: "error", + snapshot: undefined, + message: "Account snapshot expired. Open /usage account to refresh.", + }; + if ( + this.current.snapshot && + this.current.status === "ready" && + (this.monotonic() - (this.receivedTick ?? -Infinity) >= 30000 || + crossedSnapshotBoundary(this.current, this.clock())) + ) + this.current = { ...this.current, status: "stale" }; + return { ...this.current, account: context.projection }; + } + cancel(): void { + this.generation++; + this.abort?.abort(); + this.abort = undefined; + this.pending = undefined; + if (this.current.status === "loading") + this.current = { ...this.current, status: this.current.snapshot ? "stale" : "idle" }; + } + query(force = false): Promise { + const state = this.readCached(), + context = this.resolve(); + if (!context.projection.eligible || !context.key || !context.projection.endpoint) return Promise.resolve(state); + if (this.pending) return this.pending; + if (this.retryTick !== undefined && this.retryTick > this.monotonic()) return Promise.resolve(state); + if ( + !force && + state.snapshot && + this.monotonic() - (this.receivedTick ?? -Infinity) < 30000 && + !crossedSnapshotBoundary(state, this.clock()) + ) + return Promise.resolve(state); + const attempt = ++this.generation, + identity = context.identity, + controller = new AbortController(); + this.abort = controller; + this.current = { ...state, status: "loading", message: undefined }; + const valid = () => + attempt === this.generation && this.resolve().identity === identity && !controller.signal.aborted; + const work = async (): Promise => { + const response = await this.fetchFn(context.projection.endpoint!, { + method: "GET", + headers: { accept: "application/json", authorization: `Bearer ${context.key}` }, + redirect: "error", + signal: controller.signal, + }); + if (!valid()) return this.readCached(); + if (!response.ok) { + void response.body?.cancel().catch(() => {}); + let retryAt: number | undefined; + if (response.status === 429) { + const header = response.headers.get("Retry-After") ?? "1"; + const seconds = /^\d+$/u.test(header) ? Number(header) : (Date.parse(header) - this.clock()) / 1000; + const duration = Math.max(1, Number.isFinite(seconds) ? seconds : 1) * 1000; + retryAt = this.clock() + duration; + this.retryTick = this.monotonic() + duration; + } + const clear = response.status === 401 || response.status === 403; + if (clear) { + this.receivedTick = undefined; + this.retryTick = undefined; + } + const messages: Record = { + 401: "Authentication could not be verified. Retry or check your key.", + 403: "Personal Plan usage does not support this credential scope.", + 404: "Usage API is not available on this platform yet.", + 429: "Too many queries. Retry after the cooldown.", + }; + if (valid()) + this.current = { + ...state, + status: !clear && state.snapshot ? "stale" : "error", + snapshot: clear ? undefined : state.snapshot, + retryAt, + message: messages[response.status] ?? "Account quota temporarily unavailable.", + }; + return this.readCached(); + } + const snapshot = decodePlanUsage(await readUsageBody(response), this.clock()); + if (!valid()) return this.readCached(); + if (context.projection.accountId && BigInt(snapshot.account.id) !== BigInt(context.projection.accountId)) { + this.current = { + status: "error", + account: state.account, + message: "Account identity mismatch. Sign in again.", + }; + return this.readCached(); + } + this.receivedTick = this.monotonic(); + this.retryTick = undefined; + this.current = { status: "ready", account: state.account, snapshot, receivedAt: this.clock() }; + return this.readCached(); + }; + let timer!: ReturnType; + let onAbort: () => void; + const cancelled = new Promise((_, reject) => { + onAbort = () => reject(new Error("Usage cancelled")); + controller.signal.addEventListener("abort", onAbort, { once: true }); + }); + const timeout = new Promise((_, reject) => { + timer = setTimeout(() => { + controller.abort(); + reject(new Error("Usage timeout")); + }, this.timeoutMS); + }); + this.pending = Promise.race([work(), timeout, cancelled]) + .catch(() => { + if (attempt === this.generation && this.resolve().identity === identity) + this.current = { + ...state, + status: state.snapshot ? "stale" : "error", + message: "Account quota unavailable. Retry when online.", + }; + return this.readCached(); + }) + .finally(() => { + clearTimeout(timer); + controller.signal.removeEventListener("abort", onAbort); + if (attempt === this.generation) { + this.pending = undefined; + this.abort = undefined; + } + }); + return this.pending; + } +} diff --git a/packages/coding-agent/src/step/usage/activity-engine.ts b/packages/coding-agent/src/step/usage/activity-engine.ts new file mode 100644 index 00000000..bb14c5a7 --- /dev/null +++ b/packages/coding-agent/src/step/usage/activity-engine.ts @@ -0,0 +1,694 @@ +import { createHash, randomUUID } from "node:crypto"; +import { constants, type Dir } from "node:fs"; +import { lstat, mkdir, open, opendir, realpath, rename, stat, unlink } from "node:fs/promises"; +import { dirname, join, relative, resolve } from "node:path"; +import lockfile from "proper-lockfile"; +import { entryUsage, normalizeUsage } from "../../core/usage-totals.ts"; +import { projectUsageJson, UsageProjectionLimit } from "./json-projection.ts"; +import type { UsageOverview } from "./types.ts"; + +type Row = Record; +interface FileIndex { + path: string; + aliases?: string[]; + ino: string; + dev: string; + size: number; + mtime: string; + parsedBytes: number; + digest?: string; + header?: Row; + entries: Row[]; + limited: string[]; +} +export interface ActivityScanOptions { + roots: string[]; + cwd: string; + scope: "all" | "project"; + dates: string[]; + timezone: string; + cachePath: string; + key: string; +} +const MAX_ENTRIES = 25000, + MAX_FILES = 10000, + MAX_LINE = 16 * 1024 * 1024; + +async function inventory( + roots: string[], + signal?: AbortSignal, +): Promise & { limited?: boolean }> { + const result: Map & { limited?: boolean } = new Map(); + const seen = new Set(); + const visit = async (path: string, depth: number, canonicalRoot: string, logicalRoot: string): Promise => { + if (depth > 8) { + result.limited = true; + return; + } + try { + const metadata = await lstat(path); + if (!metadata.isDirectory() || metadata.isSymbolicLink() || (await realpath(path)) !== path) return; + } catch { + return; + } + let dir: Dir; + try { + dir = await opendir(path); + } catch { + return; + } + for await (const entry of dir) { + signal?.throwIfAborted(); + const file = join(path, entry.name); + if (entry.isDirectory()) await visit(file, depth + 1, canonicalRoot, logicalRoot); + else if (entry.isFile() && entry.name.endsWith(".jsonl")) { + const alias = join(logicalRoot, relative(canonicalRoot, file)); + if (seen.has(file)) { + result.get(file)?.aliases?.push(alias); + continue; + } + if (result.size >= MAX_FILES) { + result.limited = true; + continue; + } + seen.add(file); + try { + const s = await stat(file, { bigint: true }); + if (s.size > BigInt(Number.MAX_SAFE_INTEGER)) continue; + result.set(file, { + path: file, + aliases: [alias], + ino: String(s.ino), + dev: String(s.dev), + size: Number(s.size), + mtime: String(s.mtimeNs), + parsedBytes: 0, + entries: [], + limited: [], + }); + } catch { + /* disappeared */ + } + } + } + }; + for (const root of roots) { + try { + const path = await realpath(root); + await visit(path, 0, path, resolve(root)); + } catch { + /* new installation */ + } + } + return result; +} + +async function readLines( + path: string, + start: number, + end: number, + onLine: (row: Row, through: number) => void, + limited: Set, + expected?: FileIndex, + signal?: AbortSignal, +): Promise { + const fd = await open(path, constants.O_RDONLY | (constants.O_NOFOLLOW ?? 0)); + let chunks: Buffer[] = [], + length = 0, + oversized = false, + cursor = start, + parsed = start; + const hash = createHash("sha256"); + let hashValid = true; + try { + if (expected) { + const actual = await fd.stat({ bigint: true }); + if (String(actual.ino) !== expected.ino || String(actual.dev) !== expected.dev) + throw new Error("Session file replaced"); + } + const buffer = Buffer.alloc(32768); + for (let offset = 0; offset < start; ) { + const { bytesRead } = await fd.read(buffer, 0, Math.min(buffer.length, start - offset), offset); + if (!bytesRead) throw new Error("Session prefix truncated"); + hash.update(buffer.subarray(0, bytesRead)); + offset += bytesRead; + } + if (start > 0 && expected?.digest && hash.copy().digest("hex") !== expected.digest) + throw new Error("Session prefix changed"); + while (cursor < end) { + signal?.throwIfAborted(); + const { bytesRead } = await fd.read(buffer, 0, Math.min(buffer.length, end - cursor), cursor); + if (bytesRead === 0) break; + let offset = 0; + for (let i = 0; i < bytesRead; i++) + if (buffer[i] === 10) { + const slice = buffer.subarray(offset, i); + length += slice.length; + const through = cursor + i + 1; + if (!oversized && length <= MAX_LINE) { + chunks.push(Buffer.from(slice)); + const rawBytes = Buffer.concat(chunks, length); + hash.update(rawBytes); + hash.update("\n"); + try { + const raw = new TextDecoder("utf-8", { fatal: true }).decode(rawBytes); + if (raw.trim()) onLine(projectUsageJson(raw), through); + } catch (error) { + limited.add( + error instanceof UsageProjectionLimit ? "Resource-limited records" : "Invalid records", + ); + } + } else { + limited.add("Resource-limited records"); + hashValid = false; + } + parsed = through; + chunks = []; + length = 0; + oversized = false; + offset = i + 1; + } + if (offset < bytesRead) { + const slice = buffer.subarray(offset, bytesRead); + length += slice.length; + if (length > MAX_LINE) { + oversized = true; + chunks = []; + } else if (!oversized) chunks.push(Buffer.from(slice)); + } + cursor += bytesRead; + } + if (length > 0 || oversized) limited.add("Incomplete trailing records"); + if (expected) expected.digest = hashValid ? hash.digest("hex") : ""; + return parsed; + } finally { + await fd.close(); + } +} + +async function readCache(options: ActivityScanOptions, signal?: AbortSignal): Promise> { + const files = new Map(), + limited = new Set(), + chunks = new Map(); + let matching = false, + complete = false, + entries = 0; + try { + const size = (await stat(options.cachePath)).size; + if (size > 32 * 1024 * 1024) return files; + await readLines( + options.cachePath, + 0, + size, + (row) => { + const d = row.data as Row | undefined; + if (row.type === "usage-cache-meta") { + matching = d?.key === options.key && d.version === 1; + complete = d?.complete === true; + return; + } + if ( + !matching || + row.type !== "usage-cache-file" || + !d || + typeof d.path !== "string" || + !Array.isArray(d.entries) || + entries + d.entries.length > MAX_ENTRIES || + typeof d.ino !== "string" || + typeof d.dev !== "string" || + typeof d.mtime !== "string" || + typeof d.size !== "number" || + typeof d.parsedBytes !== "number" || + !Number.isSafeInteger(d.size) || + !Number.isSafeInteger(d.parsedBytes) || + d.parsedBytes < 0 || + d.parsedBytes > d.size + ) + throw new Error("cache"); + entries += d.entries.length; + const old = files.get(d.path); + if ( + typeof d.chunk !== "number" || + d.chunk !== (chunks.get(d.path) ?? 0) || + (old && (d.ino !== old.ino || d.parsedBytes !== old.parsedBytes)) + ) + throw new Error("cache chunk"); + chunks.set(d.path, d.chunk + 1); + if (old) old.entries.push(...(d.entries as Row[])); + else { + if ( + d.chunk !== 0 || + typeof d.digest !== "string" || + !Array.isArray(d.limited) || + !d.limited.every((item) => typeof item === "string") + ) + throw new Error("cache chunk"); + files.set(d.path, { + path: d.path, + ino: d.ino, + dev: d.dev, + size: d.size, + mtime: d.mtime, + parsedBytes: d.parsedBytes, + digest: d.digest, + header: d.header as Row, + entries: d.entries as Row[], + limited: d.limited, + }); + } + }, + limited, + undefined, + signal, + ); + if (!matching || !complete || limited.size) files.clear(); + } catch { + files.clear(); + } + return files; +} + +function sameSource(a: FileIndex, b: FileIndex): boolean { + return a.ino === b.ino && a.dev === b.dev; +} + +async function prefixDigest(file: FileIndex, size: number): Promise { + const fd = await open(file.path, constants.O_RDONLY | (constants.O_NOFOLLOW ?? 0)); + try { + const actual = await fd.stat({ bigint: true }); + if (String(actual.ino) !== file.ino || String(actual.dev) !== file.dev) throw new Error("Session replaced"); + const hash = createHash("sha256"), + buffer = Buffer.alloc(32768); + for (let offset = 0; offset < size; ) { + const { bytesRead } = await fd.read(buffer, 0, Math.min(buffer.length, size - offset), offset); + if (!bytesRead) throw new Error("Session truncated"); + hash.update(buffer.subarray(0, bytesRead)); + offset += bytesRead; + } + return hash.digest("hex"); + } finally { + await fd.close(); + } +} + +async function commitCache( + options: ActivityScanOptions, + files: Map, + signal?: AbortSignal, +): Promise { + let release: (() => Promise) | undefined; + const temp = `${options.cachePath}.${randomUUID()}.tmp`; + try { + await mkdir(dirname(options.cachePath), { recursive: true, mode: 0o700 }); + release = await lockfile.lock(options.cachePath, { + realpath: false, + stale: 5000, + retries: { retries: 4, factor: 1, minTimeout: 50, maxTimeout: 50 }, + }); + signal?.throwIfAborted(); + const latest = await readCache(options, signal), + current = await inventory(options.roots, signal); + for (const [path, live] of current) { + const candidate = files.get(path), + previous = latest.get(path); + if ( + !candidate || + !sameSource(candidate, live) || + live.size < candidate.size || + (live.size === candidate.size && live.mtime !== candidate.mtime) || + (previous && sameSource(previous, candidate) && previous.parsedBytes > candidate.parsedBytes) + ) + return false; + if (candidate.digest && (await prefixDigest(live, candidate.parsedBytes)) !== candidate.digest) return false; + } + for (const path of files.keys()) if (!current.has(path)) return false; + const fd = await open(temp, "wx", 0o600); + try { + await fd.write(`${JSON.stringify({ type: "usage-cache-meta", data: { key: options.key, version: 1 } })}\n`); + let bytes = 0; + for (const file of files.values()) { + signal?.throwIfAborted(); + let index = 0, + chunk = 0; + do { + const selected: Row[] = []; + let frameBytes = 0; + while (index < file.entries.length && selected.length < 100) { + const row = file.entries[index], + rowBytes = Buffer.byteLength(JSON.stringify(row)); + if (selected.length && frameBytes + rowBytes > 96 * 1024) break; + selected.push(row); + index++; + frameBytes += rowBytes; + } + const frame = { type: "usage-cache-file", data: { ...file, chunk: chunk++, entries: selected } }; + const line = `${JSON.stringify(frame)}\n`; + bytes += Buffer.byteLength(line); + if (bytes > 32 * 1024 * 1024) throw new Error("cache budget"); + await fd.write(line); + } while (index < file.entries.length); + } + await fd.write( + `${JSON.stringify({ type: "usage-cache-meta", data: { key: options.key, version: 1, complete: true } })}\n`, + ); + } finally { + await fd.close(); + } + signal?.throwIfAborted(); + await rename(temp, options.cachePath); + return true; + } catch { + return true; /* No disk cache is preferable to failing the read-only view. */ + } finally { + await unlink(temp).catch(() => {}); + await release?.().catch(() => {}); + } +} + +interface Origin { + sessionId: string; + cwd: string; + role: "user" | "subagent"; +} +function validID(id: unknown): id is string { + return typeof id === "string" && /^[a-zA-Z0-9][a-zA-Z0-9_-]{0,127}$/u.test(id); +} + +function buildOverview( + options: ActivityScanOptions, + files: Map, + limited: Set, +): UsageOverview { + const dateFormat = new Intl.DateTimeFormat("en-CA", { + timeZone: options.timezone, + year: "numeric", + month: "2-digit", + day: "2-digit", + }); + const dayKey = (time: unknown): string | undefined => { + const d = new Date(typeof time === "number" || typeof time === "string" ? time : NaN); + if (!Number.isFinite(d.getTime())) return undefined; + const parts = dateFormat.formatToParts(d); + return ["year", "month", "day"].map((k) => parts.find((p) => p.type === k)?.value).join("-"); + }; + const days = new Map(options.dates.map((date) => [date, { date, tokens: 0, chats: new Set() }])); + const dedup = new Set(), + chats = new Set(), + reasoning = new Map(), + skills = new Set(), + usedSkills = new Set(), + resources = new Map>(); + let coverage = false, + unrecorded = false; + let globalTokens = 0; + const fileMetadata = new Map< + FileIndex, + { cap?: Row; ids: Set; manifests: Map; owned: boolean } + >(); + for (const file of files.values()) { + const sessionId = file.header?.id, + manifests = new Map(); + let owned = false; + for (const entry of file.entries) + if (entry.customType === "step.usage_lineage") { + const d = entry.data as Row; + if (d?.version === 1 && d.sessionId === sessionId && d.origins && typeof d.origins === "object") { + owned = true; + for (const [id, raw] of Object.entries(d.origins)) { + const source = raw as Origin; + if ( + source && + validID(source.sessionId) && + typeof source.cwd === "string" && + ["user", "subagent"].includes(source.role) + ) + manifests.set(id, source); + } + } + } + fileMetadata.set(file, { + cap: file.entries.find( + (e) => e.customType === "step.activity_capabilities" && (e.data as Row)?.sessionId === sessionId, + )?.data as Row | undefined, + ids: new Set(file.entries.map((e) => String(e.id))), + manifests, + owned, + }); + } + const parents = new Map(); + for (const file of files.values()) { + parents.set(file.path, file); + for (const alias of file.aliases ?? []) parents.set(alias, file); + } + const origins = new Map(), + resolving = new Set(); + const originFor = (file: FileIndex, entryID: string): Origin => { + const key = `${file.path}|${entryID}`, + known = origins.get(key); + if (known) return known; + const header = file.header!, + sessionId = String(header.id), + cwd = typeof header.cwd === "string" ? header.cwd : ""; + const metadata = fileMetadata.get(file)!, + cap = metadata.cap; + const role = + cap?.sessionRole === "subagent" || sessionId.startsWith("subagent-") || sessionId.startsWith("workflow-") + ? "subagent" + : "user"; + let result: Origin = { sessionId, cwd: typeof cap?.parentCwd === "string" ? cap.parentCwd : cwd, role }; + if (resolving.has(key)) { + limited.add("Lineage cycle"); + return result; + } + resolving.add(key); + const source = metadata.manifests.get(entryID); + if (source) { + result = source; + origins.set(key, result); + resolving.delete(key); + return result; + } + if (!metadata.owned && typeof header.parentSession === "string") { + const parent = parents.get(resolve(header.parentSession)); + if (parent?.header && fileMetadata.get(parent)?.ids.has(entryID)) result = originFor(parent, entryID); + else if (!parent) limited.add("Missing legacy lineage source"); + } + origins.set(key, result); + resolving.delete(key); + return result; + }; + const recordedSince = new Map(); + for (const file of files.values()) + if (file.header && validID(file.header.id)) + for (const entry of file.entries) + if (entry.customType === "step.activity_capabilities") { + const d = entry.data as Row, + time = Date.parse(String(entry.timestamp)); + if (d?.sessionId === originFor(file, String(entry.id)).sessionId && Number.isFinite(time)) + recordedSince.set( + String(d.sessionId), + Math.min(time, recordedSince.get(String(d.sessionId)) ?? Infinity), + ); + } + for (const file of files.values()) { + if (!file.header || !validID(file.header.id)) { + limited.add("Missing session identity"); + continue; + } + const levels = new Map(), + provenance = new Map(); + for (const entry of file.entries) + if (entry.customType === "step.usage_provenance") { + const d = entry.data as Row; + if ( + typeof d?.usageEntryId === "string" && + typeof d.reasoningAtDispatch === "string" && + d.originSessionId === originFor(file, d.usageEntryId).sessionId + ) + provenance.set(d.usageEntryId, d.reasoningAtDispatch); + } + for (let index = 0; index < file.entries.length; index++) { + const entry = file.entries[index], + id = typeof entry.id === "string" ? entry.id : `legacy-${index}`, + origin = originFor(file, id), + key = `${origin.sessionId}|${id}`; + if (options.scope === "project" && resolve(origin.cwd || ".") !== resolve(options.cwd)) continue; + if (entry.type === "thinking_level_change" && typeof entry.thinkingLevel === "string") + levels.set(id, entry.thinkingLevel); + else if (typeof entry.parentId === "string" && levels.has(entry.parentId)) + levels.set(id, levels.get(entry.parentId)!); + const data = entry.data as Row | undefined, + message = entry.message as Row | undefined; + const date = dayKey(data?.startedAt ?? entry.timestamp ?? message?.timestamp), + day = date ? days.get(date) : undefined; + if (!day) continue; + if (entry.customType === "step.activity_capabilities") { + coverage = true; + continue; + } + const picked = entryUsage(entry), + usage = normalizeUsage(picked.raw); + if (picked.expected && !usage) limited.add(picked.raw === undefined ? "Missing usage" : "Invalid usage"); + if (usage && !dedup.has(key)) { + if (entry.customType === "step-auto-clm-usage" && data?.sessionId !== origin.sessionId) { + limited.add("Usage origin conflict"); + continue; + } + if (!Number.isSafeInteger(globalTokens + usage.total)) { + limited.add("Token sum exceeds exact numeric range"); + continue; + } + dedup.add(key); + day.tokens += usage.total; + globalTokens += usage.total; + if (picked.assistant && origin.role === "user") { + day.chats.add(origin.sessionId); + chats.add(origin.sessionId); + } + if (picked.assistant) { + const level = provenance.get(id) ?? levels.get(id) ?? "unknown"; + if (!provenance.has(id)) limited.add("Reasoning coverage is limited"); + reasoning.set(level, (reasoning.get(level) ?? 0) + 1); + } + } + if (message?.role === "user" && origin.role === "user") { + day.chats.add(origin.sessionId); + chats.add(origin.sessionId); + } + if ( + entry.customType === "step.activity" && + data?.sessionId === origin.sessionId && + typeof data.runId === "string" && + typeof data.resourceId === "string" && + ["skill", "plugin"].includes(String(data.kind)) + ) { + const run = `${origin.sessionId}|${data.runId}|${data.kind}|${data.resourceId}`, + name = typeof data.name === "string" ? data.name : data.resourceId; + if (data.kind === "skill" && data.stage === "discovered") skills.add(data.resourceId); + else { + if (data.kind === "skill") usedSkills.add(data.resourceId); + const runs = resources.get(name) ?? new Set(); + runs.add(run); + resources.set(name, runs); + if (typeof data.pluginId === "string") { + const plugins = resources.get(`@${data.pluginId}`) ?? new Set(); + plugins.add(`${origin.sessionId}|${data.runId}|plugin|${data.pluginId}`); + resources.set(`@${data.pluginId}`, plugins); + } + } + } + if ( + picked.assistant && + Date.parse(String(entry.timestamp)) < (recordedSince.get(origin.sessionId) ?? Infinity) + ) + unrecorded = true; + } + } + if (unrecorded) limited.add("Older sessions did not record activity"); + const result = Array.from(days.values(), (day) => ({ date: day.date, tokens: day.tokens, chats: day.chats.size })); + return { + days: result, + tokens: globalTokens, + chats: chats.size, + activeDays: result.filter((d) => d.tokens > 0 || d.chats > 0).length, + reasoning: Array.from(reasoning).sort((a, b) => b[1] - a[1]), + discoveredSkills: coverage && (!unrecorded || skills.size > 0) ? skills.size : null, + usedSkills: coverage && (!unrecorded || usedSkills.size > 0) ? usedSkills.size : null, + topResources: Array.from(resources, ([name, runs]) => [name, runs.size] as [string, number]) + .sort((a, b) => b[1] - a[1] || a[0].localeCompare(b[0])) + .slice(0, 8), + limited: Array.from(limited), + scope: options.scope, + }; +} + +export async function scanActivity(options: ActivityScanOptions, signal?: AbortSignal): Promise { + for (let attempt = 0; attempt < 2; attempt++) { + signal?.throwIfAborted(); + const limited = new Set(), + files = await inventory(options.roots, signal), + cached = await readCache(options, signal); + let entries = 0, + changed = files.size !== cached.size; + if (files.limited) limited.add("File inventory budget exceeded"); + for (const file of files.values()) { + signal?.throwIfAborted(); + const old = cached.get(file.path); + let reusable = + old && + !!old.digest && + sameSource(old, file) && + file.size >= old.size && + (file.size > old.size || file.mtime === old.mtime); + if (reusable && old && file.size > old.size) { + try { + reusable = (await prefixDigest(file, old.parsedBytes)) === old.digest; + } catch { + reusable = false; + } + } + if (!reusable || old?.size !== file.size || old?.mtime !== file.mtime) changed = true; + if (reusable && old) { + file.header = old.header; + file.entries = old.entries; + file.parsedBytes = old.parsedBytes; + file.digest = old.digest; + file.limited = old.limited.filter((reason) => reason !== "Incomplete trailing records"); + } + const fileLimited = new Set(file.limited); + entries += file.entries.length; + if (entries >= MAX_ENTRIES) { + file.limited.push("Record inventory budget exceeded"); + for (const reason of file.limited) limited.add(reason); + continue; + } + if (reusable && old && old.size === file.size && old.mtime === file.mtime && old.parsedBytes === old.size) { + for (const reason of file.limited) limited.add(reason); + continue; + } + try { + file.parsedBytes = await readLines( + file.path, + file.parsedBytes, + file.size, + (row) => { + if (row.type === "session") file.header = row; + else if ( + row.type === "thinking_level_change" || + row.type === "message" || + entryUsage(row).expected || + (row.type === "custom" && + [ + "step.activity", + "step.activity_capabilities", + "step.usage_lineage", + "step.usage_provenance", + ].includes(String(row.customType))) + ) { + if (entries < MAX_ENTRIES) { + file.entries.push(row); + entries++; + } else fileLimited.add("Record inventory budget exceeded"); + } + }, + fileLimited, + file, + signal, + ); + } catch { + fileLimited.add("Unreadable sessions"); + } + file.limited = Array.from(fileLimited); + for (const reason of file.limited) limited.add(reason); + } + signal?.throwIfAborted(); + if (changed && !(await commitCache(options, files, signal))) continue; + signal?.throwIfAborted(); + return buildOverview(options, files, limited); + } + throw new Error("Activity changed during scanning; retry the view."); +} + +export function activityCacheKey(value: unknown): string { + return createHash("sha256").update(JSON.stringify(value)).digest("hex"); +} diff --git a/packages/coding-agent/src/step/usage/activity-worker.ts b/packages/coding-agent/src/step/usage/activity-worker.ts new file mode 100644 index 00000000..780116f9 --- /dev/null +++ b/packages/coding-agent/src/step/usage/activity-worker.ts @@ -0,0 +1,12 @@ +import { parentPort, workerData } from "node:worker_threads"; +import { type ActivityScanOptions, scanActivity } from "./activity-engine.ts"; + +if (parentPort) { + const controller = new AbortController(); + parentPort.on("message", (message) => { + if (message?.cancel === true) controller.abort(); + }); + void scanActivity(workerData as ActivityScanOptions, controller.signal) + .then((overview) => parentPort!.postMessage({ overview })) + .catch(() => parentPort!.postMessage({ error: "Local activity scan unavailable. Retry to rebuild its cache." })); +} diff --git a/packages/coding-agent/src/step/usage/activity.ts b/packages/coding-agent/src/step/usage/activity.ts new file mode 100644 index 00000000..1b15f4c0 --- /dev/null +++ b/packages/coding-agent/src/step/usage/activity.ts @@ -0,0 +1,72 @@ +import { createHash } from "node:crypto"; +import { join } from "node:path"; +import { Worker } from "node:worker_threads"; +import type { UsageOverview } from "./types.ts"; + +export function overviewDates(now = new Date()): string[] { + const day = new Date(now.getFullYear(), now.getMonth(), now.getDate()); + const dates: string[] = []; + for (let i = 364; i >= 0; i--) { + const d = new Date(day.getFullYear(), day.getMonth(), day.getDate() - i); + dates.push( + `${d.getFullYear()}-${String(d.getMonth() + 1).padStart(2, "0")}-${String(d.getDate()).padStart(2, "0")}`, + ); + } + return dates; +} + +export function scanUsageOverview( + roots: string[], + cacheDir: string, + cwd: string, + scope: "all" | "project", + signal: AbortSignal, +): Promise { + const dates = overviewDates(), + timezone = Intl.DateTimeFormat().resolvedOptions().timeZone; + const key = createHash("sha256") + .update(JSON.stringify([roots, cwd, scope, dates[0], dates.at(-1), timezone, 1])) + .digest("hex"); + const workerData = { roots, cwd, scope, dates, timezone, key, cachePath: join(cacheDir, `activity-${key}.jsonl`) }; + return new Promise((resolve, reject) => { + if (signal.aborted) { + reject(new Error("Activity scan cancelled")); + return; + } + const url = new URL( + import.meta.url.endsWith(".ts") ? "./activity-worker.ts" : "./activity-worker.js", + import.meta.url, + ); + const worker = new Worker( + typeof process.versions.bun === "string" ? "./src/step/usage/activity-worker.ts" : url, + { workerData, resourceLimits: { maxOldGenerationSizeMb: 64, maxYoungGenerationSizeMb: 16 } }, + ); + let settled = false; + const finish = (error?: Error, overview?: UsageOverview, cooperative = false): void => { + if (settled) return; + settled = true; + clearTimeout(timer); + signal.removeEventListener("abort", abort); + if (cooperative) { + worker.postMessage({ cancel: true }); + const kill = setTimeout(() => { + void worker.terminate(); + }, 200); + worker.once("exit", () => clearTimeout(kill)); + } else void worker.terminate(); + if (error) reject(error); + else if (overview) resolve(overview); + else reject(new Error("Invalid activity worker result")); + }; + const abort = (): void => finish(new Error("Activity scan cancelled"), undefined, true); + const timer = setTimeout(() => finish(new Error("Activity scan resource limit reached")), 30000); + signal.addEventListener("abort", abort, { once: true }); + worker.once("message", (message: { overview?: UsageOverview; error?: string }) => + finish(message.error ? new Error(message.error) : undefined, message.overview), + ); + worker.once("error", () => finish(new Error("Activity scan unavailable or resource limited"))); + worker.once("exit", () => { + if (!settled) finish(new Error("Activity worker exited before completion")); + }); + }); +} diff --git a/packages/coding-agent/src/step/usage/commands.ts b/packages/coding-agent/src/step/usage/commands.ts new file mode 100644 index 00000000..01f7aeeb --- /dev/null +++ b/packages/coding-agent/src/step/usage/commands.ts @@ -0,0 +1,221 @@ +import type { ExtensionAPI, ExtensionCommandContext, ExtensionContext } from "../../core/extensions/types.ts"; +import { collectSessionUsage } from "../../core/usage-totals.ts"; +import { getStepPlanMode } from "../../features/step-plan.ts"; +import { STEPCODE_VERSION } from "../version.ts"; +import { type PlanUsageController, usageText } from "./account.ts"; +import { overviewDates } from "./activity.ts"; +import { registerUsageRecording } from "./recording.ts"; +import type { UsageOverview, UsagePresentation, UsageResources, UsageStatus, UsageView } from "./types.ts"; + +export const USAGE_COMMAND_HELP = "Usage: /usage [overview|account] | /usage --refresh | /usage account --refresh"; +export function parseUsageCommand(args: string): { view: UsageView; refresh: boolean } | undefined { + const text = args.trim().replace(/\s+/gu, " "); + if (text === "" || text === "overview") return { view: "overview", refresh: false }; + if (text === "account") return { view: "account", refresh: false }; + if (["--refresh", "account --refresh"].includes(text)) return { view: "account", refresh: true }; + return undefined; +} +export interface StepUsageHost { + account: PlanUsageController; + resources(): UsageResources; + scan(cwd: string, scope: "all" | "project", signal: AbortSignal): Promise; + present(ctx: ExtensionCommandContext, presentation: UsagePresentation): Promise; + output?(ctx: ExtensionCommandContext, presentation: UsagePresentation): void; +} + +export function registerStepUsageCommands( + pi: ExtensionAPI, + host: StepUsageHost | undefined, + permissions: () => string, +): { + usage(args: string, ctx: ExtensionCommandContext): Promise; + status(args: string, ctx: ExtensionCommandContext): Promise; +} { + let generation = 0, + closeView: (() => void) | undefined; + let updateSession: ((ctx: ExtensionContext) => void) | undefined; + pi.on("agent_end", (_event, ctx) => updateSession?.(ctx)); + const close = (): void => { + generation++; + closeView?.(); + host?.account.cancel(); + }; + pi.on("session_start", close); + pi.on("session_shutdown", close); + registerUsageRecording(pi, host ? () => host.resources() : undefined); + + const capture = (ctx: ExtensionCommandContext): UsageStatus => { + const model = ctx.model, + context = ctx.getContextUsage(), + cached = host?.account.readCached() ?? { + status: "unsupported" as const, + account: { presence: "missing" as const, eligible: false, source: "none" }, + message: "Account service is not configured by this host.", + }; + let endpoint: string | undefined; + try { + const url = new URL(model?.baseUrl ?? ""); + endpoint = `${url.origin}${url.pathname}`; + } catch { + /* no endpoint */ + } + return { + model: usageText(model?.id ?? "Not configured"), + provider: usageText(model?.provider ?? "Not configured"), + reasoning: ctx.thinkingLevel ?? "Unknown", + directory: usageText(ctx.cwd), + permissions: usageText(permissions()), + trusted: ctx.isProjectTrusted(), + agents: (host?.resources().agents ?? ctx.getSystemPromptOptions().contextFiles?.map((f) => f.path) ?? []).map( + usageText, + ), + collaboration: getStepPlanMode(ctx.sessionManager), + session: ctx.sessionManager.getSessionId(), + state: `${ctx.isIdle() ? "idle" : "busy"}${ctx.hasPendingMessages() ? " · queued messages" : ""}`, + context: + context?.tokens === null || !context + ? "Unknown" + : `${context.tokens ?? "?"} / ${context.contextWindow ?? "?"} tokens`, + account: cached.account, + cached, + version: STEPCODE_VERSION.value, + endpoint, + }; + }; + const show = async (ctx: ExtensionCommandContext, view: UsageView | "status", force = false): Promise => { + closeView?.(); + const token = ++generation, + status = capture(ctx), + session = collectSessionUsage(ctx.sessionManager.getEntries()), + cwd = ctx.cwd; + let scan: AbortController | undefined, + closed = false; + const valid = (): boolean => !closed && generation === token; + const model: UsagePresentation = { + view, + session, + status, + account: status.cached, + scope: "all", + selectedDate: overviewDates().at(-1)!, + select: () => {}, + refresh: () => {}, + selectScope: () => {}, + }; + const changed = (): void => { + if (valid()) model.onChange?.(); + }; + updateSession = (current) => { + if (valid() && current.sessionManager.getSessionId() === status.session && model.view === "account") { + model.session = collectSessionUsage(current.sessionManager.getEntries()); + changed(); + } + }; + const query = (refresh: boolean): void => { + if (!host || !valid() || model.view !== "account") return; + const promise = host.account.query(refresh); + model.account = host.account.readCached(); + changed(); + void promise.then((state) => { + if (valid() && model.view === "account") { + model.account = state; + changed(); + } + }); + }; + const overview = (): void => { + if (!host || !valid()) return; + scan?.abort(); + const task = new AbortController(); + scan = task; + model.overview = undefined; + model.overviewError = undefined; + changed(); + void host + .scan(cwd, model.scope, task.signal) + .then((result) => { + if (valid() && scan === task && model.view === "overview") { + model.overview = result; + model.overviewError = undefined; + changed(); + } + }) + .catch((error: unknown) => { + if (valid() && scan === task && !task.signal.aborted) { + model.overviewError = error instanceof Error ? error.message : "Activity unavailable"; + changed(); + } + }); + }; + model.select = (next) => { + if (!valid() || model.view === "status") return; + scan?.abort(); + host?.account.cancel(); + model.view = next; + if (next === "account") query(false); + else overview(); + changed(); + }; + model.refresh = () => query(true); + model.selectScope = (scope) => { + if (!valid() || model.view !== "overview") return; + model.scope = scope; + overview(); + }; + const close = (): void => { + if (closed) return; + closed = true; + updateSession = undefined; + scan?.abort(); + host?.account.cancel(); + model.onClose?.(); + if (closeView === close) closeView = undefined; + }; + closeView = close; + if (ctx.mode !== "tui") { + try { + if (view === "account" && host) model.account = await host.account.query(force); + if (view === "overview" && host) { + scan = new AbortController(); + model.overview = await host.scan(cwd, "all", scan.signal); + } + if (valid()) { + if (host?.output) host.output(ctx, model); + else ctx.ui.notify("Usage output is unavailable in this host; use the exported data APIs.", "warning"); + } + } finally { + close(); + } + return; + } + if (!host) { + ctx.ui.notify("Usage UI is not configured by this host.", "warning"); + close(); + return; + } + if (view === "account") query(force); + else if (view === "overview") overview(); + try { + await host.present(ctx, model); + } finally { + close(); + } + }; + return { + usage: async (args, ctx) => { + const command = parseUsageCommand(args); + if (!command) { + ctx.ui.notify(USAGE_COMMAND_HELP, "warning"); + return; + } + await show(ctx, command.view, command.refresh); + }, + status: async (args, ctx) => { + if (args.trim()) { + ctx.ui.notify("Usage: /status", "warning"); + return; + } + await show(ctx, "status"); + }, + }; +} diff --git a/packages/coding-agent/src/step/usage/json-projection.ts b/packages/coding-agent/src/step/usage/json-projection.ts new file mode 100644 index 00000000..6326d3d8 --- /dev/null +++ b/packages/coding-agent/src/step/usage/json-projection.ts @@ -0,0 +1,212 @@ +/** A bounded JSON parser that validates ignored content without allocating it. */ +export class UsageProjectionLimit extends Error {} + +const top = new Set([ + "type", + "id", + "parentId", + "timestamp", + "version", + "cwd", + "parentSession", + "message", + "usage", + "customType", + "data", + "thinkingLevel", +]); +const usageKeys = new Set(["input", "output", "cacheRead", "cacheWrite"]); +const dataKeys = new Set([ + "version", + "usage", + "missingUsage", + "attemptId", + "requestId", + "sessionId", + "originSessionId", + "sourceLeafId", + "usageEntryId", + "reasoningAtDispatch", + "requestedProvider", + "requestedModel", + "runId", + "kind", + "stage", + "name", + "resourceId", + "pluginId", + "startedAt", + "sessionRole", + "skillDiscovery", + "skillActivation", + "pluginInvocation", + "parentSessionId", + "parentCwd", + "origins", + "key", + "complete", + "path", + "ino", + "dev", + "size", + "mtime", + "parsedBytes", + "header", + "entries", + "chunk", + "limited", +]); + +function keepPath(path: string[]): boolean { + if (path[0] === "data" && path[1] === "digest") return path.length === 2; + if (path[0] === "data" && ["entries", "header"].includes(path[1])) return keepPath(path.slice(2)); + if (path.length === 0) return true; + if (!top.has(path[0])) return false; + if (path.length === 1) return true; + if (path[0] === "usage") return path.length === 2 && usageKeys.has(path[1]); + if (path[0] === "message") { + if (path.length === 2) + return ["role", "usage", "provider", "model", "responseModel", "timestamp"].includes(path[1]); + return path[1] === "usage" && path.length === 3 && usageKeys.has(path[2]); + } + if (path[0] === "data") { + if (!dataKeys.has(path[1])) return false; + if (path[1] === "usage") return path.length === 2 || (path.length === 3 && usageKeys.has(path[2])); + if (path[1] === "origins") + return path.length <= 3 || (path.length === 4 && ["sessionId", "cwd", "role"].includes(path[3])); + return path.length === 2; + } + return false; +} + +export function projectUsageJson(source: string): Record { + let cursor = 0, + tokens = 0, + retained = 0; + const deadline = performance.now() + 1000; + const timeBudget = (): void => { + if (performance.now() > deadline) throw new UsageProjectionLimit("Session parsing time budget exceeded"); + }; + const fail = (): never => { + throw new SyntaxError("Invalid session JSON"); + }; + const budget = (n = 1): void => { + tokens += n; + if ((tokens & 255) === 0) timeBudget(); + if (tokens > 200000 || retained > 256 * 1024) throw new UsageProjectionLimit("Session record budget exceeded"); + }; + const space = (): void => { + while (cursor < source.length && /[\t\n\r ]/u.test(source[cursor])) cursor++; + }; + const string = (keep: boolean): string => { + if (source[cursor++] !== '"') return fail(); + let result = ""; + while (cursor < source.length) { + if ((cursor & 32767) === 0) timeBudget(); + let ch = source[cursor++]; + if (ch === '"') { + if (keep) { + retained += Buffer.byteLength(result); + budget(0); + } + return result; + } + if (ch.charCodeAt(0) < 32) return fail(); + if (ch === "\\") { + const escaped = source[cursor++]; + if (escaped === "u") { + const code = source.slice(cursor, cursor + 4); + if (!/^[0-9a-fA-F]{4}$/u.test(code)) return fail(); + cursor += 4; + ch = String.fromCharCode(Number.parseInt(code, 16)); + } else { + const escapes: Record = { + '"': '"', + "\\": "\\", + "/": "/", + b: "\b", + f: "\f", + n: "\n", + r: "\r", + t: "\t", + }; + if (!(escaped in escapes)) return fail(); + ch = escapes[escaped]; + } + } + if (keep) { + result += ch; + if (result.length > 4096) throw new UsageProjectionLimit("Session string budget exceeded"); + } + } + return fail(); + }; + const value = (path: string[] | undefined, depth: number): unknown => { + budget(); + if (depth > 64) throw new UsageProjectionLimit("Session nesting budget exceeded"); + space(); + const keep = path !== undefined && keepPath(path), + ch = source[cursor]; + if (keep) { + retained += 16; + budget(0); + } + if (ch === '"') return string(keep); + if (ch === "{" || ch === "[") { + const object = ch === "{", + end = object ? "}" : "]", + result: Record | unknown[] = object ? Object.create(null) : []; + cursor++; + space(); + if (source[cursor] === end) { + cursor++; + return keep ? result : undefined; + } + for (;;) { + let key = ""; + if (object) { + space(); + key = string(keep); + space(); + if (source[cursor++] !== ":") return fail(); + } + const nextPath = keep && path ? (object ? [...path, key] : path) : undefined; + const selected = nextPath !== undefined && keepPath(nextPath), + parsed = value(selected ? nextPath : undefined, depth + 1); + if (keep && selected) { + if (object) { + if (Object.hasOwn(result, key)) return fail(); + (result as Record)[key] = parsed; + } else (result as unknown[]).push(parsed); + } + space(); + const separator = source[cursor++]; + if (separator === end) break; + if (separator !== ",") return fail(); + } + return keep ? result : undefined; + } + for (const [literal, parsed] of [ + ["true", true], + ["false", false], + ["null", null], + ] as const) { + if (source.startsWith(literal, cursor)) { + cursor += literal.length; + return keep ? parsed : undefined; + } + } + const start = cursor; + while (cursor < source.length && !/[\t\n\r ,}\]]/u.test(source[cursor])) cursor++; + const number = source.slice(start, cursor); + if (number.length > 128 || !/^-?(?:0|[1-9]\d*)(?:\.\d+)?(?:[eE][+-]?\d+)?$/u.test(number)) return fail(); + const parsed = Number(number); + if (keep && !Number.isFinite(parsed)) return fail(); + return keep ? parsed : undefined; + }; + const parsed = value([], 0); + budget(0); + space(); + if (cursor !== source.length || !parsed || typeof parsed !== "object" || Array.isArray(parsed)) return fail(); + return parsed as Record; +} diff --git a/packages/coding-agent/src/step/usage/recording.ts b/packages/coding-agent/src/step/usage/recording.ts new file mode 100644 index 00000000..188855c3 --- /dev/null +++ b/packages/coding-agent/src/step/usage/recording.ts @@ -0,0 +1,108 @@ +import { realpathSync } from "node:fs"; +import { resolve } from "node:path"; +import type { ExtensionAPI, ExtensionContext } from "../../core/extensions/types.ts"; +import { getStepMcpToolOwner } from "../mcp.ts"; +import { getStepPluginSkillOwner } from "../plugins.ts"; +import type { UsageResources } from "./types.ts"; + +/** Record through existing extension hooks; never send activity metadata to the model. */ +export function registerUsageRecording(pi: ExtensionAPI, resources: (() => UsageResources) | undefined): void { + if (!resources) return; + const discovered = new Set(); + const append = (type: string, data: unknown): void => { + try { + pi.appendEntry(type, data); + } catch { + /* Local recording never changes the task result. */ + } + }; + const collectResources = (ctx: ExtensionContext): void => { + const id = ctx.sessionManager.getSessionId(); + for (const skill of resources().skills) { + const resourceId = resolve(skill.filePath); + if (discovered.has(resourceId)) continue; + discovered.add(resourceId); + append("step.activity", { + version: 1, + sessionId: id, + kind: "skill", + stage: "discovered", + resourceId, + name: `$${skill.name}`, + runId: `discovery:${resourceId}`, + startedAt: Date.now(), + }); + } + }; + pi.on("session_start", (_event, ctx) => { + discovered.clear(); + append("step.activity_capabilities", { + version: 1, + sessionId: ctx.sessionManager.getSessionId(), + sessionRole: + process.env.STEPCODE_SUBAGENT_CHILD === "1" || + ctx.sessionManager.getSessionId().startsWith("subagent-") || + ctx.sessionManager.getSessionId().startsWith("workflow-") + ? "subagent" + : "user", + parentCwd: process.env.STEPCODE_USAGE_PARENT_CWD, + parentSessionId: process.env.STEPCODE_USAGE_PARENT_SESSION_ID, + skillDiscovery: true, + skillActivation: true, + pluginInvocation: true, + }); + collectResources(ctx); + }); + pi.on("before_agent_start", (_event, ctx) => { + collectResources(ctx); + }); + pi.on("usage_recorded", (event, ctx) => { + if (event.provenance.sessionId === ctx.sessionManager.getSessionId()) + append("step.usage_provenance", { + version: 1, + ...event.provenance, + originSessionId: event.provenance.sessionId, + usageEntryId: event.entryId, + }); + }); + pi.on("tool_execution_start", (event, ctx) => { + const sessionId = ctx.sessionManager.getSessionId(), + startedAt = Date.now(), + owner = getStepMcpToolOwner(event.toolName, ctx.sessionManager); + if (owner) + append("step.activity", { + version: 1, + sessionId, + runId: event.toolCallId, + startedAt, + kind: "plugin", + resourceId: owner, + name: `@${owner}`, + }); + if (event.toolName !== "read" || typeof event.args?.path !== "string") return; + let path = resolve(ctx.cwd, event.args.path); + try { + path = realpathSync(path); + } catch { + /* An attempted invocation still counts as a run. */ + } + const skill = resources().skills.find((s) => { + try { + return realpathSync(s.filePath) === path; + } catch { + return resolve(s.filePath) === path; + } + }); + if (skill) + append("step.activity", { + version: 1, + sessionId, + runId: event.toolCallId, + startedAt, + kind: "skill", + resourceId: resolve(skill.filePath), + name: `$${skill.name}`, + pluginId: getStepPluginSkillOwner(path, ctx.sessionManager), + }); + }); +} diff --git a/packages/coding-agent/src/step/usage/types.ts b/packages/coding-agent/src/step/usage/types.ts new file mode 100644 index 00000000..fb97401e --- /dev/null +++ b/packages/coding-agent/src/step/usage/types.ts @@ -0,0 +1,117 @@ +export type UsageView = "overview" | "account"; +export interface UsageWindow { + id: "five_hour" | "weekly"; + limited: boolean; + used_percent: number | null; + resets_at_ms: number | null; +} +export interface CreditBucket { + type: string; + total: string; + remaining: string; + used_percent: number | null; + resets_at_ms: number | null; + expires_at_ms: number | null; + unsupported?: boolean; +} +export interface PlanUsageSnapshot { + schema_version: 1; + observed_at_ms: number; + scope: "personal_account"; + account: { id: string }; + plan: { + state: "known" | "unavailable"; + family: "coding" | "token" | "none"; + name: string | null; + subscription: { id: string; expires_at_ms: number | null } | null; + }; + quota: { + state: "known" | "partial"; + allowed: boolean; + reason: string | null; + windows: UsageWindow[]; + credit_buckets: CreditBucket[] | null; + }; +} +export interface LocalPlanAccount { + presence: "present" | "missing"; + eligible: boolean; + source: string; + profile?: string; + region?: string; + accountId?: string; + endpoint?: string; + reason?: string; +} +/** Private controller input. Never pass it to a renderer, log or telemetry. */ +export interface PlanAccountContext { + projection: LocalPlanAccount; + key?: string; + identity: string; +} +export interface AccountUsageState { + status: "idle" | "loading" | "ready" | "stale" | "error" | "unsupported"; + account: LocalPlanAccount; + snapshot?: PlanUsageSnapshot; + message?: string; + receivedAt?: number; + retryAt?: number; +} + +import type { SessionUsage } from "../../core/usage-totals.ts"; + +export type { SessionUsage } from "../../core/usage-totals.ts"; +export interface OverviewDay { + date: string; + tokens: number; + chats: number; +} +export interface UsageOverview { + days: OverviewDay[]; + tokens: number; + chats: number; + activeDays: number; + reasoning: [string, number][]; + discoveredSkills: number | null; + usedSkills: number | null; + topResources: [string, number][]; + limited: string[]; + scope: "all" | "project"; +} +export interface UsageStatus { + model: string; + provider: string; + reasoning: string; + directory: string; + permissions: string; + trusted: boolean; + agents: string[]; + collaboration: string; + session: string; + state: string; + context: string; + account: LocalPlanAccount; + cached: AccountUsageState; + version: string; + endpoint?: string; +} +export interface UsagePresentation { + view: UsageView | "status"; + session: SessionUsage; + status: UsageStatus; + account: AccountUsageState; + overview?: UsageOverview; + overviewError?: string; + scope: "all" | "project"; + selectedDate: string; + onChange?: () => void; + onClose?: () => void; + select(view: UsageView): void; + refresh(): void; + selectScope(scope: "all" | "project"): void; +} + +export interface UsageResources { + agents: string[]; + skills: { name: string; filePath: string }[]; +} diff --git a/packages/coding-agent/test/agent-session-stats.test.ts b/packages/coding-agent/test/agent-session-stats.test.ts index c363865d..191e7c4b 100644 --- a/packages/coding-agent/test/agent-session-stats.test.ts +++ b/packages/coding-agent/test/agent-session-stats.test.ts @@ -10,7 +10,7 @@ import { AgentSession } from "../src/core/agent-session.ts"; import { AuthStorage } from "../src/core/auth-storage.ts"; import { SessionManager } from "../src/core/session-manager.ts"; import { SettingsManager } from "../src/core/settings-manager.ts"; -import { getUsageCostBreakdown } from "../src/core/usage-totals.ts"; +import { collectSessionUsage, getUsageCostBreakdown } from "../src/core/usage-totals.ts"; import { createInMemoryModelRegistry, getModelRuntime } from "./model-runtime-test-utils.ts"; import { createTestResourceLoader, stepModel } from "./utilities.ts"; @@ -323,6 +323,38 @@ describe("AgentSession.getSessionStats", () => { ]); }); + it("shares validated increments across session stats, Account, and the cost breakdown", async () => { + const { session, sessionManager } = await createSession(); + try { + const usage: Usage = { + ...createUsage(10), + output: 2, + cacheRead: 3, + cacheWrite: 4, + totalTokens: 999, + cost: { ...createUsage(0).cost, total: 0.5 }, + }; + sessionManager.appendMessage({ ...createAssistantMessage("valid", 0, 1), usage }); + sessionManager.appendCustomEntry("step-auto-clm-usage", { usage }); + sessionManager.appendCustomEntry("step-auto-clm-result", { usage }); + sessionManager.appendMessage({ ...createAssistantMessage("invalid", 0, 2), usage: { ...usage, input: -1 } }); + sessionManager.appendMessage({ + ...createAssistantMessage("missing cost", 0, 3), + usage: { ...usage, cost: undefined } as unknown as Usage, + }); + const stats = session.getSessionStats(), + account = collectSessionUsage(sessionManager.getEntries()), + breakdown = getUsageCostBreakdown(sessionManager.getEntries()); + expect(stats.tokens).toEqual({ input: 30, output: 6, cacheRead: 9, cacheWrite: 12, total: 57 }); + expect(stats.cost).toBe(1); + expect(account).toMatchObject({ input: 51, total: 57, invalid: 1, missing: 0, cacheHit: 9 / 51 }); + expect(breakdown.reduce((sum, row) => sum + row.tokens, 0)).toBe(stats.tokens.total); + expect(breakdown.reduce((sum, row) => sum + row.cost, 0)).toBe(stats.cost); + } finally { + session.dispose(); + } + }); + it("ignores zero-usage messages when checking for post-compaction context usage", async () => { const { session, sessionManager } = await createSession("off"); diff --git a/packages/coding-agent/test/step-onboarding.test.ts b/packages/coding-agent/test/step-onboarding.test.ts index 8d530e5e..4b933a73 100644 --- a/packages/coding-agent/test/step-onboarding.test.ts +++ b/packages/coding-agent/test/step-onboarding.test.ts @@ -117,4 +117,49 @@ describe("Step login onboarding", () => { expect(rows.map(stripAnsi).join("")).toContain(authUrl); expect(rows.every((row) => visibleWidth(row) <= width)).toBe(true); }); + + it.each(["step_plan", "step_plan_oversea"] as const)( + "renders and navigates missing-plan actions for %s", + (choice) => { + const selected: string[] = []; + let retries = 0; + let purchases = 0; + let changed = 0; + let quit = 0; + const view = new StepOnboardingView(resolveStepLoginProfiles(), { + onChoose: (profile) => selected.push(profile), + onSubmitApiKey: () => {}, + onType: () => {}, + onBackspace: () => {}, + onBack: () => { + changed++; + }, + onQuit: () => { + quit++; + }, + onRetryPlan: () => { + retries++; + }, + onOpenPlan: () => { + purchases++; + }, + requestRender: () => {}, + }); + view.setStep({ kind: "planRequired", choice }); + for (const width of [12, 40, 80, 160]) { + const rows = view.render(width); + expect(rows.every((row) => visibleWidth(row) <= width)).toBe(true); + expect(rows.map(stripAnsi).join("")).toContain( + `https://platform.stepfun.${choice === "step_plan" ? "com" : "ai"}/step-plan`, + ); + } + view.handleInput("\r"); + expect(selected).toEqual([choice === "step_plan" ? "platform_cn" : "platform_oversea"]); + view.handleInput("2"); + view.handleInput("3"); + view.handleInput("4"); + view.handleInput("\x1b"); + expect([purchases, retries, changed, quit]).toEqual([1, 1, 1, 1]); + }, + ); }); diff --git a/packages/coding-agent/test/step-plan-login.test.ts b/packages/coding-agent/test/step-plan-login.test.ts new file mode 100644 index 00000000..243b56be --- /dev/null +++ b/packages/coding-agent/test/step-plan-login.test.ts @@ -0,0 +1,311 @@ +import { mkdtemp, readFile, rm } from "node:fs/promises"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; +import { loginStepOAuth } from "../src/features/step-provider/index.ts"; +import { + checkStepPlanBeforeInteractive, + readStepLoginCredential, + runStepLogin, + writeStepLoginCredential, +} from "../src/step/login-flow.ts"; +import type { StepLoginProfileId } from "../src/step/onboarding.ts"; +import type { StepOnboardingView } from "../src/step/onboarding-view.ts"; +import { PlanUsageController, resolvePlanUsageAccount } from "../src/step/usage/account.ts"; +import type { PlanUsageSnapshot } from "../src/step/usage/types.ts"; +import { stripAnsi } from "../src/utils/ansi.ts"; +import * as browser from "../src/utils/open-browser.ts"; + +vi.mock("../src/features/step-provider/index.ts", async (original) => ({ + ...(await original()), + loginStepOAuth: vi.fn(async () => ({ + access: "plan-key", + refresh: "static", + expires: Number.MAX_SAFE_INTEGER, + uid: "123", + })), +})); + +function snapshot(family: "none" | "coding" | "token" = "none"): PlanUsageSnapshot { + return { + schema_version: 1, + observed_at_ms: Date.now(), + scope: "personal_account", + account: { id: "123" }, + plan: { + state: "known", + family, + name: family === "none" ? null : "Pro", + subscription: family === "coding" ? { id: "7", expires_at_ms: null } : null, + }, + quota: { + state: "known", + allowed: family !== "none", + reason: family === "none" ? "no_active_subscription" : null, + windows: + family === "coding" + ? [ + { id: "five_hour", limited: true, used_percent: 25, resets_at_ms: null }, + { id: "weekly", limited: true, used_percent: 50, resets_at_ms: null }, + ] + : [], + credit_buckets: + family === "token" + ? [ + { + type: "topup", + total: "10", + remaining: "5", + used_percent: 50, + resets_at_ms: null, + expires_at_ms: null, + }, + ] + : [], + }, + }; +} + +describe("Step Plan login availability", () => { + let root: string; + let authPath: string; + beforeEach(async () => { + root = await mkdtemp(join(tmpdir(), "step-plan-login-")); + authPath = join(root, "auth.json"); + vi.mocked(loginStepOAuth).mockClear(); + vi.spyOn(browser, "openBrowser").mockImplementation(() => {}); + }); + afterEach(async () => { + vi.restoreAllMocks(); + await rm(root, { recursive: true, force: true }); + }); + function controller(fetcher: typeof fetch) { + return new PlanUsageController(() => resolvePlanUsageAccount(authPath, {}), fetcher, Date.now, 100); + } + function start(planUsage: PlanUsageController, requiredPlan?: "step_plan" | "step_plan_oversea") { + let view!: StepOnboardingView; + const stop = vi.fn(); + const login = runStepLogin({ + authPath, + env: {}, + noBrowser: true, + planUsage, + requiredPlan, + createHost: () => ({ + addChild: (child) => { + view = child as StepOnboardingView; + }, + setFocus: () => {}, + requestRender: () => {}, + start: () => {}, + stop, + }), + }); + return { login, view, stop }; + } + + it.each(["step_plan", "step_plan_oversea"] as const)( + "checks the stored %s account in its own region", + async (profile) => { + await writeStepLoginCredential({ authPath, profile, apiKey: "plan-key", uid: "123" }); + const before = await readFile(authPath, "utf8"); + const fetcher = vi.fn(async () => Response.json(snapshot())); + const result = await checkStepPlanBeforeInteractive({ + authPath, + interactive: true, + env: {}, + planUsage: controller(fetcher), + }); + expect(result).toEqual({ requiredPlan: profile, warning: undefined }); + expect(fetcher.mock.calls[0][0]).toBe( + `https://api.stepfun.${profile === "step_plan" ? "com" : "ai"}/step_plan/v1/usage`, + ); + expect(await readFile(authPath, "utf8")).toBe(before); + }, + ); + + it.each(["coding", "token"] as const)( + "keeps a usable %s account, including top-up-only accounts", + async (family) => { + await writeStepLoginCredential({ authPath, profile: "step_plan", apiKey: "plan-key", uid: "123" }); + const fetcher = vi.fn(async () => Response.json(snapshot(family))); + expect( + ( + await checkStepPlanBeforeInteractive({ + authPath, + interactive: true, + env: {}, + planUsage: controller(fetcher), + }) + ).requiredPlan, + ).toBeUndefined(); + }, + ); + + it("does not treat an exhausted plan as a missing subscription", async () => { + await writeStepLoginCredential({ authPath, profile: "step_plan", apiKey: "plan-key", uid: "123" }); + const data = snapshot("coding"); + data.quota.allowed = false; + data.quota.reason = "week_limit_exceeded"; + const planUsage = controller(vi.fn(async () => Response.json(data))); + expect( + (await checkStepPlanBeforeInteractive({ authPath, interactive: true, env: {}, planUsage })).requiredPlan, + ).toBeUndefined(); + }); + + it.each([401, 403, 404, 429, 500, 503])("keeps credentials and reports uncertainty for HTTP %s", async (status) => { + await writeStepLoginCredential({ authPath, profile: "step_plan_oversea", apiKey: "plan-key", uid: "123" }); + const before = await readFile(authPath, "utf8"); + const result = await checkStepPlanBeforeInteractive({ + authPath, + interactive: true, + env: {}, + planUsage: controller(vi.fn(async () => new Response("", { status }))), + }); + expect(result.requiredPlan).toBeUndefined(); + expect(result.warning).toContain("Could not verify Step Plan availability"); + expect(await readFile(authPath, "utf8")).toBe(before); + }); + + it.each(["network", "invalid", "mismatch", "timeout", "partial", "unconfirmed"])( + "does not redirect on %s", + async (kind) => { + await writeStepLoginCredential({ authPath, profile: "step_plan", apiKey: "plan-key", uid: "123" }); + const fetcher = vi.fn(async () => { + if (kind === "network") throw new Error("offline"); + if (kind === "timeout") return new Promise(() => {}); + if (kind === "invalid") return Response.json({ data: [] }); + const data = snapshot(kind === "partial" ? "token" : "none"); + if (kind === "mismatch") data.account.id = "456"; + if (kind === "unconfirmed") data.plan.state = "unavailable"; + if (kind === "partial") { + data.quota.state = "partial"; + data.quota.credit_buckets = null; + } + return Response.json(data); + }); + const result = await checkStepPlanBeforeInteractive({ + authPath, + interactive: true, + env: {}, + planUsage: controller(fetcher), + }); + expect(result.requiredPlan).toBeUndefined(); + if (kind !== "partial") expect(result.warning).toBeTruthy(); + }, + ); + + it.each(["platform_cn", "platform_oversea", "headless", "environment", "missing"])( + "skips Plan requests for %s", + async (kind) => { + if (kind !== "missing") + await writeStepLoginCredential({ + authPath, + profile: kind.startsWith("platform") ? (kind as StepLoginProfileId) : "step_plan", + apiKey: "key", + }); + const fetcher = vi.fn(); + expect( + await checkStepPlanBeforeInteractive({ + authPath, + interactive: kind !== "headless", + env: kind === "environment" ? { STEP_API_KEY: "env-key" } : {}, + planUsage: controller(fetcher), + }), + ).toEqual({}); + expect(fetcher).not.toHaveBeenCalled(); + }, + ); + + it.each([ + ["step_plan", "1", "platform_cn", "api.stepfun.com"], + ["step_plan_oversea", "2", "platform_oversea", "api.stepfun.ai"], + ] as const)( + "offers an API key after %s browser login and preserves the region", + async (profile, choice, platform, host) => { + const fetcher = vi.fn(async () => Response.json(snapshot())); + const { login, view, stop } = start(controller(fetcher)); + try { + view.handleInput(choice); + await vi.waitFor(() => expect(view.getStep()).toEqual({ kind: "planRequired", choice: profile })); + expect(view.render(100).map(stripAnsi).join("\n")).toContain( + "No active Step Plan subscription or usable top-up credits", + ); + expect(fetcher.mock.calls[0][0]).toContain(host); + view.handleInput("1"); + expect(view.getStep()).toMatchObject({ kind: "apiKeyEntry", choice: platform }); + view.handleInput("platform-test-key"); + view.handleInput("\r"); + await expect(login).resolves.toMatchObject({ kind: "completed", profile: { id: platform } }); + expect(readStepLoginCredential(authPath)).toMatchObject({ access: "platform-test-key", profile: platform }); + expect(readStepLoginCredential(authPath)).not.toHaveProperty("uid"); + expect(fetcher).toHaveBeenCalledTimes(1); + expect(stop).toHaveBeenCalledOnce(); + } finally { + view.handleInput("\x03"); + await login; + } + }, + ); + + it("rechecks subscription purchase without another browser login", async () => { + await writeStepLoginCredential({ authPath, profile: "step_plan", apiKey: "plan-key", uid: "123" }); + const fetcher = vi.fn(async () => Response.json(snapshot("coding"))); + const { login, view } = start(controller(fetcher), "step_plan"); + try { + view.handleInput("3"); + await expect(login).resolves.toMatchObject({ kind: "completed", profile: { id: "step_plan" } }); + expect(loginStepOAuth).not.toHaveBeenCalled(); + } finally { + view.handleInput("\x03"); + await login; + } + }); + + it("cancels a pending check and ignores its late response", async () => { + await writeStepLoginCredential({ authPath, profile: "step_plan", apiKey: "plan-key", uid: "123" }); + let respond!: (value: Response) => void; + const fetcher = vi.fn( + () => + new Promise((resolve) => { + respond = resolve; + }), + ); + const { login, view, stop } = start(controller(fetcher), "step_plan"); + view.handleInput("3"); + expect(view.getStep().kind).toBe("checkingPlan"); + view.handleInput("\x03"); + await expect(login).resolves.toEqual({ kind: "exit" }); + expect(fetcher.mock.calls[0][1]?.signal?.aborted).toBe(true); + respond(Response.json(snapshot("coding"))); + await new Promise((resolve) => setTimeout(resolve, 0)); + expect(view.getStep().kind).toBe("exit"); + expect(stop).toHaveBeenCalledOnce(); + }); + + it("preserves the Plan credential when API-key entry is cancelled", async () => { + await writeStepLoginCredential({ authPath, profile: "step_plan_oversea", apiKey: "plan-key", uid: "123" }); + const before = await readFile(authPath, "utf8"); + const { login, view } = start(controller(vi.fn()), "step_plan_oversea"); + view.handleInput("1"); + view.handleInput("unsubmitted-key"); + view.handleInput("\x03"); + await expect(login).resolves.toEqual({ kind: "exit" }); + expect(await readFile(authPath, "utf8")).toBe(before); + }); + + it("reports unavailable quota after login without replacing a valid credential", async () => { + const { login, view } = start(controller(vi.fn(async () => new Response("", { status: 404 })))); + try { + view.handleInput("2"); + await expect(login).resolves.toMatchObject({ + kind: "completed", + warning: expect.stringContaining("Could not verify Step Plan availability"), + }); + expect(readStepLoginCredential(authPath)).toMatchObject({ access: "plan-key", profile: "step_plan_oversea" }); + } finally { + view.handleInput("\x03"); + await login; + } + }); +}); diff --git a/packages/coding-agent/test/step-usage-lifecycle.test.ts b/packages/coding-agent/test/step-usage-lifecycle.test.ts new file mode 100644 index 00000000..6920f196 --- /dev/null +++ b/packages/coding-agent/test/step-usage-lifecycle.test.ts @@ -0,0 +1,183 @@ +import { describe, expect, it, vi } from "vitest"; +import type { ExtensionAPI, ExtensionCommandContext, ExtensionContext } from "../src/core/extensions/types.ts"; +import { PlanUsageController } from "../src/step/usage/account.ts"; +import { registerStepUsageCommands, type StepUsageHost } from "../src/step/usage/commands.ts"; +import type { UsageOverview, UsagePresentation } from "../src/step/usage/types.ts"; + +const empty: UsageOverview = { + days: [{ date: "2026-10-13", tokens: 0, chats: 0 }], + tokens: 0, + chats: 0, + activeDays: 0, + reasoning: [], + discoveredSkills: null, + usedSkills: null, + topResources: [], + limited: [], + scope: "all", +}; +function fixture() { + const handlers = new Map unknown>(); + const entries: Record[] = []; + let sessionId = "first"; + const ctx = { + mode: "tui", + cwd: "/work", + model: { id: "step-model", provider: "step", baseUrl: "https://api.stepfun.com/v1" }, + thinkingLevel: "high", + getContextUsage: () => ({ tokens: 10, contextWindow: 100 }), + getSystemPromptOptions: () => ({}), + isProjectTrusted: () => true, + isIdle: () => true, + hasPendingMessages: () => false, + sessionManager: { getSessionId: () => sessionId, getEntries: () => entries }, + ui: { notify: vi.fn() }, + } as unknown as ExtensionCommandContext; + const pi = { + on: (event: string, callback: (event: unknown, ctx: ExtensionContext) => unknown) => { + const previous = handlers.get(event); + handlers.set(event, async (value, context) => { + await previous?.(value, context); + return callback(value, context); + }); + }, + appendEntry: vi.fn(), + } as unknown as ExtensionAPI; + const account = new PlanUsageController(() => ({ + identity: "signed-out", + projection: { presence: "missing", eligible: false, source: "none", reason: "Sign in" }, + })); + let presentation!: UsagePresentation; + let finish!: () => void; + const host: StepUsageHost = { + account, + resources: () => ({ agents: ["/work/AGENTS.md"], skills: [] }), + scan: vi.fn().mockResolvedValue(empty), + present: (_ctx, model) => { + presentation = model; + return new Promise((resolve) => { + finish = resolve; + model.onClose = resolve; + }); + }, + output: vi.fn(), + }; + const commands = registerStepUsageCommands(pi, host, () => "Ask"); + return { + ctx, + commands, + host, + handlers, + entries, + pi, + presentation: () => presentation, + finish: () => finish(), + changeSession: () => { + sessionId = "second"; + }, + }; +} + +describe("usage command lifecycle", () => { + it("captures actual local status without querying the account or scanning files", async () => { + const f = fixture(), + query = vi.spyOn(f.host.account, "query"); + const pending = f.commands.status("", f.ctx); + expect(f.presentation().status).toMatchObject({ + permissions: "Ask", + session: "first", + agents: ["/work/AGENTS.md"], + reasoning: "high", + }); + expect(query).not.toHaveBeenCalled(); + expect(f.host.scan).not.toHaveBeenCalled(); + f.finish(); + await pending; + }); + it("rejects retired commands before creating UI or starting requests", async () => { + const f = fixture(); + await f.commands.usage("history", f.ctx); + expect(f.ctx.ui.notify).toHaveBeenCalledWith(expect.stringContaining("overview|account"), "warning"); + expect(f.host.scan).not.toHaveBeenCalled(); + }); + it("closes the old view on session replacement and suppresses a late scan result", async () => { + const f = fixture(); + let resolve!: (value: UsageOverview) => void; + f.host.scan = vi.fn().mockImplementation( + () => + new Promise((done) => { + resolve = done; + }), + ); + const pending = f.commands.usage("", f.ctx), + model = f.presentation(); + f.changeSession(); + await f.handlers.get("session_start")!({}, f.ctx); + await pending; + resolve(empty); + await Promise.resolve(); + expect(model.overview).toBeUndefined(); + expect(model.status.session).toBe("first"); + }); + it("cancels old tab scans, only publishes the selected scope, and updates main-session totals", async () => { + const f = fixture(); + const resolvers: ((v: UsageOverview) => void)[] = []; + f.host.scan = vi.fn().mockImplementation(() => new Promise((r) => resolvers.push(r))); + const pending = f.commands.usage("", f.ctx), + model = f.presentation(); + model.selectScope("project"); + resolvers[0]({ ...empty, tokens: 999 }); + resolvers[1]({ ...empty, scope: "project", tokens: 42 }); + await Promise.resolve(); + expect(model.overview?.tokens).toBe(42); + expect(model.overview?.scope).toBe("project"); + model.select("account"); + f.entries.push({ + type: "message", + message: { role: "assistant", usage: { input: 10, output: 2, cacheRead: 3, cacheWrite: 0 } }, + }); + await f.handlers.get("agent_end")!({}, f.ctx); + expect(model.session.total).toBe(15); + f.finish(); + await pending; + }); + it("keeps recorder hooks and command teardown active in the same session lifecycle", async () => { + const f = fixture(); + let replacing = false, + closed = false; + f.host.resources = () => { + if (replacing) expect(closed).toBe(true); + return { agents: [], skills: [{ name: "example", filePath: "/work/SKILL.md" }] }; + }; + await f.handlers.get("session_start")!({}, f.ctx); + await f.handlers.get("before_agent_start")!({}, f.ctx); + expect(f.pi.appendEntry).toHaveBeenCalledWith( + "step.activity_capabilities", + expect.objectContaining({ sessionId: "first" }), + ); + expect(vi.mocked(f.pi.appendEntry).mock.calls.filter(([type]) => type === "step.activity")).toHaveLength(1); + const pending = f.commands.usage("account", f.ctx); + const close = f.presentation().onClose; + f.presentation().onClose = () => { + closed = true; + close?.(); + }; + replacing = true; + f.changeSession(); + await f.handlers.get("session_start")!({}, f.ctx); + await pending; + expect(f.pi.appendEntry).toHaveBeenCalledWith( + "step.activity_capabilities", + expect.objectContaining({ sessionId: "second" }), + ); + expect(vi.mocked(f.pi.appendEntry).mock.calls.filter(([type]) => type === "step.activity")).toHaveLength(2); + }); + + it("uses the explicit RPC sink and does not send command output as an assistant message", async () => { + const f = fixture(); + const rpc = { ...f.ctx, mode: "rpc" } as ExtensionCommandContext; + await f.commands.usage("overview", rpc); + expect(f.host.output).toHaveBeenCalledWith(rpc, expect.objectContaining({ overview: empty })); + expect(f.pi.appendEntry).not.toHaveBeenCalled(); + }); +}); diff --git a/packages/coding-agent/test/step-usage.test.ts b/packages/coding-agent/test/step-usage.test.ts new file mode 100644 index 00000000..1929dd91 --- /dev/null +++ b/packages/coding-agent/test/step-usage.test.ts @@ -0,0 +1,535 @@ +import { createHash } from "node:crypto"; +import { appendFileSync, mkdirSync, mkdtempSync, readFileSync, rmSync, symlinkSync, writeFileSync } from "node:fs"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; +import { afterEach, describe, expect, it, vi } from "vitest"; +import type { SessionEntry, SessionHeader } from "../src/core/session-manager.ts"; +import { createUsageLineage } from "../src/core/usage-lineage.ts"; +import { collectSessionUsage } from "../src/core/usage-totals.ts"; +import { decodePlanUsage, PlanUsageController, resolvePlanUsageAccount, usageText } from "../src/step/usage/account.ts"; +import { type ActivityScanOptions, scanActivity } from "../src/step/usage/activity-engine.ts"; +import { parseUsageCommand } from "../src/step/usage/commands.ts"; +import { projectUsageJson, UsageProjectionLimit } from "../src/step/usage/json-projection.ts"; +import type { PlanAccountContext } from "../src/step/usage/types.ts"; + +const roots: string[] = []; +afterEach(() => { + for (const root of roots.splice(0)) rmSync(root, { recursive: true, force: true }); +}); +function root(): string { + const path = mkdtempSync(join(tmpdir(), "step-usage-test-")); + roots.push(path); + return path; +} +const timestamp = "2026-10-13T10:00:00.000Z"; +function header(id: string, cwd = "/work/A", parentSession?: string): SessionHeader { + return { type: "session", version: 3, id, cwd, timestamp, parentSession }; +} +function row( + id: string, + input: number, + output: number, + cacheRead = 0, + cacheWrite = 0, + role = "assistant", +): Record { + return { + type: "message", + id, + parentId: null, + timestamp, + message: { + role, + content: [{ type: "text", text: "PRIVATE_MESSAGE_MUST_NOT_BE_CACHED" }], + usage: { input, output, cacheRead, cacheWrite }, + }, + }; +} +function caps(id: string, role = "user", parentCwd?: string): Record { + return { + type: "custom", + id: `cap-${id}`, + parentId: null, + timestamp: "2026-10-13T00:00:00.000Z", + customType: "step.activity_capabilities", + data: { + version: 1, + sessionId: id, + sessionRole: role, + parentCwd, + skillDiscovery: true, + skillActivation: true, + pluginInvocation: true, + }, + }; +} +function mainRows(): Record[] { + return [ + caps("main"), + row("a", 100, 5, 20, 10), + row("b", 50, 10, 20), + { + type: "compaction", + id: "c", + parentId: null, + timestamp, + usage: { input: 10, cacheRead: 2, cacheWrite: 0, output: 3 }, + }, + row("d", 4, 1, 0, 0, "toolResult"), + ]; +} +function write(path: string, rows: readonly unknown[]): void { + mkdirSync(join(path, ".."), { recursive: true }); + writeFileSync(path, `${rows.map((r) => JSON.stringify(r)).join("\n")}\n`); +} +function options(path: string, scope: "all" | "project" = "all"): ActivityScanOptions { + return { + roots: [join(path, "sessions")], + cachePath: join(path, "cache", `index-${scope}.jsonl`), + key: `test-${scope}`, + cwd: "/work/A", + scope, + dates: ["2026-10-12", "2026-10-13"], + timezone: "Asia/Shanghai", + }; +} + +describe("usage totals and bounded projection", () => { + it("includes all main entries, AutoCLM increments, and does not sum cumulative custom totals", () => { + const entries = mainRows(); + expect(collectSessionUsage(entries).total).toBe(235); + for (const [id, input, read, write, output] of [ + ["clm1", 10, 2, 0, 3], + ["clm2", 20, 3, 1, 4], + ] as const) + entries.push({ + type: "custom", + customType: "step-auto-clm-usage", + id, + timestamp, + data: { sessionId: "main", usage: { input, cacheRead: read, cacheWrite: write, output } }, + }); + entries.push({ + type: "custom", + customType: "step-auto-clm-result", + data: { usage: { input: 30, cacheRead: 5, cacheWrite: 1, output: 7 } }, + }); + const got = collectSessionUsage(entries); + expect(got.total).toBe(278); + expect(got.input).toBe(252); + expect(got.output).toBe(26); + expect(got.cacheHit).toBeCloseTo(47 / 252); + }); + it("reuses ANSI stripping before bounding account and status labels", () => { + expect(usageText("Plan \x1b[31mPro\x1b[0m\n")).toBe("Plan Pro"); + expect(usageText("\x1b]8;;https://example.com\x07Pro\x1b]8;;\x07")).toBe("Pro"); + expect(usageText("x".repeat(300))).toHaveLength(256); + }); + + it("keeps invalid and missing usage separate from real zero", () => { + const entries = [row("valid", 0, 0), { type: "message", message: { role: "assistant" } }, row("bad", -1, 0)]; + expect(collectSessionUsage(entries)).toMatchObject({ total: 0, missing: 1, invalid: 1, cacheHit: null }); + }); + it("validates ignored content while retaining only statistics", () => { + const projected = projectUsageJson(JSON.stringify(row("a", 10, 2))); + expect(projected.message).not.toHaveProperty("content"); + expect(collectSessionUsage([projected]).total).toBe(12); + expect(() => projectUsageJson('{"type":"message","message":{"content":[tru]}}')).toThrow(SyntaxError); + expect(() => projectUsageJson('{"id":"a","id":"b"}')).toThrow(SyntaxError); + expect(() => projectUsageJson(`{"unknown":${"[".repeat(70)}0${"]".repeat(70)}}`)).toThrow(UsageProjectionLimit); + expect(() => projectUsageJson(JSON.stringify({ id: "a".repeat(5000) }))).toThrow(UsageProjectionLimit); + }); + it("accepts escaped JSON strings and ignores large unrelated strings", () => { + const raw = JSON.stringify({ + ...row("a", 10, 2), + ignored: "x".repeat(1024 * 1024), + data: { name: 'skill\n"name\u4e2d', ignored: { deep: [1, false, null] } }, + }); + expect(projectUsageJson(raw).data).toMatchObject({ name: 'skill\n"name中' }); + }); +}); + +describe("annual local activity and cache", () => { + it("retains dispatch provenance when a fork selects the earlier assistant leaf", async () => { + const path = root(), + file = join(path, "sessions", "fork.jsonl"), + source = mainRows() as unknown as SessionEntry[]; + const provenance = { + type: "custom", + id: "provenance", + timestamp, + customType: "step.usage_provenance", + data: { + version: 1, + sessionId: "main", + originSessionId: "main", + usageEntryId: "a", + reasoningAtDispatch: "high", + }, + } as SessionEntry; + const copied = source.slice(0, 2), + lineage = createUsageLineage(header("main"), [...source, provenance], copied, "fork", timestamp, "a"); + write(file, [header("fork", "/work/B"), ...copied, ...lineage]); + const got = await scanActivity(options(path, "project")); + expect(got.tokens).toBe(135); + expect(got.reasoning).toContainEqual(["high", 1]); + }); + it("skips directory/file symlinks and concurrent cache readers do not duplicate totals", async () => { + const path = root(), + source = join(path, "sessions", "main.jsonl"), + outside = join(path, "outside"); + write(source, [header("main"), ...mainRows()]); + write(join(outside, "extra.jsonl"), [header("extra"), row("extra", 999, 0)]); + symlinkSync(outside, join(path, "sessions", "linked"), "dir"); + symlinkSync(join(outside, "extra.jsonl"), join(path, "sessions", "linked.jsonl")); + const results = await Promise.all([scanActivity(options(path)), scanActivity(options(path))]); + expect(results.map((x) => x.tokens)).toEqual([235, 235]); + expect((await scanActivity(options(path))).tokens).toBe(235); + }); + + it("deduplicates forks and excludes child chats, including project attribution", async () => { + const path = root(), + main = join(path, "sessions", "A", "main.jsonl"), + fork = join(path, "sessions", "B", "fork.jsonl"), + child = join(path, "sessions", "C", "child.jsonl"); + write(main, [header("main"), ...mainRows()]); + write(fork, [header("fork", "/work/B", main), ...mainRows(), caps("fork"), row("new", 20, 2)]); + write(child, [header("child", "/work/isolated"), caps("child", "subagent", "/work/A"), row("child-call", 30, 3)]); + const checksum = createHash("sha256").update(readFileSync(main)).digest("hex"); + const all = await scanActivity(options(path)); + expect(all).toMatchObject({ tokens: 290, chats: 2 }); + expect(all.days[1]).toMatchObject({ tokens: 290, chats: 2 }); + const project = await scanActivity(options(path, "project")); + expect(project).toMatchObject({ tokens: 268, chats: 1 }); + expect(createHash("sha256").update(readFileSync(main)).digest("hex")).toBe(checksum); + expect(readFileSync(options(path).cachePath, "utf8")).not.toContain("PRIVATE_MESSAGE"); + expect((await scanActivity(options(path))).tokens).toBe(290); + }); + it("reuses complete prefixes, includes a completed tail once, and detects prefix rewrite plus growth", async () => { + const path = root(), + file = join(path, "sessions", "main.jsonl"), + opts = options(path); + write(file, [header("main"), ...mainRows()]); + expect((await scanActivity(opts)).tokens).toBe(235); + const tail = JSON.stringify(row("tail", 20, 2)); + appendFileSync(file, tail.slice(0, -1)); + expect((await scanActivity(opts)).tokens).toBe(235); + appendFileSync(file, "}\n"); + expect((await scanActivity(opts)).tokens).toBe(257); + expect((await scanActivity(opts)).tokens).toBe(257); + const changed = mainRows(); + changed[1] = row("a", 101, 5, 20, 10); + write(file, [header("main"), ...changed, row("tail", 20, 2), row("more", 30, 3)]); + expect((await scanActivity(opts)).tokens).toBe(291); + }); + it("keeps coverage limitations through a cache hit and rebuilds corrupted cache", async () => { + const path = root(), + file = join(path, "sessions", "main.jsonl"), + opts = options(path); + write(file, [header("main"), ...mainRows()]); + appendFileSync(file, "{broken}\n"); + expect((await scanActivity(opts)).limited).toContain("Invalid records"); + expect((await scanActivity(opts)).limited).toContain("Invalid records"); + writeFileSync(opts.cachePath, "bad cache\n"); + expect((await scanActivity(opts)).tokens).toBe(235); + }); + it("preserves a multi-hop fork's original project after source files are removed", async () => { + const path = root(), + file = join(path, "sessions", "leaf.jsonl"), + main = mainRows() as unknown as SessionEntry[], + forkHeader = header("fork", "/work/B", "/removed/main.jsonl"); + const first = createUsageLineage(header("main"), main, main, "fork", timestamp, "d"); + const branch = [...main, ...first, row("new", 20, 2), caps("fork")] as unknown as SessionEntry[]; + const second = createUsageLineage(forkHeader, branch, branch, "leaf", timestamp, "cap-fork"); + write(file, [header("leaf", "/work/C", "/removed/fork.jsonl"), ...branch, ...second]); + expect((await scanActivity(options(path))).tokens).toBe(257); + expect((await scanActivity(options(path, "project"))).tokens).toBe(235); + }); + it("uses started date and independent resource dimensions, not completion time", async () => { + const path = root(), + file = join(path, "sessions", "main.jsonl"); + write(file, [ + header("main"), + caps("main"), + { + type: "custom", + id: "run1", + timestamp, + customType: "step.activity", + data: { + sessionId: "main", + kind: "skill", + stage: "started", + runId: "run", + resourceId: "skill", + name: "$skill", + startedAt: Date.parse("2026-10-12T12:00:00Z"), + pluginId: "plugin", + }, + }, + { + type: "custom", + id: "run2", + timestamp, + customType: "step.activity", + data: { + sessionId: "main", + kind: "plugin", + stage: "started", + runId: "run", + resourceId: "plugin", + name: "@plugin", + startedAt: Date.parse("2026-10-12T12:00:00Z"), + }, + }, + ]); + const got = await scanActivity(options(path)); + expect(got.topResources).toContainEqual(["$skill", 1]); + expect(got.topResources).toContainEqual(["@plugin", 1]); + expect(got.chats).toBe(0); + }); +}); + +function snapshot(): Record { + return { + schema_version: 1, + observed_at_ms: Date.now(), + scope: "personal_account", + account: { id: "123" }, + plan: { state: "known", family: "coding", name: "Pro", subscription: { id: "7", expires_at_ms: null } }, + quota: { + state: "known", + allowed: true, + reason: null, + windows: [ + { id: "five_hour", limited: true, used_percent: 0, resets_at_ms: null }, + { id: "weekly", limited: false, used_percent: null, resets_at_ms: null }, + ], + credit_buckets: [], + }, + }; +} +const accountContext = (identity = "a"): PlanAccountContext => ({ + identity, + key: "private-key", + projection: { + presence: "present", + eligible: true, + source: "plan", + accountId: "123", + endpoint: "https://api.stepfun.com/step_plan/v1/usage", + }, +}); +const response = (body: unknown = snapshot(), status = 200, headers?: Record): Response => + new Response(JSON.stringify(body), { status, headers }); + +describe("account schema and lifecycle", () => { + it("refreshes a cached snapshot after its reset boundary before the normal TTL", async () => { + let wall = Date.now(); + const body = snapshot(); + (body.quota as Record).windows = [ + { id: "five_hour", limited: true, used_percent: 0, resets_at_ms: wall + 100 }, + { id: "weekly", limited: false, used_percent: null, resets_at_ms: null }, + ]; + const fetchFn = vi.fn().mockImplementation(async () => response(body)); + const controller = new PlanUsageController( + () => accountContext(), + fetchFn, + () => wall, + ); + await controller.query(); + wall += 101; + expect(controller.readCached().status).toBe("stale"); + await controller.query(); + expect(fetchFn).toHaveBeenCalledTimes(2); + await controller.query(); + expect(fetchFn).toHaveBeenCalledTimes(2); + }); + + it("uses monotonic TTL/cooldown when wall time changes and preserves long Retry-After", async () => { + let wall = Date.now(), + tick = 0; + const fetchFn = vi.fn().mockResolvedValue(response()); + const controller = new PlanUsageController( + () => accountContext(), + fetchFn, + () => wall, + 12000, + () => tick, + ); + await controller.query(); + wall -= 3600000; + tick += 31000; + expect(controller.readCached().status).toBe("stale"); + fetchFn.mockResolvedValue(response({}, 429, { "Retry-After": "604800" })); + await controller.query(); + wall += 36000000; + tick += 86401000; + await controller.query(true); + expect(fetchFn).toHaveBeenCalledTimes(2); + tick += 7 * 86400000; + await controller.query(true); + expect(fetchFn).toHaveBeenCalledTimes(3); + }); + + it("supports genuine zero and unknown future bucket types without losing integer precision", () => { + expect(decodePlanUsage(snapshot()).quota.windows[0].used_percent).toBe(0); + const body = snapshot(); + body.plan = { state: "known", family: "token", name: null, subscription: null }; + body.quota = { + state: "known", + allowed: true, + reason: null, + windows: [], + credit_buckets: [ + { + type: "topup", + total: "9007199254740993", + remaining: "42", + used_percent: 99, + resets_at_ms: null, + expires_at_ms: null, + }, + { type: "future-type" }, + ], + }; + const got = decodePlanUsage(body); + expect(got.quota.credit_buckets?.[0].total).toBe("9007199254740993"); + expect(got.quota.credit_buckets?.[1].unsupported).toBe(true); + }); + it("rejects invalid cross-family, identity and numeric snapshots", () => { + for (const change of [ + (x: Record) => { + x.schema_version = 2; + }, + (x: Record) => { + x.account = { id: "0" }; + }, + (x: Record) => { + (x.quota as Record).windows = [ + { id: "five_hour", limited: true, used_percent: NaN, resets_at_ms: null }, + ]; + }, + ]) { + const body = snapshot(); + change(body); + expect(() => decodePlanUsage(body)).toThrow(); + } + }); + it("does not send gateway or platform keys to a personal endpoint", () => { + const path = root(), + auth = join(path, "auth.json"); + writeFileSync( + auth, + JSON.stringify({ step: { type: "oauth", access: "private-key", uid: "123", profile: "platform_cn" } }), + ); + expect(resolvePlanUsageAccount(auth, {}).projection.eligible).toBe(false); + expect(resolvePlanUsageAccount(auth, { STEP_API_KEY: "gateway" }).key).toBeUndefined(); + expect(resolvePlanUsageAccount(auth, { STEP_API_KEY: "gateway" }).projection.accountId).toBeUndefined(); + writeFileSync( + auth, + JSON.stringify({ step: { type: "oauth", access: "private-key", uid: "123", profile: "step_plan_oversea" } }), + ); + expect(resolvePlanUsageAccount(auth, {}).projection.endpoint).toContain("stepfun.ai/step_plan/v1/usage"); + }); + it("shares requests, caches results and retains snapshots on visible failure", async () => { + let now = Date.now(); + const fetchFn = vi.fn().mockResolvedValue(response()); + const controller = new PlanUsageController( + () => accountContext(), + fetchFn, + () => now, + ); + await Promise.all([controller.query(), controller.query()]); + expect(fetchFn).toHaveBeenCalledTimes(1); + await controller.query(); + expect(fetchFn).toHaveBeenCalledTimes(1); + now += 31000; + fetchFn.mockResolvedValue(response({}, 503)); + expect((await controller.query()).status).toBe("stale"); + expect(controller.readCached().snapshot).toBeDefined(); + }); + it("settles timeout even when fetch ignores abort and late 401 cannot erase a newer snapshot", async () => { + let finish!: (r: Response) => void; + let calls = 0; + const fetchFn = vi.fn().mockImplementation(() => + ++calls === 1 + ? new Promise((r) => { + finish = r; + }) + : Promise.resolve(response()), + ); + const controller = new PlanUsageController(() => accountContext(), fetchFn, Date.now, 10); + await controller.query(); + expect(controller.readCached().status).toBe("error"); + await controller.query(true); + expect(controller.readCached().status).toBe("ready"); + finish(response({}, 401)); + await new Promise((r) => setTimeout(r, 0)); + expect(controller.readCached().status).toBe("ready"); + }); + it("drops old identity results and cancel settles promptly without waiting for fetch", async () => { + let identity = "a", + finish!: (r: Response) => void; + const fetchFn = vi + .fn() + .mockImplementationOnce( + () => + new Promise((r) => { + finish = r; + }), + ) + .mockResolvedValue(response()); + const controller = new PlanUsageController(() => accountContext(identity), fetchFn); + const old = controller.query(); + identity = "b"; + await controller.query(); + finish(response({}, 401)); + await old; + expect(controller.readCached().status).toBe("ready"); + const blocked = new PlanUsageController( + () => accountContext(), + vi.fn().mockImplementation(() => new Promise(() => {})), + ); + const pending = blocked.query(); + blocked.cancel(); + await pending; + expect(blocked.readCached().status).toBe("idle"); + }); + it("enforces Retry-After without changing stored credentials", async () => { + let now = Date.now(); + const fetchFn = vi.fn().mockResolvedValue(response({}, 429, { "Retry-After": "30" })); + const controller = new PlanUsageController( + () => accountContext(), + fetchFn, + () => now, + ); + await controller.query(); + await controller.query(true); + expect(fetchFn).toHaveBeenCalledTimes(1); + now += 31000; + await controller.query(); + expect(fetchFn).toHaveBeenCalledTimes(2); + }); + it("bounds response body reads and rejects cross-account responses", async () => { + const large = new PlanUsageController( + () => accountContext(), + vi.fn().mockResolvedValue(new Response("x".repeat(256 * 1024 + 1))), + ); + expect((await large.query()).status).toBe("error"); + const body = snapshot(); + body.account = { id: "999" }; + const mismatched = new PlanUsageController( + () => accountContext(), + vi.fn().mockResolvedValue(response(body)), + ); + expect((await mismatched.query()).snapshot).toBeUndefined(); + }); +}); + +it("accepts only the two supported usage views and refresh forms", () => { + for (const command of ["", "overview", "account", "--refresh", "account --refresh"]) + expect(parseUsageCommand(command)).toBeDefined(); + for (const command of ["session", "history", "overview --refresh", "account extra", "--refresh extra"]) + expect(parseUsageCommand(command)).toBeUndefined(); +}); diff --git a/scripts/build-binaries.sh b/scripts/build-binaries.sh index c4ccf3a8..4ba6fc5e 100755 --- a/scripts/build-binaries.sh +++ b/scripts/build-binaries.sh @@ -259,9 +259,9 @@ for platform in "${PLATFORMS[@]}"; do # Disable cwd bunfig.toml autoload so project preload scripts cannot crash the # standalone binary before step starts (see #7684). if [[ "$platform" == windows-* ]]; then - bun build --compile --no-compile-autoload-bunfig '--env=STEPCODE_BUILD_*' --target="$bun_target" ../../apps/cli/dist/bun/stepcode.js ./src/utils/image-resize-worker.ts --outfile "$OUTPUT_DIR/$platform/step.exe" + bun build --compile --no-compile-autoload-bunfig '--env=STEPCODE_BUILD_*' --target="$bun_target" ../../apps/cli/dist/bun/stepcode.js ./src/utils/image-resize-worker.ts ./src/step/usage/activity-worker.ts --outfile "$OUTPUT_DIR/$platform/step.exe" else - bun build --compile --no-compile-autoload-bunfig '--env=STEPCODE_BUILD_*' --target="$bun_target" ../../apps/cli/dist/bun/stepcode.js ./src/utils/image-resize-worker.ts --outfile "$OUTPUT_DIR/$platform/step" + bun build --compile --no-compile-autoload-bunfig '--env=STEPCODE_BUILD_*' --target="$bun_target" ../../apps/cli/dist/bun/stepcode.js ./src/utils/image-resize-worker.ts ./src/step/usage/activity-worker.ts --outfile "$OUTPUT_DIR/$platform/step" fi done diff --git a/scripts/build-coding-agent-bundle.mjs b/scripts/build-coding-agent-bundle.mjs index 59fdc3eb..2a3c71d0 100644 --- a/scripts/build-coding-agent-bundle.mjs +++ b/scripts/build-coding-agent-bundle.mjs @@ -204,6 +204,7 @@ const mainResult = await build({ }); const imageResizeOutput = findContainingOutput(mainResult.metafile, "packages/coding-agent/dist/utils/image-resize.js"); +const activityOutput = findContainingOutput(mainResult.metafile, "packages/coding-agent/dist/step/usage/activity.js"); // The image-resize worker is reached through a worker URL, so the main bundle // cannot follow it. Emit a self-contained worker file beside the code that @@ -219,8 +220,10 @@ const lazyResult = await build({ }); validateExternalImports([mainResult.metafile, lazyResult.metafile]); +const activityResult = await build({ ...commonBuildOptions(), entryNames: "[name]", entryPoints: { "activity-worker": join(codingAgentDistDir, "step", "usage", "activity-worker.js") }, outdir: dirname(activityOutput), splitting: false }); +validateExternalImports([activityResult.metafile]); chmodSync(join(bundleDir, "step.js"), 0o755); -const files = new Set([...Object.keys(mainResult.metafile.outputs), ...Object.keys(lazyResult.metafile.outputs)]).size; -const mib = outputBytes([mainResult.metafile, lazyResult.metafile]) / (1024 * 1024); +const files = new Set([...Object.keys(mainResult.metafile.outputs), ...Object.keys(lazyResult.metafile.outputs), ...Object.keys(activityResult.metafile.outputs)]).size; +const mib = outputBytes([mainResult.metafile, lazyResult.metafile, activityResult.metafile]) / (1024 * 1024); console.log(`Built ${relative(repoRoot, bundleDir)} (${files} files, ${mib.toFixed(1)} MiB)`); diff --git a/scripts/smoke-usage-worker.mjs b/scripts/smoke-usage-worker.mjs new file mode 100644 index 00000000..801fe9cf --- /dev/null +++ b/scripts/smoke-usage-worker.mjs @@ -0,0 +1,44 @@ +#!/usr/bin/env node +import assert from "node:assert/strict"; +import { mkdirSync, mkdtempSync, rmSync, writeFileSync } from "node:fs"; +import { tmpdir } from "node:os"; +import { join } from "node:path"; +import { scanUsageOverview } from "../packages/coding-agent/dist/bundle/index.js"; + +// Run after build:offline. --benchmark uses the annual-index acceptance fixture. +const files = process.argv.includes("--benchmark") ? 10000 : 4; +const root = mkdtempSync(join(tmpdir(), "step-usage-worker-")); +const sessions = join(root, "sessions"); +mkdirSync(sessions); +let ticker; +try { + const timestamp = new Date().toISOString(); + let bytes = 0; + for (let i = 0; i < files; i++) { + const sessionId = `session-${i}`; + const rows = [ + { type: "session", id: sessionId, cwd: root, timestamp }, + { type: "custom", id: "cap", customType: "step.activity_capabilities", timestamp, data: { version: 1, sessionId, sessionRole: "user" } }, + { type: "message", id: "usage", timestamp, message: { role: "assistant", content: [{ type: "text", text: "x".repeat(32768) }], usage: { input: 100, output: 5, cacheRead: 20, cacheWrite: 10 } } }, + ]; + const line = `${rows.map((row) => JSON.stringify(row)).join("\n")}\n`; + bytes += Buffer.byteLength(line); + writeFileSync(join(sessions, `${i}.jsonl`), line); + } + const gaps = []; + let last = performance.now(); + ticker = setInterval(() => { const now = performance.now(); gaps.push(now - last); last = now; }, 10); + const scan = (signal = new AbortController().signal) => scanUsageOverview([sessions], join(root, "cache"), root, "all", signal); + const start = performance.now(), first = await scan(), fullMs = performance.now() - start; + const cachedStart = performance.now(), cached = await scan(), cacheMs = performance.now() - cachedStart; + assert.equal(first.tokens, files * 135); assert.equal(cached.tokens, first.tokens); assert.equal(first.chats, files); + assert(!first.limited.includes("File inventory budget exceeded")); + const controller = new AbortController(), cancelStart = performance.now(), pending = scan(controller.signal); + setTimeout(() => controller.abort(), files === 4 ? 0 : 20); + await assert.rejects(pending, /cancelled/u); + gaps.sort((a, b) => a - b); + console.log(JSON.stringify({ files, MiB: bytes / 1024 / 1024, fullMs, cacheMs, cancelMs: performance.now() - cancelStart, mainThreadTickP95: gaps[Math.floor(gaps.length * 0.95)], rssMiB: process.memoryUsage().rss / 1024 / 1024, tokens: first.tokens, chats: first.chats, coverage: first.limited }, null, 2)); +} finally { + clearInterval(ticker); + rmSync(root, { recursive: true, force: true }); +}