Zen Budget
ultra-budgetSet and monitor conversation budgets to control costs, tokens, and duration for AI sessions. Manage spending limits and track usage across multiple AI providers.
Instructions
Set and monitor conversation budgets for cost and token control
Input Schema
TableJSON Schema
| Name | Required | Description | Default |
|---|---|---|---|
| action | Yes | Action to perform | |
| sessionId | Yes | Session ID to manage budget for | |
| maxTokens | No | Maximum tokens allowed for the session | |
| maxCostUsd | No | Maximum cost in USD allowed for the session | |
| maxDurationMs | No | Maximum duration in milliseconds allowed for the session |
Implementation Reference
- src/server.ts:994-1000 (registration)Registration of the ultra-budget tool using server.registerTool, specifying title, description, ZenBudgetSchema for input validation, and delegating to handleBudget function.
server.registerTool("ultra-budget", { title: "Zen Budget", description: "Set and monitor conversation budgets for cost and token control", inputSchema: ZenBudgetSchema.shape, }, async (args) => { return await handleBudget(args) as any; }); - src/server.ts:170-176 (schema)ZenBudgetSchema: Zod schema defining the input structure and validation for the ultra-budget tool, including actions (set/get/check), sessionId, and optional budget limits.
const ZenBudgetSchema = z.object({ action: z.enum(["set", "get", "check"]).describe("Action to perform"), sessionId: z.string().describe("Session ID to manage budget for"), maxTokens: z.number().optional().describe("Maximum tokens allowed for the session"), maxCostUsd: z.number().optional().describe("Maximum cost in USD allowed for the session"), maxDurationMs: z.number().optional().describe("Maximum duration in milliseconds allowed for the session"), }); - src/handlers/ultra-tools.ts:359-411 (handler)Core handler function for ultra-budget tool. Processes 'set', 'get', and 'check' actions using conversationMemory to manage session budgets for tokens, cost (USD), and duration (ms). Returns formatted text responses with status and limits.
export async function handleBudget(args: any) { const { action, sessionId, maxTokens, maxCostUsd, maxDurationMs } = args; switch (action) { case 'set': { const budget = await conversationMemory.setBudget(sessionId, maxTokens, maxCostUsd, maxDurationMs); return { content: [ { type: 'text', text: `## Budget Set\n\nSession: ${sessionId}\nMax Tokens: ${maxTokens || 'None'}\nMax Cost: $${maxCostUsd || 'None'}\nMax Duration: ${maxDurationMs || 'None'}ms` } ] }; } case 'get': { const context = await conversationMemory.getConversationContext(sessionId); const budget = context.budget; return { content: [ { type: 'text', text: `## Budget Status\n\nSession: ${sessionId}\n\n**Limits:**\n- Max Tokens: ${budget?.maxTokens || 'None'}\n- Max Cost: $${budget?.maxCostUsd || 'None'}\n- Max Duration: ${budget?.maxDurationMs || 'None'}ms\n\n**Usage:**\n- Used Tokens: ${budget?.usedTokens || 0}\n- Used Cost: $${budget?.usedCostUsd || 0}\n- Used Duration: ${budget?.usedDurationMs || 0}ms` } ] }; } case 'check': { const limits = await conversationMemory.checkBudgetLimits(sessionId); const status = limits.withinLimits ? '✅ Within Limits' : '⚠️ Limits Exceeded'; const details = [ `Token Limit: ${limits.tokenLimitExceeded ? '❌ Exceeded' : '✅ OK'}`, `Cost Limit: ${limits.costLimitExceeded ? '❌ Exceeded' : '✅ OK'}`, `Duration Limit: ${limits.durationLimitExceeded ? '❌ Exceeded' : '✅ OK'}` ].join('\n'); return { content: [ { type: 'text', text: `## Budget Check\n\nSession: ${sessionId}\nStatus: ${status}\n\n${details}` } ] }; } default: throw new Error(`Unknown action: ${action}`); } }