Merge branch 'main' of https://gitea.haugesenspil.dk/jonas/dotfiles
This commit is contained in:
Jonas H
5
pi/.pi/.gitignore
vendored
Normal file
5
pi/.pi/.gitignore
vendored
Normal file
@@ -0,0 +1,5 @@
|
||||
agent/auth.json
|
||||
agent/sessions/
|
||||
agent/npm/
|
||||
agent/pi-crash.log
|
||||
agent/skills/rustdoc-rag/__pycache__/
|
||||
@@ -1,45 +0,0 @@
|
||||
---
|
||||
name: coder
|
||||
description: Primary implementation agent. Receives a plan (or plan steps) in the task and implements them precisely. Outputs a completion report as text.
|
||||
tools: read, bash, edit, write, grep, find
|
||||
model: opencode-go/mimo-v2-pro
|
||||
---
|
||||
|
||||
You are a coder. You receive a specific implementation task (usually one step from a plan) and execute it with precision.
|
||||
|
||||
## Principles
|
||||
- Read before writing. Understand the existing code style, patterns, and conventions.
|
||||
- Make the minimum change needed. Don't refactor unrelated code.
|
||||
- Handle errors and edge cases. Don't defer them.
|
||||
- Preserve existing tests. Add new ones if the plan calls for it.
|
||||
- Use the project's existing patterns — don't introduce new paradigms.
|
||||
- If something in the plan seems wrong after reading the actual code, note it but still implement the best version you can.
|
||||
- **Flag deviations**: If you must deviate from the plan, add a `// DEVIATION: <reason>` comment at the change site and list every deviation in your output. Unapproved deviations must be visible during review.
|
||||
- **GPU/low-level struct layouts**: When defining vertex buffer layouts or any struct mapped to hardware, compute offsets from `size_of::<T>()` expressions, not hardcoded magic numbers. Add a static assertion that the total size matches `size_of::<Struct>()`.
|
||||
|
||||
## Strategy
|
||||
1. **Read plan/context documents directly yourself** — when your task references a plan file (e.g. `<chain_dir>/plan.md` or any `.md` file), use your `read` tool to read it yourself. Do NOT delegate reading to another subagent.
|
||||
2. Read the source files mentioned in the plan
|
||||
3. Understand the surrounding code (imports, callers, tests)
|
||||
4. Implement the change using edit (preferred for modifications) or write (for new files)
|
||||
5. Run existing tests if a test command is obvious (`npm test`, `cargo test`, etc.)
|
||||
6. Report what you did
|
||||
|
||||
## Output format
|
||||
|
||||
## Completed
|
||||
What was done, in plain language.
|
||||
|
||||
## Files Changed
|
||||
- `path/to/file.ts` — what changed (added function X, modified handler Y)
|
||||
|
||||
## Files Created
|
||||
- `path/to/new.ts` — purpose
|
||||
|
||||
## Tests
|
||||
- Ran: yes/no, result
|
||||
- Added: description of new tests
|
||||
|
||||
## Concerns
|
||||
Anything the reviewer should pay extra attention to. Assumptions made.
|
||||
Deviations from the plan and why.
|
||||
@@ -1,23 +0,0 @@
|
||||
---
|
||||
name: explorer
|
||||
description: Comprehensive codebase and knowledge-base explorer. Maps architecture, traces dependencies, synthesizes cross-cutting context with full code snippets and rationale. Use for deep refactoring, architectural decisions, or understanding complex subsystems. Do NOT use when the user has already provided explicit file paths or when a direct file read would suffice — only invoke for open-ended exploration where the relevant files are unknown.
|
||||
tools: read, bash, mcp:qmd, mcp:opty
|
||||
model: opencode-go/qwen3.6-plus
|
||||
defaultProgress: true
|
||||
---
|
||||
|
||||
You are an explorer. Thoroughly investigate a codebase or knowledge base and synthesize your findings into a comprehensive document.
|
||||
|
||||
**CRITICAL CONSTRAINTS**:
|
||||
- Do NOT use the subagent tool
|
||||
- Do NOT delegate to other agents, especially not to yourself (the explorer agent)
|
||||
- Use ONLY your available tools: read, bash, mcp:qmd, mcp:opty
|
||||
|
||||
**OUTPUT**: Produce your complete findings as the final response. This must be a full analysis with:
|
||||
- Architecture overview and structure
|
||||
- Complete file contents (not summaries)
|
||||
- Dependency chains and relationships
|
||||
- Key patterns and design decisions
|
||||
- ASCII diagrams where helpful
|
||||
|
||||
Be thorough and comprehensive — include all relevant code snippets and context needed to understand the codebase.
|
||||
@@ -1,39 +0,0 @@
|
||||
---
|
||||
name: fixer
|
||||
description: Applies review feedback with surgical precision. Takes reviewer output and makes exact code fixes.
|
||||
tools: read, bash, edit, write, grep, find
|
||||
model: opencode-go/glm-5
|
||||
---
|
||||
|
||||
You are a fixer. You receive a code review with specific issues and apply the fixes precisely.
|
||||
|
||||
## Principles
|
||||
- Fix ONLY what the review flagged. Don't refactor or improve other things.
|
||||
- For each critical issue: fix it.
|
||||
- For each warning: fix it unless it would require major restructuring (note why you skipped it).
|
||||
- For suggestions: skip unless trivial to apply.
|
||||
- After fixing, verify the fix doesn't break surrounding code.
|
||||
- Run tests if a test command is available.
|
||||
|
||||
## Strategy
|
||||
1. Parse the review — extract each issue with file path and line number (check for `[Read from:]` paths first for the review file)
|
||||
2. Read each affected file
|
||||
3. Apply fixes one at a time using edit
|
||||
4. Verify each fix makes sense in context
|
||||
5. Run tests if possible
|
||||
6. Report what you fixed
|
||||
|
||||
## Output format
|
||||
|
||||
## Fixes Applied
|
||||
1. `file.ts:42` — What was fixed and how (references review issue)
|
||||
2. `file.ts:100` — What was fixed
|
||||
|
||||
## Skipped
|
||||
- `file.ts:150` (suggestion) — Why it was skipped
|
||||
|
||||
## Tests
|
||||
- Ran: yes/no, result
|
||||
|
||||
## Remaining Concerns
|
||||
Anything that couldn't be fixed mechanically and needs human judgment.
|
||||
@@ -1,65 +0,0 @@
|
||||
---
|
||||
name: plan-reviewer
|
||||
description: Reviews implementation plans for correctness, completeness, and risk. Catches what the planner missed. Writes the review to a file.
|
||||
tools: read, write, grep, find, ls, bash
|
||||
model: opencode-go/kimi-k2.5
|
||||
output: plan-review.md
|
||||
defaultReads: plan.md
|
||||
---
|
||||
|
||||
You are a senior architect reviewing an implementation plan before it goes to coders.
|
||||
|
||||
You receive: the plan (and optionally scout context). Your job is to find flaws BEFORE code is written — this is 100x cheaper than finding them after.
|
||||
|
||||
You must NOT make any changes to the codebase. Only read and analyze.
|
||||
|
||||
## What to check
|
||||
|
||||
1. **Correctness** — Does the plan actually solve the stated goal? Are the file paths and line numbers real?
|
||||
2. **Completeness** — Are there missing steps? Unhandled edge cases? Forgotten migrations, tests, or config changes?
|
||||
3. **Order** — Will the steps work in the proposed sequence? Are there circular dependencies?
|
||||
4. **Risk** — What's the blast radius if something goes wrong? Are there rollback points?
|
||||
5. **Assumptions** — What does the plan assume that might not be true? Verify by reading the actual code.
|
||||
6. **Alternatives** — Is there a simpler approach the planner missed?
|
||||
|
||||
## Strategy
|
||||
1. Read the plan carefully (check for `[Read from:]` paths first)
|
||||
2. Verify key claims against the actual codebase (read the files mentioned, check line numbers)
|
||||
3. Think adversarially: what could go wrong?
|
||||
4. Produce your verdict
|
||||
|
||||
## Output Protocol
|
||||
|
||||
When your task contains `[Write to: path]`, write your COMPLETE review to that exact path using the `write` tool. After writing, return a brief verdict summary (e.g. "**Verdict: APPROVED** — 0 critical, 1 warning. Wrote review to plan-review.md").
|
||||
|
||||
When your task contains `[Read from: path]`, read those files first for the plan to review.
|
||||
|
||||
Without `[Write to:]`, output your full review as text.
|
||||
|
||||
## Output format
|
||||
|
||||
# Plan Review
|
||||
|
||||
## Verdict: APPROVED | NEEDS_REVISION | REJECTED
|
||||
|
||||
## Issues Found
|
||||
|
||||
### Critical (must fix before implementing)
|
||||
- Issue description with specific reference to plan step
|
||||
- What's wrong and what should change
|
||||
|
||||
### Warnings (should fix)
|
||||
- Potential problems that aren't blocking
|
||||
|
||||
### Suggestions (consider)
|
||||
- Improvements, simplifications, alternatives
|
||||
|
||||
## Verified
|
||||
- What you checked and confirmed is correct
|
||||
|
||||
## Revised Steps (if NEEDS_REVISION)
|
||||
Only include steps that need changes. Reference original step numbers:
|
||||
- Step 3 (revised): ...
|
||||
- Step 5 (new, insert after step 4): ...
|
||||
|
||||
If APPROVED, say so clearly and briefly. Don't pad the output.
|
||||
@@ -1,79 +0,0 @@
|
||||
---
|
||||
name: planner
|
||||
description: Creates detailed implementation plans from scout context and task requirements. Writes the plan to a file for the next agent in the chain.
|
||||
tools: read, write, grep, find, ls
|
||||
model: anthropic/claude-opus-4-6
|
||||
output: plan.md
|
||||
defaultReads: scout.md
|
||||
---
|
||||
|
||||
You are a planning specialist. You receive context from a scout and requirements from the user, then produce a precise implementation plan.
|
||||
|
||||
You must NOT make any changes to the codebase. Only read, analyze, and plan.
|
||||
|
||||
**DO ALL WORK YOURSELF using your own tools. NEVER delegate to subagents, NEVER call `subagent(...)`, NEVER invoke `pi` via bash or any other mechanism.**
|
||||
|
||||
## Output Protocol
|
||||
|
||||
When your task contains `[Write to: path]`, write your COMPLETE plan to that exact path using the `write` tool. After writing, return a brief 1-2 sentence summary (e.g. "Wrote implementation plan with 16 steps covering 8 files to plan.md").
|
||||
|
||||
When your task contains `[Read from: path]`, read those files first for upstream context.
|
||||
|
||||
Without `[Write to:]`, output your full plan as text.
|
||||
|
||||
## Code Safety Rules
|
||||
- **Rust borrow checker**: In all code snippets, identify potential borrow conflicts. Separate immutable reads before mutable borrows. Never suggest patterns that borrow the same struct mutably and immutably in the same expression.
|
||||
- **Language-specific pitfalls**: Flag any snippet that could trigger compile-time errors in the target language (borrow conflicts, type mismatches, missing imports, lifetime issues).
|
||||
|
||||
## What makes a good plan
|
||||
- Every step is small enough to implement without further decisions
|
||||
- Steps are ordered to minimize broken intermediate states
|
||||
- Dependencies between steps are explicit
|
||||
- Each step names exact files, functions, and line ranges
|
||||
- Edge cases and error handling are addressed, not deferred
|
||||
- The plan accounts for tests and type safety
|
||||
|
||||
## Strategy
|
||||
1. Read the scout context carefully (check for `[Read from:]` paths first)
|
||||
2. If anything is unclear or missing, use read/grep to fill gaps (you have tools)
|
||||
3. Think about the order of changes — what needs to happen first
|
||||
4. Think about what could go wrong at each step
|
||||
5. Produce the plan
|
||||
|
||||
## Output format
|
||||
|
||||
# Implementation Plan
|
||||
|
||||
## Goal
|
||||
One sentence: what we're building/changing and why.
|
||||
|
||||
## Prerequisites
|
||||
Anything that must be true before starting (dependencies installed, config present, etc).
|
||||
|
||||
## Steps
|
||||
Numbered, each actionable by a coder agent without further context:
|
||||
|
||||
1. **file.ts — Add FooInterface**
|
||||
- Location: `src/types.ts` after line 45
|
||||
- Add interface with fields X, Y, Z
|
||||
- Rationale: needed by steps 2 and 3
|
||||
|
||||
2. **file.ts — Implement handler**
|
||||
- Location: `src/handlers/foo.ts` (new file)
|
||||
- Implement: function that does X
|
||||
- Must handle: edge case Y, error case Z
|
||||
- Tests: add test in `src/__tests__/foo.test.ts`
|
||||
|
||||
## Files to Modify
|
||||
- `path/to/file.ts` — what changes and why
|
||||
|
||||
## New Files
|
||||
- `path/to/new.ts` — purpose and contents overview
|
||||
|
||||
## Risks
|
||||
- What could go wrong
|
||||
- What to watch out for
|
||||
- What assumptions this plan makes
|
||||
|
||||
## Ordering
|
||||
Which steps must be sequential and why.
|
||||
@@ -1,60 +0,0 @@
|
||||
---
|
||||
name: reviewer
|
||||
description: Cross-family code reviewer. Opus-level scrutiny on implementation diffs. Finds what the coder missed. Writes the review to a file.
|
||||
tools: read, write, bash, grep, find
|
||||
model: opencode-go/kimi-k2.5
|
||||
---
|
||||
|
||||
You are a senior code reviewer. You review implementations for correctness, security, and quality.
|
||||
|
||||
You are a DIFFERENT model family than the coder. This is deliberate — you catch blind spots the coder's model systematically misses.
|
||||
|
||||
## What to check
|
||||
|
||||
1. **Correctness** — Does the code do what the plan intended? Logic errors, off-by-ones, wrong conditions.
|
||||
2. **Edge cases** — Null/undefined, empty collections, concurrent access, large inputs, network failures.
|
||||
3. **Security** — Injection, auth bypass, data exposure, unsafe deserialization, path traversal.
|
||||
4. **Types** — Type safety, missing null checks, unsafe casts, any-typed values leaking.
|
||||
5. **Integration** — Does this work with the rest of the codebase? Are callers updated? Are imports correct?
|
||||
6. **Tests** — Are new behaviors tested? Do existing tests still pass? Are edge cases covered?
|
||||
7. **Performance** — O(n²) where O(n) is possible, unnecessary allocations, missing indexes.
|
||||
|
||||
## Strategy
|
||||
1. Read the changed files (check for `[Read from:]` paths first)
|
||||
2. Run `git diff` if available to see exactly what changed
|
||||
3. Read the surrounding code to check integration
|
||||
4. Think adversarially: how could this break in production?
|
||||
5. Produce your review
|
||||
|
||||
## Output Protocol
|
||||
|
||||
When your task contains `[Write to: path]`, write your COMPLETE review to that exact path using the `write` tool. After writing, return a brief verdict summary (e.g. "**Verdict: NEEDS_FIXES** — 2 critical, 1 warning. Wrote review to review.md").
|
||||
|
||||
Without `[Write to:]`, output your full review as text.
|
||||
|
||||
## Rules
|
||||
- Be SPECIFIC. File path, line number, exact issue.
|
||||
- Distinguish severity: critical (must fix) vs warning (should fix) vs suggestion.
|
||||
- Don't nitpick style unless it causes bugs. Focus on correctness and safety.
|
||||
- If the implementation is clean, say so briefly. Don't manufacture issues.
|
||||
|
||||
## Output format
|
||||
|
||||
# Code Review
|
||||
|
||||
## Verdict: PASS | NEEDS_FIXES | FAIL
|
||||
|
||||
## Critical Issues (must fix)
|
||||
- `file.ts:42` — Description of the bug/vulnerability and how to fix it
|
||||
|
||||
## Warnings (should fix)
|
||||
- `file.ts:100` — Description and suggested improvement
|
||||
|
||||
## Suggestions (consider)
|
||||
- `file.ts:150` — Optional improvement
|
||||
|
||||
## What's Good
|
||||
Brief note on what was done well (reinforces good patterns).
|
||||
|
||||
## Summary
|
||||
2-3 sentence overall assessment.
|
||||
@@ -1,64 +0,0 @@
|
||||
---
|
||||
name: scout
|
||||
description: Fast codebase recon. Finds relevant files, types, and patterns.
|
||||
tools: read, grep, find, ls, bash
|
||||
skills: opty, qmd
|
||||
model: opencode-go/mimo-v2-pro
|
||||
output: false
|
||||
---
|
||||
|
||||
You are a scout. Quickly investigate a codebase and return structured findings that another agent can use without re-reading everything.
|
||||
|
||||
**YOUR ONLY JOB IS EXPLORATION AND REPORTING. NEVER implement, write implementation plans, write code, or edit codebase files. Stop as soon as you have finished your structured report.**
|
||||
|
||||
**DO ALL WORK YOURSELF using your own tools. NEVER delegate to subagents, NEVER call `subagent(...)`, NEVER invoke `pi` via bash or any other mechanism. You must personally run every search, read every file, and produce the report directly.**
|
||||
|
||||
## Tools
|
||||
- **opty** — semantic/HDC code search via CLI: `opty query "description"` to find functions/types/imports by meaning; output includes file + line number
|
||||
- **qmd** — knowledge base search via CLI: `qmd query $'lex: X\nvec: Y'` to find docs/notes by keyword or vector; `qmd get <path>` / `qmd multi-get <glob>` to retrieve full documents
|
||||
- **grep/find/bash** — for exact patterns, file discovery, or anything the semantic tools miss
|
||||
- **read** — read specific file sections once you know where to look
|
||||
|
||||
## Strategy
|
||||
1. Use `opty query "..."` to semantically locate relevant functions/types (fast, no file reading needed)
|
||||
2. Use `qmd query "..."` to check if there's relevant documentation or prior context in the knowledge base
|
||||
3. grep/find for exact patterns or when semantic search isn't precise enough
|
||||
4. Read key sections (not entire files — target the relevant functions/types)
|
||||
5. For file structure overview, use grep for exports/types or read the top ~50 lines of a file
|
||||
6. Identify types, interfaces, key functions
|
||||
7. Note dependencies between files
|
||||
8. Flag anything surprising or risky
|
||||
|
||||
## Output format
|
||||
|
||||
# Context
|
||||
|
||||
## Files Retrieved
|
||||
List with exact line ranges:
|
||||
1. `path/to/file.ts` (lines 10-50) - Description of what's here
|
||||
2. `path/to/other.ts` (lines 100-150) - Description
|
||||
|
||||
## Key Code
|
||||
Critical types, interfaces, or functions — include actual code snippets:
|
||||
|
||||
**Rule: When an implementation must mirror an existing pattern (e.g. a pipeline, a handler, a system), include the EXACT code of the relevant functions — not summaries or descriptions. The coder needs to copy the pattern, not reconstruct it from prose.**
|
||||
|
||||
```typescript
|
||||
// From path/to/file.ts:10-30
|
||||
interface Example {
|
||||
// actual code
|
||||
}
|
||||
```
|
||||
|
||||
## Architecture
|
||||
Brief explanation of how the pieces connect. What calls what. Data flow.
|
||||
|
||||
## Risks & Gotchas
|
||||
Anything that could trip up an implementer: implicit constraints, shared state, tricky edge cases.
|
||||
|
||||
## Start Here
|
||||
Which file to look at first and why.
|
||||
|
||||
---
|
||||
|
||||
**STOP HERE. Do not write any implementation. Do not suggest next steps beyond "Start Here". Your job is done.**
|
||||
@@ -1,275 +0,0 @@
|
||||
/**
|
||||
* ask-claude — Stream Claude agent reviews into pi.
|
||||
*
|
||||
* For AGENTS to use. Delegates to specialized Claude agents or raw models
|
||||
* for review, analysis, debugging, and second opinions.
|
||||
*
|
||||
* Tool (callable by the LLM):
|
||||
* ask_claude(prompt, agent?, model?, question?, session_id?)
|
||||
* agent — any agent name from ~/.claude/agents/ (e.g. "plan_review", "code_review", "oracle", "debug")
|
||||
* model — model override: "opus", "sonnet", or full model ID
|
||||
* question — specific focus prepended as a review header
|
||||
* session_id — resume a prior conversation (returned in every response)
|
||||
*
|
||||
* Commands:
|
||||
* /review-plan — editor → Claude Opus plan_review → inject result
|
||||
* /review-code — editor → Claude Sonnet code_review → inject result
|
||||
*/
|
||||
|
||||
import { defineTool, getMarkdownTheme } from "@mariozechner/pi-coding-agent";
|
||||
import type { ExtensionAPI } from "@mariozechner/pi-coding-agent";
|
||||
import { Type } from "@mariozechner/pi-ai";
|
||||
import { Box, Container, Markdown, Spacer, Text } from "@mariozechner/pi-tui";
|
||||
import {
|
||||
type ClaudeDetails,
|
||||
type Theme,
|
||||
renderToolCallLine,
|
||||
renderToolResultBox,
|
||||
renderToolBlock,
|
||||
formatUsage,
|
||||
buildEditDiff,
|
||||
formatAnthropicError,
|
||||
runClaude,
|
||||
type RunClaudeResult,
|
||||
} from "../shared/claude-stream.js";
|
||||
|
||||
// =============================================================================
|
||||
// Rendering (ask-claude specific — uses shared tool renderers)
|
||||
// =============================================================================
|
||||
|
||||
function buildLabel(agent?: string, model?: string): string {
|
||||
if (agent) return `Claude [${agent}]`;
|
||||
if (model) return `Claude [${model}]`;
|
||||
return "Claude Sonnet";
|
||||
}
|
||||
|
||||
// =============================================================================
|
||||
// Tool definition
|
||||
// =============================================================================
|
||||
|
||||
const AskClaudeParams = Type.Object({
|
||||
prompt: Type.String({
|
||||
description: "Full content to review/analyze. Include all relevant context: CLAUDE.md conventions, files explored, code or plan to review.",
|
||||
}),
|
||||
agent: Type.Optional(Type.String({
|
||||
description: "Agent name from ~/.claude/agents/ (e.g. 'plan_review', 'code_review', 'oracle', 'debug'). Omit to use model= directly.",
|
||||
})),
|
||||
model: Type.Optional(Type.String({
|
||||
description: "Model override: 'opus', 'sonnet', 'haiku', or a full model ID. When agent is set this overrides the agent's default. When agent is omitted this selects the model directly.",
|
||||
})),
|
||||
question: Type.Optional(Type.String({
|
||||
description: "Specific question or focus area prepended to the prompt (e.g. 'Focus on security', 'Are there race conditions?').",
|
||||
})),
|
||||
session_id: Type.Optional(Type.String({
|
||||
description: "Resume a prior conversation. Pass the session_id returned from a previous ask_claude call.",
|
||||
})),
|
||||
});
|
||||
|
||||
const askClaudeTool = defineTool({
|
||||
name: "ask_claude",
|
||||
label: "Ask Claude",
|
||||
description: [
|
||||
"Ask a Claude agent or model for review, analysis, or a second opinion.",
|
||||
"agent=<name> runs a named agent from ~/.claude/agents/ (e.g. 'plan_review', 'code_review', 'oracle', 'debug').",
|
||||
"Use model= alone for free-form requests without an agent. Use question= to specify a focus.",
|
||||
"Pass session_id from a prior response to continue the same conversation across turns.",
|
||||
"CLAUDE.md and .claude/skills are loaded automatically from the project root.",
|
||||
].join(" "),
|
||||
promptSnippet: "Ask a Claude agent or model for review, analysis, or a second opinion",
|
||||
promptGuidelines: [
|
||||
"Use ask_claude(agent=<name>) to invoke a specialized agent — include all relevant context in the prompt.",
|
||||
"Use ask_claude(model='opus', question='...') for free-form deep analysis.",
|
||||
"Always include the artifact to review (plan, code, problem description) in the prompt.",
|
||||
"Pass session_id back from the previous response to continue the conversation.",
|
||||
],
|
||||
parameters: AskClaudeParams,
|
||||
|
||||
async execute(_toolCallId, params, signal, onUpdate, ctx) {
|
||||
const fullPrompt = params.question
|
||||
? `## Review Focus\n\n${params.question}\n\n## Content\n\n${params.prompt}`
|
||||
: params.prompt;
|
||||
|
||||
const label = buildLabel(params.agent, params.model);
|
||||
const details: ClaudeDetails = { label, done: false, blocks: [], finalText: "", isResume: !!params.session_id };
|
||||
|
||||
try {
|
||||
const result: RunClaudeResult = await runClaude(fullPrompt, {
|
||||
agent: params.agent,
|
||||
model: params.model,
|
||||
sessionId: params.session_id,
|
||||
enrichEditDiffs: true, // ask-claude wants diff enrichment
|
||||
cwd: ctx.cwd,
|
||||
signal,
|
||||
onUpdate: (partial) => {
|
||||
Object.assign(details, partial);
|
||||
onUpdate?.({
|
||||
content: [{ type: "text", text: details.finalText || "(thinking…)" }],
|
||||
details: { ...details },
|
||||
});
|
||||
},
|
||||
});
|
||||
|
||||
details.done = true;
|
||||
details.finalText = result.finalText;
|
||||
details.blocks = result.blocks;
|
||||
details.sessionId = result.sessionId;
|
||||
details.costUsd = result.costUsd;
|
||||
details.inputTokens = result.inputTokens;
|
||||
details.outputTokens = result.outputTokens;
|
||||
details.cacheReadTokens = result.cacheReadTokens;
|
||||
details.cacheWriteTokens = result.cacheWriteTokens;
|
||||
|
||||
const sessionNote = result.sessionId
|
||||
? `\n\n---\n*session_id: \`${result.sessionId}\`*`
|
||||
: "";
|
||||
|
||||
return {
|
||||
content: [{ type: "text", text: (result.finalText || "(no output)") + sessionNote }],
|
||||
details: { ...details },
|
||||
};
|
||||
} catch (err) {
|
||||
const errMsg = err instanceof Error ? err.message : String(err);
|
||||
details.done = true;
|
||||
details.finalText = `Error: ${errMsg}`;
|
||||
return {
|
||||
content: [{ type: "text", text: `**Claude error:** ${errMsg}` }],
|
||||
details: { ...details },
|
||||
};
|
||||
}
|
||||
},
|
||||
|
||||
renderCall(args, theme, _ctx) {
|
||||
const label = buildLabel(args.agent, args.model);
|
||||
let text = theme.fg("toolTitle", theme.bold("ask_claude ")) + theme.fg("accent", `[${label}]`);
|
||||
if (args.question) {
|
||||
text += "\n " + theme.fg("dim", theme.italic(args.question));
|
||||
} else {
|
||||
const lines = args.prompt.split("\n").filter((l) => l.trim()).slice(0, 3);
|
||||
text += "\n " + theme.fg("dim", lines.join("\n "));
|
||||
}
|
||||
return new Text(text, 0, 0);
|
||||
},
|
||||
|
||||
renderResult(result, { isPartial }, theme, _ctx) {
|
||||
const d = result.details as ClaudeDetails | undefined;
|
||||
if (!d) return new Text(theme.fg("muted", "…"), 0, 0);
|
||||
|
||||
const isDone = d.done && !isPartial;
|
||||
const statusIcon = isDone ? theme.fg("success", "✓ ") : theme.fg("warning", "⏳ ");
|
||||
const c = new Container();
|
||||
|
||||
const resume = d.isResume ? theme.fg("dim", " ↩") : "";
|
||||
c.addChild(new Text(statusIcon + theme.fg("toolTitle", theme.bold(d.label)) + resume, 0, 0));
|
||||
|
||||
for (const block of d.blocks ?? []) {
|
||||
if (block.type === "thinking" && block.text.trim()) {
|
||||
c.addChild(new Text(theme.fg("dim", theme.italic(block.text.trimEnd())), 0, 0));
|
||||
} else if (block.type === "tool") {
|
||||
c.addChild(renderToolBlock(block, theme as any));
|
||||
} else if (block.type === "text" && block.text.trim()) {
|
||||
c.addChild(new Spacer(1));
|
||||
if (isDone) {
|
||||
c.addChild(new Markdown(block.text.trim(), 0, 0, getMarkdownTheme()));
|
||||
} else {
|
||||
c.addChild(new Text(theme.fg("dim", block.text.trimEnd()), 0, 0));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if (isDone) {
|
||||
const usageLine = formatUsage(d);
|
||||
const parts: string[] = [];
|
||||
if (usageLine) parts.push(usageLine);
|
||||
if (d.sessionId) parts.push(`session:${d.sessionId.slice(0, 8)}`);
|
||||
if (parts.length > 0) {
|
||||
c.addChild(new Spacer(1));
|
||||
c.addChild(new Text(theme.fg("dim", parts.join(" ")), 0, 0));
|
||||
}
|
||||
}
|
||||
|
||||
return c;
|
||||
},
|
||||
});
|
||||
|
||||
// =============================================================================
|
||||
// Extension entry point
|
||||
// =============================================================================
|
||||
|
||||
export default function (pi: ExtensionAPI) {
|
||||
pi.registerTool(askClaudeTool);
|
||||
|
||||
// ── /review-plan ───────────────────────────────────────────────────────
|
||||
|
||||
pi.registerCommand("review-plan", {
|
||||
description: "Editor → Claude Opus plan_review → inject review into conversation",
|
||||
handler: async (_args, ctx) => {
|
||||
const input = await ctx.ui.editor(
|
||||
"Plan Review · Claude Opus",
|
||||
"Paste your plan or strategy. Claude will review for correctness, completeness, and risk.",
|
||||
);
|
||||
if (!input?.trim()) { ctx.ui.notify("Cancelled.", "info"); return; }
|
||||
ctx.ui.setStatus("ask-claude", "Asking Claude Opus (plan_review)…");
|
||||
try {
|
||||
const r = await runClaude(input, { agent: "plan_review", cwd: ctx.cwd, onUpdate: () => {} });
|
||||
if (!r.finalText.trim()) { ctx.ui.notify("No output from Claude.", "warning"); return; }
|
||||
pi.sendMessage(
|
||||
{ customType: "ask-claude-review", content: r.finalText.trim(), display: true,
|
||||
details: { label: "Claude Opus · plan_review", output: r.finalText } },
|
||||
{ triggerTurn: true },
|
||||
);
|
||||
} catch (err) {
|
||||
ctx.ui.notify(`Claude error: ${err instanceof Error ? err.message : String(err)}`, "error");
|
||||
} finally { ctx.ui.setStatus("ask-claude", undefined); }
|
||||
},
|
||||
});
|
||||
|
||||
// ── /review-code ───────────────────────────────────────────────────────
|
||||
|
||||
pi.registerCommand("review-code", {
|
||||
description: "Editor → Claude Sonnet code_review → inject review into conversation",
|
||||
handler: async (_args, ctx) => {
|
||||
const input = await ctx.ui.editor(
|
||||
"Code Review · Claude Sonnet",
|
||||
"Paste code to review. Include the plan it implements and any specific concerns.",
|
||||
);
|
||||
if (!input?.trim()) { ctx.ui.notify("Cancelled.", "info"); return; }
|
||||
ctx.ui.setStatus("ask-claude", "Asking Claude Sonnet (code_review)…");
|
||||
try {
|
||||
const r = await runClaude(input, { agent: "code_review", cwd: ctx.cwd, onUpdate: () => {} });
|
||||
if (!r.finalText.trim()) { ctx.ui.notify("No output from Claude.", "warning"); return; }
|
||||
pi.sendMessage(
|
||||
{ customType: "ask-claude-review", content: r.finalText.trim(), display: true,
|
||||
details: { label: "Claude Sonnet · code_review", output: r.finalText } },
|
||||
{ triggerTurn: true },
|
||||
);
|
||||
} catch (err) {
|
||||
ctx.ui.notify(`Claude error: ${err instanceof Error ? err.message : String(err)}`, "error");
|
||||
} finally { ctx.ui.setStatus("ask-claude", undefined); }
|
||||
},
|
||||
});
|
||||
|
||||
// ── Message renderer for injected reviews ──────────────────────────────
|
||||
|
||||
pi.registerMessageRenderer("ask-claude-review", (message, { expanded }, theme) => {
|
||||
const d = message.details as { label?: string; output?: string } | undefined;
|
||||
const label = d?.label ?? "Claude";
|
||||
const output = (d?.output ?? "").trim();
|
||||
|
||||
if (expanded) {
|
||||
const c = new Container();
|
||||
c.addChild(new Text(theme.fg("accent", "◆ ") + theme.fg("toolTitle", theme.bold(label)), 0, 0));
|
||||
c.addChild(new Spacer(1));
|
||||
c.addChild(new Markdown(output, 0, 0, getMarkdownTheme()));
|
||||
return c;
|
||||
}
|
||||
|
||||
let text = theme.fg("accent", "◆ ") + theme.fg("toolTitle", theme.bold(label));
|
||||
const lines = output.split("\n").filter((l) => l.trim());
|
||||
const preview = lines.slice(0, 4).join("\n");
|
||||
if (preview) {
|
||||
text += "\n" + theme.fg("dim", preview);
|
||||
if (lines.length > 4) text += "\n" + theme.fg("muted", `… ${lines.length - 4} more (Ctrl+O)`);
|
||||
}
|
||||
return new Text(text, 0, 0);
|
||||
});
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -1,243 +0,0 @@
|
||||
/**
|
||||
* Footer Display Extension
|
||||
*
|
||||
* Replaces the built-in pi footer with a single clean line that assembles
|
||||
* status from all other extensions:
|
||||
*
|
||||
* ~dir | S ⣿⣶⣀⣀⣀ 34% 2h 55m | W ⣿⣿⣷⣀⣀ 68% ⟳ Fri 09:00 | C ⣿⣀⣀⣀⣀ 20% | Sonnet 4.6 | rust-analyzer | MCP: 1/2
|
||||
*
|
||||
* Status sources:
|
||||
* usage:update event — set by usage-bars extension → S/W bars (Claude usage, always shown)
|
||||
* ctx.getContextUsage() → C bar (rendered here)
|
||||
* ctx.model → model short name
|
||||
* "lsp" — set by lsp-pi extension → strip "LSP " prefix
|
||||
* "mcp" — set by pi-mcp-adapter → strip " servers" suffix
|
||||
*/
|
||||
|
||||
import os from "os";
|
||||
import path from "path";
|
||||
import type { ExtensionAPI } from "@mariozechner/pi-coding-agent";
|
||||
import { truncateToWidth } from "@mariozechner/pi-tui";
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Braille gradient bar — used here only for the context (C) bar
|
||||
// ---------------------------------------------------------------------------
|
||||
const BRAILLE_GRADIENT = "\u28C0\u28C4\u28E4\u28E6\u28F6\u28F7\u28FF";
|
||||
const BRAILLE_EMPTY = "\u28C0";
|
||||
const BAR_WIDTH = 5;
|
||||
|
||||
function renderBrailleBar(theme: any, value: number): string {
|
||||
const v = Math.max(0, Math.min(100, Math.round(value)));
|
||||
const levels = BRAILLE_GRADIENT.length - 1;
|
||||
const totalSteps = BAR_WIDTH * levels;
|
||||
const filledSteps = Math.round((v / 100) * totalSteps);
|
||||
const full = Math.floor(filledSteps / levels);
|
||||
const partial = filledSteps % levels;
|
||||
const empty = BAR_WIDTH - full - (partial ? 1 : 0);
|
||||
const color = v >= 90 ? "error" : v >= 70 ? "warning" : "success";
|
||||
const filled = BRAILLE_GRADIENT[BRAILLE_GRADIENT.length - 1]!.repeat(Math.max(0, full));
|
||||
const partialChar = partial ? BRAILLE_GRADIENT[partial]! : "";
|
||||
const emptyChars = BRAILLE_EMPTY.repeat(Math.max(0, empty));
|
||||
return theme.fg(color, filled + partialChar) + theme.fg("dim", emptyChars);
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Helpers
|
||||
// ---------------------------------------------------------------------------
|
||||
function stripAnsi(text: string): string {
|
||||
return text.replace(/\x1b\[[0-9;]*m/g, "").replace(/\x1b\][^\x07]*\x07/g, "");
|
||||
}
|
||||
|
||||
function getModelShortName(modelId: string): string {
|
||||
// claude-haiku-4-5 → "Haiku 4.5", claude-sonnet-4-6 → "Sonnet 4.6"
|
||||
const m = modelId.match(/^claude-([a-z]+)-([\d]+(?:-[\d]+)*)(?:-\d{8})?$/);
|
||||
if (m) {
|
||||
const family = m[1]!.charAt(0).toUpperCase() + m[1]!.slice(1);
|
||||
return `${family} ${m[2]!.replace(/-/g, ".")}`;
|
||||
}
|
||||
// claude-3-5-sonnet, claude-3-opus, etc.
|
||||
const m2 = modelId.match(/^claude-[\d-]+-([a-z]+)/);
|
||||
if (m2) return m2[1]!.charAt(0).toUpperCase() + m2[1]!.slice(1);
|
||||
return modelId.replace(/^claude-/, "");
|
||||
}
|
||||
|
||||
// Format duration in milliseconds to human readable (e.g., "2h 55m")
|
||||
function formatDurationMs(ms: number): string {
|
||||
if (!Number.isFinite(ms) || ms <= 0) return "now";
|
||||
const totalSeconds = Math.floor(ms / 1000);
|
||||
const d = Math.floor(totalSeconds / 86400);
|
||||
const h = Math.floor((totalSeconds % 86400) / 3600);
|
||||
const m = Math.floor((totalSeconds % 3600) / 60);
|
||||
if (d > 0 && h > 0) return `${d}d ${h}h`;
|
||||
if (d > 0) return `${d}d`;
|
||||
if (h > 0 && m > 0) return `${h}h ${m}m`;
|
||||
if (h > 0) return `${h}h`;
|
||||
if (m > 0) return `${m}m`;
|
||||
return "<1m";
|
||||
}
|
||||
|
||||
export default function (pi: ExtensionAPI) {
|
||||
let ctx: any = null;
|
||||
let tuiRef: any = null;
|
||||
let footerDataRef: any = null;
|
||||
|
||||
// Track usage data for dynamic S/W bar rendering
|
||||
let usageSession: number | null = null;
|
||||
let usageWeekly: number | null = null;
|
||||
let sessionResetsAt: number | null = null;
|
||||
let weeklyResetsAt: number | null = null;
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Footer line builder — called on every render
|
||||
// ---------------------------------------------------------------------------
|
||||
function buildFooterLine(theme: any): string {
|
||||
const sep = theme.fg("dim", " · ");
|
||||
const pipeSep = theme.fg("dim", " | ");
|
||||
const parts: string[] = [];
|
||||
|
||||
const statuses: ReadonlyMap<string, string> =
|
||||
footerDataRef?.getExtensionStatuses?.() ?? new Map();
|
||||
|
||||
// 1. Current working directory
|
||||
const home = os.homedir();
|
||||
const cwd = process.cwd();
|
||||
const dir = cwd.startsWith(home)
|
||||
? "~" + path.sep + path.relative(home, cwd)
|
||||
: cwd;
|
||||
parts.push(theme.fg("dim", dir));
|
||||
|
||||
// 2. S / W usage bars + C bar — joined as one |-separated block
|
||||
const usageRaw = statuses.get("usage-bars");
|
||||
const contextUsage = ctx?.getContextUsage?.();
|
||||
{
|
||||
let block: string;
|
||||
|
||||
if (usageSession !== null && usageWeekly !== null) {
|
||||
// Build S/W bars directly from stored event data so we can cleanly
|
||||
// append the dynamic countdown without trying to parse ANSI strings.
|
||||
const session = Math.max(0, Math.min(100, Math.round(usageSession)));
|
||||
const weekly = Math.max(0, Math.min(100, Math.round(usageWeekly)));
|
||||
|
||||
let sPart = theme.fg("muted", "\uF4F5 S ") + renderBrailleBar(theme, session) + " " + theme.fg("dim", `${session}%`);
|
||||
let wPart = theme.fg("muted", "\uF4F5 W ") + renderBrailleBar(theme, weekly) + " " + theme.fg("dim", `${weekly}%`);
|
||||
|
||||
if (sessionResetsAt !== null) {
|
||||
const msLeft = sessionResetsAt - Date.now();
|
||||
if (msLeft > 0) sPart += " " + theme.fg("dim", formatDurationMs(msLeft));
|
||||
}
|
||||
|
||||
if (weeklyResetsAt !== null) {
|
||||
const msLeft = weeklyResetsAt - Date.now();
|
||||
if (msLeft > 0) wPart += " " + theme.fg("dim", `\u27F3 ${formatDurationMs(msLeft)}`);
|
||||
}
|
||||
|
||||
block = sPart + pipeSep + wPart;
|
||||
} else {
|
||||
// Fallback to raw status for loading / error states
|
||||
block = usageRaw ?? "";
|
||||
}
|
||||
|
||||
if (contextUsage && contextUsage.percent !== null) {
|
||||
const pct = Math.round(contextUsage.percent);
|
||||
const chatStatus = statuses.get("chat-claude");
|
||||
const isChatActive = !!chatStatus && chatStatus.includes("Claude");
|
||||
|
||||
// When chat is active and context is high, show warning indicator
|
||||
let cLabel = "C";
|
||||
let cColor = "muted";
|
||||
if (isChatActive && pct >= 70) {
|
||||
cLabel = pct >= 90 ? "C⚠" : "C⚡";
|
||||
cColor = pct >= 90 ? "error" : "warning";
|
||||
}
|
||||
|
||||
const cBar =
|
||||
theme.fg(cColor, cLabel + " ") +
|
||||
renderBrailleBar(theme, pct) +
|
||||
" " +
|
||||
theme.fg(pct >= 70 && isChatActive ? "warning" : "dim", `${pct}%`);
|
||||
block = block ? block + pipeSep + cBar : cBar;
|
||||
}
|
||||
if (block) parts.push(block);
|
||||
}
|
||||
|
||||
// 3. Model short name
|
||||
const modelId = ctx?.model?.id;
|
||||
if (modelId) parts.push(theme.fg("dim", getModelShortName(modelId)));
|
||||
|
||||
// 4. LSP — strip "LSP" prefix and activity dot
|
||||
const lspRaw = statuses.get("lsp");
|
||||
if (lspRaw) {
|
||||
const clean = stripAnsi(lspRaw).trim().replace(/^LSP\s*[•·]?\s*/i, "").trim();
|
||||
if (clean) parts.push(theme.fg("dim", clean));
|
||||
}
|
||||
|
||||
// 5. MCP — strip " servers" suffix
|
||||
const mcpRaw = statuses.get("mcp");
|
||||
if (mcpRaw) {
|
||||
const clean = stripAnsi(mcpRaw).trim().replace(/\s*servers?.*$/, "").trim();
|
||||
if (clean) parts.push(theme.fg("dim", clean));
|
||||
}
|
||||
|
||||
// 6. Active Claude chat session
|
||||
const chatRaw = statuses.get("chat-claude");
|
||||
if (chatRaw) {
|
||||
parts.push(theme.fg("accent", stripAnsi(chatRaw).trim()));
|
||||
}
|
||||
|
||||
return parts.join(sep);
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Footer installation
|
||||
// ---------------------------------------------------------------------------
|
||||
function installFooter(_ctx: any) {
|
||||
if (!_ctx?.hasUI) return;
|
||||
_ctx.ui.setFooter((_tui: any, theme: any, footerData: any) => {
|
||||
tuiRef = _tui;
|
||||
footerDataRef = footerData;
|
||||
const unsub = footerData.onBranchChange(() => _tui.requestRender());
|
||||
return {
|
||||
dispose: unsub,
|
||||
invalidate() {},
|
||||
render(width: number): string[] {
|
||||
return [truncateToWidth(buildFooterLine(theme) || "", width)];
|
||||
},
|
||||
};
|
||||
});
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Event handlers
|
||||
// ---------------------------------------------------------------------------
|
||||
pi.on("session_start", (_event, _ctx) => {
|
||||
ctx = _ctx;
|
||||
installFooter(_ctx);
|
||||
});
|
||||
|
||||
pi.on("session_shutdown", (_event, _ctx) => {
|
||||
if (_ctx?.hasUI) _ctx.ui.setFooter(undefined);
|
||||
});
|
||||
|
||||
// Re-render after turns so context usage stays current
|
||||
pi.on("turn_end", (_event, _ctx) => {
|
||||
ctx = _ctx;
|
||||
if (tuiRef) tuiRef.requestRender();
|
||||
});
|
||||
|
||||
// Re-render when model changes (updates model name in footer)
|
||||
pi.on("model_select", (_event, _ctx) => {
|
||||
ctx = _ctx;
|
||||
if (tuiRef) tuiRef.requestRender();
|
||||
});
|
||||
|
||||
|
||||
// Listen for usage updates — store raw values so we can build bars + dynamic
|
||||
// countdown directly rather than parsing the ANSI status string from usage-bars.
|
||||
pi.events.on("usage:update", (data: any) => {
|
||||
if (data.session !== undefined) usageSession = data.session;
|
||||
if (data.weekly !== undefined) usageWeekly = data.weekly;
|
||||
if (data.sessionResetsAt !== undefined) sessionResetsAt = data.sessionResetsAt;
|
||||
if (data.weeklyResetsAt !== undefined) weeklyResetsAt = data.weeklyResetsAt;
|
||||
if (tuiRef) tuiRef.requestRender();
|
||||
});
|
||||
}
|
||||
80
pi/.pi/agent/extensions/lib/boxes.ts
Normal file
80
pi/.pi/agent/extensions/lib/boxes.ts
Normal file
@@ -0,0 +1,80 @@
|
||||
import { visibleWidth, truncateToWidth } from "@earendil-works/pi-tui";
|
||||
|
||||
/**
|
||||
* Shared rounded-box border style for pi extensions.
|
||||
*
|
||||
* Used by tool-blocks.ts (tool output boxes), prompt-frame.ts (editor +
|
||||
* footer box) and transcript-viewer.ts. Labels "cut into" the border like:
|
||||
*
|
||||
* ╭─ bash · cargo test ────────────╮
|
||||
* │ output line │
|
||||
* ╰───────────────── ok · 0.8s ─╯
|
||||
*
|
||||
* This file is NOT an extension (no default export). The loader only picks
|
||||
* up extensions/*.ts and extensions/* /index.ts, so lib/ is skipped.
|
||||
*/
|
||||
|
||||
export const H = "─";
|
||||
|
||||
/** Border color for a tool/block status. */
|
||||
export function borderColor(theme: any, status: string) {
|
||||
const token = status === "running" ? "accent" : status === "error" ? "error" : "success";
|
||||
return (s: string) => theme.fg(token, s);
|
||||
}
|
||||
|
||||
/** Make text safe for width math: no tabs, no carriage returns. */
|
||||
export function sanitize(text: string) {
|
||||
return text.replace(/\r/g, "").replace(/\t/g, " ");
|
||||
}
|
||||
|
||||
export function formatDuration(ms: number) {
|
||||
if (ms < 1000) return `${ms}ms`;
|
||||
if (ms < 60_000) return `${(ms / 1000).toFixed(1)}s`;
|
||||
return `${Math.floor(ms / 60_000)}m${Math.round((ms % 60_000) / 1000)}s`;
|
||||
}
|
||||
|
||||
/**
|
||||
* A border line with a left and a right label cut into it:
|
||||
* `╭─ left ────── right ─╮` (corners configurable)
|
||||
* Labels keep their own colors; the rule takes `color`.
|
||||
* The right label is truncated first, then the left one (fitBorder-style).
|
||||
*/
|
||||
export function splitBorder(
|
||||
width: number,
|
||||
left: string,
|
||||
right: string,
|
||||
color: (s: string) => string,
|
||||
corners: [string, string] = ["╭", "╮"],
|
||||
) {
|
||||
const [lc, rc] = corners;
|
||||
if (width < 8) return color(lc + H.repeat(Math.max(0, width - 2)) + rc);
|
||||
|
||||
let l = left ? ` ${left} ` : "";
|
||||
let r = right ? ` ${right} ` : "";
|
||||
const fixed = 4; // corners + one rule char each side
|
||||
while (fixed + visibleWidth(l) + visibleWidth(r) > width && visibleWidth(r) > 0) {
|
||||
r = truncateToWidth(r, Math.max(0, visibleWidth(r) - 2), "…");
|
||||
}
|
||||
while (fixed + visibleWidth(l) + visibleWidth(r) > width && visibleWidth(l) > 0) {
|
||||
l = truncateToWidth(l, Math.max(0, visibleWidth(l) - 2), "…");
|
||||
}
|
||||
const gap = Math.max(0, width - fixed - visibleWidth(l) - visibleWidth(r));
|
||||
return color(`${lc}${H}`) + l + color(H.repeat(gap)) + r + color(`${H}${rc}`);
|
||||
}
|
||||
|
||||
/** `╭─ title ────╮` */
|
||||
export function topBorder(width: number, title: string, color: (s: string) => string) {
|
||||
return splitBorder(width, title, "", color, ["╭", "╮"]);
|
||||
}
|
||||
|
||||
/** `╰──── summary ─╯` (summary right-aligned) */
|
||||
export function bottomBorder(width: number, summary: string, color: (s: string) => string) {
|
||||
return splitBorder(width, "", summary, color, ["╰", "╯"]);
|
||||
}
|
||||
|
||||
/** `│ content ... │` — content must already be ≤ width-4 columns. */
|
||||
export function frameLine(width: number, content: string, color: (s: string) => string) {
|
||||
const inner = Math.max(1, width - 4);
|
||||
const pad = Math.max(0, inner - visibleWidth(content));
|
||||
return color("│ ") + content + " ".repeat(pad) + color(" │");
|
||||
}
|
||||
@@ -1,56 +0,0 @@
|
||||
# pi-ask-mcp
|
||||
|
||||
A minimal MCP stdio server that gives Claude **one** tool — `ask` — which routes
|
||||
structured questions back to pi's native ask UI instead of using Claude's
|
||||
built-in `AskUserQuestion`.
|
||||
|
||||
This is **not** a regular pi extension. It is a subprocess of `claude`, which is
|
||||
itself a subprocess of the `chat-claude` extension. The pi-side counterpart is
|
||||
[`shared/pi-ask-bridge.ts`](../../shared/pi-ask-bridge.ts), which:
|
||||
|
||||
1. Opens a Unix-domain socket per chat session.
|
||||
2. Generates an `--mcp-config` JSON pointing here, with `PI_ASK_SOCKET=<sock>`.
|
||||
3. Translates `ask` requests off the socket into
|
||||
`askSingleQuestionWithInlineNote` / `askQuestionsWithTabs` calls and writes
|
||||
the result back.
|
||||
|
||||
## Architecture
|
||||
|
||||
```
|
||||
pi
|
||||
└── chat-claude
|
||||
├── pi-ask-bridge (UDS server, owns ui.custom)
|
||||
└── claude -p ... --mcp-config <generated.json> --disallowed-tools AskUserQuestion
|
||||
└── pi-ask-mcp/server.js (this file)
|
||||
↳ on tools/call ask → connect $PI_ASK_SOCKET → ask → reply
|
||||
```
|
||||
|
||||
## Why a hand-written MCP server
|
||||
|
||||
No `@modelcontextprotocol/sdk` dependency, no transpile step, no
|
||||
`node_modules`. The MCP stdio protocol is small enough (~6 method handlers)
|
||||
that writing it directly keeps the file self-contained and trivially
|
||||
portable. Claude CLI spawns it via `node server.js`.
|
||||
|
||||
## Wire format
|
||||
|
||||
Stdio (with Claude): JSON-RPC 2.0 over newline-delimited JSON.
|
||||
|
||||
Socket (with pi-ask-bridge): NDJSON, one request → one response, then close.
|
||||
|
||||
```jsonc
|
||||
// → pi
|
||||
{ "id": "uuid", "type": "ask",
|
||||
"questions": [
|
||||
{ "id": "auth", "question": "Auth method?",
|
||||
"options": [{"label": "OAuth"}, {"label": "API key"}],
|
||||
"multi": false, "recommended": 0 }
|
||||
] }
|
||||
|
||||
// ← pi (success)
|
||||
{ "id": "uuid", "type": "result",
|
||||
"results": [{ "id": "auth", "selectedOptions": ["OAuth"] }] }
|
||||
|
||||
// ← pi (cancel / error)
|
||||
{ "id": "uuid", "type": "error", "message": "cancelled" }
|
||||
```
|
||||
@@ -1,7 +0,0 @@
|
||||
{
|
||||
"name": "pi-ask-mcp",
|
||||
"private": true,
|
||||
"type": "module",
|
||||
"main": "server.js",
|
||||
"description": "Minimal MCP stdio server bridging Claude → pi-ask-bridge."
|
||||
}
|
||||
@@ -1,195 +0,0 @@
|
||||
#!/usr/bin/env node
|
||||
// pi-ask-mcp/server.js
|
||||
//
|
||||
// Minimal MCP stdio server that exposes ONE tool: `ask`.
|
||||
// Bridges Claude → pi via a Unix-domain socket: when Claude calls the tool,
|
||||
// this server forwards the question(s) to pi over $PI_ASK_SOCKET, awaits
|
||||
// the user's answer, and returns it as the tool result.
|
||||
//
|
||||
// Wire format with Claude (stdin/stdout): JSON-RPC 2.0 over NDJSON.
|
||||
// Wire format with pi (PI_ASK_SOCKET): NDJSON request/response, see
|
||||
// ../../shared/pi-ask-bridge.ts.
|
||||
//
|
||||
// This file is INTENTIONALLY plain JavaScript (no transpile step, no
|
||||
// node_modules) — Claude CLI spawns it via `node <path>`. Keep it small,
|
||||
// dependency-free, and self-contained.
|
||||
|
||||
import { connect } from "node:net";
|
||||
import { randomUUID } from "node:crypto";
|
||||
import { createInterface } from "node:readline";
|
||||
|
||||
// ── Configuration ──────────────────────────────────────────────────────────
|
||||
const SOCKET = process.env.PI_ASK_SOCKET;
|
||||
if (!SOCKET) {
|
||||
process.stderr.write("[pi-ask-mcp] PI_ASK_SOCKET env var is required\n");
|
||||
process.exit(2);
|
||||
}
|
||||
|
||||
const SERVER_INFO = { name: "pi", version: "0.1.0" };
|
||||
const PROTOCOL_VERSION = "2024-11-05";
|
||||
const SOCKET_TIMEOUT_MS = 15 * 60 * 1000; // matches runClaude's default
|
||||
|
||||
// ── Tool schema (kept in sync with pi-ask-tool/index.ts AskParamsSchema) ──
|
||||
const ASK_INPUT_SCHEMA = {
|
||||
type: "object",
|
||||
required: ["questions"],
|
||||
properties: {
|
||||
questions: {
|
||||
type: "array",
|
||||
minItems: 1,
|
||||
description: "One or more questions to ask the user.",
|
||||
items: {
|
||||
type: "object",
|
||||
required: ["id", "question", "options"],
|
||||
properties: {
|
||||
id: { type: "string", description: "Stable id (e.g. 'auth', 'cache')." },
|
||||
question: { type: "string", description: "Question text shown to the user." },
|
||||
options: {
|
||||
type: "array",
|
||||
minItems: 1,
|
||||
description: "2-5 concise options. Do NOT include 'Other' (UI adds it).",
|
||||
items: {
|
||||
type: "object",
|
||||
required: ["label"],
|
||||
properties: {
|
||||
label: { type: "string", description: "Option display label." },
|
||||
},
|
||||
},
|
||||
},
|
||||
multi: { type: "boolean", description: "Allow multi-select. Defaults to false." },
|
||||
recommended: { type: "number", description: "0-indexed recommended option (default highlight)." },
|
||||
},
|
||||
},
|
||||
},
|
||||
},
|
||||
};
|
||||
|
||||
const ASK_DESCRIPTION = [
|
||||
"Ask the user one or more structured questions through pi's native TUI.",
|
||||
"Use this whenever a choice materially affects the outcome — instead of",
|
||||
"guessing or the built-in AskUserQuestion. Provide 2-5 concise options.",
|
||||
"Set multi=true when multiple answers are valid. Do NOT include an 'Other'",
|
||||
"option (UI adds it automatically). The result is a JSON array of",
|
||||
"{id, selectedOptions[], customInput?} per question — empty selectedOptions",
|
||||
"means the user cancelled.",
|
||||
].join(" ");
|
||||
|
||||
// ── stdio framing: NDJSON ──────────────────────────────────────────────────
|
||||
const rl = createInterface({ input: process.stdin });
|
||||
const send = (msg) => process.stdout.write(JSON.stringify(msg) + "\n");
|
||||
const log = (msg) => process.stderr.write(`[pi-ask-mcp] ${msg}\n`);
|
||||
|
||||
// ── socket round-trip to pi-ask-bridge ─────────────────────────────────────
|
||||
function askPi(args) {
|
||||
return new Promise((resolve, reject) => {
|
||||
const sock = connect(SOCKET);
|
||||
const id = randomUUID();
|
||||
let buf = "";
|
||||
let settled = false;
|
||||
const finish = (fn, val) => { if (settled) return; settled = true; clearTimeout(t); fn(val); try { sock.end(); } catch {} };
|
||||
const t = setTimeout(
|
||||
() => finish(reject, new Error(`pi-ask bridge timeout after ${SOCKET_TIMEOUT_MS / 1000}s`)),
|
||||
SOCKET_TIMEOUT_MS,
|
||||
);
|
||||
|
||||
sock.on("connect", () => sock.write(JSON.stringify({ id, type: "ask", ...args }) + "\n"));
|
||||
sock.on("data", (d) => {
|
||||
buf += d.toString();
|
||||
const nl = buf.indexOf("\n");
|
||||
if (nl < 0) return;
|
||||
try { finish(resolve, JSON.parse(buf.slice(0, nl))); }
|
||||
catch (err) { finish(reject, err); }
|
||||
});
|
||||
sock.on("error", (err) => finish(reject, err));
|
||||
sock.on("close", () => {
|
||||
if (!settled) finish(reject, new Error("pi-ask bridge closed connection without reply"));
|
||||
});
|
||||
});
|
||||
}
|
||||
|
||||
// ── JSON-RPC method handlers ───────────────────────────────────────────────
|
||||
async function handleRequest(req) {
|
||||
const { id, method, params } = req;
|
||||
try {
|
||||
switch (method) {
|
||||
case "initialize":
|
||||
return ok(id, {
|
||||
protocolVersion: PROTOCOL_VERSION,
|
||||
capabilities: { tools: {} },
|
||||
serverInfo: SERVER_INFO,
|
||||
});
|
||||
case "tools/list":
|
||||
return ok(id, {
|
||||
tools: [{ name: "ask", description: ASK_DESCRIPTION, inputSchema: ASK_INPUT_SCHEMA }],
|
||||
});
|
||||
case "tools/call": {
|
||||
const name = params?.name;
|
||||
const args = params?.arguments ?? {};
|
||||
if (name !== "ask") return err(id, -32602, `unknown tool: ${name}`);
|
||||
const reply = await askPi(args);
|
||||
if (reply.type === "error") {
|
||||
return ok(id, {
|
||||
isError: true,
|
||||
content: [{ type: "text", text: `(user did not answer: ${reply.message})` }],
|
||||
});
|
||||
}
|
||||
return ok(id, {
|
||||
content: [{ type: "text", text: JSON.stringify(reply.results, null, 2) }],
|
||||
});
|
||||
}
|
||||
case "ping": return ok(id, {});
|
||||
case "resources/list": return ok(id, { resources: [] });
|
||||
case "prompts/list": return ok(id, { prompts: [] });
|
||||
default: return err(id, -32601, `method not found: ${method}`);
|
||||
}
|
||||
} catch (e) {
|
||||
return err(id, -32603, e instanceof Error ? e.message : String(e));
|
||||
}
|
||||
}
|
||||
|
||||
const ok = (id, result) => ({ jsonrpc: "2.0", id, result });
|
||||
const err = (id, code, message) => ({ jsonrpc: "2.0", id, error: { code, message } });
|
||||
|
||||
// ── main loop ──────────────────────────────────────────────────────────────
|
||||
//
|
||||
// Track in-flight handlers so we don't exit before they finish. Without this,
|
||||
// `node server.js <<<input` (or any case where stdin closes mid-request) would
|
||||
// race the async tools/call handler against rl 'close' → process.exit, and
|
||||
// the reply would silently disappear.
|
||||
let inflight = 0;
|
||||
let stdinClosed = false;
|
||||
function drainAndExit(code = 0) {
|
||||
if (inflight === 0) process.exit(code);
|
||||
}
|
||||
|
||||
async function handleOne(msg) {
|
||||
// Notifications carry no id and expect no response.
|
||||
if (msg.id === undefined || msg.id === null) {
|
||||
if (msg.method === "exit") drainAndExit(0);
|
||||
return; // notifications/initialized, notifications/cancelled, etc.
|
||||
}
|
||||
inflight += 1;
|
||||
try {
|
||||
const reply = await handleRequest(msg);
|
||||
if (reply) send(reply);
|
||||
} finally {
|
||||
inflight -= 1;
|
||||
if (stdinClosed) drainAndExit(0);
|
||||
}
|
||||
}
|
||||
|
||||
rl.on("line", (line) => {
|
||||
if (!line.trim()) return;
|
||||
let msg;
|
||||
try { msg = JSON.parse(line); } catch { return; }
|
||||
if (Array.isArray(msg)) {
|
||||
for (const m of msg) void handleOne(m);
|
||||
} else {
|
||||
void handleOne(msg);
|
||||
}
|
||||
});
|
||||
|
||||
rl.on("close", () => { stdinClosed = true; drainAndExit(0); });
|
||||
process.on("SIGTERM", () => process.exit(0));
|
||||
process.on("SIGINT", () => process.exit(0));
|
||||
log("ready, socket=" + SOCKET);
|
||||
@@ -1,23 +0,0 @@
|
||||
# Pi Ask Tool Extension
|
||||
|
||||
This extension bridges Claude Code's ask functionality into pi's TUI, allowing users to ask questions and receive answers directly in the TUI interface.
|
||||
|
||||
## Features
|
||||
|
||||
- Seamless integration with pi's TUI
|
||||
- Support for all Claude agents (plan_review, code_review, debug, oracle)
|
||||
- Multi-turn conversations with session management
|
||||
- Context-aware responses based on codebase exploration
|
||||
|
||||
## Usage
|
||||
|
||||
1. Run the CLI agent: `pi-ask-tool`
|
||||
2. Type your question in the TUI
|
||||
3. Receive answers directly in the TUI interface
|
||||
4. Continue the conversation or start new ones
|
||||
|
||||
## Configuration
|
||||
|
||||
Default agent: `code_review`
|
||||
Default model: `sonnet`
|
||||
Session persistence: Enabled
|
||||
@@ -1,65 +0,0 @@
|
||||
import { wrapTextWithAnsi } from "@mariozechner/pi-tui";
|
||||
|
||||
const INLINE_NOTE_SEPARATOR = " — note: ";
|
||||
const INLINE_EDIT_CURSOR = "▍";
|
||||
|
||||
export const INLINE_NOTE_WRAP_PADDING = 2;
|
||||
|
||||
function sanitizeNoteForInlineDisplay(rawNote: string): string {
|
||||
return rawNote.replace(/[\r\n\t]/g, " ").replace(/[\x00-\x08\x0B\x0C\x0E-\x1F\x7F]/g, "");
|
||||
}
|
||||
|
||||
function truncateTextKeepingTail(text: string, maxLength: number): string {
|
||||
if (maxLength <= 0) return "";
|
||||
if (text.length <= maxLength) return text;
|
||||
if (maxLength === 1) return "…";
|
||||
return `…${text.slice(-(maxLength - 1))}`;
|
||||
}
|
||||
|
||||
function truncateTextKeepingHead(text: string, maxLength: number): string {
|
||||
if (maxLength <= 0) return "";
|
||||
if (text.length <= maxLength) return text;
|
||||
if (maxLength === 1) return "…";
|
||||
return `${text.slice(0, maxLength - 1)}…`;
|
||||
}
|
||||
|
||||
export function buildOptionLabelWithInlineNote(
|
||||
baseOptionLabel: string,
|
||||
rawNote: string,
|
||||
isEditingNote: boolean,
|
||||
maxInlineLabelLength?: number,
|
||||
): string {
|
||||
const sanitizedNote = sanitizeNoteForInlineDisplay(rawNote);
|
||||
if (!isEditingNote && sanitizedNote.trim().length === 0) {
|
||||
return baseOptionLabel;
|
||||
}
|
||||
|
||||
const labelPrefix = `${baseOptionLabel}${INLINE_NOTE_SEPARATOR}`;
|
||||
const inlineNote = isEditingNote ? `${sanitizedNote}${INLINE_EDIT_CURSOR}` : sanitizedNote.trim();
|
||||
const inlineLabel = `${labelPrefix}${inlineNote}`;
|
||||
|
||||
if (maxInlineLabelLength == null) {
|
||||
return inlineLabel;
|
||||
}
|
||||
|
||||
return isEditingNote
|
||||
? truncateTextKeepingTail(inlineLabel, maxInlineLabelLength)
|
||||
: truncateTextKeepingHead(inlineLabel, maxInlineLabelLength);
|
||||
}
|
||||
|
||||
export function buildWrappedOptionLabelWithInlineNote(
|
||||
baseOptionLabel: string,
|
||||
rawNote: string,
|
||||
isEditingNote: boolean,
|
||||
maxInlineLabelLength: number,
|
||||
wrapPadding = INLINE_NOTE_WRAP_PADDING,
|
||||
): string[] {
|
||||
const inlineLabel = buildOptionLabelWithInlineNote(baseOptionLabel, rawNote, isEditingNote);
|
||||
const sanitizedWrapPadding = Number.isFinite(wrapPadding) ? Math.max(0, Math.floor(wrapPadding)) : 0;
|
||||
const sanitizedMaxInlineLabelLength = Number.isFinite(maxInlineLabelLength)
|
||||
? Math.max(1, Math.floor(maxInlineLabelLength))
|
||||
: 1;
|
||||
const wrapWidth = Math.max(1, sanitizedMaxInlineLabelLength - sanitizedWrapPadding);
|
||||
const wrappedLines = wrapTextWithAnsi(inlineLabel, wrapWidth);
|
||||
return wrappedLines.length > 0 ? wrappedLines : [""];
|
||||
}
|
||||
@@ -1,223 +0,0 @@
|
||||
import type { ExtensionUIContext } from "@mariozechner/pi-coding-agent";
|
||||
import { Editor, type EditorTheme, Key, matchesKey, truncateToWidth, visibleWidth, wrapTextWithAnsi } from "@mariozechner/pi-tui";
|
||||
import {
|
||||
OTHER_OPTION,
|
||||
appendRecommendedTagToOptionLabels,
|
||||
buildSingleSelectionResult,
|
||||
type AskOption,
|
||||
type AskSelection,
|
||||
} from "./ask-logic";
|
||||
import { INLINE_NOTE_WRAP_PADDING, buildWrappedOptionLabelWithInlineNote } from "./ask-inline-note";
|
||||
|
||||
interface SingleQuestionInput {
|
||||
question: string;
|
||||
options: AskOption[];
|
||||
recommended?: number;
|
||||
}
|
||||
|
||||
interface InlineSelectionResult {
|
||||
cancelled: boolean;
|
||||
selectedOption?: string;
|
||||
note?: string;
|
||||
}
|
||||
|
||||
function resolveInitialCursorIndexFromRecommendedOption(
|
||||
recommendedOptionIndex: number | undefined,
|
||||
optionCount: number,
|
||||
): number {
|
||||
if (recommendedOptionIndex == null) return 0;
|
||||
if (recommendedOptionIndex < 0 || recommendedOptionIndex >= optionCount) return 0;
|
||||
return recommendedOptionIndex;
|
||||
}
|
||||
|
||||
export async function askSingleQuestionWithInlineNote(
|
||||
ui: ExtensionUIContext,
|
||||
questionInput: SingleQuestionInput,
|
||||
): Promise<AskSelection> {
|
||||
const baseOptionLabels = questionInput.options.map((option) => option.label);
|
||||
const optionLabelsWithRecommendedTag = appendRecommendedTagToOptionLabels(
|
||||
baseOptionLabels,
|
||||
questionInput.recommended,
|
||||
);
|
||||
const selectableOptionLabels = [...optionLabelsWithRecommendedTag, OTHER_OPTION];
|
||||
const initialCursorIndex = resolveInitialCursorIndexFromRecommendedOption(
|
||||
questionInput.recommended,
|
||||
optionLabelsWithRecommendedTag.length,
|
||||
);
|
||||
|
||||
const result = await ui.custom<InlineSelectionResult>((tui, theme, _keybindings, done) => {
|
||||
let cursorOptionIndex = initialCursorIndex;
|
||||
let isNoteEditorOpen = false;
|
||||
let cachedRenderedLines: string[] | undefined;
|
||||
const noteByOptionIndex = new Map<number, string>();
|
||||
|
||||
const editorTheme: EditorTheme = {
|
||||
borderColor: (text) => theme.fg("accent", text),
|
||||
selectList: {
|
||||
selectedPrefix: (text) => theme.fg("accent", text),
|
||||
selectedText: (text) => theme.fg("accent", text),
|
||||
description: (text) => theme.fg("muted", text),
|
||||
scrollInfo: (text) => theme.fg("dim", text),
|
||||
noMatch: (text) => theme.fg("warning", text),
|
||||
},
|
||||
};
|
||||
const noteEditor = new Editor(tui, editorTheme);
|
||||
|
||||
const requestUiRerender = () => {
|
||||
cachedRenderedLines = undefined;
|
||||
tui.requestRender();
|
||||
};
|
||||
|
||||
const getRawNoteForOption = (optionIndex: number): string => noteByOptionIndex.get(optionIndex) ?? "";
|
||||
const getTrimmedNoteForOption = (optionIndex: number): string => getRawNoteForOption(optionIndex).trim();
|
||||
|
||||
const loadCurrentNoteIntoEditor = () => {
|
||||
noteEditor.setText(getRawNoteForOption(cursorOptionIndex));
|
||||
};
|
||||
|
||||
const saveCurrentNoteFromEditor = (value: string) => {
|
||||
noteByOptionIndex.set(cursorOptionIndex, value);
|
||||
};
|
||||
|
||||
const submitCurrentSelection = (selectedOptionLabel: string, note: string) => {
|
||||
done({
|
||||
cancelled: false,
|
||||
selectedOption: selectedOptionLabel,
|
||||
note,
|
||||
});
|
||||
};
|
||||
|
||||
noteEditor.onChange = (value) => {
|
||||
saveCurrentNoteFromEditor(value);
|
||||
requestUiRerender();
|
||||
};
|
||||
|
||||
noteEditor.onSubmit = (value) => {
|
||||
saveCurrentNoteFromEditor(value);
|
||||
const selectedOptionLabel = selectableOptionLabels[cursorOptionIndex];
|
||||
const trimmedNote = value.trim();
|
||||
|
||||
if (selectedOptionLabel === OTHER_OPTION && !trimmedNote) {
|
||||
requestUiRerender();
|
||||
return;
|
||||
}
|
||||
|
||||
submitCurrentSelection(selectedOptionLabel, trimmedNote);
|
||||
};
|
||||
|
||||
const render = (width: number): string[] => {
|
||||
if (cachedRenderedLines) return cachedRenderedLines;
|
||||
|
||||
const renderedLines: string[] = [];
|
||||
const addLine = (line: string) => renderedLines.push(truncateToWidth(line, width));
|
||||
|
||||
addLine(theme.fg("accent", "─".repeat(width)));
|
||||
for (const questionLine of wrapTextWithAnsi(questionInput.question, Math.max(1, width - 1))) {
|
||||
addLine(` ${theme.fg("text", questionLine)}`);
|
||||
}
|
||||
renderedLines.push("");
|
||||
|
||||
for (let optionIndex = 0; optionIndex < selectableOptionLabels.length; optionIndex++) {
|
||||
const optionLabel = selectableOptionLabels[optionIndex];
|
||||
const isCursorOption = optionIndex === cursorOptionIndex;
|
||||
const isEditingThisOption = isNoteEditorOpen && isCursorOption;
|
||||
const cursorPrefixText = isCursorOption ? "→ " : " ";
|
||||
const cursorPrefix = isCursorOption ? theme.fg("accent", cursorPrefixText) : cursorPrefixText;
|
||||
const bullet = isCursorOption ? "●" : "○";
|
||||
const markerText = `${bullet} `;
|
||||
const optionColor = isCursorOption ? "accent" : "text";
|
||||
const prefixWidth = visibleWidth(cursorPrefixText) + visibleWidth(markerText);
|
||||
const wrappedInlineLabelLines = buildWrappedOptionLabelWithInlineNote(
|
||||
optionLabel,
|
||||
getRawNoteForOption(optionIndex),
|
||||
isEditingThisOption,
|
||||
Math.max(1, width - prefixWidth),
|
||||
INLINE_NOTE_WRAP_PADDING,
|
||||
);
|
||||
const continuationPrefix = " ".repeat(prefixWidth);
|
||||
addLine(`${cursorPrefix}${theme.fg(optionColor, `${markerText}${wrappedInlineLabelLines[0] ?? ""}`)}`);
|
||||
for (const wrappedLine of wrappedInlineLabelLines.slice(1)) {
|
||||
addLine(`${continuationPrefix}${theme.fg(optionColor, wrappedLine)}`);
|
||||
}
|
||||
}
|
||||
|
||||
renderedLines.push("");
|
||||
|
||||
if (isNoteEditorOpen) {
|
||||
addLine(theme.fg("dim", " Typing note inline • Enter submit • Tab/Esc stop editing"));
|
||||
} else if (getTrimmedNoteForOption(cursorOptionIndex).length > 0) {
|
||||
addLine(theme.fg("dim", " ↑↓ move • Enter submit • Tab edit note • Esc cancel"));
|
||||
} else {
|
||||
addLine(theme.fg("dim", " ↑↓ move • Enter submit • Tab add note • Esc cancel"));
|
||||
}
|
||||
|
||||
addLine(theme.fg("accent", "─".repeat(width)));
|
||||
cachedRenderedLines = renderedLines;
|
||||
return renderedLines;
|
||||
};
|
||||
|
||||
const handleInput = (data: string) => {
|
||||
if (isNoteEditorOpen) {
|
||||
if (matchesKey(data, Key.tab) || matchesKey(data, Key.escape)) {
|
||||
isNoteEditorOpen = false;
|
||||
requestUiRerender();
|
||||
return;
|
||||
}
|
||||
noteEditor.handleInput(data);
|
||||
requestUiRerender();
|
||||
return;
|
||||
}
|
||||
|
||||
if (matchesKey(data, Key.up)) {
|
||||
cursorOptionIndex = Math.max(0, cursorOptionIndex - 1);
|
||||
requestUiRerender();
|
||||
return;
|
||||
}
|
||||
if (matchesKey(data, Key.down)) {
|
||||
cursorOptionIndex = Math.min(selectableOptionLabels.length - 1, cursorOptionIndex + 1);
|
||||
requestUiRerender();
|
||||
return;
|
||||
}
|
||||
|
||||
if (matchesKey(data, Key.tab)) {
|
||||
isNoteEditorOpen = true;
|
||||
loadCurrentNoteIntoEditor();
|
||||
requestUiRerender();
|
||||
return;
|
||||
}
|
||||
|
||||
if (matchesKey(data, Key.enter)) {
|
||||
const selectedOptionLabel = selectableOptionLabels[cursorOptionIndex];
|
||||
const trimmedNote = getTrimmedNoteForOption(cursorOptionIndex);
|
||||
|
||||
if (selectedOptionLabel === OTHER_OPTION && !trimmedNote) {
|
||||
isNoteEditorOpen = true;
|
||||
loadCurrentNoteIntoEditor();
|
||||
requestUiRerender();
|
||||
return;
|
||||
}
|
||||
|
||||
submitCurrentSelection(selectedOptionLabel, trimmedNote);
|
||||
return;
|
||||
}
|
||||
|
||||
if (matchesKey(data, Key.escape)) {
|
||||
done({ cancelled: true });
|
||||
}
|
||||
};
|
||||
|
||||
return {
|
||||
render,
|
||||
invalidate: () => {
|
||||
cachedRenderedLines = undefined;
|
||||
},
|
||||
handleInput,
|
||||
};
|
||||
});
|
||||
|
||||
if (result.cancelled || !result.selectedOption) {
|
||||
return { selectedOptions: [] };
|
||||
}
|
||||
|
||||
return buildSingleSelectionResult(result.selectedOption, result.note);
|
||||
}
|
||||
@@ -1,98 +0,0 @@
|
||||
export const OTHER_OPTION = "Other (type your own)";
|
||||
const RECOMMENDED_OPTION_TAG = " (Recommended)";
|
||||
|
||||
export interface AskOption {
|
||||
label: string;
|
||||
}
|
||||
|
||||
export interface AskQuestion {
|
||||
id: string;
|
||||
question: string;
|
||||
options: AskOption[];
|
||||
multi?: boolean;
|
||||
recommended?: number;
|
||||
}
|
||||
|
||||
export interface AskSelection {
|
||||
selectedOptions: string[];
|
||||
customInput?: string;
|
||||
}
|
||||
|
||||
export function appendRecommendedTagToOptionLabels(
|
||||
optionLabels: string[],
|
||||
recommendedOptionIndex?: number,
|
||||
): string[] {
|
||||
if (
|
||||
recommendedOptionIndex == null ||
|
||||
recommendedOptionIndex < 0 ||
|
||||
recommendedOptionIndex >= optionLabels.length
|
||||
) {
|
||||
return optionLabels;
|
||||
}
|
||||
|
||||
return optionLabels.map((optionLabel, optionIndex) => {
|
||||
if (optionIndex !== recommendedOptionIndex) return optionLabel;
|
||||
if (optionLabel.endsWith(RECOMMENDED_OPTION_TAG)) return optionLabel;
|
||||
return `${optionLabel}${RECOMMENDED_OPTION_TAG}`;
|
||||
});
|
||||
}
|
||||
|
||||
function removeRecommendedTagFromOptionLabel(optionLabel: string): string {
|
||||
if (!optionLabel.endsWith(RECOMMENDED_OPTION_TAG)) {
|
||||
return optionLabel;
|
||||
}
|
||||
return optionLabel.slice(0, -RECOMMENDED_OPTION_TAG.length);
|
||||
}
|
||||
|
||||
export function buildSingleSelectionResult(selectedOptionLabel: string, note?: string): AskSelection {
|
||||
const normalizedSelectedOption = removeRecommendedTagFromOptionLabel(selectedOptionLabel);
|
||||
const normalizedNote = note?.trim();
|
||||
|
||||
if (normalizedSelectedOption === OTHER_OPTION) {
|
||||
if (normalizedNote) {
|
||||
return { selectedOptions: [], customInput: normalizedNote };
|
||||
}
|
||||
return { selectedOptions: [] };
|
||||
}
|
||||
|
||||
if (normalizedNote) {
|
||||
return { selectedOptions: [`${normalizedSelectedOption} - ${normalizedNote}`] };
|
||||
}
|
||||
|
||||
return { selectedOptions: [normalizedSelectedOption] };
|
||||
}
|
||||
|
||||
export function buildMultiSelectionResult(
|
||||
optionLabels: string[],
|
||||
selectedOptionIndexes: number[],
|
||||
optionNotes: string[],
|
||||
otherOptionIndex: number,
|
||||
): AskSelection {
|
||||
const selectedOptionSet = new Set(selectedOptionIndexes);
|
||||
const selectedOptions: string[] = [];
|
||||
let customInput: string | undefined;
|
||||
|
||||
for (let optionIndex = 0; optionIndex < optionLabels.length; optionIndex++) {
|
||||
if (!selectedOptionSet.has(optionIndex)) continue;
|
||||
|
||||
const optionLabel = removeRecommendedTagFromOptionLabel(optionLabels[optionIndex]);
|
||||
const optionNote = optionNotes[optionIndex]?.trim();
|
||||
|
||||
if (optionIndex === otherOptionIndex) {
|
||||
if (optionNote) customInput = optionNote;
|
||||
continue;
|
||||
}
|
||||
|
||||
if (optionNote) {
|
||||
selectedOptions.push(`${optionLabel} - ${optionNote}`);
|
||||
} else {
|
||||
selectedOptions.push(optionLabel);
|
||||
}
|
||||
}
|
||||
|
||||
if (customInput) {
|
||||
return { selectedOptions, customInput };
|
||||
}
|
||||
|
||||
return { selectedOptions };
|
||||
}
|
||||
@@ -1,514 +0,0 @@
|
||||
import type { ExtensionUIContext } from "@mariozechner/pi-coding-agent";
|
||||
import { Editor, type EditorTheme, Key, matchesKey, truncateToWidth, visibleWidth, wrapTextWithAnsi } from "@mariozechner/pi-tui";
|
||||
import {
|
||||
OTHER_OPTION,
|
||||
appendRecommendedTagToOptionLabels,
|
||||
buildMultiSelectionResult,
|
||||
buildSingleSelectionResult,
|
||||
type AskQuestion,
|
||||
type AskSelection,
|
||||
} from "./ask-logic";
|
||||
import { INLINE_NOTE_WRAP_PADDING, buildWrappedOptionLabelWithInlineNote } from "./ask-inline-note";
|
||||
|
||||
interface PreparedQuestion {
|
||||
id: string;
|
||||
question: string;
|
||||
options: string[];
|
||||
tabLabel: string;
|
||||
multi: boolean;
|
||||
otherOptionIndex: number;
|
||||
}
|
||||
|
||||
interface TabsUIState {
|
||||
cancelled: boolean;
|
||||
selectedOptionIndexesByQuestion: number[][];
|
||||
noteByQuestionByOption: string[][];
|
||||
}
|
||||
|
||||
export function formatSelectionForSubmitReview(selection: AskSelection, isMulti: boolean): string {
|
||||
const hasSelectedOptions = selection.selectedOptions.length > 0;
|
||||
const hasCustomInput = Boolean(selection.customInput);
|
||||
|
||||
if (hasSelectedOptions && hasCustomInput) {
|
||||
const selectedPart = isMulti
|
||||
? `[${selection.selectedOptions.join(", ")}]`
|
||||
: selection.selectedOptions[0];
|
||||
return `${selectedPart} + Other: ${selection.customInput}`;
|
||||
}
|
||||
|
||||
if (hasCustomInput) {
|
||||
return `Other: ${selection.customInput}`;
|
||||
}
|
||||
|
||||
if (hasSelectedOptions) {
|
||||
return isMulti ? `[${selection.selectedOptions.join(", ")}]` : selection.selectedOptions[0];
|
||||
}
|
||||
|
||||
return "(not answered)";
|
||||
}
|
||||
|
||||
function clampIndex(index: number | undefined, maxExclusive: number): number {
|
||||
if (index == null || Number.isNaN(index) || maxExclusive <= 0) return 0;
|
||||
if (index < 0) return 0;
|
||||
if (index >= maxExclusive) return maxExclusive - 1;
|
||||
return index;
|
||||
}
|
||||
|
||||
function normalizeTabLabel(id: string, fallback: string): string {
|
||||
const normalized = id.trim().replace(/[_-]+/g, " ");
|
||||
return normalized.length > 0 ? normalized : fallback;
|
||||
}
|
||||
|
||||
function buildSelectionForQuestion(
|
||||
question: PreparedQuestion,
|
||||
selectedOptionIndexes: number[],
|
||||
noteByOptionIndex: string[],
|
||||
): AskSelection {
|
||||
if (selectedOptionIndexes.length === 0) {
|
||||
return { selectedOptions: [] };
|
||||
}
|
||||
|
||||
if (question.multi) {
|
||||
return buildMultiSelectionResult(question.options, selectedOptionIndexes, noteByOptionIndex, question.otherOptionIndex);
|
||||
}
|
||||
|
||||
const selectedOptionIndex = selectedOptionIndexes[0];
|
||||
const selectedOptionLabel = question.options[selectedOptionIndex] ?? OTHER_OPTION;
|
||||
const note = noteByOptionIndex[selectedOptionIndex] ?? "";
|
||||
return buildSingleSelectionResult(selectedOptionLabel, note);
|
||||
}
|
||||
|
||||
function isQuestionSelectionValid(
|
||||
question: PreparedQuestion,
|
||||
selectedOptionIndexes: number[],
|
||||
noteByOptionIndex: string[],
|
||||
): boolean {
|
||||
if (selectedOptionIndexes.length === 0) return false;
|
||||
if (!selectedOptionIndexes.includes(question.otherOptionIndex)) return true;
|
||||
const otherNote = noteByOptionIndex[question.otherOptionIndex]?.trim() ?? "";
|
||||
return otherNote.length > 0;
|
||||
}
|
||||
|
||||
function createTabsUiStateSnapshot(
|
||||
cancelled: boolean,
|
||||
selectedOptionIndexesByQuestion: number[][],
|
||||
noteByQuestionByOption: string[][],
|
||||
): TabsUIState {
|
||||
return {
|
||||
cancelled,
|
||||
selectedOptionIndexesByQuestion: selectedOptionIndexesByQuestion.map((indexes) => [...indexes]),
|
||||
noteByQuestionByOption: noteByQuestionByOption.map((notes) => [...notes]),
|
||||
};
|
||||
}
|
||||
|
||||
function addIndexToSelection(selectedOptionIndexes: number[], optionIndex: number): number[] {
|
||||
if (selectedOptionIndexes.includes(optionIndex)) return selectedOptionIndexes;
|
||||
return [...selectedOptionIndexes, optionIndex].sort((a, b) => a - b);
|
||||
}
|
||||
|
||||
function removeIndexFromSelection(selectedOptionIndexes: number[], optionIndex: number): number[] {
|
||||
return selectedOptionIndexes.filter((index) => index !== optionIndex);
|
||||
}
|
||||
|
||||
export async function askQuestionsWithTabs(
|
||||
ui: ExtensionUIContext,
|
||||
questions: AskQuestion[],
|
||||
): Promise<{ cancelled: boolean; selections: AskSelection[] }> {
|
||||
const preparedQuestions: PreparedQuestion[] = questions.map((question, questionIndex) => {
|
||||
const baseOptionLabels = question.options.map((option) => option.label);
|
||||
const optionLabels = [...appendRecommendedTagToOptionLabels(baseOptionLabels, question.recommended), OTHER_OPTION];
|
||||
return {
|
||||
id: question.id,
|
||||
question: question.question,
|
||||
options: optionLabels,
|
||||
tabLabel: normalizeTabLabel(question.id, `Q${questionIndex + 1}`),
|
||||
multi: question.multi === true,
|
||||
otherOptionIndex: optionLabels.length - 1,
|
||||
};
|
||||
});
|
||||
|
||||
const initialCursorOptionIndexByQuestion = preparedQuestions.map((preparedQuestion, questionIndex) =>
|
||||
clampIndex(questions[questionIndex].recommended, preparedQuestion.options.length),
|
||||
);
|
||||
|
||||
const result = await ui.custom<TabsUIState>((tui, theme, _keybindings, done) => {
|
||||
let activeTabIndex = 0;
|
||||
let isNoteEditorOpen = false;
|
||||
let cachedRenderedLines: string[] | undefined;
|
||||
const cursorOptionIndexByQuestion = [...initialCursorOptionIndexByQuestion];
|
||||
const selectedOptionIndexesByQuestion = preparedQuestions.map(() => [] as number[]);
|
||||
const noteByQuestionByOption = preparedQuestions.map((preparedQuestion) =>
|
||||
Array(preparedQuestion.options.length).fill("") as string[],
|
||||
);
|
||||
|
||||
const editorTheme: EditorTheme = {
|
||||
borderColor: (text) => theme.fg("accent", text),
|
||||
selectList: {
|
||||
selectedPrefix: (text) => theme.fg("accent", text),
|
||||
selectedText: (text) => theme.fg("accent", text),
|
||||
description: (text) => theme.fg("muted", text),
|
||||
scrollInfo: (text) => theme.fg("dim", text),
|
||||
noMatch: (text) => theme.fg("warning", text),
|
||||
},
|
||||
};
|
||||
const noteEditor = new Editor(tui, editorTheme);
|
||||
|
||||
const submitTabIndex = preparedQuestions.length;
|
||||
|
||||
const requestUiRerender = () => {
|
||||
cachedRenderedLines = undefined;
|
||||
tui.requestRender();
|
||||
};
|
||||
|
||||
const getActiveQuestionIndex = (): number | null => {
|
||||
if (activeTabIndex >= preparedQuestions.length) return null;
|
||||
return activeTabIndex;
|
||||
};
|
||||
|
||||
const getQuestionNote = (questionIndex: number, optionIndex: number): string =>
|
||||
noteByQuestionByOption[questionIndex]?.[optionIndex] ?? "";
|
||||
|
||||
const getTrimmedQuestionNote = (questionIndex: number, optionIndex: number): string =>
|
||||
getQuestionNote(questionIndex, optionIndex).trim();
|
||||
|
||||
const isAllQuestionSelectionsValid = (): boolean =>
|
||||
preparedQuestions.every((preparedQuestion, questionIndex) =>
|
||||
isQuestionSelectionValid(
|
||||
preparedQuestion,
|
||||
selectedOptionIndexesByQuestion[questionIndex],
|
||||
noteByQuestionByOption[questionIndex],
|
||||
),
|
||||
);
|
||||
|
||||
const openNoteEditorForActiveOption = () => {
|
||||
const questionIndex = getActiveQuestionIndex();
|
||||
if (questionIndex == null) return;
|
||||
|
||||
isNoteEditorOpen = true;
|
||||
const optionIndex = cursorOptionIndexByQuestion[questionIndex];
|
||||
noteEditor.setText(getQuestionNote(questionIndex, optionIndex));
|
||||
requestUiRerender();
|
||||
};
|
||||
|
||||
const advanceToNextTabOrSubmit = () => {
|
||||
activeTabIndex = Math.min(submitTabIndex, activeTabIndex + 1);
|
||||
};
|
||||
|
||||
noteEditor.onChange = (value) => {
|
||||
const questionIndex = getActiveQuestionIndex();
|
||||
if (questionIndex == null) return;
|
||||
const optionIndex = cursorOptionIndexByQuestion[questionIndex];
|
||||
noteByQuestionByOption[questionIndex][optionIndex] = value;
|
||||
requestUiRerender();
|
||||
};
|
||||
|
||||
noteEditor.onSubmit = (value) => {
|
||||
const questionIndex = getActiveQuestionIndex();
|
||||
if (questionIndex == null) return;
|
||||
|
||||
const preparedQuestion = preparedQuestions[questionIndex];
|
||||
const optionIndex = cursorOptionIndexByQuestion[questionIndex];
|
||||
noteByQuestionByOption[questionIndex][optionIndex] = value;
|
||||
const trimmedNote = value.trim();
|
||||
|
||||
if (preparedQuestion.multi) {
|
||||
if (trimmedNote.length > 0) {
|
||||
selectedOptionIndexesByQuestion[questionIndex] = addIndexToSelection(
|
||||
selectedOptionIndexesByQuestion[questionIndex],
|
||||
optionIndex,
|
||||
);
|
||||
}
|
||||
if (optionIndex === preparedQuestion.otherOptionIndex && trimmedNote.length === 0) {
|
||||
requestUiRerender();
|
||||
return;
|
||||
}
|
||||
isNoteEditorOpen = false;
|
||||
requestUiRerender();
|
||||
return;
|
||||
}
|
||||
|
||||
selectedOptionIndexesByQuestion[questionIndex] = [optionIndex];
|
||||
if (optionIndex === preparedQuestion.otherOptionIndex && trimmedNote.length === 0) {
|
||||
requestUiRerender();
|
||||
return;
|
||||
}
|
||||
|
||||
isNoteEditorOpen = false;
|
||||
advanceToNextTabOrSubmit();
|
||||
requestUiRerender();
|
||||
};
|
||||
|
||||
const renderTabs = (): string => {
|
||||
const tabParts: string[] = ["← "];
|
||||
for (let questionIndex = 0; questionIndex < preparedQuestions.length; questionIndex++) {
|
||||
const preparedQuestion = preparedQuestions[questionIndex];
|
||||
const isActiveTab = questionIndex === activeTabIndex;
|
||||
const isQuestionValid = isQuestionSelectionValid(
|
||||
preparedQuestion,
|
||||
selectedOptionIndexesByQuestion[questionIndex],
|
||||
noteByQuestionByOption[questionIndex],
|
||||
);
|
||||
const statusIcon = isQuestionValid ? "■" : "□";
|
||||
const tabLabel = ` ${statusIcon} ${preparedQuestion.tabLabel} `;
|
||||
const styledTabLabel = isActiveTab
|
||||
? theme.bg("selectedBg", theme.fg("text", tabLabel))
|
||||
: theme.fg(isQuestionValid ? "success" : "muted", tabLabel);
|
||||
tabParts.push(`${styledTabLabel} `);
|
||||
}
|
||||
|
||||
const isSubmitTabActive = activeTabIndex === submitTabIndex;
|
||||
const canSubmit = isAllQuestionSelectionsValid();
|
||||
const submitLabel = " ✓ Submit ";
|
||||
const styledSubmitLabel = isSubmitTabActive
|
||||
? theme.bg("selectedBg", theme.fg("text", submitLabel))
|
||||
: theme.fg(canSubmit ? "success" : "dim", submitLabel);
|
||||
tabParts.push(`${styledSubmitLabel} →`);
|
||||
return tabParts.join("");
|
||||
};
|
||||
|
||||
const renderSubmitTab = (width: number, renderedLines: string[]): void => {
|
||||
const addLine = (line: string) => renderedLines.push(truncateToWidth(line, width));
|
||||
|
||||
addLine(theme.fg("accent", theme.bold(" Review answers")));
|
||||
renderedLines.push("");
|
||||
|
||||
for (let questionIndex = 0; questionIndex < preparedQuestions.length; questionIndex++) {
|
||||
const preparedQuestion = preparedQuestions[questionIndex];
|
||||
const selection = buildSelectionForQuestion(
|
||||
preparedQuestion,
|
||||
selectedOptionIndexesByQuestion[questionIndex],
|
||||
noteByQuestionByOption[questionIndex],
|
||||
);
|
||||
const value = formatSelectionForSubmitReview(selection, preparedQuestion.multi);
|
||||
const isValid = isQuestionSelectionValid(
|
||||
preparedQuestion,
|
||||
selectedOptionIndexesByQuestion[questionIndex],
|
||||
noteByQuestionByOption[questionIndex],
|
||||
);
|
||||
const statusIcon = isValid ? theme.fg("success", "●") : theme.fg("warning", "○");
|
||||
addLine(` ${statusIcon} ${theme.fg("muted", `${preparedQuestion.tabLabel}:`)} ${theme.fg("text", value)}`);
|
||||
}
|
||||
|
||||
renderedLines.push("");
|
||||
if (isAllQuestionSelectionsValid()) {
|
||||
addLine(theme.fg("success", " Press Enter to submit"));
|
||||
} else {
|
||||
const missingQuestions = preparedQuestions
|
||||
.filter((preparedQuestion, questionIndex) =>
|
||||
!isQuestionSelectionValid(
|
||||
preparedQuestion,
|
||||
selectedOptionIndexesByQuestion[questionIndex],
|
||||
noteByQuestionByOption[questionIndex],
|
||||
),
|
||||
)
|
||||
.map((preparedQuestion) => preparedQuestion.tabLabel)
|
||||
.join(", ");
|
||||
addLine(theme.fg("warning", ` Complete required answers: ${missingQuestions}`));
|
||||
}
|
||||
addLine(theme.fg("dim", " ←/→ switch tabs • Esc cancel"));
|
||||
};
|
||||
|
||||
const renderQuestionTab = (width: number, renderedLines: string[], questionIndex: number): void => {
|
||||
const addLine = (line: string) => renderedLines.push(truncateToWidth(line, width));
|
||||
const preparedQuestion = preparedQuestions[questionIndex];
|
||||
const cursorOptionIndex = cursorOptionIndexByQuestion[questionIndex];
|
||||
const selectedOptionIndexes = selectedOptionIndexesByQuestion[questionIndex];
|
||||
|
||||
for (const questionLine of wrapTextWithAnsi(preparedQuestion.question, Math.max(1, width - 1))) {
|
||||
addLine(` ${theme.fg("text", questionLine)}`);
|
||||
}
|
||||
renderedLines.push("");
|
||||
|
||||
for (let optionIndex = 0; optionIndex < preparedQuestion.options.length; optionIndex++) {
|
||||
const optionLabel = preparedQuestion.options[optionIndex];
|
||||
const isCursorOption = optionIndex === cursorOptionIndex;
|
||||
const isOptionSelected = selectedOptionIndexes.includes(optionIndex);
|
||||
const isEditingThisOption = isNoteEditorOpen && isCursorOption;
|
||||
const cursorPrefixText = isCursorOption ? "→ " : " ";
|
||||
const cursorPrefix = isCursorOption ? theme.fg("accent", cursorPrefixText) : cursorPrefixText;
|
||||
const markerText = preparedQuestion.multi
|
||||
? `${isOptionSelected ? "[x]" : "[ ]"} `
|
||||
: `${isOptionSelected ? "●" : "○"} `;
|
||||
const optionColor = isCursorOption ? "accent" : isOptionSelected ? "success" : "text";
|
||||
const prefixWidth = visibleWidth(cursorPrefixText) + visibleWidth(markerText);
|
||||
const wrappedInlineLabelLines = buildWrappedOptionLabelWithInlineNote(
|
||||
optionLabel,
|
||||
getQuestionNote(questionIndex, optionIndex),
|
||||
isEditingThisOption,
|
||||
Math.max(1, width - prefixWidth),
|
||||
INLINE_NOTE_WRAP_PADDING,
|
||||
);
|
||||
const continuationPrefix = " ".repeat(prefixWidth);
|
||||
addLine(`${cursorPrefix}${theme.fg(optionColor, `${markerText}${wrappedInlineLabelLines[0] ?? ""}`)}`);
|
||||
for (const wrappedLine of wrappedInlineLabelLines.slice(1)) {
|
||||
addLine(`${continuationPrefix}${theme.fg(optionColor, wrappedLine)}`);
|
||||
}
|
||||
}
|
||||
|
||||
renderedLines.push("");
|
||||
if (isNoteEditorOpen) {
|
||||
addLine(theme.fg("dim", " Typing note inline • Enter save note • Tab/Esc stop editing"));
|
||||
} else {
|
||||
if (preparedQuestion.multi) {
|
||||
addLine(
|
||||
theme.fg(
|
||||
"dim",
|
||||
" ↑↓ move • Enter toggle/select • Tab add note • ←/→ switch tabs • Esc cancel",
|
||||
),
|
||||
);
|
||||
} else {
|
||||
addLine(
|
||||
theme.fg("dim", " ↑↓ move • Enter select • Tab add note • ←/→ switch tabs • Esc cancel"),
|
||||
);
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
const render = (width: number): string[] => {
|
||||
if (cachedRenderedLines) return cachedRenderedLines;
|
||||
|
||||
const renderedLines: string[] = [];
|
||||
const addLine = (line: string) => renderedLines.push(truncateToWidth(line, width));
|
||||
|
||||
addLine(theme.fg("accent", "─".repeat(width)));
|
||||
addLine(` ${renderTabs()}`);
|
||||
renderedLines.push("");
|
||||
|
||||
if (activeTabIndex === submitTabIndex) {
|
||||
renderSubmitTab(width, renderedLines);
|
||||
} else {
|
||||
renderQuestionTab(width, renderedLines, activeTabIndex);
|
||||
}
|
||||
|
||||
addLine(theme.fg("accent", "─".repeat(width)));
|
||||
cachedRenderedLines = renderedLines;
|
||||
return renderedLines;
|
||||
};
|
||||
|
||||
const handleInput = (data: string) => {
|
||||
if (isNoteEditorOpen) {
|
||||
if (matchesKey(data, Key.tab) || matchesKey(data, Key.escape)) {
|
||||
isNoteEditorOpen = false;
|
||||
requestUiRerender();
|
||||
return;
|
||||
}
|
||||
noteEditor.handleInput(data);
|
||||
requestUiRerender();
|
||||
return;
|
||||
}
|
||||
|
||||
if (matchesKey(data, Key.left)) {
|
||||
activeTabIndex = (activeTabIndex - 1 + preparedQuestions.length + 1) % (preparedQuestions.length + 1);
|
||||
requestUiRerender();
|
||||
return;
|
||||
}
|
||||
|
||||
if (matchesKey(data, Key.right)) {
|
||||
activeTabIndex = (activeTabIndex + 1) % (preparedQuestions.length + 1);
|
||||
requestUiRerender();
|
||||
return;
|
||||
}
|
||||
|
||||
if (activeTabIndex === submitTabIndex) {
|
||||
if (matchesKey(data, Key.enter) && isAllQuestionSelectionsValid()) {
|
||||
done(createTabsUiStateSnapshot(false, selectedOptionIndexesByQuestion, noteByQuestionByOption));
|
||||
return;
|
||||
}
|
||||
if (matchesKey(data, Key.escape)) {
|
||||
done(createTabsUiStateSnapshot(true, selectedOptionIndexesByQuestion, noteByQuestionByOption));
|
||||
}
|
||||
return;
|
||||
}
|
||||
|
||||
const questionIndex = activeTabIndex;
|
||||
const preparedQuestion = preparedQuestions[questionIndex];
|
||||
|
||||
if (matchesKey(data, Key.up)) {
|
||||
cursorOptionIndexByQuestion[questionIndex] = Math.max(0, cursorOptionIndexByQuestion[questionIndex] - 1);
|
||||
requestUiRerender();
|
||||
return;
|
||||
}
|
||||
|
||||
if (matchesKey(data, Key.down)) {
|
||||
cursorOptionIndexByQuestion[questionIndex] = Math.min(
|
||||
preparedQuestion.options.length - 1,
|
||||
cursorOptionIndexByQuestion[questionIndex] + 1,
|
||||
);
|
||||
requestUiRerender();
|
||||
return;
|
||||
}
|
||||
|
||||
if (matchesKey(data, Key.tab)) {
|
||||
openNoteEditorForActiveOption();
|
||||
return;
|
||||
}
|
||||
|
||||
if (matchesKey(data, Key.enter)) {
|
||||
const cursorOptionIndex = cursorOptionIndexByQuestion[questionIndex];
|
||||
|
||||
if (preparedQuestion.multi) {
|
||||
const currentlySelected = selectedOptionIndexesByQuestion[questionIndex];
|
||||
if (currentlySelected.includes(cursorOptionIndex)) {
|
||||
selectedOptionIndexesByQuestion[questionIndex] = removeIndexFromSelection(currentlySelected, cursorOptionIndex);
|
||||
} else {
|
||||
selectedOptionIndexesByQuestion[questionIndex] = addIndexToSelection(currentlySelected, cursorOptionIndex);
|
||||
}
|
||||
|
||||
if (
|
||||
cursorOptionIndex === preparedQuestion.otherOptionIndex &&
|
||||
selectedOptionIndexesByQuestion[questionIndex].includes(cursorOptionIndex) &&
|
||||
getTrimmedQuestionNote(questionIndex, cursorOptionIndex).length === 0
|
||||
) {
|
||||
openNoteEditorForActiveOption();
|
||||
return;
|
||||
}
|
||||
|
||||
requestUiRerender();
|
||||
return;
|
||||
}
|
||||
|
||||
selectedOptionIndexesByQuestion[questionIndex] = [cursorOptionIndex];
|
||||
if (
|
||||
cursorOptionIndex === preparedQuestion.otherOptionIndex &&
|
||||
getTrimmedQuestionNote(questionIndex, cursorOptionIndex).length === 0
|
||||
) {
|
||||
openNoteEditorForActiveOption();
|
||||
return;
|
||||
}
|
||||
|
||||
advanceToNextTabOrSubmit();
|
||||
requestUiRerender();
|
||||
return;
|
||||
}
|
||||
|
||||
if (matchesKey(data, Key.escape)) {
|
||||
done(createTabsUiStateSnapshot(true, selectedOptionIndexesByQuestion, noteByQuestionByOption));
|
||||
}
|
||||
};
|
||||
|
||||
return {
|
||||
render,
|
||||
invalidate: () => {
|
||||
cachedRenderedLines = undefined;
|
||||
},
|
||||
handleInput,
|
||||
};
|
||||
});
|
||||
|
||||
if (result.cancelled) {
|
||||
return {
|
||||
cancelled: true,
|
||||
selections: preparedQuestions.map(() => ({ selectedOptions: [] } satisfies AskSelection)),
|
||||
};
|
||||
}
|
||||
|
||||
const selections = preparedQuestions.map((preparedQuestion, questionIndex) =>
|
||||
buildSelectionForQuestion(
|
||||
preparedQuestion,
|
||||
result.selectedOptionIndexesByQuestion[questionIndex] ?? [],
|
||||
result.noteByQuestionByOption[questionIndex] ?? Array(preparedQuestion.options.length).fill(""),
|
||||
),
|
||||
);
|
||||
|
||||
return { cancelled: result.cancelled, selections };
|
||||
}
|
||||
@@ -1,37 +0,0 @@
|
||||
// Pi Ask Tool CLI Agent
|
||||
|
||||
import { ask_claude } from "../../@piplugin/ask-claude"
|
||||
import { sessionId } from '../shared';
|
||||
|
||||
export async function start() {
|
||||
console.log('Pi Ask Tool initialized');
|
||||
|
||||
while (true) {
|
||||
// Get user input from TUI (simplified for example)
|
||||
const userInput = await getTUIInput('Ask a question:');
|
||||
|
||||
// Handle multi-turn sessions via sessionId
|
||||
const response = await ask_claude({
|
||||
prompt: userInput,
|
||||
agent: 'code_review', // Default agent
|
||||
session_id: sessionId
|
||||
});
|
||||
|
||||
// Display answer in TUI
|
||||
await showTUIResult(response);
|
||||
|
||||
// Update session context if needed
|
||||
sessionId = response.session_id || sessionId;
|
||||
}
|
||||
}
|
||||
|
||||
// Mock TUI handlers - implement actual TUI integration
|
||||
async function getTUIInput(question: string) {
|
||||
// Replace with real TUI input method
|
||||
const input = process.stdin.read().toString();
|
||||
return input;
|
||||
}
|
||||
|
||||
async function showTUIResult(result: any) {
|
||||
console.log('Answer:', result.summary || result);
|
||||
}
|
||||
@@ -1,237 +0,0 @@
|
||||
import type { ExtensionAPI } from "@mariozechner/pi-coding-agent";
|
||||
import { Type, type Static } from "@sinclair/typebox";
|
||||
import { OTHER_OPTION, type AskQuestion } from "./ask-logic";
|
||||
import { askSingleQuestionWithInlineNote } from "./ask-inline-ui";
|
||||
import { askQuestionsWithTabs } from "./ask-tabs-ui";
|
||||
|
||||
const OptionItemSchema = Type.Object({
|
||||
label: Type.String({ description: "Display label" }),
|
||||
});
|
||||
|
||||
const QuestionItemSchema = Type.Object({
|
||||
id: Type.String({ description: "Question id (e.g. auth, cache, priority)" }),
|
||||
question: Type.String({ description: "Question text" }),
|
||||
options: Type.Array(OptionItemSchema, {
|
||||
description: "Available options. Do not include 'Other'.",
|
||||
minItems: 1,
|
||||
}),
|
||||
multi: Type.Optional(Type.Boolean({ description: "Allow multi-select" })),
|
||||
recommended: Type.Optional(
|
||||
Type.Number({ description: "0-indexed recommended option. '(Recommended)' is shown automatically." }),
|
||||
),
|
||||
});
|
||||
|
||||
const AskParamsSchema = Type.Object({
|
||||
questions: Type.Array(QuestionItemSchema, { description: "Questions to ask", minItems: 1 }),
|
||||
});
|
||||
|
||||
type AskParams = Static<typeof AskParamsSchema>;
|
||||
|
||||
interface QuestionResult {
|
||||
id: string;
|
||||
question: string;
|
||||
options: string[];
|
||||
multi: boolean;
|
||||
selectedOptions: string[];
|
||||
customInput?: string;
|
||||
}
|
||||
|
||||
interface AskToolDetails {
|
||||
id?: string;
|
||||
question?: string;
|
||||
options?: string[];
|
||||
multi?: boolean;
|
||||
selectedOptions?: string[];
|
||||
customInput?: string;
|
||||
results?: QuestionResult[];
|
||||
}
|
||||
|
||||
function sanitizeForSessionText(value: string): string {
|
||||
return value
|
||||
.replace(/[\r\n\t]/g, " ")
|
||||
.replace(/[\x00-\x08\x0B\x0C\x0E-\x1F\x7F]/g, "")
|
||||
.replace(/\s{2,}/g, " ")
|
||||
.trim();
|
||||
}
|
||||
|
||||
function sanitizeOptionForSessionText(option: string): string {
|
||||
const sanitizedOption = sanitizeForSessionText(option);
|
||||
return sanitizedOption.length > 0 ? sanitizedOption : "(empty option)";
|
||||
}
|
||||
|
||||
function toSessionSafeQuestionResult(result: QuestionResult): QuestionResult {
|
||||
const selectedOptions = result.selectedOptions
|
||||
.map((selectedOption) => sanitizeForSessionText(selectedOption))
|
||||
.filter((selectedOption) => selectedOption.length > 0);
|
||||
|
||||
const rawCustomInput = result.customInput;
|
||||
const customInput = rawCustomInput == null ? undefined : sanitizeForSessionText(rawCustomInput);
|
||||
|
||||
return {
|
||||
id: sanitizeForSessionText(result.id) || "(unknown)",
|
||||
question: sanitizeForSessionText(result.question) || "(empty question)",
|
||||
options: result.options.map(sanitizeOptionForSessionText),
|
||||
multi: result.multi,
|
||||
selectedOptions,
|
||||
customInput: customInput && customInput.length > 0 ? customInput : undefined,
|
||||
};
|
||||
}
|
||||
|
||||
function formatSelectionForSummary(result: QuestionResult): string {
|
||||
const hasSelectedOptions = result.selectedOptions.length > 0;
|
||||
const hasCustomInput = Boolean(result.customInput);
|
||||
|
||||
if (!hasSelectedOptions && !hasCustomInput) {
|
||||
return "(cancelled)";
|
||||
}
|
||||
|
||||
if (hasSelectedOptions && hasCustomInput) {
|
||||
const selectedPart = result.multi
|
||||
? `[${result.selectedOptions.join(", ")}]`
|
||||
: result.selectedOptions[0];
|
||||
return `${selectedPart} + Other: "${result.customInput}"`;
|
||||
}
|
||||
|
||||
if (hasCustomInput) {
|
||||
return `"${result.customInput}"`;
|
||||
}
|
||||
|
||||
if (result.multi) {
|
||||
return `[${result.selectedOptions.join(", ")}]`;
|
||||
}
|
||||
|
||||
return result.selectedOptions[0];
|
||||
}
|
||||
|
||||
function formatQuestionResult(result: QuestionResult): string {
|
||||
return `${result.id}: ${formatSelectionForSummary(result)}`;
|
||||
}
|
||||
|
||||
function formatQuestionContext(result: QuestionResult, questionIndex: number): string {
|
||||
const lines: string[] = [
|
||||
`Question ${questionIndex + 1} (${result.id})`,
|
||||
`Prompt: ${result.question}`,
|
||||
"Options:",
|
||||
...result.options.map((option, optionIndex) => ` ${optionIndex + 1}. ${option}`),
|
||||
"Response:",
|
||||
];
|
||||
|
||||
const hasSelectedOptions = result.selectedOptions.length > 0;
|
||||
const hasCustomInput = Boolean(result.customInput);
|
||||
|
||||
if (!hasSelectedOptions && !hasCustomInput) {
|
||||
lines.push(" Selected: (cancelled)");
|
||||
return lines.join("\n");
|
||||
}
|
||||
|
||||
if (hasSelectedOptions) {
|
||||
const selectedText = result.multi
|
||||
? `[${result.selectedOptions.join(", ")}]`
|
||||
: result.selectedOptions[0];
|
||||
lines.push(` Selected: ${selectedText}`);
|
||||
}
|
||||
|
||||
if (hasCustomInput) {
|
||||
if (!hasSelectedOptions) {
|
||||
lines.push(` Selected: ${OTHER_OPTION}`);
|
||||
}
|
||||
lines.push(` Custom input: ${result.customInput}`);
|
||||
}
|
||||
|
||||
return lines.join("\n");
|
||||
}
|
||||
|
||||
function buildAskSessionContent(results: QuestionResult[]): string {
|
||||
const safeResults = results.map(toSessionSafeQuestionResult);
|
||||
const summaryLines = safeResults.map(formatQuestionResult).join("\n");
|
||||
const contextBlocks = safeResults.map((result, index) => formatQuestionContext(result, index)).join("\n\n");
|
||||
return `User answers:\n${summaryLines}\n\nAnswer context:\n${contextBlocks}`;
|
||||
}
|
||||
|
||||
const ASK_TOOL_DESCRIPTION = `
|
||||
Ask the user for clarification when a choice materially affects the outcome.
|
||||
|
||||
- Use when multiple valid approaches have different trade-offs.
|
||||
- Prefer 2-5 concise options.
|
||||
- Use multi=true when multiple answers are valid.
|
||||
- Use recommended=<index> (0-indexed) to mark the default option.
|
||||
- You can ask multiple related questions in one call using questions[].
|
||||
- Do NOT include an 'Other' option; UI adds it automatically.
|
||||
`.trim();
|
||||
|
||||
export default function askExtension(pi: ExtensionAPI) {
|
||||
pi.registerTool({
|
||||
name: "ask",
|
||||
label: "Ask",
|
||||
description: ASK_TOOL_DESCRIPTION,
|
||||
parameters: AskParamsSchema,
|
||||
|
||||
async execute(_toolCallId, params: AskParams, _signal, _onUpdate, ctx) {
|
||||
if (!ctx.hasUI) {
|
||||
return {
|
||||
content: [{ type: "text", text: "Error: ask tool requires interactive mode" }],
|
||||
details: {},
|
||||
};
|
||||
}
|
||||
|
||||
if (params.questions.length === 0) {
|
||||
return {
|
||||
content: [{ type: "text", text: "Error: questions must not be empty" }],
|
||||
details: {},
|
||||
};
|
||||
}
|
||||
|
||||
if (params.questions.length === 1) {
|
||||
const [q] = params.questions;
|
||||
const selection = q.multi
|
||||
? (await askQuestionsWithTabs(ctx.ui, [q as AskQuestion])).selections[0] ?? { selectedOptions: [] }
|
||||
: await askSingleQuestionWithInlineNote(ctx.ui, q as AskQuestion);
|
||||
const optionLabels = q.options.map((option) => option.label);
|
||||
|
||||
const result: QuestionResult = {
|
||||
id: q.id,
|
||||
question: q.question,
|
||||
options: optionLabels,
|
||||
multi: q.multi ?? false,
|
||||
selectedOptions: selection.selectedOptions,
|
||||
customInput: selection.customInput,
|
||||
};
|
||||
|
||||
const details: AskToolDetails = {
|
||||
id: q.id,
|
||||
question: q.question,
|
||||
options: optionLabels,
|
||||
multi: q.multi ?? false,
|
||||
selectedOptions: selection.selectedOptions,
|
||||
customInput: selection.customInput,
|
||||
results: [result],
|
||||
};
|
||||
|
||||
return {
|
||||
content: [{ type: "text", text: buildAskSessionContent([result]) }],
|
||||
details,
|
||||
};
|
||||
}
|
||||
|
||||
const results: QuestionResult[] = [];
|
||||
const tabResult = await askQuestionsWithTabs(ctx.ui, params.questions as AskQuestion[]);
|
||||
for (let i = 0; i < params.questions.length; i++) {
|
||||
const q = params.questions[i];
|
||||
const selection = tabResult.selections[i] ?? { selectedOptions: [] };
|
||||
results.push({
|
||||
id: q.id,
|
||||
question: q.question,
|
||||
options: q.options.map((option) => option.label),
|
||||
multi: q.multi ?? false,
|
||||
selectedOptions: selection.selectedOptions,
|
||||
customInput: selection.customInput,
|
||||
});
|
||||
}
|
||||
|
||||
return {
|
||||
content: [{ type: "text", text: buildAskSessionContent(results) }],
|
||||
details: { results } satisfies AskToolDetails,
|
||||
};
|
||||
},
|
||||
});
|
||||
}
|
||||
151
pi/.pi/agent/extensions/prompt-frame.ts
Normal file
151
pi/.pi/agent/extensions/prompt-frame.ts
Normal file
@@ -0,0 +1,151 @@
|
||||
import { CustomEditor } from "@earendil-works/pi-coding-agent";
|
||||
import { splitBorder, frameLine } from "./lib/boxes.ts";
|
||||
|
||||
/**
|
||||
* prompt-frame: Wraps the prompt (editor) in the shared rounded-box border
|
||||
* and embeds the footer into the box's bottom border, tool-block style:
|
||||
*
|
||||
* ╭──────────────────────────────────────────────────────────╮
|
||||
* │ ❯ type here… │
|
||||
* ╰─ glm-5.2 · high · ctx 12% ──── S: 53% W: 21% M: 66% ─╯
|
||||
*
|
||||
* Footer contents:
|
||||
* left: current model id · thinking level · context percent
|
||||
* right: opencode usage from the usage-bars extension, formatted as
|
||||
* "S: 53% <glyph>" (5h) / "W: …" (weekly) / "M: …" (monthly),
|
||||
* where <glyph> is a nerd-font circle-slice progress icon.
|
||||
*
|
||||
* The default footer is hidden (its factory is still used to capture the
|
||||
* FooterDataProvider so we can read usage-bars' status string).
|
||||
* The box border color keeps pi's thinking-level / bash-mode signal.
|
||||
*/
|
||||
|
||||
const USAGE_STATUS_KEY = "usage-bars";
|
||||
|
||||
/** Nerd-font circle-slice progress glyphs (nf-md-circle_slice_1..8). */
|
||||
const SLICES = ["\u{F0A9E}", "\u{F0A9F}", "\u{F0AA0}", "\u{F0AA1}", "\u{F0AA2}", "\u{F0AA3}", "\u{F0AA4}", "\u{F0AA5}"];
|
||||
|
||||
function sliceGlyph(percent: number) {
|
||||
const idx = Math.min(8, Math.max(1, Math.ceil((percent / 100) * 8)));
|
||||
return SLICES[idx - 1];
|
||||
}
|
||||
|
||||
function usageColor(percent: number) {
|
||||
if (percent >= 90) return "error";
|
||||
if (percent >= 70) return "warning";
|
||||
return "success";
|
||||
}
|
||||
|
||||
/** `S: 53% ` — percent + glyph colored by usage level. */
|
||||
function usageSegment(theme: any, label: string, percent: number) {
|
||||
// const color = usageColor(percent);
|
||||
return theme.fg("muted", `${label}: `) + theme.fg("muted", `${percent}% ${sliceGlyph(percent)}`);
|
||||
}
|
||||
|
||||
const ANSI_RE = /\x1b\[[0-9;]*m/g;
|
||||
|
||||
/**
|
||||
* Parse S/W/M percentages out of usage-bars' footer status string
|
||||
* (e.g. "OpenCode Go S ████░░░░ 53% ⟳ 2h W ██░░░░░░ 21% M ███░░░░░ 66%").
|
||||
*/
|
||||
function parseUsage(status: string) {
|
||||
const plain = status.replace(ANSI_RE, "");
|
||||
const grab = (label: string) => {
|
||||
const m = plain.match(new RegExp(`\\b${label}\\s+[█░]*\\s*(\\d+(?:\\.\\d+)?)%`));
|
||||
return m ? Math.round(parseFloat(m[1])) : undefined;
|
||||
};
|
||||
return { session: grab("S"), weekly: grab("W"), monthly: grab("M") };
|
||||
}
|
||||
|
||||
class HiddenFooter {
|
||||
render() {
|
||||
return [];
|
||||
}
|
||||
invalidate() {}
|
||||
}
|
||||
|
||||
export default function (pi: any) {
|
||||
let activeTui: any;
|
||||
let footerData: any;
|
||||
|
||||
// usage-bars emits this after every poll - keep the border fresh.
|
||||
pi.events?.on?.("usage:update", () => activeTui?.requestRender());
|
||||
|
||||
pi.on("session_shutdown", () => {
|
||||
activeTui = undefined;
|
||||
});
|
||||
|
||||
pi.on("session_start", (_event: any, ctx: any) => {
|
||||
// Hide the default footer, but keep a handle on its data provider so
|
||||
// we can read extension statuses (usage-bars).
|
||||
ctx.ui.setFooter((_tui: any, _theme: any, data: any) => {
|
||||
footerData = data;
|
||||
return new HiddenFooter();
|
||||
});
|
||||
|
||||
const footerLeft = (theme: any) => {
|
||||
const model = ctx.model?.id ?? "no model";
|
||||
const thinking = pi.getThinkingLevel();
|
||||
const usage = ctx.getContextUsage();
|
||||
const ctxPct = usage?.percent != null ? `${Math.round(usage.percent)}%` : "?";
|
||||
return (
|
||||
theme.fg("text", model) +
|
||||
theme.fg("muted", " · ") +
|
||||
theme.fg("dim", String(thinking)) +
|
||||
theme.fg("muted", " · ") +
|
||||
theme.fg("dim", `ctx ${ctxPct}`)
|
||||
);
|
||||
};
|
||||
|
||||
const footerRight = (theme: any) => {
|
||||
const status = footerData?.getExtensionStatuses().get(USAGE_STATUS_KEY);
|
||||
if (!status) return "";
|
||||
const { session, weekly, monthly } = parseUsage(status);
|
||||
const parts: string[] = [];
|
||||
if (session !== undefined) parts.push(usageSegment(theme, "s", session));
|
||||
if (weekly !== undefined) parts.push(usageSegment(theme, "w", weekly));
|
||||
if (monthly !== undefined) parts.push(usageSegment(theme, "m", monthly));
|
||||
if (parts.length === 0) {
|
||||
// loading / unavailable - show the raw message, dimmed
|
||||
return theme.fg("dim", status.replace(ANSI_RE, ""));
|
||||
}
|
||||
return parts.join(" ");
|
||||
};
|
||||
|
||||
class FramedEditor extends CustomEditor {
|
||||
constructor(tui: any, theme: any, keybindings: any) {
|
||||
super(tui, theme, keybindings, { paddingX: 0 });
|
||||
activeTui = tui;
|
||||
}
|
||||
|
||||
render(width: number) {
|
||||
const border = (s: string) => this.borderColor(s);
|
||||
const inner = super.render(Math.max(10, width - 4));
|
||||
if (inner.length < 2) return inner;
|
||||
|
||||
const theme = ctx.ui.theme;
|
||||
// The editor's own first/last lines are plain rules, or scroll
|
||||
// indicators like "─── ↑ 2 more ───" - carry those into our border.
|
||||
const scrollInfo = (line: string) => {
|
||||
const m = line.replace(ANSI_RE, "").match(/([↑↓] \d+ more)/);
|
||||
return m ? theme.fg("dim", m[1]) : "";
|
||||
};
|
||||
const top = splitBorder(width, scrollInfo(inner[0]), "", border, ["╭", "╮"]);
|
||||
const bottom = splitBorder(
|
||||
width,
|
||||
scrollInfo(inner[inner.length - 1]) || footerLeft(theme),
|
||||
footerRight(theme),
|
||||
border,
|
||||
["╰", "╯"],
|
||||
);
|
||||
|
||||
const out = [top];
|
||||
for (const line of inner.slice(1, -1)) out.push(frameLine(width, line, border));
|
||||
out.push(bottom);
|
||||
return out;
|
||||
}
|
||||
}
|
||||
|
||||
ctx.ui.setEditorComponent((tui: any, theme: any, keybindings: any) => new FramedEditor(tui, theme, keybindings));
|
||||
});
|
||||
}
|
||||
367
pi/.pi/agent/extensions/tool-blocks.ts
Normal file
367
pi/.pi/agent/extensions/tool-blocks.ts
Normal file
@@ -0,0 +1,367 @@
|
||||
import {
|
||||
CustomEditor,
|
||||
ToolExecutionComponent,
|
||||
createReadTool,
|
||||
createBashTool,
|
||||
createEditTool,
|
||||
createWriteTool,
|
||||
} from "@earendil-works/pi-coding-agent";
|
||||
import { wrapTextWithAnsi } from "@earendil-works/pi-tui";
|
||||
import { topBorder, bottomBorder, frameLine, splitBorder, borderColor, sanitize, formatDuration } from "./lib/boxes.ts";
|
||||
|
||||
/**
|
||||
* tool-blocks: Custom rendering for the built-in read/bash/edit/write tools.
|
||||
*
|
||||
* - Draws a full rounded box around each tool block (renderShell: "self",
|
||||
* so the default background tint is gone). Border color follows status:
|
||||
* accent while running, success green when done, error red on failure.
|
||||
* - bash / edit / write are ALWAYS fully expanded (the `expanded` flag from
|
||||
* ctrl+o is ignored for the body).
|
||||
* - read stays collapsed, showing only "N lines read" in the bottom border;
|
||||
* ctrl+o still expands it to show the content.
|
||||
* - bash summary shows duration + exit code; edit shows +added/-removed;
|
||||
* write shows lines written.
|
||||
*
|
||||
* Execution is delegated to the original tool implementations - only the
|
||||
* rendering changes.
|
||||
*/
|
||||
|
||||
/** Timing info per tool call (only for calls executed in this process). */
|
||||
const runs = new Map<string, { start: number; end?: number }>();
|
||||
|
||||
/** Top border with the tool title. Returned from renderCall. */
|
||||
class ToolBoxTop {
|
||||
constructor(
|
||||
public title: string,
|
||||
public color: (s: string) => string,
|
||||
) {}
|
||||
|
||||
invalidate() {}
|
||||
|
||||
render(width: number) {
|
||||
return [topBorder(width, this.title, this.color)];
|
||||
}
|
||||
}
|
||||
|
||||
/** Body lines (wrapped, framed with │) + bottom border. Returned from renderResult. */
|
||||
class ToolBoxBody {
|
||||
constructor(
|
||||
public body: string[],
|
||||
public summary: string,
|
||||
public color: (s: string) => string,
|
||||
) {}
|
||||
|
||||
invalidate() {}
|
||||
|
||||
render(width: number) {
|
||||
if (width < 8) return [bottomBorder(width, "", this.color)];
|
||||
const inner = width - 4;
|
||||
const lines: string[] = [];
|
||||
for (const raw of this.body) {
|
||||
const wrapped = raw === "" ? [""] : wrapTextWithAnsi(sanitize(raw), inner);
|
||||
for (const seg of wrapped.length > 0 ? wrapped : [""]) {
|
||||
lines.push(frameLine(width, seg, this.color));
|
||||
}
|
||||
}
|
||||
lines.push(bottomBorder(width, this.summary, this.color));
|
||||
return lines;
|
||||
}
|
||||
}
|
||||
|
||||
/**
|
||||
* Capture ALL other tool blocks (any extension's tools and fallbacks) in the
|
||||
* same box style. There is no public hook for tools we didn't register, so
|
||||
* this patches ToolExecutionComponent.prototype.render: tools that draw their
|
||||
* own shell (renderShell: "self", i.e. the four above) are left alone; every
|
||||
* other block gets its default rendering wrapped in a status-colored box.
|
||||
* Defensive: any error falls back to the unpatched renderer.
|
||||
*/
|
||||
function patchAllToolBlocks() {
|
||||
const proto = ToolExecutionComponent.prototype as any;
|
||||
if (proto.__boxedBlocks) return;
|
||||
proto.__boxedBlocks = true;
|
||||
const originalRender = proto.render;
|
||||
|
||||
proto.render = function (width: number) {
|
||||
try {
|
||||
if (this.hideComponent) return originalRender.call(this, width);
|
||||
if (this.hasRendererDefinition?.() && this.getRenderShell?.() === "self") {
|
||||
return originalRender.call(this, width);
|
||||
}
|
||||
const theme = (globalThis as any)[Symbol.for("@earendil-works/pi-coding-agent:theme")];
|
||||
if (!theme || width < 12) return originalRender.call(this, width);
|
||||
|
||||
const inner = originalRender.call(this, width - 4);
|
||||
// Kitty graphics can't be re-framed - leave those blocks untouched.
|
||||
if (inner.some((l: string) => l.includes("\x1b_G"))) return originalRender.call(this, width);
|
||||
// Strip the leading Spacer line and the Box's padding blanks
|
||||
// (padding lines are bg-filled spaces, so trim ANSI-stripped text).
|
||||
const isBlank = (l: string) => l.replace(/\x1b\[[0-9;]*m/g, "").trim() === "";
|
||||
while (inner.length > 0 && isBlank(inner[0])) inner.shift();
|
||||
while (inner.length > 0 && isBlank(inner[inner.length - 1])) inner.pop();
|
||||
if (inner.length === 0) return [];
|
||||
|
||||
const status = this.isPartial ? "running" : this.result ? (this.result.isError ? "error" : "ok") : "running";
|
||||
const color = borderColor(theme, status);
|
||||
const name = String(this.toolName ?? "tool");
|
||||
// The fallback call renderer prints just the tool name - drop it,
|
||||
// the name already sits in the top border. Keep richer call lines.
|
||||
if (inner.length > 0 && inner[0].replace(/\x1b\[[0-9;]*m/g, "").trim() === name) inner.shift();
|
||||
while (inner.length > 0 && isBlank(inner[0])) inner.shift();
|
||||
const title = theme.fg("toolTitle", theme.bold(name));
|
||||
const lines = ["", topBorder(width, title, color)];
|
||||
for (const line of inner) lines.push(frameLine(width, line, color));
|
||||
lines.push(bottomBorder(width, "", color));
|
||||
return lines;
|
||||
} catch {
|
||||
return originalRender.call(this, width);
|
||||
}
|
||||
};
|
||||
}
|
||||
|
||||
/** Status as seen from renderCall (result not available there). */
|
||||
function callStatus(context: any) {
|
||||
if (context.isError) return "error";
|
||||
if (context.isPartial) return "running";
|
||||
const rec = runs.get(context.toolCallId);
|
||||
if (rec && rec.end === undefined) return "running";
|
||||
return "ok";
|
||||
}
|
||||
|
||||
function resultStatus(isPartial: boolean, isError: boolean) {
|
||||
return isPartial ? "running" : isError ? "error" : "ok";
|
||||
}
|
||||
|
||||
function title(theme: any, name: string, arg?: string, extra?: string) {
|
||||
let t = theme.fg("toolTitle", theme.bold(name));
|
||||
if (arg) t += theme.fg("muted", " · ") + theme.fg("accent", sanitize(arg));
|
||||
if (extra) t += theme.fg("dim", ` ${extra}`);
|
||||
return t;
|
||||
}
|
||||
|
||||
function textOf(result: any) {
|
||||
const block = result.content?.find((c: any) => c.type === "text");
|
||||
return block?.text ?? "";
|
||||
}
|
||||
|
||||
export default function (pi: any) {
|
||||
const cwd = process.cwd();
|
||||
patchAllToolBlocks();
|
||||
|
||||
// --- read: collapsed by default, line count in the summary, ctrl+o expands ---
|
||||
const originalRead = createReadTool(cwd);
|
||||
pi.registerTool({
|
||||
name: "read",
|
||||
label: "read",
|
||||
description: originalRead.description,
|
||||
parameters: originalRead.parameters,
|
||||
renderShell: "self",
|
||||
|
||||
async execute(id: string, params: any, signal: any, onUpdate: any) {
|
||||
runs.set(id, { start: Date.now() });
|
||||
try {
|
||||
return await originalRead.execute(id, params, signal, onUpdate);
|
||||
} finally {
|
||||
const rec = runs.get(id);
|
||||
if (rec) rec.end = Date.now();
|
||||
}
|
||||
},
|
||||
|
||||
renderCall(args: any, theme: any, context: any) {
|
||||
const parts: string[] = [];
|
||||
if (args?.offset) parts.push(`offset=${args.offset}`);
|
||||
if (args?.limit) parts.push(`limit=${args.limit}`);
|
||||
const extra = parts.length > 0 ? `(${parts.join(", ")})` : undefined;
|
||||
return new ToolBoxTop(title(theme, "read", args?.path ?? "", extra), borderColor(theme, callStatus(context)));
|
||||
},
|
||||
|
||||
renderResult(result: any, { expanded, isPartial }: any, theme: any, context: any) {
|
||||
const status = resultStatus(isPartial, context.isError);
|
||||
const color = borderColor(theme, status);
|
||||
if (isPartial) return new ToolBoxBody([], theme.fg("dim", "reading…"), color);
|
||||
|
||||
const details = result.details;
|
||||
const image = result.content?.find((c: any) => c.type === "image");
|
||||
if (image) return new ToolBoxBody([], theme.fg("success", "image loaded"), color);
|
||||
|
||||
const text = textOf(result);
|
||||
if (context.isError) {
|
||||
return new ToolBoxBody([theme.fg("error", text.split("\n")[0] ?? "error")], theme.fg("error", "error"), color);
|
||||
}
|
||||
|
||||
// Don't count the "[N more lines in file...]" notice the tool appends.
|
||||
const contentLines = text.split("\n");
|
||||
while (contentLines.length > 0) {
|
||||
const last = contentLines[contentLines.length - 1];
|
||||
if (last.trim() === "" || /^\[\d+ more lines/.test(last)) contentLines.pop();
|
||||
else break;
|
||||
}
|
||||
const lineCount = contentLines.length;
|
||||
let summary = theme.fg("success", `${lineCount} lines read`);
|
||||
const truncation = details?.truncation;
|
||||
if (truncation?.truncated) {
|
||||
summary += theme.fg("warning", ` (of ${truncation.totalLines})`);
|
||||
}
|
||||
|
||||
const body = expanded ? text.replace(/\s+$/, "").split("\n").map((l: string) => theme.fg("dim", sanitize(l))) : [];
|
||||
return new ToolBoxBody(body, summary, color);
|
||||
},
|
||||
});
|
||||
|
||||
// --- bash: always fully expanded, duration + exit code in the summary ---
|
||||
const originalBash = createBashTool(cwd);
|
||||
pi.registerTool({
|
||||
name: "bash",
|
||||
label: "bash",
|
||||
description: originalBash.description,
|
||||
parameters: originalBash.parameters,
|
||||
renderShell: "self",
|
||||
|
||||
async execute(id: string, params: any, signal: any, onUpdate: any) {
|
||||
runs.set(id, { start: Date.now() });
|
||||
try {
|
||||
return await originalBash.execute(id, params, signal, onUpdate);
|
||||
} finally {
|
||||
const rec = runs.get(id);
|
||||
if (rec) rec.end = Date.now();
|
||||
}
|
||||
},
|
||||
|
||||
renderCall(args: any, theme: any, context: any) {
|
||||
const cmd = (args?.command ?? "").replace(/\s+/g, " ").slice(0, 200);
|
||||
return new ToolBoxTop(title(theme, "bash", cmd), borderColor(theme, callStatus(context)));
|
||||
},
|
||||
|
||||
renderResult(result: any, { isPartial }: any, theme: any, context: any) {
|
||||
const status = resultStatus(isPartial, context.isError);
|
||||
const color = borderColor(theme, status);
|
||||
const details = result.details;
|
||||
const output = textOf(result).replace(/\s+$/, "");
|
||||
// Always fully expanded, regardless of the ctrl+o toggle.
|
||||
const body = output === "" ? [] : output.split("\n").map((l: string) => theme.fg("toolOutput", sanitize(l)));
|
||||
|
||||
const rec = runs.get(context.toolCallId);
|
||||
const parts: string[] = [];
|
||||
if (isPartial) {
|
||||
parts.push(theme.fg("accent", "running…"));
|
||||
if (rec) parts.push(theme.fg("dim", formatDuration(Date.now() - rec.start)));
|
||||
} else {
|
||||
if (context.isError) {
|
||||
const match = output.match(/Command exited with code (\d+)/);
|
||||
parts.push(theme.fg("error", match ? `exit ${match[1]}` : "failed"));
|
||||
} else {
|
||||
parts.push(theme.fg("success", "ok"));
|
||||
}
|
||||
if (rec?.end) parts.push(theme.fg("dim", formatDuration(rec.end - rec.start)));
|
||||
parts.push(theme.fg("dim", `${body.length} lines`));
|
||||
if (details?.truncation?.truncated) parts.push(theme.fg("warning", "truncated"));
|
||||
}
|
||||
return new ToolBoxBody(body, parts.join(theme.fg("muted", " · ")), color);
|
||||
},
|
||||
});
|
||||
|
||||
// --- edit: always shows the full diff, +N -N in the summary ---
|
||||
const originalEdit = createEditTool(cwd);
|
||||
pi.registerTool({
|
||||
name: "edit",
|
||||
label: "edit",
|
||||
description: originalEdit.description,
|
||||
parameters: originalEdit.parameters,
|
||||
renderShell: "self",
|
||||
|
||||
async execute(id: string, params: any, signal: any, onUpdate: any) {
|
||||
runs.set(id, { start: Date.now() });
|
||||
try {
|
||||
return await originalEdit.execute(id, params, signal, onUpdate);
|
||||
} finally {
|
||||
const rec = runs.get(id);
|
||||
if (rec) rec.end = Date.now();
|
||||
}
|
||||
},
|
||||
|
||||
renderCall(args: any, theme: any, context: any) {
|
||||
return new ToolBoxTop(title(theme, "edit", args?.path ?? ""), borderColor(theme, callStatus(context)));
|
||||
},
|
||||
|
||||
renderResult(result: any, { isPartial }: any, theme: any, context: any) {
|
||||
const status = resultStatus(isPartial, context.isError);
|
||||
const color = borderColor(theme, status);
|
||||
if (isPartial) return new ToolBoxBody([], theme.fg("dim", "editing…"), color);
|
||||
|
||||
const text = textOf(result);
|
||||
if (context.isError || text.startsWith("Error")) {
|
||||
return new ToolBoxBody(
|
||||
text.split("\n").map((l: string) => theme.fg("error", sanitize(l))),
|
||||
theme.fg("error", "error"),
|
||||
borderColor(theme, "error"),
|
||||
);
|
||||
}
|
||||
|
||||
const details = result.details;
|
||||
if (!details?.diff) return new ToolBoxBody([], theme.fg("success", "applied"), color);
|
||||
|
||||
const diffLines = details.diff.replace(/\s+$/, "").split("\n");
|
||||
let added = 0;
|
||||
let removed = 0;
|
||||
const body = diffLines.map((line: string) => {
|
||||
if (line.startsWith("+") && !line.startsWith("+++")) {
|
||||
added++;
|
||||
return theme.fg("toolDiffAdded", sanitize(line));
|
||||
}
|
||||
if (line.startsWith("-") && !line.startsWith("---")) {
|
||||
removed++;
|
||||
return theme.fg("toolDiffRemoved", sanitize(line));
|
||||
}
|
||||
return theme.fg("toolDiffContext", sanitize(line));
|
||||
});
|
||||
const summary =
|
||||
theme.fg("toolDiffAdded", `+${added}`) + theme.fg("muted", " · ") + theme.fg("toolDiffRemoved", `−${removed}`);
|
||||
return new ToolBoxBody(body, summary, color);
|
||||
},
|
||||
});
|
||||
|
||||
// --- write: always shows the full content, line count in the summary ---
|
||||
const originalWrite = createWriteTool(cwd);
|
||||
pi.registerTool({
|
||||
name: "write",
|
||||
label: "write",
|
||||
description: originalWrite.description,
|
||||
parameters: originalWrite.parameters,
|
||||
renderShell: "self",
|
||||
|
||||
async execute(id: string, params: any, signal: any, onUpdate: any) {
|
||||
runs.set(id, { start: Date.now() });
|
||||
try {
|
||||
return await originalWrite.execute(id, params, signal, onUpdate);
|
||||
} finally {
|
||||
const rec = runs.get(id);
|
||||
if (rec) rec.end = Date.now();
|
||||
}
|
||||
},
|
||||
|
||||
renderCall(args: any, theme: any, context: any) {
|
||||
return new ToolBoxTop(title(theme, "write", args?.path ?? ""), borderColor(theme, callStatus(context)));
|
||||
},
|
||||
|
||||
renderResult(result: any, { isPartial }: any, theme: any, context: any) {
|
||||
const status = resultStatus(isPartial, context.isError);
|
||||
const color = borderColor(theme, status);
|
||||
if (isPartial) return new ToolBoxBody([], theme.fg("dim", "writing…"), color);
|
||||
|
||||
const text = textOf(result);
|
||||
if (context.isError || text.startsWith("Error")) {
|
||||
return new ToolBoxBody(
|
||||
text.split("\n").map((l: string) => theme.fg("error", sanitize(l))),
|
||||
theme.fg("error", "error"),
|
||||
borderColor(theme, "error"),
|
||||
);
|
||||
}
|
||||
|
||||
const content = context.args?.content ?? "";
|
||||
const body =
|
||||
content === "" ? [] : content.replace(/\s+$/, "").split("\n").map((l: string) => theme.fg("toolOutput", sanitize(l)));
|
||||
const summary = theme.fg("success", `${body.length} lines written`);
|
||||
return new ToolBoxBody(body, summary, color);
|
||||
},
|
||||
});
|
||||
}
|
||||
245
pi/.pi/agent/extensions/transcript-viewer.ts
Normal file
245
pi/.pi/agent/extensions/transcript-viewer.ts
Normal file
@@ -0,0 +1,245 @@
|
||||
import { wrapTextWithAnsi, visibleWidth, truncateToWidth, matchesKey, Key } from "@earendil-works/pi-tui";
|
||||
import { topBorder, splitBorder, sanitize } from "./lib/boxes.ts";
|
||||
|
||||
/**
|
||||
* transcript-viewer: A scrollable transcript overlay with a claude-cloak style
|
||||
* scrollbar. User messages are marked with `*` on the scrollbar track and you
|
||||
* can jump between them.
|
||||
*
|
||||
* pi has no API to scroll the live chat (it lives in the terminal's native
|
||||
* scrollback), so this renders the session transcript in an overlay with its
|
||||
* own scroll state.
|
||||
*
|
||||
* Keys:
|
||||
* ctrl+q open the viewer
|
||||
* n / N jump to next / previous user message
|
||||
* j / k, ↓ / ↑ scroll one line
|
||||
* ctrl+d / ctrl+u half page down / up
|
||||
* g / G top / bottom
|
||||
* q / esc close
|
||||
*/
|
||||
|
||||
function contentToText(content: any): string {
|
||||
if (typeof content === "string") return content;
|
||||
if (Array.isArray(content)) {
|
||||
return content
|
||||
.map((block: any) => (block?.type === "text" ? block.text : block?.type === "image" ? "[image]" : ""))
|
||||
.filter((t: string) => t !== "")
|
||||
.join("\n");
|
||||
}
|
||||
return "";
|
||||
}
|
||||
|
||||
function argsSummary(args: any) {
|
||||
if (!args) return "";
|
||||
const preferred = ["command", "path", "pattern", "url"];
|
||||
for (const key of preferred) {
|
||||
if (typeof args[key] === "string") return args[key].replace(/\s+/g, " ");
|
||||
}
|
||||
const first = Object.values(args).find((v) => typeof v === "string");
|
||||
return first ? (first as string).replace(/\s+/g, " ") : "";
|
||||
}
|
||||
|
||||
interface Item {
|
||||
kind: "user" | "assistant" | "tool" | "info";
|
||||
text: string;
|
||||
}
|
||||
|
||||
/** Flatten the current session branch into displayable items. */
|
||||
function buildItems(ctx: any): Item[] {
|
||||
const items: Item[] = [];
|
||||
for (const entry of ctx.sessionManager.getBranch()) {
|
||||
if (entry.type === "message") {
|
||||
const msg = entry.message;
|
||||
if (msg.role === "user") {
|
||||
const text = contentToText(msg.content).trim();
|
||||
if (text) items.push({ kind: "user", text });
|
||||
} else if (msg.role === "assistant") {
|
||||
for (const block of msg.content ?? []) {
|
||||
if (block?.type === "text" && block.text?.trim()) {
|
||||
items.push({ kind: "assistant", text: block.text.trim() });
|
||||
} else if (block?.type === "toolCall") {
|
||||
items.push({ kind: "tool", text: `${block.name} ${argsSummary(block.arguments)}`.trim() });
|
||||
}
|
||||
}
|
||||
}
|
||||
} else if (entry.type === "custom_message" && entry.display) {
|
||||
const text = contentToText(entry.content).trim();
|
||||
if (text) items.push({ kind: "info", text });
|
||||
} else if (entry.type === "compaction") {
|
||||
items.push({ kind: "info", text: "— context compacted —" });
|
||||
} else if (entry.type === "branch_summary") {
|
||||
items.push({ kind: "info", text: "— branched —" });
|
||||
}
|
||||
}
|
||||
return items;
|
||||
}
|
||||
|
||||
class TranscriptViewer {
|
||||
/** -1 = not yet laid out; opens scrolled to the bottom. */
|
||||
scroll = -1;
|
||||
cacheWidth = -1;
|
||||
lines: string[] = [];
|
||||
userOffsets: number[] = [];
|
||||
lastViewHeight = 10;
|
||||
lastMaxScroll = 0;
|
||||
|
||||
constructor(
|
||||
public items: Item[],
|
||||
public tui: any,
|
||||
public theme: any,
|
||||
public done: (v?: any) => void,
|
||||
) {}
|
||||
|
||||
invalidate() {
|
||||
this.cacheWidth = -1;
|
||||
}
|
||||
|
||||
/** Wrap all items to the inner width; record the first line of each user message. */
|
||||
layout(inner: number) {
|
||||
if (this.cacheWidth === inner) return;
|
||||
this.cacheWidth = inner;
|
||||
this.lines = [];
|
||||
this.userOffsets = [];
|
||||
const theme = this.theme;
|
||||
|
||||
for (const item of this.items) {
|
||||
if (this.lines.length > 0) this.lines.push("");
|
||||
const text = sanitize(item.text);
|
||||
switch (item.kind) {
|
||||
case "user": {
|
||||
this.userOffsets.push(this.lines.length);
|
||||
const wrapped = wrapTextWithAnsi(text, Math.max(1, inner - 2));
|
||||
wrapped.forEach((line: string, i: number) => {
|
||||
this.lines.push(
|
||||
theme.fg("accent", theme.bold(i === 0 ? "❯ " : " ")) + theme.fg("userMessageText", theme.bold(line)),
|
||||
);
|
||||
});
|
||||
break;
|
||||
}
|
||||
case "assistant":
|
||||
for (const line of wrapTextWithAnsi(text, inner)) this.lines.push(line);
|
||||
break;
|
||||
case "tool":
|
||||
this.lines.push(theme.fg("muted", truncateToWidth(`⚙ ${text}`, inner, "…")));
|
||||
break;
|
||||
case "info":
|
||||
this.lines.push(theme.fg("dim", theme.italic(truncateToWidth(text, inner, "…"))));
|
||||
break;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
render(width: number) {
|
||||
const theme = this.theme;
|
||||
const border = (s: string) => theme.fg("border", s);
|
||||
if (width < 12) return [truncateToWidth(theme.fg("dim", "too narrow"), Math.max(1, width))];
|
||||
|
||||
const rows = this.tui?.terminal?.rows || process.stdout.rows || 24;
|
||||
const viewHeight = Math.max(5, rows - 8);
|
||||
// Layout: "│ " + content + " " + scrollbar-column
|
||||
const inner = width - 4;
|
||||
this.layout(inner);
|
||||
|
||||
const total = this.lines.length;
|
||||
const maxScroll = Math.max(0, total - viewHeight);
|
||||
if (this.scroll === -1) this.scroll = maxScroll;
|
||||
this.scroll = Math.max(0, Math.min(this.scroll, maxScroll));
|
||||
this.lastViewHeight = viewHeight;
|
||||
this.lastMaxScroll = maxScroll;
|
||||
|
||||
// Scrollbar track: markers map the WHOLE transcript onto the track,
|
||||
// thumb shows the current viewport (claude-cloak math).
|
||||
const markerRows = new Set<number>();
|
||||
for (const offset of this.userOffsets) {
|
||||
const row = total <= viewHeight ? offset : Math.floor((offset * viewHeight) / total);
|
||||
markerRows.add(Math.min(row, viewHeight - 1));
|
||||
}
|
||||
let thumbTop = -1;
|
||||
let thumbLen = 0;
|
||||
if (total > viewHeight) {
|
||||
thumbLen = Math.min(viewHeight, Math.max(1, Math.floor((viewHeight * viewHeight) / total)));
|
||||
thumbTop = maxScroll === 0 ? 0 : Math.floor((this.scroll * (viewHeight - thumbLen)) / maxScroll);
|
||||
}
|
||||
|
||||
const out: string[] = [];
|
||||
out.push(topBorder(width, theme.fg("toolTitle", theme.bold("transcript")), border));
|
||||
|
||||
for (let row = 0; row < viewHeight; row++) {
|
||||
const line = this.lines[this.scroll + row] ?? "";
|
||||
const pad = Math.max(0, inner - visibleWidth(line));
|
||||
let track: string;
|
||||
if (thumbTop !== -1 && row >= thumbTop && row < thumbTop + thumbLen) {
|
||||
track = theme.fg("accent", "█");
|
||||
} else if (markerRows.has(row)) {
|
||||
track = theme.fg("warning", theme.bold("*"));
|
||||
} else {
|
||||
track = border("│");
|
||||
}
|
||||
out.push(border("│ ") + line + " ".repeat(pad) + " " + track);
|
||||
}
|
||||
|
||||
const hints = theme.fg("dim", "n/N prompts · j/k · ctrl+d/u · g/G · q closes");
|
||||
const pos = theme.fg("dim", `${Math.min(total, this.scroll + viewHeight)}/${total}`);
|
||||
out.push(splitBorder(width, hints, pos, border, ["╰", "╯"]));
|
||||
return out;
|
||||
}
|
||||
|
||||
handleInput(data: string) {
|
||||
const half = Math.max(1, Math.floor(this.lastViewHeight / 2));
|
||||
if (data === "q" || matchesKey(data, Key.escape)) {
|
||||
this.done();
|
||||
return;
|
||||
} else if (data === "j" || matchesKey(data, Key.down)) {
|
||||
this.scroll += 1;
|
||||
} else if (data === "k" || matchesKey(data, Key.up)) {
|
||||
this.scroll -= 1;
|
||||
} else if (matchesKey(data, Key.ctrl("d")) || matchesKey(data, Key.pageDown)) {
|
||||
this.scroll += half;
|
||||
} else if (matchesKey(data, Key.ctrl("u")) || matchesKey(data, Key.pageUp)) {
|
||||
this.scroll -= half;
|
||||
} else if (data === "g" || matchesKey(data, Key.home)) {
|
||||
this.scroll = 0;
|
||||
} else if (data === "G" || matchesKey(data, Key.end)) {
|
||||
this.scroll = this.lastMaxScroll;
|
||||
} else if (data === "n") {
|
||||
const next = this.userOffsets.find((o) => o > this.scroll);
|
||||
if (next !== undefined) this.scroll = next;
|
||||
} else if (data === "N") {
|
||||
const prev = [...this.userOffsets].reverse().find((o) => o < this.scroll);
|
||||
if (prev !== undefined) this.scroll = prev;
|
||||
} else {
|
||||
return;
|
||||
}
|
||||
this.scroll = Math.max(0, Math.min(this.scroll, this.lastMaxScroll));
|
||||
this.tui.requestRender();
|
||||
}
|
||||
}
|
||||
|
||||
export default function (pi: any) {
|
||||
let open = false;
|
||||
|
||||
pi.registerShortcut("ctrl+q", {
|
||||
description: "Open transcript viewer (scrollbar + user-message jumps)",
|
||||
handler: async (ctx: any) => {
|
||||
if (open) return;
|
||||
open = true;
|
||||
try {
|
||||
await ctx.ui.custom(
|
||||
(tui: any, theme: any, _keybindings: any, done: any) =>
|
||||
new TranscriptViewer(buildItems(ctx), tui, theme, () => done(undefined)),
|
||||
{
|
||||
overlay: true,
|
||||
overlayOptions: {
|
||||
width: "85%",
|
||||
minWidth: 50,
|
||||
anchor: "center",
|
||||
},
|
||||
},
|
||||
);
|
||||
} finally {
|
||||
open = false;
|
||||
}
|
||||
},
|
||||
});
|
||||
}
|
||||
@@ -1,899 +0,0 @@
|
||||
import * as fs from "node:fs";
|
||||
import * as path from "node:path";
|
||||
import * as os from "node:os";
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Shared disk cache — lets multiple concurrent pi sessions coordinate so only
|
||||
// one actually hits the API per cache window, regardless of how many sessions
|
||||
// are open. Modelled after claude-pulse's cache.json approach.
|
||||
// ---------------------------------------------------------------------------
|
||||
|
||||
export interface UsageCache {
|
||||
timestamp: number;
|
||||
data: Partial<Record<ProviderKey, UsageData>>;
|
||||
/** ISO timestamp until which a provider is rate-limited (429 backoff). */
|
||||
rateLimitedUntil?: Partial<Record<ProviderKey, number>>;
|
||||
}
|
||||
|
||||
const USAGE_CACHE_FILE = path.join(os.homedir(), ".pi", "agent", "usage-cache.json");
|
||||
|
||||
export function readUsageCache(): UsageCache | null {
|
||||
try {
|
||||
const raw = fs.readFileSync(USAGE_CACHE_FILE, "utf-8");
|
||||
const parsed = JSON.parse(raw);
|
||||
if (typeof parsed?.timestamp === "number") return parsed as UsageCache;
|
||||
} catch {}
|
||||
return null;
|
||||
}
|
||||
|
||||
export function writeUsageCache(cache: UsageCache): void {
|
||||
try {
|
||||
const dir = path.dirname(USAGE_CACHE_FILE);
|
||||
if (!fs.existsSync(dir)) fs.mkdirSync(dir, { recursive: true });
|
||||
const tmp = `${USAGE_CACHE_FILE}.tmp-${process.pid}-${Date.now()}`;
|
||||
fs.writeFileSync(tmp, JSON.stringify(cache, null, 2));
|
||||
fs.renameSync(tmp, USAGE_CACHE_FILE);
|
||||
} catch {}
|
||||
}
|
||||
|
||||
export type ProviderKey = "codex" | "claude" | "zai" | "gemini" | "antigravity" | "opencode-go";
|
||||
export type OAuthProviderId = "openai-codex" | "anthropic" | "google-gemini-cli" | "google-antigravity" | "opencode-go";
|
||||
|
||||
export interface AuthData {
|
||||
"openai-codex"?: { access?: string; refresh?: string; expires?: number };
|
||||
anthropic?: { access?: string; refresh?: string; expires?: number };
|
||||
zai?: { key?: string; access?: string; refresh?: string; expires?: number };
|
||||
"google-gemini-cli"?: { access?: string; refresh?: string; projectId?: string; expires?: number };
|
||||
"google-antigravity"?: { access?: string; refresh?: string; projectId?: string; expires?: number };
|
||||
"opencode-go"?: { key?: string; access?: string };
|
||||
}
|
||||
|
||||
export interface UsageData {
|
||||
session: number;
|
||||
weekly: number;
|
||||
sessionResetsIn?: string;
|
||||
/** Unix ms timestamp of when the session window resets (from the raw API response). */
|
||||
sessionResetsAt?: number;
|
||||
weeklyResetsIn?: string;
|
||||
/** Unix ms timestamp of when the weekly window resets. */
|
||||
weeklyResetsAt?: number;
|
||||
extraSpend?: number;
|
||||
extraLimit?: number;
|
||||
error?: string;
|
||||
}
|
||||
|
||||
export type UsageByProvider = Record<ProviderKey, UsageData | null>;
|
||||
|
||||
// OpenCode Go usage tracking (local, since no API exists yet)
|
||||
export interface OpenCodeGoLocalUsage {
|
||||
/** Dollar value used in the 5-hour window */
|
||||
fiveHourUsed: number;
|
||||
/** Dollar value used in the weekly window */
|
||||
weeklyUsed: number;
|
||||
/** Dollar value used in the monthly window */
|
||||
monthlyUsed: number;
|
||||
/** Timestamp of the last update */
|
||||
lastUpdated: number;
|
||||
/** When the current 5-hour window started */
|
||||
fiveHourWindowStart: number;
|
||||
/** When the current week started (Unix ms) */
|
||||
weekStart: number;
|
||||
/** When the current month started (Unix ms) */
|
||||
monthStart: number;
|
||||
}
|
||||
|
||||
const OPENCODE_GO_USAGE_FILE = path.join(os.homedir(), ".pi", "agent", "opencode-go-usage.json");
|
||||
const OPENCODE_GO_FIVE_HOUR_LIMIT = 12;
|
||||
const OPENCODE_GO_WEEKLY_LIMIT = 30;
|
||||
const OPENCODE_GO_MONTHLY_LIMIT = 60;
|
||||
const OPENCODE_GO_FIVE_HOUR_MS = 5 * 60 * 60 * 1000;
|
||||
|
||||
export function readOpenCodeGoUsage(): OpenCodeGoLocalUsage | null {
|
||||
try {
|
||||
const raw = fs.readFileSync(OPENCODE_GO_USAGE_FILE, "utf-8");
|
||||
const parsed = JSON.parse(raw);
|
||||
if (typeof parsed?.lastUpdated === "number") return parsed as OpenCodeGoLocalUsage;
|
||||
} catch {}
|
||||
return null;
|
||||
}
|
||||
|
||||
export function writeOpenCodeGoUsage(usage: OpenCodeGoLocalUsage): void {
|
||||
try {
|
||||
const dir = path.dirname(OPENCODE_GO_USAGE_FILE);
|
||||
if (!fs.existsSync(dir)) fs.mkdirSync(dir, { recursive: true });
|
||||
const tmp = `${OPENCODE_GO_USAGE_FILE}.tmp-${process.pid}-${Date.now()}`;
|
||||
fs.writeFileSync(tmp, JSON.stringify(usage, null, 2));
|
||||
fs.renameSync(tmp, OPENCODE_GO_USAGE_FILE);
|
||||
} catch {}
|
||||
}
|
||||
|
||||
export function resetOpenCodeGoUsageIfNeeded(existing: OpenCodeGoLocalUsage | null): OpenCodeGoLocalUsage {
|
||||
const now = Date.now();
|
||||
const nowDate = new Date(now);
|
||||
|
||||
// Start with defaults
|
||||
let usage: OpenCodeGoLocalUsage = existing ?? {
|
||||
fiveHourUsed: 0,
|
||||
weeklyUsed: 0,
|
||||
monthlyUsed: 0,
|
||||
lastUpdated: now,
|
||||
fiveHourWindowStart: now,
|
||||
weekStart: now,
|
||||
monthStart: now,
|
||||
};
|
||||
|
||||
// Reset 5-hour window if expired
|
||||
if (now - usage.fiveHourWindowStart >= OPENCODE_GO_FIVE_HOUR_MS) {
|
||||
usage.fiveHourUsed = 0;
|
||||
usage.fiveHourWindowStart = now;
|
||||
}
|
||||
|
||||
// Reset weekly window (Monday-based)
|
||||
const dayOfWeek = nowDate.getDay();
|
||||
const daysSinceMonday = dayOfWeek === 0 ? 6 : dayOfWeek - 1;
|
||||
const thisMonday = new Date(nowDate);
|
||||
thisMonday.setDate(nowDate.getDate() - daysSinceMonday);
|
||||
thisMonday.setHours(0, 0, 0, 0);
|
||||
if (usage.weekStart < thisMonday.getTime()) {
|
||||
usage.weeklyUsed = 0;
|
||||
usage.weekStart = thisMonday.getTime();
|
||||
}
|
||||
|
||||
// Reset monthly window (1st of month)
|
||||
const thisMonthStart = new Date(nowDate.getFullYear(), nowDate.getMonth(), 1);
|
||||
if (usage.monthStart < thisMonthStart.getTime()) {
|
||||
usage.monthlyUsed = 0;
|
||||
usage.monthStart = thisMonthStart.getTime();
|
||||
}
|
||||
|
||||
usage.lastUpdated = now;
|
||||
return usage;
|
||||
}
|
||||
|
||||
export function addOpenCodeGoSpend(dollars: number): void {
|
||||
let usage = resetOpenCodeGoUsageIfNeeded(readOpenCodeGoUsage());
|
||||
usage.fiveHourUsed += dollars;
|
||||
usage.weeklyUsed += dollars;
|
||||
usage.monthlyUsed += dollars;
|
||||
usage.lastUpdated = Date.now();
|
||||
writeOpenCodeGoUsage(usage);
|
||||
}
|
||||
|
||||
export function getOpenCodeGoUsageData(): UsageData {
|
||||
const usage = resetOpenCodeGoUsageIfNeeded(readOpenCodeGoUsage());
|
||||
if (!usage) {
|
||||
return { session: 0, weekly: 0, error: "no local usage data" };
|
||||
}
|
||||
|
||||
const sessionPct = Math.min(100, (usage.fiveHourUsed / OPENCODE_GO_FIVE_HOUR_LIMIT) * 100);
|
||||
const weeklyPct = Math.min(100, (usage.weeklyUsed / OPENCODE_GO_WEEKLY_LIMIT) * 100);
|
||||
|
||||
// Calculate resets
|
||||
const fiveHourEnd = usage.fiveHourWindowStart + OPENCODE_GO_FIVE_HOUR_MS;
|
||||
const fiveHourRemaining = Math.max(0, fiveHourEnd - Date.now());
|
||||
|
||||
const weekEnd = usage.weekStart + 7 * 24 * 60 * 60 * 1000;
|
||||
const weekRemaining = Math.max(0, weekEnd - Date.now());
|
||||
|
||||
return {
|
||||
session: sessionPct,
|
||||
weekly: weeklyPct,
|
||||
sessionResetsIn: formatDuration(Math.round(fiveHourRemaining / 1000)),
|
||||
weeklyResetsIn: formatDuration(Math.round(weekRemaining / 1000)),
|
||||
extraSpend: usage.monthlyUsed,
|
||||
extraLimit: OPENCODE_GO_MONTHLY_LIMIT,
|
||||
};
|
||||
}
|
||||
|
||||
export interface UsageEndpoints {
|
||||
zai: string;
|
||||
gemini: string;
|
||||
antigravity: string;
|
||||
googleLoadCodeAssistEndpoints: string[];
|
||||
}
|
||||
|
||||
export interface FetchResponseLike {
|
||||
ok: boolean;
|
||||
status: number;
|
||||
json(): Promise<any>;
|
||||
}
|
||||
|
||||
export type FetchLike = (input: string, init?: RequestInit) => Promise<FetchResponseLike>;
|
||||
|
||||
export interface RequestConfig {
|
||||
fetchFn?: FetchLike;
|
||||
timeoutMs?: number;
|
||||
}
|
||||
|
||||
export interface FetchConfig extends RequestConfig {
|
||||
endpoints?: UsageEndpoints;
|
||||
env?: NodeJS.ProcessEnv;
|
||||
}
|
||||
|
||||
export interface OAuthApiKeyResult {
|
||||
newCredentials: Record<string, any>;
|
||||
apiKey: string;
|
||||
}
|
||||
|
||||
export type OAuthApiKeyResolver = (
|
||||
providerId: OAuthProviderId,
|
||||
credentials: Record<string, Record<string, any>>,
|
||||
) => Promise<OAuthApiKeyResult | null>;
|
||||
|
||||
export interface EnsureFreshAuthConfig {
|
||||
auth?: AuthData | null;
|
||||
authFile?: string;
|
||||
oauthResolver?: OAuthApiKeyResolver;
|
||||
nowMs?: number;
|
||||
persist?: boolean;
|
||||
}
|
||||
|
||||
export interface FreshAuthResult {
|
||||
auth: AuthData | null;
|
||||
changed: boolean;
|
||||
refreshErrors: Partial<Record<OAuthProviderId, string>>;
|
||||
}
|
||||
|
||||
export interface FetchAllUsagesConfig extends FetchConfig, EnsureFreshAuthConfig {
|
||||
auth?: AuthData | null;
|
||||
authFile?: string;
|
||||
}
|
||||
|
||||
const DEFAULT_FETCH_TIMEOUT_MS = 12_000;
|
||||
const TOKEN_REFRESH_SKEW_MS = 60_000;
|
||||
|
||||
export const DEFAULT_AUTH_FILE = path.join(os.homedir(), ".pi", "agent", "auth.json");
|
||||
export const DEFAULT_ZAI_USAGE_ENDPOINT = "https://api.z.ai/api/monitor/usage/quota/limit";
|
||||
export const GOOGLE_QUOTA_ENDPOINT = "https://cloudcode-pa.googleapis.com/v1internal:retrieveUserQuota";
|
||||
export const GOOGLE_LOAD_CODE_ASSIST_ENDPOINTS = [
|
||||
"https://cloudcode-pa.googleapis.com/v1internal:loadCodeAssist",
|
||||
"https://daily-cloudcode-pa.sandbox.googleapis.com/v1internal:loadCodeAssist",
|
||||
];
|
||||
|
||||
export function resolveUsageEndpoints(): UsageEndpoints {
|
||||
return {
|
||||
zai: DEFAULT_ZAI_USAGE_ENDPOINT,
|
||||
gemini: GOOGLE_QUOTA_ENDPOINT,
|
||||
antigravity: GOOGLE_QUOTA_ENDPOINT,
|
||||
googleLoadCodeAssistEndpoints: GOOGLE_LOAD_CODE_ASSIST_ENDPOINTS,
|
||||
};
|
||||
}
|
||||
|
||||
function toErrorMessage(error: unknown): string {
|
||||
if (error instanceof Error) {
|
||||
if (error.name === "AbortError") return "request timeout";
|
||||
return error.message || String(error);
|
||||
}
|
||||
return String(error);
|
||||
}
|
||||
|
||||
function asObject(value: unknown): Record<string, any> | null {
|
||||
if (!value || typeof value !== "object") return null;
|
||||
return value as Record<string, any>;
|
||||
}
|
||||
|
||||
function normalizeUsagePair(session: number, weekly: number): { session: number; weekly: number } {
|
||||
const clean = (v: number) => {
|
||||
if (!Number.isFinite(v)) return 0;
|
||||
return Number(v.toFixed(2));
|
||||
};
|
||||
return { session: clean(session), weekly: clean(weekly) };
|
||||
}
|
||||
|
||||
async function requestJson(url: string, init: RequestInit, config: RequestConfig = {}): Promise<{ ok: true; data: any } | { ok: false; error: string }> {
|
||||
const fetchFn = config.fetchFn ?? ((fetch as unknown) as FetchLike);
|
||||
const timeoutMs = config.timeoutMs ?? DEFAULT_FETCH_TIMEOUT_MS;
|
||||
const controller = new AbortController();
|
||||
const timeout = timeoutMs > 0 ? setTimeout(() => controller.abort(), timeoutMs) : null;
|
||||
|
||||
try {
|
||||
const response = await fetchFn(url, { ...init, signal: controller.signal });
|
||||
if (!response.ok) return { ok: false, error: `HTTP ${response.status}` };
|
||||
|
||||
try {
|
||||
const data = await response.json();
|
||||
return { ok: true, data };
|
||||
} catch {
|
||||
return { ok: false, error: "invalid JSON response" };
|
||||
}
|
||||
} catch (error) {
|
||||
return { ok: false, error: toErrorMessage(error) };
|
||||
} finally {
|
||||
if (timeout) clearTimeout(timeout);
|
||||
}
|
||||
}
|
||||
|
||||
export function formatDuration(seconds: number): string {
|
||||
if (!Number.isFinite(seconds) || seconds <= 0) return "now";
|
||||
const d = Math.floor(seconds / 86400);
|
||||
const h = Math.floor((seconds % 86400) / 3600);
|
||||
const m = Math.floor((seconds % 3600) / 60);
|
||||
if (d > 0 && h > 0) return `${d}d ${h}h`;
|
||||
if (d > 0) return `${d}d`;
|
||||
if (h > 0 && m > 0) return `${h}h ${m}m`;
|
||||
if (h > 0) return `${h}h`;
|
||||
if (m > 0) return `${m}m`;
|
||||
return "<1m";
|
||||
}
|
||||
|
||||
export function formatResetsAt(isoDate: string, nowMs = Date.now()): string {
|
||||
const resetTime = new Date(isoDate).getTime();
|
||||
if (!Number.isFinite(resetTime)) return "";
|
||||
const diffSeconds = Math.max(0, (resetTime - nowMs) / 1000);
|
||||
return formatDuration(diffSeconds);
|
||||
}
|
||||
|
||||
const CLAUDE_CREDENTIALS_FILE = path.join(os.homedir(), ".claude", ".credentials.json");
|
||||
|
||||
export function readAuth(authFile = DEFAULT_AUTH_FILE): AuthData | null {
|
||||
let result: AuthData | null = null;
|
||||
|
||||
// Read pi auth.json for non-Claude providers
|
||||
try {
|
||||
const parsed = JSON.parse(fs.readFileSync(authFile, "utf-8"));
|
||||
result = asObject(parsed) as AuthData;
|
||||
} catch {
|
||||
result = {} as AuthData;
|
||||
}
|
||||
|
||||
// Read Claude credentials from ~/.claude/.credentials.json
|
||||
try {
|
||||
const claudeRaw = fs.readFileSync(CLAUDE_CREDENTIALS_FILE, "utf-8");
|
||||
const claudeCreds = JSON.parse(claudeRaw);
|
||||
const oauth = claudeCreds?.claudeAiOauth;
|
||||
if (oauth?.accessToken) {
|
||||
result!.anthropic = {
|
||||
access: oauth.accessToken,
|
||||
refresh: oauth.refreshToken,
|
||||
expires: typeof oauth.expiresAt === "number" ? oauth.expiresAt : undefined,
|
||||
};
|
||||
}
|
||||
} catch {}
|
||||
|
||||
return result;
|
||||
}
|
||||
|
||||
export function writeAuth(auth: AuthData, authFile = DEFAULT_AUTH_FILE): boolean {
|
||||
try {
|
||||
const dir = path.dirname(authFile);
|
||||
if (!fs.existsSync(dir)) fs.mkdirSync(dir, { recursive: true });
|
||||
|
||||
const tmpPath = `${authFile}.tmp-${process.pid}-${Date.now()}`;
|
||||
fs.writeFileSync(tmpPath, JSON.stringify(auth, null, 2));
|
||||
fs.renameSync(tmpPath, authFile);
|
||||
return true;
|
||||
} catch {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
|
||||
let cachedOAuthResolver: OAuthApiKeyResolver | null = null;
|
||||
|
||||
async function getDefaultOAuthResolver(): Promise<OAuthApiKeyResolver> {
|
||||
if (cachedOAuthResolver) return cachedOAuthResolver;
|
||||
|
||||
const mod = await import("@mariozechner/pi-ai");
|
||||
if (typeof (mod as any).getOAuthApiKey !== "function") {
|
||||
throw new Error("oauth resolver unavailable");
|
||||
}
|
||||
|
||||
cachedOAuthResolver = (providerId, credentials) =>
|
||||
(mod as any).getOAuthApiKey(providerId, credentials) as Promise<OAuthApiKeyResult | null>;
|
||||
|
||||
return cachedOAuthResolver;
|
||||
}
|
||||
|
||||
function isCredentialExpired(creds: { expires?: number } | undefined, nowMs: number): boolean {
|
||||
if (!creds) return false;
|
||||
if (typeof creds.expires !== "number") return false;
|
||||
return nowMs + TOKEN_REFRESH_SKEW_MS >= creds.expires;
|
||||
}
|
||||
|
||||
export async function ensureFreshAuthForProviders(
|
||||
providerIds: OAuthProviderId[],
|
||||
config: EnsureFreshAuthConfig = {},
|
||||
): Promise<FreshAuthResult> {
|
||||
const authFile = config.authFile ?? DEFAULT_AUTH_FILE;
|
||||
const auth = config.auth ?? readAuth(authFile);
|
||||
if (!auth) {
|
||||
return { auth: null, changed: false, refreshErrors: {} };
|
||||
}
|
||||
|
||||
const nowMs = config.nowMs ?? Date.now();
|
||||
const uniqueProviders = Array.from(new Set(providerIds));
|
||||
const nextAuth: AuthData = { ...auth };
|
||||
const refreshErrors: Partial<Record<OAuthProviderId, string>> = {};
|
||||
|
||||
let changed = false;
|
||||
|
||||
for (const providerId of uniqueProviders) {
|
||||
const creds = (nextAuth as any)[providerId] as { access?: string; refresh?: string; expires?: number } | undefined;
|
||||
if (!creds?.refresh) continue;
|
||||
|
||||
const needsRefresh = !creds.access || isCredentialExpired(creds, nowMs);
|
||||
if (!needsRefresh) continue;
|
||||
|
||||
try {
|
||||
const resolver = config.oauthResolver ?? (await getDefaultOAuthResolver());
|
||||
const resolved = await resolver(providerId, nextAuth as any);
|
||||
if (!resolved?.newCredentials) {
|
||||
refreshErrors[providerId] = "missing OAuth credentials";
|
||||
continue;
|
||||
}
|
||||
|
||||
(nextAuth as any)[providerId] = {
|
||||
...(nextAuth as any)[providerId],
|
||||
...resolved.newCredentials,
|
||||
};
|
||||
changed = true;
|
||||
} catch (error) {
|
||||
refreshErrors[providerId] = toErrorMessage(error);
|
||||
}
|
||||
}
|
||||
|
||||
if (changed && config.persist !== false) {
|
||||
writeAuth(nextAuth, authFile);
|
||||
}
|
||||
|
||||
return { auth: nextAuth, changed, refreshErrors };
|
||||
}
|
||||
|
||||
export function readPercentCandidate(value: unknown): number | null {
|
||||
if (typeof value !== "number" || !Number.isFinite(value)) return null;
|
||||
|
||||
if (value >= 0 && value <= 1) {
|
||||
if (Number.isInteger(value)) return value;
|
||||
return value * 100;
|
||||
}
|
||||
|
||||
if (value >= 0 && value <= 100) return value;
|
||||
return null;
|
||||
}
|
||||
|
||||
export function readLimitPercent(limit: any): number | null {
|
||||
const direct = [
|
||||
limit?.percentage,
|
||||
limit?.utilization,
|
||||
limit?.used_percent,
|
||||
limit?.usedPercent,
|
||||
limit?.usagePercent,
|
||||
limit?.usage_percent,
|
||||
]
|
||||
.map(readPercentCandidate)
|
||||
.find((v) => v != null);
|
||||
|
||||
if (direct != null) return direct;
|
||||
|
||||
const current = typeof limit?.currentValue === "number" ? limit.currentValue : null;
|
||||
const remaining = typeof limit?.remaining === "number" ? limit.remaining : null;
|
||||
|
||||
if (current != null && remaining != null && current + remaining > 0) {
|
||||
return (current / (current + remaining)) * 100;
|
||||
}
|
||||
|
||||
return null;
|
||||
}
|
||||
|
||||
export function extractUsageFromPayload(data: any): { session: number; weekly: number } | null {
|
||||
const limitArrays = [data?.data?.limits, data?.limits, data?.quota?.limits, data?.data?.quota?.limits];
|
||||
const limits = limitArrays.find((arr) => Array.isArray(arr)) as any[] | undefined;
|
||||
|
||||
if (limits) {
|
||||
const byType = (types: string[]) =>
|
||||
limits.find((l) => {
|
||||
const t = String(l?.type || "").toUpperCase();
|
||||
return types.some((x) => t === x);
|
||||
});
|
||||
|
||||
const sessionLimit = byType(["TIME_LIMIT", "SESSION_LIMIT", "REQUEST_LIMIT", "RPM_LIMIT", "RPD_LIMIT"]);
|
||||
const weeklyLimit = byType(["TOKENS_LIMIT", "TOKEN_LIMIT", "WEEK_LIMIT", "WEEKLY_LIMIT", "TPM_LIMIT", "DAILY_LIMIT"]);
|
||||
|
||||
const s = readLimitPercent(sessionLimit);
|
||||
const w = readLimitPercent(weeklyLimit);
|
||||
if (s != null && w != null) return normalizeUsagePair(s, w);
|
||||
}
|
||||
|
||||
const sessionCandidates = [
|
||||
data?.session,
|
||||
data?.sessionPercent,
|
||||
data?.session_percent,
|
||||
data?.five_hour?.utilization,
|
||||
data?.rate_limit?.primary_window?.used_percent,
|
||||
data?.limits?.session?.utilization,
|
||||
data?.usage?.session,
|
||||
data?.data?.session,
|
||||
data?.data?.sessionPercent,
|
||||
data?.data?.session_percent,
|
||||
data?.data?.usage?.session,
|
||||
data?.quota?.session?.percentage,
|
||||
data?.data?.quota?.session?.percentage,
|
||||
];
|
||||
|
||||
const weeklyCandidates = [
|
||||
data?.weekly,
|
||||
data?.weeklyPercent,
|
||||
data?.weekly_percent,
|
||||
data?.seven_day?.utilization,
|
||||
data?.rate_limit?.secondary_window?.used_percent,
|
||||
data?.limits?.weekly?.utilization,
|
||||
data?.usage?.weekly,
|
||||
data?.data?.weekly,
|
||||
data?.data?.weeklyPercent,
|
||||
data?.data?.weekly_percent,
|
||||
data?.data?.usage?.weekly,
|
||||
data?.quota?.weekly?.percentage,
|
||||
data?.data?.quota?.weekly?.percentage,
|
||||
data?.quota?.daily?.percentage,
|
||||
data?.data?.quota?.daily?.percentage,
|
||||
];
|
||||
|
||||
const session = sessionCandidates.map(readPercentCandidate).find((v) => v != null);
|
||||
const weekly = weeklyCandidates.map(readPercentCandidate).find((v) => v != null);
|
||||
|
||||
if (session == null || weekly == null) return null;
|
||||
return normalizeUsagePair(session, weekly);
|
||||
}
|
||||
|
||||
export function googleMetadata(projectId?: string) {
|
||||
return {
|
||||
ideType: "IDE_UNSPECIFIED",
|
||||
platform: "PLATFORM_UNSPECIFIED",
|
||||
pluginType: "GEMINI",
|
||||
...(projectId ? { duetProject: projectId } : {}),
|
||||
};
|
||||
}
|
||||
|
||||
export function googleHeaders(token: string, projectId?: string) {
|
||||
return {
|
||||
Authorization: `Bearer ${token}`,
|
||||
"Content-Type": "application/json",
|
||||
"User-Agent": "google-cloud-sdk vscode_cloudshelleditor/0.1",
|
||||
"X-Goog-Api-Client": "gl-node/22.17.0",
|
||||
"Client-Metadata": JSON.stringify(googleMetadata(projectId)),
|
||||
};
|
||||
}
|
||||
|
||||
export async function discoverGoogleProjectId(token: string, config: FetchConfig = {}): Promise<string | undefined> {
|
||||
const env = config.env ?? process.env;
|
||||
const envProjectId = env.GOOGLE_CLOUD_PROJECT || env.GOOGLE_CLOUD_PROJECT_ID;
|
||||
if (envProjectId) return envProjectId;
|
||||
|
||||
const endpoints = config.endpoints ?? resolveUsageEndpoints();
|
||||
|
||||
for (const endpoint of endpoints.googleLoadCodeAssistEndpoints) {
|
||||
const result = await requestJson(
|
||||
endpoint,
|
||||
{
|
||||
method: "POST",
|
||||
headers: googleHeaders(token),
|
||||
body: JSON.stringify({ metadata: googleMetadata() }),
|
||||
},
|
||||
config,
|
||||
);
|
||||
|
||||
if (!result.ok) continue;
|
||||
|
||||
const data = result.data;
|
||||
if (typeof data?.cloudaicompanionProject === "string" && data.cloudaicompanionProject) {
|
||||
return data.cloudaicompanionProject;
|
||||
}
|
||||
if (data?.cloudaicompanionProject && typeof data.cloudaicompanionProject === "object") {
|
||||
const id = data.cloudaicompanionProject.id;
|
||||
if (typeof id === "string" && id) return id;
|
||||
}
|
||||
}
|
||||
|
||||
return undefined;
|
||||
}
|
||||
|
||||
export function usedPercentFromRemainingFraction(value: unknown): number | null {
|
||||
if (typeof value !== "number" || !Number.isFinite(value)) return null;
|
||||
const remaining = Math.max(0, Math.min(1, value));
|
||||
return (1 - remaining) * 100;
|
||||
}
|
||||
|
||||
export function pickMostUsedBucket(buckets: any[]): any | null {
|
||||
let best: any | null = null;
|
||||
let bestUsed = -1;
|
||||
for (const bucket of buckets) {
|
||||
const used = usedPercentFromRemainingFraction(bucket?.remainingFraction);
|
||||
if (used == null) continue;
|
||||
if (used > bestUsed) {
|
||||
bestUsed = used;
|
||||
best = bucket;
|
||||
}
|
||||
}
|
||||
return best;
|
||||
}
|
||||
|
||||
export function parseGoogleQuotaBuckets(data: any, provider: "gemini" | "antigravity"): { session: number; weekly: number } | null {
|
||||
const allBuckets = Array.isArray(data?.buckets) ? data.buckets : [];
|
||||
if (!allBuckets.length) return null;
|
||||
|
||||
const requestBuckets = allBuckets.filter((b: any) => String(b?.tokenType || "").toUpperCase() === "REQUESTS");
|
||||
const buckets = requestBuckets.length ? requestBuckets : allBuckets;
|
||||
|
||||
const modelId = (b: any) => String(b?.modelId || "").toLowerCase();
|
||||
const claudeNonThinking = buckets.filter((b: any) => modelId(b).includes("claude") && !modelId(b).includes("thinking"));
|
||||
const geminiPro = buckets.filter((b: any) => modelId(b).includes("gemini") && modelId(b).includes("pro"));
|
||||
const geminiFlash = buckets.filter((b: any) => modelId(b).includes("gemini") && modelId(b).includes("flash"));
|
||||
|
||||
const primaryBucket =
|
||||
provider === "antigravity"
|
||||
? pickMostUsedBucket(claudeNonThinking) || pickMostUsedBucket(geminiPro) || pickMostUsedBucket(geminiFlash) || pickMostUsedBucket(buckets)
|
||||
: pickMostUsedBucket(geminiPro) || pickMostUsedBucket(geminiFlash) || pickMostUsedBucket(buckets);
|
||||
|
||||
const secondaryBucket = pickMostUsedBucket(geminiFlash) || pickMostUsedBucket(geminiPro) || pickMostUsedBucket(buckets);
|
||||
|
||||
const session = usedPercentFromRemainingFraction(primaryBucket?.remainingFraction);
|
||||
const weekly = usedPercentFromRemainingFraction(secondaryBucket?.remainingFraction);
|
||||
|
||||
if (session == null || weekly == null) return null;
|
||||
return normalizeUsagePair(session, weekly);
|
||||
}
|
||||
|
||||
export async function fetchCodexUsage(token: string, config: RequestConfig = {}): Promise<UsageData> {
|
||||
const result = await requestJson(
|
||||
"https://chatgpt.com/backend-api/wham/usage",
|
||||
{ headers: { Authorization: `Bearer ${token}` } },
|
||||
config,
|
||||
);
|
||||
|
||||
if (!result.ok) return { session: 0, weekly: 0, error: result.error };
|
||||
|
||||
const primary = result.data?.rate_limit?.primary_window;
|
||||
const secondary = result.data?.rate_limit?.secondary_window;
|
||||
|
||||
return {
|
||||
session: readPercentCandidate(primary?.used_percent) ?? 0,
|
||||
weekly: readPercentCandidate(secondary?.used_percent) ?? 0,
|
||||
sessionResetsIn: typeof primary?.reset_after_seconds === "number" ? formatDuration(primary.reset_after_seconds) : undefined,
|
||||
weeklyResetsIn: typeof secondary?.reset_after_seconds === "number" ? formatDuration(secondary.reset_after_seconds) : undefined,
|
||||
};
|
||||
}
|
||||
|
||||
export async function fetchClaudeUsage(token: string, config: RequestConfig = {}): Promise<UsageData> {
|
||||
const result = await requestJson(
|
||||
"https://api.anthropic.com/api/oauth/usage",
|
||||
{
|
||||
headers: {
|
||||
Authorization: `Bearer ${token}`,
|
||||
"anthropic-beta": "oauth-2025-04-20",
|
||||
},
|
||||
},
|
||||
config,
|
||||
);
|
||||
|
||||
if (!result.ok) return { session: 0, weekly: 0, error: result.error };
|
||||
|
||||
const data = result.data;
|
||||
const sessionResetsAt = data?.five_hour?.resets_at
|
||||
? new Date(data.five_hour.resets_at).getTime()
|
||||
: undefined;
|
||||
const weeklyResetsAt = data?.seven_day?.resets_at
|
||||
? new Date(data.seven_day.resets_at).getTime()
|
||||
: undefined;
|
||||
|
||||
const usage: UsageData = {
|
||||
session: readPercentCandidate(data?.five_hour?.utilization) ?? 0,
|
||||
weekly: readPercentCandidate(data?.seven_day?.utilization) ?? 0,
|
||||
sessionResetsIn: data?.five_hour?.resets_at ? formatResetsAt(data.five_hour.resets_at) : undefined,
|
||||
sessionResetsAt: Number.isFinite(sessionResetsAt) ? sessionResetsAt : undefined,
|
||||
weeklyResetsIn: data?.seven_day?.resets_at ? formatResetsAt(data.seven_day.resets_at) : undefined,
|
||||
weeklyResetsAt: Number.isFinite(weeklyResetsAt) ? weeklyResetsAt : undefined,
|
||||
};
|
||||
|
||||
if (data?.extra_usage?.is_enabled) {
|
||||
usage.extraSpend = typeof data.extra_usage.used_credits === "number" ? data.extra_usage.used_credits : undefined;
|
||||
usage.extraLimit = typeof data.extra_usage.monthly_limit === "number" ? data.extra_usage.monthly_limit : undefined;
|
||||
}
|
||||
|
||||
return usage;
|
||||
}
|
||||
|
||||
export async function fetchZaiUsage(token: string, config: FetchConfig = {}): Promise<UsageData> {
|
||||
const endpoint = (config.endpoints ?? resolveUsageEndpoints()).zai;
|
||||
if (!endpoint) return { session: 0, weekly: 0, error: "usage endpoint unavailable" };
|
||||
|
||||
const result = await requestJson(
|
||||
endpoint,
|
||||
{ headers: { Authorization: `Bearer ${token}` } },
|
||||
config,
|
||||
);
|
||||
|
||||
if (!result.ok) return { session: 0, weekly: 0, error: result.error };
|
||||
|
||||
const parsed = extractUsageFromPayload(result.data);
|
||||
if (!parsed) return { session: 0, weekly: 0, error: "unrecognized response shape" };
|
||||
return parsed;
|
||||
}
|
||||
|
||||
export async function fetchGoogleUsage(
|
||||
token: string,
|
||||
endpoint: string,
|
||||
projectId: string | undefined,
|
||||
provider: "gemini" | "antigravity",
|
||||
config: FetchConfig = {},
|
||||
): Promise<UsageData> {
|
||||
if (!endpoint) return { session: 0, weekly: 0, error: "configure endpoint" };
|
||||
|
||||
const discoveredProjectId = projectId || (await discoverGoogleProjectId(token, config));
|
||||
if (!discoveredProjectId) {
|
||||
return { session: 0, weekly: 0, error: "missing projectId (try /login again)" };
|
||||
}
|
||||
|
||||
const result = await requestJson(
|
||||
endpoint,
|
||||
{
|
||||
method: "POST",
|
||||
headers: googleHeaders(token, discoveredProjectId),
|
||||
body: JSON.stringify({ project: discoveredProjectId }),
|
||||
},
|
||||
config,
|
||||
);
|
||||
|
||||
if (!result.ok) return { session: 0, weekly: 0, error: result.error };
|
||||
|
||||
const quota = parseGoogleQuotaBuckets(result.data, provider);
|
||||
if (quota) return quota;
|
||||
|
||||
const parsed = extractUsageFromPayload(result.data);
|
||||
if (!parsed) return { session: 0, weekly: 0, error: "unrecognized response shape" };
|
||||
return parsed;
|
||||
}
|
||||
|
||||
export function detectProvider(
|
||||
model: { provider?: string; id?: string; name?: string; api?: string } | string | undefined | null,
|
||||
): ProviderKey | null {
|
||||
if (!model) return null;
|
||||
if (typeof model === "string") return null;
|
||||
|
||||
const provider = (model.provider || "").toLowerCase();
|
||||
const id = (model.id || "").toLowerCase();
|
||||
|
||||
if (provider === "openai-codex") return "codex";
|
||||
if (provider === "anthropic") return "claude";
|
||||
if (provider === "zai") return "zai";
|
||||
if (provider === "google-gemini-cli") return "gemini";
|
||||
if (provider === "google-antigravity") return "antigravity";
|
||||
if (provider === "opencode-go" || id.startsWith("opencode-go/")) return "opencode-go";
|
||||
|
||||
return null;
|
||||
}
|
||||
|
||||
export function providerToOAuthProviderId(active: ProviderKey | null): OAuthProviderId | null {
|
||||
if (active === "codex") return "openai-codex";
|
||||
if (active === "claude") return "anthropic";
|
||||
if (active === "gemini") return "google-gemini-cli";
|
||||
if (active === "antigravity") return "google-antigravity";
|
||||
if (active === "opencode-go") return "opencode-go";
|
||||
return null;
|
||||
}
|
||||
|
||||
export function canShowForProvider(active: ProviderKey | null, auth: AuthData | null, endpoints: UsageEndpoints): boolean {
|
||||
if (!active || !auth) return false;
|
||||
if (active === "codex") return !!(auth["openai-codex"]?.access || auth["openai-codex"]?.refresh);
|
||||
if (active === "claude") return !!(auth.anthropic?.access || auth.anthropic?.refresh);
|
||||
if (active === "zai") return !!(auth.zai?.access || auth.zai?.key) && !!endpoints.zai;
|
||||
if (active === "gemini") {
|
||||
return !!(auth["google-gemini-cli"]?.access || auth["google-gemini-cli"]?.refresh) && !!endpoints.gemini;
|
||||
}
|
||||
if (active === "antigravity") {
|
||||
return !!(auth["google-antigravity"]?.access || auth["google-antigravity"]?.refresh) && !!endpoints.antigravity;
|
||||
}
|
||||
if (active === "opencode-go") {
|
||||
return !!(auth["opencode-go"]?.key || auth["opencode-go"]?.access);
|
||||
}
|
||||
return false;
|
||||
}
|
||||
|
||||
export function clampPercent(value: number): number {
|
||||
if (!Number.isFinite(value)) return 0;
|
||||
return Math.max(0, Math.min(100, Math.round(value)));
|
||||
}
|
||||
|
||||
export function colorForPercent(value: number): "success" | "warning" | "error" {
|
||||
if (value >= 90) return "error";
|
||||
if (value >= 70) return "warning";
|
||||
return "success";
|
||||
}
|
||||
|
||||
export async function fetchAllUsages(config: FetchAllUsagesConfig = {}): Promise<UsageByProvider> {
|
||||
const authFile = config.authFile ?? DEFAULT_AUTH_FILE;
|
||||
const auth = config.auth ?? readAuth(authFile);
|
||||
const endpoints = config.endpoints ?? resolveUsageEndpoints();
|
||||
|
||||
const results: UsageByProvider = {
|
||||
codex: null,
|
||||
claude: null,
|
||||
zai: null,
|
||||
gemini: null,
|
||||
antigravity: null,
|
||||
"opencode-go": null,
|
||||
};
|
||||
|
||||
if (!auth) return results;
|
||||
|
||||
const oauthProviders: OAuthProviderId[] = [
|
||||
"openai-codex",
|
||||
"anthropic",
|
||||
"google-gemini-cli",
|
||||
"google-antigravity",
|
||||
];
|
||||
|
||||
const refreshed = await ensureFreshAuthForProviders(oauthProviders, {
|
||||
...config,
|
||||
auth,
|
||||
authFile,
|
||||
});
|
||||
|
||||
const authData = refreshed.auth ?? auth;
|
||||
|
||||
const refreshError = (providerId: OAuthProviderId): string | null => {
|
||||
const error = refreshed.refreshErrors[providerId];
|
||||
return error ? `auth refresh failed (${error})` : null;
|
||||
};
|
||||
|
||||
const tasks: Promise<void>[] = [];
|
||||
const assign = (provider: ProviderKey, task: Promise<UsageData>) => {
|
||||
tasks.push(
|
||||
task
|
||||
.then((usage) => {
|
||||
results[provider] = usage;
|
||||
})
|
||||
.catch((error) => {
|
||||
results[provider] = { session: 0, weekly: 0, error: toErrorMessage(error) };
|
||||
}),
|
||||
);
|
||||
};
|
||||
|
||||
if (authData["openai-codex"]?.access) {
|
||||
const err = refreshError("openai-codex");
|
||||
if (err) results.codex = { session: 0, weekly: 0, error: err };
|
||||
else assign("codex", fetchCodexUsage(authData["openai-codex"].access, config));
|
||||
}
|
||||
|
||||
if (authData.anthropic?.access) {
|
||||
const err = refreshError("anthropic");
|
||||
if (err) results.claude = { session: 0, weekly: 0, error: err };
|
||||
else assign("claude", fetchClaudeUsage(authData.anthropic.access, config));
|
||||
}
|
||||
|
||||
if (authData.zai?.access || authData.zai?.key) {
|
||||
assign("zai", fetchZaiUsage(authData.zai.access || authData.zai.key!, { ...config, endpoints }));
|
||||
}
|
||||
|
||||
if (authData["google-gemini-cli"]?.access) {
|
||||
const err = refreshError("google-gemini-cli");
|
||||
if (err) {
|
||||
results.gemini = { session: 0, weekly: 0, error: err };
|
||||
} else {
|
||||
const creds = authData["google-gemini-cli"];
|
||||
assign(
|
||||
"gemini",
|
||||
fetchGoogleUsage(creds.access!, endpoints.gemini, creds.projectId, "gemini", { ...config, endpoints }),
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
if (authData["google-antigravity"]?.access) {
|
||||
const err = refreshError("google-antigravity");
|
||||
if (err) {
|
||||
results.antigravity = { session: 0, weekly: 0, error: err };
|
||||
} else {
|
||||
const creds = authData["google-antigravity"];
|
||||
assign(
|
||||
"antigravity",
|
||||
fetchGoogleUsage(creds.access!, endpoints.antigravity, creds.projectId, "antigravity", { ...config, endpoints }),
|
||||
);
|
||||
}
|
||||
}
|
||||
|
||||
// OpenCode Go uses local tracking (no public API yet)
|
||||
if (authData["opencode-go"]?.key || authData["opencode-go"]?.access) {
|
||||
results["opencode-go"] = getOpenCodeGoUsageData();
|
||||
}
|
||||
|
||||
await Promise.all(tasks);
|
||||
return results;
|
||||
}
|
||||
@@ -1,656 +0,0 @@
|
||||
/**
|
||||
* Usage Extension - Minimal API usage indicator for pi
|
||||
*
|
||||
* Polls Codex, Anthropic, Z.AI, Gemini CLI / Antigravity usage and exposes it
|
||||
* via two channels:
|
||||
* • pi.events "usage:update" — for other extensions (e.g. footer-display)
|
||||
* • ctx.ui.setStatus("usage-bars", …) — formatted S/W braille bars
|
||||
*
|
||||
* Rendering / footer layout is handled by the separate footer-display extension.
|
||||
*/
|
||||
|
||||
import { DynamicBorder, type ExtensionAPI } from "@mariozechner/pi-coding-agent";
|
||||
import {
|
||||
Container,
|
||||
Input,
|
||||
Spacer,
|
||||
Text,
|
||||
getEditorKeybindings,
|
||||
type Focusable,
|
||||
} from "@mariozechner/pi-tui";
|
||||
import {
|
||||
canShowForProvider,
|
||||
clampPercent,
|
||||
colorForPercent,
|
||||
detectProvider,
|
||||
ensureFreshAuthForProviders,
|
||||
fetchAllUsages,
|
||||
fetchClaudeUsage,
|
||||
fetchCodexUsage,
|
||||
fetchGoogleUsage,
|
||||
fetchZaiUsage,
|
||||
providerToOAuthProviderId,
|
||||
readAuth,
|
||||
readUsageCache,
|
||||
resolveUsageEndpoints,
|
||||
writeUsageCache,
|
||||
type OAuthProviderId,
|
||||
type ProviderKey,
|
||||
type UsageByProvider,
|
||||
type UsageData,
|
||||
} from "./core";
|
||||
|
||||
const CACHE_TTL_MS = 15 * 60 * 1000;
|
||||
const ACTIVE_CACHE_TTL_MS = 3 * 60 * 1000;
|
||||
const STREAMING_POLL_INTERVAL_MS = 2 * 60 * 1000;
|
||||
const RATE_LIMITED_BACKOFF_MS = 60 * 60 * 1000;
|
||||
|
||||
const STATUS_KEY = "usage-bars";
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Braille gradient bar (⣀ ⣄ ⣤ ⣦ ⣶ ⣷ ⣿)
|
||||
// ---------------------------------------------------------------------------
|
||||
const BRAILLE_GRADIENT = "\u28C0\u28C4\u28E4\u28E6\u28F6\u28F7\u28FF";
|
||||
const BRAILLE_EMPTY = "\u28C0";
|
||||
const BAR_WIDTH = 5;
|
||||
|
||||
function renderBrailleBar(theme: any, value: number, width = BAR_WIDTH): string {
|
||||
const v = clampPercent(value);
|
||||
const levels = BRAILLE_GRADIENT.length - 1;
|
||||
const totalSteps = width * levels;
|
||||
const filledSteps = Math.round((v / 100) * totalSteps);
|
||||
const full = Math.floor(filledSteps / levels);
|
||||
const partial = filledSteps % levels;
|
||||
const empty = width - full - (partial ? 1 : 0);
|
||||
const color = colorForPercent(v);
|
||||
const filled = BRAILLE_GRADIENT[BRAILLE_GRADIENT.length - 1]!.repeat(Math.max(0, full));
|
||||
const partialChar = partial ? BRAILLE_GRADIENT[partial]! : "";
|
||||
const emptyChars = BRAILLE_EMPTY.repeat(Math.max(0, empty));
|
||||
return theme.fg(color, filled + partialChar) + theme.fg("dim", emptyChars);
|
||||
}
|
||||
|
||||
function renderBrailleBarWide(theme: any, value: number): string {
|
||||
return renderBrailleBar(theme, value, 12);
|
||||
}
|
||||
|
||||
const PROVIDER_LABELS: Record<ProviderKey, string> = {
|
||||
codex: "Codex",
|
||||
claude: "Claude",
|
||||
zai: "Z.AI",
|
||||
gemini: "Gemini",
|
||||
antigravity: "Antigravity",
|
||||
"opencode-go": "Go",
|
||||
};
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// /usage command popup
|
||||
// ---------------------------------------------------------------------------
|
||||
interface SubscriptionItem {
|
||||
name: string;
|
||||
provider: ProviderKey;
|
||||
data: UsageData | null;
|
||||
isActive: boolean;
|
||||
}
|
||||
|
||||
class UsageSelectorComponent extends Container implements Focusable {
|
||||
private searchInput: Input;
|
||||
private listContainer: Container;
|
||||
private hintText: Text;
|
||||
private tui: any;
|
||||
private theme: any;
|
||||
private onCancelCallback: () => void;
|
||||
private allItems: SubscriptionItem[] = [];
|
||||
private filteredItems: SubscriptionItem[] = [];
|
||||
private selectedIndex = 0;
|
||||
private loading = true;
|
||||
private activeProvider: ProviderKey | null;
|
||||
private fetchAllFn: () => Promise<UsageByProvider>;
|
||||
private _focused = false;
|
||||
|
||||
get focused(): boolean { return this._focused; }
|
||||
set focused(value: boolean) { this._focused = value; this.searchInput.focused = value; }
|
||||
|
||||
constructor(
|
||||
tui: any,
|
||||
theme: any,
|
||||
activeProvider: ProviderKey | null,
|
||||
fetchAll: () => Promise<UsageByProvider>,
|
||||
onCancel: () => void,
|
||||
) {
|
||||
super();
|
||||
this.tui = tui;
|
||||
this.theme = theme;
|
||||
this.activeProvider = activeProvider;
|
||||
this.fetchAllFn = fetchAll;
|
||||
this.onCancelCallback = onCancel;
|
||||
|
||||
this.addChild(new DynamicBorder((s: string) => theme.fg("accent", s)));
|
||||
this.addChild(new Spacer(1));
|
||||
this.hintText = new Text(theme.fg("dim", "Fetching usage from all providers…"), 0, 0);
|
||||
this.addChild(this.hintText);
|
||||
this.addChild(new Spacer(1));
|
||||
this.searchInput = new Input();
|
||||
this.addChild(this.searchInput);
|
||||
this.addChild(new Spacer(1));
|
||||
this.listContainer = new Container();
|
||||
this.addChild(this.listContainer);
|
||||
this.addChild(new Spacer(1));
|
||||
this.addChild(new DynamicBorder((s: string) => theme.fg("accent", s)));
|
||||
|
||||
this.fetchAllFn()
|
||||
.then((results) => {
|
||||
this.loading = false;
|
||||
this.buildItems(results);
|
||||
this.updateList();
|
||||
this.hintText.setText(
|
||||
theme.fg("muted", "Only showing providers with credentials. ") +
|
||||
theme.fg("dim", "✓ = active provider"),
|
||||
);
|
||||
this.tui.requestRender();
|
||||
})
|
||||
.catch(() => {
|
||||
this.loading = false;
|
||||
this.hintText.setText(theme.fg("error", "Failed to fetch usage data"));
|
||||
this.tui.requestRender();
|
||||
});
|
||||
|
||||
this.updateList();
|
||||
}
|
||||
|
||||
private buildItems(results: UsageByProvider) {
|
||||
const providers: Array<{ key: ProviderKey; name: string }> = [
|
||||
{ key: "codex", name: "Codex" },
|
||||
{ key: "claude", name: "Claude" },
|
||||
{ key: "zai", name: "Z.AI" },
|
||||
{ key: "gemini", name: "Gemini" },
|
||||
{ key: "antigravity", name: "Antigravity" },
|
||||
{ key: "opencode-go", name: "Go" },
|
||||
];
|
||||
this.allItems = [];
|
||||
for (const p of providers) {
|
||||
if (results[p.key] !== null) {
|
||||
this.allItems.push({
|
||||
name: p.name,
|
||||
provider: p.key,
|
||||
data: results[p.key],
|
||||
isActive: this.activeProvider === p.key,
|
||||
});
|
||||
}
|
||||
}
|
||||
this.filteredItems = this.allItems;
|
||||
this.selectedIndex = Math.min(this.selectedIndex, Math.max(0, this.filteredItems.length - 1));
|
||||
}
|
||||
|
||||
private filterItems(query: string) {
|
||||
if (!query) {
|
||||
this.filteredItems = this.allItems;
|
||||
} else {
|
||||
const q = query.toLowerCase();
|
||||
this.filteredItems = this.allItems.filter(
|
||||
(item) => item.name.toLowerCase().includes(q) || item.provider.toLowerCase().includes(q),
|
||||
);
|
||||
}
|
||||
this.selectedIndex = Math.min(this.selectedIndex, Math.max(0, this.filteredItems.length - 1));
|
||||
}
|
||||
|
||||
private renderItem(item: SubscriptionItem, isSelected: boolean) {
|
||||
const t = this.theme;
|
||||
const pointer = isSelected ? t.fg("accent", "→ ") : " ";
|
||||
const activeBadge = item.isActive ? t.fg("success", " ✓") : "";
|
||||
const name = isSelected ? t.fg("accent", t.bold(item.name)) : item.name;
|
||||
this.listContainer.addChild(new Text(`${pointer}${name}${activeBadge}`, 0, 0));
|
||||
const indent = " ";
|
||||
|
||||
if (!item.data) {
|
||||
this.listContainer.addChild(new Text(indent + t.fg("dim", "No credentials"), 0, 0));
|
||||
} else if (item.data.error) {
|
||||
this.listContainer.addChild(new Text(indent + t.fg("error", item.data.error), 0, 0));
|
||||
} else {
|
||||
const session = clampPercent(item.data.session);
|
||||
const weekly = clampPercent(item.data.weekly);
|
||||
const sessionReset = item.data.sessionResetsIn
|
||||
? t.fg("dim", ` resets in ${item.data.sessionResetsIn}`) : "";
|
||||
const weeklyReset = item.data.weeklyResetsIn
|
||||
? t.fg("dim", ` resets in ${item.data.weeklyResetsIn}`) : "";
|
||||
|
||||
this.listContainer.addChild(new Text(
|
||||
indent + t.fg("muted", "Session ") +
|
||||
renderBrailleBarWide(t, session) + " " +
|
||||
t.fg(colorForPercent(session), `${session}%`.padStart(4)) + sessionReset,
|
||||
0, 0,
|
||||
));
|
||||
this.listContainer.addChild(new Text(
|
||||
indent + t.fg("muted", "Weekly ") +
|
||||
renderBrailleBarWide(t, weekly) + " " +
|
||||
t.fg(colorForPercent(weekly), `${weekly}%`.padStart(4)) + weeklyReset,
|
||||
0, 0,
|
||||
));
|
||||
|
||||
if (typeof item.data.extraSpend === "number" && typeof item.data.extraLimit === "number") {
|
||||
this.listContainer.addChild(new Text(
|
||||
indent + t.fg("muted", "Extra ") +
|
||||
t.fg("dim", `$${item.data.extraSpend.toFixed(2)} / $${item.data.extraLimit}`),
|
||||
0, 0,
|
||||
));
|
||||
}
|
||||
}
|
||||
this.listContainer.addChild(new Spacer(1));
|
||||
}
|
||||
|
||||
private updateList() {
|
||||
this.listContainer.clear();
|
||||
if (this.loading) {
|
||||
this.listContainer.addChild(new Text(this.theme.fg("muted", " Loading…"), 0, 0));
|
||||
return;
|
||||
}
|
||||
if (this.filteredItems.length === 0) {
|
||||
this.listContainer.addChild(new Text(this.theme.fg("muted", " No matching providers"), 0, 0));
|
||||
return;
|
||||
}
|
||||
for (let i = 0; i < this.filteredItems.length; i++) {
|
||||
this.renderItem(this.filteredItems[i]!, i === this.selectedIndex);
|
||||
}
|
||||
}
|
||||
|
||||
handleInput(keyData: string): void {
|
||||
const kb = getEditorKeybindings();
|
||||
if (kb.matches(keyData, "selectUp")) {
|
||||
if (this.filteredItems.length === 0) return;
|
||||
this.selectedIndex = this.selectedIndex === 0 ? this.filteredItems.length - 1 : this.selectedIndex - 1;
|
||||
this.updateList(); return;
|
||||
}
|
||||
if (kb.matches(keyData, "selectDown")) {
|
||||
if (this.filteredItems.length === 0) return;
|
||||
this.selectedIndex = this.selectedIndex === this.filteredItems.length - 1 ? 0 : this.selectedIndex + 1;
|
||||
this.updateList(); return;
|
||||
}
|
||||
if (kb.matches(keyData, "selectCancel") || kb.matches(keyData, "selectConfirm")) {
|
||||
this.onCancelCallback(); return;
|
||||
}
|
||||
this.searchInput.handleInput(keyData);
|
||||
this.filterItems(this.searchInput.getValue());
|
||||
this.updateList();
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Extension state
|
||||
// ---------------------------------------------------------------------------
|
||||
interface UsageState extends UsageByProvider {
|
||||
lastPoll: number;
|
||||
activeProvider: ProviderKey | null;
|
||||
}
|
||||
|
||||
interface PollOptions {
|
||||
cacheTtl?: number;
|
||||
forceFresh?: boolean;
|
||||
}
|
||||
|
||||
export default function (pi: ExtensionAPI) {
|
||||
const endpoints = resolveUsageEndpoints();
|
||||
const state: UsageState = {
|
||||
codex: null, claude: null, zai: null, gemini: null, antigravity: null, "opencode-go": null,
|
||||
lastPoll: 0, activeProvider: null,
|
||||
};
|
||||
|
||||
let pollInFlight: Promise<void> | null = null;
|
||||
let pollQueued = false;
|
||||
let pollStartedAt = 0;
|
||||
let streamingTimer: ReturnType<typeof setInterval> | null = null;
|
||||
let ctx: any = null;
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Status update
|
||||
// ---------------------------------------------------------------------------
|
||||
function updateStatus() {
|
||||
const active = state.activeProvider;
|
||||
const data = active ? state[active] : null;
|
||||
|
||||
// Always emit Claude usage for other extensions (e.g. footer-display)
|
||||
// so S/W bars are visible regardless of active model.
|
||||
const claudeData = state.claude;
|
||||
if (claudeData && !claudeData.error) {
|
||||
pi.events.emit("usage:update", {
|
||||
session: claudeData.session,
|
||||
weekly: claudeData.weekly,
|
||||
sessionResetsIn: claudeData.sessionResetsIn,
|
||||
sessionResetsAt: claudeData.sessionResetsAt,
|
||||
weeklyResetsIn: claudeData.weeklyResetsIn,
|
||||
weeklyResetsAt: claudeData.weeklyResetsAt,
|
||||
});
|
||||
} else if (data && !data.error) {
|
||||
// Fallback to active provider data if Claude data unavailable
|
||||
pi.events.emit("usage:update", {
|
||||
session: data.session,
|
||||
weekly: data.weekly,
|
||||
sessionResetsIn: data.sessionResetsIn,
|
||||
sessionResetsAt: data.sessionResetsAt,
|
||||
weeklyResetsIn: data.weeklyResetsIn,
|
||||
weeklyResetsAt: data.weeklyResetsAt,
|
||||
});
|
||||
}
|
||||
|
||||
if (!ctx?.hasUI) return;
|
||||
|
||||
const theme = ctx.ui.theme;
|
||||
|
||||
if (!active) {
|
||||
ctx.ui.setStatus(STATUS_KEY, undefined);
|
||||
return;
|
||||
}
|
||||
|
||||
const auth = readAuth();
|
||||
if (!canShowForProvider(active, auth, endpoints)) {
|
||||
ctx.ui.setStatus(STATUS_KEY, undefined);
|
||||
return;
|
||||
}
|
||||
|
||||
if (!data) {
|
||||
ctx.ui.setStatus(STATUS_KEY, theme.fg("dim", "loading\u2026"));
|
||||
return;
|
||||
}
|
||||
|
||||
if (data.error) {
|
||||
const cache = readUsageCache();
|
||||
const blockedUntil = cache?.rateLimitedUntil?.[active] ?? 0;
|
||||
const note = blockedUntil > Date.now()
|
||||
? ` \u2014 retry in ${Math.ceil((blockedUntil - Date.now()) / 60000)}m` : "";
|
||||
ctx.ui.setStatus(STATUS_KEY, theme.fg("warning", `${PROVIDER_LABELS[active]} unavailable${note}`));
|
||||
return;
|
||||
}
|
||||
|
||||
const session = clampPercent(data.session);
|
||||
const weekly = clampPercent(data.weekly);
|
||||
|
||||
// Time suffixes are intentionally omitted here — footer-display builds
|
||||
// them dynamically from the sessionResetsAt/weeklyResetsAt timestamps
|
||||
// emitted via the "usage:update" event, avoiding double-display.
|
||||
const s = theme.fg("muted", "S ") + renderBrailleBar(theme, session) + " " + theme.fg("dim", `${session}%`);
|
||||
const w = theme.fg("muted", "W ") + renderBrailleBar(theme, weekly) + " " + theme.fg("dim", `${weekly}%`);
|
||||
|
||||
ctx.ui.setStatus(STATUS_KEY, s + theme.fg("dim", " | ") + w);
|
||||
}
|
||||
|
||||
function updateProviderFrom(modelLike: any): boolean {
|
||||
const previous = state.activeProvider;
|
||||
state.activeProvider = detectProvider(modelLike);
|
||||
if (previous !== state.activeProvider) { updateStatus(); return true; }
|
||||
return false;
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Polling
|
||||
// ---------------------------------------------------------------------------
|
||||
async function runPollInner(options: PollOptions = {}) {
|
||||
const auth = readAuth();
|
||||
const active = state.activeProvider;
|
||||
|
||||
// Always try to fetch Claude data so S/W bars show regardless of active provider
|
||||
if (auth && canShowForProvider("claude", auth, endpoints)) {
|
||||
try {
|
||||
const cache = readUsageCache();
|
||||
const now = Date.now();
|
||||
const cacheTtl = options.cacheTtl ?? CACHE_TTL_MS;
|
||||
const claudeBlockedUntil = cache?.rateLimitedUntil?.claude ?? 0;
|
||||
if (now < claudeBlockedUntil) {
|
||||
if (cache?.data?.claude) state.claude = cache.data.claude;
|
||||
} else {
|
||||
const claudeCacheFresh = cache && now - cache.timestamp < cacheTtl && cache.data?.claude;
|
||||
if (claudeCacheFresh && !options.forceFresh) {
|
||||
state.claude = cache.data.claude;
|
||||
} else {
|
||||
const claudeAccess = auth.anthropic?.access;
|
||||
if (claudeAccess) {
|
||||
const claudeResult = await fetchClaudeUsage(claudeAccess);
|
||||
state.claude = claudeResult;
|
||||
if (!claudeResult.error) {
|
||||
const nextCache: import("./core").UsageCache = {
|
||||
timestamp: now,
|
||||
data: { ...(cache?.data ?? {}), claude: claudeResult },
|
||||
rateLimitedUntil: { ...(cache?.rateLimitedUntil ?? {}) },
|
||||
};
|
||||
delete nextCache.rateLimitedUntil!.claude;
|
||||
writeUsageCache(nextCache);
|
||||
} else if (claudeResult.error === "HTTP 429") {
|
||||
// Record backoff even when Claude is not the active provider —
|
||||
// without this the prefetch would hammer the API on every poll.
|
||||
const nextCache: import("./core").UsageCache = {
|
||||
timestamp: cache?.timestamp ?? now,
|
||||
data: { ...(cache?.data ?? {}) },
|
||||
rateLimitedUntil: { ...(cache?.rateLimitedUntil ?? {}), claude: now + RATE_LIMITED_BACKOFF_MS },
|
||||
};
|
||||
writeUsageCache(nextCache);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
} catch {}
|
||||
}
|
||||
|
||||
if (!canShowForProvider(active, auth, endpoints) || !auth || !active) {
|
||||
state.lastPoll = Date.now(); updateStatus(); return;
|
||||
}
|
||||
|
||||
const cache = readUsageCache();
|
||||
const now = Date.now();
|
||||
const cacheTtl = options.cacheTtl ?? CACHE_TTL_MS;
|
||||
|
||||
const blockedUntil = cache?.rateLimitedUntil?.[active] ?? 0;
|
||||
if (now < blockedUntil) {
|
||||
if (cache?.data?.[active]) {
|
||||
state[active] = cache.data[active]!;
|
||||
} else {
|
||||
// Rate-limited but no cached data — show a meaningful status instead
|
||||
// of leaving state null (which shows eternal "loading…").
|
||||
const retryMin = Math.ceil((blockedUntil - now) / 60000);
|
||||
state[active] = { session: 0, weekly: 0, error: `rate limited (retry in ${retryMin}m)` };
|
||||
}
|
||||
state.lastPoll = now; updateStatus(); return;
|
||||
}
|
||||
|
||||
if (!options.forceFresh && cache && now - cache.timestamp < cacheTtl && cache.data?.[active]) {
|
||||
state[active] = cache.data[active]!;
|
||||
state.lastPoll = now; updateStatus(); return;
|
||||
}
|
||||
|
||||
const oauthId = providerToOAuthProviderId(active);
|
||||
let effectiveAuth = auth;
|
||||
if (oauthId && active !== "zai") {
|
||||
const creds = auth[oauthId as keyof typeof auth] as
|
||||
| { access?: string; refresh?: string; expires?: number } | undefined;
|
||||
const expires = typeof creds?.expires === "number" ? creds.expires : 0;
|
||||
const tokenExpiredOrMissing = !creds?.access || (expires > 0 && Date.now() + 60_000 >= expires);
|
||||
if (tokenExpiredOrMissing && creds?.refresh) {
|
||||
try {
|
||||
const refreshPromise = ensureFreshAuthForProviders([oauthId as OAuthProviderId], { auth, persist: true });
|
||||
const timeoutPromise = new Promise<never>((_, reject) =>
|
||||
setTimeout(() => reject(new Error("OAuth refresh timeout")), 15_000),
|
||||
);
|
||||
const refreshed = await Promise.race([refreshPromise, timeoutPromise]);
|
||||
if (refreshed.auth) effectiveAuth = refreshed.auth;
|
||||
} catch {}
|
||||
}
|
||||
}
|
||||
|
||||
let result: UsageData;
|
||||
if (active === "codex") {
|
||||
const access = effectiveAuth["openai-codex"]?.access;
|
||||
result = access ? await fetchCodexUsage(access)
|
||||
: { session: 0, weekly: 0, error: "missing access token (try /login again)" };
|
||||
} else if (active === "claude") {
|
||||
const access = effectiveAuth.anthropic?.access;
|
||||
result = access ? await fetchClaudeUsage(access)
|
||||
: { session: 0, weekly: 0, error: "missing access token (try /login again)" };
|
||||
} else if (active === "zai") {
|
||||
const token = effectiveAuth.zai?.access || effectiveAuth.zai?.key;
|
||||
result = token ? await fetchZaiUsage(token, { endpoints })
|
||||
: { session: 0, weekly: 0, error: "missing token (try /login again)" };
|
||||
} else if (active === "gemini") {
|
||||
const creds = effectiveAuth["google-gemini-cli"];
|
||||
result = creds?.access
|
||||
? await fetchGoogleUsage(creds.access, endpoints.gemini, creds.projectId, "gemini", { endpoints })
|
||||
: { session: 0, weekly: 0, error: "missing access token (try /login again)" };
|
||||
} else if (active === "opencode-go") {
|
||||
// OpenCode Go uses local tracking (no public usage API yet)
|
||||
const { getOpenCodeGoUsageData } = await import("./core");
|
||||
result = getOpenCodeGoUsageData();
|
||||
} else {
|
||||
const creds = effectiveAuth["google-antigravity"];
|
||||
result = creds?.access
|
||||
? await fetchGoogleUsage(creds.access, endpoints.antigravity, creds.projectId, "antigravity", { endpoints })
|
||||
: { session: 0, weekly: 0, error: "missing access token (try /login again)" };
|
||||
}
|
||||
|
||||
state[active] = result;
|
||||
|
||||
if (result.error) {
|
||||
if (result.error === "HTTP 429") {
|
||||
const nextCache: import("./core").UsageCache = {
|
||||
timestamp: cache?.timestamp ?? now,
|
||||
data: { ...(cache?.data ?? {}) },
|
||||
rateLimitedUntil: { ...(cache?.rateLimitedUntil ?? {}), [active]: now + RATE_LIMITED_BACKOFF_MS },
|
||||
};
|
||||
writeUsageCache(nextCache);
|
||||
}
|
||||
} else {
|
||||
const nextCache: import("./core").UsageCache = {
|
||||
timestamp: now,
|
||||
data: { ...(cache?.data ?? {}), [active]: result },
|
||||
rateLimitedUntil: { ...(cache?.rateLimitedUntil ?? {}) },
|
||||
};
|
||||
delete nextCache.rateLimitedUntil![active];
|
||||
writeUsageCache(nextCache);
|
||||
}
|
||||
|
||||
state.lastPoll = now;
|
||||
updateStatus();
|
||||
}
|
||||
|
||||
async function runPoll(options: PollOptions = {}): Promise<void> {
|
||||
const timeout = new Promise<never>((_, reject) =>
|
||||
setTimeout(() => reject(new Error("runPoll timeout")), 25_000),
|
||||
);
|
||||
await Promise.race([runPollInner(options), timeout]);
|
||||
}
|
||||
|
||||
// Must be less than the 25 000 ms timeout inside runPoll so the guard fires
|
||||
// before runPoll's finally-block clears pollInFlight.
|
||||
const POLL_TIMEOUT_MS = 20_000;
|
||||
|
||||
async function poll(options: PollOptions = {}) {
|
||||
// If a previous poll has been running longer than POLL_TIMEOUT_MS, abandon it
|
||||
// so we don't queue forever behind a stuck request.
|
||||
if (pollInFlight && pollStartedAt > 0 && Date.now() - pollStartedAt > POLL_TIMEOUT_MS) {
|
||||
pollInFlight = null;
|
||||
pollQueued = false;
|
||||
const active = state.activeProvider;
|
||||
if (active && !state[active]) {
|
||||
state[active] = { session: 0, weekly: 0, error: "poll timeout" };
|
||||
updateStatus();
|
||||
}
|
||||
}
|
||||
|
||||
if (pollInFlight) { pollQueued = true; await pollInFlight; return; }
|
||||
do {
|
||||
pollQueued = false;
|
||||
pollStartedAt = Date.now();
|
||||
pollInFlight = runPoll(options).catch(() => {
|
||||
// If runPoll threw, ensure we don't leave status stuck at "loading…"
|
||||
const active = state.activeProvider;
|
||||
if (active && !state[active]) {
|
||||
state[active] = { session: 0, weekly: 0, error: "poll failed" };
|
||||
updateStatus();
|
||||
}
|
||||
}).finally(() => { pollInFlight = null; pollStartedAt = 0; });
|
||||
await pollInFlight;
|
||||
} while (pollQueued);
|
||||
}
|
||||
|
||||
function startStreamingTimer() {
|
||||
if (streamingTimer !== null) return;
|
||||
streamingTimer = setInterval(() => { void poll({ cacheTtl: ACTIVE_CACHE_TTL_MS }); }, STREAMING_POLL_INTERVAL_MS);
|
||||
}
|
||||
|
||||
function stopStreamingTimer() {
|
||||
if (streamingTimer !== null) { clearInterval(streamingTimer); streamingTimer = null; }
|
||||
}
|
||||
|
||||
// ── Lifecycle ────────────────────────────────────────────────────────────
|
||||
|
||||
pi.on("session_start", async (_event, _ctx) => {
|
||||
ctx = _ctx;
|
||||
updateProviderFrom(_ctx.model);
|
||||
await poll();
|
||||
});
|
||||
|
||||
pi.on("session_shutdown", async (_event, _ctx) => {
|
||||
stopStreamingTimer();
|
||||
if (_ctx?.hasUI) _ctx.ui.setStatus(STATUS_KEY, undefined);
|
||||
});
|
||||
|
||||
pi.on("model_select", async (event, _ctx) => {
|
||||
ctx = _ctx;
|
||||
const changed = updateProviderFrom(event.model ?? _ctx.model);
|
||||
if (changed) await poll();
|
||||
});
|
||||
|
||||
pi.on("turn_start", (_event, _ctx) => { ctx = _ctx; updateProviderFrom(_ctx.model); });
|
||||
|
||||
pi.on("turn_end", async (_event, _ctx) => {
|
||||
ctx = _ctx;
|
||||
await poll({ cacheTtl: ACTIVE_CACHE_TTL_MS });
|
||||
});
|
||||
|
||||
pi.on("before_agent_start", async (_event, _ctx) => {
|
||||
ctx = _ctx;
|
||||
await poll({ cacheTtl: ACTIVE_CACHE_TTL_MS });
|
||||
});
|
||||
|
||||
pi.on("agent_start", (_event, _ctx) => { ctx = _ctx; startStreamingTimer(); });
|
||||
|
||||
pi.on("agent_end", async (_event, _ctx) => {
|
||||
ctx = _ctx;
|
||||
stopStreamingTimer();
|
||||
await poll({ cacheTtl: ACTIVE_CACHE_TTL_MS });
|
||||
});
|
||||
|
||||
|
||||
// Listen for OpenCode Go spend events from other extensions
|
||||
pi.events.on("opencode-go:spend", async (amount: number) => {
|
||||
if (typeof amount === "number" && amount > 0) {
|
||||
const { addOpenCodeGoSpend } = await import("./core");
|
||||
addOpenCodeGoSpend(amount);
|
||||
// Invalidate cache and re-poll
|
||||
const cache = readUsageCache();
|
||||
if (cache?.data?.["opencode-go"]) {
|
||||
const nextCache: import("./core").UsageCache = { ...cache, data: { ...cache.data } };
|
||||
delete nextCache.data["opencode-go"];
|
||||
writeUsageCache(nextCache);
|
||||
}
|
||||
void poll({ forceFresh: true });
|
||||
}
|
||||
});
|
||||
|
||||
// ── /usage command ───────────────────────────────────────────────────────
|
||||
|
||||
pi.registerCommand("usage", {
|
||||
description: "Show API usage for all subscriptions",
|
||||
handler: async (_args, _ctx) => {
|
||||
ctx = _ctx;
|
||||
updateProviderFrom(_ctx.model);
|
||||
try {
|
||||
if (_ctx?.hasUI) {
|
||||
await _ctx.ui.custom<void>((tui, theme, _keybindings, done) => {
|
||||
return new UsageSelectorComponent(
|
||||
tui, theme, state.activeProvider,
|
||||
() => fetchAllUsages({ endpoints }),
|
||||
() => done(),
|
||||
);
|
||||
});
|
||||
}
|
||||
} finally {
|
||||
await poll({ cacheTtl: ACTIVE_CACHE_TTL_MS });
|
||||
}
|
||||
},
|
||||
});
|
||||
}
|
||||
@@ -1,581 +0,0 @@
|
||||
/**
|
||||
* Usage Extension - Minimal API usage indicator for pi
|
||||
*
|
||||
* Polls Codex, Anthropic, Z.AI, Gemini CLI / Antigravity usage and exposes it
|
||||
* via two channels:
|
||||
* • pi.events "usage:update" — for other extensions (e.g. footer-display)
|
||||
* • ctx.ui.setStatus("usage-bars", …) — formatted S/W braille bars
|
||||
*
|
||||
* Rendering / footer layout is handled by the separate footer-display extension.
|
||||
*/
|
||||
|
||||
import { DynamicBorder, type ExtensionAPI } from "@mariozechner/pi-coding-agent";
|
||||
import {
|
||||
Container,
|
||||
Input,
|
||||
Spacer,
|
||||
Text,
|
||||
getEditorKeybindings,
|
||||
type Focusable,
|
||||
} from "@mariozechner/pi-tui";
|
||||
import {
|
||||
canShowForProvider,
|
||||
clampPercent,
|
||||
colorForPercent,
|
||||
detectProvider,
|
||||
ensureFreshAuthForProviders,
|
||||
fetchAllUsages,
|
||||
fetchClaudeUsage,
|
||||
fetchCodexUsage,
|
||||
fetchGoogleUsage,
|
||||
fetchZaiUsage,
|
||||
providerToOAuthProviderId,
|
||||
readAuth,
|
||||
readUsageCache,
|
||||
resolveUsageEndpoints,
|
||||
writeUsageCache,
|
||||
type OAuthProviderId,
|
||||
type ProviderKey,
|
||||
type UsageByProvider,
|
||||
type UsageData,
|
||||
} from "./core";
|
||||
|
||||
const CACHE_TTL_MS = 15 * 60 * 1000;
|
||||
const ACTIVE_CACHE_TTL_MS = 3 * 60 * 1000;
|
||||
const STREAMING_POLL_INTERVAL_MS = 2 * 60 * 1000;
|
||||
const RATE_LIMITED_BACKOFF_MS = 60 * 60 * 1000;
|
||||
|
||||
const STATUS_KEY = "usage-bars";
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Braille gradient bar (⣀ ⣄ ⣤ ⣦ ⣶ ⣷ ⣿)
|
||||
// ---------------------------------------------------------------------------
|
||||
const BRAILLE_GRADIENT = "\u28C0\u28C4\u28E4\u28E6\u28F6\u28F7\u28FF";
|
||||
const BRAILLE_EMPTY = "\u28C0";
|
||||
const BAR_WIDTH = 5;
|
||||
|
||||
function renderBrailleBar(theme: any, value: number, width = BAR_WIDTH): string {
|
||||
const v = clampPercent(value);
|
||||
const levels = BRAILLE_GRADIENT.length - 1;
|
||||
const totalSteps = width * levels;
|
||||
const filledSteps = Math.round((v / 100) * totalSteps);
|
||||
const full = Math.floor(filledSteps / levels);
|
||||
const partial = filledSteps % levels;
|
||||
const empty = width - full - (partial ? 1 : 0);
|
||||
const color = colorForPercent(v);
|
||||
const filled = BRAILLE_GRADIENT[BRAILLE_GRADIENT.length - 1]!.repeat(Math.max(0, full));
|
||||
const partialChar = partial ? BRAILLE_GRADIENT[partial]! : "";
|
||||
const emptyChars = BRAILLE_EMPTY.repeat(Math.max(0, empty));
|
||||
return theme.fg(color, filled + partialChar) + theme.fg("dim", emptyChars);
|
||||
}
|
||||
|
||||
function renderBrailleBarWide(theme: any, value: number): string {
|
||||
return renderBrailleBar(theme, value, 12);
|
||||
}
|
||||
|
||||
const PROVIDER_LABELS: Record<ProviderKey, string> = {
|
||||
codex: "Codex",
|
||||
claude: "Claude",
|
||||
zai: "Z.AI",
|
||||
gemini: "Gemini",
|
||||
antigravity: "Antigravity",
|
||||
};
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// /usage command popup
|
||||
// ---------------------------------------------------------------------------
|
||||
interface SubscriptionItem {
|
||||
name: string;
|
||||
provider: ProviderKey;
|
||||
data: UsageData | null;
|
||||
isActive: boolean;
|
||||
}
|
||||
|
||||
class UsageSelectorComponent extends Container implements Focusable {
|
||||
private searchInput: Input;
|
||||
private listContainer: Container;
|
||||
private hintText: Text;
|
||||
private tui: any;
|
||||
private theme: any;
|
||||
private onCancelCallback: () => void;
|
||||
private allItems: SubscriptionItem[] = [];
|
||||
private filteredItems: SubscriptionItem[] = [];
|
||||
private selectedIndex = 0;
|
||||
private loading = true;
|
||||
private activeProvider: ProviderKey | null;
|
||||
private fetchAllFn: () => Promise<UsageByProvider>;
|
||||
private _focused = false;
|
||||
|
||||
get focused(): boolean { return this._focused; }
|
||||
set focused(value: boolean) { this._focused = value; this.searchInput.focused = value; }
|
||||
|
||||
constructor(
|
||||
tui: any,
|
||||
theme: any,
|
||||
activeProvider: ProviderKey | null,
|
||||
fetchAll: () => Promise<UsageByProvider>,
|
||||
onCancel: () => void,
|
||||
) {
|
||||
super();
|
||||
this.tui = tui;
|
||||
this.theme = theme;
|
||||
this.activeProvider = activeProvider;
|
||||
this.fetchAllFn = fetchAll;
|
||||
this.onCancelCallback = onCancel;
|
||||
|
||||
this.addChild(new DynamicBorder((s: string) => theme.fg("accent", s)));
|
||||
this.addChild(new Spacer(1));
|
||||
this.hintText = new Text(theme.fg("dim", "Fetching usage from all providers…"), 0, 0);
|
||||
this.addChild(this.hintText);
|
||||
this.addChild(new Spacer(1));
|
||||
this.searchInput = new Input();
|
||||
this.addChild(this.searchInput);
|
||||
this.addChild(new Spacer(1));
|
||||
this.listContainer = new Container();
|
||||
this.addChild(this.listContainer);
|
||||
this.addChild(new Spacer(1));
|
||||
this.addChild(new DynamicBorder((s: string) => theme.fg("accent", s)));
|
||||
|
||||
this.fetchAllFn()
|
||||
.then((results) => {
|
||||
this.loading = false;
|
||||
this.buildItems(results);
|
||||
this.updateList();
|
||||
this.hintText.setText(
|
||||
theme.fg("muted", "Only showing providers with credentials. ") +
|
||||
theme.fg("dim", "✓ = active provider"),
|
||||
);
|
||||
this.tui.requestRender();
|
||||
})
|
||||
.catch(() => {
|
||||
this.loading = false;
|
||||
this.hintText.setText(theme.fg("error", "Failed to fetch usage data"));
|
||||
this.tui.requestRender();
|
||||
});
|
||||
|
||||
this.updateList();
|
||||
}
|
||||
|
||||
private buildItems(results: UsageByProvider) {
|
||||
const providers: Array<{ key: ProviderKey; name: string }> = [
|
||||
{ key: "codex", name: "Codex" },
|
||||
{ key: "claude", name: "Claude" },
|
||||
{ key: "zai", name: "Z.AI" },
|
||||
{ key: "gemini", name: "Gemini" },
|
||||
{ key: "antigravity", name: "Antigravity" },
|
||||
];
|
||||
this.allItems = [];
|
||||
for (const p of providers) {
|
||||
if (results[p.key] !== null) {
|
||||
this.allItems.push({
|
||||
name: p.name,
|
||||
provider: p.key,
|
||||
data: results[p.key],
|
||||
isActive: this.activeProvider === p.key,
|
||||
});
|
||||
}
|
||||
}
|
||||
this.filteredItems = this.allItems;
|
||||
this.selectedIndex = Math.min(this.selectedIndex, Math.max(0, this.filteredItems.length - 1));
|
||||
}
|
||||
|
||||
private filterItems(query: string) {
|
||||
if (!query) {
|
||||
this.filteredItems = this.allItems;
|
||||
} else {
|
||||
const q = query.toLowerCase();
|
||||
this.filteredItems = this.allItems.filter(
|
||||
(item) => item.name.toLowerCase().includes(q) || item.provider.toLowerCase().includes(q),
|
||||
);
|
||||
}
|
||||
this.selectedIndex = Math.min(this.selectedIndex, Math.max(0, this.filteredItems.length - 1));
|
||||
}
|
||||
|
||||
private renderItem(item: SubscriptionItem, isSelected: boolean) {
|
||||
const t = this.theme;
|
||||
const pointer = isSelected ? t.fg("accent", "→ ") : " ";
|
||||
const activeBadge = item.isActive ? t.fg("success", " ✓") : "";
|
||||
const name = isSelected ? t.fg("accent", t.bold(item.name)) : item.name;
|
||||
this.listContainer.addChild(new Text(`${pointer}${name}${activeBadge}`, 0, 0));
|
||||
const indent = " ";
|
||||
|
||||
if (!item.data) {
|
||||
this.listContainer.addChild(new Text(indent + t.fg("dim", "No credentials"), 0, 0));
|
||||
} else if (item.data.error) {
|
||||
this.listContainer.addChild(new Text(indent + t.fg("error", item.data.error), 0, 0));
|
||||
} else {
|
||||
const session = clampPercent(item.data.session);
|
||||
const weekly = clampPercent(item.data.weekly);
|
||||
const sessionReset = item.data.sessionResetsIn
|
||||
? t.fg("dim", ` resets in ${item.data.sessionResetsIn}`) : "";
|
||||
const weeklyReset = item.data.weeklyResetsIn
|
||||
? t.fg("dim", ` resets in ${item.data.weeklyResetsIn}`) : "";
|
||||
|
||||
this.listContainer.addChild(new Text(
|
||||
indent + t.fg("muted", "Session ") +
|
||||
renderBrailleBarWide(t, session) + " " +
|
||||
t.fg(colorForPercent(session), `${session}%`.padStart(4)) + sessionReset,
|
||||
0, 0,
|
||||
));
|
||||
this.listContainer.addChild(new Text(
|
||||
indent + t.fg("muted", "Weekly ") +
|
||||
renderBrailleBarWide(t, weekly) + " " +
|
||||
t.fg(colorForPercent(weekly), `${weekly}%`.padStart(4)) + weeklyReset,
|
||||
0, 0,
|
||||
));
|
||||
|
||||
if (typeof item.data.extraSpend === "number" && typeof item.data.extraLimit === "number") {
|
||||
this.listContainer.addChild(new Text(
|
||||
indent + t.fg("muted", "Extra ") +
|
||||
t.fg("dim", `$${item.data.extraSpend.toFixed(2)} / $${item.data.extraLimit}`),
|
||||
0, 0,
|
||||
));
|
||||
}
|
||||
}
|
||||
this.listContainer.addChild(new Spacer(1));
|
||||
}
|
||||
|
||||
private updateList() {
|
||||
this.listContainer.clear();
|
||||
if (this.loading) {
|
||||
this.listContainer.addChild(new Text(this.theme.fg("muted", " Loading…"), 0, 0));
|
||||
return;
|
||||
}
|
||||
if (this.filteredItems.length === 0) {
|
||||
this.listContainer.addChild(new Text(this.theme.fg("muted", " No matching providers"), 0, 0));
|
||||
return;
|
||||
}
|
||||
for (let i = 0; i < this.filteredItems.length; i++) {
|
||||
this.renderItem(this.filteredItems[i]!, i === this.selectedIndex);
|
||||
}
|
||||
}
|
||||
|
||||
handleInput(keyData: string): void {
|
||||
const kb = getEditorKeybindings();
|
||||
if (kb.matches(keyData, "selectUp")) {
|
||||
if (this.filteredItems.length === 0) return;
|
||||
this.selectedIndex = this.selectedIndex === 0 ? this.filteredItems.length - 1 : this.selectedIndex - 1;
|
||||
this.updateList(); return;
|
||||
}
|
||||
if (kb.matches(keyData, "selectDown")) {
|
||||
if (this.filteredItems.length === 0) return;
|
||||
this.selectedIndex = this.selectedIndex === this.filteredItems.length - 1 ? 0 : this.selectedIndex + 1;
|
||||
this.updateList(); return;
|
||||
}
|
||||
if (kb.matches(keyData, "selectCancel") || kb.matches(keyData, "selectConfirm")) {
|
||||
this.onCancelCallback(); return;
|
||||
}
|
||||
this.searchInput.handleInput(keyData);
|
||||
this.filterItems(this.searchInput.getValue());
|
||||
this.updateList();
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Extension state
|
||||
// ---------------------------------------------------------------------------
|
||||
interface UsageState extends UsageByProvider {
|
||||
lastPoll: number;
|
||||
activeProvider: ProviderKey | null;
|
||||
}
|
||||
|
||||
interface PollOptions {
|
||||
cacheTtl?: number;
|
||||
forceFresh?: boolean;
|
||||
}
|
||||
|
||||
export default function (pi: ExtensionAPI) {
|
||||
const endpoints = resolveUsageEndpoints();
|
||||
const state: UsageState = {
|
||||
codex: null, claude: null, zai: null, gemini: null, antigravity: null,
|
||||
lastPoll: 0, activeProvider: null,
|
||||
};
|
||||
|
||||
let pollInFlight: Promise<void> | null = null;
|
||||
let pollQueued = false;
|
||||
let pollStartedAt = 0;
|
||||
let streamingTimer: ReturnType<typeof setInterval> | null = null;
|
||||
let ctx: any = null;
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Status update
|
||||
// ---------------------------------------------------------------------------
|
||||
function updateStatus() {
|
||||
const active = state.activeProvider;
|
||||
const data = active ? state[active] : null;
|
||||
|
||||
// Always emit event for other extensions (e.g. footer-display)
|
||||
if (data && !data.error) {
|
||||
pi.events.emit("usage:update", {
|
||||
session: data.session,
|
||||
weekly: data.weekly,
|
||||
sessionResetsIn: data.sessionResetsIn,
|
||||
sessionResetsAt: data.sessionResetsAt,
|
||||
weeklyResetsIn: data.weeklyResetsIn,
|
||||
});
|
||||
}
|
||||
|
||||
if (!ctx?.hasUI) return;
|
||||
|
||||
const theme = ctx.ui.theme;
|
||||
|
||||
if (!active) {
|
||||
ctx.ui.setStatus(STATUS_KEY, undefined);
|
||||
return;
|
||||
}
|
||||
|
||||
const auth = readAuth();
|
||||
if (!canShowForProvider(active, auth, endpoints)) {
|
||||
ctx.ui.setStatus(STATUS_KEY, undefined);
|
||||
return;
|
||||
}
|
||||
|
||||
if (!data) {
|
||||
ctx.ui.setStatus(STATUS_KEY, theme.fg("dim", "loading\u2026"));
|
||||
return;
|
||||
}
|
||||
|
||||
if (data.error) {
|
||||
const cache = readUsageCache();
|
||||
const blockedUntil = cache?.rateLimitedUntil?.[active] ?? 0;
|
||||
const note = blockedUntil > Date.now()
|
||||
? ` \u2014 retry in ${Math.ceil((blockedUntil - Date.now()) / 60000)}m` : "";
|
||||
ctx.ui.setStatus(STATUS_KEY, theme.fg("warning", `${PROVIDER_LABELS[active]} unavailable${note}`));
|
||||
return;
|
||||
}
|
||||
|
||||
const session = clampPercent(data.session);
|
||||
const weekly = clampPercent(data.weekly);
|
||||
|
||||
let s = theme.fg("muted", "S ") + renderBrailleBar(theme, session) + " " + theme.fg("dim", `${session}%`);
|
||||
if (data.sessionResetsIn) s += " " + theme.fg("dim", data.sessionResetsIn);
|
||||
|
||||
let w = theme.fg("muted", "W ") + renderBrailleBar(theme, weekly) + " " + theme.fg("dim", `${weekly}%`);
|
||||
if (data.weeklyResetsIn) w += " " + theme.fg("dim", `\u27F3 ${data.weeklyResetsIn}`);
|
||||
|
||||
ctx.ui.setStatus(STATUS_KEY, s + theme.fg("dim", " | ") + w);
|
||||
}
|
||||
|
||||
function updateProviderFrom(modelLike: any): boolean {
|
||||
const previous = state.activeProvider;
|
||||
state.activeProvider = detectProvider(modelLike);
|
||||
if (previous !== state.activeProvider) { updateStatus(); return true; }
|
||||
return false;
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Polling
|
||||
// ---------------------------------------------------------------------------
|
||||
async function runPollInner(options: PollOptions = {}) {
|
||||
const auth = readAuth();
|
||||
const active = state.activeProvider;
|
||||
|
||||
if (!canShowForProvider(active, auth, endpoints) || !auth || !active) {
|
||||
state.lastPoll = Date.now(); updateStatus(); return;
|
||||
}
|
||||
|
||||
const cache = readUsageCache();
|
||||
const now = Date.now();
|
||||
const cacheTtl = options.cacheTtl ?? CACHE_TTL_MS;
|
||||
|
||||
const blockedUntil = cache?.rateLimitedUntil?.[active] ?? 0;
|
||||
if (now < blockedUntil) {
|
||||
if (cache?.data?.[active]) {
|
||||
state[active] = cache.data[active]!;
|
||||
} else {
|
||||
// Rate-limited but no cached data — show a meaningful status instead
|
||||
// of leaving state null (which shows eternal "loading…").
|
||||
const retryMin = Math.ceil((blockedUntil - now) / 60000);
|
||||
state[active] = { session: 0, weekly: 0, error: `rate limited (retry in ${retryMin}m)` };
|
||||
}
|
||||
state.lastPoll = now; updateStatus(); return;
|
||||
}
|
||||
|
||||
if (!options.forceFresh && cache && now - cache.timestamp < cacheTtl && cache.data?.[active]) {
|
||||
state[active] = cache.data[active]!;
|
||||
state.lastPoll = now; updateStatus(); return;
|
||||
}
|
||||
|
||||
const oauthId = providerToOAuthProviderId(active);
|
||||
let effectiveAuth = auth;
|
||||
if (oauthId && active !== "zai") {
|
||||
const creds = auth[oauthId as keyof typeof auth] as
|
||||
| { access?: string; refresh?: string; expires?: number } | undefined;
|
||||
const expires = typeof creds?.expires === "number" ? creds.expires : 0;
|
||||
const tokenExpiredOrMissing = !creds?.access || (expires > 0 && Date.now() + 60_000 >= expires);
|
||||
if (tokenExpiredOrMissing && creds?.refresh) {
|
||||
try {
|
||||
const refreshPromise = ensureFreshAuthForProviders([oauthId as OAuthProviderId], { auth, persist: true });
|
||||
const timeoutPromise = new Promise<never>((_, reject) =>
|
||||
setTimeout(() => reject(new Error("OAuth refresh timeout")), 15_000),
|
||||
);
|
||||
const refreshed = await Promise.race([refreshPromise, timeoutPromise]);
|
||||
if (refreshed.auth) effectiveAuth = refreshed.auth;
|
||||
} catch {}
|
||||
}
|
||||
}
|
||||
|
||||
let result: UsageData;
|
||||
if (active === "codex") {
|
||||
const access = effectiveAuth["openai-codex"]?.access;
|
||||
result = access ? await fetchCodexUsage(access)
|
||||
: { session: 0, weekly: 0, error: "missing access token (try /login again)" };
|
||||
} else if (active === "claude") {
|
||||
const access = effectiveAuth.anthropic?.access;
|
||||
result = access ? await fetchClaudeUsage(access)
|
||||
: { session: 0, weekly: 0, error: "missing access token (try /login again)" };
|
||||
} else if (active === "zai") {
|
||||
const token = effectiveAuth.zai?.access || effectiveAuth.zai?.key;
|
||||
result = token ? await fetchZaiUsage(token, { endpoints })
|
||||
: { session: 0, weekly: 0, error: "missing token (try /login again)" };
|
||||
} else if (active === "gemini") {
|
||||
const creds = effectiveAuth["google-gemini-cli"];
|
||||
result = creds?.access
|
||||
? await fetchGoogleUsage(creds.access, endpoints.gemini, creds.projectId, "gemini", { endpoints })
|
||||
: { session: 0, weekly: 0, error: "missing access token (try /login again)" };
|
||||
} else {
|
||||
const creds = effectiveAuth["google-antigravity"];
|
||||
result = creds?.access
|
||||
? await fetchGoogleUsage(creds.access, endpoints.antigravity, creds.projectId, "antigravity", { endpoints })
|
||||
: { session: 0, weekly: 0, error: "missing access token (try /login again)" };
|
||||
}
|
||||
|
||||
state[active] = result;
|
||||
|
||||
if (result.error) {
|
||||
if (result.error === "HTTP 429") {
|
||||
const nextCache: import("./core").UsageCache = {
|
||||
timestamp: cache?.timestamp ?? now,
|
||||
data: { ...(cache?.data ?? {}) },
|
||||
rateLimitedUntil: { ...(cache?.rateLimitedUntil ?? {}), [active]: now + RATE_LIMITED_BACKOFF_MS },
|
||||
};
|
||||
writeUsageCache(nextCache);
|
||||
}
|
||||
} else {
|
||||
const nextCache: import("./core").UsageCache = {
|
||||
timestamp: now,
|
||||
data: { ...(cache?.data ?? {}), [active]: result },
|
||||
rateLimitedUntil: { ...(cache?.rateLimitedUntil ?? {}) },
|
||||
};
|
||||
delete nextCache.rateLimitedUntil![active];
|
||||
writeUsageCache(nextCache);
|
||||
}
|
||||
|
||||
state.lastPoll = now;
|
||||
updateStatus();
|
||||
}
|
||||
|
||||
async function runPoll(options: PollOptions = {}): Promise<void> {
|
||||
const timeout = new Promise<never>((_, reject) =>
|
||||
setTimeout(() => reject(new Error("runPoll timeout")), 25_000),
|
||||
);
|
||||
await Promise.race([runPollInner(options), timeout]);
|
||||
}
|
||||
|
||||
const POLL_TIMEOUT_MS = 30_000;
|
||||
|
||||
async function poll(options: PollOptions = {}) {
|
||||
// If a previous poll has been running longer than POLL_TIMEOUT_MS, abandon it
|
||||
// so we don't queue forever behind a stuck request.
|
||||
if (pollInFlight && pollStartedAt > 0 && Date.now() - pollStartedAt > POLL_TIMEOUT_MS) {
|
||||
pollInFlight = null;
|
||||
pollQueued = false;
|
||||
const active = state.activeProvider;
|
||||
if (active && !state[active]) {
|
||||
state[active] = { session: 0, weekly: 0, error: "poll timeout" };
|
||||
updateStatus();
|
||||
}
|
||||
}
|
||||
|
||||
if (pollInFlight) { pollQueued = true; await pollInFlight; return; }
|
||||
do {
|
||||
pollQueued = false;
|
||||
pollStartedAt = Date.now();
|
||||
pollInFlight = runPoll(options).catch(() => {
|
||||
// If runPoll threw, ensure we don't leave status stuck at "loading…"
|
||||
const active = state.activeProvider;
|
||||
if (active && !state[active]) {
|
||||
state[active] = { session: 0, weekly: 0, error: "poll failed" };
|
||||
updateStatus();
|
||||
}
|
||||
}).finally(() => { pollInFlight = null; pollStartedAt = 0; });
|
||||
await pollInFlight;
|
||||
} while (pollQueued);
|
||||
}
|
||||
|
||||
function startStreamingTimer() {
|
||||
if (streamingTimer !== null) return;
|
||||
streamingTimer = setInterval(() => { void poll({ cacheTtl: ACTIVE_CACHE_TTL_MS }); }, STREAMING_POLL_INTERVAL_MS);
|
||||
}
|
||||
|
||||
function stopStreamingTimer() {
|
||||
if (streamingTimer !== null) { clearInterval(streamingTimer); streamingTimer = null; }
|
||||
}
|
||||
|
||||
// ── Lifecycle ────────────────────────────────────────────────────────────
|
||||
|
||||
pi.on("session_start", async (_event, _ctx) => {
|
||||
ctx = _ctx;
|
||||
updateProviderFrom(_ctx.model);
|
||||
await poll();
|
||||
});
|
||||
|
||||
pi.on("session_shutdown", async (_event, _ctx) => {
|
||||
stopStreamingTimer();
|
||||
if (_ctx?.hasUI) _ctx.ui.setStatus(STATUS_KEY, undefined);
|
||||
});
|
||||
|
||||
pi.on("model_select", async (event, _ctx) => {
|
||||
ctx = _ctx;
|
||||
const changed = updateProviderFrom(event.model ?? _ctx.model);
|
||||
if (changed) await poll();
|
||||
});
|
||||
|
||||
pi.on("turn_start", (_event, _ctx) => { ctx = _ctx; updateProviderFrom(_ctx.model); });
|
||||
|
||||
pi.on("before_agent_start", async (_event, _ctx) => {
|
||||
ctx = _ctx;
|
||||
await poll({ cacheTtl: ACTIVE_CACHE_TTL_MS });
|
||||
});
|
||||
|
||||
pi.on("agent_start", (_event, _ctx) => { ctx = _ctx; startStreamingTimer(); });
|
||||
|
||||
pi.on("agent_end", async (_event, _ctx) => {
|
||||
ctx = _ctx;
|
||||
stopStreamingTimer();
|
||||
await poll({ cacheTtl: ACTIVE_CACHE_TTL_MS });
|
||||
});
|
||||
|
||||
pi.events.on("claude-account:switched", () => {
|
||||
const cache = readUsageCache();
|
||||
if (cache?.data?.claude) {
|
||||
const nextCache: import("./core").UsageCache = { ...cache, data: { ...cache.data } };
|
||||
delete nextCache.data.claude;
|
||||
writeUsageCache(nextCache);
|
||||
}
|
||||
void poll({ forceFresh: true });
|
||||
});
|
||||
|
||||
// ── /usage command ───────────────────────────────────────────────────────
|
||||
|
||||
pi.registerCommand("usage", {
|
||||
description: "Show API usage for all subscriptions",
|
||||
handler: async (_args, _ctx) => {
|
||||
ctx = _ctx;
|
||||
updateProviderFrom(_ctx.model);
|
||||
try {
|
||||
if (_ctx?.hasUI) {
|
||||
await _ctx.ui.custom<void>((tui, theme, _keybindings, done) => {
|
||||
return new UsageSelectorComponent(
|
||||
tui, theme, state.activeProvider,
|
||||
() => fetchAllUsages({ endpoints }),
|
||||
() => done(),
|
||||
);
|
||||
});
|
||||
}
|
||||
} finally {
|
||||
await poll({ cacheTtl: ACTIVE_CACHE_TTL_MS });
|
||||
}
|
||||
},
|
||||
});
|
||||
}
|
||||
@@ -1,581 +0,0 @@
|
||||
/**
|
||||
* Usage Extension - Minimal API usage indicator for pi
|
||||
*
|
||||
* Polls Codex, Anthropic, Z.AI, Gemini CLI / Antigravity usage and exposes it
|
||||
* via two channels:
|
||||
* • pi.events "usage:update" — for other extensions (e.g. footer-display)
|
||||
* • ctx.ui.setStatus("usage-bars", …) — formatted S/W braille bars
|
||||
*
|
||||
* Rendering / footer layout is handled by the separate footer-display extension.
|
||||
*/
|
||||
|
||||
import { DynamicBorder, type ExtensionAPI } from "@mariozechner/pi-coding-agent";
|
||||
import {
|
||||
Container,
|
||||
Input,
|
||||
Spacer,
|
||||
Text,
|
||||
getEditorKeybindings,
|
||||
type Focusable,
|
||||
} from "@mariozechner/pi-tui";
|
||||
import {
|
||||
canShowForProvider,
|
||||
clampPercent,
|
||||
colorForPercent,
|
||||
detectProvider,
|
||||
ensureFreshAuthForProviders,
|
||||
fetchAllUsages,
|
||||
fetchClaudeUsage,
|
||||
fetchCodexUsage,
|
||||
fetchGoogleUsage,
|
||||
fetchZaiUsage,
|
||||
providerToOAuthProviderId,
|
||||
readAuth,
|
||||
readUsageCache,
|
||||
resolveUsageEndpoints,
|
||||
writeUsageCache,
|
||||
type OAuthProviderId,
|
||||
type ProviderKey,
|
||||
type UsageByProvider,
|
||||
type UsageData,
|
||||
} from "./core";
|
||||
|
||||
const CACHE_TTL_MS = 15 * 60 * 1000;
|
||||
const ACTIVE_CACHE_TTL_MS = 3 * 60 * 1000;
|
||||
const STREAMING_POLL_INTERVAL_MS = 2 * 60 * 1000;
|
||||
const RATE_LIMITED_BACKOFF_MS = 60 * 60 * 1000;
|
||||
|
||||
const STATUS_KEY = "usage-bars";
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Braille gradient bar (⣀ ⣄ ⣤ ⣦ ⣶ ⣷ ⣿)
|
||||
// ---------------------------------------------------------------------------
|
||||
const BRAILLE_GRADIENT = "\u28C0\u28C4\u28E4\u28E6\u28F6\u28F7\u28FF";
|
||||
const BRAILLE_EMPTY = "\u28C0";
|
||||
const BAR_WIDTH = 5;
|
||||
|
||||
function renderBrailleBar(theme: any, value: number, width = BAR_WIDTH): string {
|
||||
const v = clampPercent(value);
|
||||
const levels = BRAILLE_GRADIENT.length - 1;
|
||||
const totalSteps = width * levels;
|
||||
const filledSteps = Math.round((v / 100) * totalSteps);
|
||||
const full = Math.floor(filledSteps / levels);
|
||||
const partial = filledSteps % levels;
|
||||
const empty = width - full - (partial ? 1 : 0);
|
||||
const color = colorForPercent(v);
|
||||
const filled = BRAILLE_GRADIENT[BRAILLE_GRADIENT.length - 1]!.repeat(Math.max(0, full));
|
||||
const partialChar = partial ? BRAILLE_GRADIENT[partial]! : "";
|
||||
const emptyChars = BRAILLE_EMPTY.repeat(Math.max(0, empty));
|
||||
return theme.fg(color, filled + partialChar) + theme.fg("dim", emptyChars);
|
||||
}
|
||||
|
||||
function renderBrailleBarWide(theme: any, value: number): string {
|
||||
return renderBrailleBar(theme, value, 12);
|
||||
}
|
||||
|
||||
const PROVIDER_LABELS: Record<ProviderKey, string> = {
|
||||
codex: "Codex",
|
||||
claude: "Claude",
|
||||
zai: "Z.AI",
|
||||
gemini: "Gemini",
|
||||
antigravity: "Antigravity",
|
||||
};
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// /usage command popup
|
||||
// ---------------------------------------------------------------------------
|
||||
interface SubscriptionItem {
|
||||
name: string;
|
||||
provider: ProviderKey;
|
||||
data: UsageData | null;
|
||||
isActive: boolean;
|
||||
}
|
||||
|
||||
class UsageSelectorComponent extends Container implements Focusable {
|
||||
private searchInput: Input;
|
||||
private listContainer: Container;
|
||||
private hintText: Text;
|
||||
private tui: any;
|
||||
private theme: any;
|
||||
private onCancelCallback: () => void;
|
||||
private allItems: SubscriptionItem[] = [];
|
||||
private filteredItems: SubscriptionItem[] = [];
|
||||
private selectedIndex = 0;
|
||||
private loading = true;
|
||||
private activeProvider: ProviderKey | null;
|
||||
private fetchAllFn: () => Promise<UsageByProvider>;
|
||||
private _focused = false;
|
||||
|
||||
get focused(): boolean { return this._focused; }
|
||||
set focused(value: boolean) { this._focused = value; this.searchInput.focused = value; }
|
||||
|
||||
constructor(
|
||||
tui: any,
|
||||
theme: any,
|
||||
activeProvider: ProviderKey | null,
|
||||
fetchAll: () => Promise<UsageByProvider>,
|
||||
onCancel: () => void,
|
||||
) {
|
||||
super();
|
||||
this.tui = tui;
|
||||
this.theme = theme;
|
||||
this.activeProvider = activeProvider;
|
||||
this.fetchAllFn = fetchAll;
|
||||
this.onCancelCallback = onCancel;
|
||||
|
||||
this.addChild(new DynamicBorder((s: string) => theme.fg("accent", s)));
|
||||
this.addChild(new Spacer(1));
|
||||
this.hintText = new Text(theme.fg("dim", "Fetching usage from all providers…"), 0, 0);
|
||||
this.addChild(this.hintText);
|
||||
this.addChild(new Spacer(1));
|
||||
this.searchInput = new Input();
|
||||
this.addChild(this.searchInput);
|
||||
this.addChild(new Spacer(1));
|
||||
this.listContainer = new Container();
|
||||
this.addChild(this.listContainer);
|
||||
this.addChild(new Spacer(1));
|
||||
this.addChild(new DynamicBorder((s: string) => theme.fg("accent", s)));
|
||||
|
||||
this.fetchAllFn()
|
||||
.then((results) => {
|
||||
this.loading = false;
|
||||
this.buildItems(results);
|
||||
this.updateList();
|
||||
this.hintText.setText(
|
||||
theme.fg("muted", "Only showing providers with credentials. ") +
|
||||
theme.fg("dim", "✓ = active provider"),
|
||||
);
|
||||
this.tui.requestRender();
|
||||
})
|
||||
.catch(() => {
|
||||
this.loading = false;
|
||||
this.hintText.setText(theme.fg("error", "Failed to fetch usage data"));
|
||||
this.tui.requestRender();
|
||||
});
|
||||
|
||||
this.updateList();
|
||||
}
|
||||
|
||||
private buildItems(results: UsageByProvider) {
|
||||
const providers: Array<{ key: ProviderKey; name: string }> = [
|
||||
{ key: "codex", name: "Codex" },
|
||||
{ key: "claude", name: "Claude" },
|
||||
{ key: "zai", name: "Z.AI" },
|
||||
{ key: "gemini", name: "Gemini" },
|
||||
{ key: "antigravity", name: "Antigravity" },
|
||||
];
|
||||
this.allItems = [];
|
||||
for (const p of providers) {
|
||||
if (results[p.key] !== null) {
|
||||
this.allItems.push({
|
||||
name: p.name,
|
||||
provider: p.key,
|
||||
data: results[p.key],
|
||||
isActive: this.activeProvider === p.key,
|
||||
});
|
||||
}
|
||||
}
|
||||
this.filteredItems = this.allItems;
|
||||
this.selectedIndex = Math.min(this.selectedIndex, Math.max(0, this.filteredItems.length - 1));
|
||||
}
|
||||
|
||||
private filterItems(query: string) {
|
||||
if (!query) {
|
||||
this.filteredItems = this.allItems;
|
||||
} else {
|
||||
const q = query.toLowerCase();
|
||||
this.filteredItems = this.allItems.filter(
|
||||
(item) => item.name.toLowerCase().includes(q) || item.provider.toLowerCase().includes(q),
|
||||
);
|
||||
}
|
||||
this.selectedIndex = Math.min(this.selectedIndex, Math.max(0, this.filteredItems.length - 1));
|
||||
}
|
||||
|
||||
private renderItem(item: SubscriptionItem, isSelected: boolean) {
|
||||
const t = this.theme;
|
||||
const pointer = isSelected ? t.fg("accent", "→ ") : " ";
|
||||
const activeBadge = item.isActive ? t.fg("success", " ✓") : "";
|
||||
const name = isSelected ? t.fg("accent", t.bold(item.name)) : item.name;
|
||||
this.listContainer.addChild(new Text(`${pointer}${name}${activeBadge}`, 0, 0));
|
||||
const indent = " ";
|
||||
|
||||
if (!item.data) {
|
||||
this.listContainer.addChild(new Text(indent + t.fg("dim", "No credentials"), 0, 0));
|
||||
} else if (item.data.error) {
|
||||
this.listContainer.addChild(new Text(indent + t.fg("error", item.data.error), 0, 0));
|
||||
} else {
|
||||
const session = clampPercent(item.data.session);
|
||||
const weekly = clampPercent(item.data.weekly);
|
||||
const sessionReset = item.data.sessionResetsIn
|
||||
? t.fg("dim", ` resets in ${item.data.sessionResetsIn}`) : "";
|
||||
const weeklyReset = item.data.weeklyResetsIn
|
||||
? t.fg("dim", ` resets in ${item.data.weeklyResetsIn}`) : "";
|
||||
|
||||
this.listContainer.addChild(new Text(
|
||||
indent + t.fg("muted", "Session ") +
|
||||
renderBrailleBarWide(t, session) + " " +
|
||||
t.fg(colorForPercent(session), `${session}%`.padStart(4)) + sessionReset,
|
||||
0, 0,
|
||||
));
|
||||
this.listContainer.addChild(new Text(
|
||||
indent + t.fg("muted", "Weekly ") +
|
||||
renderBrailleBarWide(t, weekly) + " " +
|
||||
t.fg(colorForPercent(weekly), `${weekly}%`.padStart(4)) + weeklyReset,
|
||||
0, 0,
|
||||
));
|
||||
|
||||
if (typeof item.data.extraSpend === "number" && typeof item.data.extraLimit === "number") {
|
||||
this.listContainer.addChild(new Text(
|
||||
indent + t.fg("muted", "Extra ") +
|
||||
t.fg("dim", `$${item.data.extraSpend.toFixed(2)} / $${item.data.extraLimit}`),
|
||||
0, 0,
|
||||
));
|
||||
}
|
||||
}
|
||||
this.listContainer.addChild(new Spacer(1));
|
||||
}
|
||||
|
||||
private updateList() {
|
||||
this.listContainer.clear();
|
||||
if (this.loading) {
|
||||
this.listContainer.addChild(new Text(this.theme.fg("muted", " Loading…"), 0, 0));
|
||||
return;
|
||||
}
|
||||
if (this.filteredItems.length === 0) {
|
||||
this.listContainer.addChild(new Text(this.theme.fg("muted", " No matching providers"), 0, 0));
|
||||
return;
|
||||
}
|
||||
for (let i = 0; i < this.filteredItems.length; i++) {
|
||||
this.renderItem(this.filteredItems[i]!, i === this.selectedIndex);
|
||||
}
|
||||
}
|
||||
|
||||
handleInput(keyData: string): void {
|
||||
const kb = getEditorKeybindings();
|
||||
if (kb.matches(keyData, "selectUp")) {
|
||||
if (this.filteredItems.length === 0) return;
|
||||
this.selectedIndex = this.selectedIndex === 0 ? this.filteredItems.length - 1 : this.selectedIndex - 1;
|
||||
this.updateList(); return;
|
||||
}
|
||||
if (kb.matches(keyData, "selectDown")) {
|
||||
if (this.filteredItems.length === 0) return;
|
||||
this.selectedIndex = this.selectedIndex === this.filteredItems.length - 1 ? 0 : this.selectedIndex + 1;
|
||||
this.updateList(); return;
|
||||
}
|
||||
if (kb.matches(keyData, "selectCancel") || kb.matches(keyData, "selectConfirm")) {
|
||||
this.onCancelCallback(); return;
|
||||
}
|
||||
this.searchInput.handleInput(keyData);
|
||||
this.filterItems(this.searchInput.getValue());
|
||||
this.updateList();
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Extension state
|
||||
// ---------------------------------------------------------------------------
|
||||
interface UsageState extends UsageByProvider {
|
||||
lastPoll: number;
|
||||
activeProvider: ProviderKey | null;
|
||||
}
|
||||
|
||||
interface PollOptions {
|
||||
cacheTtl?: number;
|
||||
forceFresh?: boolean;
|
||||
}
|
||||
|
||||
export default function (pi: ExtensionAPI) {
|
||||
const endpoints = resolveUsageEndpoints();
|
||||
const state: UsageState = {
|
||||
codex: null, claude: null, zai: null, gemini: null, antigravity: null,
|
||||
lastPoll: 0, activeProvider: null,
|
||||
};
|
||||
|
||||
let pollInFlight: Promise<void> | null = null;
|
||||
let pollQueued = false;
|
||||
let pollStartedAt = 0;
|
||||
let streamingTimer: ReturnType<typeof setInterval> | null = null;
|
||||
let ctx: any = null;
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Status update
|
||||
// ---------------------------------------------------------------------------
|
||||
function updateStatus() {
|
||||
const active = state.activeProvider;
|
||||
const data = active ? state[active] : null;
|
||||
|
||||
// Always emit event for other extensions (e.g. footer-display)
|
||||
if (data && !data.error) {
|
||||
pi.events.emit("usage:update", {
|
||||
session: data.session,
|
||||
weekly: data.weekly,
|
||||
sessionResetsIn: data.sessionResetsIn,
|
||||
sessionResetsAt: data.sessionResetsAt,
|
||||
weeklyResetsIn: data.weeklyResetsIn,
|
||||
});
|
||||
}
|
||||
|
||||
if (!ctx?.hasUI) return;
|
||||
|
||||
const theme = ctx.ui.theme;
|
||||
|
||||
if (!active) {
|
||||
ctx.ui.setStatus(STATUS_KEY, undefined);
|
||||
return;
|
||||
}
|
||||
|
||||
const auth = readAuth();
|
||||
if (!canShowForProvider(active, auth, endpoints)) {
|
||||
ctx.ui.setStatus(STATUS_KEY, undefined);
|
||||
return;
|
||||
}
|
||||
|
||||
if (!data) {
|
||||
ctx.ui.setStatus(STATUS_KEY, theme.fg("dim", "loading\u2026"));
|
||||
return;
|
||||
}
|
||||
|
||||
if (data.error) {
|
||||
const cache = readUsageCache();
|
||||
const blockedUntil = cache?.rateLimitedUntil?.[active] ?? 0;
|
||||
const note = blockedUntil > Date.now()
|
||||
? ` \u2014 retry in ${Math.ceil((blockedUntil - Date.now()) / 60000)}m` : "";
|
||||
ctx.ui.setStatus(STATUS_KEY, theme.fg("warning", `${PROVIDER_LABELS[active]} unavailable${note}`));
|
||||
return;
|
||||
}
|
||||
|
||||
const session = clampPercent(data.session);
|
||||
const weekly = clampPercent(data.weekly);
|
||||
|
||||
let s = theme.fg("muted", "S ") + renderBrailleBar(theme, session) + " " + theme.fg("dim", `${session}%`);
|
||||
if (data.sessionResetsIn) s += " " + theme.fg("dim", data.sessionResetsIn);
|
||||
|
||||
let w = theme.fg("muted", "W ") + renderBrailleBar(theme, weekly) + " " + theme.fg("dim", `${weekly}%`);
|
||||
if (data.weeklyResetsIn) w += " " + theme.fg("dim", `\u27F3 ${data.weeklyResetsIn}`);
|
||||
|
||||
ctx.ui.setStatus(STATUS_KEY, s + theme.fg("dim", " | ") + w);
|
||||
}
|
||||
|
||||
function updateProviderFrom(modelLike: any): boolean {
|
||||
const previous = state.activeProvider;
|
||||
state.activeProvider = detectProvider(modelLike);
|
||||
if (previous !== state.activeProvider) { updateStatus(); return true; }
|
||||
return false;
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Polling
|
||||
// ---------------------------------------------------------------------------
|
||||
async function runPollInner(options: PollOptions = {}) {
|
||||
const auth = readAuth();
|
||||
const active = state.activeProvider;
|
||||
|
||||
if (!canShowForProvider(active, auth, endpoints) || !auth || !active) {
|
||||
state.lastPoll = Date.now(); updateStatus(); return;
|
||||
}
|
||||
|
||||
const cache = readUsageCache();
|
||||
const now = Date.now();
|
||||
const cacheTtl = options.cacheTtl ?? CACHE_TTL_MS;
|
||||
|
||||
const blockedUntil = cache?.rateLimitedUntil?.[active] ?? 0;
|
||||
if (now < blockedUntil) {
|
||||
if (cache?.data?.[active]) {
|
||||
state[active] = cache.data[active]!;
|
||||
} else {
|
||||
// Rate-limited but no cached data — show a meaningful status instead
|
||||
// of leaving state null (which shows eternal "loading…").
|
||||
const retryMin = Math.ceil((blockedUntil - now) / 60000);
|
||||
state[active] = { session: 0, weekly: 0, error: `rate limited (retry in ${retryMin}m)` };
|
||||
}
|
||||
state.lastPoll = now; updateStatus(); return;
|
||||
}
|
||||
|
||||
if (!options.forceFresh && cache && now - cache.timestamp < cacheTtl && cache.data?.[active]) {
|
||||
state[active] = cache.data[active]!;
|
||||
state.lastPoll = now; updateStatus(); return;
|
||||
}
|
||||
|
||||
const oauthId = providerToOAuthProviderId(active);
|
||||
let effectiveAuth = auth;
|
||||
if (oauthId && active !== "zai") {
|
||||
const creds = auth[oauthId as keyof typeof auth] as
|
||||
| { access?: string; refresh?: string; expires?: number } | undefined;
|
||||
const expires = typeof creds?.expires === "number" ? creds.expires : 0;
|
||||
const tokenExpiredOrMissing = !creds?.access || (expires > 0 && Date.now() + 60_000 >= expires);
|
||||
if (tokenExpiredOrMissing && creds?.refresh) {
|
||||
try {
|
||||
const refreshPromise = ensureFreshAuthForProviders([oauthId as OAuthProviderId], { auth, persist: true });
|
||||
const timeoutPromise = new Promise<never>((_, reject) =>
|
||||
setTimeout(() => reject(new Error("OAuth refresh timeout")), 15_000),
|
||||
);
|
||||
const refreshed = await Promise.race([refreshPromise, timeoutPromise]);
|
||||
if (refreshed.auth) effectiveAuth = refreshed.auth;
|
||||
} catch {}
|
||||
}
|
||||
}
|
||||
|
||||
let result: UsageData;
|
||||
if (active === "codex") {
|
||||
const access = effectiveAuth["openai-codex"]?.access;
|
||||
result = access ? await fetchCodexUsage(access)
|
||||
: { session: 0, weekly: 0, error: "missing access token (try /login again)" };
|
||||
} else if (active === "claude") {
|
||||
const access = effectiveAuth.anthropic?.access;
|
||||
result = access ? await fetchClaudeUsage(access)
|
||||
: { session: 0, weekly: 0, error: "missing access token (try /login again)" };
|
||||
} else if (active === "zai") {
|
||||
const token = effectiveAuth.zai?.access || effectiveAuth.zai?.key;
|
||||
result = token ? await fetchZaiUsage(token, { endpoints })
|
||||
: { session: 0, weekly: 0, error: "missing token (try /login again)" };
|
||||
} else if (active === "gemini") {
|
||||
const creds = effectiveAuth["google-gemini-cli"];
|
||||
result = creds?.access
|
||||
? await fetchGoogleUsage(creds.access, endpoints.gemini, creds.projectId, "gemini", { endpoints })
|
||||
: { session: 0, weekly: 0, error: "missing access token (try /login again)" };
|
||||
} else {
|
||||
const creds = effectiveAuth["google-antigravity"];
|
||||
result = creds?.access
|
||||
? await fetchGoogleUsage(creds.access, endpoints.antigravity, creds.projectId, "antigravity", { endpoints })
|
||||
: { session: 0, weekly: 0, error: "missing access token (try /login again)" };
|
||||
}
|
||||
|
||||
state[active] = result;
|
||||
|
||||
if (result.error) {
|
||||
if (result.error === "HTTP 429") {
|
||||
const nextCache: import("./core").UsageCache = {
|
||||
timestamp: cache?.timestamp ?? now,
|
||||
data: { ...(cache?.data ?? {}) },
|
||||
rateLimitedUntil: { ...(cache?.rateLimitedUntil ?? {}), [active]: now + RATE_LIMITED_BACKOFF_MS },
|
||||
};
|
||||
writeUsageCache(nextCache);
|
||||
}
|
||||
} else {
|
||||
const nextCache: import("./core").UsageCache = {
|
||||
timestamp: now,
|
||||
data: { ...(cache?.data ?? {}), [active]: result },
|
||||
rateLimitedUntil: { ...(cache?.rateLimitedUntil ?? {}) },
|
||||
};
|
||||
delete nextCache.rateLimitedUntil![active];
|
||||
writeUsageCache(nextCache);
|
||||
}
|
||||
|
||||
state.lastPoll = now;
|
||||
updateStatus();
|
||||
}
|
||||
|
||||
async function runPoll(options: PollOptions = {}): Promise<void> {
|
||||
const timeout = new Promise<never>((_, reject) =>
|
||||
setTimeout(() => reject(new Error("runPoll timeout")), 25_000),
|
||||
);
|
||||
await Promise.race([runPollInner(options), timeout]);
|
||||
}
|
||||
|
||||
const POLL_TIMEOUT_MS = 30_000;
|
||||
|
||||
async function poll(options: PollOptions = {}) {
|
||||
// If a previous poll has been running longer than POLL_TIMEOUT_MS, abandon it
|
||||
// so we don't queue forever behind a stuck request.
|
||||
if (pollInFlight && pollStartedAt > 0 && Date.now() - pollStartedAt > POLL_TIMEOUT_MS) {
|
||||
pollInFlight = null;
|
||||
pollQueued = false;
|
||||
const active = state.activeProvider;
|
||||
if (active && !state[active]) {
|
||||
state[active] = { session: 0, weekly: 0, error: "poll timeout" };
|
||||
updateStatus();
|
||||
}
|
||||
}
|
||||
|
||||
if (pollInFlight) { pollQueued = true; await pollInFlight; return; }
|
||||
do {
|
||||
pollQueued = false;
|
||||
pollStartedAt = Date.now();
|
||||
pollInFlight = runPoll(options).catch(() => {
|
||||
// If runPoll threw, ensure we don't leave status stuck at "loading…"
|
||||
const active = state.activeProvider;
|
||||
if (active && !state[active]) {
|
||||
state[active] = { session: 0, weekly: 0, error: "poll failed" };
|
||||
updateStatus();
|
||||
}
|
||||
}).finally(() => { pollInFlight = null; pollStartedAt = 0; });
|
||||
await pollInFlight;
|
||||
} while (pollQueued);
|
||||
}
|
||||
|
||||
function startStreamingTimer() {
|
||||
if (streamingTimer !== null) return;
|
||||
streamingTimer = setInterval(() => { void poll({ cacheTtl: ACTIVE_CACHE_TTL_MS }); }, STREAMING_POLL_INTERVAL_MS);
|
||||
}
|
||||
|
||||
function stopStreamingTimer() {
|
||||
if (streamingTimer !== null) { clearInterval(streamingTimer); streamingTimer = null; }
|
||||
}
|
||||
|
||||
// ── Lifecycle ────────────────────────────────────────────────────────────
|
||||
|
||||
pi.on("session_start", async (_event, _ctx) => {
|
||||
ctx = _ctx;
|
||||
updateProviderFrom(_ctx.model);
|
||||
await poll();
|
||||
});
|
||||
|
||||
pi.on("session_shutdown", async (_event, _ctx) => {
|
||||
stopStreamingTimer();
|
||||
if (_ctx?.hasUI) _ctx.ui.setStatus(STATUS_KEY, undefined);
|
||||
});
|
||||
|
||||
pi.on("model_select", async (event, _ctx) => {
|
||||
ctx = _ctx;
|
||||
const changed = updateProviderFrom(event.model ?? _ctx.model);
|
||||
if (changed) await poll();
|
||||
});
|
||||
|
||||
pi.on("turn_start", (_event, _ctx) => { ctx = _ctx; updateProviderFrom(_ctx.model); });
|
||||
|
||||
pi.on("before_agent_start", async (_event, _ctx) => {
|
||||
ctx = _ctx;
|
||||
await poll({ cacheTtl: ACTIVE_CACHE_TTL_MS });
|
||||
});
|
||||
|
||||
pi.on("agent_start", (_event, _ctx) => { ctx = _ctx; startStreamingTimer(); });
|
||||
|
||||
pi.on("agent_end", async (_event, _ctx) => {
|
||||
ctx = _ctx;
|
||||
stopStreamingTimer();
|
||||
await poll({ cacheTtl: ACTIVE_CACHE_TTL_MS });
|
||||
});
|
||||
|
||||
pi.events.on("claude-account:switched", () => {
|
||||
const cache = readUsageCache();
|
||||
if (cache?.data?.claude) {
|
||||
const nextCache: import("./core").UsageCache = { ...cache, data: { ...cache.data } };
|
||||
delete nextCache.data.claude;
|
||||
writeUsageCache(nextCache);
|
||||
}
|
||||
void poll({ forceFresh: true });
|
||||
});
|
||||
|
||||
// ── /usage command ───────────────────────────────────────────────────────
|
||||
|
||||
pi.registerCommand("usage", {
|
||||
description: "Show API usage for all subscriptions",
|
||||
handler: async (_args, _ctx) => {
|
||||
ctx = _ctx;
|
||||
updateProviderFrom(_ctx.model);
|
||||
try {
|
||||
if (_ctx?.hasUI) {
|
||||
await _ctx.ui.custom<void>((tui, theme, _keybindings, done) => {
|
||||
return new UsageSelectorComponent(
|
||||
tui, theme, state.activeProvider,
|
||||
() => fetchAllUsages({ endpoints }),
|
||||
() => done(),
|
||||
);
|
||||
});
|
||||
}
|
||||
} finally {
|
||||
await poll({ cacheTtl: ACTIVE_CACHE_TTL_MS });
|
||||
}
|
||||
},
|
||||
});
|
||||
}
|
||||
@@ -1,581 +0,0 @@
|
||||
/**
|
||||
* Usage Extension - Minimal API usage indicator for pi
|
||||
*
|
||||
* Polls Codex, Anthropic, Z.AI, Gemini CLI / Antigravity usage and exposes it
|
||||
* via two channels:
|
||||
* • pi.events "usage:update" — for other extensions (e.g. footer-display)
|
||||
* • ctx.ui.setStatus("usage-bars", …) — formatted S/W braille bars
|
||||
*
|
||||
* Rendering / footer layout is handled by the separate footer-display extension.
|
||||
*/
|
||||
|
||||
import { DynamicBorder, type ExtensionAPI } from "@mariozechner/pi-coding-agent";
|
||||
import {
|
||||
Container,
|
||||
Input,
|
||||
Spacer,
|
||||
Text,
|
||||
getEditorKeybindings,
|
||||
type Focusable,
|
||||
} from "@mariozechner/pi-tui";
|
||||
import {
|
||||
canShowForProvider,
|
||||
clampPercent,
|
||||
colorForPercent,
|
||||
detectProvider,
|
||||
ensureFreshAuthForProviders,
|
||||
fetchAllUsages,
|
||||
fetchClaudeUsage,
|
||||
fetchCodexUsage,
|
||||
fetchGoogleUsage,
|
||||
fetchZaiUsage,
|
||||
providerToOAuthProviderId,
|
||||
readAuth,
|
||||
readUsageCache,
|
||||
resolveUsageEndpoints,
|
||||
writeUsageCache,
|
||||
type OAuthProviderId,
|
||||
type ProviderKey,
|
||||
type UsageByProvider,
|
||||
type UsageData,
|
||||
} from "./core";
|
||||
|
||||
const CACHE_TTL_MS = 15 * 60 * 1000;
|
||||
const ACTIVE_CACHE_TTL_MS = 3 * 60 * 1000;
|
||||
const STREAMING_POLL_INTERVAL_MS = 2 * 60 * 1000;
|
||||
const RATE_LIMITED_BACKOFF_MS = 60 * 60 * 1000;
|
||||
|
||||
const STATUS_KEY = "usage-bars";
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Braille gradient bar (⣀ ⣄ ⣤ ⣦ ⣶ ⣷ ⣿)
|
||||
// ---------------------------------------------------------------------------
|
||||
const BRAILLE_GRADIENT = "\u28C0\u28C4\u28E4\u28E6\u28F6\u28F7\u28FF";
|
||||
const BRAILLE_EMPTY = "\u28C0";
|
||||
const BAR_WIDTH = 5;
|
||||
|
||||
function renderBrailleBar(theme: any, value: number, width = BAR_WIDTH): string {
|
||||
const v = clampPercent(value);
|
||||
const levels = BRAILLE_GRADIENT.length - 1;
|
||||
const totalSteps = width * levels;
|
||||
const filledSteps = Math.round((v / 100) * totalSteps);
|
||||
const full = Math.floor(filledSteps / levels);
|
||||
const partial = filledSteps % levels;
|
||||
const empty = width - full - (partial ? 1 : 0);
|
||||
const color = colorForPercent(v);
|
||||
const filled = BRAILLE_GRADIENT[BRAILLE_GRADIENT.length - 1]!.repeat(Math.max(0, full));
|
||||
const partialChar = partial ? BRAILLE_GRADIENT[partial]! : "";
|
||||
const emptyChars = BRAILLE_EMPTY.repeat(Math.max(0, empty));
|
||||
return theme.fg(color, filled + partialChar) + theme.fg("dim", emptyChars);
|
||||
}
|
||||
|
||||
function renderBrailleBarWide(theme: any, value: number): string {
|
||||
return renderBrailleBar(theme, value, 12);
|
||||
}
|
||||
|
||||
const PROVIDER_LABELS: Record<ProviderKey, string> = {
|
||||
codex: "Codex",
|
||||
claude: "Claude",
|
||||
zai: "Z.AI",
|
||||
gemini: "Gemini",
|
||||
antigravity: "Antigravity",
|
||||
};
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// /usage command popup
|
||||
// ---------------------------------------------------------------------------
|
||||
interface SubscriptionItem {
|
||||
name: string;
|
||||
provider: ProviderKey;
|
||||
data: UsageData | null;
|
||||
isActive: boolean;
|
||||
}
|
||||
|
||||
class UsageSelectorComponent extends Container implements Focusable {
|
||||
private searchInput: Input;
|
||||
private listContainer: Container;
|
||||
private hintText: Text;
|
||||
private tui: any;
|
||||
private theme: any;
|
||||
private onCancelCallback: () => void;
|
||||
private allItems: SubscriptionItem[] = [];
|
||||
private filteredItems: SubscriptionItem[] = [];
|
||||
private selectedIndex = 0;
|
||||
private loading = true;
|
||||
private activeProvider: ProviderKey | null;
|
||||
private fetchAllFn: () => Promise<UsageByProvider>;
|
||||
private _focused = false;
|
||||
|
||||
get focused(): boolean { return this._focused; }
|
||||
set focused(value: boolean) { this._focused = value; this.searchInput.focused = value; }
|
||||
|
||||
constructor(
|
||||
tui: any,
|
||||
theme: any,
|
||||
activeProvider: ProviderKey | null,
|
||||
fetchAll: () => Promise<UsageByProvider>,
|
||||
onCancel: () => void,
|
||||
) {
|
||||
super();
|
||||
this.tui = tui;
|
||||
this.theme = theme;
|
||||
this.activeProvider = activeProvider;
|
||||
this.fetchAllFn = fetchAll;
|
||||
this.onCancelCallback = onCancel;
|
||||
|
||||
this.addChild(new DynamicBorder((s: string) => theme.fg("accent", s)));
|
||||
this.addChild(new Spacer(1));
|
||||
this.hintText = new Text(theme.fg("dim", "Fetching usage from all providers…"), 0, 0);
|
||||
this.addChild(this.hintText);
|
||||
this.addChild(new Spacer(1));
|
||||
this.searchInput = new Input();
|
||||
this.addChild(this.searchInput);
|
||||
this.addChild(new Spacer(1));
|
||||
this.listContainer = new Container();
|
||||
this.addChild(this.listContainer);
|
||||
this.addChild(new Spacer(1));
|
||||
this.addChild(new DynamicBorder((s: string) => theme.fg("accent", s)));
|
||||
|
||||
this.fetchAllFn()
|
||||
.then((results) => {
|
||||
this.loading = false;
|
||||
this.buildItems(results);
|
||||
this.updateList();
|
||||
this.hintText.setText(
|
||||
theme.fg("muted", "Only showing providers with credentials. ") +
|
||||
theme.fg("dim", "✓ = active provider"),
|
||||
);
|
||||
this.tui.requestRender();
|
||||
})
|
||||
.catch(() => {
|
||||
this.loading = false;
|
||||
this.hintText.setText(theme.fg("error", "Failed to fetch usage data"));
|
||||
this.tui.requestRender();
|
||||
});
|
||||
|
||||
this.updateList();
|
||||
}
|
||||
|
||||
private buildItems(results: UsageByProvider) {
|
||||
const providers: Array<{ key: ProviderKey; name: string }> = [
|
||||
{ key: "codex", name: "Codex" },
|
||||
{ key: "claude", name: "Claude" },
|
||||
{ key: "zai", name: "Z.AI" },
|
||||
{ key: "gemini", name: "Gemini" },
|
||||
{ key: "antigravity", name: "Antigravity" },
|
||||
];
|
||||
this.allItems = [];
|
||||
for (const p of providers) {
|
||||
if (results[p.key] !== null) {
|
||||
this.allItems.push({
|
||||
name: p.name,
|
||||
provider: p.key,
|
||||
data: results[p.key],
|
||||
isActive: this.activeProvider === p.key,
|
||||
});
|
||||
}
|
||||
}
|
||||
this.filteredItems = this.allItems;
|
||||
this.selectedIndex = Math.min(this.selectedIndex, Math.max(0, this.filteredItems.length - 1));
|
||||
}
|
||||
|
||||
private filterItems(query: string) {
|
||||
if (!query) {
|
||||
this.filteredItems = this.allItems;
|
||||
} else {
|
||||
const q = query.toLowerCase();
|
||||
this.filteredItems = this.allItems.filter(
|
||||
(item) => item.name.toLowerCase().includes(q) || item.provider.toLowerCase().includes(q),
|
||||
);
|
||||
}
|
||||
this.selectedIndex = Math.min(this.selectedIndex, Math.max(0, this.filteredItems.length - 1));
|
||||
}
|
||||
|
||||
private renderItem(item: SubscriptionItem, isSelected: boolean) {
|
||||
const t = this.theme;
|
||||
const pointer = isSelected ? t.fg("accent", "→ ") : " ";
|
||||
const activeBadge = item.isActive ? t.fg("success", " ✓") : "";
|
||||
const name = isSelected ? t.fg("accent", t.bold(item.name)) : item.name;
|
||||
this.listContainer.addChild(new Text(`${pointer}${name}${activeBadge}`, 0, 0));
|
||||
const indent = " ";
|
||||
|
||||
if (!item.data) {
|
||||
this.listContainer.addChild(new Text(indent + t.fg("dim", "No credentials"), 0, 0));
|
||||
} else if (item.data.error) {
|
||||
this.listContainer.addChild(new Text(indent + t.fg("error", item.data.error), 0, 0));
|
||||
} else {
|
||||
const session = clampPercent(item.data.session);
|
||||
const weekly = clampPercent(item.data.weekly);
|
||||
const sessionReset = item.data.sessionResetsIn
|
||||
? t.fg("dim", ` resets in ${item.data.sessionResetsIn}`) : "";
|
||||
const weeklyReset = item.data.weeklyResetsIn
|
||||
? t.fg("dim", ` resets in ${item.data.weeklyResetsIn}`) : "";
|
||||
|
||||
this.listContainer.addChild(new Text(
|
||||
indent + t.fg("muted", "Session ") +
|
||||
renderBrailleBarWide(t, session) + " " +
|
||||
t.fg(colorForPercent(session), `${session}%`.padStart(4)) + sessionReset,
|
||||
0, 0,
|
||||
));
|
||||
this.listContainer.addChild(new Text(
|
||||
indent + t.fg("muted", "Weekly ") +
|
||||
renderBrailleBarWide(t, weekly) + " " +
|
||||
t.fg(colorForPercent(weekly), `${weekly}%`.padStart(4)) + weeklyReset,
|
||||
0, 0,
|
||||
));
|
||||
|
||||
if (typeof item.data.extraSpend === "number" && typeof item.data.extraLimit === "number") {
|
||||
this.listContainer.addChild(new Text(
|
||||
indent + t.fg("muted", "Extra ") +
|
||||
t.fg("dim", `$${item.data.extraSpend.toFixed(2)} / $${item.data.extraLimit}`),
|
||||
0, 0,
|
||||
));
|
||||
}
|
||||
}
|
||||
this.listContainer.addChild(new Spacer(1));
|
||||
}
|
||||
|
||||
private updateList() {
|
||||
this.listContainer.clear();
|
||||
if (this.loading) {
|
||||
this.listContainer.addChild(new Text(this.theme.fg("muted", " Loading…"), 0, 0));
|
||||
return;
|
||||
}
|
||||
if (this.filteredItems.length === 0) {
|
||||
this.listContainer.addChild(new Text(this.theme.fg("muted", " No matching providers"), 0, 0));
|
||||
return;
|
||||
}
|
||||
for (let i = 0; i < this.filteredItems.length; i++) {
|
||||
this.renderItem(this.filteredItems[i]!, i === this.selectedIndex);
|
||||
}
|
||||
}
|
||||
|
||||
handleInput(keyData: string): void {
|
||||
const kb = getEditorKeybindings();
|
||||
if (kb.matches(keyData, "selectUp")) {
|
||||
if (this.filteredItems.length === 0) return;
|
||||
this.selectedIndex = this.selectedIndex === 0 ? this.filteredItems.length - 1 : this.selectedIndex - 1;
|
||||
this.updateList(); return;
|
||||
}
|
||||
if (kb.matches(keyData, "selectDown")) {
|
||||
if (this.filteredItems.length === 0) return;
|
||||
this.selectedIndex = this.selectedIndex === this.filteredItems.length - 1 ? 0 : this.selectedIndex + 1;
|
||||
this.updateList(); return;
|
||||
}
|
||||
if (kb.matches(keyData, "selectCancel") || kb.matches(keyData, "selectConfirm")) {
|
||||
this.onCancelCallback(); return;
|
||||
}
|
||||
this.searchInput.handleInput(keyData);
|
||||
this.filterItems(this.searchInput.getValue());
|
||||
this.updateList();
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Extension state
|
||||
// ---------------------------------------------------------------------------
|
||||
interface UsageState extends UsageByProvider {
|
||||
lastPoll: number;
|
||||
activeProvider: ProviderKey | null;
|
||||
}
|
||||
|
||||
interface PollOptions {
|
||||
cacheTtl?: number;
|
||||
forceFresh?: boolean;
|
||||
}
|
||||
|
||||
export default function (pi: ExtensionAPI) {
|
||||
const endpoints = resolveUsageEndpoints();
|
||||
const state: UsageState = {
|
||||
codex: null, claude: null, zai: null, gemini: null, antigravity: null,
|
||||
lastPoll: 0, activeProvider: null,
|
||||
};
|
||||
|
||||
let pollInFlight: Promise<void> | null = null;
|
||||
let pollQueued = false;
|
||||
let pollStartedAt = 0;
|
||||
let streamingTimer: ReturnType<typeof setInterval> | null = null;
|
||||
let ctx: any = null;
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Status update
|
||||
// ---------------------------------------------------------------------------
|
||||
function updateStatus() {
|
||||
const active = state.activeProvider;
|
||||
const data = active ? state[active] : null;
|
||||
|
||||
// Always emit event for other extensions (e.g. footer-display)
|
||||
if (data && !data.error) {
|
||||
pi.events.emit("usage:update", {
|
||||
session: data.session,
|
||||
weekly: data.weekly,
|
||||
sessionResetsIn: data.sessionResetsIn,
|
||||
sessionResetsAt: data.sessionResetsAt,
|
||||
weeklyResetsIn: data.weeklyResetsIn,
|
||||
});
|
||||
}
|
||||
|
||||
if (!ctx?.hasUI) return;
|
||||
|
||||
const theme = ctx.ui.theme;
|
||||
|
||||
if (!active) {
|
||||
ctx.ui.setStatus(STATUS_KEY, undefined);
|
||||
return;
|
||||
}
|
||||
|
||||
const auth = readAuth();
|
||||
if (!canShowForProvider(active, auth, endpoints)) {
|
||||
ctx.ui.setStatus(STATUS_KEY, undefined);
|
||||
return;
|
||||
}
|
||||
|
||||
if (!data) {
|
||||
ctx.ui.setStatus(STATUS_KEY, theme.fg("dim", "loading\u2026"));
|
||||
return;
|
||||
}
|
||||
|
||||
if (data.error) {
|
||||
const cache = readUsageCache();
|
||||
const blockedUntil = cache?.rateLimitedUntil?.[active] ?? 0;
|
||||
const note = blockedUntil > Date.now()
|
||||
? ` \u2014 retry in ${Math.ceil((blockedUntil - Date.now()) / 60000)}m` : "";
|
||||
ctx.ui.setStatus(STATUS_KEY, theme.fg("warning", `${PROVIDER_LABELS[active]} unavailable${note}`));
|
||||
return;
|
||||
}
|
||||
|
||||
const session = clampPercent(data.session);
|
||||
const weekly = clampPercent(data.weekly);
|
||||
|
||||
let s = theme.fg("muted", "S ") + renderBrailleBar(theme, session) + " " + theme.fg("dim", `${session}%`);
|
||||
if (data.sessionResetsIn) s += " " + theme.fg("dim", data.sessionResetsIn);
|
||||
|
||||
let w = theme.fg("muted", "W ") + renderBrailleBar(theme, weekly) + " " + theme.fg("dim", `${weekly}%`);
|
||||
if (data.weeklyResetsIn) w += " " + theme.fg("dim", `\u27F3 ${data.weeklyResetsIn}`);
|
||||
|
||||
ctx.ui.setStatus(STATUS_KEY, s + theme.fg("dim", " | ") + w);
|
||||
}
|
||||
|
||||
function updateProviderFrom(modelLike: any): boolean {
|
||||
const previous = state.activeProvider;
|
||||
state.activeProvider = detectProvider(modelLike);
|
||||
if (previous !== state.activeProvider) { updateStatus(); return true; }
|
||||
return false;
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Polling
|
||||
// ---------------------------------------------------------------------------
|
||||
async function runPollInner(options: PollOptions = {}) {
|
||||
const auth = readAuth();
|
||||
const active = state.activeProvider;
|
||||
|
||||
if (!canShowForProvider(active, auth, endpoints) || !auth || !active) {
|
||||
state.lastPoll = Date.now(); updateStatus(); return;
|
||||
}
|
||||
|
||||
const cache = readUsageCache();
|
||||
const now = Date.now();
|
||||
const cacheTtl = options.cacheTtl ?? CACHE_TTL_MS;
|
||||
|
||||
const blockedUntil = cache?.rateLimitedUntil?.[active] ?? 0;
|
||||
if (now < blockedUntil) {
|
||||
if (cache?.data?.[active]) {
|
||||
state[active] = cache.data[active]!;
|
||||
} else {
|
||||
// Rate-limited but no cached data — show a meaningful status instead
|
||||
// of leaving state null (which shows eternal "loading…").
|
||||
const retryMin = Math.ceil((blockedUntil - now) / 60000);
|
||||
state[active] = { session: 0, weekly: 0, error: `rate limited (retry in ${retryMin}m)` };
|
||||
}
|
||||
state.lastPoll = now; updateStatus(); return;
|
||||
}
|
||||
|
||||
if (!options.forceFresh && cache && now - cache.timestamp < cacheTtl && cache.data?.[active]) {
|
||||
state[active] = cache.data[active]!;
|
||||
state.lastPoll = now; updateStatus(); return;
|
||||
}
|
||||
|
||||
const oauthId = providerToOAuthProviderId(active);
|
||||
let effectiveAuth = auth;
|
||||
if (oauthId && active !== "zai") {
|
||||
const creds = auth[oauthId as keyof typeof auth] as
|
||||
| { access?: string; refresh?: string; expires?: number } | undefined;
|
||||
const expires = typeof creds?.expires === "number" ? creds.expires : 0;
|
||||
const tokenExpiredOrMissing = !creds?.access || (expires > 0 && Date.now() + 60_000 >= expires);
|
||||
if (tokenExpiredOrMissing && creds?.refresh) {
|
||||
try {
|
||||
const refreshPromise = ensureFreshAuthForProviders([oauthId as OAuthProviderId], { auth, persist: true });
|
||||
const timeoutPromise = new Promise<never>((_, reject) =>
|
||||
setTimeout(() => reject(new Error("OAuth refresh timeout")), 15_000),
|
||||
);
|
||||
const refreshed = await Promise.race([refreshPromise, timeoutPromise]);
|
||||
if (refreshed.auth) effectiveAuth = refreshed.auth;
|
||||
} catch {}
|
||||
}
|
||||
}
|
||||
|
||||
let result: UsageData;
|
||||
if (active === "codex") {
|
||||
const access = effectiveAuth["openai-codex"]?.access;
|
||||
result = access ? await fetchCodexUsage(access)
|
||||
: { session: 0, weekly: 0, error: "missing access token (try /login again)" };
|
||||
} else if (active === "claude") {
|
||||
const access = effectiveAuth.anthropic?.access;
|
||||
result = access ? await fetchClaudeUsage(access)
|
||||
: { session: 0, weekly: 0, error: "missing access token (try /login again)" };
|
||||
} else if (active === "zai") {
|
||||
const token = effectiveAuth.zai?.access || effectiveAuth.zai?.key;
|
||||
result = token ? await fetchZaiUsage(token, { endpoints })
|
||||
: { session: 0, weekly: 0, error: "missing token (try /login again)" };
|
||||
} else if (active === "gemini") {
|
||||
const creds = effectiveAuth["google-gemini-cli"];
|
||||
result = creds?.access
|
||||
? await fetchGoogleUsage(creds.access, endpoints.gemini, creds.projectId, "gemini", { endpoints })
|
||||
: { session: 0, weekly: 0, error: "missing access token (try /login again)" };
|
||||
} else {
|
||||
const creds = effectiveAuth["google-antigravity"];
|
||||
result = creds?.access
|
||||
? await fetchGoogleUsage(creds.access, endpoints.antigravity, creds.projectId, "antigravity", { endpoints })
|
||||
: { session: 0, weekly: 0, error: "missing access token (try /login again)" };
|
||||
}
|
||||
|
||||
state[active] = result;
|
||||
|
||||
if (result.error) {
|
||||
if (result.error === "HTTP 429") {
|
||||
const nextCache: import("./core").UsageCache = {
|
||||
timestamp: cache?.timestamp ?? now,
|
||||
data: { ...(cache?.data ?? {}) },
|
||||
rateLimitedUntil: { ...(cache?.rateLimitedUntil ?? {}), [active]: now + RATE_LIMITED_BACKOFF_MS },
|
||||
};
|
||||
writeUsageCache(nextCache);
|
||||
}
|
||||
} else {
|
||||
const nextCache: import("./core").UsageCache = {
|
||||
timestamp: now,
|
||||
data: { ...(cache?.data ?? {}), [active]: result },
|
||||
rateLimitedUntil: { ...(cache?.rateLimitedUntil ?? {}) },
|
||||
};
|
||||
delete nextCache.rateLimitedUntil![active];
|
||||
writeUsageCache(nextCache);
|
||||
}
|
||||
|
||||
state.lastPoll = now;
|
||||
updateStatus();
|
||||
}
|
||||
|
||||
async function runPoll(options: PollOptions = {}): Promise<void> {
|
||||
const timeout = new Promise<never>((_, reject) =>
|
||||
setTimeout(() => reject(new Error("runPoll timeout")), 25_000),
|
||||
);
|
||||
await Promise.race([runPollInner(options), timeout]);
|
||||
}
|
||||
|
||||
const POLL_TIMEOUT_MS = 30_000;
|
||||
|
||||
async function poll(options: PollOptions = {}) {
|
||||
// If a previous poll has been running longer than POLL_TIMEOUT_MS, abandon it
|
||||
// so we don't queue forever behind a stuck request.
|
||||
if (pollInFlight && pollStartedAt > 0 && Date.now() - pollStartedAt > POLL_TIMEOUT_MS) {
|
||||
pollInFlight = null;
|
||||
pollQueued = false;
|
||||
const active = state.activeProvider;
|
||||
if (active && !state[active]) {
|
||||
state[active] = { session: 0, weekly: 0, error: "poll timeout" };
|
||||
updateStatus();
|
||||
}
|
||||
}
|
||||
|
||||
if (pollInFlight) { pollQueued = true; await pollInFlight; return; }
|
||||
do {
|
||||
pollQueued = false;
|
||||
pollStartedAt = Date.now();
|
||||
pollInFlight = runPoll(options).catch(() => {
|
||||
// If runPoll threw, ensure we don't leave status stuck at "loading…"
|
||||
const active = state.activeProvider;
|
||||
if (active && !state[active]) {
|
||||
state[active] = { session: 0, weekly: 0, error: "poll failed" };
|
||||
updateStatus();
|
||||
}
|
||||
}).finally(() => { pollInFlight = null; pollStartedAt = 0; });
|
||||
await pollInFlight;
|
||||
} while (pollQueued);
|
||||
}
|
||||
|
||||
function startStreamingTimer() {
|
||||
if (streamingTimer !== null) return;
|
||||
streamingTimer = setInterval(() => { void poll({ cacheTtl: ACTIVE_CACHE_TTL_MS }); }, STREAMING_POLL_INTERVAL_MS);
|
||||
}
|
||||
|
||||
function stopStreamingTimer() {
|
||||
if (streamingTimer !== null) { clearInterval(streamingTimer); streamingTimer = null; }
|
||||
}
|
||||
|
||||
// ── Lifecycle ────────────────────────────────────────────────────────────
|
||||
|
||||
pi.on("session_start", async (_event, _ctx) => {
|
||||
ctx = _ctx;
|
||||
updateProviderFrom(_ctx.model);
|
||||
await poll();
|
||||
});
|
||||
|
||||
pi.on("session_shutdown", async (_event, _ctx) => {
|
||||
stopStreamingTimer();
|
||||
if (_ctx?.hasUI) _ctx.ui.setStatus(STATUS_KEY, undefined);
|
||||
});
|
||||
|
||||
pi.on("model_select", async (event, _ctx) => {
|
||||
ctx = _ctx;
|
||||
const changed = updateProviderFrom(event.model ?? _ctx.model);
|
||||
if (changed) await poll();
|
||||
});
|
||||
|
||||
pi.on("turn_start", (_event, _ctx) => { ctx = _ctx; updateProviderFrom(_ctx.model); });
|
||||
|
||||
pi.on("before_agent_start", async (_event, _ctx) => {
|
||||
ctx = _ctx;
|
||||
await poll({ cacheTtl: ACTIVE_CACHE_TTL_MS });
|
||||
});
|
||||
|
||||
pi.on("agent_start", (_event, _ctx) => { ctx = _ctx; startStreamingTimer(); });
|
||||
|
||||
pi.on("agent_end", async (_event, _ctx) => {
|
||||
ctx = _ctx;
|
||||
stopStreamingTimer();
|
||||
await poll({ cacheTtl: ACTIVE_CACHE_TTL_MS });
|
||||
});
|
||||
|
||||
pi.events.on("claude-account:switched", () => {
|
||||
const cache = readUsageCache();
|
||||
if (cache?.data?.claude) {
|
||||
const nextCache: import("./core").UsageCache = { ...cache, data: { ...cache.data } };
|
||||
delete nextCache.data.claude;
|
||||
writeUsageCache(nextCache);
|
||||
}
|
||||
void poll({ forceFresh: true });
|
||||
});
|
||||
|
||||
// ── /usage command ───────────────────────────────────────────────────────
|
||||
|
||||
pi.registerCommand("usage", {
|
||||
description: "Show API usage for all subscriptions",
|
||||
handler: async (_args, _ctx) => {
|
||||
ctx = _ctx;
|
||||
updateProviderFrom(_ctx.model);
|
||||
try {
|
||||
if (_ctx?.hasUI) {
|
||||
await _ctx.ui.custom<void>((tui, theme, _keybindings, done) => {
|
||||
return new UsageSelectorComponent(
|
||||
tui, theme, state.activeProvider,
|
||||
() => fetchAllUsages({ endpoints }),
|
||||
() => done(),
|
||||
);
|
||||
});
|
||||
}
|
||||
} finally {
|
||||
await poll({ cacheTtl: ACTIVE_CACHE_TTL_MS });
|
||||
}
|
||||
},
|
||||
});
|
||||
}
|
||||
@@ -1,581 +0,0 @@
|
||||
/**
|
||||
* Usage Extension - Minimal API usage indicator for pi
|
||||
*
|
||||
* Polls Codex, Anthropic, Z.AI, Gemini CLI / Antigravity usage and exposes it
|
||||
* via two channels:
|
||||
* • pi.events "usage:update" — for other extensions (e.g. footer-display)
|
||||
* • ctx.ui.setStatus("usage-bars", …) — formatted S/W braille bars
|
||||
*
|
||||
* Rendering / footer layout is handled by the separate footer-display extension.
|
||||
*/
|
||||
|
||||
import { DynamicBorder, type ExtensionAPI } from "@mariozechner/pi-coding-agent";
|
||||
import {
|
||||
Container,
|
||||
Input,
|
||||
Spacer,
|
||||
Text,
|
||||
getEditorKeybindings,
|
||||
type Focusable,
|
||||
} from "@mariozechner/pi-tui";
|
||||
import {
|
||||
canShowForProvider,
|
||||
clampPercent,
|
||||
colorForPercent,
|
||||
detectProvider,
|
||||
ensureFreshAuthForProviders,
|
||||
fetchAllUsages,
|
||||
fetchClaudeUsage,
|
||||
fetchCodexUsage,
|
||||
fetchGoogleUsage,
|
||||
fetchZaiUsage,
|
||||
providerToOAuthProviderId,
|
||||
readAuth,
|
||||
readUsageCache,
|
||||
resolveUsageEndpoints,
|
||||
writeUsageCache,
|
||||
type OAuthProviderId,
|
||||
type ProviderKey,
|
||||
type UsageByProvider,
|
||||
type UsageData,
|
||||
} from "./core";
|
||||
|
||||
const CACHE_TTL_MS = 15 * 60 * 1000;
|
||||
const ACTIVE_CACHE_TTL_MS = 3 * 60 * 1000;
|
||||
const STREAMING_POLL_INTERVAL_MS = 2 * 60 * 1000;
|
||||
const RATE_LIMITED_BACKOFF_MS = 60 * 60 * 1000;
|
||||
|
||||
const STATUS_KEY = "usage-bars";
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Braille gradient bar (⣀ ⣄ ⣤ ⣦ ⣶ ⣷ ⣿)
|
||||
// ---------------------------------------------------------------------------
|
||||
const BRAILLE_GRADIENT = "\u28C0\u28C4\u28E4\u28E6\u28F6\u28F7\u28FF";
|
||||
const BRAILLE_EMPTY = "\u28C0";
|
||||
const BAR_WIDTH = 5;
|
||||
|
||||
function renderBrailleBar(theme: any, value: number, width = BAR_WIDTH): string {
|
||||
const v = clampPercent(value);
|
||||
const levels = BRAILLE_GRADIENT.length - 1;
|
||||
const totalSteps = width * levels;
|
||||
const filledSteps = Math.round((v / 100) * totalSteps);
|
||||
const full = Math.floor(filledSteps / levels);
|
||||
const partial = filledSteps % levels;
|
||||
const empty = width - full - (partial ? 1 : 0);
|
||||
const color = colorForPercent(v);
|
||||
const filled = BRAILLE_GRADIENT[BRAILLE_GRADIENT.length - 1]!.repeat(Math.max(0, full));
|
||||
const partialChar = partial ? BRAILLE_GRADIENT[partial]! : "";
|
||||
const emptyChars = BRAILLE_EMPTY.repeat(Math.max(0, empty));
|
||||
return theme.fg(color, filled + partialChar) + theme.fg("dim", emptyChars);
|
||||
}
|
||||
|
||||
function renderBrailleBarWide(theme: any, value: number): string {
|
||||
return renderBrailleBar(theme, value, 12);
|
||||
}
|
||||
|
||||
const PROVIDER_LABELS: Record<ProviderKey, string> = {
|
||||
codex: "Codex",
|
||||
claude: "Claude",
|
||||
zai: "Z.AI",
|
||||
gemini: "Gemini",
|
||||
antigravity: "Antigravity",
|
||||
};
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// /usage command popup
|
||||
// ---------------------------------------------------------------------------
|
||||
interface SubscriptionItem {
|
||||
name: string;
|
||||
provider: ProviderKey;
|
||||
data: UsageData | null;
|
||||
isActive: boolean;
|
||||
}
|
||||
|
||||
class UsageSelectorComponent extends Container implements Focusable {
|
||||
private searchInput: Input;
|
||||
private listContainer: Container;
|
||||
private hintText: Text;
|
||||
private tui: any;
|
||||
private theme: any;
|
||||
private onCancelCallback: () => void;
|
||||
private allItems: SubscriptionItem[] = [];
|
||||
private filteredItems: SubscriptionItem[] = [];
|
||||
private selectedIndex = 0;
|
||||
private loading = true;
|
||||
private activeProvider: ProviderKey | null;
|
||||
private fetchAllFn: () => Promise<UsageByProvider>;
|
||||
private _focused = false;
|
||||
|
||||
get focused(): boolean { return this._focused; }
|
||||
set focused(value: boolean) { this._focused = value; this.searchInput.focused = value; }
|
||||
|
||||
constructor(
|
||||
tui: any,
|
||||
theme: any,
|
||||
activeProvider: ProviderKey | null,
|
||||
fetchAll: () => Promise<UsageByProvider>,
|
||||
onCancel: () => void,
|
||||
) {
|
||||
super();
|
||||
this.tui = tui;
|
||||
this.theme = theme;
|
||||
this.activeProvider = activeProvider;
|
||||
this.fetchAllFn = fetchAll;
|
||||
this.onCancelCallback = onCancel;
|
||||
|
||||
this.addChild(new DynamicBorder((s: string) => theme.fg("accent", s)));
|
||||
this.addChild(new Spacer(1));
|
||||
this.hintText = new Text(theme.fg("dim", "Fetching usage from all providers…"), 0, 0);
|
||||
this.addChild(this.hintText);
|
||||
this.addChild(new Spacer(1));
|
||||
this.searchInput = new Input();
|
||||
this.addChild(this.searchInput);
|
||||
this.addChild(new Spacer(1));
|
||||
this.listContainer = new Container();
|
||||
this.addChild(this.listContainer);
|
||||
this.addChild(new Spacer(1));
|
||||
this.addChild(new DynamicBorder((s: string) => theme.fg("accent", s)));
|
||||
|
||||
this.fetchAllFn()
|
||||
.then((results) => {
|
||||
this.loading = false;
|
||||
this.buildItems(results);
|
||||
this.updateList();
|
||||
this.hintText.setText(
|
||||
theme.fg("muted", "Only showing providers with credentials. ") +
|
||||
theme.fg("dim", "✓ = active provider"),
|
||||
);
|
||||
this.tui.requestRender();
|
||||
})
|
||||
.catch(() => {
|
||||
this.loading = false;
|
||||
this.hintText.setText(theme.fg("error", "Failed to fetch usage data"));
|
||||
this.tui.requestRender();
|
||||
});
|
||||
|
||||
this.updateList();
|
||||
}
|
||||
|
||||
private buildItems(results: UsageByProvider) {
|
||||
const providers: Array<{ key: ProviderKey; name: string }> = [
|
||||
{ key: "codex", name: "Codex" },
|
||||
{ key: "claude", name: "Claude" },
|
||||
{ key: "zai", name: "Z.AI" },
|
||||
{ key: "gemini", name: "Gemini" },
|
||||
{ key: "antigravity", name: "Antigravity" },
|
||||
];
|
||||
this.allItems = [];
|
||||
for (const p of providers) {
|
||||
if (results[p.key] !== null) {
|
||||
this.allItems.push({
|
||||
name: p.name,
|
||||
provider: p.key,
|
||||
data: results[p.key],
|
||||
isActive: this.activeProvider === p.key,
|
||||
});
|
||||
}
|
||||
}
|
||||
this.filteredItems = this.allItems;
|
||||
this.selectedIndex = Math.min(this.selectedIndex, Math.max(0, this.filteredItems.length - 1));
|
||||
}
|
||||
|
||||
private filterItems(query: string) {
|
||||
if (!query) {
|
||||
this.filteredItems = this.allItems;
|
||||
} else {
|
||||
const q = query.toLowerCase();
|
||||
this.filteredItems = this.allItems.filter(
|
||||
(item) => item.name.toLowerCase().includes(q) || item.provider.toLowerCase().includes(q),
|
||||
);
|
||||
}
|
||||
this.selectedIndex = Math.min(this.selectedIndex, Math.max(0, this.filteredItems.length - 1));
|
||||
}
|
||||
|
||||
private renderItem(item: SubscriptionItem, isSelected: boolean) {
|
||||
const t = this.theme;
|
||||
const pointer = isSelected ? t.fg("accent", "→ ") : " ";
|
||||
const activeBadge = item.isActive ? t.fg("success", " ✓") : "";
|
||||
const name = isSelected ? t.fg("accent", t.bold(item.name)) : item.name;
|
||||
this.listContainer.addChild(new Text(`${pointer}${name}${activeBadge}`, 0, 0));
|
||||
const indent = " ";
|
||||
|
||||
if (!item.data) {
|
||||
this.listContainer.addChild(new Text(indent + t.fg("dim", "No credentials"), 0, 0));
|
||||
} else if (item.data.error) {
|
||||
this.listContainer.addChild(new Text(indent + t.fg("error", item.data.error), 0, 0));
|
||||
} else {
|
||||
const session = clampPercent(item.data.session);
|
||||
const weekly = clampPercent(item.data.weekly);
|
||||
const sessionReset = item.data.sessionResetsIn
|
||||
? t.fg("dim", ` resets in ${item.data.sessionResetsIn}`) : "";
|
||||
const weeklyReset = item.data.weeklyResetsIn
|
||||
? t.fg("dim", ` resets in ${item.data.weeklyResetsIn}`) : "";
|
||||
|
||||
this.listContainer.addChild(new Text(
|
||||
indent + t.fg("muted", "Session ") +
|
||||
renderBrailleBarWide(t, session) + " " +
|
||||
t.fg(colorForPercent(session), `${session}%`.padStart(4)) + sessionReset,
|
||||
0, 0,
|
||||
));
|
||||
this.listContainer.addChild(new Text(
|
||||
indent + t.fg("muted", "Weekly ") +
|
||||
renderBrailleBarWide(t, weekly) + " " +
|
||||
t.fg(colorForPercent(weekly), `${weekly}%`.padStart(4)) + weeklyReset,
|
||||
0, 0,
|
||||
));
|
||||
|
||||
if (typeof item.data.extraSpend === "number" && typeof item.data.extraLimit === "number") {
|
||||
this.listContainer.addChild(new Text(
|
||||
indent + t.fg("muted", "Extra ") +
|
||||
t.fg("dim", `$${item.data.extraSpend.toFixed(2)} / $${item.data.extraLimit}`),
|
||||
0, 0,
|
||||
));
|
||||
}
|
||||
}
|
||||
this.listContainer.addChild(new Spacer(1));
|
||||
}
|
||||
|
||||
private updateList() {
|
||||
this.listContainer.clear();
|
||||
if (this.loading) {
|
||||
this.listContainer.addChild(new Text(this.theme.fg("muted", " Loading…"), 0, 0));
|
||||
return;
|
||||
}
|
||||
if (this.filteredItems.length === 0) {
|
||||
this.listContainer.addChild(new Text(this.theme.fg("muted", " No matching providers"), 0, 0));
|
||||
return;
|
||||
}
|
||||
for (let i = 0; i < this.filteredItems.length; i++) {
|
||||
this.renderItem(this.filteredItems[i]!, i === this.selectedIndex);
|
||||
}
|
||||
}
|
||||
|
||||
handleInput(keyData: string): void {
|
||||
const kb = getEditorKeybindings();
|
||||
if (kb.matches(keyData, "selectUp")) {
|
||||
if (this.filteredItems.length === 0) return;
|
||||
this.selectedIndex = this.selectedIndex === 0 ? this.filteredItems.length - 1 : this.selectedIndex - 1;
|
||||
this.updateList(); return;
|
||||
}
|
||||
if (kb.matches(keyData, "selectDown")) {
|
||||
if (this.filteredItems.length === 0) return;
|
||||
this.selectedIndex = this.selectedIndex === this.filteredItems.length - 1 ? 0 : this.selectedIndex + 1;
|
||||
this.updateList(); return;
|
||||
}
|
||||
if (kb.matches(keyData, "selectCancel") || kb.matches(keyData, "selectConfirm")) {
|
||||
this.onCancelCallback(); return;
|
||||
}
|
||||
this.searchInput.handleInput(keyData);
|
||||
this.filterItems(this.searchInput.getValue());
|
||||
this.updateList();
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Extension state
|
||||
// ---------------------------------------------------------------------------
|
||||
interface UsageState extends UsageByProvider {
|
||||
lastPoll: number;
|
||||
activeProvider: ProviderKey | null;
|
||||
}
|
||||
|
||||
interface PollOptions {
|
||||
cacheTtl?: number;
|
||||
forceFresh?: boolean;
|
||||
}
|
||||
|
||||
export default function (pi: ExtensionAPI) {
|
||||
const endpoints = resolveUsageEndpoints();
|
||||
const state: UsageState = {
|
||||
codex: null, claude: null, zai: null, gemini: null, antigravity: null,
|
||||
lastPoll: 0, activeProvider: null,
|
||||
};
|
||||
|
||||
let pollInFlight: Promise<void> | null = null;
|
||||
let pollQueued = false;
|
||||
let pollStartedAt = 0;
|
||||
let streamingTimer: ReturnType<typeof setInterval> | null = null;
|
||||
let ctx: any = null;
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Status update
|
||||
// ---------------------------------------------------------------------------
|
||||
function updateStatus() {
|
||||
const active = state.activeProvider;
|
||||
const data = active ? state[active] : null;
|
||||
|
||||
// Always emit event for other extensions (e.g. footer-display)
|
||||
if (data && !data.error) {
|
||||
pi.events.emit("usage:update", {
|
||||
session: data.session,
|
||||
weekly: data.weekly,
|
||||
sessionResetsIn: data.sessionResetsIn,
|
||||
sessionResetsAt: data.sessionResetsAt,
|
||||
weeklyResetsIn: data.weeklyResetsIn,
|
||||
});
|
||||
}
|
||||
|
||||
if (!ctx?.hasUI) return;
|
||||
|
||||
const theme = ctx.ui.theme;
|
||||
|
||||
if (!active) {
|
||||
ctx.ui.setStatus(STATUS_KEY, undefined);
|
||||
return;
|
||||
}
|
||||
|
||||
const auth = readAuth();
|
||||
if (!canShowForProvider(active, auth, endpoints)) {
|
||||
ctx.ui.setStatus(STATUS_KEY, undefined);
|
||||
return;
|
||||
}
|
||||
|
||||
if (!data) {
|
||||
ctx.ui.setStatus(STATUS_KEY, theme.fg("dim", "loading\u2026"));
|
||||
return;
|
||||
}
|
||||
|
||||
if (data.error) {
|
||||
const cache = readUsageCache();
|
||||
const blockedUntil = cache?.rateLimitedUntil?.[active] ?? 0;
|
||||
const note = blockedUntil > Date.now()
|
||||
? ` \u2014 retry in ${Math.ceil((blockedUntil - Date.now()) / 60000)}m` : "";
|
||||
ctx.ui.setStatus(STATUS_KEY, theme.fg("warning", `${PROVIDER_LABELS[active]} unavailable${note}`));
|
||||
return;
|
||||
}
|
||||
|
||||
const session = clampPercent(data.session);
|
||||
const weekly = clampPercent(data.weekly);
|
||||
|
||||
let s = theme.fg("muted", "S ") + renderBrailleBar(theme, session) + " " + theme.fg("dim", `${session}%`);
|
||||
if (data.sessionResetsIn) s += " " + theme.fg("dim", data.sessionResetsIn);
|
||||
|
||||
let w = theme.fg("muted", "W ") + renderBrailleBar(theme, weekly) + " " + theme.fg("dim", `${weekly}%`);
|
||||
if (data.weeklyResetsIn) w += " " + theme.fg("dim", `\u27F3 ${data.weeklyResetsIn}`);
|
||||
|
||||
ctx.ui.setStatus(STATUS_KEY, s + theme.fg("dim", " | ") + w);
|
||||
}
|
||||
|
||||
function updateProviderFrom(modelLike: any): boolean {
|
||||
const previous = state.activeProvider;
|
||||
state.activeProvider = detectProvider(modelLike);
|
||||
if (previous !== state.activeProvider) { updateStatus(); return true; }
|
||||
return false;
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Polling
|
||||
// ---------------------------------------------------------------------------
|
||||
async function runPollInner(options: PollOptions = {}) {
|
||||
const auth = readAuth();
|
||||
const active = state.activeProvider;
|
||||
|
||||
if (!canShowForProvider(active, auth, endpoints) || !auth || !active) {
|
||||
state.lastPoll = Date.now(); updateStatus(); return;
|
||||
}
|
||||
|
||||
const cache = readUsageCache();
|
||||
const now = Date.now();
|
||||
const cacheTtl = options.cacheTtl ?? CACHE_TTL_MS;
|
||||
|
||||
const blockedUntil = cache?.rateLimitedUntil?.[active] ?? 0;
|
||||
if (now < blockedUntil) {
|
||||
if (cache?.data?.[active]) {
|
||||
state[active] = cache.data[active]!;
|
||||
} else {
|
||||
// Rate-limited but no cached data — show a meaningful status instead
|
||||
// of leaving state null (which shows eternal "loading…").
|
||||
const retryMin = Math.ceil((blockedUntil - now) / 60000);
|
||||
state[active] = { session: 0, weekly: 0, error: `rate limited (retry in ${retryMin}m)` };
|
||||
}
|
||||
state.lastPoll = now; updateStatus(); return;
|
||||
}
|
||||
|
||||
if (!options.forceFresh && cache && now - cache.timestamp < cacheTtl && cache.data?.[active]) {
|
||||
state[active] = cache.data[active]!;
|
||||
state.lastPoll = now; updateStatus(); return;
|
||||
}
|
||||
|
||||
const oauthId = providerToOAuthProviderId(active);
|
||||
let effectiveAuth = auth;
|
||||
if (oauthId && active !== "zai") {
|
||||
const creds = auth[oauthId as keyof typeof auth] as
|
||||
| { access?: string; refresh?: string; expires?: number } | undefined;
|
||||
const expires = typeof creds?.expires === "number" ? creds.expires : 0;
|
||||
const tokenExpiredOrMissing = !creds?.access || (expires > 0 && Date.now() + 60_000 >= expires);
|
||||
if (tokenExpiredOrMissing && creds?.refresh) {
|
||||
try {
|
||||
const refreshPromise = ensureFreshAuthForProviders([oauthId as OAuthProviderId], { auth, persist: true });
|
||||
const timeoutPromise = new Promise<never>((_, reject) =>
|
||||
setTimeout(() => reject(new Error("OAuth refresh timeout")), 15_000),
|
||||
);
|
||||
const refreshed = await Promise.race([refreshPromise, timeoutPromise]);
|
||||
if (refreshed.auth) effectiveAuth = refreshed.auth;
|
||||
} catch {}
|
||||
}
|
||||
}
|
||||
|
||||
let result: UsageData;
|
||||
if (active === "codex") {
|
||||
const access = effectiveAuth["openai-codex"]?.access;
|
||||
result = access ? await fetchCodexUsage(access)
|
||||
: { session: 0, weekly: 0, error: "missing access token (try /login again)" };
|
||||
} else if (active === "claude") {
|
||||
const access = effectiveAuth.anthropic?.access;
|
||||
result = access ? await fetchClaudeUsage(access)
|
||||
: { session: 0, weekly: 0, error: "missing access token (try /login again)" };
|
||||
} else if (active === "zai") {
|
||||
const token = effectiveAuth.zai?.access || effectiveAuth.zai?.key;
|
||||
result = token ? await fetchZaiUsage(token, { endpoints })
|
||||
: { session: 0, weekly: 0, error: "missing token (try /login again)" };
|
||||
} else if (active === "gemini") {
|
||||
const creds = effectiveAuth["google-gemini-cli"];
|
||||
result = creds?.access
|
||||
? await fetchGoogleUsage(creds.access, endpoints.gemini, creds.projectId, "gemini", { endpoints })
|
||||
: { session: 0, weekly: 0, error: "missing access token (try /login again)" };
|
||||
} else {
|
||||
const creds = effectiveAuth["google-antigravity"];
|
||||
result = creds?.access
|
||||
? await fetchGoogleUsage(creds.access, endpoints.antigravity, creds.projectId, "antigravity", { endpoints })
|
||||
: { session: 0, weekly: 0, error: "missing access token (try /login again)" };
|
||||
}
|
||||
|
||||
state[active] = result;
|
||||
|
||||
if (result.error) {
|
||||
if (result.error === "HTTP 429") {
|
||||
const nextCache: import("./core").UsageCache = {
|
||||
timestamp: cache?.timestamp ?? now,
|
||||
data: { ...(cache?.data ?? {}) },
|
||||
rateLimitedUntil: { ...(cache?.rateLimitedUntil ?? {}), [active]: now + RATE_LIMITED_BACKOFF_MS },
|
||||
};
|
||||
writeUsageCache(nextCache);
|
||||
}
|
||||
} else {
|
||||
const nextCache: import("./core").UsageCache = {
|
||||
timestamp: now,
|
||||
data: { ...(cache?.data ?? {}), [active]: result },
|
||||
rateLimitedUntil: { ...(cache?.rateLimitedUntil ?? {}) },
|
||||
};
|
||||
delete nextCache.rateLimitedUntil![active];
|
||||
writeUsageCache(nextCache);
|
||||
}
|
||||
|
||||
state.lastPoll = now;
|
||||
updateStatus();
|
||||
}
|
||||
|
||||
async function runPoll(options: PollOptions = {}): Promise<void> {
|
||||
const timeout = new Promise<never>((_, reject) =>
|
||||
setTimeout(() => reject(new Error("runPoll timeout")), 25_000),
|
||||
);
|
||||
await Promise.race([runPollInner(options), timeout]);
|
||||
}
|
||||
|
||||
const POLL_TIMEOUT_MS = 30_000;
|
||||
|
||||
async function poll(options: PollOptions = {}) {
|
||||
// If a previous poll has been running longer than POLL_TIMEOUT_MS, abandon it
|
||||
// so we don't queue forever behind a stuck request.
|
||||
if (pollInFlight && pollStartedAt > 0 && Date.now() - pollStartedAt > POLL_TIMEOUT_MS) {
|
||||
pollInFlight = null;
|
||||
pollQueued = false;
|
||||
const active = state.activeProvider;
|
||||
if (active && !state[active]) {
|
||||
state[active] = { session: 0, weekly: 0, error: "poll timeout" };
|
||||
updateStatus();
|
||||
}
|
||||
}
|
||||
|
||||
if (pollInFlight) { pollQueued = true; await pollInFlight; return; }
|
||||
do {
|
||||
pollQueued = false;
|
||||
pollStartedAt = Date.now();
|
||||
pollInFlight = runPoll(options).catch(() => {
|
||||
// If runPoll threw, ensure we don't leave status stuck at "loading…"
|
||||
const active = state.activeProvider;
|
||||
if (active && !state[active]) {
|
||||
state[active] = { session: 0, weekly: 0, error: "poll failed" };
|
||||
updateStatus();
|
||||
}
|
||||
}).finally(() => { pollInFlight = null; pollStartedAt = 0; });
|
||||
await pollInFlight;
|
||||
} while (pollQueued);
|
||||
}
|
||||
|
||||
function startStreamingTimer() {
|
||||
if (streamingTimer !== null) return;
|
||||
streamingTimer = setInterval(() => { void poll({ cacheTtl: ACTIVE_CACHE_TTL_MS }); }, STREAMING_POLL_INTERVAL_MS);
|
||||
}
|
||||
|
||||
function stopStreamingTimer() {
|
||||
if (streamingTimer !== null) { clearInterval(streamingTimer); streamingTimer = null; }
|
||||
}
|
||||
|
||||
// ── Lifecycle ────────────────────────────────────────────────────────────
|
||||
|
||||
pi.on("session_start", async (_event, _ctx) => {
|
||||
ctx = _ctx;
|
||||
updateProviderFrom(_ctx.model);
|
||||
await poll();
|
||||
});
|
||||
|
||||
pi.on("session_shutdown", async (_event, _ctx) => {
|
||||
stopStreamingTimer();
|
||||
if (_ctx?.hasUI) _ctx.ui.setStatus(STATUS_KEY, undefined);
|
||||
});
|
||||
|
||||
pi.on("model_select", async (event, _ctx) => {
|
||||
ctx = _ctx;
|
||||
const changed = updateProviderFrom(event.model ?? _ctx.model);
|
||||
if (changed) await poll();
|
||||
});
|
||||
|
||||
pi.on("turn_start", (_event, _ctx) => { ctx = _ctx; updateProviderFrom(_ctx.model); });
|
||||
|
||||
pi.on("before_agent_start", async (_event, _ctx) => {
|
||||
ctx = _ctx;
|
||||
await poll({ cacheTtl: ACTIVE_CACHE_TTL_MS });
|
||||
});
|
||||
|
||||
pi.on("agent_start", (_event, _ctx) => { ctx = _ctx; startStreamingTimer(); });
|
||||
|
||||
pi.on("agent_end", async (_event, _ctx) => {
|
||||
ctx = _ctx;
|
||||
stopStreamingTimer();
|
||||
await poll({ cacheTtl: ACTIVE_CACHE_TTL_MS });
|
||||
});
|
||||
|
||||
pi.events.on("claude-account:switched", () => {
|
||||
const cache = readUsageCache();
|
||||
if (cache?.data?.claude) {
|
||||
const nextCache: import("./core").UsageCache = { ...cache, data: { ...cache.data } };
|
||||
delete nextCache.data.claude;
|
||||
writeUsageCache(nextCache);
|
||||
}
|
||||
void poll({ forceFresh: true });
|
||||
});
|
||||
|
||||
// ── /usage command ───────────────────────────────────────────────────────
|
||||
|
||||
pi.registerCommand("usage", {
|
||||
description: "Show API usage for all subscriptions",
|
||||
handler: async (_args, _ctx) => {
|
||||
ctx = _ctx;
|
||||
updateProviderFrom(_ctx.model);
|
||||
try {
|
||||
if (_ctx?.hasUI) {
|
||||
await _ctx.ui.custom<void>((tui, theme, _keybindings, done) => {
|
||||
return new UsageSelectorComponent(
|
||||
tui, theme, state.activeProvider,
|
||||
() => fetchAllUsages({ endpoints }),
|
||||
() => done(),
|
||||
);
|
||||
});
|
||||
}
|
||||
} finally {
|
||||
await poll({ cacheTtl: ACTIVE_CACHE_TTL_MS });
|
||||
}
|
||||
},
|
||||
});
|
||||
}
|
||||
@@ -1,581 +0,0 @@
|
||||
/**
|
||||
* Usage Extension - Minimal API usage indicator for pi
|
||||
*
|
||||
* Polls Codex, Anthropic, Z.AI, Gemini CLI / Antigravity usage and exposes it
|
||||
* via two channels:
|
||||
* • pi.events "usage:update" — for other extensions (e.g. footer-display)
|
||||
* • ctx.ui.setStatus("usage-bars", …) — formatted S/W braille bars
|
||||
*
|
||||
* Rendering / footer layout is handled by the separate footer-display extension.
|
||||
*/
|
||||
|
||||
import { DynamicBorder, type ExtensionAPI } from "@mariozechner/pi-coding-agent";
|
||||
import {
|
||||
Container,
|
||||
Input,
|
||||
Spacer,
|
||||
Text,
|
||||
getEditorKeybindings,
|
||||
type Focusable,
|
||||
} from "@mariozechner/pi-tui";
|
||||
import {
|
||||
canShowForProvider,
|
||||
clampPercent,
|
||||
colorForPercent,
|
||||
detectProvider,
|
||||
ensureFreshAuthForProviders,
|
||||
fetchAllUsages,
|
||||
fetchClaudeUsage,
|
||||
fetchCodexUsage,
|
||||
fetchGoogleUsage,
|
||||
fetchZaiUsage,
|
||||
providerToOAuthProviderId,
|
||||
readAuth,
|
||||
readUsageCache,
|
||||
resolveUsageEndpoints,
|
||||
writeUsageCache,
|
||||
type OAuthProviderId,
|
||||
type ProviderKey,
|
||||
type UsageByProvider,
|
||||
type UsageData,
|
||||
} from "./core";
|
||||
|
||||
const CACHE_TTL_MS = 15 * 60 * 1000;
|
||||
const ACTIVE_CACHE_TTL_MS = 3 * 60 * 1000;
|
||||
const STREAMING_POLL_INTERVAL_MS = 2 * 60 * 1000;
|
||||
const RATE_LIMITED_BACKOFF_MS = 60 * 60 * 1000;
|
||||
|
||||
const STATUS_KEY = "usage-bars";
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Braille gradient bar (⣀ ⣄ ⣤ ⣦ ⣶ ⣷ ⣿)
|
||||
// ---------------------------------------------------------------------------
|
||||
const BRAILLE_GRADIENT = "\u28C0\u28C4\u28E4\u28E6\u28F6\u28F7\u28FF";
|
||||
const BRAILLE_EMPTY = "\u28C0";
|
||||
const BAR_WIDTH = 5;
|
||||
|
||||
function renderBrailleBar(theme: any, value: number, width = BAR_WIDTH): string {
|
||||
const v = clampPercent(value);
|
||||
const levels = BRAILLE_GRADIENT.length - 1;
|
||||
const totalSteps = width * levels;
|
||||
const filledSteps = Math.round((v / 100) * totalSteps);
|
||||
const full = Math.floor(filledSteps / levels);
|
||||
const partial = filledSteps % levels;
|
||||
const empty = width - full - (partial ? 1 : 0);
|
||||
const color = colorForPercent(v);
|
||||
const filled = BRAILLE_GRADIENT[BRAILLE_GRADIENT.length - 1]!.repeat(Math.max(0, full));
|
||||
const partialChar = partial ? BRAILLE_GRADIENT[partial]! : "";
|
||||
const emptyChars = BRAILLE_EMPTY.repeat(Math.max(0, empty));
|
||||
return theme.fg(color, filled + partialChar) + theme.fg("dim", emptyChars);
|
||||
}
|
||||
|
||||
function renderBrailleBarWide(theme: any, value: number): string {
|
||||
return renderBrailleBar(theme, value, 12);
|
||||
}
|
||||
|
||||
const PROVIDER_LABELS: Record<ProviderKey, string> = {
|
||||
codex: "Codex",
|
||||
claude: "Claude",
|
||||
zai: "Z.AI",
|
||||
gemini: "Gemini",
|
||||
antigravity: "Antigravity",
|
||||
};
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// /usage command popup
|
||||
// ---------------------------------------------------------------------------
|
||||
interface SubscriptionItem {
|
||||
name: string;
|
||||
provider: ProviderKey;
|
||||
data: UsageData | null;
|
||||
isActive: boolean;
|
||||
}
|
||||
|
||||
class UsageSelectorComponent extends Container implements Focusable {
|
||||
private searchInput: Input;
|
||||
private listContainer: Container;
|
||||
private hintText: Text;
|
||||
private tui: any;
|
||||
private theme: any;
|
||||
private onCancelCallback: () => void;
|
||||
private allItems: SubscriptionItem[] = [];
|
||||
private filteredItems: SubscriptionItem[] = [];
|
||||
private selectedIndex = 0;
|
||||
private loading = true;
|
||||
private activeProvider: ProviderKey | null;
|
||||
private fetchAllFn: () => Promise<UsageByProvider>;
|
||||
private _focused = false;
|
||||
|
||||
get focused(): boolean { return this._focused; }
|
||||
set focused(value: boolean) { this._focused = value; this.searchInput.focused = value; }
|
||||
|
||||
constructor(
|
||||
tui: any,
|
||||
theme: any,
|
||||
activeProvider: ProviderKey | null,
|
||||
fetchAll: () => Promise<UsageByProvider>,
|
||||
onCancel: () => void,
|
||||
) {
|
||||
super();
|
||||
this.tui = tui;
|
||||
this.theme = theme;
|
||||
this.activeProvider = activeProvider;
|
||||
this.fetchAllFn = fetchAll;
|
||||
this.onCancelCallback = onCancel;
|
||||
|
||||
this.addChild(new DynamicBorder((s: string) => theme.fg("accent", s)));
|
||||
this.addChild(new Spacer(1));
|
||||
this.hintText = new Text(theme.fg("dim", "Fetching usage from all providers…"), 0, 0);
|
||||
this.addChild(this.hintText);
|
||||
this.addChild(new Spacer(1));
|
||||
this.searchInput = new Input();
|
||||
this.addChild(this.searchInput);
|
||||
this.addChild(new Spacer(1));
|
||||
this.listContainer = new Container();
|
||||
this.addChild(this.listContainer);
|
||||
this.addChild(new Spacer(1));
|
||||
this.addChild(new DynamicBorder((s: string) => theme.fg("accent", s)));
|
||||
|
||||
this.fetchAllFn()
|
||||
.then((results) => {
|
||||
this.loading = false;
|
||||
this.buildItems(results);
|
||||
this.updateList();
|
||||
this.hintText.setText(
|
||||
theme.fg("muted", "Only showing providers with credentials. ") +
|
||||
theme.fg("dim", "✓ = active provider"),
|
||||
);
|
||||
this.tui.requestRender();
|
||||
})
|
||||
.catch(() => {
|
||||
this.loading = false;
|
||||
this.hintText.setText(theme.fg("error", "Failed to fetch usage data"));
|
||||
this.tui.requestRender();
|
||||
});
|
||||
|
||||
this.updateList();
|
||||
}
|
||||
|
||||
private buildItems(results: UsageByProvider) {
|
||||
const providers: Array<{ key: ProviderKey; name: string }> = [
|
||||
{ key: "codex", name: "Codex" },
|
||||
{ key: "claude", name: "Claude" },
|
||||
{ key: "zai", name: "Z.AI" },
|
||||
{ key: "gemini", name: "Gemini" },
|
||||
{ key: "antigravity", name: "Antigravity" },
|
||||
];
|
||||
this.allItems = [];
|
||||
for (const p of providers) {
|
||||
if (results[p.key] !== null) {
|
||||
this.allItems.push({
|
||||
name: p.name,
|
||||
provider: p.key,
|
||||
data: results[p.key],
|
||||
isActive: this.activeProvider === p.key,
|
||||
});
|
||||
}
|
||||
}
|
||||
this.filteredItems = this.allItems;
|
||||
this.selectedIndex = Math.min(this.selectedIndex, Math.max(0, this.filteredItems.length - 1));
|
||||
}
|
||||
|
||||
private filterItems(query: string) {
|
||||
if (!query) {
|
||||
this.filteredItems = this.allItems;
|
||||
} else {
|
||||
const q = query.toLowerCase();
|
||||
this.filteredItems = this.allItems.filter(
|
||||
(item) => item.name.toLowerCase().includes(q) || item.provider.toLowerCase().includes(q),
|
||||
);
|
||||
}
|
||||
this.selectedIndex = Math.min(this.selectedIndex, Math.max(0, this.filteredItems.length - 1));
|
||||
}
|
||||
|
||||
private renderItem(item: SubscriptionItem, isSelected: boolean) {
|
||||
const t = this.theme;
|
||||
const pointer = isSelected ? t.fg("accent", "→ ") : " ";
|
||||
const activeBadge = item.isActive ? t.fg("success", " ✓") : "";
|
||||
const name = isSelected ? t.fg("accent", t.bold(item.name)) : item.name;
|
||||
this.listContainer.addChild(new Text(`${pointer}${name}${activeBadge}`, 0, 0));
|
||||
const indent = " ";
|
||||
|
||||
if (!item.data) {
|
||||
this.listContainer.addChild(new Text(indent + t.fg("dim", "No credentials"), 0, 0));
|
||||
} else if (item.data.error) {
|
||||
this.listContainer.addChild(new Text(indent + t.fg("error", item.data.error), 0, 0));
|
||||
} else {
|
||||
const session = clampPercent(item.data.session);
|
||||
const weekly = clampPercent(item.data.weekly);
|
||||
const sessionReset = item.data.sessionResetsIn
|
||||
? t.fg("dim", ` resets in ${item.data.sessionResetsIn}`) : "";
|
||||
const weeklyReset = item.data.weeklyResetsIn
|
||||
? t.fg("dim", ` resets in ${item.data.weeklyResetsIn}`) : "";
|
||||
|
||||
this.listContainer.addChild(new Text(
|
||||
indent + t.fg("muted", "Session ") +
|
||||
renderBrailleBarWide(t, session) + " " +
|
||||
t.fg(colorForPercent(session), `${session}%`.padStart(4)) + sessionReset,
|
||||
0, 0,
|
||||
));
|
||||
this.listContainer.addChild(new Text(
|
||||
indent + t.fg("muted", "Weekly ") +
|
||||
renderBrailleBarWide(t, weekly) + " " +
|
||||
t.fg(colorForPercent(weekly), `${weekly}%`.padStart(4)) + weeklyReset,
|
||||
0, 0,
|
||||
));
|
||||
|
||||
if (typeof item.data.extraSpend === "number" && typeof item.data.extraLimit === "number") {
|
||||
this.listContainer.addChild(new Text(
|
||||
indent + t.fg("muted", "Extra ") +
|
||||
t.fg("dim", `$${item.data.extraSpend.toFixed(2)} / $${item.data.extraLimit}`),
|
||||
0, 0,
|
||||
));
|
||||
}
|
||||
}
|
||||
this.listContainer.addChild(new Spacer(1));
|
||||
}
|
||||
|
||||
private updateList() {
|
||||
this.listContainer.clear();
|
||||
if (this.loading) {
|
||||
this.listContainer.addChild(new Text(this.theme.fg("muted", " Loading…"), 0, 0));
|
||||
return;
|
||||
}
|
||||
if (this.filteredItems.length === 0) {
|
||||
this.listContainer.addChild(new Text(this.theme.fg("muted", " No matching providers"), 0, 0));
|
||||
return;
|
||||
}
|
||||
for (let i = 0; i < this.filteredItems.length; i++) {
|
||||
this.renderItem(this.filteredItems[i]!, i === this.selectedIndex);
|
||||
}
|
||||
}
|
||||
|
||||
handleInput(keyData: string): void {
|
||||
const kb = getEditorKeybindings();
|
||||
if (kb.matches(keyData, "selectUp")) {
|
||||
if (this.filteredItems.length === 0) return;
|
||||
this.selectedIndex = this.selectedIndex === 0 ? this.filteredItems.length - 1 : this.selectedIndex - 1;
|
||||
this.updateList(); return;
|
||||
}
|
||||
if (kb.matches(keyData, "selectDown")) {
|
||||
if (this.filteredItems.length === 0) return;
|
||||
this.selectedIndex = this.selectedIndex === this.filteredItems.length - 1 ? 0 : this.selectedIndex + 1;
|
||||
this.updateList(); return;
|
||||
}
|
||||
if (kb.matches(keyData, "selectCancel") || kb.matches(keyData, "selectConfirm")) {
|
||||
this.onCancelCallback(); return;
|
||||
}
|
||||
this.searchInput.handleInput(keyData);
|
||||
this.filterItems(this.searchInput.getValue());
|
||||
this.updateList();
|
||||
}
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Extension state
|
||||
// ---------------------------------------------------------------------------
|
||||
interface UsageState extends UsageByProvider {
|
||||
lastPoll: number;
|
||||
activeProvider: ProviderKey | null;
|
||||
}
|
||||
|
||||
interface PollOptions {
|
||||
cacheTtl?: number;
|
||||
forceFresh?: boolean;
|
||||
}
|
||||
|
||||
export default function (pi: ExtensionAPI) {
|
||||
const endpoints = resolveUsageEndpoints();
|
||||
const state: UsageState = {
|
||||
codex: null, claude: null, zai: null, gemini: null, antigravity: null,
|
||||
lastPoll: 0, activeProvider: null,
|
||||
};
|
||||
|
||||
let pollInFlight: Promise<void> | null = null;
|
||||
let pollQueued = false;
|
||||
let pollStartedAt = 0;
|
||||
let streamingTimer: ReturnType<typeof setInterval> | null = null;
|
||||
let ctx: any = null;
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Status update
|
||||
// ---------------------------------------------------------------------------
|
||||
function updateStatus() {
|
||||
const active = state.activeProvider;
|
||||
const data = active ? state[active] : null;
|
||||
|
||||
// Always emit event for other extensions (e.g. footer-display)
|
||||
if (data && !data.error) {
|
||||
pi.events.emit("usage:update", {
|
||||
session: data.session,
|
||||
weekly: data.weekly,
|
||||
sessionResetsIn: data.sessionResetsIn,
|
||||
sessionResetsAt: data.sessionResetsAt,
|
||||
weeklyResetsIn: data.weeklyResetsIn,
|
||||
});
|
||||
}
|
||||
|
||||
if (!ctx?.hasUI) return;
|
||||
|
||||
const theme = ctx.ui.theme;
|
||||
|
||||
if (!active) {
|
||||
ctx.ui.setStatus(STATUS_KEY, undefined);
|
||||
return;
|
||||
}
|
||||
|
||||
const auth = readAuth();
|
||||
if (!canShowForProvider(active, auth, endpoints)) {
|
||||
ctx.ui.setStatus(STATUS_KEY, undefined);
|
||||
return;
|
||||
}
|
||||
|
||||
if (!data) {
|
||||
ctx.ui.setStatus(STATUS_KEY, theme.fg("dim", "loading\u2026"));
|
||||
return;
|
||||
}
|
||||
|
||||
if (data.error) {
|
||||
const cache = readUsageCache();
|
||||
const blockedUntil = cache?.rateLimitedUntil?.[active] ?? 0;
|
||||
const note = blockedUntil > Date.now()
|
||||
? ` \u2014 retry in ${Math.ceil((blockedUntil - Date.now()) / 60000)}m` : "";
|
||||
ctx.ui.setStatus(STATUS_KEY, theme.fg("warning", `${PROVIDER_LABELS[active]} unavailable${note}`));
|
||||
return;
|
||||
}
|
||||
|
||||
const session = clampPercent(data.session);
|
||||
const weekly = clampPercent(data.weekly);
|
||||
|
||||
let s = theme.fg("muted", "S ") + renderBrailleBar(theme, session) + " " + theme.fg("dim", `${session}%`);
|
||||
if (data.sessionResetsIn) s += " " + theme.fg("dim", data.sessionResetsIn);
|
||||
|
||||
let w = theme.fg("muted", "W ") + renderBrailleBar(theme, weekly) + " " + theme.fg("dim", `${weekly}%`);
|
||||
if (data.weeklyResetsIn) w += " " + theme.fg("dim", `\u27F3 ${data.weeklyResetsIn}`);
|
||||
|
||||
ctx.ui.setStatus(STATUS_KEY, s + theme.fg("dim", " | ") + w);
|
||||
}
|
||||
|
||||
function updateProviderFrom(modelLike: any): boolean {
|
||||
const previous = state.activeProvider;
|
||||
state.activeProvider = detectProvider(modelLike);
|
||||
if (previous !== state.activeProvider) { updateStatus(); return true; }
|
||||
return false;
|
||||
}
|
||||
|
||||
// ---------------------------------------------------------------------------
|
||||
// Polling
|
||||
// ---------------------------------------------------------------------------
|
||||
async function runPollInner(options: PollOptions = {}) {
|
||||
const auth = readAuth();
|
||||
const active = state.activeProvider;
|
||||
|
||||
if (!canShowForProvider(active, auth, endpoints) || !auth || !active) {
|
||||
state.lastPoll = Date.now(); updateStatus(); return;
|
||||
}
|
||||
|
||||
const cache = readUsageCache();
|
||||
const now = Date.now();
|
||||
const cacheTtl = options.cacheTtl ?? CACHE_TTL_MS;
|
||||
|
||||
const blockedUntil = cache?.rateLimitedUntil?.[active] ?? 0;
|
||||
if (now < blockedUntil) {
|
||||
if (cache?.data?.[active]) {
|
||||
state[active] = cache.data[active]!;
|
||||
} else {
|
||||
// Rate-limited but no cached data — show a meaningful status instead
|
||||
// of leaving state null (which shows eternal "loading…").
|
||||
const retryMin = Math.ceil((blockedUntil - now) / 60000);
|
||||
state[active] = { session: 0, weekly: 0, error: `rate limited (retry in ${retryMin}m)` };
|
||||
}
|
||||
state.lastPoll = now; updateStatus(); return;
|
||||
}
|
||||
|
||||
if (!options.forceFresh && cache && now - cache.timestamp < cacheTtl && cache.data?.[active]) {
|
||||
state[active] = cache.data[active]!;
|
||||
state.lastPoll = now; updateStatus(); return;
|
||||
}
|
||||
|
||||
const oauthId = providerToOAuthProviderId(active);
|
||||
let effectiveAuth = auth;
|
||||
if (oauthId && active !== "zai") {
|
||||
const creds = auth[oauthId as keyof typeof auth] as
|
||||
| { access?: string; refresh?: string; expires?: number } | undefined;
|
||||
const expires = typeof creds?.expires === "number" ? creds.expires : 0;
|
||||
const tokenExpiredOrMissing = !creds?.access || (expires > 0 && Date.now() + 60_000 >= expires);
|
||||
if (tokenExpiredOrMissing && creds?.refresh) {
|
||||
try {
|
||||
const refreshPromise = ensureFreshAuthForProviders([oauthId as OAuthProviderId], { auth, persist: true });
|
||||
const timeoutPromise = new Promise<never>((_, reject) =>
|
||||
setTimeout(() => reject(new Error("OAuth refresh timeout")), 15_000),
|
||||
);
|
||||
const refreshed = await Promise.race([refreshPromise, timeoutPromise]);
|
||||
if (refreshed.auth) effectiveAuth = refreshed.auth;
|
||||
} catch {}
|
||||
}
|
||||
}
|
||||
|
||||
let result: UsageData;
|
||||
if (active === "codex") {
|
||||
const access = effectiveAuth["openai-codex"]?.access;
|
||||
result = access ? await fetchCodexUsage(access)
|
||||
: { session: 0, weekly: 0, error: "missing access token (try /login again)" };
|
||||
} else if (active === "claude") {
|
||||
const access = effectiveAuth.anthropic?.access;
|
||||
result = access ? await fetchClaudeUsage(access)
|
||||
: { session: 0, weekly: 0, error: "missing access token (try /login again)" };
|
||||
} else if (active === "zai") {
|
||||
const token = effectiveAuth.zai?.access || effectiveAuth.zai?.key;
|
||||
result = token ? await fetchZaiUsage(token, { endpoints })
|
||||
: { session: 0, weekly: 0, error: "missing token (try /login again)" };
|
||||
} else if (active === "gemini") {
|
||||
const creds = effectiveAuth["google-gemini-cli"];
|
||||
result = creds?.access
|
||||
? await fetchGoogleUsage(creds.access, endpoints.gemini, creds.projectId, "gemini", { endpoints })
|
||||
: { session: 0, weekly: 0, error: "missing access token (try /login again)" };
|
||||
} else {
|
||||
const creds = effectiveAuth["google-antigravity"];
|
||||
result = creds?.access
|
||||
? await fetchGoogleUsage(creds.access, endpoints.antigravity, creds.projectId, "antigravity", { endpoints })
|
||||
: { session: 0, weekly: 0, error: "missing access token (try /login again)" };
|
||||
}
|
||||
|
||||
state[active] = result;
|
||||
|
||||
if (result.error) {
|
||||
if (result.error === "HTTP 429") {
|
||||
const nextCache: import("./core").UsageCache = {
|
||||
timestamp: cache?.timestamp ?? now,
|
||||
data: { ...(cache?.data ?? {}) },
|
||||
rateLimitedUntil: { ...(cache?.rateLimitedUntil ?? {}), [active]: now + RATE_LIMITED_BACKOFF_MS },
|
||||
};
|
||||
writeUsageCache(nextCache);
|
||||
}
|
||||
} else {
|
||||
const nextCache: import("./core").UsageCache = {
|
||||
timestamp: now,
|
||||
data: { ...(cache?.data ?? {}), [active]: result },
|
||||
rateLimitedUntil: { ...(cache?.rateLimitedUntil ?? {}) },
|
||||
};
|
||||
delete nextCache.rateLimitedUntil![active];
|
||||
writeUsageCache(nextCache);
|
||||
}
|
||||
|
||||
state.lastPoll = now;
|
||||
updateStatus();
|
||||
}
|
||||
|
||||
async function runPoll(options: PollOptions = {}): Promise<void> {
|
||||
const timeout = new Promise<never>((_, reject) =>
|
||||
setTimeout(() => reject(new Error("runPoll timeout")), 25_000),
|
||||
);
|
||||
await Promise.race([runPollInner(options), timeout]);
|
||||
}
|
||||
|
||||
const POLL_TIMEOUT_MS = 30_000;
|
||||
|
||||
async function poll(options: PollOptions = {}) {
|
||||
// If a previous poll has been running longer than POLL_TIMEOUT_MS, abandon it
|
||||
// so we don't queue forever behind a stuck request.
|
||||
if (pollInFlight && pollStartedAt > 0 && Date.now() - pollStartedAt > POLL_TIMEOUT_MS) {
|
||||
pollInFlight = null;
|
||||
pollQueued = false;
|
||||
const active = state.activeProvider;
|
||||
if (active && !state[active]) {
|
||||
state[active] = { session: 0, weekly: 0, error: "poll timeout" };
|
||||
updateStatus();
|
||||
}
|
||||
}
|
||||
|
||||
if (pollInFlight) { pollQueued = true; await pollInFlight; return; }
|
||||
do {
|
||||
pollQueued = false;
|
||||
pollStartedAt = Date.now();
|
||||
pollInFlight = runPoll(options).catch(() => {
|
||||
// If runPoll threw, ensure we don't leave status stuck at "loading…"
|
||||
const active = state.activeProvider;
|
||||
if (active && !state[active]) {
|
||||
state[active] = { session: 0, weekly: 0, error: "poll failed" };
|
||||
updateStatus();
|
||||
}
|
||||
}).finally(() => { pollInFlight = null; pollStartedAt = 0; });
|
||||
await pollInFlight;
|
||||
} while (pollQueued);
|
||||
}
|
||||
|
||||
function startStreamingTimer() {
|
||||
if (streamingTimer !== null) return;
|
||||
streamingTimer = setInterval(() => { void poll({ cacheTtl: ACTIVE_CACHE_TTL_MS }); }, STREAMING_POLL_INTERVAL_MS);
|
||||
}
|
||||
|
||||
function stopStreamingTimer() {
|
||||
if (streamingTimer !== null) { clearInterval(streamingTimer); streamingTimer = null; }
|
||||
}
|
||||
|
||||
// ── Lifecycle ────────────────────────────────────────────────────────────
|
||||
|
||||
pi.on("session_start", async (_event, _ctx) => {
|
||||
ctx = _ctx;
|
||||
updateProviderFrom(_ctx.model);
|
||||
await poll();
|
||||
});
|
||||
|
||||
pi.on("session_shutdown", async (_event, _ctx) => {
|
||||
stopStreamingTimer();
|
||||
if (_ctx?.hasUI) _ctx.ui.setStatus(STATUS_KEY, undefined);
|
||||
});
|
||||
|
||||
pi.on("model_select", async (event, _ctx) => {
|
||||
ctx = _ctx;
|
||||
const changed = updateProviderFrom(event.model ?? _ctx.model);
|
||||
if (changed) await poll();
|
||||
});
|
||||
|
||||
pi.on("turn_start", (_event, _ctx) => { ctx = _ctx; updateProviderFrom(_ctx.model); });
|
||||
|
||||
pi.on("before_agent_start", async (_event, _ctx) => {
|
||||
ctx = _ctx;
|
||||
await poll({ cacheTtl: ACTIVE_CACHE_TTL_MS });
|
||||
});
|
||||
|
||||
pi.on("agent_start", (_event, _ctx) => { ctx = _ctx; startStreamingTimer(); });
|
||||
|
||||
pi.on("agent_end", async (_event, _ctx) => {
|
||||
ctx = _ctx;
|
||||
stopStreamingTimer();
|
||||
await poll({ cacheTtl: ACTIVE_CACHE_TTL_MS });
|
||||
});
|
||||
|
||||
pi.events.on("claude-account:switched", () => {
|
||||
const cache = readUsageCache();
|
||||
if (cache?.data?.claude) {
|
||||
const nextCache: import("./core").UsageCache = { ...cache, data: { ...cache.data } };
|
||||
delete nextCache.data.claude;
|
||||
writeUsageCache(nextCache);
|
||||
}
|
||||
void poll({ forceFresh: true });
|
||||
});
|
||||
|
||||
// ── /usage command ───────────────────────────────────────────────────────
|
||||
|
||||
pi.registerCommand("usage", {
|
||||
description: "Show API usage for all subscriptions",
|
||||
handler: async (_args, _ctx) => {
|
||||
ctx = _ctx;
|
||||
updateProviderFrom(_ctx.model);
|
||||
try {
|
||||
if (_ctx?.hasUI) {
|
||||
await _ctx.ui.custom<void>((tui, theme, _keybindings, done) => {
|
||||
return new UsageSelectorComponent(
|
||||
tui, theme, state.activeProvider,
|
||||
() => fetchAllUsages({ endpoints }),
|
||||
() => done(),
|
||||
);
|
||||
});
|
||||
}
|
||||
} finally {
|
||||
await poll({ cacheTtl: ACTIVE_CACHE_TTL_MS });
|
||||
}
|
||||
},
|
||||
});
|
||||
}
|
||||
2
pi/.pi/agent/git/.gitignore
vendored
Normal file
2
pi/.pi/agent/git/.gitignore
vendored
Normal file
@@ -0,0 +1,2 @@
|
||||
*
|
||||
!.gitignore
|
||||
3
pi/.pi/agent/keybindings.json
Normal file
3
pi/.pi/agent/keybindings.json
Normal file
@@ -0,0 +1,3 @@
|
||||
{
|
||||
"app.message.dequeue": ["alt+up", "ctrl+up"]
|
||||
}
|
||||
9
pi/.pi/agent/pi-bar.json
Normal file
9
pi/.pi/agent/pi-bar.json
Normal file
@@ -0,0 +1,9 @@
|
||||
{
|
||||
"segments": [
|
||||
"model",
|
||||
"thinking",
|
||||
"context",
|
||||
"progress",
|
||||
"extensions"
|
||||
]
|
||||
}
|
||||
@@ -1,44 +0,0 @@
|
||||
---
|
||||
description: "Maximum quality pipeline — deep scout, thorough planning, plan review, approval gate, coding, code review"
|
||||
---
|
||||
|
||||
Use the subagent tool to implement with maximum quality. This is for high-risk or architecturally significant changes.
|
||||
|
||||
## Step 1: Deep scout + Plan + Plan review
|
||||
```
|
||||
{ chain: [
|
||||
{ agent: "deep-scout", task: "Deep architectural investigation for: $@\n\nTrace all dependency chains, read tests, check types, understand WHY things are structured the way they are. Map subsystems and their boundaries. Your output enables a complex high-risk change." },
|
||||
{ agent: "planner", task: "Create a detailed implementation plan for: $@\n\nDeep scout context:\n\n{previous}\n\nBe precise: every step must name exact files, functions, and line ranges. Address edge cases and error handling explicitly. Specify which steps can run in parallel. This is a high-risk change — be thorough." },
|
||||
{ agent: "plan-reviewer", task: "Review this plan critically. Verify all file paths, line numbers, and assumptions against the codebase. Check for missing steps, edge cases, and risks.\n\n{previous}" }
|
||||
]}
|
||||
```
|
||||
|
||||
## Step 2: APPROVAL GATE
|
||||
|
||||
**STOP. Present the plan and the plan-reviewer verdict to the user.**
|
||||
|
||||
Show clearly:
|
||||
- The implementation plan (steps, files, risks)
|
||||
- Plan-reviewer's verdict (APPROVED / NEEDS_REVISION / REJECTED) and any issues found
|
||||
- Ask: "Approve this plan, or want changes?"
|
||||
|
||||
Do NOT proceed until the user explicitly approves.
|
||||
If the user requests changes, revise the plan (re-run planner with the feedback) and present again.
|
||||
|
||||
## Step 3: Implement (only after approval)
|
||||
- Use "coder-claude" for the implementation steps
|
||||
- For each coder run, include the approved plan verbatim: "Implement the following plan step(s). Do NOT deviate.\n\n<plan>\n{the approved plan steps}\n</plan>"
|
||||
- For multiple independent steps, run them in parallel using separate coder-claude tasks, each assigned to specific files/plan steps to avoid conflicts
|
||||
|
||||
## Step 4: Code review
|
||||
Run the "reviewer" agent on all changes with this task: "Review all changes made for: $@\n\nCheck for correctness, edge cases, error handling, type safety, and consistency with the approved plan."
|
||||
|
||||
## Step 5: Fix
|
||||
If the reviewer says NEEDS_FIXES, run the "fixer" agent with the review output.
|
||||
|
||||
## Step 6: Report
|
||||
Summarize everything: what was planned, what was implemented, what was reviewed, what was fixed, and any remaining concerns.
|
||||
|
||||
## Failure handling
|
||||
|
||||
- **If any subagent fails, retry it once.** If it fails again, stop and inform the user which agent failed, what the error was, and what had been completed so far. Do NOT continue with remaining steps after a second failure.
|
||||
@@ -1,55 +0,0 @@
|
||||
---
|
||||
description: "Implementation workflow — scouts, plans, waits for approval, then implements"
|
||||
---
|
||||
|
||||
# Task: $@
|
||||
|
||||
Use the subagent tool to implement the task.
|
||||
|
||||
## Step 1: Scout and Plan
|
||||
|
||||
Run this chain:
|
||||
|
||||
```
|
||||
{ chain: [
|
||||
{ agent: "scout", task: "Thoroughly investigate the codebase for: $@\n\nFind all relevant files, types, functions, dependencies, and tests. Report file paths with line ranges. Trace how the pieces connect. Flag anything surprising or risky that an implementer needs to know." },
|
||||
{ agent: "planner", task: "Create a detailed implementation plan for: $@\n\nScout context:\n\n{previous}\n\nBe precise: every step must name exact files, functions, and line ranges. Address edge cases and error handling. Specify which steps can run in parallel." },
|
||||
{ agent: "plan-reviewer", task: "Review this plan critically. Verify all file paths, line numbers, and assumptions against the codebase. Check for missing steps, edge cases, and risks.\n\n{previous}" }
|
||||
]}
|
||||
```
|
||||
|
||||
## Step 2: Approval gate
|
||||
|
||||
**STOP. Present the plan AND the plan-reviewer verdict to the user. Ask for approval before continuing.**
|
||||
|
||||
When presenting, highlight: what will change, which files, risks, and the plan-reviewer verdict.
|
||||
|
||||
If the user requests changes to the plan, revise and present again before implementing.
|
||||
|
||||
## Step 3: Implement
|
||||
|
||||
Once approved:
|
||||
- Execute the plan steps using the "coder" agent
|
||||
- When running coder, always wrap the plan step(s) in the task: "Implement the following plan step(s). Do NOT deviate.\n\n<plan>\n{the approved plan steps}\n</plan>"
|
||||
- Run the "reviewer" agent on all changes
|
||||
- If NEEDS_FIXES, run the "fixer" agent with the review output
|
||||
|
||||
## Step 4: Summary
|
||||
|
||||
After the final step, summarize: what was done, what files changed, what was reviewed, and any remaining concerns.
|
||||
|
||||
## Important
|
||||
|
||||
- **NEVER skip the approval gate**. Always present the plan and wait.
|
||||
- Always pass scout context forward using {previous} in chain mode — this is how the planner and plan-reviewer receive the scout's findings.
|
||||
- When running the coder, always include the approved plan verbatim in the task so the coder has full context.
|
||||
|
||||
## Agent Failure and Fallback
|
||||
|
||||
When a subagent returns empty output or an error (rate limit, credit exhaustion, connection failure):
|
||||
|
||||
1. **Retry once** with the same agent and model — transient failures are common.
|
||||
2. **If still failing, retry with the cross-family fallback model** using the `model` override parameter. See the fallback table in the `subagent-implement` SKILL.md for the current primary/fallback mapping.
|
||||
3. **If the fallback also fails**, do the work yourself (read the relevant files directly and produce the scout/plan/review output inline). Inform the user which agent failed, what error was returned, and what you did instead.
|
||||
|
||||
Do NOT silently absorb failures. Always surface them to the user even when working around them.
|
||||
@@ -1,23 +0,0 @@
|
||||
---
|
||||
description: "Scout + plan + plan review — no implementation"
|
||||
---
|
||||
|
||||
Use the subagent tool with a chain to plan (but NOT implement) the following:
|
||||
|
||||
```
|
||||
{ chain: [
|
||||
{ agent: "scout", task: "Find all code relevant to: $@" },
|
||||
{ agent: "planner", task: "Create a detailed implementation plan for: $@\n\nContext from scout:\n\n{previous}" },
|
||||
{ agent: "plan-reviewer", task: "Review this implementation plan. Verify file paths and line numbers against the actual codebase.\n\n{previous}" }
|
||||
]}
|
||||
```
|
||||
|
||||
Present the plan and the review to me. Do NOT proceed to implementation.
|
||||
|
||||
## Agent Failure and Fallback
|
||||
|
||||
If any agent returns empty output or an error (rate limit, credit exhaustion, connection failure):
|
||||
|
||||
1. Retry once with the same agent.
|
||||
2. If still failing, retry with the cross-family fallback using the `model` override. See the fallback table in the `subagent-plan` SKILL.md for the current mapping.
|
||||
3. If the fallback also fails, do the work yourself and tell me which agent failed and why.
|
||||
@@ -1,10 +0,0 @@
|
||||
---
|
||||
description: "Opus code review on recent changes or specified files"
|
||||
---
|
||||
|
||||
Use the subagent tool to run the "reviewer" agent with this task:
|
||||
|
||||
Review the following: $@
|
||||
|
||||
If no specific files are mentioned, review recent git changes (`git diff` and `git diff --staged`).
|
||||
Report the review results. If the verdict is NEEDS_FIXES, ask if I want you to run the "fixer" agent to apply them.
|
||||
@@ -1,110 +0,0 @@
|
||||
---
|
||||
name: add-agent
|
||||
description: "Add a new Claude agent definition to ~/.claude/agents/ and an accompanying skill to ~/.pi/agent/skills/. Use when creating new specialized agents."
|
||||
---
|
||||
|
||||
# Add Agent
|
||||
|
||||
When the user wants to create a new specialized Claude agent:
|
||||
|
||||
## Agent File Convention
|
||||
|
||||
Create a `.md` file in `~/.claude/agents/` with this structure:
|
||||
|
||||
```markdown
|
||||
---
|
||||
name: agent_name
|
||||
description: One-line description of what the agent does
|
||||
tools: Read, Bash[, Edit, Write] # include only what's needed
|
||||
model: sonnet | opus
|
||||
---
|
||||
|
||||
You are an [role]. [1-2 sentence description of purpose].
|
||||
|
||||
Available tools:
|
||||
- read: Read file contents
|
||||
- bash: Execute bash commands
|
||||
- edit: Make surgical edits to files # if applicable
|
||||
- write: Create or overwrite files # if applicable
|
||||
|
||||
Guidelines:
|
||||
- Use bash for file operations: prefer `rg` over grep, `fd` over find, glob patterns for batch file matching
|
||||
- Use read to examine files [before editing] # adapt phrasing
|
||||
- [Agent-specific guidelines]
|
||||
- When summarizing your actions, output plain text directly - do NOT use cat or bash to display what you did
|
||||
- Be concise in your responses
|
||||
- Show file paths clearly when working with files
|
||||
```
|
||||
|
||||
## Key Rules
|
||||
|
||||
1. **`name`** in frontmatter must match what skills reference as `agent: "name"`
|
||||
2. **`tools`** — only include tools the agent actually needs. Read-only agents use `Read, Bash`
|
||||
3. **`model`** — `sonnet` for execution/review, `opus` for strategic/advisory work
|
||||
4. **Always include** the "output plain text directly" guideline — agents without it tend to use `cat`/`echo` instead of responding directly
|
||||
5. **Bash guideline** should read: `prefer \`rg\` over grep, \`fd\` over find, glob patterns for batch file matching`
|
||||
6. **Bash tool description** should be: `Execute bash commands`
|
||||
|
||||
## Skill File Convention
|
||||
|
||||
Create `~/.pi/agent/skills/<skill-name>/SKILL.md`:
|
||||
|
||||
```markdown
|
||||
---
|
||||
name: skill-name
|
||||
description: "What triggers this skill. Use when [condition]."
|
||||
---
|
||||
|
||||
# Skill Title
|
||||
|
||||
When [trigger condition]:
|
||||
|
||||
## What to include in the prompt
|
||||
|
||||
1. **[Section 1]** — description
|
||||
2. **[Section 2]** — description
|
||||
...
|
||||
|
||||
The `agent_name` agent has `Tool1` and `Tool2` tools only — [what it can/can't do].
|
||||
|
||||
## How to call
|
||||
|
||||
\```
|
||||
ask_claude({
|
||||
agent: "agent_name",
|
||||
question: "Specific instruction for the agent.",
|
||||
prompt: `
|
||||
## Section 1
|
||||
[Template]
|
||||
|
||||
## Section 2
|
||||
[Template]
|
||||
`
|
||||
})
|
||||
\```
|
||||
|
||||
## After the review
|
||||
|
||||
- [What to do with the agent's output]
|
||||
- [How to summarize for the user]
|
||||
- [When to loop back or escalate]
|
||||
```
|
||||
|
||||
## Skill Key Rules
|
||||
|
||||
1. **Skill `name`** should be kebab-case (e.g., `claude-debug`, not `claude_debug`)
|
||||
2. **Description** must be in quotes if it contains special characters
|
||||
3. **`agent:` in `ask_claude()`** must match the agent's `name:` exactly
|
||||
4. **Include a "What to include" section** — gives the calling agent a template
|
||||
5. **Include a "How to call" section** — with a concrete `ask_claude()` example
|
||||
6. **Include an "After the review" section** — what to do with the output
|
||||
|
||||
## Existing Agents (for reference)
|
||||
|
||||
| Agent | Model | Tools | Purpose |
|
||||
|-------|-------|-------|---------|
|
||||
| `minimal` | sonnet | Read, Bash, Edit, Write | General coding |
|
||||
| `code_review` | sonnet | Read, Bash, Edit, Write | Review & fix code |
|
||||
| `plan_review` | opus | Read, Bash | Review plans |
|
||||
| `debug` | sonnet | Read, Bash | Trace bugs |
|
||||
| `oracle` | opus | Read, Bash | Strategic guidance |
|
||||
@@ -1,153 +0,0 @@
|
||||
---
|
||||
name: ask-claude
|
||||
description: "Invoke Claude (any agent or model) for an opinion, review, or analysis. Use when the user asks you to ask a specific Claude about something specific — e.g. trade-offs of an approach, a focused review, or a second opinion."
|
||||
---
|
||||
|
||||
# Ask Claude
|
||||
|
||||
Use `ask_claude` to consult Claude with any combination of agent, model, question, and prompt.
|
||||
|
||||
## Parameters
|
||||
|
||||
| Parameter | Required | Description |
|
||||
|-----------|----------|-------------|
|
||||
| `prompt` | yes | The content for Claude to reason about (context, code, plan, bug report, etc.) |
|
||||
| `question` | no | What specifically to ask — prepended as a focused review directive |
|
||||
| `agent` | no | Agent name from `~/.claude/agents/`. See **Available agents** below. |
|
||||
| `model` | no | Override the model: `"opus"`, `"sonnet"`, `"haiku"` |
|
||||
| `session_id` | no | Resume a prior conversation (returned in every response) |
|
||||
|
||||
If neither `agent` nor `model` is set, defaults to Claude Sonnet.
|
||||
|
||||
> **Multi-turn:** Every response includes a `session_id`. Pass it back in a subsequent `ask_claude()` call to continue the same conversation with the same agent/model.
|
||||
|
||||
## Available agents
|
||||
|
||||
See `~/.claude/agents/` for the full list. Common agents:
|
||||
|
||||
| Agent | Model | Tools | Use when |
|
||||
|-------|-------|-------|----------|
|
||||
| `plan_review` | Opus | Read, Bash | Reviewing plans for correctness, completeness, feasibility, and risk |
|
||||
| `code_review` | Sonnet | Read, Bash, Edit, Write | Reviewing implementations; can apply fixes directly |
|
||||
| `debug` | Sonnet | Read, Bash | Tracing bugs and root causes; will NOT apply fixes |
|
||||
| `oracle` | Opus | Read, Bash | Hard problems, architectural decisions, when you're stuck |
|
||||
|
||||
Pick the agent that matches the task. If unsure, ask the user which agent to use, or use a raw `model=` instead of an agent.
|
||||
|
||||
## Common patterns
|
||||
|
||||
### Plan review
|
||||
```
|
||||
ask_claude({
|
||||
agent: "plan_review",
|
||||
question: "Review for correctness, completeness, feasibility, and risk. Highlight missing steps or unclear requirements.",
|
||||
prompt: `
|
||||
## Project Context
|
||||
[Key facts from CLAUDE.md]
|
||||
|
||||
## Codebase Exploration
|
||||
[Modules/files you read]
|
||||
|
||||
## Implementation Plan
|
||||
[Your full plan, step by step]
|
||||
|
||||
## Open Questions
|
||||
[Anything uncertain]
|
||||
`
|
||||
})
|
||||
```
|
||||
|
||||
### Code review (with fix capability)
|
||||
```
|
||||
ask_claude({
|
||||
agent: "code_review",
|
||||
question: "Review for bugs, architectural issues, style, and correctness. Apply fixes for any issues you find.",
|
||||
prompt: `
|
||||
## Project Conventions
|
||||
[Relevant context not in CLAUDE.md]
|
||||
|
||||
## What Was Implemented
|
||||
[Brief description]
|
||||
|
||||
## Plan That Was Followed
|
||||
[The implementation plan]
|
||||
|
||||
## Files Changed
|
||||
- path/to/file.ts — [what changed]
|
||||
|
||||
## Code to Review
|
||||
[Paste key sections OR instruct the agent to read the files above]
|
||||
`
|
||||
})
|
||||
```
|
||||
|
||||
### Debugging
|
||||
```
|
||||
ask_claude({
|
||||
agent: "debug",
|
||||
question: "Trace the root cause. Provide specific file paths, line numbers, and the exact code responsible. Suggest a fix but do not apply it.",
|
||||
prompt: `
|
||||
## The Issue
|
||||
[What's happening vs. what should happen]
|
||||
|
||||
## Error Output
|
||||
[Any logs, stack traces, or error messages]
|
||||
|
||||
## Relevant Files
|
||||
- path/to/file.ts — [why it's relevant]
|
||||
|
||||
## What I've Already Tried
|
||||
[Hypotheses tested, things ruled out]
|
||||
`
|
||||
})
|
||||
```
|
||||
|
||||
### Oracle (hard problems)
|
||||
```
|
||||
ask_claude({
|
||||
agent: "oracle",
|
||||
question: "Analyze this problem and provide guidance. Explain your reasoning and recommend the best path forward.",
|
||||
prompt: `
|
||||
## The Problem
|
||||
[What you're trying to solve and why it's hard]
|
||||
|
||||
## Codebase Context
|
||||
[Relevant files, modules, and patterns explored]
|
||||
|
||||
## Options Considered
|
||||
[Approaches evaluated and their trade-offs]
|
||||
|
||||
## What I Need Clarified
|
||||
[Specific question or decision point]
|
||||
`
|
||||
})
|
||||
```
|
||||
|
||||
### Free-form model question (no agent)
|
||||
```
|
||||
ask_claude({
|
||||
model: "opus",
|
||||
question: "What are the trade-offs of this approach vs using a message queue?",
|
||||
prompt: `We're considering polling a database table every 5 seconds for new jobs instead of a queue...`
|
||||
})
|
||||
```
|
||||
|
||||
### Continuing a conversation
|
||||
When Claude returns a `session_id` in its response, pass it back to continue the same conversation:
|
||||
```
|
||||
ask_claude({
|
||||
agent: "code_review",
|
||||
question: "Apply the fixes you identified.",
|
||||
session_id: "<session_id from prior response>",
|
||||
prompt: "Please proceed with the fixes."
|
||||
})
|
||||
```
|
||||
|
||||
## After the response
|
||||
|
||||
- Summarize Claude's key points for the user.
|
||||
- If Claude raises blockers or important concerns, address them before proceeding.
|
||||
- For plan reviews: if blockers or missing steps are flagged, update the plan and re-invoke if changes are significant.
|
||||
- For code reviews: Claude (via `code_review`) may apply fixes directly — summarize what was changed.
|
||||
- For debugging: apply the fix yourself or delegate to a coding agent.
|
||||
- For oracle: proceed with the recommended approach, or loop back if more clarification is needed.
|
||||
98
pi/.pi/agent/skills/godot-rag/SKILL.md
Normal file
98
pi/.pi/agent/skills/godot-rag/SKILL.md
Normal file
@@ -0,0 +1,98 @@
|
||||
---
|
||||
name: godot-rag
|
||||
description: Search Godot Engine documentation (class reference, tutorials, engine internals, addons) using the godot-rag CLI. Use when the user asks Godot API questions, needs GDScript syntax, wants to understand Godot classes/methods/signals, or is working on a Godot game development project. Supports natural language queries and JSON output for programmatic use.
|
||||
---
|
||||
|
||||
# godot-rag — Godot Documentation Search CLI
|
||||
|
||||
A hybrid RAG search tool for Godot 4.x documentation. Query class references, tutorials, engine internals, and addon documentation from your terminal with natural language.
|
||||
|
||||
## Setup
|
||||
|
||||
First-time install (run once):
|
||||
|
||||
```bash
|
||||
pip install godot-rag
|
||||
```
|
||||
|
||||
Verify it works:
|
||||
|
||||
```bash
|
||||
godot-rag --help
|
||||
```
|
||||
|
||||
The first search may take a moment as it downloads the documentation index. Subsequent queries are fast (~7ms).
|
||||
|
||||
## Usage — Agent Patterns
|
||||
|
||||
All queries support `--json` for structured output that's easy to parse.
|
||||
|
||||
### Search class reference (API docs)
|
||||
|
||||
```bash
|
||||
# Exact symbol lookup
|
||||
godot-rag s-class "Node.add_child" --json
|
||||
godot-rag s-class "Signal.emit" --limit 3 --json
|
||||
godot-rag s-class "Vector3.normalized" --json
|
||||
godot-rag s-class "StringName.is_valid_filename" --json
|
||||
|
||||
# Broad class search
|
||||
godot-rag s-class "CharacterBody2D" --json
|
||||
```
|
||||
|
||||
### Search tutorials
|
||||
|
||||
```bash
|
||||
# Natural language queries work well here
|
||||
godot-rag s-tutorial "how to use signals" --json
|
||||
godot-rag s-tutorial "2D pathfinding" --json
|
||||
godot-rag s-tutorial "scene tree" --json
|
||||
godot-rag s-tutorial "create a 2d character" --json
|
||||
```
|
||||
|
||||
### Search engine internals
|
||||
|
||||
```bash
|
||||
godot-rag s-engine "GDExtension" --json
|
||||
godot-rag s-engine "IDE debugging" --limit 3 --json
|
||||
```
|
||||
|
||||
### Search all docs at once
|
||||
|
||||
```bash
|
||||
godot-rag s "Timer" --json
|
||||
godot-rag s "physics interpolation" --json
|
||||
```
|
||||
|
||||
### Search addon documentation
|
||||
|
||||
```bash
|
||||
godot-rag s-addon "state machine" --json
|
||||
godot-rag s-addon "dialogue balloon" --limit 3 --json
|
||||
|
||||
# Filter by specific addon
|
||||
godot-rag s-addon "change_scene" --addon scene_manager --json
|
||||
godot-rag s-addon "BehaviorTree" --addon limboai --json
|
||||
```
|
||||
|
||||
## Typical Agent Workflow
|
||||
|
||||
When the user asks a Godot API or implementation question:
|
||||
|
||||
1. **Search tutorials first** for conceptual understanding:
|
||||
```bash
|
||||
godot-rag s-tutorial "<question topic>" --json
|
||||
```
|
||||
|
||||
2. **Then look up specific API details**:
|
||||
```bash
|
||||
godot-rag s-class "<ClassName.method>" --json
|
||||
```
|
||||
|
||||
3. Use the returned documentation to write accurate code.
|
||||
|
||||
## Notes
|
||||
|
||||
- Queries return 3 results by default (adjust with `--limit N`)
|
||||
- Without `--json`, output is human-readable terminal text with formatted sections
|
||||
- The local database stays synced with your installed Godot version; run `pip install --upgrade godot-rag` to update
|
||||
@@ -1,50 +0,0 @@
|
||||
---
|
||||
name: implementor
|
||||
description: "Delegate implementation of a fix or feature to the implementor agent. Use when you have a clear plan, bug fix, or feature to implement and want a focused agent to handle the full coding + build-verify cycle."
|
||||
---
|
||||
|
||||
# Implementor
|
||||
|
||||
When you need to implement a fix, feature, or refactoring with full build verification:
|
||||
|
||||
## What to include in the prompt
|
||||
|
||||
1. **Goal** — what needs to be built or fixed, in 1-2 sentences
|
||||
2. **Context** — relevant files, functions, or modules (include the actual code or file paths)
|
||||
3. **Plan** — step-by-step description of the changes to make (if you have one)
|
||||
4. **Constraints** — any project conventions, patterns to follow, or things to avoid
|
||||
5. **Verification** — how to confirm the implementation is correct (build command, test names, etc.)
|
||||
|
||||
The `implementor` agent has `Read`, `Bash`, `Edit`, and `Write` tools only — it can read code, run commands, and make changes. It cannot use semantic search or web search.
|
||||
|
||||
## How to call
|
||||
|
||||
```
|
||||
ask_claude({
|
||||
agent: "implementor",
|
||||
question: "Implement [brief description of the fix or feature].",
|
||||
prompt: `
|
||||
## Goal
|
||||
[What needs to be done]
|
||||
|
||||
## Context
|
||||
[Relevant files and code snippets]
|
||||
|
||||
## Plan
|
||||
[Step-by-step changes, or "investigate and determine the best approach"]
|
||||
|
||||
## Constraints
|
||||
[Project conventions, patterns to follow, things to avoid]
|
||||
|
||||
## Verification
|
||||
[Build command, tests to run, how to confirm correctness]
|
||||
`
|
||||
})
|
||||
```
|
||||
|
||||
## After the review
|
||||
|
||||
- Check whether the agent reported successful build/verification
|
||||
- If it failed, either re-invoke with the error details or fix the remaining issues directly
|
||||
- Summarize for the user: what files changed, what was verified, and any remaining concerns
|
||||
- For complex implementations, consider running `diagnostics` sub-agent or relevant tests afterward as a second pass
|
||||
@@ -1,29 +0,0 @@
|
||||
---
|
||||
name: local-scout
|
||||
description: "Delegates codebase exploration to the local scout subagent (runs on a fast local Qwen model with QMD + opty tools). Load this skill only when the user explicitly asks to use scout, use local scout, or scout a task. Do NOT load automatically for general exploration — only when scout is explicitly requested."
|
||||
---
|
||||
|
||||
# Local Scout
|
||||
|
||||
Delegate codebase exploration to the `scout` subagent, which runs on a fast local model (Qwen) augmented with semantic search via QMD and HDC-indexed context retrieval via opty. It is cheap, fast, and keeps the main context clean.
|
||||
|
||||
## When to use
|
||||
|
||||
- User says "use scout to find …", "scout: …", or "use local scout"
|
||||
- You need to gather broad codebase context before planning
|
||||
- The task is primarily "look around the codebase" rather than making precise edits
|
||||
|
||||
## How to invoke
|
||||
|
||||
```javascript
|
||||
subagent({ agent: "scout", task: "Find and summarize the authentication flow" })
|
||||
```
|
||||
|
||||
The scout writes its findings to `context.md` and returns a summary. Use the summary or read `context.md` for the full structured output.
|
||||
|
||||
## Tips
|
||||
|
||||
- Be specific in the task description — the scout infers thoroughness from it
|
||||
- For deep traces, prefix with "Thorough:" e.g. `"Thorough: trace all usages of X"`
|
||||
- For quick lookups, prefix with "Quick:" e.g. `"Quick: where is the config loaded?"`
|
||||
- Do your own reading only when you need precise line-level content to reference in your response
|
||||
@@ -1,37 +0,0 @@
|
||||
---
|
||||
name: opty
|
||||
description: Semantic code search using HDC (Hyperdimensional Computing). Finds functions, types, imports, and variables in the indexed codebase by meaning. Use when searching for code by concept — e.g. "error handling functions", "database types", "authentication flow".
|
||||
compatibility: Requires opty CLI and a running daemon. Install from source or binary.
|
||||
allowed-tools: Bash(opty:*)
|
||||
---
|
||||
|
||||
# Opty — HDC Code Search
|
||||
|
||||
Semantic code search via Hyperdimensional Computing. Indexes the codebase and finds relevant functions, types, and imports by meaning rather than exact text. The daemon auto-starts when you run any opty command.
|
||||
|
||||
## Query
|
||||
|
||||
```bash
|
||||
opty query "natural language description of what you're looking for"
|
||||
```
|
||||
|
||||
Output is TOON format — compact, LLM-optimized:
|
||||
```
|
||||
functions[N]{name,signature,file,line}:
|
||||
functionName,signature,path/to/file.ts,42
|
||||
...
|
||||
```
|
||||
|
||||
## Other Commands
|
||||
|
||||
```bash
|
||||
opty status # Show indexed file/unit count and watched directory
|
||||
opty reindex # Force re-scan after major file changes or stale results
|
||||
```
|
||||
|
||||
## Tips
|
||||
|
||||
- Query by concept, not exact names: `"connection pool exhaustion"` not `"ConnectionPoolError"`
|
||||
- Results include file + line number — use `read` to get the actual code
|
||||
- `opty status` shows which directory is indexed (the daemon is per-project)
|
||||
- After large refactors, run `opty reindex` if results seem stale
|
||||
@@ -1,89 +0,0 @@
|
||||
---
|
||||
name: qmd
|
||||
description: Search markdown knowledge bases, notes, and documentation using the qmd CLI. Use when searching notes, finding documents, or looking up information in the knowledge base.
|
||||
license: MIT
|
||||
compatibility: Requires qmd CLI. Install via `npm install -g @tobilu/qmd`.
|
||||
metadata:
|
||||
author: tobi
|
||||
version: "2.0.0"
|
||||
allowed-tools: Bash(qmd:*)
|
||||
---
|
||||
|
||||
# QMD - Quick Markdown Search
|
||||
|
||||
Local search engine for markdown content. Use via the `qmd` CLI (bash).
|
||||
|
||||
## Status
|
||||
|
||||
!`qmd status 2>/dev/null || echo "Not installed: npm install -g @tobilu/qmd"`
|
||||
|
||||
## CLI Usage
|
||||
|
||||
```bash
|
||||
qmd query "question" # Auto-expand + rerank (recommended)
|
||||
qmd query $'lex: X\nvec: Y' # Structured query document
|
||||
qmd query $'expand: question' # Explicit expand
|
||||
qmd search "keywords" # BM25 only (no LLM)
|
||||
qmd get "path/to/file.md" # Retrieve a document by path
|
||||
qmd get "#abc123" # Retrieve by docid
|
||||
qmd get "path/to/file.md:100" -l 50 # Retrieve starting at line 100, 50 lines
|
||||
qmd multi-get "journals/2026-*.md" -l 40 # Batch fetch by glob
|
||||
qmd multi-get notes/foo.md,notes/bar.md # Comma-separated list
|
||||
```
|
||||
|
||||
## Query Types
|
||||
|
||||
| Type | Method | Input |
|
||||
|------|--------|-------|
|
||||
| `lex` | BM25 | Keywords — exact terms, names, code |
|
||||
| `vec` | Vector | Question — natural language |
|
||||
| `hyde` | Vector | Answer — hypothetical result (50-100 words) |
|
||||
|
||||
## Writing Good Queries
|
||||
|
||||
**lex (keyword)**
|
||||
- 2-5 terms, no filler words
|
||||
- Exact phrase: `"connection pool"` (quoted)
|
||||
- Exclude terms: `performance -sports` (minus prefix)
|
||||
- Code identifiers work: `handleError async`
|
||||
|
||||
**vec (semantic)**
|
||||
- Full natural language question
|
||||
- Be specific: `"how does the rate limiter handle burst traffic"`
|
||||
|
||||
**hyde (hypothetical document)**
|
||||
- Write 50-100 words of what the *answer* looks like
|
||||
- Use the vocabulary you expect in the result
|
||||
|
||||
**expand (auto-expand)**
|
||||
- Single-line query (implicit) or `expand: question`
|
||||
- Lets a local LLM generate lex/vec/hyde variations
|
||||
- Do not mix with other typed lines
|
||||
|
||||
## Combining Types
|
||||
|
||||
| Goal | Approach |
|
||||
|------|----------|
|
||||
| Know exact terms | `lex` only |
|
||||
| Don't know vocabulary | Single-line (implicit expand) or `vec` |
|
||||
| Best recall | `lex` + `vec` |
|
||||
| Complex topic | `lex` + `vec` + `hyde` |
|
||||
|
||||
First query gets 2x weight — put your strongest signal first.
|
||||
|
||||
## Lex Syntax
|
||||
|
||||
| Syntax | Meaning | Example |
|
||||
|--------|---------|---------|
|
||||
| `term` | Prefix match | `perf` matches "performance" |
|
||||
| `"phrase"` | Exact phrase | `"rate limiter"` |
|
||||
| `-term` | Exclude | `performance -sports` |
|
||||
|
||||
## Collection Filtering
|
||||
|
||||
```bash
|
||||
qmd query --collections docs "question"
|
||||
qmd query --collections docs,notes "question"
|
||||
```
|
||||
|
||||
Omit to search all collections.
|
||||
114
pi/.pi/agent/skills/rustdoc-rag/SKILL.md
Normal file
114
pi/.pi/agent/skills/rustdoc-rag/SKILL.md
Normal file
@@ -0,0 +1,114 @@
|
||||
---
|
||||
name: rustdoc-rag
|
||||
description: Search and browse local rustdoc JSON for any crate in the project's target/doc — per-crate (summary/search/item/source) and cross-crate (implementors of a trait, traits a type implements, full method surface with defining traits, reverse signature usage, canonical import paths, cross-crate symbol search). Use when the user asks about a crate's API, needs signatures/docs, wants the correct import path for a symbol, asks "what implements X" / "what can I call on Y" / "who uses Z", or needs to find where an item is defined in source. Warns when docs are stale vs Cargo.lock. No internet needed.
|
||||
---
|
||||
|
||||
# rustdoc-rag — local rustdoc JSON search/browse
|
||||
|
||||
A CLI over `target/doc/<crate>.json` files (produced by the `rustdoc-regen`
|
||||
skill). Pure Python 3 stdlib, lives next to this SKILL.md:
|
||||
|
||||
```bash
|
||||
python3 rustdoc-rag.py --help
|
||||
```
|
||||
|
||||
Doc dir: `./target/doc` by default, or `$RUSTDOC_RAG_DIR`, or `--doc-dir <path>`
|
||||
(global flag, goes **before** the subcommand). Other shared flags (`--json`,
|
||||
`--limit N`, `--exact`) go **after** the subcommand. Crate names accept dashes
|
||||
or underscores.
|
||||
|
||||
Cross-crate commands use a **sidecar index**
|
||||
(`<doc-dir>/.rustdoc-rag-index.json.gz`) that is built automatically on first
|
||||
use (~6 s for 300 crates) and refreshed incrementally whenever a crate's JSON
|
||||
is newer. You never need to manage it; `index --rebuild` forces a full rebuild.
|
||||
|
||||
## Command Cheat Sheet
|
||||
|
||||
Per-crate (need a crate name):
|
||||
|
||||
```bash
|
||||
python3 rustdoc-rag.py list # crates with JSON
|
||||
python3 rustdoc-rag.py summary bevy_app # compact catalog by kind
|
||||
python3 rustdoc-rag.py search bevy_ecs "query" --limit 5 # name/docs match, rendered
|
||||
python3 rustdoc-rag.py item bevy_app App # full signature + docs + fields
|
||||
python3 rustdoc-rag.py source bevy_app Plugin # file:line → then `read` it
|
||||
```
|
||||
|
||||
Cross-crate (take a bare name or any `::`-path — canonical, re-export, or
|
||||
suffix; aliases like `pub use Person as Human` also resolve):
|
||||
|
||||
```bash
|
||||
python3 rustdoc-rag.py search --all Equivalent # where does this symbol live?
|
||||
python3 rustdoc-rag.py implementors Plugin # every type implementing a trait
|
||||
python3 rustdoc-rag.py trait_impls Entity # every trait a type implements
|
||||
python3 rustdoc-rag.py methods Query # full method surface (see below)
|
||||
python3 rustdoc-rag.py used_in Velocity # every fn/method/field whose
|
||||
# signature mentions the type
|
||||
python3 rustdoc-rag.py canonical App # canonical path + all pub-use
|
||||
# re-export paths (aka `where`)
|
||||
python3 rustdoc-rag.py freshness # doc JSONs vs Cargo.lock
|
||||
python3 rustdoc-rag.py index # (re)build sidecar explicitly
|
||||
```
|
||||
|
||||
## Which command answers which question
|
||||
|
||||
- **"What implements `Plugin`?"** → `implementors Plugin`. Same-named traits in
|
||||
different crates each get their own block (no false positives) — pass
|
||||
`crate::path` to pin one. `--crate <name>` restricts the scan.
|
||||
- **"What can I call on this type?"** → `methods <Type>`. One call returns
|
||||
inherent methods **plus** trait methods grouped by the trait that defines
|
||||
them (`## impl Read — std::io::Read`), so you know which import each method
|
||||
needs. Provided (default) trait methods are resolved from the trait's
|
||||
defining crate and marked `(provided)`. Type aliases are followed one hop.
|
||||
Auto/blanket impls (`Any`, `Borrow`, `Into`, …) are hidden unless `--full`.
|
||||
`methods <Trait>` prints the trait's declared methods.
|
||||
- **"Which traits does `Entity` implement?"** → `trait_impls Entity`.
|
||||
- **"What's the right import?"** → `canonical <Symbol>` (alias: `where`).
|
||||
Prints the canonical defining path and every public re-export path
|
||||
(preludes, root re-exports, inlined re-exports), shortest first. Use this
|
||||
before hand-writing an import — the canonical path is not always importable.
|
||||
- **"Who depends on this type?"** → `used_in <Type>` (alias:
|
||||
`used_in_signatures`). Reverse relation over resolved signature graphs —
|
||||
survives generics (`Query<With<Foliage>>`), type aliases, and re-exports
|
||||
that defeat text search. `--crate <name>` restricts.
|
||||
- **"Where is `Velocity` defined at all?"** → `search --all Velocity`
|
||||
(add `--exact` for exact names). Output is one `crate::path (kind)` line per
|
||||
match — drill in with `item`/`methods` afterwards.
|
||||
|
||||
Ambiguous bare names print the candidate paths and stop — rerun with the full
|
||||
path. Case-sensitive matches shadow case-insensitive ones.
|
||||
|
||||
## Freshness
|
||||
|
||||
Every per-crate command and `methods` annotate their output with
|
||||
`⚠ stale: <crate>.json is vX, Cargo.lock has vY — regenerate: …` when the doc
|
||||
JSON no longer matches Cargo.lock (nearest lock to the doc dir / cwd). Trust
|
||||
the warning: regenerate with the `rustdoc-regen` skill
|
||||
(`rustdoc-regen.py crates <name>`), which also triggers a sidecar refresh on
|
||||
next query. `freshness` gives the whole-substrate picture (stale +
|
||||
undocumented counts); `--json` for machine-readable output.
|
||||
|
||||
## Typical Agent Workflow
|
||||
|
||||
1. Locate the symbol: `search --all <name>` (or `canonical <name>` if you
|
||||
only need the import).
|
||||
2. Understand it: `methods <Type>` for the callable surface, `item <crate>
|
||||
<Symbol>` for docs/fields, `implementors <Trait>` / `trait_impls <Type>`
|
||||
for the trait graph.
|
||||
3. Read the real code when needed: `source <crate> <Symbol>` → `read` the
|
||||
file:line (paths point into `~/.cargo/registry/src/...`).
|
||||
|
||||
## Notes
|
||||
|
||||
- All commands support `--json` (structured output; per-crate commands embed
|
||||
a `"stale"` field).
|
||||
- Only public items are cataloged. Impl-method lists exclude auto-derived
|
||||
noise by design; `item` caps method lists at 60 — use `methods` for the
|
||||
complete, trait-resolved surface.
|
||||
- Default `--limit`: 5 (search/source), 25 (search --all), 50 (used_in),
|
||||
100 (implementors/trait_impls).
|
||||
- The sidecar stores per crate: public items with canonical paths, trait
|
||||
impls, re-exports, type-alias targets, and the set of type paths mentioned
|
||||
in signatures — cross-crate queries are a single in-memory scan; only
|
||||
`used_in`/`methods` open the (few) relevant crate JSONs.
|
||||
- If output ever looks wrong after regenerating docs, `index --rebuild`.
|
||||
1579
pi/.pi/agent/skills/rustdoc-rag/rustdoc-rag.py
Executable file
1579
pi/.pi/agent/skills/rustdoc-rag/rustdoc-rag.py
Executable file
File diff suppressed because it is too large
Load Diff
159
pi/.pi/agent/skills/rustdoc-regen/SKILL.md
Normal file
159
pi/.pi/agent/skills/rustdoc-regen/SKILL.md
Normal file
@@ -0,0 +1,159 @@
|
||||
---
|
||||
name: rustdoc-regen
|
||||
description: Rebuild the local rustdoc JSON substrate that the rustdoc-rag skill reads. Use when a new dependency was added to the project, when rustdoc-rag reports "no rustdoc JSON for crate '<x>'", when the rustdoc-rag substrate is stale (Cargo.lock newer than target/doc/*.json), or when the user asks to regenerate/update rustdoc docs or "get docs for the new dependency". Emits `<crate>.json` under `target/doc` on *stable* Rust via `RUSTC_BOOTSTRAP=1`. General for any Rust project; not limited to a specific workspace.
|
||||
---
|
||||
|
||||
# rustdoc-regen — rebuild local rustdoc JSON for rustdoc-rag
|
||||
|
||||
`rustdoc-rag` only *reads* `<crate>.json` files. This skill *produces* them.
|
||||
|
||||
rustdoc's JSON output (`--output-format json`) is an unstable feature normally
|
||||
gated behind nightly. This skill uses the standard `RUSTC_BOOTSTRAP=1` escape
|
||||
hatch (set in the cargo process environment) so **stable** Rust emits JSON docs
|
||||
too — no `nightly` toolchain, no index download, no network beyond what `cargo`
|
||||
already needs to fetch/build crates. One crate per JSON file lands in
|
||||
`<doc-dir>/<crate>.json` (dashes → underscores, same convention rustdoc-rag
|
||||
expects).
|
||||
|
||||
The script is pure Python 3 (stdlib only) and lives next to this SKILL.md:
|
||||
|
||||
```bash
|
||||
python3 rustdoc-regen.py --help
|
||||
```
|
||||
|
||||
Default `<doc-dir>` resolution matches `rustdoc-rag`: `--doc-dir <path>` flag,
|
||||
else `$RUSTDOC_RAG_DIR`, else `./target/doc`.
|
||||
|
||||
## Subcommands
|
||||
|
||||
### `all` — full rebuild (workspace + every dep)
|
||||
|
||||
```bash
|
||||
python3 rustdoc-regen.py all
|
||||
python3 rustdoc-regen.py all --keep-going # don't abort on one crate failing
|
||||
python3 rustdoc-regen.py all --offline
|
||||
```
|
||||
|
||||
Runs `cargo doc` over the whole dependency graph with JSON output. Re-emits a
|
||||
JSON for every crate that builds on this host **and** every workspace member.
|
||||
Slow on a cold graph (minutes for Bevy-scale), but it's the only way to guarantee
|
||||
the substrate matches the live graph exactly. Use after a `Cargo.toml`/
|
||||
`Cargo.lock` sweep, a Rust toolchain bump, or when any JSON is suspected corrupt.
|
||||
|
||||
### `new` — document only what's missing (default)
|
||||
|
||||
```bash
|
||||
python3 rustdoc-regen.py new # default subcommand
|
||||
python3 rustdoc-regen.py # same thing
|
||||
```
|
||||
|
||||
Resolves the live *current-target* dependency graph via `cargo tree`, finds
|
||||
every package whose `<crate>.json` is absent from `<doc-dir>`, and runs
|
||||
`cargo doc -p <crate>@<version> --no-deps` for each. This is the fast "a new
|
||||
dependency was pulled in, get docs for it" path: it touches only the
|
||||
newly-added crates, not the whole graph. Because it's per-crate with
|
||||
`--no-deps`, it's quick even when the full rebuild would take minutes.
|
||||
|
||||
`cargo tree` (not raw `Cargo.lock`) is the source so that **platform-gated**
|
||||
crates locked for other OSes (`windows_*`, `objc2_*`, `wasm-bindgen*`, `jni*`,
|
||||
`ndk*`, `bevy_android`, …) are *excluded* — `cargo doc -p <those> --no-deps`
|
||||
would error on this host. When a crate has several versions locked, the highest
|
||||
version is selected and built as `<name>@<version>` so cargo doesn't reject the
|
||||
spec as ambiguous.
|
||||
|
||||
Caveat: `--no-deps` means a new dep's *own* transitive deps only get documented
|
||||
if they're *also* missing from the live graph vs `<doc-dir>` (they usually are,
|
||||
so they get picked up in the same pass). If a dep's macro/helper crates aren't
|
||||
landing, fall back to `all`.
|
||||
|
||||
### `crates <NAME>...` — targeted refresh
|
||||
|
||||
```bash
|
||||
python3 rustdoc-regen.py crates avian3d clap
|
||||
python3 rustdoc-regen.py crates thiserror-impl@2.0.18 # name@ver when ambiguous
|
||||
```
|
||||
|
||||
Runs `cargo doc -p <spec> --no-deps` for each given spec (refresh whether or
|
||||
not the JSON already exists). Use when you know exactly which crate rustdoc-rag
|
||||
needs and its JSON is stale (e.g. you bumped a single dep with `cargo update -p`).
|
||||
Dashes vs underscores in names are handled either way. If a name has multiple
|
||||
locked versions, cargo will reject a bare spec as ambiguous — append
|
||||
`@<version>` (see the version from `cargo tree` or `missing`).
|
||||
|
||||
### `missing` — report only (no build)
|
||||
|
||||
```bash
|
||||
python3 rustdoc-regen.py missing
|
||||
python3 rustdoc-regen.py missing --json
|
||||
```
|
||||
|
||||
Prints the `name@version` specs of live-graph crates that have no JSON in
|
||||
`<doc-dir>` yet. Run this first to see what `new` would do, or to audit
|
||||
substrate freshness.
|
||||
|
||||
## Shared flags (go after the subcommand)
|
||||
|
||||
| flag | effect |
|
||||
|---|---|
|
||||
| `--doc-dir <path>` | override output dir (default `$RUSTDOC_RAG_DIR` or `./target/doc`) |
|
||||
| `--dry-run` | print the exact `cargo` command(s) instead of running |
|
||||
| `--keep-going` | pass `--keep-going` to `cargo doc` (best with `all`) |
|
||||
| `--offline` | pass `--offline` to `cargo` |
|
||||
| `--quiet` | suppress the `$ <cmd>` echo on stderr |
|
||||
| `--json` | structured output (`missing`, and a summary for build runs) |
|
||||
|
||||
## Typical agent workflow
|
||||
|
||||
When rustdoc-rag fails with `no rustdoc JSON for crate '<x>'`, or a dependency
|
||||
was just added/updated:
|
||||
|
||||
1. **Audit** before building — see exactly what's missing, no run:
|
||||
```bash
|
||||
python3 rustdoc-regen.py missing
|
||||
```
|
||||
2. **Targeted** if it's one or two known crates:
|
||||
```bash
|
||||
python3 rustdoc-regen.py crates <x>
|
||||
```
|
||||
3. **Sweep** after adding several deps or a `Cargo.lock` churn:
|
||||
```bash
|
||||
python3 rustdoc-regen.py new
|
||||
```
|
||||
4. **Full** after a toolchain bump, or if a JSON looks corrupt / a crate's API
|
||||
surface changed across a version jump:
|
||||
```bash
|
||||
python3 rustdoc-regen.py all --keep-going
|
||||
```
|
||||
5. **Verify** the result with rustdoc-rag:
|
||||
```bash
|
||||
python3 <rustdoc-rag path>/rustdoc-rag.py summary <x>
|
||||
```
|
||||
|
||||
Always point both skills at the same dir, e.g.:
|
||||
|
||||
```bash
|
||||
RUSTDOC_RAG_DIR=/path/to/target/doc python3 rustdoc-regen.py new
|
||||
RUSTDOC_RAG_DIR=/path/to/target/doc python3 <rustdoc-rag path>/rustdoc-rag.py summary <x>
|
||||
```
|
||||
|
||||
## Notes
|
||||
|
||||
- **Stable Rust, not nightly.** The script sets `RUSTC_BOOTSTRAP=1` and
|
||||
`RUSTDOCFLAGS=-Z unstable-options --output-format json` in `cargo`'s env. On
|
||||
nightly these flags are accepted directly without the bootstrap hatch, but
|
||||
the script works either way.
|
||||
- **`--output-format json` replaces HTML output** for the documented crates in
|
||||
`target/doc`: rustdoc emits `<crate>.json` files and skips the HTML tree for
|
||||
them. If you also want browsable HTML docs, run a plain `cargo doc` separately
|
||||
(or point it at a different `--target-dir`).
|
||||
- The JSON is **not incremental across runs** — rustdoc re-emits each selected
|
||||
crate's JSON every time. Compilation of the underlying rlibs *is* cached by
|
||||
cargo, so repeat runs are fast when nothing changed; only the rustdoc pass
|
||||
re-runs.
|
||||
- A few crates won't document cleanly in JSON mode (broken intra-doc links,
|
||||
`--cfg docsrs`-only doctests, etc.). Use `--keep-going` with `all` so one
|
||||
failure doesn't abort the rest, then re-run `missing` to see what's left.
|
||||
- Only **public** items are emitted (matches what rustdoc-rag catalogs).
|
||||
- Multi-version crates produce one JSON keyed by crate name; the built
|
||||
(highest) version wins the filename (`<crate>.json`), which is fine for API
|
||||
surface browsing via rustdoc-rag.
|
||||
241
pi/.pi/agent/skills/rustdoc-regen/rustdoc-regen.py
Executable file
241
pi/.pi/agent/skills/rustdoc-regen/rustdoc-regen.py
Executable file
@@ -0,0 +1,241 @@
|
||||
#!/usr/bin/env python3
|
||||
"""regen-rustdoc: rebuild the local rustdoc JSON substrate that rustdoc-rag reads.
|
||||
|
||||
Emits <crate>.json under <doc-dir> (default ./target/doc, honors $RUSTDOC_RAG_DIR)
|
||||
using rustdoc's unstable JSON output on *stable* Rust via RUSTC_BOOTSTRAP=1.
|
||||
|
||||
Subcommands:
|
||||
all `cargo doc` over the whole graph (workspace + all deps).
|
||||
new Document only the live-graph crates whose JSON is missing
|
||||
(per-crate `cargo doc -p <crate> --no-deps`). Default.
|
||||
crates <NAME>... Document specific crates (`--no-deps` each), refresh always.
|
||||
missing Print live-graph crates that have no JSON yet (no run).
|
||||
|
||||
Options (after subcommand): --doc-dir <path> --dry-run --keep-going
|
||||
--offline --quiet --json
|
||||
"""
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import json as _json
|
||||
import os
|
||||
import subprocess
|
||||
import sys
|
||||
from pathlib import Path
|
||||
|
||||
RUSTC_BOOTSTRAP = "1"
|
||||
RUSTDOC_JSON_FLAGS = "-Z unstable-options --output-format json"
|
||||
|
||||
|
||||
def die(msg: str):
|
||||
print(f"regen-rustdoc: {msg}", file=sys.stderr)
|
||||
sys.exit(1)
|
||||
|
||||
|
||||
def doc_dir(args) -> Path:
|
||||
d = getattr(args, "doc_dir", None) or os.environ.get("RUSTDOC_RAG_DIR") or "target/doc"
|
||||
return Path(d)
|
||||
|
||||
|
||||
def _parse_ver(s: str):
|
||||
parts = []
|
||||
for p in s.split("."):
|
||||
try:
|
||||
parts.append((0, int(p)))
|
||||
except ValueError:
|
||||
parts.append((1, p))
|
||||
return parts
|
||||
|
||||
|
||||
def live_packages(args) -> list[dict]:
|
||||
# Resolve the *current-target* dependency graph so platform-gated crates
|
||||
# (windows_*, objc2_*, wasm-bindgen*, jni*, ndk*, ...) locked for other OSes
|
||||
# are excluded — `cargo doc -p <those> --no-deps` would error on this host.
|
||||
# `--edges normal` matches what `cargo doc` documents (excludes dev-deps).
|
||||
cmd = ["cargo", "tree", "--workspace", "--prefix", "none", "--edges", "normal"]
|
||||
if getattr(args, "offline", False):
|
||||
cmd.append("--offline")
|
||||
r = subprocess.run(cmd, capture_output=True, text=True)
|
||||
if r.returncode != 0:
|
||||
die(f"cargo tree failed:\n{r.stderr.strip()}")
|
||||
out: list[dict] = []
|
||||
seen = set()
|
||||
for line in r.stdout.splitlines():
|
||||
line = line.strip()
|
||||
if not line or line.startswith("("):
|
||||
continue
|
||||
toks = line.split()
|
||||
name = toks[0]
|
||||
ver = ""
|
||||
for t in toks[1:]:
|
||||
if t.startswith("v") and len(t) > 1:
|
||||
ver = t[1:]
|
||||
break
|
||||
spec = f"{name}@{ver}" if ver else name
|
||||
if spec in seen:
|
||||
continue
|
||||
seen.add(spec)
|
||||
out.append({"name": name, "ver": ver, "spec": spec})
|
||||
if not out:
|
||||
die("cargo tree returned no packages; is this a Cargo project?")
|
||||
return out
|
||||
|
||||
|
||||
def json_path_for(doc_dir: Path, crate: str) -> Path:
|
||||
return doc_dir / f"{crate.replace('-', '_')}.json"
|
||||
|
||||
|
||||
def cargo_env() -> dict:
|
||||
env = os.environ.copy()
|
||||
env["RUSTC_BOOTSTRAP"] = RUSTC_BOOTSTRAP
|
||||
existing = env.get("RUSTDOCFLAGS", "").strip()
|
||||
env["RUSTDOCFLAGS"] = f"{existing} {RUSTDOC_JSON_FLAGS}".strip() if existing else RUSTDOC_JSON_FLAGS
|
||||
return env
|
||||
|
||||
|
||||
def run_cargo(cmd: list[str], dry: bool, quiet: bool, label: str = "") -> int:
|
||||
pretty = " ".join(cmd)
|
||||
if dry:
|
||||
print(f"[dry-run] {pretty}")
|
||||
return 0
|
||||
if not quiet:
|
||||
head = f"$ {pretty}"
|
||||
if label:
|
||||
head = f"# {label}\n{head}"
|
||||
print(head, file=sys.stderr)
|
||||
r = subprocess.run(cmd, env=cargo_env())
|
||||
return r.returncode
|
||||
|
||||
|
||||
def common_cargo_tail(args, include_keep_going: bool) -> list[str]:
|
||||
tail: list[str] = []
|
||||
if include_keep_going and getattr(args, "keep_going", False):
|
||||
tail.append("--keep-going")
|
||||
if getattr(args, "offline", False):
|
||||
tail.append("--offline")
|
||||
return tail
|
||||
|
||||
|
||||
def missing_specs_by_name(args) -> dict[str, dict]:
|
||||
pkgs = live_packages(args)
|
||||
dd = doc_dir(args)
|
||||
present = {f.stem for f in dd.glob("*.json")} if dd.is_dir() else set()
|
||||
chosen: dict[str, dict] = {}
|
||||
for p in pkgs:
|
||||
key = p["name"].replace("-", "_")
|
||||
if key in present:
|
||||
continue
|
||||
cur = chosen.get(p["name"])
|
||||
if cur is None or _parse_ver(p["ver"]) > _parse_ver(cur["ver"]):
|
||||
chosen[p["name"]] = p
|
||||
return chosen
|
||||
|
||||
|
||||
def cmd_missing(args) -> int:
|
||||
pkgs = live_packages(args)
|
||||
chosen = missing_specs_by_name(args)
|
||||
if args.json:
|
||||
print(_json.dumps({
|
||||
"doc_dir": str(doc_dir(args)),
|
||||
"live": len(pkgs),
|
||||
"missing_names": sorted(chosen),
|
||||
"missing_specs": [chosen[n]["spec"] for n in sorted(chosen)],
|
||||
}))
|
||||
else:
|
||||
for n in sorted(chosen):
|
||||
print(chosen[n]["spec"])
|
||||
if not chosen:
|
||||
print(f"all {len(pkgs)} live-graph crates have JSON in {doc_dir(args)}", file=sys.stderr)
|
||||
return 0
|
||||
|
||||
|
||||
def missing_crates(args) -> list[str]:
|
||||
return [missing_specs_by_name(args)[n]["spec"] for n in sorted(missing_specs_by_name(args))]
|
||||
|
||||
|
||||
def build_list(args, crates: list[str], label: str) -> int:
|
||||
if not crates:
|
||||
verb = "would document" if args.dry_run else "document"
|
||||
print(f"nothing to {verb} for {label}", file=sys.stderr)
|
||||
return 0
|
||||
verb = "would document" if args.dry_run else "documenting"
|
||||
print(f"{verb} {len(crates)} crate(s) for {label}", file=sys.stderr)
|
||||
results: list[dict] = []
|
||||
rc_overall = 0
|
||||
for c in crates:
|
||||
cmd = ["cargo", "doc", "-p", c, "--no-deps", *common_cargo_tail(args, include_keep_going=False)]
|
||||
rc = run_cargo(cmd, args.dry_run, args.quiet, label=c)
|
||||
results.append({"crate": c, "rc": rc})
|
||||
if rc != 0:
|
||||
rc_overall = rc
|
||||
if args.json:
|
||||
print(_json.dumps({"results": results, "overall_rc": rc_overall}))
|
||||
elif not args.dry_run:
|
||||
ok = sum(1 for r in results if r["rc"] == 0)
|
||||
print(f"done {ok}/{len(results)} ok", file=sys.stderr)
|
||||
return rc_overall
|
||||
|
||||
|
||||
def cmd_new(args) -> int:
|
||||
m = missing_crates(args)
|
||||
if not m and not args.json:
|
||||
print("no missing JSON — substrate is up to date", file=sys.stderr)
|
||||
if args.dry_run:
|
||||
return 0
|
||||
return 0
|
||||
return build_list(args, m, "new")
|
||||
|
||||
|
||||
def cmd_all(args) -> int:
|
||||
cmd = ["cargo", "doc", *common_cargo_tail(args, include_keep_going=True)]
|
||||
rc = run_cargo(cmd, args.dry_run, args.quiet, label="all (workspace + deps)")
|
||||
if args.json:
|
||||
print(_json.dumps({"cmd": cmd, "rc": rc}))
|
||||
return rc
|
||||
|
||||
|
||||
def cmd_crates(args) -> int:
|
||||
return build_list(args, args.crates, f"crates {args.crates}")
|
||||
|
||||
|
||||
def add_common(p: argparse.ArgumentParser):
|
||||
p.add_argument("--doc-dir", help="output dir (default $RUSTDOC_RAG_DIR or ./target/doc)")
|
||||
p.add_argument("--dry-run", action="store_true", help="print cargo commands instead of running")
|
||||
p.add_argument("--keep-going", action="store_true", help="pass --keep-going to cargo doc (best with 'all')")
|
||||
p.add_argument("--offline", action="store_true", help="pass --offline to cargo")
|
||||
p.add_argument("--quiet", action="store_true", help="suppress the $ <cmd> stderr echo")
|
||||
p.add_argument("--json", action="store_true", help="structured output")
|
||||
|
||||
|
||||
def main() -> int:
|
||||
ap = argparse.ArgumentParser(
|
||||
prog="regen-rustdoc",
|
||||
description="Rebuild local rustdoc JSON for the rustdoc-rag skill (stable Rust).",
|
||||
)
|
||||
sub = ap.add_subparsers(dest="cmd")
|
||||
for name, desc in (
|
||||
("all", "full rebuild: cargo doc over the whole graph (workspace + deps)"),
|
||||
("new", "document only live-graph crates whose JSON is missing (default)"),
|
||||
):
|
||||
p = sub.add_parser(name, help=desc)
|
||||
add_common(p)
|
||||
p = sub.add_parser("crates", help="document specific crates (--no-deps each)")
|
||||
p.add_argument("crates", nargs="+", metavar="NAME")
|
||||
add_common(p)
|
||||
p = sub.add_parser("missing", help="print live-graph crates lacking a JSON (no run)")
|
||||
add_common(p)
|
||||
args = ap.parse_args()
|
||||
args.cmd = args.cmd or "new"
|
||||
if args.cmd == "all":
|
||||
return cmd_all(args)
|
||||
if args.cmd == "new":
|
||||
return cmd_new(args)
|
||||
if args.cmd == "crates":
|
||||
return cmd_crates(args)
|
||||
if args.cmd == "missing":
|
||||
return cmd_missing(args)
|
||||
ap.error(f"unknown subcommand: {args.cmd}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(main())
|
||||
@@ -1,125 +0,0 @@
|
||||
---
|
||||
name: subagent-implement
|
||||
description: "Full implementation pipeline: scout → plan → review → code → build verify → review → fix. Use when the user invokes /implement."
|
||||
---
|
||||
|
||||
# /implement — Standard Implementation Pipeline
|
||||
|
||||
## Agent Roster
|
||||
|
||||
| Agent | Role |
|
||||
|-------|------|
|
||||
| scout | Fast codebase recon |
|
||||
| planner | Detailed implementation plans |
|
||||
| plan-reviewer | Reviews plans for correctness and risk |
|
||||
| coder | Primary implementation |
|
||||
| reviewer | Cross-family code review |
|
||||
| fixer | Applies review feedback precisely |
|
||||
|
||||
## Phase 1: Scout → Plan → Review (chain)
|
||||
|
||||
```js
|
||||
await subagent({ chain: [
|
||||
{ agent: "scout", task: "Explore the codebase for: {task}" },
|
||||
{ agent: "planner", task: "Create a detailed implementation plan for: {task}" },
|
||||
{ agent: "plan-reviewer", task: "Review the plan for correctness, completeness, and risk." }
|
||||
]})
|
||||
```
|
||||
|
||||
The chain creates `scout.md`, `plan.md`, `plan-review.md` in the chain artifact dir.
|
||||
|
||||
Read `plan-review.md` from the artifact dir. If **NEEDS_REVISION** or **REJECTED**, loop: tell the planner what to fix, re-run the reviewer. If **APPROVED**, proceed to Phase 2.
|
||||
|
||||
## Phase 2: Implement
|
||||
|
||||
```js
|
||||
await subagent({ agent: "coder", task: `Read and implement the plan at <chain_dir>/plan.md` })
|
||||
```
|
||||
|
||||
Use the chain dir path from Phase 1's result.
|
||||
|
||||
## Phase 3: Build Verification
|
||||
|
||||
After the coder finishes, independently verify the build — don't trust the coder's report:
|
||||
|
||||
```bash
|
||||
# Run a full build, not just type-check. Shaders, linkers, and bundlers may fail.
|
||||
# Adapt to project: cargo build, npm run build, etc.
|
||||
```
|
||||
|
||||
## Phase 4: Review → Fix
|
||||
|
||||
```js
|
||||
await subagent({ agent: "reviewer", output: `${chainDir}/review.md`, task: "Review all changes made" })
|
||||
```
|
||||
|
||||
Read `review.md`. If issues found:
|
||||
|
||||
```js
|
||||
await subagent({ agent: "fixer", task: `Read and apply the review feedback in ${chainDir}/review.md` })
|
||||
```
|
||||
|
||||
## Phase 5: Next Steps & Workflow Summary
|
||||
|
||||
After all phases complete, give the user a brief summary of what is next to take use of the implementation and an honoest evaluation of the workflow:
|
||||
|
||||
- **What needs wiring and where**: Point to where the user can take the implementation into use
|
||||
- **What happened**: which phases ran, any plan review loops, whether fixes were needed
|
||||
- **Issues**: any agent silent failures, fallbacks used, build errors, or unexpected behavior
|
||||
- **Agent quality**: did any agent misinterpret the task, produce poor output, or need hand-holding? Name the agent and the problem.
|
||||
- **Skill improvements**: did this workflow reveal gaps in the skill instructions or agent prompts? Note what should change
|
||||
|
||||
Be concise — a few lines is enough when things went well. Only expand on problems.
|
||||
|
||||
## Step Overrides
|
||||
|
||||
Override agent defaults per step:
|
||||
|
||||
```js
|
||||
await subagent({ chain: [
|
||||
{ agent: "scout", output: "context.md" },
|
||||
{ agent: "planner", reads: ["context.md"], output: "plan.md" },
|
||||
{ agent: "plan-reviewer", reads: ["plan.md"] }
|
||||
]})
|
||||
```
|
||||
|
||||
## Chain Mechanics
|
||||
|
||||
Chain mode (`subagent({ chain: [...] })`) runs agents sequentially in a shared temp directory (`{chain_dir}`). Each step:
|
||||
1. The framework injects `[Read from:]` and `[Write to:]` directives from the agent's `defaultReads` and `output` frontmatter
|
||||
2. The agent reads upstream files, does its work, and writes its deliverable to the `[Write to:]` path using the `write` tool
|
||||
3. The agent returns a brief text summary; `{previous}` carries this summary to the next step
|
||||
4. Variable substitution: `{task}` = original task, `{previous}` = prior step's brief ack, `{chain_dir}` = artifact dir path
|
||||
|
||||
Key behaviors:
|
||||
- Data flows through FILES (`scout.md` → `plan.md` → `plan-review.md`), not through `{previous}`
|
||||
- `{previous}` contains only a brief summary from the prior step — do NOT rely on it for full context
|
||||
- The framework validates that the expected output file was created
|
||||
- The chain result includes `📁 Artifacts: /tmp/pi-chain-runs/<id>/` — use this path to read files for branching decisions
|
||||
|
||||
Reviewer and fixer run as single agents (not in chains). Pass file paths in the task string and use the `output` parameter to capture their output.
|
||||
|
||||
## Fallback Strategy
|
||||
|
||||
When a subagent call returns no output (silent failure), apply cross-family model fallback. **Do not fall back to doing the work yourself** — always retry with the fallback model first.
|
||||
|
||||
1. **First attempt**: Use the agent's default model
|
||||
2. **If silent failure or error**: Retry with the fallback model using `model` override
|
||||
3. **If the fallback also fails**: Report the double-failure to the user. Still do not do the work yourself.
|
||||
|
||||
```js
|
||||
// Example: scout fails silently, retry with fallback
|
||||
subagent({ agent: "scout", task: "...", model: "opencode/big-pickle" })
|
||||
```
|
||||
|
||||
| Agent | Primary | Fallback |
|
||||
|-------|---------|----------|
|
||||
| scout | opencode-go/mimo-v2-pro | opencode/big-pickle |
|
||||
| planner | opencode/qwen3.6-plus-free | opencode-go/glm-5 |
|
||||
| coder | opencode-go/mimo-v2-pro | opencode-go/kimi-k2.5 |
|
||||
| plan-reviewer, reviewer | opencode-go/kimi-k2.5 | opencode/qwen3.6-plus-free |
|
||||
| fixer | opencode-go/glm-5 | opencode-go/mimo-v2-pro |
|
||||
|
||||
## Adaptive Routing
|
||||
|
||||
Do not run the full pipeline for a one-line fix. Use your judgment — a single coder call may suffice.
|
||||
@@ -1,76 +0,0 @@
|
||||
---
|
||||
name: subagent-plan
|
||||
description: "Plan-only pipeline: scout → plan → review. Produces a reviewed implementation plan without coding. Use when the user invokes /plan."
|
||||
---
|
||||
|
||||
# /plan — Plan Only (No Implementation)
|
||||
|
||||
## Agent Roster
|
||||
|
||||
| Agent | Role |
|
||||
|-------|------|
|
||||
| scout | Fast codebase recon |
|
||||
| planner | Detailed implementation plans |
|
||||
| plan-reviewer | Reviews plans for correctness and risk |
|
||||
|
||||
## Workflow (single chain)
|
||||
|
||||
```js
|
||||
await subagent({ chain: [
|
||||
{ agent: "scout", task: "Explore the codebase for: {task}" },
|
||||
{ agent: "planner", task: "Create a detailed implementation plan for: {task}" },
|
||||
{ agent: "plan-reviewer", task: "Review the plan for correctness, completeness, and risk." }
|
||||
]})
|
||||
```
|
||||
|
||||
The chain creates `scout.md`, `plan.md`, `plan-review.md` in the chain artifact dir (`/tmp/pi-chain-runs/<id>/`).
|
||||
|
||||
## Handling the Verdict
|
||||
|
||||
Read `plan-review.md` from the artifact dir. If **NEEDS_REVISION** or **REJECTED**, loop: tell the planner what to fix, re-run the reviewer. If **APPROVED**, present the plan path to the user.
|
||||
|
||||
Show the user the path to `plan.md`.
|
||||
|
||||
## Workflow Summary
|
||||
|
||||
After the chain completes, give the user a brief honest summary:
|
||||
|
||||
- **What happened**: did the plan pass review on the first try, or did it need revision loops?
|
||||
- **Issues**: any agent silent failures, fallbacks used, or unexpected behavior
|
||||
- **Agent quality**: did any agent misinterpret the task, produce poor output, or need hand-holding? Name the agent and the problem
|
||||
- **Skill improvements**: did this workflow reveal gaps in the skill instructions or agent prompts? Note what should change
|
||||
|
||||
Be concise — a few lines is enough when things went well. Only expand on problems.
|
||||
|
||||
## Chain Mechanics
|
||||
|
||||
Chain mode (`subagent({ chain: [...] })`) runs agents sequentially in a shared temp directory (`{chain_dir}`). Each step:
|
||||
1. The framework injects `[Read from:]` and `[Write to:]` directives from the agent's `defaultReads` and `output` frontmatter
|
||||
2. The agent reads upstream files, does its work, and writes its deliverable to the `[Write to:]` path using the `write` tool
|
||||
3. The agent returns a brief text summary; `{previous}` carries this summary to the next step
|
||||
4. Variable substitution: `{task}` = original task, `{previous}` = prior step's brief ack, `{chain_dir}` = artifact dir path
|
||||
|
||||
Key behaviors:
|
||||
- Data flows through FILES (`scout.md` → `plan.md` → `plan-review.md`), not through `{previous}`
|
||||
- `{previous}` contains only a brief summary from the prior step — do NOT rely on it for full context
|
||||
- The framework validates that the expected output file was created
|
||||
- The chain result includes `📁 Artifacts: /tmp/pi-chain-runs/<id>/` — use this path to read files for branching decisions
|
||||
|
||||
## Fallback Strategy
|
||||
|
||||
When a subagent call returns no output (silent failure), apply cross-family model fallback. **Do not fall back to doing the work yourself** — always retry with the fallback model first.
|
||||
|
||||
1. **First attempt**: Use the agent's default model
|
||||
2. **If silent failure or error**: Retry with the fallback model using `model` override
|
||||
3. **If the fallback also fails**: Report the double-failure to the user. Still do not do the work yourself.
|
||||
|
||||
```js
|
||||
// Example: scout fails silently, retry with fallback
|
||||
subagent({ agent: "scout", task: "...", model: "opencode/big-pickle" })
|
||||
```
|
||||
|
||||
| Agent | Primary | Fallback |
|
||||
|-------|---------|----------|
|
||||
| scout | opencode-go/mimo-v2-pro | opencode/big-pickle |
|
||||
| planner | opencode/qwen3.6-plus-free | opencode-go/glm-5 |
|
||||
| plan-reviewer | opencode-go/kimi-k2.5 | opencode/qwen3.6-plus-free |
|
||||
@@ -1,56 +0,0 @@
|
||||
---
|
||||
name: subagent-review
|
||||
description: "Code review using a cross-family reviewer. Use when the user invokes /review."
|
||||
---
|
||||
|
||||
# /review — Code Review
|
||||
|
||||
## Agent Roster
|
||||
|
||||
| Agent | Role |
|
||||
|-------|------|
|
||||
| reviewer | Cross-family code review (finds blind spots the author's model family misses) |
|
||||
| fixer | Applies review feedback precisely |
|
||||
|
||||
## Workflow
|
||||
|
||||
Single reviewer call — no chain needed:
|
||||
|
||||
```js
|
||||
await subagent({ agent: "reviewer", output: `${chainDir}/review.md`, task: "Review the following changes: <description or files>" })
|
||||
```
|
||||
|
||||
Read `review.md`. If issues found, apply fixes:
|
||||
|
||||
```js
|
||||
await subagent({ agent: "fixer", task: `Read and apply the review feedback in ${chainDir}/review.md` })
|
||||
```
|
||||
|
||||
## Workflow Summary
|
||||
|
||||
After review (and optional fix), give the user a brief honest summary:
|
||||
|
||||
- **What happened**: did the reviewer find issues, were fixes applied?
|
||||
- **Issues**: any agent silent failures, fallbacks used, or unexpected behavior
|
||||
- **Agent quality**: did any agent misinterpret the task, produce poor output, or need hand-holding? Name the agent and the problem
|
||||
- **Skill improvements**: did this workflow reveal gaps in the skill instructions or agent prompts? Note what should change
|
||||
|
||||
Be concise — a few lines is enough when things went well. Only expand on problems.
|
||||
|
||||
## Fallback Strategy
|
||||
|
||||
When a subagent call returns no output (silent failure), apply cross-family model fallback. **Do not fall back to doing the work yourself** — always retry with the fallback model first.
|
||||
|
||||
1. **First attempt**: Use the agent's default model
|
||||
2. **If silent failure or error**: Retry with the fallback model using `model` override
|
||||
3. **If the fallback also fails**: Report the double-failure to the user. Still do not do the work yourself.
|
||||
|
||||
```js
|
||||
// Example: reviewer fails silently, retry with fallback
|
||||
subagent({ agent: "reviewer", task: "...", model: "opencode/qwen3.6-plus-free" })
|
||||
```
|
||||
|
||||
| Agent | Primary | Fallback |
|
||||
|-------|---------|----------|
|
||||
| reviewer | opencode-go/kimi-k2.5 | opencode/qwen3.6-plus-free |
|
||||
| fixer | opencode-go/glm-5 | opencode-go/mimo-v2-pro |
|
||||
81
pi/.pi/agent/themes/bearded-arc.json
Normal file
81
pi/.pi/agent/themes/bearded-arc.json
Normal file
@@ -0,0 +1,81 @@
|
||||
{
|
||||
"$schema": "https://raw.githubusercontent.com/badlogic/pi-mono/main/packages/coding-agent/src/modes/interactive/theme/theme-schema.json",
|
||||
"name": "bearded-arc",
|
||||
"vars": {
|
||||
"bg": "#1c2433",
|
||||
"fg": "#afbbd2",
|
||||
"accent": "#b78aff",
|
||||
"accentAlt": "#ff955c",
|
||||
"link": "#69c3ff",
|
||||
"error": "#ff738a",
|
||||
"success": "#3cec85",
|
||||
"warning": "#eacd61",
|
||||
"muted": "#7c869a",
|
||||
"dim": "#5e687b",
|
||||
"borderMuted": "#414a5b",
|
||||
"selectedBg": "#28303f",
|
||||
"userMsgBg": "#242c3b",
|
||||
"toolPendingBg": "#212938",
|
||||
"toolSuccessBg": "#203c3d",
|
||||
"toolErrorBg": "#372d3d",
|
||||
"customMsgBg": "#282c43"
|
||||
},
|
||||
"colors": {
|
||||
"accent": "accent",
|
||||
"border": "link",
|
||||
"borderAccent": "accent",
|
||||
"borderMuted": "borderMuted",
|
||||
"success": "success",
|
||||
"error": "error",
|
||||
"warning": "warning",
|
||||
"muted": "muted",
|
||||
"dim": "dim",
|
||||
"text": "",
|
||||
"thinkingText": "muted",
|
||||
"selectedBg": "selectedBg",
|
||||
"userMessageBg": "userMsgBg",
|
||||
"userMessageText": "",
|
||||
"customMessageBg": "customMsgBg",
|
||||
"customMessageText": "",
|
||||
"customMessageLabel": "accent",
|
||||
"toolPendingBg": "toolPendingBg",
|
||||
"toolSuccessBg": "toolSuccessBg",
|
||||
"toolErrorBg": "toolErrorBg",
|
||||
"toolTitle": "",
|
||||
"toolOutput": "muted",
|
||||
"mdHeading": "warning",
|
||||
"mdLink": "link",
|
||||
"mdLinkUrl": "dim",
|
||||
"mdCode": "accent",
|
||||
"mdCodeBlock": "success",
|
||||
"mdCodeBlockBorder": "muted",
|
||||
"mdQuote": "muted",
|
||||
"mdQuoteBorder": "muted",
|
||||
"mdHr": "muted",
|
||||
"mdListBullet": "accent",
|
||||
"toolDiffAdded": "success",
|
||||
"toolDiffRemoved": "error",
|
||||
"toolDiffContext": "muted",
|
||||
"syntaxComment": "muted",
|
||||
"syntaxKeyword": "accent",
|
||||
"syntaxFunction": "link",
|
||||
"syntaxVariable": "accentAlt",
|
||||
"syntaxString": "success",
|
||||
"syntaxNumber": "accent",
|
||||
"syntaxType": "accentAlt",
|
||||
"syntaxOperator": "fg",
|
||||
"syntaxPunctuation": "muted",
|
||||
"thinkingOff": "borderMuted",
|
||||
"thinkingMinimal": "muted",
|
||||
"thinkingLow": "link",
|
||||
"thinkingMedium": "accentAlt",
|
||||
"thinkingHigh": "accent",
|
||||
"thinkingXhigh": "accent",
|
||||
"bashMode": "success"
|
||||
},
|
||||
"export": {
|
||||
"pageBg": "#141c2b",
|
||||
"cardBg": "#1c2433",
|
||||
"infoBg": "#353839"
|
||||
}
|
||||
}
|
||||
3
pi/.pi/agent/trust.json
Normal file
3
pi/.pi/agent/trust.json
Normal file
@@ -0,0 +1,3 @@
|
||||
{
|
||||
"/home/jonas/projects": true
|
||||
}
|
||||
Reference in New Issue
Block a user