diff --git a/README.md b/README.md index 44bee0df..e45f11a7 100644 --- a/README.md +++ b/README.md @@ -16,13 +16,14 @@ https://github.com/user-attachments/assets/c859872f-ca5e-4f8b-b6a0-7cc7461fe62a - **๐Ÿ“Š Request-Level Analytics** - Track latency, token usage, and costs in real-time - **๐Ÿ” Deep Debugging** - Full request/response logging and error traces - **โšก <10ms Overhead** - Minimal performance impact on your API calls +- **Supports z.ai coder plan** - Setup Claude and z.ai accounts and prioritize in which order they are used - **๐Ÿ’ธ Free & Open Source** - Run it yourself, modify it, own your infrastructure ## Quick Start ```bash # Clone and install -git clone https://github.com/snipeship/ccflare +git clone https://github.com/tombii/ccflare cd ccflare bun install @@ -87,7 +88,7 @@ Full documentation available in [`docs/`](docs/): ## Requirements - [Bun](https://bun.sh) >= 1.2.8 -- Claude API accounts (Free, Pro, or Team) +- Claude API accounts (Free, Pro, or Team) or z.ai code plan accounts ## Contributing diff --git a/apps/tui/src/components/AccountsScreen.tsx b/apps/tui/src/components/AccountsScreen.tsx index ec3c1c56..37ecdedb 100644 --- a/apps/tui/src/components/AccountsScreen.tsx +++ b/apps/tui/src/components/AccountsScreen.tsx @@ -10,7 +10,13 @@ interface AccountsScreenProps { onBack: () => void; } -type Mode = "list" | "add" | "remove" | "confirmRemove" | "waitingForCode"; +type Mode = + | "list" + | "add" + | "remove" + | "confirmRemove" + | "waitingForCode" + | "setPriority"; export function AccountsScreen({ onBack }: AccountsScreenProps) { const [mode, setMode] = useState("list"); @@ -27,6 +33,9 @@ export function AccountsScreen({ onBack }: AccountsScreenProps) { const [error, setError] = useState(null); const [accountToRemove, setAccountToRemove] = useState(""); const [confirmInput, setConfirmInput] = useState(""); + const [accountForPriority, setAccountForPriority] = + useState(null); + const [priorityInput, setPriorityInput] = useState(""); useInput((input, key) => { if (key.escape) { @@ -34,6 +43,10 @@ export function AccountsScreen({ onBack }: AccountsScreenProps) { setMode("list"); setAccountToRemove(""); setConfirmInput(""); + } else if (mode === "setPriority") { + setMode("list"); + setAccountForPriority(null); + setPriorityInput(""); } else if (mode === "add" || mode === "waitingForCode") { setMode("list"); setNewAccountName(""); @@ -64,6 +77,7 @@ export function AccountsScreen({ onBack }: AccountsScreenProps) { name: newAccountName, mode: selectedMode, tier: selectedTier, + priority: 0, // Default priority }); setOauthFlowData(flowData); setMode("waitingForCode"); @@ -83,6 +97,7 @@ export function AccountsScreen({ onBack }: AccountsScreenProps) { name: newAccountName, mode: selectedMode, tier: selectedTier, + priority: 0, // Default priority code: authCode, flowData: oauthFlowData, }); @@ -100,12 +115,42 @@ export function AccountsScreen({ onBack }: AccountsScreenProps) { } }; - const handleRemoveAccount = (name: string) => { + const _handleRemoveAccount = (name: string) => { setAccountToRemove(name); setConfirmInput(""); setMode("confirmRemove"); }; + const handleSetPriority = (account: AccountDisplay) => { + setAccountForPriority(account); + setPriorityInput(account.priority.toString()); + setMode("setPriority"); + }; + + const handleUpdatePriority = async () => { + if (!accountForPriority || !priorityInput) return; + + const priority = parseInt(priorityInput, 10); + if (Number.isNaN(priority) || priority < 0 || priority > 100) { + setError("Priority must be a number between 0 and 100"); + return; + } + + try { + // Using tuiCore for priority update + await tuiCore.updateAccountPriority(accountForPriority.name, priority); + await loadAccounts(); + setMode("list"); + setAccountForPriority(null); + setPriorityInput(""); + setError(null); + } catch (error) { + setError( + error instanceof Error ? error.message : "Failed to update priority", + ); + } + }; + const handleConfirmRemove = async () => { if (confirmInput !== accountToRemove) { return; @@ -262,11 +307,55 @@ export function AccountsScreen({ onBack }: AccountsScreenProps) { ); } + if (mode === "setPriority") { + return ( + + + ๐Ÿ”ข Set Account Priority + + + + + Setting priority for account:{" "} + {accountForPriority?.name} + + + Current priority: {accountForPriority?.priority} + + + Priority values range from 0 (lowest) to 100 (highest) + + + + + New priority (0-100): + { + if (priorityInput) handleUpdatePriority(); + }} + /> + + + {error && ( + + {error} + + )} + + + Press ENTER to confirm, ESC to cancel + + + ); + } + const menuItems = [ ...accounts.map((acc) => { const presenter = new AccountPresenter(acc); return { - label: `${acc.name} (${presenter.tierDisplay})`, + label: `${acc.name} (${presenter.tierDisplay}) - Priority: ${acc.priority}`, value: `account:${acc.name}`, }; }), @@ -298,7 +387,10 @@ export function AccountsScreen({ onBack }: AccountsScreenProps) { setMode("add"); } else if (item.value.startsWith("account:")) { const accountName = item.value.replace("account:", ""); - handleRemoveAccount(accountName); + const account = accounts.find((acc) => acc.name === accountName); + if (account) { + handleSetPriority(account); + } } }} /> diff --git a/apps/tui/src/components/ServerScreen.tsx b/apps/tui/src/components/ServerScreen.tsx index 91d71a2e..1b54ccac 100644 --- a/apps/tui/src/components/ServerScreen.tsx +++ b/apps/tui/src/components/ServerScreen.tsx @@ -1,5 +1,5 @@ -import { NETWORK } from "@ccflare/core"; import { Config } from "@ccflare/config"; +import { NETWORK } from "@ccflare/core"; import { Box, Text, useInput } from "ink"; interface ServerScreenProps { diff --git a/apps/tui/src/main.ts b/apps/tui/src/main.ts index 4c7c89ea..9d85ee8b 100644 --- a/apps/tui/src/main.ts +++ b/apps/tui/src/main.ts @@ -48,12 +48,14 @@ Options: --logs [N] Stream latest N lines then follow --stats Show statistics (JSON output) --add-account Add a new account - --mode Account mode (default: max) + --mode Account mode (default: max) --tier <1|5|20> Account tier (default: 1) + --priority Account priority (default: 0) --list List all accounts --remove Remove an account --pause Pause an account --resume Resume an account + --set-priority Set account priority --analyze Analyze database performance --reset-stats Reset usage statistics --clear-history Clear request history @@ -116,6 +118,7 @@ Examples: name: parsed.addAccount, mode: parsed.mode || "max", tier: parsed.tier || 1, + priority: parsed.priority || 0, }); console.log(`โœ… Account "${parsed.addAccount}" added successfully`); return; @@ -128,7 +131,9 @@ Examples: } else { console.log("\nAccounts:"); accounts.forEach((acc) => { - console.log(` - ${acc.name} (${acc.mode} mode, tier ${acc.tier})`); + console.log( + ` - ${acc.name} (${acc.mode} mode, tier ${acc.tier}, priority ${acc.priority})`, + ); }); } return; @@ -170,6 +175,23 @@ Examples: return; } + if (parsed.setPriority) { + const [name, priorityStr] = parsed.setPriority; + const priority = parseInt(priorityStr, 10); + + if (Number.isNaN(priority)) { + console.error(`โŒ Invalid priority value: ${priorityStr}`); + process.exit(1); + } + + const result = await tuiCore.updateAccountPriority(name, priority); + console.log(result.message); + if (!result.success) { + process.exit(1); + } + return; + } + if (parsed.analyze) { await tuiCore.analyzePerformance(); return; @@ -189,6 +211,7 @@ Examples: "opus-4": CLAUDE_MODEL_IDS.OPUS_4, "sonnet-4": CLAUDE_MODEL_IDS.SONNET_4, "opus-4.1": CLAUDE_MODEL_IDS.OPUS_4_1, + "sonnet-4.5": CLAUDE_MODEL_IDS.SONNET_4_5, }; const fullModel = modelMap[parsed.setModel]; diff --git a/docs/api-http.md b/docs/api-http.md index 20615b07..adc7ffd8 100644 --- a/docs/api-http.md +++ b/docs/api-http.md @@ -120,6 +120,7 @@ List all configured accounts with their current status. "lastUsed": "2024-12-17T10:25:30.123Z", "created": "2024-12-01T08:00:00.000Z", "tier": 5, + "priority": 50, "paused": false, "tokenStatus": "valid", "rateLimitStatus": "allowed_warning (5m)", @@ -148,7 +149,8 @@ Initialize OAuth flow for adding a new account. { "name": "myaccount", "mode": "max", // "max" or "console" (default: "max") - "tier": 5 // 1, 5, or 20 (default: 1) + "tier": 5, // 1, 5, or 20 (default: 1) + "priority": 50 // 0-100, lower value = higher priority (optional, defaults: 0) } ``` @@ -187,7 +189,8 @@ Complete OAuth flow after user authorization. "success": true, "message": "Account 'myaccount' added successfully!", "mode": "Claude Max", - "tier": 5 + "tier": 5, + "priority": 50 } ``` @@ -254,6 +257,39 @@ curl -X POST http://localhost:8080/api/accounts/uuid-here/tier \ -d '{"tier": 20}' ``` +#### POST /api/accounts/:accountId/priority + +Update account priority. Lower priority values increase the likelihood of the account being selected by the load balancer. + +**Request:** +```json +{ + "priority": 75 // 0-100, lower value = higher priority +} +``` + +**Response:** +```json +{ + "success": true, + "priority": 75 +} +``` + +**Example:** +```bash +curl -X POST http://localhost:8080/api/accounts/uuid-here/priority \ + -H "Content-Type: application/json" \ + -d '{"priority": 75}' +``` + +**How Priorities Work:** +- Priority values range from 0 to 100 +- Lower values indicate higher priority in load balancing +- Priority is optional and defaults to 0 (highest priority) if not specified +- Priority affects both primary account selection and fallback order +- Changes take effect immediately without restarting the server + #### POST /api/accounts/:accountId/pause Pause an account temporarily. diff --git a/docs/cli.md b/docs/cli.md index 631ed28c..03b46baf 100644 --- a/docs/cli.md +++ b/docs/cli.md @@ -8,6 +8,7 @@ The ccflare CLI provides a command-line interface for managing OAuth accounts, m - [Global Options and Help](#global-options-and-help) - [Command Reference](#command-reference) - [Account Management](#account-management) + - [Account Priorities](#account-priorities) - [Statistics and History](#statistics-and-history) - [System Commands](#system-commands) - [Server and Monitoring](#server-and-monitoring) @@ -27,7 +28,7 @@ The ccflare CLI provides a command-line interface for managing OAuth accounts, m 1. Clone the repository: ```bash -git clone https://github.com/snipe-code/ccflare.git +git clone https://github.com/tombii/ccflare.git cd ccflare ``` @@ -89,10 +90,12 @@ Options: --add-account Add a new account --mode Account mode (default: max) --tier <1|5|20> Account tier (default: 1) + --priority Account priority (0-100, default: 0) --list List all accounts --remove Remove an account --pause Pause an account --resume Resume an account + --set-priority Set account priority (0-100) --analyze Analyze database performance --reset-stats Reset usage statistics --clear-history Clear request history @@ -112,7 +115,7 @@ Add a new OAuth account to the load balancer pool. **Syntax:** ```bash -ccflare --add-account [--mode ] [--tier <1|5|20>] +ccflare --add-account [--mode ] [--tier <1|5|20>] [--priority ] ``` **Options:** @@ -123,13 +126,18 @@ ccflare --add-account [--mode ] [--tier <1|5|20>] - `1`: Tier 1 account - `5`: Tier 5 account - `20`: Tier 20 account +- `--priority`: Account priority (optional, defaults to 0) + - Range: 0-100 + - Lower numbers indicate higher priority in load balancing + - If not specified, defaults to 0 (highest priority) **Interactive Flow:** 1. If mode not provided, defaults to "max" 2. If tier not provided (Max accounts only), defaults to 1 -3. Opens browser for OAuth authentication -4. Waits for OAuth callback on localhost:7856 -5. Stores account credentials securely in the database +3. If priority not provided, defaults to 0 +4. Opens browser for OAuth authentication +5. Waits for OAuth callback on localhost:7856 +6. Stores account credentials securely in the database #### `--list` @@ -143,8 +151,8 @@ ccflare --list **Output Format:** ``` Accounts: - - account1 (max mode, tier 5) - - account2 (console mode, tier 1) + - account1 (max mode, tier 5, priority 10) + - account2 (console mode, tier 1, priority 5) ``` #### `--remove ` @@ -184,6 +192,39 @@ Re-enable a paused account for load balancing. ccflare --resume ``` +### Account Priorities + +#### `--set-priority ` + +Set or update the priority of an account. Accounts with lower priority numbers are preferred in the load balancing algorithm. + +**Syntax:** +```bash +ccflare --set-priority +``` + +**Parameters:** +- `name`: Account name to update +- `priority`: Priority value (0-100, where lower numbers indicate higher priority) + +**How Priorities Work:** +- Accounts with lower priority numbers are selected first +- Default priority is 0 if not specified +- Priority affects both primary account selection and fallback order +- Changes take effect immediately without restarting the server + +**Example:** +```bash +# Set account to high priority (low number) +ccflare --set-priority production-account 10 + +# Set account to medium priority +ccflare --set-priority development-account 50 + +# Set account to low priority (high number) +ccflare --set-priority backup-account 90 +``` + ### Statistics and History #### `--stats` @@ -300,15 +341,21 @@ ccflare --logs 50 ### Basic Account Setup ```bash -# Add a Claude Max account with tier 5 -ccflare --add-account work-account --mode max --tier 5 +# Add a Claude Max account with tier 5 and high priority (low number) +ccflare --add-account work-account --mode max --tier 5 --priority 10 + +# Add a Console account with medium priority +ccflare --add-account personal-account --mode console --priority 50 -# Add a Console account -ccflare --add-account personal-account --mode console +# Add a backup account with low priority (high number) +ccflare --add-account backup-account --mode max --tier 1 --priority 90 # List all accounts ccflare --list +# Update account priority +ccflare --set-priority backup-account 20 + # View statistics ccflare --stats ``` @@ -368,10 +415,10 @@ ccflare ### Automation Examples ```bash -# Add multiple accounts via script -for i in {1..3}; do - ccflare --add-account "account-$i" --mode max --tier 5 -done +# Add multiple accounts with different priorities +ccflare --add-account "primary-account" --mode max --tier 20 --priority 10 +ccflare --add-account "secondary-account" --mode max --tier 5 --priority 50 +ccflare --add-account "backup-account" --mode max --tier 1 --priority 90 # Monitor account status watch -n 5 'ccflare --list' @@ -381,6 +428,11 @@ ccflare --clear-history && ccflare --reset-stats # Export statistics for monitoring ccflare --stats > stats.json + +# Prioritize specific account temporarily +ccflare --set-priority primary-account 5 +# ... run important workload ... +ccflare --set-priority primary-account 10 # Restore normal priority ``` ## Configuration @@ -556,6 +608,10 @@ ccflare --logs - Use descriptive account names - Distribute load across multiple accounts - Keep accounts of similar tiers for consistent performance + - Use account priorities to control load distribution: + - Set lower priority numbers for premium or preferred accounts + - Use higher priority numbers for backup or development accounts + - Adjust priorities temporarily for specific workloads - Pause accounts proactively when approaching rate limits 3. **Security** diff --git a/docs/index.md b/docs/index.md index 7db1bd8a..889ade9e 100644 --- a/docs/index.md +++ b/docs/index.md @@ -56,6 +56,7 @@ When working with Claude API at scale, rate limits can become a significant bott - [Load Balancing Strategy](./load-balancing.md) - Session-based strategy for safe account usage - [Provider System](./providers.md) - Provider abstraction and OAuth implementation - [Database Schema](./database.md) - SQLite structure, migrations, and maintenance +- [Model Mappings](./models.md) - Claude AI model definitions and constants ### User Interfaces - [HTTP API Reference](./api-http.md) - Complete REST API documentation diff --git a/docs/load-balancing.md b/docs/load-balancing.md index cbb30d67..1199d416 100644 --- a/docs/load-balancing.md +++ b/docs/load-balancing.md @@ -3,10 +3,11 @@ ## Table of Contents 1. [Overview](#overview) 2. [Session-Based Strategy](#session-based-strategy) -3. [Configuration](#configuration) -4. [Account Selection Process](#account-selection-process) -5. [Performance Considerations](#performance-considerations) -6. [Important: Why Only Session-Based Strategy](#important-why-only-session-based-strategy) +3. [Account Priorities](#account-priorities) +4. [Configuration](#configuration) +5. [Account Selection Process](#account-selection-process) +6. [Performance Considerations](#performance-considerations) +7. [Important: Why Only Session-Based Strategy](#important-why-only-session-based-strategy) ## Overview @@ -16,6 +17,7 @@ ccflare implements a session-based load balancing system to distribute requests - **Account Health Monitoring**: Automatically filters out rate-limited or paused accounts - **Failover Support**: Returns ordered lists of accounts for automatic failover - **Session Persistence**: Maintains configurable sessions on specific accounts +- **Account Priorities**: Supports prioritized account selection for better control over load distribution - **Real-time Configuration**: Change settings without restarting the server - **Provider Filtering**: Accounts are filtered by provider compatibility @@ -84,6 +86,53 @@ export class SessionStrategy implements LoadBalancingStrategy { - โš ๏ธ **Uneven Load Distribution**: May concentrate load on fewer accounts - โš ๏ธ **Session Dependency**: Performance tied to specific account availability +## Account Priorities + +Account priorities allow you to control which accounts are preferred when multiple accounts are available. This feature gives you fine-grained control over load distribution and account selection. + +### How Priorities Work + +- **Priority Range**: Accounts can have a priority value from 0-100 (default: 0) +- **Lower Value = Higher Priority**: Accounts with lower priority values are selected first +- **Optional Parameter**: Priority is optional when adding accounts and defaults to 0 (highest priority) +- **Affects Both Primary and Fallback Selection**: Priorities determine both the primary account and the order of fallback accounts +- **Real-time Updates**: Priority changes take effect immediately without restarting the server + +### Setting Account Priorities + +Priorities can be set when adding an account or updated later: + +```bash +# Add account with priority +ccflare --add-account myaccount --mode max --tier 5 --priority 10 + +# Update account priority +ccflare set-priority myaccount 20 +``` + +### Priority in Load Balancing + +The SessionStrategy considers priorities when selecting accounts: + +1. **Active Session Check**: First looks for an account with an active session +2. **Priority Sorting**: If no active session or the active account is unavailable, available accounts are sorted by priority (descending) +3. **Fallback Order**: Remaining accounts are also ordered by priority for failover scenarios + +```typescript +// From load-balancer/src/strategies/index.ts +// Filter available accounts and sort by priority (lower value = higher priority) +const available = accounts + .filter((a) => isAccountAvailable(a, now)) + .sort((a, b) => a.priority - b.priority); // Ascending sort +``` + +### Use Cases for Priorities + +1. **Primary/Backup Setup**: Assign higher priorities to preferred accounts +2. **Cost Management**: Prioritize free or lower-cost accounts +3. **Performance Optimization**: Prioritize accounts with better performance characteristics +4. **Tiered Access**: Create hierarchical access patterns based on account capabilities + ## Configuration ccflare uses a hierarchical configuration system where environment variables take precedence over configuration file settings. @@ -190,7 +239,7 @@ The SessionStrategy manages account sessions through the following process: 2. **Session Validation**: Checks if the session is within the configured duration 3. **Account Ordering**: Returns accounts in priority order: - Active session account (if available) comes first - - Other available accounts follow as fallback options + - Other available accounts are sorted by priority (lower values first) as fallback options ### 4. Session Reset Sessions are reset when: diff --git a/docs/models.md b/docs/models.md new file mode 100644 index 00000000..80b7a4b5 --- /dev/null +++ b/docs/models.md @@ -0,0 +1,165 @@ +# Model Mappings in ccflare + +## Overview + +ccflare includes centralized model definitions and mappings for Claude AI models. These mappings are defined in `packages/core/src/models.ts` and are used throughout the system for consistent model identification and display. + +## Model Constants + +### CLAUDE_MODEL_IDS + +Full model IDs as used by the Anthropic API: + +```typescript +export const CLAUDE_MODEL_IDS = { + // Claude 3.5 models + HAIKU_3_5: "claude-3-5-haiku-20241022", + SONNET_3_5: "claude-3-5-sonnet-20241022", + + // Claude 4 models + SONNET_4: "claude-sonnet-4-20250514", + OPUS_4: "claude-opus-4-20250514", + OPUS_4_1: "claude-opus-4-1-20250805", + + // Legacy Claude 3 models (for documentation/API examples) + OPUS_3: "claude-3-opus-20240229", + SONNET_3: "claude-3-sonnet-20240229", +} as const; +``` + +### MODEL_DISPLAY_NAMES + +Human-readable display names for models: + +```typescript +export const MODEL_DISPLAY_NAMES: Record = { + [CLAUDE_MODEL_IDS.HAIKU_3_5]: "Claude Haiku 3.5", + [CLAUDE_MODEL_IDS.SONNET_3_5]: "Claude Sonnet 3.5 v2", + [CLAUDE_MODEL_IDS.SONNET_4]: "Claude Sonnet 4", + [CLAUDE_MODEL_IDS.OPUS_4]: "Claude Opus 4", + [CLAUDE_MODEL_IDS.OPUS_4_1]: "Claude Opus 4.1", + [CLAUDE_MODEL_IDS.OPUS_3]: "Claude Opus 3", + [CLAUDE_MODEL_IDS.SONNET_3]: "Claude Sonnet 3", +}; +``` + +### MODEL_SHORT_NAMES + +Short names used in UI components (for color mapping, etc.): + +```typescript +export const MODEL_SHORT_NAMES: Record = { + [CLAUDE_MODEL_IDS.HAIKU_3_5]: "claude-3.5-haiku", + [CLAUDE_MODEL_IDS.SONNET_3_5]: "claude-3.5-sonnet", + [CLAUDE_MODEL_IDS.SONNET_4]: "claude-sonnet-4", + [CLAUDE_MODEL_IDS.OPUS_4]: "claude-opus-4", + [CLAUDE_MODEL_IDS.OPUS_4_1]: "claude-opus-4.1", + [CLAUDE_MODEL_IDS.OPUS_3]: "claude-3-opus", + [CLAUDE_MODEL_IDS.SONNET_3]: "claude-3-sonnet", +}; +``` + +## Default Models + +```typescript +// Default model for general use +export const DEFAULT_MODEL = CLAUDE_MODEL_IDS.SONNET_4; + +// Default model for agents +export const DEFAULT_AGENT_MODEL = CLAUDE_MODEL_IDS.SONNET_4; +``` + +## Helper Functions + +### getModelShortName(modelId: string): string + +Returns the short name for a given model ID, or the model ID itself if no mapping exists. + +```typescript +export function getModelShortName(modelId: string): string { + return MODEL_SHORT_NAMES[modelId] || modelId; +} +``` + +### getModelDisplayName(modelId: string): string + +Returns the display name for a given model ID, or the model ID itself if no mapping exists. + +```typescript +export function getModelDisplayName(modelId: string): string { + return MODEL_DISPLAY_NAMES[modelId] || modelId; +} +``` + +### isValidModelId(modelId: string): modelId is ClaudeModelId + +Validates if a string is a recognized Claude model ID. + +```typescript +export function isValidModelId(modelId: string): modelId is ClaudeModelId { + return Object.values(CLAUDE_MODEL_IDS).includes(modelId as ClaudeModelId); +} +``` + +## Type Definitions + +```typescript +// Type for all valid model IDs +export type ClaudeModelId = (typeof CLAUDE_MODEL_IDS)[keyof typeof CLAUDE_MODEL_IDS]; +``` + +## Usage Examples + +### Getting a Model's Display Name + +```typescript +import { getModelDisplayName, CLAUDE_MODEL_IDS } from "@ccflare/core"; + +const modelId = CLAUDE_MODEL_IDS.SONNET_4; +const displayName = getModelDisplayName(modelId); +// Returns: "Claude Sonnet 4" +``` + +### Validating a Model ID + +```typescript +import { isValidModelId } from "@ccflare/core"; + +if (isValidModelId("claude-3-5-sonnet-20241022")) { + // This is a valid model ID +} +``` + +### Using the Model Type + +```typescript +import type { ClaudeModelId } from "@ccflare/core"; + +function processModel(modelId: ClaudeModelId) { + // TypeScript knows this is a valid model ID + console.log(getModelDisplayName(modelId)); +} +``` + +## Model Version Information + +The model mappings include both current and legacy model versions: + +- **Claude 4 Models** (Latest): + - `claude-sonnet-4-20250514` + - `claude-opus-4-20250514` + - `claude-opus-4-1-20250805` + +- **Claude 3.5 Models**: + - `claude-3-5-haiku-20241022` + - `claude-3-5-sonnet-20241022` + +- **Legacy Claude 3 Models**: + - `claude-3-opus-20240229` + - `claude-3-sonnet-20240229` + +## Agent Model Preferences + +ccflare supports agent-specific model preferences through the agent system. When an agent is detected in a request, the system can automatically override the model selection based on the agent's configured preference. + +See [Agent Documentation](./providers.md#agent-system) for more details on how agents work with model preferences. \ No newline at end of file diff --git a/docs/providers.md b/docs/providers.md index adeeb91a..570c38d6 100644 --- a/docs/providers.md +++ b/docs/providers.md @@ -6,6 +6,9 @@ - **Anthropic** - Single provider with two modes: - **console mode**: Standard Claude API (console.anthropic.com) - **max mode**: Claude Code (claude.ai) +- **Z.ai** - Claude proxy service with API key authentication: + - Lite, Pro, and Max plans with higher rate limits than direct Claude API + - Uses API key authentication (no OAuth support) ### Key Points - All API requests route to `https://api.anthropic.com` @@ -35,8 +38,14 @@ The ccflare providers system is a modular architecture designed to support multi - **Claude API** (console mode) - Standard API access via console.anthropic.com - **Claude Code** (max mode) - Enhanced access via claude.ai +2. **Z.ai Provider** - Provides access to: + - **Z.ai API** - Claude proxy service with enhanced rate limits + - Uses API key authentication instead of OAuth + - Supports Lite, Pro, and Max plans with ~3ร— the usage quota of equivalent Claude plans + The providers system handles: -- OAuth authentication flows with PKCE security +- OAuth authentication flows with PKCE security (Anthropic) +- API key authentication (Z.ai) - Token lifecycle management (refresh, expiration) - Provider-specific request routing and header management - Rate limit detection and handling @@ -179,9 +188,52 @@ The OAuth client ID can be configured in multiple ways (in order of precedence): The AnthropicProvider extends the BaseProvider class and implements Anthropic-specific functionality. -### Request Routing +## ZaiProvider Implementation + +The ZaiProvider extends the BaseProvider class and implements Z.ai-specific functionality. + +### Key Features + +1. **API Key Authentication**: Uses `x-api-key` header instead of OAuth +2. **No Tier Detection**: Returns null from `extractTierInfo()` - tiers must be set manually +3. **Usage Extraction**: Parses token usage from both streaming and non-streaming responses (similar to Anthropic) +4. **Request Routing**: Routes all requests to `https://api.z.ai/api/anthropic` +5. **Compatible Response Format**: Z.ai responses follow the same format as Anthropic's API + +### Z.ai Request Routing + +The Z.ai provider routes all requests to the Z.ai API endpoint: + +```typescript +buildUrl(path: string, query: string): string { + return `https://api.z.ai/api/anthropic${path}${query}`; +} +``` + +### Z.ai Authentication + +Z.ai uses API key authentication via the `x-api-key` header: -The provider handles all request paths and routes them to the standard Anthropic API endpoint: +```typescript +prepareHeaders(headers: Headers, accessToken?: string, apiKey?: string): Headers { + const newHeaders = new Headers(headers); + + // z.ai expects the API key in x-api-key header + if (accessToken) { + newHeaders.set("x-api-key", accessToken); + } else if (apiKey) { + newHeaders.set("x-api-key", apiKey); + } + + return newHeaders; +} +``` + +The API key is stored in the `refresh_token` field of the account record for consistency with the authentication system. + +## Anthropic Request Routing + +The Anthropic provider handles all request paths and routes them to the standard Anthropic API endpoint: ```typescript canHandle(_path: string): boolean { @@ -301,9 +353,27 @@ ccflare supports three account tiers based on Anthropic's subscription levels: | Pro | 5 | 200,000 tokens/min | Individual pro subscriptions | | Team | 20 | 800,000+ tokens/min | Team/enterprise accounts | -### Automatic Tier Detection +### Z.ai Tier Mapping + +Z.ai plans map to ccflare tiers as follows: + +| Z.ai Plan | ccflare Tier | Usage Limit | Description | +|-----------|--------------|-------------|-------------| +| Lite | 1 | ~120 prompts/5hrs | ~3ร— Claude Pro usage quota | +| Pro | 5 | ~600 prompts/5hrs | ~3ร— Claude Max (5ร—) usage quota | +| Max | 20 | ~2400 prompts/5hrs | ~3ร— Claude Max (20ร—) usage quota | + +**Important**: Z.ai does not provide automatic tier detection. You must manually specify the tier when adding z.ai accounts: + +```bash +ccflare --add-account my-zai-account --tier 1 # For Lite plan +ccflare --add-account my-zai-account --tier 5 # For Pro plan +ccflare --add-account my-zai-account --tier 20 # For Max plan +``` + +### Automatic Tier Detection (Anthropic Only) -The system automatically detects account tiers from API responses: +The system automatically detects account tiers from Anthropic API responses: ```typescript async extractTierInfo(response: Response): Promise { @@ -410,7 +480,7 @@ interface Account { } ``` -**Note**: The current implementation prioritizes OAuth authentication. API key support is maintained for backward compatibility but OAuth is the preferred method. +**Note**: The current implementation prioritizes OAuth authentication for Anthropic. Z.ai uses API key authentication exclusively as it does not support OAuth. ### Token Refresh Strategy diff --git a/packages/cli-commands/src/commands/account.ts b/packages/cli-commands/src/commands/account.ts index fc5cc383..d7f07b6c 100644 --- a/packages/cli-commands/src/commands/account.ts +++ b/packages/cli-commands/src/commands/account.ts @@ -12,8 +12,9 @@ import { openBrowser } from "../utils/browser"; // Re-export types with adapter extension for CLI-specific options export interface AddAccountOptions { name: string; - mode?: "max" | "console"; + mode?: "max" | "console" | "zai"; tier?: 1 | 5 | 20; + priority?: number; adapter?: PromptAdapter; } @@ -22,7 +23,7 @@ export type { AccountListItem } from "@ccflare/types"; // Add mode property to AccountListItem for CLI display export interface AccountListItemWithMode extends AccountListItem { - mode: "max" | "console"; + mode: "max" | "console" | "zai"; } /** @@ -37,6 +38,7 @@ export async function addAccount( name, mode: providedMode, tier: providedTier, + priority: providedPriority, adapter = stdPromptAdapter, } = options; @@ -49,52 +51,99 @@ export async function addAccount( (await adapter.select("What type of account would you like to add?", [ { label: "Claude Max account", value: "max" }, { label: "Claude Console account", value: "console" }, + { label: "z.ai account (API key)", value: "zai" }, ])); - // Begin OAuth flow - const flowResult = await oauthFlow.begin({ - name, - mode: mode as "max" | "console", - }); - const { authUrl, sessionId } = flowResult; - - // Open browser and prompt for code - console.log(`\nOpening browser to authenticate...`); - console.log(`URL: ${authUrl}`); - const browserOpened = await openBrowser(authUrl); - if (!browserOpened) { - console.log( - `\nFailed to open browser automatically. Please manually open the URL above.`, + if (mode === "zai") { + // Handle z.ai accounts with API keys + const apiKey = await adapter.input("\nEnter your z.ai API key: "); + + // Get tier for z.ai accounts + const tier = + providedTier || + (await adapter.select( + "Select the tier for this account (used for weighted load balancing):", + [ + { label: "1x tier (default)", value: 1 }, + { label: "5x tier (higher priority)", value: 5 }, + { label: "20x tier (highest priority)", value: 20 }, + ], + )); + + // Create z.ai account directly in database + const accountId = crypto.randomUUID(); + const now = Date.now(); + + dbOps.getDatabase().run( + `INSERT INTO accounts ( + id, name, provider, api_key, refresh_token, access_token, + expires_at, created_at, account_tier, request_count, total_requests, priority + ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`, + [ + accountId, + name, + "zai", + apiKey, + apiKey, // Store API key as refresh_token for consistency + apiKey, // Store API key as access_token + now + 365 * 24 * 60 * 60 * 1000, // 1 year expiry + now, + tier, + 0, + 0, + providedPriority || 0, + ], ); - } - // Get authorization code - const code = await adapter.input("\nEnter the authorization code: "); - - // Get tier for Max accounts - const tier = - mode === "max" - ? providedTier || - (await adapter.select( - "Select the tier for this account (used for weighted load balancing):", - [ - { label: "1x tier (default free account)", value: 1 }, - { label: "5x tier (paid account)", value: 5 }, - { label: "20x tier (enterprise account)", value: 20 }, - ], - )) - : 1; - - // Complete OAuth flow - console.log("\nExchanging code for tokens..."); - const _account = await oauthFlow.complete( - { sessionId, code, tier, name }, - flowResult, - ); - - console.log(`\nAccount '${name}' added successfully!`); - console.log(`Type: ${mode === "max" ? "Claude Max" : "Claude Console"}`); - console.log(`Tier: ${tier}x`); + console.log(`\nAccount '${name}' added successfully!`); + console.log("Type: z.ai (API key)"); + console.log(`Tier: ${tier}x`); + } else { + // Handle OAuth accounts (Anthropic) + const flowResult = await oauthFlow.begin({ + name, + mode: mode as "max" | "console", + }); + const { authUrl, sessionId } = flowResult; + + // Open browser and prompt for code + console.log(`\nOpening browser to authenticate...`); + console.log(`URL: ${authUrl}`); + const browserOpened = await openBrowser(authUrl); + if (!browserOpened) { + console.log( + `\nFailed to open browser automatically. Please manually open the URL above.`, + ); + } + + // Get authorization code + const code = await adapter.input("\nEnter the authorization code: "); + + // Get tier for Max accounts + const tier = + mode === "max" + ? providedTier || + (await adapter.select( + "Select the tier for this account (used for weighted load balancing):", + [ + { label: "1x tier (default free account)", value: 1 }, + { label: "5x tier (paid account)", value: 5 }, + { label: "20x tier (enterprise account)", value: 20 }, + ], + )) + : 1; + + // Complete OAuth flow + console.log("\nExchanging code for tokens..."); + const _account = await oauthFlow.complete( + { sessionId, code, tier, name, priority: providedPriority || 0 }, + flowResult, + ); + + console.log(`\nAccount '${name}' added successfully!`); + console.log(`Type: ${mode === "max" ? "Claude Max" : "Claude Console"}`); + console.log(`Tier: ${tier}x`); + } } /** @@ -137,7 +186,13 @@ export function getAccountsList(dbOps: DatabaseOperations): AccountListItem[] { rateLimitStatus, sessionInfo, tier: account.account_tier || 1, - mode: account.account_tier > 1 ? "max" : "console", + mode: + account.provider === "zai" + ? "zai" + : account.account_tier > 1 + ? "max" + : "console", + priority: account.priority || 0, }; }); } @@ -264,3 +319,37 @@ export function resumeAccount( ): { success: boolean; message: string } { return toggleAccountPause(dbOps, name, false); } + +/** + * Set the priority of an account by name + */ +export function setAccountPriority( + dbOps: DatabaseOperations, + name: string, + priority: number, +): { success: boolean; message: string } { + const db = dbOps.getDatabase(); + + // Get account ID by name + const account = db + .query<{ id: string }, [string]>("SELECT id FROM accounts WHERE name = ?") + .get(name); + + if (!account) { + return { + success: false, + message: `Account '${name}' not found`, + }; + } + + // Update the account priority + db.run("UPDATE accounts SET priority = ? WHERE id = ?", [ + priority, + account.id, + ]); + + return { + success: true, + message: `Account '${name}' priority set to ${priority}`, + }; +} diff --git a/packages/cli-commands/src/commands/help.ts b/packages/cli-commands/src/commands/help.ts index c789ae41..1c341355 100644 --- a/packages/cli-commands/src/commands/help.ts +++ b/packages/cli-commands/src/commands/help.ts @@ -6,10 +6,11 @@ export function getHelpText(): string { Usage: ccflare-cli [options] Commands: - add [--mode ] [--tier <1|5|20>] + add [--mode ] [--tier <1|5|20>] [--priority ] Add a new account using OAuth --mode: Account type (optional, will prompt if not provided) --tier: Account tier (1, 5, or 20) (optional, will prompt for Max accounts) + --priority: Account priority (0-100, default 0, lower numbers = higher priority) list List all accounts with their details @@ -24,6 +25,10 @@ Commands: resume Resume a paused account to include it in load balancing + set-priority + Set the priority of an account + --priority: Account priority (0-100, lower numbers = higher priority) + reset-stats Reset request counts for all accounts @@ -37,10 +42,11 @@ Commands: Show this help message Examples: - ccflare-cli add myaccount --mode max --tier 5 + ccflare-cli add myaccount --mode max --tier 5 --priority 10 ccflare-cli list ccflare-cli remove myaccount ccflare-cli pause myaccount ccflare-cli resume myaccount + ccflare-cli set-priority myaccount 20 `; } diff --git a/packages/cli-commands/src/runner.ts b/packages/cli-commands/src/runner.ts index a35c7587..b2598e88 100644 --- a/packages/cli-commands/src/runner.ts +++ b/packages/cli-commands/src/runner.ts @@ -9,6 +9,7 @@ import { pauseAccount, removeAccountWithConfirmation, resumeAccount, + setAccountPriority, } from "./commands/account"; import { analyzePerformance } from "./commands/analyze"; import { getHelpText } from "./commands/help"; @@ -33,6 +34,7 @@ export async function runCli(argv: string[]): Promise { options: { mode: { type: "string" }, tier: { type: "string" }, + priority: { type: "string" }, force: { type: "boolean" }, }, }); @@ -45,7 +47,7 @@ export async function runCli(argv: string[]): Promise { if (!name) { console.error("Error: Account name is required"); console.log( - "Usage: ccflare-cli add [--mode ] [--tier <1|5|20>]", + "Usage: ccflare-cli add [--mode ] [--tier <1|5|20>] [--priority ]", ); process.exit(1); } @@ -59,8 +61,15 @@ export async function runCli(argv: string[]): Promise { tierValue === 1 || tierValue === 5 || tierValue === 20 ? tierValue : undefined; + const priorityValue = values.priority + ? parseInt(values.priority as string) + : undefined; + const priority = + typeof priorityValue === "number" && !Number.isNaN(priorityValue) + ? priorityValue + : undefined; - await addAccount(dbOps, config, { name, mode, tier }); + await addAccount(dbOps, config, { name, mode, tier, priority }); break; } @@ -78,6 +87,7 @@ export async function runCli(argv: string[]): Promise { "Name".padEnd(20) + "Type".padEnd(10) + "Tier".padEnd(6) + + "Priority".padEnd(9) + "Requests".padEnd(12) + "Token".padEnd(10) + "Status".padEnd(20) + @@ -91,6 +101,7 @@ export async function runCli(argv: string[]): Promise { account.name.padEnd(20) + account.provider.padEnd(10) + account.tierDisplay.padEnd(6) + + account.priority.toString().padEnd(9) + `${account.requestCount}/${account.totalRequests}`.padEnd(12) + account.tokenStatus.padEnd(10) + account.rateLimitStatus.padEnd(20) + @@ -167,6 +178,36 @@ export async function runCli(argv: string[]): Promise { break; } + case "set-priority": { + const name = positionals[1]; + const priorityValue = positionals[2]; + + if (!name) { + console.error("Error: Account name is required"); + console.log("Usage: ccflare-cli set-priority "); + process.exit(1); + } + + if (priorityValue === undefined) { + console.error("Error: Priority value is required"); + console.log("Usage: ccflare-cli set-priority "); + process.exit(1); + } + + const priority = parseInt(priorityValue); + if (Number.isNaN(priority)) { + console.error("Error: Priority must be a number"); + process.exit(1); + } + + const result = setAccountPriority(dbOps, name, priority); + console.log(result.message); + if (!result.success) { + process.exit(1); + } + break; + } + case "analyze": { const db = dbOps.getDatabase(); analyzePerformance(db); diff --git a/packages/config/src/index.ts b/packages/config/src/index.ts index 3b05dba2..c0469a1d 100644 --- a/packages/config/src/index.ts +++ b/packages/config/src/index.ts @@ -8,6 +8,9 @@ import { NETWORK, type StrategyName, TIME_CONSTANTS, + ValidationError, + validateNumber, + validateString, } from "@ccflare/core"; import { Logger } from "@ccflare/logger"; import { resolveConfigPath } from "./paths"; @@ -19,6 +22,19 @@ export interface RuntimeConfig { retry: { attempts: number; delayMs: number; backoff: number }; sessionDurationMs: number; port: number; + database?: { + walMode?: boolean; + busyTimeoutMs?: number; + cacheSize?: number; + synchronous?: "OFF" | "NORMAL" | "FULL"; + mmapSize?: number; + retry?: { + attempts?: number; + delayMs?: number; + backoff?: number; + maxDelayMs?: number; + }; + }; } export interface ConfigData { @@ -32,9 +48,106 @@ export interface ConfigData { default_agent_model?: string; data_retention_days?: number; request_retention_days?: number; + // Database configuration + db_wal_mode?: boolean; + db_busy_timeout_ms?: number; + db_cache_size?: number; + db_synchronous?: "OFF" | "NORMAL" | "FULL"; + db_mmap_size?: number; + db_retry_attempts?: number; + db_retry_delay_ms?: number; + db_retry_backoff?: number; + db_retry_max_delay_ms?: number; [key: string]: string | number | boolean | undefined; } +/** + * Validates database configuration parameters + */ +function validateDatabaseConfig( + config: Partial, +): void { + if (!config) return; + + // Validate synchronous mode + if (config.synchronous !== undefined) { + validateString(config.synchronous, "db_synchronous", { + allowedValues: ["OFF", "NORMAL", "FULL"], + }); + } + + // Validate numeric parameters with reasonable bounds + if (config.busyTimeoutMs !== undefined) { + validateNumber(config.busyTimeoutMs, "db_busy_timeout_ms", { + min: 0, + max: 300000, // 5 minutes max + integer: true, + }); + } + + if (config.cacheSize !== undefined) { + validateNumber(config.cacheSize, "db_cache_size", { + min: -2000000, // -2GB max negative (KB) + max: 1000000, // 1M pages max positive + integer: true, + }); + } + + if (config.mmapSize !== undefined) { + validateNumber(config.mmapSize, "db_mmap_size", { + min: 0, + max: 1073741824, // 1GB max + integer: true, + }); + } + + // Validate retry configuration consistency + if (config.retry) { + const retry = config.retry; + + if (retry.attempts !== undefined) { + validateNumber(retry.attempts, "db_retry_attempts", { + min: 1, + max: 10, + integer: true, + }); + } + + if (retry.delayMs !== undefined) { + validateNumber(retry.delayMs, "db_retry_delay_ms", { + min: 1, + max: 60000, // 1 minute max + integer: true, + }); + } + + if (retry.backoff !== undefined) { + validateNumber(retry.backoff, "db_retry_backoff", { + min: 1, + max: 10, + }); + } + + if (retry.maxDelayMs !== undefined) { + validateNumber(retry.maxDelayMs, "db_retry_max_delay_ms", { + min: 1, + max: 300000, // 5 minutes max + integer: true, + }); + } + + // Ensure maxDelayMs is greater than delayMs if both are specified + if (retry.delayMs !== undefined && retry.maxDelayMs !== undefined) { + if (retry.maxDelayMs < retry.delayMs) { + throw new ValidationError( + "db_retry_max_delay_ms must be greater than or equal to db_retry_delay_ms", + "db_retry_max_delay_ms", + ); + } + } + } +} + export class Config extends EventEmitter { private configPath: string; private data: ConfigData = {}; @@ -203,6 +316,19 @@ export class Config extends EventEmitter { }, sessionDurationMs: TIME_CONSTANTS.SESSION_DURATION_DEFAULT, port: NETWORK.DEFAULT_PORT, + database: { + walMode: true, + busyTimeoutMs: 5000, + cacheSize: -20000, // 20MB cache + synchronous: "NORMAL", + mmapSize: 268435456, // 256MB + retry: { + attempts: 3, + delayMs: 100, + backoff: 2, + maxDelayMs: 5000, + }, + }, }; // Override with environment variables if present @@ -245,6 +371,76 @@ export class Config extends EventEmitter { defaults.port = this.data.port; } + // Database configuration overrides + // Ensure database configuration object exists + if (!defaults.database) { + defaults.database = { + walMode: true, + busyTimeoutMs: 5000, + cacheSize: -20000, + synchronous: "NORMAL", + mmapSize: 268435456, + retry: { + attempts: 3, + delayMs: 100, + backoff: 2, + maxDelayMs: 5000, + }, + }; + } + + // Ensure retry configuration object exists + if (!defaults.database.retry) { + defaults.database.retry = { + attempts: 3, + delayMs: 100, + backoff: 2, + maxDelayMs: 5000, + }; + } + + if (typeof this.data.db_wal_mode === "boolean") { + defaults.database.walMode = this.data.db_wal_mode; + } + if (typeof this.data.db_busy_timeout_ms === "number") { + defaults.database.busyTimeoutMs = this.data.db_busy_timeout_ms; + } + if (typeof this.data.db_cache_size === "number") { + defaults.database.cacheSize = this.data.db_cache_size; + } + if (typeof this.data.db_synchronous === "string") { + defaults.database.synchronous = this.data.db_synchronous as + | "OFF" + | "NORMAL" + | "FULL"; + } + if (typeof this.data.db_mmap_size === "number") { + defaults.database.mmapSize = this.data.db_mmap_size; + } + if (typeof this.data.db_retry_attempts === "number") { + defaults.database.retry.attempts = this.data.db_retry_attempts; + } + if (typeof this.data.db_retry_delay_ms === "number") { + defaults.database.retry.delayMs = this.data.db_retry_delay_ms; + } + if (typeof this.data.db_retry_backoff === "number") { + defaults.database.retry.backoff = this.data.db_retry_backoff; + } + if (typeof this.data.db_retry_max_delay_ms === "number") { + defaults.database.retry.maxDelayMs = this.data.db_retry_max_delay_ms; + } + + // Validate the final database configuration + try { + validateDatabaseConfig(defaults.database); + } catch (error) { + if (error instanceof ValidationError) { + log.error(`Database configuration validation failed: ${error.message}`); + throw error; + } + throw error; + } + return defaults; } } diff --git a/packages/core/src/models.ts b/packages/core/src/models.ts index 61d402d4..8c242593 100644 --- a/packages/core/src/models.ts +++ b/packages/core/src/models.ts @@ -11,6 +11,7 @@ export const CLAUDE_MODEL_IDS = { // Claude 4 models SONNET_4: "claude-sonnet-4-20250514", + SONNET_4_5: "claude-sonnet-4-5-20250929", OPUS_4: "claude-opus-4-20250514", OPUS_4_1: "claude-opus-4-1-20250805", @@ -24,6 +25,7 @@ export const MODEL_DISPLAY_NAMES: Record = { [CLAUDE_MODEL_IDS.HAIKU_3_5]: "Claude Haiku 3.5", [CLAUDE_MODEL_IDS.SONNET_3_5]: "Claude Sonnet 3.5 v2", [CLAUDE_MODEL_IDS.SONNET_4]: "Claude Sonnet 4", + [CLAUDE_MODEL_IDS.SONNET_4_5]: "Claude Sonnet 4.5", [CLAUDE_MODEL_IDS.OPUS_4]: "Claude Opus 4", [CLAUDE_MODEL_IDS.OPUS_4_1]: "Claude Opus 4.1", [CLAUDE_MODEL_IDS.OPUS_3]: "Claude Opus 3", @@ -35,6 +37,7 @@ export const MODEL_SHORT_NAMES: Record = { [CLAUDE_MODEL_IDS.HAIKU_3_5]: "claude-3.5-haiku", [CLAUDE_MODEL_IDS.SONNET_3_5]: "claude-3.5-sonnet", [CLAUDE_MODEL_IDS.SONNET_4]: "claude-sonnet-4", + [CLAUDE_MODEL_IDS.SONNET_4_5]: "claude-sonnet-4.5", [CLAUDE_MODEL_IDS.OPUS_4]: "claude-opus-4", [CLAUDE_MODEL_IDS.OPUS_4_1]: "claude-opus-4.1", [CLAUDE_MODEL_IDS.OPUS_3]: "claude-3-opus", diff --git a/packages/core/src/pricing.ts b/packages/core/src/pricing.ts index cfe7f75b..c410ddfb 100644 --- a/packages/core/src/pricing.ts +++ b/packages/core/src/pricing.ts @@ -66,6 +66,16 @@ const BUNDLED_PRICING: ApiResponse = { cache_write: 3.75, }, }, + [CLAUDE_MODEL_IDS.SONNET_4_5]: { + id: CLAUDE_MODEL_IDS.SONNET_4_5, + name: MODEL_DISPLAY_NAMES[CLAUDE_MODEL_IDS.SONNET_4_5], + cost: { + input: 3, + output: 15, + cache_read: 0.3, + cache_write: 3.75, + }, + }, [CLAUDE_MODEL_IDS.OPUS_4]: { id: CLAUDE_MODEL_IDS.OPUS_4, name: MODEL_DISPLAY_NAMES[CLAUDE_MODEL_IDS.OPUS_4], @@ -90,6 +100,52 @@ const BUNDLED_PRICING: ApiResponse = { }, }; +// Pricing for Zhipu AI models (GLM models) +BUNDLED_PRICING.zai = { + models: { + "glm-4.5": { + id: "glm-4.5", + name: "GLM-4.5", + cost: { + input: 0.6, + output: 2.2, + cache_read: 0.11, + cache_write: 0, + }, + }, + "glm-4.5-air": { + id: "glm-4.5-air", + name: "GLM-4.5-Air", + cost: { + input: 0.2, + output: 1.1, + cache_read: 0.03, + cache_write: 0, + }, + }, + "glm-4.6": { + id: "glm-4.6", + name: "GLM-4.6", + cost: { + input: 0.6, + output: 2.2, + cache_read: 0.11, + cache_write: 0, + }, + }, + "glm-4.6-air": { + id: "glm-4.6-air", + name: "GLM-4.6-Air", + cost: { + input: 0.2, + output: 1.1, + cache_read: 0.03, + cache_write: 0, + }, + }, + }, +}; + interface Logger { warn(message: string, ...args: unknown[]): void; } @@ -135,6 +191,127 @@ class PriceCatalogue { } } + /** + * Merge remote pricing data with bundled pricing data to ensure all models are included + */ + private mergePricingData( + remote: ApiResponse, + bundled: ApiResponse, + ): ApiResponse { + const merged: ApiResponse = {}; + + // List of preferred providers in priority order + const preferredProviders = ["zai", "anthropic"]; + + // First, add preferred providers from remote data + for (const providerName of preferredProviders) { + if (remote[providerName]) { + merged[providerName] = remote[providerName]; + } + } + + // Then add remaining providers from remote data, filtering out problematic ones + for (const [providerName, providerData] of Object.entries(remote)) { + if ( + !merged[providerName] && + !this.shouldFilterProvider(providerName, providerData) + ) { + merged[providerName] = providerData; + } + } + + // For each provider in bundled pricing, ensure it exists in merged data + for (const [providerName, providerData] of Object.entries(bundled)) { + if (!merged[providerName]) { + this.logger?.warn( + "Provider %s not found in remote pricing, using bundled data", + providerName, + ); + merged[providerName] = providerData; + } else if (providerData.models) { + // Merge models from bundled into remote data + if (!merged[providerName].models) { + merged[providerName].models = {}; + } + + // Add any missing models from bundled data + let addedModels = 0; + for (const [modelId, modelData] of Object.entries( + providerData.models, + )) { + if (!merged[providerName].models?.[modelId]) { + merged[providerName].models[modelId] = modelData; + addedModels++; + } + } + + if (addedModels > 0) { + this.logger?.warn( + "Added %d missing models for provider %s from bundled pricing", + addedModels, + providerName, + ); + } + } + } + + return merged; + } + + /** + * Determine if a provider should be filtered out (e.g., zero-cost duplicates) + */ + private shouldFilterProvider( + providerName: string, + providerData: { models?: Record }, + ): boolean { + // Filter out providers with names that suggest they're coding plans or special variants + const problematicPatterns = [ + /-coding-plan$/, + /-special$/, + /-demo$/, + /-free$/, + /-trial$/, + ]; + + if (problematicPatterns.some((pattern) => pattern.test(providerName))) { + this.logger?.warn( + "Filtering out provider %s due to problematic name pattern", + providerName, + ); + return true; + } + + // Filter out providers that have models with all zero costs + if (providerData.models) { + const modelEntries = Object.entries(providerData.models); + if (modelEntries.length > 0) { + const allZeroCost = modelEntries.every(([, model]) => { + if (!model.cost) return true; + const { + input = 0, + output = 0, + cache_read = 0, + cache_write = 0, + } = model.cost; + return ( + input === 0 && output === 0 && cache_read === 0 && cache_write === 0 + ); + }); + + if (allZeroCost) { + this.logger?.warn( + "Filtering out provider %s because all models have zero cost", + providerName, + ); + return true; + } + } + } + + return false; + } + private async loadFromCache(): Promise { try { const cachePath = this.getCachePath(); @@ -197,8 +374,11 @@ class PriceCatalogue { data = await this.loadFromCache(); } - // Fall back to bundled pricing - if (!data) { + // If we have remote data, merge it with bundled pricing to ensure we have all models + if (data) { + data = this.mergePricingData(data, BUNDLED_PRICING); + } else { + // Fall back to bundled pricing data = BUNDLED_PRICING; } @@ -207,13 +387,21 @@ class PriceCatalogue { return data; } - warnOnce(modelId: string): void { + warnOnce(modelId: string, error?: Error | string): void { if (!this.warnedModels.has(modelId)) { this.warnedModels.add(modelId); - this.logger?.warn( - "Price for model %s not found - cost set to 0", - modelId, - ); + if (error) { + this.logger?.warn( + "Price for model %s not found - cost set to 0 (reason: %s)", + modelId, + error instanceof Error ? error.message : error, + ); + } else { + this.logger?.warn( + "Price for model %s not found - cost set to 0", + modelId, + ); + } } } } @@ -299,8 +487,8 @@ export async function estimateCostUSD( } return totalCost; - } catch (_error) { - catalogue.warnOnce(modelId); + } catch (error) { + catalogue.warnOnce(modelId, error instanceof Error ? error : String(error)); return 0; } } diff --git a/packages/dashboard-web/src/api.ts b/packages/dashboard-web/src/api.ts index 429d7e77..85ac1f95 100644 --- a/packages/dashboard-web/src/api.ts +++ b/packages/dashboard-web/src/api.ts @@ -56,8 +56,10 @@ class API extends HttpClient { async initAddAccount(data: { name: string; - mode: "max" | "console"; + mode: "max" | "console" | "zai"; + apiKey?: string; tier: number; + priority: number; }): Promise<{ authUrl: string; sessionId: string }> { try { return await this.post<{ authUrl: string; sessionId: string }>( @@ -89,6 +91,25 @@ class API extends HttpClient { } } + async addZaiAccount(data: { + name: string; + apiKey: string; + tier: number; + priority: number; + }): Promise<{ message: string; account: Account }> { + try { + return await this.post<{ message: string; account: Account }>( + "/api/accounts/zai", + data, + ); + } catch (error) { + if (error instanceof HttpError) { + throw new Error(error.message); + } + throw error; + } + } + async removeAccount(name: string, confirm: string): Promise { try { await this.delete(`/api/accounts/${name}`, { @@ -211,6 +232,20 @@ class API extends HttpClient { } } + async updateAccountPriority( + accountId: string, + priority: number, + ): Promise { + try { + await this.post(`/api/accounts/${accountId}/priority`, { priority }); + } catch (error) { + if (error instanceof HttpError) { + throw new Error(error.message); + } + throw error; + } + } + async getStrategy(): Promise { const data = await this.get<{ strategy: string }>("/api/config/strategy"); return data.strategy; diff --git a/packages/dashboard-web/src/components/AccountsTab.tsx b/packages/dashboard-web/src/components/AccountsTab.tsx index d6e3b0c1..2c2a2036 100644 --- a/packages/dashboard-web/src/components/AccountsTab.tsx +++ b/packages/dashboard-web/src/components/AccountsTab.tsx @@ -6,6 +6,7 @@ import { useApiError } from "../hooks/useApiError"; import { AccountAddForm, AccountList, + AccountPriorityDialog, DeleteConfirmationDialog, RenameAccountDialog, } from "./accounts"; @@ -45,12 +46,20 @@ export function AccountsTab() { isOpen: false, account: null, }); + const [priorityDialog, setPriorityDialog] = useState<{ + isOpen: boolean; + account: Account | null; + }>({ + isOpen: false, + account: null, + }); const [actionError, setActionError] = useState(null); const handleAddAccount = async (params: { name: string; - mode: "max" | "console"; + mode: "max" | "console" | "zai"; tier: number; + priority: number; }) => { try { const result = await api.initAddAccount(params); @@ -77,6 +86,23 @@ export function AccountsTab() { } }; + const handleAddZaiAccount = async (params: { + name: string; + apiKey: string; + tier: number; + priority: number; + }) => { + try { + await api.addZaiAccount(params); + await loadAccounts(); + setAdding(false); + setActionError(null); + } catch (err) { + setActionError(formatError(err)); + throw err; + } + }; + const handleRemoveAccount = (name: string) => { setConfirmDelete({ show: true, accountName: name, confirmInput: "" }); }; @@ -134,6 +160,20 @@ export function AccountsTab() { } }; + const handlePriorityChange = (account: Account) => { + setPriorityDialog({ isOpen: true, account }); + }; + + const handleUpdatePriority = async (accountId: string, priority: number) => { + try { + await api.updateAccountPriority(accountId, priority); + await loadAccounts(); + } catch (err) { + setActionError(formatError(err)); + throw err; + } + }; + if (loading) { return ( @@ -179,6 +219,7 @@ export function AccountsTab() { { setAdding(false); setActionError(null); @@ -195,6 +236,7 @@ export function AccountsTab() { onPauseToggle={handlePauseToggle} onRemove={handleRemoveAccount} onRename={handleRename} + onPriorityChange={handlePriorityChange} /> @@ -230,6 +272,20 @@ export function AccountsTab() { isLoading={renameAccount.isPending} /> )} + + {priorityDialog.isOpen && priorityDialog.account && ( + + setPriorityDialog({ + isOpen: open, + account: open ? priorityDialog.account : null, + }) + } + onUpdatePriority={handleUpdatePriority} + /> + )} ); } diff --git a/packages/dashboard-web/src/components/accounts/AccountAddForm.tsx b/packages/dashboard-web/src/components/accounts/AccountAddForm.tsx index 16d32c3b..f36fe1df 100644 --- a/packages/dashboard-web/src/components/accounts/AccountAddForm.tsx +++ b/packages/dashboard-web/src/components/accounts/AccountAddForm.tsx @@ -13,13 +13,20 @@ import { interface AccountAddFormProps { onAddAccount: (params: { name: string; - mode: "max" | "console"; + mode: "max" | "console" | "zai"; tier: number; + priority: number; }) => Promise<{ authUrl: string; sessionId: string }>; onCompleteAccount: (params: { sessionId: string; code: string; }) => Promise; + onAddZaiAccount: (params: { + name: string; + apiKey: string; + tier: number; + priority: number; + }) => Promise; onCancel: () => void; onSuccess: () => void; onError: (error: string) => void; @@ -28,6 +35,7 @@ interface AccountAddFormProps { export function AccountAddForm({ onAddAccount, onCompleteAccount, + onAddZaiAccount, onCancel, onSuccess, onError, @@ -37,8 +45,10 @@ export function AccountAddForm({ const [sessionId, setSessionId] = useState(""); const [newAccount, setNewAccount] = useState({ name: "", - mode: "max" as "max" | "console", + mode: "max" as "max" | "console" | "zai", tier: 1, + priority: 0, + apiKey: "", }); const handleAddAccount = async () => { @@ -46,7 +56,32 @@ export function AccountAddForm({ onError("Account name is required"); return; } - // Step 1: Initialize OAuth flow + + if (newAccount.mode === "zai") { + if (!newAccount.apiKey) { + onError("API key is required for z.ai accounts"); + return; + } + // For z.ai accounts, we don't need OAuth flow + await onAddZaiAccount({ + name: newAccount.name, + apiKey: newAccount.apiKey, + tier: newAccount.tier, + priority: newAccount.priority, + }); + // Reset form and signal success + setNewAccount({ + name: "", + mode: "max", + tier: 1, + priority: 0, + apiKey: "", + }); + onSuccess(); + return; + } + + // Step 1: Initialize OAuth flow for Max/Console accounts const { authUrl, sessionId } = await onAddAccount(newAccount); setSessionId(sessionId); @@ -74,7 +109,7 @@ export function AccountAddForm({ setAuthStep("form"); setAuthCode(""); setSessionId(""); - setNewAccount({ name: "", mode: "max", tier: 1 }); + setNewAccount({ name: "", mode: "max", tier: 1, priority: 0, apiKey: "" }); onSuccess(); }; @@ -82,7 +117,7 @@ export function AccountAddForm({ setAuthStep("form"); setAuthCode(""); setSessionId(""); - setNewAccount({ name: "", mode: "max", tier: 1 }); + setNewAccount({ name: "", mode: "max", tier: 1, priority: 0, apiKey: "" }); onCancel(); }; @@ -111,7 +146,7 @@ export function AccountAddForm({ + {newAccount.mode === "zai" && ( +
+ + ) => + setNewAccount({ + ...newAccount, + apiKey: (e.target as HTMLInputElement).value, + }) + } + placeholder="Enter your z.ai API key" + /> +
+ )}
+
+ + +
)} {authStep === "form" ? ( diff --git a/packages/dashboard-web/src/components/accounts/AccountList.tsx b/packages/dashboard-web/src/components/accounts/AccountList.tsx index 2c6c0f4b..b36b29d4 100644 --- a/packages/dashboard-web/src/components/accounts/AccountList.tsx +++ b/packages/dashboard-web/src/components/accounts/AccountList.tsx @@ -6,6 +6,7 @@ interface AccountListProps { onPauseToggle: (account: Account) => void; onRemove: (name: string) => void; onRename: (account: Account) => void; + onPriorityChange: (account: Account) => void; } export function AccountList({ @@ -13,6 +14,7 @@ export function AccountList({ onPauseToggle, onRemove, onRename, + onPriorityChange, }: AccountListProps) { if (!accounts || accounts.length === 0) { return

No accounts configured

; @@ -45,6 +47,7 @@ export function AccountList({ onPauseToggle={onPauseToggle} onRemove={onRemove} onRename={onRename} + onPriorityChange={onPriorityChange} /> ))} diff --git a/packages/dashboard-web/src/components/accounts/AccountListItem.tsx b/packages/dashboard-web/src/components/accounts/AccountListItem.tsx index 35e07aa6..e9a4a4bb 100644 --- a/packages/dashboard-web/src/components/accounts/AccountListItem.tsx +++ b/packages/dashboard-web/src/components/accounts/AccountListItem.tsx @@ -6,6 +6,7 @@ import { Pause, Play, Trash2, + Zap, } from "lucide-react"; import type { Account } from "../../api"; import { Button } from "../ui/button"; @@ -17,6 +18,7 @@ interface AccountListItemProps { onPauseToggle: (account: Account) => void; onRemove: (name: string) => void; onRename: (account: Account) => void; + onPriorityChange: (account: Account) => void; } export function AccountListItem({ @@ -25,6 +27,7 @@ export function AccountListItem({ onPauseToggle, onRemove, onRename, + onPriorityChange, }: AccountListItemProps) { const presenter = new AccountPresenter(account); @@ -47,6 +50,9 @@ export function AccountListItem({ Active )} + + Priority: {account.priority} +

{account.provider} โ€ข {presenter.tierDisplay} @@ -78,6 +84,14 @@ export function AccountListItem({ > + + + + + + ); +} diff --git a/packages/dashboard-web/src/components/accounts/index.ts b/packages/dashboard-web/src/components/accounts/index.ts index 7b5232db..51251017 100644 --- a/packages/dashboard-web/src/components/accounts/index.ts +++ b/packages/dashboard-web/src/components/accounts/index.ts @@ -1,6 +1,7 @@ export { AccountAddForm } from "./AccountAddForm"; export { AccountList } from "./AccountList"; export { AccountListItem } from "./AccountListItem"; +export { AccountPriorityDialog } from "./AccountPriorityDialog"; export { DeleteConfirmationDialog } from "./DeleteConfirmationDialog"; export { RateLimitProgress } from "./RateLimitProgress"; export { RenameAccountDialog } from "./RenameAccountDialog"; diff --git a/packages/dashboard-web/src/components/charts/ModelPerformanceComparison.tsx b/packages/dashboard-web/src/components/charts/ModelPerformanceComparison.tsx index 03a49156..5400b0f4 100644 --- a/packages/dashboard-web/src/components/charts/ModelPerformanceComparison.tsx +++ b/packages/dashboard-web/src/components/charts/ModelPerformanceComparison.tsx @@ -40,6 +40,9 @@ const MODEL_COLORS: Record = { "claude-3.5-haiku": COLORS.success, "claude-3-opus": COLORS.blue, "claude-opus-4": COLORS.pink, + "claude-opus-4.1": COLORS.indigo, + "claude-sonnet-4": COLORS.cyan, + "claude-sonnet-4.5": COLORS.purple, }; function getModelColor(model: string): string { diff --git a/packages/dashboard-web/src/components/charts/ModelTokenSpeedChart.tsx b/packages/dashboard-web/src/components/charts/ModelTokenSpeedChart.tsx index 4a68622f..9fec577c 100644 --- a/packages/dashboard-web/src/components/charts/ModelTokenSpeedChart.tsx +++ b/packages/dashboard-web/src/components/charts/ModelTokenSpeedChart.tsx @@ -34,7 +34,9 @@ const MODEL_COLORS: Record = { "claude-3.5-haiku": COLORS.success, "claude-3-opus": COLORS.blue, "claude-opus-4": COLORS.pink, - // Add more models as needed + "claude-opus-4.1": COLORS.indigo, + "claude-sonnet-4": COLORS.cyan, + "claude-sonnet-4.5": COLORS.purple, }; function getModelColor(model: string): string { diff --git a/packages/dashboard-web/src/components/charts/MultiModelChart.tsx b/packages/dashboard-web/src/components/charts/MultiModelChart.tsx index 9d384c7b..a393b19d 100644 --- a/packages/dashboard-web/src/components/charts/MultiModelChart.tsx +++ b/packages/dashboard-web/src/components/charts/MultiModelChart.tsx @@ -53,6 +53,9 @@ const MODEL_COLORS: Record = { "claude-3.5-haiku": COLORS.success, "claude-3-opus": COLORS.blue, "claude-opus-4": COLORS.pink, + "claude-opus-4.1": COLORS.indigo, + "claude-sonnet-4": COLORS.cyan, + "claude-sonnet-4.5": COLORS.purple, }; function getModelColor(model: string, index: number): string { diff --git a/packages/dashboard-web/src/components/conversation/Message.tsx b/packages/dashboard-web/src/components/conversation/Message.tsx index 0e9e919f..5db16d93 100644 --- a/packages/dashboard-web/src/components/conversation/Message.tsx +++ b/packages/dashboard-web/src/components/conversation/Message.tsx @@ -43,7 +43,8 @@ function MessageComponent({ ); const thinkingText = typeof thinkingBlock?.thinking === "string" ? thinkingBlock.thinking : ""; - const hasThinking = thinkingText && cleanLineNumbers(thinkingText).trim().length > 0; + const hasThinking = + thinkingText && cleanLineNumbers(thinkingText).trim().length > 0; const cleanedContent = typeof content === "string" ? cleanLineNumbers(content).trim() : ""; const hasTools = tools?.length || 0; diff --git a/packages/dashboard-web/src/components/navigation.tsx b/packages/dashboard-web/src/components/navigation.tsx index 317ee65b..2fbe9d8a 100644 --- a/packages/dashboard-web/src/components/navigation.tsx +++ b/packages/dashboard-web/src/components/navigation.tsx @@ -148,7 +148,7 @@ export function Navigation() {

- v1.0.0 + v1.0.1
diff --git a/packages/database/src/async-writer.ts b/packages/database/src/async-writer.ts index 92c0fc8a..e2604179 100644 --- a/packages/database/src/async-writer.ts +++ b/packages/database/src/async-writer.ts @@ -9,6 +9,8 @@ export class AsyncDbWriter implements Disposable { private queue: DbJob[] = []; private running = false; private intervalId: Timer | null = null; + private readonly MAX_QUEUE_SIZE = 10000; // Prevent unbounded growth + private droppedJobs = 0; constructor() { // Process queue every 100ms @@ -16,6 +18,18 @@ export class AsyncDbWriter implements Disposable { } enqueue(job: DbJob): void { + // Check queue size limit + if (this.queue.length >= this.MAX_QUEUE_SIZE) { + this.droppedJobs++; + if (this.droppedJobs % 100 === 1) { + // Log every 100 dropped jobs to avoid log spam + logger.warn( + `Queue at capacity (${this.MAX_QUEUE_SIZE}), dropping jobs. Total dropped: ${this.droppedJobs}`, + ); + } + return; + } + this.queue.push(job); // Immediately try to process if not already running void this.processQueue(); @@ -57,6 +71,7 @@ export class AsyncDbWriter implements Disposable { logger.info("Async DB writer queue flushed", { remainingJobs: this.queue.length, + droppedJobs: this.droppedJobs, }); } } diff --git a/packages/database/src/database-operations.ts b/packages/database/src/database-operations.ts index ea6178d7..2cc26627 100644 --- a/packages/database/src/database-operations.ts +++ b/packages/database/src/database-operations.ts @@ -1,6 +1,7 @@ import { Database } from "bun:sqlite"; import { mkdirSync } from "node:fs"; import { dirname } from "node:path"; +import type { RuntimeConfig } from "@ccflare/config"; import type { Disposable } from "@ccflare/core"; import type { Account, StrategyStore } from "@ccflare/types"; import { ensureSchema, runMigrations } from "./migrations"; @@ -14,9 +15,102 @@ import { } from "./repositories/request.repository"; import { StatsRepository } from "./repositories/stats.repository"; import { StrategyRepository } from "./repositories/strategy.repository"; +import { withDatabaseRetrySync } from "./retry"; + +export interface DatabaseConfig { + /** Enable WAL (Write-Ahead Logging) mode for better concurrency */ + walMode?: boolean; + /** SQLite busy timeout in milliseconds */ + busyTimeoutMs?: number; + /** Cache size in pages (negative value = KB) */ + cacheSize?: number; + /** Synchronous mode: OFF, NORMAL, FULL */ + synchronous?: "OFF" | "NORMAL" | "FULL"; + /** Memory-mapped I/O size in bytes */ + mmapSize?: number; + /** Retry configuration for database operations */ + retry?: DatabaseRetryConfig; +} -export interface RuntimeConfig { - sessionDurationMs?: number; +export interface DatabaseRetryConfig { + /** Maximum number of retry attempts for database operations */ + attempts?: number; + /** Initial delay between retries in milliseconds */ + delayMs?: number; + /** Backoff multiplier for exponential backoff */ + backoff?: number; + /** Maximum delay between retries in milliseconds */ + maxDelayMs?: number; +} + +/** + * Apply SQLite pragmas for optimal performance on distributed filesystems + * Integrates your performance improvements with the new architecture + */ +function configureSqlite(db: Database, config: DatabaseConfig): void { + try { + // Check database integrity first + const integrityResult = db.query("PRAGMA integrity_check").get() as { + integrity_check: string; + }; + if (integrityResult.integrity_check !== "ok") { + throw new Error( + `Database integrity check failed: ${integrityResult.integrity_check}`, + ); + } + + // Enable WAL mode for better concurrency (with error handling) + if (config.walMode !== false) { + try { + const result = db.query("PRAGMA journal_mode = WAL").get() as { + journal_mode: string; + }; + if (result.journal_mode !== "wal") { + console.warn( + "Failed to enable WAL mode, falling back to DELETE mode", + ); + db.run("PRAGMA journal_mode = DELETE"); + } + } catch (error) { + console.warn("WAL mode failed, using DELETE mode:", error); + db.run("PRAGMA journal_mode = DELETE"); + } + } + + // Set busy timeout for lock handling + if (config.busyTimeoutMs !== undefined) { + db.run(`PRAGMA busy_timeout = ${config.busyTimeoutMs}`); + } + + // Configure cache size + if (config.cacheSize !== undefined) { + db.run(`PRAGMA cache_size = ${config.cacheSize}`); + } + + // Set synchronous mode (more conservative for distributed filesystems) + const syncMode = config.synchronous || "FULL"; // Default to FULL for safety + db.run(`PRAGMA synchronous = ${syncMode}`); + + // Configure memory-mapped I/O (disable on distributed filesystems if problematic) + if (config.mmapSize !== undefined && config.mmapSize > 0) { + try { + db.run(`PRAGMA mmap_size = ${config.mmapSize}`); + } catch (error) { + console.warn("Memory-mapped I/O failed, disabling:", error); + db.run("PRAGMA mmap_size = 0"); + } + } + + // Additional optimizations for distributed filesystems + db.run("PRAGMA temp_store = MEMORY"); + db.run("PRAGMA foreign_keys = ON"); + + // Add checkpoint interval for WAL mode + db.run("PRAGMA wal_autocheckpoint = 1000"); + } catch (error) { + console.error("Database configuration failed:", error); + throw new Error(`Failed to configure SQLite database: ${error}`); + } } /** @@ -26,6 +120,8 @@ export interface RuntimeConfig { export class DatabaseOperations implements StrategyStore, Disposable { private db: Database; private runtime?: RuntimeConfig; + private dbConfig: DatabaseConfig; + private retryConfig: DatabaseRetryConfig; // Repositories private accounts: AccountRepository; @@ -35,19 +131,41 @@ export class DatabaseOperations implements StrategyStore, Disposable { private stats: StatsRepository; private agentPreferences: AgentPreferenceRepository; - constructor(dbPath?: string) { + constructor( + dbPath?: string, + dbConfig?: DatabaseConfig, + retryConfig?: DatabaseRetryConfig, + ) { const resolvedPath = dbPath ?? resolveDbPath(); + // Default database configuration optimized for distributed filesystems + // More conservative settings to prevent corruption on Rook Ceph + this.dbConfig = { + walMode: true, + busyTimeoutMs: 10000, // Increased timeout for distributed storage + cacheSize: -10000, // Reduced cache size (10MB) for stability + synchronous: "FULL", // Full synchronous mode for data safety + mmapSize: 0, // Disable memory-mapped I/O on distributed filesystems + ...dbConfig, + }; + + // Default retry configuration for database operations + this.retryConfig = { + attempts: 3, + delayMs: 100, + backoff: 2, + maxDelayMs: 5000, + ...retryConfig, + }; + // Ensure the directory exists const dir = dirname(resolvedPath); mkdirSync(dir, { recursive: true }); this.db = new Database(resolvedPath, { create: true }); - // Configure SQLite for better concurrency - this.db.exec("PRAGMA journal_mode = WAL"); // Enable Write-Ahead Logging - this.db.exec("PRAGMA busy_timeout = 5000"); // Wait up to 5 seconds before throwing "database is locked" - this.db.exec("PRAGMA synchronous = NORMAL"); // Better performance while maintaining safety + // Apply SQLite configuration for distributed filesystem optimization + configureSqlite(this.db, this.dbConfig); ensureSchema(this.db); runMigrations(this.db); @@ -63,19 +181,46 @@ export class DatabaseOperations implements StrategyStore, Disposable { setRuntimeConfig(runtime: RuntimeConfig): void { this.runtime = runtime; + + // Update retry config from runtime config if available + if (runtime.database?.retry) { + this.retryConfig = { + ...this.retryConfig, + ...runtime.database.retry, + }; + } } getDatabase(): Database { return this.db; } - // Account operations delegated to repository + /** + * Get the current retry configuration + */ + getRetryConfig(): DatabaseRetryConfig { + return this.retryConfig; + } + + // Account operations delegated to repository with retry logic getAllAccounts(): Account[] { - return this.accounts.findAll(); + return withDatabaseRetrySync( + () => { + return this.accounts.findAll(); + }, + this.retryConfig, + "getAllAccounts", + ); } getAccount(accountId: string): Account | null { - return this.accounts.findById(accountId); + return withDatabaseRetrySync( + () => { + return this.accounts.findById(accountId); + }, + this.retryConfig, + "getAccount", + ); } updateAccountTokens( @@ -84,17 +229,40 @@ export class DatabaseOperations implements StrategyStore, Disposable { expiresAt: number, refreshToken?: string, ): void { - this.accounts.updateTokens(accountId, accessToken, expiresAt, refreshToken); + withDatabaseRetrySync( + () => { + this.accounts.updateTokens( + accountId, + accessToken, + expiresAt, + refreshToken, + ); + }, + this.retryConfig, + "updateAccountTokens", + ); } updateAccountUsage(accountId: string): void { const sessionDuration = this.runtime?.sessionDurationMs || 5 * 60 * 60 * 1000; - this.accounts.incrementUsage(accountId, sessionDuration); + withDatabaseRetrySync( + () => { + this.accounts.incrementUsage(accountId, sessionDuration); + }, + this.retryConfig, + "updateAccountUsage", + ); } markAccountRateLimited(accountId: string, until: number): void { - this.accounts.setRateLimited(accountId, until); + withDatabaseRetrySync( + () => { + this.accounts.setRateLimited(accountId, until); + }, + this.retryConfig, + "markAccountRateLimited", + ); } updateAccountRateLimitMeta( @@ -130,6 +298,10 @@ export class DatabaseOperations implements StrategyStore, Disposable { this.accounts.updateRequestCount(accountId, count); } + updateAccountPriority(accountId: string, priority: number): void { + this.accounts.updatePriority(accountId, priority); + } + // Request operations delegated to repository saveRequestMeta( id: string, diff --git a/packages/database/src/factory.ts b/packages/database/src/factory.ts index 854e020e..108e6b72 100644 --- a/packages/database/src/factory.ts +++ b/packages/database/src/factory.ts @@ -1,5 +1,10 @@ +import type { RuntimeConfig } from "@ccflare/config"; import { registerDisposable, unregisterDisposable } from "@ccflare/core"; -import { DatabaseOperations, type RuntimeConfig } from "./index"; +import { + type DatabaseConfig, + DatabaseOperations, + type DatabaseRetryConfig, +} from "./database-operations"; let instance: DatabaseOperations | null = null; let dbPath: string | undefined; @@ -15,7 +20,31 @@ export function initialize( export function getInstance(): DatabaseOperations { if (!instance) { - instance = new DatabaseOperations(dbPath); + // Extract database configuration from runtime config + const dbConfig: DatabaseConfig | undefined = runtimeConfig?.database + ? { + ...(runtimeConfig.database.walMode !== undefined && { + walMode: runtimeConfig.database.walMode, + }), + ...(runtimeConfig.database.busyTimeoutMs !== undefined && { + busyTimeoutMs: runtimeConfig.database.busyTimeoutMs, + }), + ...(runtimeConfig.database.cacheSize !== undefined && { + cacheSize: runtimeConfig.database.cacheSize, + }), + ...(runtimeConfig.database.synchronous !== undefined && { + synchronous: runtimeConfig.database.synchronous, + }), + ...(runtimeConfig.database.mmapSize !== undefined && { + mmapSize: runtimeConfig.database.mmapSize, + }), + } + : undefined; + + const retryConfig: DatabaseRetryConfig | undefined = + runtimeConfig?.database?.retry; + + instance = new DatabaseOperations(dbPath, dbConfig, retryConfig); if (runtimeConfig) { instance.setRuntimeConfig(runtimeConfig); } diff --git a/packages/database/src/index.ts b/packages/database/src/index.ts index da488b0d..d1b3c4b7 100644 --- a/packages/database/src/index.ts +++ b/packages/database/src/index.ts @@ -2,9 +2,13 @@ import { DatabaseOperations } from "./database-operations"; export { DatabaseOperations }; +export type { RuntimeConfig } from "@ccflare/config"; // Re-export other utilities export { AsyncDbWriter } from "./async-writer"; -export type { RuntimeConfig } from "./database-operations"; +export type { + DatabaseConfig, + DatabaseRetryConfig, +} from "./database-operations"; export { DatabaseFactory } from "./factory"; export { ensureSchema, runMigrations } from "./migrations"; export { resolveDbPath } from "./paths"; @@ -12,3 +16,6 @@ export { analyzeIndexUsage } from "./performance-indexes"; // Re-export repository types export type { StatsRepository } from "./repositories/stats.repository"; + +// Re-export retry utilities for external use (from your improvements) +export { withDatabaseRetry, withDatabaseRetrySync } from "./retry"; diff --git a/packages/database/src/migrations.ts b/packages/database/src/migrations.ts index 0afb29eb..f35dcb4e 100644 --- a/packages/database/src/migrations.ts +++ b/packages/database/src/migrations.ts @@ -19,7 +19,8 @@ export function ensureSchema(db: Database): void { last_used INTEGER, request_count INTEGER DEFAULT 0, total_requests INTEGER DEFAULT 0, - account_tier INTEGER DEFAULT 1 + account_tier INTEGER DEFAULT 1, + priority INTEGER DEFAULT 0 ) `); @@ -50,11 +51,21 @@ export function ensureSchema(db: Database): void { ) `); - // Create index for faster queries + // Create indexes for faster queries db.run( `CREATE INDEX IF NOT EXISTS idx_requests_timestamp ON requests(timestamp DESC)`, ); + // Index for JOIN performance with accounts table + db.run( + `CREATE INDEX IF NOT EXISTS idx_requests_account_used ON requests(account_used)`, + ); + + // Composite index for the main requests query (timestamp DESC with account_used for JOIN) + db.run( + `CREATE INDEX IF NOT EXISTS idx_requests_timestamp_account ON requests(timestamp DESC, account_used)`, + ); + // Create request_payloads table for storing full request/response data db.run(` CREATE TABLE IF NOT EXISTS request_payloads ( @@ -170,6 +181,14 @@ export function runMigrations(db: Database): void { log.info("Added rate_limit_remaining column to accounts table"); } + // Add priority column if it doesn't exist + if (!accountsColumnNames.includes("priority")) { + db.prepare( + "ALTER TABLE accounts ADD COLUMN priority INTEGER DEFAULT 0", + ).run(); + log.info("Added priority column to accounts table"); + } + // Check columns in requests table const requestsInfo = db .prepare("PRAGMA table_info(requests)") diff --git a/packages/database/src/performance-indexes.ts b/packages/database/src/performance-indexes.ts index 17b6ec71..18653548 100644 --- a/packages/database/src/performance-indexes.ts +++ b/packages/database/src/performance-indexes.ts @@ -103,11 +103,18 @@ export function addPerformanceIndexes(db: Database): void { // Composite index for account ordering in load balancer db.run(` - CREATE INDEX IF NOT EXISTS idx_accounts_request_count + CREATE INDEX IF NOT EXISTS idx_accounts_request_count ON accounts(request_count DESC, last_used) `); log.info("Added index: idx_accounts_request_count"); + // Index for account priority in load balancer + db.run(` + CREATE INDEX IF NOT EXISTS idx_accounts_priority + ON accounts(priority ASC, request_count DESC, last_used) + `); + log.info("Added index: idx_accounts_priority"); + log.info("Performance indexes added successfully"); } diff --git a/packages/database/src/repositories/account.repository.ts b/packages/database/src/repositories/account.repository.ts index 4aedd72d..e15ca043 100644 --- a/packages/database/src/repositories/account.repository.ts +++ b/packages/database/src/repositories/account.repository.ts @@ -4,14 +4,16 @@ import { BaseRepository } from "./base.repository"; export class AccountRepository extends BaseRepository { findAll(): Account[] { const rows = this.query(` - SELECT + SELECT id, name, provider, api_key, refresh_token, access_token, expires_at, created_at, last_used, request_count, total_requests, rate_limited_until, session_start, session_request_count, COALESCE(account_tier, 1) as account_tier, COALESCE(paused, 0) as paused, - rate_limit_reset, rate_limit_status, rate_limit_remaining + rate_limit_reset, rate_limit_status, rate_limit_remaining, + COALESCE(priority, 0) as priority FROM accounts + ORDER BY priority DESC `); return rows.map(toAccount); } @@ -19,13 +21,14 @@ export class AccountRepository extends BaseRepository { findById(accountId: string): Account | null { const row = this.get( ` - SELECT + SELECT id, name, provider, api_key, refresh_token, access_token, expires_at, created_at, last_used, request_count, total_requests, rate_limited_until, session_start, session_request_count, COALESCE(account_tier, 1) as account_tier, COALESCE(paused, 0) as paused, - rate_limit_reset, rate_limit_status, rate_limit_remaining + rate_limit_reset, rate_limit_status, rate_limit_remaining, + COALESCE(priority, 0) as priority FROM accounts WHERE id = ? `, @@ -128,4 +131,11 @@ export class AccountRepository extends BaseRepository { rename(accountId: string, newName: string): void { this.run(`UPDATE accounts SET name = ? WHERE id = ?`, [newName, accountId]); } + + updatePriority(accountId: string, priority: number): void { + this.run(`UPDATE accounts SET priority = ? WHERE id = ?`, [ + priority, + accountId, + ]); + } } diff --git a/packages/database/src/retry.ts b/packages/database/src/retry.ts new file mode 100644 index 00000000..8cd592e0 --- /dev/null +++ b/packages/database/src/retry.ts @@ -0,0 +1,231 @@ +import { Logger } from "@ccflare/logger"; +import type { DatabaseRetryConfig } from "./database-operations"; + +const logger = new Logger("db-retry"); + +/** + * Error codes that indicate database lock contention and should trigger retries + */ +const RETRYABLE_SQLITE_ERRORS = [ + "SQLITE_BUSY", + "SQLITE_LOCKED", + "database is locked", + "database table is locked", +]; + +/** + * Check if an error is retryable (indicates database lock contention) + */ +function isRetryableError(error: unknown): boolean { + if (!error) return false; + + const errorMessage = error instanceof Error ? error.message : String(error); + const errorCode = (error as any)?.code; + + return RETRYABLE_SQLITE_ERRORS.some( + (retryableError) => + errorMessage.includes(retryableError) || errorCode === retryableError, + ); +} + +/** + * Calculate delay for exponential backoff with jitter + */ +function calculateDelay( + attempt: number, + config: Required, +): number { + const baseDelay = config.delayMs * config.backoff ** attempt; + const jitter = Math.random() * 0.1 * baseDelay; // Add 10% jitter + const delayWithJitter = baseDelay + jitter; + + return Math.min(delayWithJitter, config.maxDelayMs); +} + +/** + * Sleep for the specified number of milliseconds + */ +function sleep(ms: number): Promise { + return new Promise((resolve) => setTimeout(resolve, ms)); +} + +/** + * Synchronous sleep function + */ +function sleepSync(ms: number): void { + // Synchronous sleep using Bun.sleepSync if available, otherwise Node.js fallback + if (typeof Bun !== "undefined" && Bun.sleepSync) { + Bun.sleepSync(ms); + } else { + // Try Node.js child_process.spawnSync as fallback + try { + const { spawnSync } = require("node:child_process"); + const sleepCommand = process.platform === "win32" ? "timeout" : "sleep"; + const sleepArg = + process.platform === "win32" + ? `/t ${Math.ceil(ms / 1000)}` + : `${ms / 1000}`; + + spawnSync(sleepCommand, [sleepArg], { + stdio: "ignore", + shell: process.platform === "win32", + }); + } catch (error) { + // If child_process is not available or fails, throw an error instead of busy waiting + throw new Error( + `Synchronous sleep not supported in this environment. ` + + `Bun.sleepSync is not available and Node.js child_process failed: ${error instanceof Error ? error.message : String(error)}`, + ); + } + } +} + +/** + * Synchronous retry logic + */ +function executeWithRetrySync( + operation: () => T, + config: Required, + operationName: string, +): T { + let lastError: unknown; + + for (let attempt = 0; attempt < config.attempts; attempt++) { + try { + const result = operation(); + + // Log successful retry if this wasn't the first attempt + if (attempt > 0) { + logger.info(`${operationName} succeeded after ${attempt + 1} attempts`); + } + + return result; + } catch (error) { + lastError = error; + + // Check if this is a retryable error + if (!isRetryableError(error)) { + logger.debug( + `${operationName} failed with non-retryable error:`, + error, + ); + throw error; + } + + // If this was the last attempt, throw the error + if (attempt === config.attempts - 1) { + logger.error( + `${operationName} failed after ${config.attempts} attempts:`, + error, + ); + throw error; + } + + // Calculate delay and wait before retry + const delay = calculateDelay(attempt, config); + logger.warn( + `${operationName} failed (attempt ${attempt + 1}/${config.attempts}), retrying in ${delay.toFixed(0)}ms:`, + error instanceof Error ? error.message : String(error), + ); + + sleepSync(delay); + } + } + + // This should never be reached, but TypeScript requires it + throw lastError; +} + +/** + * Async retry logic - uses iterative approach to avoid recursive Promise chains + */ +async function executeWithRetryAsync( + operation: () => T | Promise, + config: Required, + operationName: string, +): Promise { + let lastError: unknown; + + for (let attempt = 0; attempt < config.attempts; attempt++) { + try { + const result = await operation(); + + // Log successful retry if this wasn't the first attempt + if (attempt > 0) { + logger.info(`${operationName} succeeded after ${attempt + 1} attempts`); + } + + return result; + } catch (error) { + lastError = error; + + // Check if this is a retryable error + if (!isRetryableError(error)) { + logger.debug( + `${operationName} failed with non-retryable error:`, + error, + ); + throw error; + } + + // If this was the last attempt, throw the error + if (attempt === config.attempts - 1) { + logger.error( + `${operationName} failed after ${config.attempts} attempts:`, + error, + ); + throw error; + } + + // Calculate delay and wait before retry + const delay = calculateDelay(attempt, config); + logger.warn( + `${operationName} failed (attempt ${attempt + 1}/${config.attempts}), retrying in ${delay.toFixed(0)}ms:`, + error instanceof Error ? error.message : String(error), + ); + + await sleep(delay); + } + } + + // This should never be reached, but TypeScript requires it + throw lastError; +} + +/** + * Retry wrapper for database operations with exponential backoff + */ +export async function withDatabaseRetry( + operation: () => T | Promise, + config: DatabaseRetryConfig = {}, + operationName = "database operation", +): Promise { + const retryConfig: Required = { + attempts: 3, + delayMs: 100, + backoff: 2, + maxDelayMs: 5000, + ...config, + }; + + return executeWithRetryAsync(operation, retryConfig, operationName); +} + +/** + * Synchronous retry wrapper for database operations + */ +export function withDatabaseRetrySync( + operation: () => T, + config: DatabaseRetryConfig = {}, + operationName = "database operation", +): T { + const retryConfig: Required = { + attempts: 3, + delayMs: 100, + backoff: 2, + maxDelayMs: 5000, + ...config, + }; + + return executeWithRetrySync(operation, retryConfig, operationName); +} diff --git a/packages/http-api/src/handlers/accounts.ts b/packages/http-api/src/handlers/accounts.ts index f1e4785d..807c31a0 100644 --- a/packages/http-api/src/handlers/accounts.ts +++ b/packages/http-api/src/handlers/accounts.ts @@ -31,7 +31,7 @@ export function createAccountsListHandler(db: Database) { const accounts = db .query( ` - SELECT + SELECT id, name, provider, @@ -48,11 +48,12 @@ export function createAccountsListHandler(db: Database) { session_request_count, COALESCE(account_tier, 1) as account_tier, COALESCE(paused, 0) as paused, - CASE - WHEN expires_at > ?1 THEN 1 - ELSE 0 + COALESCE(priority, 0) as priority, + CASE + WHEN expires_at > ?1 THEN 1 + ELSE 0 END as token_valid, - CASE + CASE WHEN rate_limited_until > ?2 THEN 1 ELSE 0 END as rate_limited, @@ -62,8 +63,8 @@ export function createAccountsListHandler(db: Database) { ELSE '-' END as session_info FROM accounts - ORDER BY request_count DESC - `, + ORDER BY priority DESC, request_count DESC + `, ) .all(now, now, now, sessionDuration) as Array<{ id: string; @@ -82,6 +83,7 @@ export function createAccountsListHandler(db: Database) { session_request_count: number; account_tier: number; paused: 0 | 1; + priority: number; token_valid: 0 | 1; rate_limited: 0 | 1; session_info: string | null; @@ -123,6 +125,7 @@ export function createAccountsListHandler(db: Database) { created: new Date(account.created_at).toISOString(), tier: account.account_tier, paused: account.paused === 1, + priority: account.priority, tokenStatus: account.token_valid ? "valid" : "expired", tokenExpiresAt: account.expires_at ? new Date(account.expires_at).toISOString() @@ -169,6 +172,47 @@ export function createAccountTierUpdateHandler(dbOps: DatabaseOperations) { }; } +/** + * Create an account priority update handler + */ +export function createAccountPriorityUpdateHandler(dbOps: DatabaseOperations) { + return async (req: Request, accountId: string): Promise => { + try { + const body = await req.json(); + + // Validate priority input + const priority = validateNumber(body.priority, "priority", { + required: true, + min: 0, + max: 100, + integer: true, + }); + + if (priority === undefined) { + return errorResponse(BadRequest("Priority is required")); + } + + // Check if account exists + const db = dbOps.getDatabase(); + const account = db + .query<{ id: string }, [string]>("SELECT id FROM accounts WHERE id = ?") + .get(accountId); + + if (!account) { + return errorResponse(NotFound("Account not found")); + } + + dbOps.updateAccountPriority(accountId, priority); + + return jsonResponse({ success: true, priority }); + } catch (_error) { + return errorResponse( + InternalServerError("Failed to update account priority"), + ); + } + }; +} + /** * Create an account add handler (manual token addition) * This is primarily used for adding accounts with existing tokens @@ -223,6 +267,14 @@ export function createAccountAddHandler( allowedValues: [1, 5, 20] as const, }) || 1) as 1 | 5 | 20; + // Validate priority + const priority = + validateNumber(body.priority, "priority", { + min: 0, + max: 100, + integer: true, + }) || 0; + try { // Add account directly to database const accountId = crypto.randomUUID(); @@ -231,15 +283,25 @@ export function createAccountAddHandler( dbOps.getDatabase().run( `INSERT INTO accounts ( id, name, provider, refresh_token, access_token, - created_at, request_count, total_requests, account_tier - ) VALUES (?, ?, ?, ?, ?, ?, 0, 0, ?)`, - [accountId, name, provider, refreshToken, accessToken, now, tier], + created_at, request_count, total_requests, account_tier, priority + ) VALUES (?, ?, ?, ?, ?, ?, 0, 0, ?, ?)`, + [ + accountId, + name, + provider, + refreshToken, + accessToken, + now, + tier, + priority, + ], ); return jsonResponse({ success: true, message: `Account ${name} added successfully`, tier, + priority, accountId, }); } catch (error) { @@ -434,3 +496,142 @@ export function createAccountRenameHandler(dbOps: DatabaseOperations) { } }; } + +/** + * Create a z.ai account add handler + */ +export function createZaiAccountAddHandler(dbOps: DatabaseOperations) { + return async (req: Request): Promise => { + try { + const body = await req.json(); + + // Validate account name + const name = validateString(body.name, "name", { + required: true, + minLength: 1, + maxLength: 100, + pattern: patterns.accountName, + transform: sanitizers.trim, + }); + + if (!name) { + return errorResponse(BadRequest("Account name is required")); + } + + // Validate API key + const apiKey = validateString(body.apiKey, "apiKey", { + required: true, + minLength: 1, + }); + + if (!apiKey) { + return errorResponse(BadRequest("API key is required")); + } + + // Validate tier + const tier = (validateNumber(body.tier, "tier", { + allowedValues: [1, 5, 20] as const, + }) || 1) as 1 | 5 | 20; + + // Validate priority + const priority = + validateNumber(body.priority, "priority", { + min: 0, + max: 100, + integer: true, + }) || 0; + + // Create z.ai account directly in database + const accountId = crypto.randomUUID(); + const now = Date.now(); + + const db = dbOps.getDatabase(); + db.run( + `INSERT INTO accounts ( + id, name, provider, api_key, refresh_token, access_token, + expires_at, created_at, account_tier, request_count, total_requests, priority + ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)`, + [ + accountId, + name, + "zai", + apiKey, + apiKey, // Use API key as refresh token for consistency with CLI + apiKey, // Use API key as access token + now + 365 * 24 * 60 * 60 * 1000, // 1 year from now + now, + tier, + 0, + 0, + priority, + ], + ); + + log.info( + `Successfully added z.ai account: ${name} (Tier ${tier}, Priority ${priority})`, + ); + + // Get the created account for response + const account = db + .query< + { + id: string; + name: string; + provider: string; + request_count: number; + total_requests: number; + last_used: number | null; + created_at: number; + expires_at: number; + account_tier: number; + paused: number; + }, + [string] + >( + `SELECT + id, name, provider, request_count, total_requests, + last_used, created_at, expires_at, account_tier, + COALESCE(paused, 0) as paused + FROM accounts WHERE id = ?`, + ) + .get(accountId); + + if (!account) { + return errorResponse( + InternalServerError("Failed to retrieve created account"), + ); + } + + return jsonResponse({ + message: `z.ai account '${name}' added successfully`, + account: { + id: account.id, + name: account.name, + provider: account.provider, + requestCount: account.request_count, + totalRequests: account.total_requests, + lastUsed: account.last_used + ? new Date(account.last_used).toISOString() + : null, + created: new Date(account.created_at).toISOString(), + tier: account.account_tier, + paused: account.paused === 1, + priority: priority, + tokenStatus: "valid" as const, + tokenExpiresAt: new Date(account.expires_at).toISOString(), + rateLimitStatus: "OK", + rateLimitReset: null, + rateLimitRemaining: null, + sessionInfo: "No active session", + }, + }); + } catch (error) { + log.error("z.ai account creation error:", error); + return errorResponse( + error instanceof Error + ? error + : new Error("Failed to create z.ai account"), + ); + } + }; +} diff --git a/packages/http-api/src/handlers/requests.ts b/packages/http-api/src/handlers/requests.ts index cc584073..57449a5d 100644 --- a/packages/http-api/src/handlers/requests.ts +++ b/packages/http-api/src/handlers/requests.ts @@ -94,3 +94,22 @@ export function createRequestsDetailHandler(dbOps: DatabaseOperations) { return jsonResponse(parsed); }; } + +/** + * Create a handler for lazy loading individual request payloads + * This endpoint supports the performance optimization that eliminates JSON parsing bottleneck + */ +export function createRequestPayloadHandler(dbOps: DatabaseOperations) { + return (requestId: string): Response => { + const payload = dbOps.getRequestPayload(requestId); + + if (!payload) { + return new Response(JSON.stringify({ error: "Request not found" }), { + status: 404, + headers: { "Content-Type": "application/json" }, + }); + } + + return jsonResponse(payload); + }; +} diff --git a/packages/http-api/src/router.ts b/packages/http-api/src/router.ts index 82fa52d0..94101ec3 100644 --- a/packages/http-api/src/router.ts +++ b/packages/http-api/src/router.ts @@ -2,11 +2,13 @@ import { validateNumber } from "@ccflare/core"; import { createAccountAddHandler, createAccountPauseHandler, + createAccountPriorityUpdateHandler, createAccountRemoveHandler, createAccountRenameHandler, createAccountResumeHandler, createAccountsListHandler, createAccountTierUpdateHandler, + createZaiAccountAddHandler, } from "./handlers/accounts"; import { createAgentPreferenceUpdateHandler, @@ -29,6 +31,7 @@ import { createOAuthInitHandler, } from "./handlers/oauth"; import { + createRequestPayloadHandler, createRequestsDetailHandler, createRequestsSummaryHandler, } from "./handlers/requests"; @@ -62,6 +65,7 @@ export class APIRouter { const statsResetHandler = createStatsResetHandler(dbOps); const accountsHandler = createAccountsListHandler(db); const accountAddHandler = createAccountAddHandler(dbOps, config); + const zaiAccountAddHandler = createZaiAccountAddHandler(dbOps); const _accountRemoveHandler = createAccountRemoveHandler(dbOps); const _accountTierHandler = createAccountTierUpdateHandler(dbOps); const requestsSummaryHandler = createRequestsSummaryHandler(db); @@ -84,6 +88,9 @@ export class APIRouter { this.handlers.set("POST:/api/stats/reset", () => statsResetHandler()); this.handlers.set("GET:/api/accounts", () => accountsHandler()); this.handlers.set("POST:/api/accounts", (req) => accountAddHandler(req)); + this.handlers.set("POST:/api/accounts/zai", (req) => + zaiAccountAddHandler(req), + ); this.handlers.set("POST:/api/oauth/init", (req) => oauthInitHandler(req)); this.handlers.set("POST:/api/oauth/callback", (req) => oauthCallbackHandler(req), @@ -179,6 +186,19 @@ export class APIRouter { return await this.wrapHandler(handler)(req, url); } + // Check for dynamic request payload endpoint + if (path.startsWith("/api/requests/payload/") && method === "GET") { + const parts = path.split("/"); + const requestId = parts[4]; + if (requestId) { + const payloadHandler = createRequestPayloadHandler(this.context.dbOps); + return await this.wrapHandler(() => payloadHandler(requestId))( + req, + url, + ); + } + } + // Check for dynamic account endpoints if (path.startsWith("/api/accounts/")) { const parts = path.split("/"); @@ -220,6 +240,17 @@ export class APIRouter { ); } + // Account priority update + if (path.endsWith("/priority") && method === "POST") { + const priorityHandler = createAccountPriorityUpdateHandler( + this.context.dbOps, + ); + return await this.wrapHandler((req) => priorityHandler(req, accountId))( + req, + url, + ); + } + // Account removal if (parts.length === 4 && method === "DELETE") { const removeHandler = createAccountRemoveHandler(this.context.dbOps); diff --git a/packages/load-balancer/src/strategies/index.ts b/packages/load-balancer/src/strategies/index.ts index 0bb1162f..a3a0ecfa 100644 --- a/packages/load-balancer/src/strategies/index.ts +++ b/packages/load-balancer/src/strategies/index.ts @@ -71,24 +71,26 @@ export class SessionStrategy implements LoadBalancingStrategy { this.log.info( `Continuing session for account ${activeAccount.name} (${activeAccount.session_request_count} requests in session)`, ); - // Return active account first, then others as fallback - const others = accounts.filter( - (a) => a.id !== activeAccount.id && isAccountAvailable(a, now), - ); + // Return active account first, then others as fallback (sorted by priority) + const others = accounts + .filter((a) => a.id !== activeAccount.id && isAccountAvailable(a, now)) + .sort((a, b) => a.priority - b.priority); return [activeAccount, ...others]; } // No active session or active account is rate limited - // Filter available accounts - const available = accounts.filter((a) => isAccountAvailable(a, now)); + // Filter available accounts and sort by priority (lower number = higher priority) + const available = accounts + .filter((a) => isAccountAvailable(a, now)) + .sort((a, b) => a.priority - b.priority); if (available.length === 0) return []; - // Pick the first available account and start a new session with it + // Pick the highest priority account (first in sorted list) and start a new session with it const chosenAccount = available[0]; this.resetSessionIfExpired(chosenAccount); - // Return chosen account first, then others as fallback + // Return chosen account first, then others as fallback (already sorted by priority) const others = available.filter((a) => a.id !== chosenAccount.id); return [chosenAccount, ...others]; } diff --git a/packages/oauth-flow/src/index.ts b/packages/oauth-flow/src/index.ts index 6fd45139..df11edbe 100644 --- a/packages/oauth-flow/src/index.ts +++ b/packages/oauth-flow/src/index.ts @@ -27,6 +27,7 @@ export interface CompleteOptions { code: string; tier?: AccountTier; name: string; // Required to properly create the account + priority?: number; } export interface AccountCreated { @@ -132,7 +133,7 @@ export class OAuthFlow { opts: CompleteOptions, flowData: BeginResult, ): Promise { - const { code, tier = 1, name } = opts; + const { code, tier = 1, name, priority = 0 } = opts; // Get OAuth provider const oauthProvider = getOAuthProvider("anthropic"); @@ -152,11 +153,17 @@ export class OAuthFlow { // Handle console mode - create API key if (flowData.mode === "console" || !tokens.refreshToken) { const apiKey = await this.createAnthropicApiKey(tokens.accessToken); - return this.createAccountWithApiKey(accountId, name, apiKey, tier); + return this.createAccountWithApiKey( + accountId, + name, + apiKey, + tier, + priority, + ); } // Handle max mode - standard OAuth flow - return this.createAccountWithOAuth(accountId, name, tokens, tier); + return this.createAccountWithOAuth(accountId, name, tokens, tier, priority); } /** @@ -206,15 +213,16 @@ export class OAuthFlow { name: string, tokens: OAuthTokens, tier: AccountTier, + priority: number, ): AccountCreated { const db = this.dbOps.getDatabase(); db.run( ` INSERT INTO accounts ( - id, name, provider, api_key, refresh_token, access_token, expires_at, - created_at, request_count, total_requests, account_tier - ) VALUES (?, ?, ?, NULL, ?, ?, ?, ?, 0, 0, ?) + id, name, provider, api_key, refresh_token, access_token, expires_at, + created_at, request_count, total_requests, account_tier, priority + ) VALUES (?, ?, ?, NULL, ?, ?, ?, ?, 0, 0, ?, ?) `, [ id, @@ -225,6 +233,7 @@ export class OAuthFlow { tokens.expiresAt, Date.now(), tier, + priority, ], ); @@ -254,17 +263,18 @@ export class OAuthFlow { name: string, apiKey: string, tier: AccountTier, + priority: number, ): AccountCreated { const db = this.dbOps.getDatabase(); db.run( ` INSERT INTO accounts ( - id, name, provider, api_key, refresh_token, access_token, expires_at, - created_at, request_count, total_requests, account_tier - ) VALUES (?, ?, ?, ?, NULL, NULL, NULL, ?, 0, 0, ?) + id, name, provider, api_key, refresh_token, access_token, expires_at, + created_at, request_count, total_requests, account_tier, priority + ) VALUES (?, ?, ?, ?, NULL, NULL, NULL, ?, 0, 0, ?, ?) `, - [id, name, "anthropic", apiKey, Date.now(), tier], + [id, name, "anthropic", apiKey, Date.now(), tier, priority], ); return { diff --git a/packages/providers/src/index.ts b/packages/providers/src/index.ts index 2d974a6c..e81cbcc4 100644 --- a/packages/providers/src/index.ts +++ b/packages/providers/src/index.ts @@ -17,7 +17,9 @@ export { export * from "./types"; import { AnthropicProvider } from "./providers/anthropic/provider"; +import { ZaiProvider } from "./providers/zai/provider"; // Auto-register built-in providers import { registry } from "./registry"; registry.registerProvider(new AnthropicProvider()); +registry.registerProvider(new ZaiProvider()); diff --git a/packages/providers/src/providers/anthropic/provider.ts b/packages/providers/src/providers/anthropic/provider.ts index b4ee000d..141dea1a 100644 --- a/packages/providers/src/providers/anthropic/provider.ts +++ b/packages/providers/src/providers/anthropic/provider.ts @@ -233,10 +233,36 @@ export class AnthropicProvider extends BaseProvider { const maxBytes = BUFFER_SIZES.ANTHROPIC_STREAM_CAP_BYTES; const decoder = new TextDecoder(); let foundMessageStart = false; + const READ_TIMEOUT_MS = 10000; // 10 second timeout for stream reads + const startTime = Date.now(); try { while (buffered.length < maxBytes) { - const { value, done } = await reader.read(); + // Check for timeout + if (Date.now() - startTime > READ_TIMEOUT_MS) { + await reader.cancel(); + throw new Error( + "Stream read timeout while extracting usage info", + ); + } + + // Read with timeout + const readPromise = reader.read(); + const timeoutPromise = new Promise<{ + value?: Uint8Array; + done: boolean; + }>((_, reject) => + setTimeout( + () => reject(new Error("Read operation timeout")), + 5000, + ), + ); + + const { value, done } = await Promise.race([ + readPromise, + timeoutPromise, + ]); + if (done) break; buffered += decoder.decode(value, { stream: true }); @@ -245,7 +271,22 @@ export class AnthropicProvider extends BaseProvider { if (buffered.includes("event: message_start")) { foundMessageStart = true; // Read a bit more to ensure we get the data line - const { value: nextValue, done: nextDone } = await reader.read(); + const nextReadPromise = reader.read(); + const nextTimeoutPromise = new Promise<{ + value?: Uint8Array; + done: boolean; + }>((_, reject) => + setTimeout( + () => reject(new Error("Read operation timeout")), + 5000, + ), + ); + + const { value: nextValue, done: nextDone } = await Promise.race([ + nextReadPromise, + nextTimeoutPromise, + ]); + if (!nextDone && nextValue) { buffered += decoder.decode(nextValue, { stream: true }); } diff --git a/packages/providers/src/providers/index.ts b/packages/providers/src/providers/index.ts index fdd25ce1..dfcbfbc8 100644 --- a/packages/providers/src/providers/index.ts +++ b/packages/providers/src/providers/index.ts @@ -2,3 +2,4 @@ export { AnthropicOAuthProvider, AnthropicProvider, } from "./anthropic/index"; +export { ZaiProvider } from "./zai/index"; diff --git a/packages/providers/src/providers/zai/index.ts b/packages/providers/src/providers/zai/index.ts new file mode 100644 index 00000000..13d87298 --- /dev/null +++ b/packages/providers/src/providers/zai/index.ts @@ -0,0 +1 @@ +export { ZaiProvider } from "./provider"; diff --git a/packages/providers/src/providers/zai/provider.ts b/packages/providers/src/providers/zai/provider.ts new file mode 100644 index 00000000..1ee3928b --- /dev/null +++ b/packages/providers/src/providers/zai/provider.ts @@ -0,0 +1,403 @@ +import { estimateCostUSD } from "@ccflare/core"; +import { sanitizeProxyHeaders } from "@ccflare/http-common"; +import { Logger } from "@ccflare/logger"; +import type { Account } from "@ccflare/types"; +import { BaseProvider } from "../../base"; +import type { RateLimitInfo, TokenRefreshResult } from "../../types"; + +const log = new Logger("ZaiProvider"); + +export class ZaiProvider extends BaseProvider { + name = "zai"; + + canHandle(_path: string): boolean { + // Handle all paths for z.ai endpoints + return true; + } + + async refreshToken( + account: Account, + _clientId: string, + ): Promise { + // z.ai uses API keys, not OAuth tokens + // If refresh_token is actually an API key, return it as the access token + if (!account.refresh_token) { + throw new Error(`No API key available for account ${account.name}`); + } + + log.info(`Using API key for z.ai account ${account.name}`); + + // For API key based authentication, we don't have token refresh + // The "refresh_token" field stores the API key + return { + accessToken: account.refresh_token, + expiresAt: Date.now() + 365 * 24 * 60 * 60 * 1000, // 1 year from now + refreshToken: account.refresh_token, + }; + } + + buildUrl(path: string, query: string): string { + return `https://api.z.ai/api/anthropic${path}${query}`; + } + + prepareHeaders( + headers: Headers, + accessToken?: string, + apiKey?: string, + ): Headers { + const newHeaders = new Headers(headers); + + // z.ai expects the API key in x-api-key header + if (accessToken) { + newHeaders.set("x-api-key", accessToken); + } else if (apiKey) { + newHeaders.set("x-api-key", apiKey); + } + + // Remove authorization header since z.ai uses x-api-key + newHeaders.delete("authorization"); + + // Remove host header + newHeaders.delete("host"); + + // Remove compression headers to avoid decompression issues + newHeaders.delete("accept-encoding"); + newHeaders.delete("content-encoding"); + + return newHeaders; + } + + parseRateLimit(response: Response): RateLimitInfo { + // Check for standard rate limit headers + if (response.status !== 429) { + return { isRateLimited: false }; + } + + // Try to extract reset time from headers + const retryAfter = response.headers.get("retry-after"); + let resetTime: number | undefined; + + if (retryAfter) { + // Retry-After can be seconds or HTTP date + const seconds = Number(retryAfter); + if (!Number.isNaN(seconds)) { + resetTime = Date.now() + seconds * 1000; + } else { + resetTime = new Date(retryAfter).getTime(); + } + } + + return { isRateLimited: true, resetTime }; + } + + async processResponse( + response: Response, + _account: Account | null, + ): Promise { + // Sanitize headers by removing hop-by-hop headers + const headers = sanitizeProxyHeaders(response.headers); + + return new Response(response.body, { + status: response.status, + statusText: response.statusText, + headers, + }); + } + + async extractTierInfo(_response: Response): Promise { + // z.ai doesn't provide tier information in responses + // We'll rely on the account tier set during account creation + return null; + } + + async extractUsageInfo(response: Response): Promise<{ + model?: string; + promptTokens?: number; + completionTokens?: number; + totalTokens?: number; + costUsd?: number; + inputTokens?: number; + cacheReadInputTokens?: number; + cacheCreationInputTokens?: number; + outputTokens?: number; + } | null> { + try { + const clone = response.clone(); + const contentType = response.headers.get("content-type"); + + // Handle streaming responses (SSE) - similar to Anthropic + if (contentType?.includes("text/event-stream")) { + // For streaming, we'll extract usage from the message_start and message_delta events + // This is similar to the Anthropic implementation but handles both events + const reader = clone.body?.getReader(); + if (!reader) return null; + + let buffered = ""; + const maxBytes = 100000; // Larger buffer to capture more of the stream + const decoder = new TextDecoder(); + + // Track usage from both message_start and message_delta + let messageStartUsage: { + input_tokens?: number; + output_tokens?: number; + cache_creation_input_tokens?: number; + cache_read_input_tokens?: number; + model?: string; + } | null = null; + + let messageDeltaUsage: { + input_tokens?: number; + output_tokens?: number; + cache_read_input_tokens?: number; + } | null = null; + + try { + while (buffered.length < maxBytes) { + const { value, done } = await reader.read(); + if (done) break; + + buffered += decoder.decode(value, { stream: true }); + + // Process all complete lines in the buffer + const lines = buffered.split("\n"); + buffered = lines.pop() || ""; // Keep incomplete line in buffer + + for (let i = 0; i < lines.length; i++) { + const line = lines[i].trim(); + if (!line) continue; + + // Parse SSE event + if (line.startsWith("event: message_start")) { + // Look for the next data line, skipping empty lines + let dataLine = null; + for (let j = i + 1; j < lines.length; j++) { + const nextLine = lines[j].trim(); + if (nextLine.startsWith("data: ")) { + dataLine = nextLine; + break; + } else if (nextLine && !nextLine.startsWith("event: ")) { + // If we encounter a non-empty line that's not an event, break + break; + } + } + + if (dataLine) { + try { + const jsonStr = dataLine.slice(6); + const data = JSON.parse(jsonStr) as { + message?: { + model?: string; + usage?: { + input_tokens?: number; + output_tokens?: number; + cache_creation_input_tokens?: number; + cache_read_input_tokens?: number; + }; + }; + }; + + if (data.message?.usage) { + messageStartUsage = { + input_tokens: data.message.usage.input_tokens, + output_tokens: data.message.usage.output_tokens, + cache_creation_input_tokens: + data.message.usage.cache_creation_input_tokens, + cache_read_input_tokens: + data.message.usage.cache_read_input_tokens, + model: data.message.model, + }; + } + } catch { + // Ignore parse errors + } + } + } else if (line.startsWith("event: message_delta")) { + // Look for the next data line, skipping empty lines + let dataLine = null; + for (let j = i + 1; j < lines.length; j++) { + const nextLine = lines[j].trim(); + if (nextLine.startsWith("data: ")) { + dataLine = nextLine; + break; + } else if (nextLine && !nextLine.startsWith("event: ")) { + // If we encounter a non-empty line that's not an event, break + break; + } + } + + if (dataLine) { + try { + const jsonStr = dataLine.slice(6); + const data = JSON.parse(jsonStr) as { + usage?: { + input_tokens?: number; + output_tokens?: number; + cache_read_input_tokens?: number; + }; + }; + + if (data.usage) { + messageDeltaUsage = { + input_tokens: data.usage.input_tokens, + output_tokens: data.usage.output_tokens, + cache_read_input_tokens: + data.usage.cache_read_input_tokens, + }; + } + } catch { + // Ignore parse errors + } + } + } + } + + // If we have both message_start and message_delta, we can return the complete usage + if (messageDeltaUsage) { + break; // We have the final usage from message_delta + } + } + } finally { + reader.cancel().catch(() => {}); + } + + // For ZAI streaming responses, message_delta always contains the final authoritative token counts + // We should always prefer message_delta when available, regardless of whether tokens are zero + const finalUsage = messageDeltaUsage || messageStartUsage; + + if (finalUsage) { + // Use the model from message_start (z.ai returns GLM model names directly) + const model = messageStartUsage?.model; + + // For message_delta, input_tokens and cache_read_input_tokens may be the final counts + // For message_start, we have all the detailed breakdown + const inputTokens = + finalUsage.input_tokens || messageStartUsage?.input_tokens || 0; + const cacheReadInputTokens = + finalUsage.cache_read_input_tokens || + messageStartUsage?.cache_read_input_tokens || + 0; + const cacheCreationInputTokens = + messageStartUsage?.cache_creation_input_tokens || 0; + const outputTokens = + finalUsage.output_tokens || messageStartUsage?.output_tokens || 0; + + const promptTokens = + (inputTokens || 0) + + cacheReadInputTokens + + cacheCreationInputTokens; + const completionTokens = outputTokens; + const totalTokens = promptTokens + completionTokens; + + // Calculate cost if we have a model + let costUsd: number | undefined; + if (model) { + log.debug(`Calculating cost for z.ai model: ${model}`); + try { + costUsd = await estimateCostUSD(model, { + inputTokens, + outputTokens, + cacheReadInputTokens, + cacheCreationInputTokens, + }); + log.debug(`Cost calculated for ${model}: $${costUsd}`); + } catch (error) { + log.warn(`Failed to calculate cost for model ${model}:`, error); + } + } else { + log.debug(`No model found in z.ai response for cost calculation`); + } + + return { + model, + promptTokens, + completionTokens, + totalTokens, + inputTokens, + cacheReadInputTokens, + cacheCreationInputTokens, + outputTokens, + costUsd, + }; + } + + return null; + } else { + // Handle non-streaming JSON responses + const json = (await clone.json()) as { + model?: string; + usage?: { + input_tokens?: number; + output_tokens?: number; + cache_creation_input_tokens?: number; + cache_read_input_tokens?: number; + }; + }; + + if (!json.usage) return null; + + const inputTokens = json.usage.input_tokens || 0; + const cacheCreationInputTokens = + json.usage.cache_creation_input_tokens || 0; + const cacheReadInputTokens = json.usage.cache_read_input_tokens || 0; + const outputTokens = json.usage.output_tokens || 0; + const promptTokens = + inputTokens + cacheCreationInputTokens + cacheReadInputTokens; + const completionTokens = outputTokens; + const totalTokens = promptTokens + completionTokens; + + // Calculate cost if we have a model (z.ai returns GLM model names directly) + const model = json.model; + let costUsd: number | undefined; + if (model) { + log.debug( + `Calculating cost for z.ai model (non-streaming): ${model}`, + ); + try { + costUsd = await estimateCostUSD(model, { + inputTokens, + outputTokens, + cacheReadInputTokens, + cacheCreationInputTokens, + }); + log.debug(`Cost calculated for ${model}: $${costUsd}`); + } catch (error) { + log.warn(`Failed to calculate cost for model ${model}:`, error); + } + } else { + log.debug( + `No model found in z.ai non-streaming response for cost calculation`, + ); + } + + return { + model, + promptTokens, + completionTokens, + totalTokens, + inputTokens, + cacheReadInputTokens, + cacheCreationInputTokens, + outputTokens, + costUsd, + }; + } + } catch { + return null; + } + } + + /** + * Check if a response is a streaming response + */ + isStreamingResponse(response: Response): boolean { + const contentType = response.headers.get("content-type"); + return contentType?.includes("text/event-stream") || false; + } + + /** + * z.ai doesn't support OAuth - uses API keys instead + */ + supportsOAuth(): boolean { + return false; + } +} diff --git a/packages/proxy/src/handlers/account-selector.ts b/packages/proxy/src/handlers/account-selector.ts index 678d8120..00cdc4c1 100644 --- a/packages/proxy/src/handlers/account-selector.ts +++ b/packages/proxy/src/handlers/account-selector.ts @@ -12,12 +12,8 @@ export function getOrderedAccounts( ctx: ProxyContext, ): Account[] { const allAccounts = ctx.dbOps.getAllAccounts(); - // Filter accounts by provider - const providerAccounts = allAccounts.filter( - (account) => - account.provider === ctx.provider.name || account.provider === null, - ); - return ctx.strategy.select(providerAccounts, meta); + // Return all accounts - the provider will be determined dynamically per account + return ctx.strategy.select(allAccounts, meta); } /** diff --git a/packages/proxy/src/handlers/agent-interceptor.ts b/packages/proxy/src/handlers/agent-interceptor.ts index 4456ad98..99828790 100644 --- a/packages/proxy/src/handlers/agent-interceptor.ts +++ b/packages/proxy/src/handlers/agent-interceptor.ts @@ -328,21 +328,20 @@ function extractAgentDirectories(systemPrompt: string): string[] { // Regex #1: Look for explicit /.claude/agents paths const agentPathRegex = /([\\/][\w\-. ]*?\/.claude\/agents)(?=[\s"'\]])/g; - let match: RegExpExecArray | null; - match = agentPathRegex.exec(systemPrompt); - while (match !== null) { + // Use matchAll to avoid infinite loop issues with exec() + const agentPathMatches = systemPrompt.matchAll(agentPathRegex); + for (const match of agentPathMatches) { const dir = resolve(match[1]); directories.add(dir); - match = agentPathRegex.exec(systemPrompt); } // Regex #2: Look for repo root pattern "Contents of (.*?)/CLAUDE.md" const repoRootRegex = /Contents of ([^\n]+?)\/CLAUDE\.md/g; let matchCount = 0; - match = repoRootRegex.exec(systemPrompt); - while (match !== null) { + const repoRootMatches = systemPrompt.matchAll(repoRootRegex); + for (const match of repoRootMatches) { matchCount++; const repoRoot = match[1]; extractDirLog.info( @@ -357,7 +356,6 @@ function extractAgentDirectories(systemPrompt: string): string[] { extractDirLog.info(`Resolved agents dir: "${resolvedDir}"`); directories.add(resolvedDir); - match = repoRootRegex.exec(systemPrompt); } if (matchCount === 0 && systemPrompt.includes("CLAUDE.md")) { diff --git a/packages/proxy/src/handlers/proxy-operations.ts b/packages/proxy/src/handlers/proxy-operations.ts index f1b69440..672fde51 100644 --- a/packages/proxy/src/handlers/proxy-operations.ts +++ b/packages/proxy/src/handlers/proxy-operations.ts @@ -1,5 +1,6 @@ import { logError, ProviderError } from "@ccflare/core"; import { Logger } from "@ccflare/logger"; +import { getProvider } from "@ccflare/providers"; import type { Account, RequestMeta } from "@ccflare/types"; import { forwardToClient } from "../response-handler"; import { ERROR_MESSAGES, type ProxyContext } from "./proxy-types"; @@ -98,18 +99,23 @@ export async function proxyWithAccount( ctx: ProxyContext, ): Promise { try { - log.info(`Attempting request with account: ${account.name}`); + log.info( + `Attempting request with account: ${account.name} (provider: ${account.provider})`, + ); + + // Get the provider for this account + const provider = getProvider(account.provider) || ctx.provider; // Get valid access token const accessToken = await getValidAccessToken(account, ctx); - // Prepare request - const headers = ctx.provider.prepareHeaders( + // Prepare request using account-specific provider + const headers = provider.prepareHeaders( req.headers, accessToken, account.api_key || undefined, ); - const targetUrl = ctx.provider.buildUrl(url.pathname, url.search); + const targetUrl = provider.buildUrl(url.pathname, url.search); // Make the request const response = await makeProxyRequest( @@ -120,8 +126,16 @@ export async function proxyWithAccount( !!req.body, ); - // Process response and check for rate limit - const isRateLimited = processProxyResponse(response, account, ctx); + // Process response and check for rate limit using account-specific provider + const isRateLimited = processProxyResponse( + response, + account, + { + ...ctx, + provider, + }, + requestMeta.id, + ); if (isRateLimited) { return null; // Signal to try next account } @@ -141,7 +155,7 @@ export async function proxyWithAccount( failoverAttempts, agentUsed: requestMeta.agentUsed, }, - ctx, + { ...ctx, provider }, ); } catch (err) { handleProxyError(err, account, log); diff --git a/packages/proxy/src/handlers/response-processor.ts b/packages/proxy/src/handlers/response-processor.ts index e134f4b5..5093b764 100644 --- a/packages/proxy/src/handlers/response-processor.ts +++ b/packages/proxy/src/handlers/response-processor.ts @@ -43,11 +43,13 @@ export function handleRateLimitResponse( * @param account - The account to update * @param response - The response to extract metadata from * @param ctx - The proxy context + * @param requestId - The request ID for usage tracking */ export function updateAccountMetadata( account: Account, response: Response, ctx: ProxyContext, + requestId?: string, ): void { // Update basic usage ctx.asyncWriter.enqueue(() => ctx.dbOps.updateAccountUsage(account.id)); @@ -82,6 +84,23 @@ export function updateAccountMetadata( } })(); } + + // Extract usage info if supported + if (ctx.provider.extractUsageInfo && requestId) { + const extractUsageInfo = ctx.provider.extractUsageInfo.bind(ctx.provider); + (async () => { + const usageInfo = await extractUsageInfo(response.clone() as Response); + if (usageInfo) { + log.debug( + `Extracted usage info for account ${account.name}: ${JSON.stringify(usageInfo)}`, + ); + // Store usage info in database + ctx.asyncWriter.enqueue(() => + ctx.dbOps.updateRequestUsage(requestId, usageInfo), + ); + } + })(); + } } /** @@ -89,12 +108,14 @@ export function updateAccountMetadata( * @param response - The provider response * @param account - The account used * @param ctx - The proxy context + * @param requestId - The request ID for usage tracking * @returns Whether the response is rate-limited */ export function processProxyResponse( response: Response, account: Account, ctx: ProxyContext, + requestId?: string, ): boolean { const isStream = ctx.provider.isStreamingResponse?.(response) ?? false; const rateLimitInfo = ctx.provider.parseRateLimit(response); @@ -103,12 +124,12 @@ export function processProxyResponse( if (!isStream && rateLimitInfo.isRateLimited && rateLimitInfo.resetTime) { handleRateLimitResponse(account, rateLimitInfo, ctx); // Also update metadata for rate-limited responses - updateAccountMetadata(account, response, ctx); + updateAccountMetadata(account, response, ctx, requestId); return true; // Signal rate limit } // Update account metadata in background - updateAccountMetadata(account, response, ctx); + updateAccountMetadata(account, response, ctx, requestId); return false; } diff --git a/packages/proxy/src/handlers/token-manager.ts b/packages/proxy/src/handlers/token-manager.ts index 22d6448d..57f48aa8 100644 --- a/packages/proxy/src/handlers/token-manager.ts +++ b/packages/proxy/src/handlers/token-manager.ts @@ -1,6 +1,6 @@ import { ServiceUnavailableError, TokenRefreshError } from "@ccflare/core"; import { Logger } from "@ccflare/logger"; -import type { TokenRefreshResult } from "@ccflare/providers"; +import { getProvider, type TokenRefreshResult } from "@ccflare/providers"; import type { Account } from "@ccflare/types"; import { TOKEN_REFRESH_BACKOFF_MS, TOKEN_SAFETY_WINDOW_MS } from "../constants"; import { ERROR_MESSAGES, type ProxyContext } from "./proxy-types"; @@ -33,8 +33,11 @@ export async function refreshAccessTokenSafe( // Check if a refresh is already in progress for this account if (!ctx.refreshInFlight.has(account.id)) { + // Get the provider for this account + const provider = getProvider(account.provider) || ctx.provider; + // Create a new refresh promise and store it - const refreshPromise = ctx.provider + const refreshPromise = provider .refreshToken(account, ctx.runtime.clientId) .then((result: TokenRefreshResult) => { // 1. Persist to database asynchronously diff --git a/packages/proxy/src/post-processor.worker.ts b/packages/proxy/src/post-processor.worker.ts index 5cc16d1e..099c1da0 100644 --- a/packages/proxy/src/post-processor.worker.ts +++ b/packages/proxy/src/post-processor.worker.ts @@ -172,15 +172,36 @@ function extractUsageFromData(data: string, state: RequestState): void { state.firstTokenTimestamp = Date.now(); } - // Handle message_delta - provider's authoritative output token count AND end time + // Handle message_delta - provider's authoritative token counts AND end time if (parsed.type === "message_delta") { state.lastTokenTimestamp = Date.now(); - - if (parsed.usage?.output_tokens !== undefined) { - state.providerFinalOutputTokens = parsed.usage.output_tokens; - state.usage.outputTokens = parsed.usage.output_tokens; + log.info(`ZAI message_delta event received with usage:`, parsed.usage); + + if (parsed.usage) { + // Update all token counts from message_delta (authoritative for zai) + if (parsed.usage.output_tokens !== undefined) { + state.providerFinalOutputTokens = parsed.usage.output_tokens; + state.usage.outputTokens = parsed.usage.output_tokens; + log.info( + `ZAI set providerFinalOutputTokens to: ${parsed.usage.output_tokens}`, + ); + } + if (parsed.usage.input_tokens !== undefined) { + state.usage.inputTokens = parsed.usage.input_tokens; + log.info(`ZAI set inputTokens to: ${parsed.usage.input_tokens}`); + } + if (parsed.usage.cache_read_input_tokens !== undefined) { + state.usage.cacheReadInputTokens = + parsed.usage.cache_read_input_tokens; + log.info( + `ZAI set cacheReadInputTokens to: ${parsed.usage.cache_read_input_tokens}`, + ); + } return; // No further processing needed + } else { + log.info(`ZAI message_delta event has no usage information`); } + // Even if no usage info, we still set the timestamp for duration calculation } // Count tokens locally as fallback (but provider's count takes precedence) @@ -370,6 +391,16 @@ async function handleEnd(msg: EndMessage): Promise { (state.usage.cacheReadInputTokens || 0) + (state.usage.cacheCreationInputTokens || 0); + // Debug: Log the values being used for calculation + log.info( + `Token calculation debug - finalOutputTokens: ${finalOutputTokens}, providerFinalOutputTokens: ${state.providerFinalOutputTokens}, usage.outputTokens: ${state.usage.outputTokens}, outputTokensComputed: ${state.usage.outputTokensComputed}, totalTokens: ${state.usage.totalTokens}`, + ); + + // Log timestamp info + log.info( + `Timestamp debug - firstTokenTimestamp: ${state.firstTokenTimestamp}, lastTokenTimestamp: ${state.lastTokenTimestamp}, responseTime: ${responseTime}`, + ); + state.usage.costUsd = await estimateCostUSD(state.usage.model, { inputTokens: state.usage.inputTokens, outputTokens: finalOutputTokens, @@ -377,19 +408,56 @@ async function handleEnd(msg: EndMessage): Promise { cacheCreationInputTokens: state.usage.cacheCreationInputTokens, }); - // Calculate tokens per second using actual streaming duration - if ( - state.firstTokenTimestamp && - state.lastTokenTimestamp && - finalOutputTokens > 0 - ) { - const durationSec = - (state.lastTokenTimestamp - state.firstTokenTimestamp) / 1000; - if (durationSec > 0) { - state.usage.tokensPerSecond = finalOutputTokens / durationSec; - } else if (finalOutputTokens > 0) { - // If tokens were generated instantly, use a very small duration + // Calculate tokens per second - zai specific vs other providers + if (finalOutputTokens > 0) { + const totalDurationSec = responseTime / 1000; + + if (totalDurationSec > 0) { + // Check if this is a zai model (glm-*) + const isZaiModel = state.usage.model?.startsWith("glm-"); + + if (isZaiModel) { + // For zai models, use total response time (more intuitive for users) + state.usage.tokensPerSecond = finalOutputTokens / totalDurationSec; + log.info( + `ZAI token/s calculation: ${finalOutputTokens} tokens / ${totalDurationSec}s = ${state.usage.tokensPerSecond} tok/s (using total response time: ${responseTime}ms)`, + ); + } else { + // For other providers (like Anthropic), use streaming duration if available + if (state.firstTokenTimestamp && state.lastTokenTimestamp) { + const streamingDurationMs = + state.lastTokenTimestamp - state.firstTokenTimestamp; + const streamingDurationSec = streamingDurationMs / 1000; + + if (streamingDurationMs > 0) { + // Use streaming duration for generation speed + state.usage.tokensPerSecond = + finalOutputTokens / streamingDurationSec; + log.info( + `Token/s calculation (streaming): ${finalOutputTokens} tokens / ${streamingDurationSec}s = ${state.usage.tokensPerSecond} tok/s (streaming duration: ${streamingDurationMs}ms)`, + ); + } else { + // Fallback to total response time + state.usage.tokensPerSecond = + finalOutputTokens / totalDurationSec; + log.info( + `Token/s calculation (fallback): ${finalOutputTokens} tokens / ${totalDurationSec}s = ${state.usage.tokensPerSecond} tok/s (total response time: ${responseTime}ms)`, + ); + } + } else { + // No streaming timestamps available, use total response time + state.usage.tokensPerSecond = finalOutputTokens / totalDurationSec; + log.info( + `Token/s calculation (no timestamps): ${finalOutputTokens} tokens / ${totalDurationSec}s = ${state.usage.tokensPerSecond} tok/s (total response time: ${responseTime}ms)`, + ); + } + } + } else { + // If response time is 0, use a very small duration state.usage.tokensPerSecond = finalOutputTokens / 0.001; + log.info( + `Token/s calculation (instant): ${finalOutputTokens} tokens / 0.001s = ${state.usage.tokensPerSecond} tok/s`, + ); } } } diff --git a/packages/proxy/src/response-handler.ts b/packages/proxy/src/response-handler.ts index e7528c4d..625f7c27 100644 --- a/packages/proxy/src/response-handler.ts +++ b/packages/proxy/src/response-handler.ts @@ -111,14 +111,55 @@ export async function forwardToClient( const analyticsClone = response.clone(); (async () => { + const STREAM_TIMEOUT_MS = 300000; // 5 minutes max stream duration + const CHUNK_TIMEOUT_MS = 30000; // 30 seconds between chunks + try { const reader = analyticsClone.body?.getReader(); if (!reader) return; // Safety check + + const startTime = Date.now(); + let lastChunkTime = Date.now(); + // eslint-disable-next-line no-constant-condition while (true) { - const { value, done } = await reader.read(); + // Check for overall stream timeout + if (Date.now() - startTime > STREAM_TIMEOUT_MS) { + await reader.cancel(); + throw new Error( + `Stream timeout: exceeded ${STREAM_TIMEOUT_MS}ms total duration`, + ); + } + + // Check for chunk timeout (no data received) + if (Date.now() - lastChunkTime > CHUNK_TIMEOUT_MS) { + await reader.cancel(); + throw new Error( + `Stream timeout: no data received for ${CHUNK_TIMEOUT_MS}ms`, + ); + } + + // Read with a timeout wrapper + const readPromise = reader.read(); + const timeoutPromise = new Promise<{ + value?: Uint8Array; + done: boolean; + }>((_, reject) => + setTimeout( + () => reject(new Error("Read operation timeout")), + CHUNK_TIMEOUT_MS, + ), + ); + + const { value, done } = await Promise.race([ + readPromise, + timeoutPromise, + ]); + if (done) break; + if (value) { + lastChunkTime = Date.now(); const chunkMsg: ChunkMessage = { type: "chunk", requestId, diff --git a/packages/tui-core/src/accounts.ts b/packages/tui-core/src/accounts.ts index 0be9ddbc..16518ee2 100644 --- a/packages/tui-core/src/accounts.ts +++ b/packages/tui-core/src/accounts.ts @@ -17,6 +17,14 @@ export async function beginAddAccount( options: AddAccountOptions, ): Promise { const { name, mode = "max" } = options; + + // z.ai accounts don't use OAuth flow + if (mode === "zai") { + throw new Error( + "z.ai accounts should be added directly with API key, not via OAuth flow", + ); + } + const config = new Config(); const dbOps = DatabaseFactory.getInstance(); @@ -24,7 +32,10 @@ export async function beginAddAccount( const oauthFlow = await createOAuthFlow(dbOps, config); // Begin OAuth flow - const flowResult = await oauthFlow.begin({ name, mode }); + const flowResult = await oauthFlow.begin({ + name, + mode: mode as "max" | "console", + }); // Open browser console.log(`\nOpening browser to authenticate...`); @@ -44,7 +55,14 @@ export async function beginAddAccount( export async function completeAddAccount( options: AddAccountOptions & { code: string; flowData: OAuthFlowResult }, ): Promise { - const { name, mode = "max", tier = 1, code, flowData } = options; + const { + name, + mode = "max", + tier = 1, + priority = 0, + code, + flowData, + } = options; const config = new Config(); const dbOps = DatabaseFactory.getInstance(); @@ -54,13 +72,14 @@ export async function completeAddAccount( // Complete OAuth flow console.log("\nExchanging code for tokens..."); const _account = await oauthFlow.complete( - { sessionId: flowData.sessionId, code, tier, name }, + { sessionId: flowData.sessionId, code, tier, name, priority }, flowData, ); console.log(`\nAccount '${name}' added successfully!`); console.log(`Type: ${mode === "max" ? "Claude Max" : "Claude Console"}`); console.log(`Tier: ${tier}x`); + console.log(`Priority: ${priority}`); } /** @@ -73,6 +92,7 @@ export async function addAccount(options: AddAccountOptions): Promise { name: options.name, mode: options.mode || "max", tier: options.tier || 1, + priority: options.priority || 0, }); } @@ -99,3 +119,11 @@ export async function resumeAccount( const dbOps = DatabaseFactory.getInstance(); return cliCommands.resumeAccount(dbOps, name); } + +export async function updateAccountPriority( + name: string, + priority: number, +): Promise<{ success: boolean; message: string }> { + const dbOps = DatabaseFactory.getInstance(); + return cliCommands.setAccountPriority(dbOps, name, priority); +} diff --git a/packages/tui-core/src/args.ts b/packages/tui-core/src/args.ts index 6c682fda..c2cf831a 100644 --- a/packages/tui-core/src/args.ts +++ b/packages/tui-core/src/args.ts @@ -7,12 +7,14 @@ export interface ParsedArgs { logs?: boolean | number; stats?: boolean; addAccount?: string; - mode?: "max" | "console"; + mode?: "max" | "console" | "zai"; tier?: 1 | 5 | 20; + priority?: number; list?: boolean; remove?: string; pause?: string; resume?: string; + setPriority?: [string, string]; analyze?: boolean; resetStats?: boolean; clearHistory?: boolean; @@ -22,7 +24,7 @@ export interface ParsedArgs { export function parseArgs(args: string[]): ParsedArgs { try { - const { values } = nodeParseArgs({ + const { values, positionals } = nodeParseArgs({ args, options: { help: { type: "boolean", short: "h" }, @@ -33,10 +35,12 @@ export function parseArgs(args: string[]): ParsedArgs { "add-account": { type: "string" }, mode: { type: "string" }, tier: { type: "string" }, + priority: { type: "string" }, list: { type: "boolean" }, remove: { type: "string" }, pause: { type: "string" }, resume: { type: "string" }, + "set-priority": { type: "boolean" }, analyze: { type: "boolean" }, "reset-stats": { type: "boolean" }, "clear-history": { type: "boolean" }, @@ -56,12 +60,16 @@ export function parseArgs(args: string[]): ParsedArgs { } if (values.stats) result.stats = true; if (values["add-account"]) result.addAccount = values["add-account"]; - if (values.mode) result.mode = values.mode as "max" | "console"; + if (values.mode) result.mode = values.mode as "max" | "console" | "zai"; if (values.tier) result.tier = parseInt(values.tier, 10) as 1 | 5 | 20; + if (values.priority) result.priority = parseInt(values.priority, 10); if (values.list) result.list = true; if (values.remove) result.remove = values.remove; if (values.pause) result.pause = values.pause; if (values.resume) result.resume = values.resume; + if (values["set-priority"] && positionals.length >= 2) { + result.setPriority = [positionals[0], positionals[1]]; + } if (values.analyze) result.analyze = true; if (values["reset-stats"]) result.resetStats = true; if (values["clear-history"]) result.clearHistory = true; diff --git a/packages/tui-core/src/requests.ts b/packages/tui-core/src/requests.ts index 9af1c13b..29952925 100644 --- a/packages/tui-core/src/requests.ts +++ b/packages/tui-core/src/requests.ts @@ -1,10 +1,14 @@ -import { DatabaseFactory } from "@ccflare/database"; +import { DatabaseFactory, withDatabaseRetrySync } from "@ccflare/database"; import type { RequestPayload } from "@ccflare/types"; export type { RequestPayload }; export interface RequestSummary { id: string; + timestamp: number; + status: number | null; + accountUsed: string | null; + accountName: string | null; model?: string; inputTokens?: number; outputTokens?: number; @@ -15,6 +19,119 @@ export interface RequestSummary { responseTimeMs?: number; } +/** + * Get request summaries without loading full payloads + * This is the optimized version that eliminates JSON parsing bottleneck + */ +export async function getRequestSummaries( + limit = 100, +): Promise> { + const dbOps = DatabaseFactory.getInstance(); + const db = dbOps.getDatabase(); + + const summaries = withDatabaseRetrySync( + () => { + return db + .query( + ` + SELECT + r.id, + r.timestamp, + r.status_code as status, + r.account_used as accountUsed, + a.name as accountName, + r.model, + r.input_tokens as inputTokens, + r.output_tokens as outputTokens, + r.total_tokens as totalTokens, + r.cache_read_input_tokens as cacheReadInputTokens, + r.cache_creation_input_tokens as cacheCreationInputTokens, + r.cost_usd as costUsd, + r.response_time_ms as responseTimeMs + FROM requests r + LEFT JOIN accounts a ON r.account_used = a.id + ORDER BY r.timestamp DESC + LIMIT ? + `, + ) + .all(limit) as Array<{ + id: string; + timestamp: number; + status: number | null; + accountUsed: string | null; + accountName: string | null; + model?: string; + inputTokens?: number; + outputTokens?: number; + totalTokens?: number; + cacheReadInputTokens?: number; + cacheCreationInputTokens?: number; + costUsd?: number; + responseTimeMs?: number; + }>; + }, + dbOps.getRetryConfig(), + "getRequestSummaries", + ); + + const summaryMap = new Map(); + for (const summary of summaries) { + summaryMap.set(summary.id, { + id: summary.id, + timestamp: summary.timestamp, + status: summary.status, + accountUsed: summary.accountUsed, + accountName: summary.accountName, + model: summary.model || undefined, + inputTokens: summary.inputTokens || undefined, + outputTokens: summary.outputTokens || undefined, + totalTokens: summary.totalTokens || undefined, + cacheReadInputTokens: summary.cacheReadInputTokens || undefined, + cacheCreationInputTokens: summary.cacheCreationInputTokens || undefined, + costUsd: summary.costUsd || undefined, + responseTimeMs: summary.responseTimeMs || undefined, + }); + } + + return summaryMap; +} + +/** + * Get a single request payload by ID (lazy loading) + */ +export async function getRequestPayload( + id: string, +): Promise { + const dbOps = DatabaseFactory.getInstance(); + + return withDatabaseRetrySync( + () => { + const payload = dbOps.getRequestPayload(id); + if (!payload) { + return null; + } + + try { + return payload as RequestPayload; + } catch { + return { + id, + error: "Failed to parse payload", + request: { headers: {}, body: null }, + response: null, + meta: { timestamp: Date.now() }, + } as RequestPayload; + } + }, + dbOps.getRetryConfig(), + "getRequestPayload", + ); +} + +/** + * Legacy function for backward compatibility - now uses optimized approach + * @deprecated Use getRequestSummaries() and getRequestPayload() instead + */ export async function getRequests(limit = 100): Promise { const dbOps = DatabaseFactory.getInstance(); const rows = dbOps.listRequestPayloads(limit); @@ -43,55 +160,3 @@ export async function getRequests(limit = 100): Promise { return parsed; } - -export async function getRequestSummaries( - limit = 100, -): Promise> { - const dbOps = DatabaseFactory.getInstance(); - const db = dbOps.getDatabase(); - - const summaries = db - .query(` - SELECT - id, - model, - input_tokens as inputTokens, - output_tokens as outputTokens, - total_tokens as totalTokens, - cache_read_input_tokens as cacheReadInputTokens, - cache_creation_input_tokens as cacheCreationInputTokens, - cost_usd as costUsd, - response_time_ms as responseTimeMs - FROM requests - ORDER BY timestamp DESC - LIMIT ? - `) - .all(limit) as Array<{ - id: string; - model?: string; - inputTokens?: number; - outputTokens?: number; - totalTokens?: number; - cacheReadInputTokens?: number; - cacheCreationInputTokens?: number; - costUsd?: number; - responseTimeMs?: number; - }>; - - const summaryMap = new Map(); - summaries.forEach((summary) => { - summaryMap.set(summary.id, { - id: summary.id, - model: summary.model || undefined, - inputTokens: summary.inputTokens || undefined, - outputTokens: summary.outputTokens || undefined, - totalTokens: summary.totalTokens || undefined, - cacheReadInputTokens: summary.cacheReadInputTokens || undefined, - cacheCreationInputTokens: summary.cacheCreationInputTokens || undefined, - costUsd: summary.costUsd || undefined, - responseTimeMs: summary.responseTimeMs || undefined, - }); - }); - - return summaryMap; -} diff --git a/packages/types/src/account.ts b/packages/types/src/account.ts index 1f2dbc29..3a54e6fb 100644 --- a/packages/types/src/account.ts +++ b/packages/types/src/account.ts @@ -22,6 +22,7 @@ export interface AccountRow { rate_limit_reset?: number | null; rate_limit_status?: string | null; rate_limit_remaining?: number | null; + priority?: number; } // Domain model - used throughout the application @@ -45,6 +46,7 @@ export interface Account { rate_limit_reset: number | null; rate_limit_status: string | null; rate_limit_remaining: number | null; + priority: number; } // API response type - what clients receive @@ -64,6 +66,7 @@ export interface AccountResponse { rateLimitReset: string | null; rateLimitRemaining: number | null; sessionInfo: string; + priority: number; } // UI display type - used in TUI and web dashboard @@ -86,6 +89,7 @@ export interface AccountDisplay { session_start?: number | null; session_request_count?: number; access_token?: string | null; + priority: number; } // CLI list item type @@ -103,14 +107,16 @@ export interface AccountListItem { tokenStatus: "valid" | "expired"; rateLimitStatus: string; sessionInfo: string; - mode: "max" | "console"; + mode: "max" | "console" | "zai"; + priority: number; } // Account creation types export interface AddAccountOptions { name: string; - mode?: "max" | "console"; + mode?: "max" | "console" | "zai"; tier?: 1 | 5 | 20; + priority?: number; } export interface AccountDeleteRequest { @@ -139,6 +145,7 @@ export function toAccount(row: AccountRow): Account { rate_limit_reset: row.rate_limit_reset || null, rate_limit_status: row.rate_limit_status || null, rate_limit_remaining: row.rate_limit_remaining || null, + priority: row.priority || 0, }; } @@ -177,6 +184,7 @@ export function toAccountResponse(account: Account): AccountResponse { : null, rateLimitRemaining: account.rate_limit_remaining, sessionInfo, + priority: account.priority, }; } @@ -212,5 +220,6 @@ export function toAccountDisplay(account: Account): AccountDisplay { session_start: account.session_start, session_request_count: account.session_request_count, access_token: account.access_token, + priority: account.priority, }; } diff --git a/packages/types/src/agent.ts b/packages/types/src/agent.ts index 4cfa135b..5442a43d 100644 --- a/packages/types/src/agent.ts +++ b/packages/types/src/agent.ts @@ -42,6 +42,7 @@ export const ALLOWED_MODELS = [ CLAUDE_MODEL_IDS.OPUS_4, CLAUDE_MODEL_IDS.OPUS_4_1, CLAUDE_MODEL_IDS.SONNET_4, + CLAUDE_MODEL_IDS.SONNET_4_5, ] as const; export type AllowedModel = (typeof ALLOWED_MODELS)[number]; diff --git a/packages/ui-common/src/parsers/parse-conversation.ts b/packages/ui-common/src/parsers/parse-conversation.ts index 2ff17885..0c809599 100644 --- a/packages/ui-common/src/parsers/parse-conversation.ts +++ b/packages/ui-common/src/parsers/parse-conversation.ts @@ -45,7 +45,7 @@ export function parseRequestMessages(body: string | null): MessageData[] { for (const item of msg.content) { if (item.type === "text") { - // Filter out system reminders + // Filter out system reminders let text = normalizeText(item.text || ""); if (text.includes("")) { text = text @@ -72,11 +72,18 @@ export function parseRequestMessages(body: string | null): MessageData[] { name: item.name, input: item.input, }); - } else if (item.type === "tool_result") { + } else if (item.type === "tool_result") { const resultContent = Array.isArray((item as any).content) - ? ((item as any).content as Array<{ type: string; text?: string }>) - .map((c) => normalizeText(typeof c.text === "string" ? c.text : "")) - .join("") + ? ( + (item as any).content as Array<{ + type: string; + text?: string; + }> + ) + .map((c) => + normalizeText(typeof c.text === "string" ? c.text : ""), + ) + .join("") : typeof (item as any).content === "string" ? normalizeText((item as any).content as string) : ""; @@ -89,7 +96,7 @@ export function parseRequestMessages(body: string | null): MessageData[] { tool_use_id: (item as any).tool_use_id, content: resultContent, }); - } else if (item.type === "thinking") { + } else if (item.type === "thinking") { const thinking = normalizeText((item as any).thinking || ""); if (thinking) { message.contentBlocks?.push({ @@ -221,54 +228,61 @@ export function parseAssistantMessage(body: string | null): MessageData | null { if (!isStreaming) { try { const parsed = JSON.parse(body); - if (parsed.content) { - if (typeof parsed.content === "string") { - currentContent = normalizeText(parsed.content); - } else if (Array.isArray(parsed.content)) { - for (const item of parsed.content) { - if (item.type === "text" && item.text) { - const norm = normalizeText(item.text); - currentContent += norm; - message.contentBlocks?.push({ - type: ContentBlockType.Text, - text: norm, - }); - } else if (item.type === "tool_use") { - message.tools?.push({ - id: item.id, - name: item.name || "unknown", - input: item.input, - }); - message.contentBlocks?.push({ - type: ContentBlockType.ToolUse, - ...item, - }); - } else if (item.type === "thinking") { - const thinking = normalizeText((item as any).thinking || ""); - if (thinking) { - message.contentBlocks?.push({ - type: ContentBlockType.Thinking, - thinking, - }); - } - } else if (item.type === "tool_result") { - const resultContent = Array.isArray((item as any).content) - ? ((item as any).content as Array<{ type: string; text?: string }>) - .map((c) => normalizeText(typeof c.text === "string" ? c.text : "")) - .join("") - : typeof (item as any).content === "string" - ? normalizeText((item as any).content as string) - : ""; - message.toolResults?.push({ - tool_use_id: (item as any).tool_use_id || "", - content: resultContent, - }); + if (parsed.content) { + if (typeof parsed.content === "string") { + currentContent = normalizeText(parsed.content); + } else if (Array.isArray(parsed.content)) { + for (const item of parsed.content) { + if (item.type === "text" && item.text) { + const norm = normalizeText(item.text); + currentContent += norm; + message.contentBlocks?.push({ + type: ContentBlockType.Text, + text: norm, + }); + } else if (item.type === "tool_use") { + message.tools?.push({ + id: item.id, + name: item.name || "unknown", + input: item.input, + }); + message.contentBlocks?.push({ + type: ContentBlockType.ToolUse, + ...item, + }); + } else if (item.type === "thinking") { + const thinking = normalizeText((item as any).thinking || ""); + if (thinking) { message.contentBlocks?.push({ - type: ContentBlockType.ToolResult, - tool_use_id: (item as any).tool_use_id, - content: resultContent, + type: ContentBlockType.Thinking, + thinking, }); } + } else if (item.type === "tool_result") { + const resultContent = Array.isArray((item as any).content) + ? ( + (item as any).content as Array<{ + type: string; + text?: string; + }> + ) + .map((c) => + normalizeText(typeof c.text === "string" ? c.text : ""), + ) + .join("") + : typeof (item as any).content === "string" + ? normalizeText((item as any).content as string) + : ""; + message.toolResults?.push({ + tool_use_id: (item as any).tool_use_id || "", + content: resultContent, + }); + message.contentBlocks?.push({ + type: ContentBlockType.ToolResult, + tool_use_id: (item as any).tool_use_id, + content: resultContent, + }); + } } } } diff --git a/packages/ui-common/src/utils/normalize-text.ts b/packages/ui-common/src/utils/normalize-text.ts index b3a4c21a..d5e000fc 100644 --- a/packages/ui-common/src/utils/normalize-text.ts +++ b/packages/ui-common/src/utils/normalize-text.ts @@ -3,47 +3,49 @@ * and optionally strip a single pair of wrapping quotes. */ export function normalizeText(input: unknown): string { - let s = typeof input === "string" ? input : ""; - if (!s) return ""; + let s = typeof input === "string" ? input : ""; + if (!s) return ""; - // 1) If it's a JSON-encoded string (e.g., "...\n..."), try to parse directly - if (s.length >= 2 && s.startsWith('"') && s.endsWith('"')) { - try { - s = JSON.parse(s); - } catch { - // If JSON.parse fails, fall back to manual unquoting - s = s.slice(1, -1); - } - } else if (/\\[nrt"\\]/.test(s)) { - // 2) If it contains escaped sequences, decode them via JSON.parse wrapper - try { - const literal = - '"' + - s - .replace(/\\/g, "\\\\") - .replace(/\n/g, "\\n") - .replace(/\r/g, "\\r") - .replace(/\t/g, "\\t") - .replace(/"/g, "\\\"") + - '"'; - s = JSON.parse(literal); - } catch { - // Ignore if decoding fails - } - } + // 1) If it's a JSON-encoded string (e.g., "...\n..."), try to parse directly + if (s.length >= 2 && s.startsWith('"') && s.endsWith('"')) { + try { + s = JSON.parse(s); + } catch { + // If JSON.parse fails, fall back to manual unquoting + s = s.slice(1, -1); + } + } else if (/\\[nrt"\\]/.test(s)) { + // 2) If it contains escaped sequences, decode them via JSON.parse wrapper + try { + const literal = + '"' + + s + .replace(/\\/g, "\\\\") + .replace(/\n/g, "\\n") + .replace(/\r/g, "\\r") + .replace(/\t/g, "\\t") + .replace(/"/g, '\\"') + + '"'; + s = JSON.parse(literal); + } catch { + // Ignore if decoding fails + } + } - // 3) Heuristic: repair mojibake (UTF-8 mis-decoded as Latin-1) - if (/[รƒร‚รข]/.test(s)) { - try { - const bytes = new Uint8Array(Array.from(s, (ch) => ch.charCodeAt(0) & 0xff)); - const recoded = new TextDecoder("utf-8", { fatal: false }).decode(bytes); - if (recoded && recoded !== s) { - s = recoded; - } - } catch { - // Ignore decoding errors - } - } + // 3) Heuristic: repair mojibake (UTF-8 mis-decoded as Latin-1) + if (/[รƒร‚รข]/.test(s)) { + try { + const bytes = new Uint8Array( + Array.from(s, (ch) => ch.charCodeAt(0) & 0xff), + ); + const recoded = new TextDecoder("utf-8", { fatal: false }).decode(bytes); + if (recoded && recoded !== s) { + s = recoded; + } + } catch { + // Ignore decoding errors + } + } - return s; + return s; } diff --git a/packages/ui-constants/src/index.ts b/packages/ui-constants/src/index.ts index e0531d1b..7c267e27 100644 --- a/packages/ui-constants/src/index.ts +++ b/packages/ui-constants/src/index.ts @@ -7,6 +7,8 @@ export const COLORS = { blue: "#3b82f6", purple: "#8b5cf6", pink: "#ec4899", + indigo: "#6366f1", + cyan: "#06b6d4", } as const; // Chart color sequence for multi-series charts