@robota-sdk/agent-command 3.0.0-beta.67 → 3.0.0-beta.69

package/src/provider/provider-command-profile-operations.ts

@@ -1,5 +1,7 @@
 import {
   buildProviderSetupPatch,
+  formatOrgPolicyViolationMessage,
+  isApiKeyPlaintext,
   setCurrentProvider,
   upsertProviderProfile,
 } from '@robota-sdk/agent-framework';
@@ -19,12 +21,6 @@ import type {
   IProviderSetupInput,
 } from '@robota-sdk/agent-framework';
 
-const YES = 'yes';
-const PROVIDER_RESTART_EFFECT = {
-  type: 'session-restart-requested',
-  reason: 'other',
-} as const;
-
 export function formatProviderChoiceLabel(
   name: string,
   profile: IProviderProfileSettings,
@@ -40,53 +36,37 @@ export function buildProviderSwitch(
   options: IProviderCommandModuleOptions,
 ): ICommandResult {
   if (!profileName) {
-    return { message: 'Usage: provider use <profile>', success: false };
+    return { message: 'Usage: /provider switch <profile>', success: false };
   }
   if (!providers?.[profileName]) {
     return { message: `Provider profile "${profileName}" was not found.`, success: false };
   }
+  const { orgPolicy } = options;
+  if (orgPolicy?.allowedProviders && !orgPolicy.allowedProviders.includes(profileName)) {
+    return {
+      message: formatOrgPolicyViolationMessage(
+        `Provider "${profileName}" is not allowed by your organization policy. Allowed: ${orgPolicy.allowedProviders.join(', ')}.`,
+        orgPolicy.adminContact,
+      ),
+      success: false,
+    };
+  }
   if (options.settings.readMergedSettings().currentProvider === profileName) {
-    return { message: `Provider profile "${profileName}" is already current.`, success: true };
+    return { message: `Already using provider "${profileName}".`, success: true };
   }
+  const profile = providers[profileName];
+  const target = options.settings.readTargetSettings();
+  const merged = options.settings.readMergedSettings();
+  const next =
+    target.providers?.[profileName] !== undefined || merged.providers?.[profileName] !== undefined
+      ? { ...target, currentProvider: profileName }
+      : setCurrentProvider(target, profileName);
+  options.settings.writeTargetSettings(next);
+  const modelLabel = profile.model ?? 'unknown model';
   return {
-    message: `Provider change requested: ${profileName}`,
+    message: `Switched to ${profileName} (${modelLabel}). History preserved.`,
     success: true,
-    interaction: createProviderSwitchInteraction(profileName, options),
-  };
-}
-
-function createProviderSwitchInteraction(
-  profileName: string,
-  options: IProviderCommandModuleOptions,
-): ICommandInteraction {
-  return {
-    prompt: {
-      kind: 'choice',
-      title: `Change provider to ${profileName}? This will restart the session.`,
-      options: [
-        { value: YES, label: 'Yes' },
-        { value: 'no', label: 'No' },
-      ],
-    },
-    submit: (value) => {
-      if (value !== YES) {
-        return { message: 'Provider change cancelled.', success: true };
-      }
-      const merged = options.settings.readMergedSettings();
-      const target = options.settings.readTargetSettings();
-      const next =
-        target.providers?.[profileName] !== undefined ||
-        merged.providers?.[profileName] !== undefined
-          ? { ...target, currentProvider: profileName }
-          : setCurrentProvider(target, profileName);
-      options.settings.writeTargetSettings(next);
-      return {
-        message: `Provider changed to ${profileName}. Restarting...`,
-        success: true,
-        effects: [{ ...PROVIDER_RESTART_EFFECT, message: 'Provider change restart' }],
-      };
-    },
-    cancel: () => ({ message: 'Provider change cancelled.', success: true }),
+    effects: [{ type: 'provider-hot-swap-requested', profileName }],
   };
 }
 
@@ -176,6 +156,16 @@ function completeProviderEdit(
   if (!currentProfile) {
     return { message: `Provider profile "${profileName}" was not found.`, success: false };
   }
+  const { orgPolicy } = options;
+  if (orgPolicy?.requireApiKeyFromEnv && isApiKeyPlaintext(input.apiKey)) {
+    return {
+      message: formatOrgPolicyViolationMessage(
+        'Your organization policy requires API keys to be stored as environment variable references ($ENV:VAR_NAME), not as plaintext.',
+        orgPolicy.adminContact,
+      ),
+      success: false,
+    };
+  }
   const target = options.settings.readTargetSettings();
   const patch = buildProviderSetupPatch(input, {
     providerDefinitions: options.providerDefinitions,
@@ -190,11 +180,9 @@ function completeProviderEdit(
   const isCurrent = merged.currentProvider === profileName;
   return {
     message: isCurrent
-      ? `Provider ${profileName} updated. Restarting...`
+      ? `Provider ${profileName} updated. Switching...`
       : `Provider ${profileName} updated.`,
     success: true,
-    ...(isCurrent
-      ? { effects: [{ ...PROVIDER_RESTART_EFFECT, message: 'Provider edit restart' }] }
-      : {}),
+    ...(isCurrent ? { effects: [{ type: 'provider-hot-swap-requested', profileName }] } : {}),
   };
 }

package/src/provider/provider-onboarding.ts

@@ -0,0 +1,64 @@
+import type { TPromptInput } from './provider-setup-flow.js';
+import type { ITerminalOutput } from '@robota-sdk/agent-core';
+
+export type TOnboardingPath = 'has-key' | 'free-key' | 'local';
+
+export interface IOnboardingResult {
+  path: TOnboardingPath;
+  preselectedType?: string;
+}
+
+const ONBOARDING_PROMPT = `
+  Do you have an API key for an AI provider?
+
+    1. Yes, I have an API key
+    2. No — get a free key (Google Gemini, takes ~2 min)
+    3. No — use a local model (LM Studio, no API key needed)
+
+  Choose [1-3] (default: 1): `;
+
+const GEMINI_GUIDE = `
+  ── Get a free Gemini API key ────────────────────────────────────────────────
+
+  1. Open  https://aistudio.google.com/apikey  in your browser
+  2. Sign in with your Google account
+  3. Click "Create API key"  (takes ~30 seconds)
+  4. Copy the key starting with "AIza..."
+
+  ─────────────────────────────────────────────────────────────────────────────
+`;
+
+const LOCAL_MODEL_GUIDE = `
+  ── Set up a local model with LM Studio ──────────────────────────────────────
+
+  1. Download LM Studio from  https://lmstudio.ai
+  2. Open LM Studio → search for a model (e.g. "llama" or "phi") → Download
+  3. Go to  Developer  tab → click  Start Server
+     (default address: http://localhost:1234)
+
+  When the server is running, come back here and press Enter to continue.
+
+  ─────────────────────────────────────────────────────────────────────────────
+`;
+
+export async function runOnboardingBranch(
+  promptInput: TPromptInput,
+  terminal: ITerminalOutput,
+): Promise<IOnboardingResult> {
+  const raw = await promptInput(ONBOARDING_PROMPT);
+  const choice = raw.trim() || '1';
+
+  if (choice === '2') {
+    terminal.writeLine(GEMINI_GUIDE);
+    await promptInput('  Press Enter when you have your API key: ');
+    return { path: 'free-key', preselectedType: 'gemini' };
+  }
+
+  if (choice === '3') {
+    terminal.writeLine(LOCAL_MODEL_GUIDE);
+    await promptInput('  Press Enter when LM Studio server is running: ');
+    return { path: 'local', preselectedType: 'gemma' };
+  }
+
+  return { path: 'has-key' };
+}

package/src/provider/provider-setup-flow.ts

@@ -177,12 +177,21 @@ export function formatProviderSetupPromptLabel(
   return `${prefix}  ${step.title}${suffix}: `;
 }
 
+const CATEGORY_BADGES: Record<string, string> = {
+  'cloud-paid': '[Cloud/Paid]',
+  'cloud-free': '[Cloud/Free]',
+  'local-free': '[Local/Free]',
+};
+
 export function formatProviderSetupChoiceLabel(definition: IProviderDefinition): string {
-  const label =
+  const badge =
+    definition.category !== undefined ? `${CATEGORY_BADGES[definition.category] ?? ''} ` : '';
+  const name =
     definition.displayName !== undefined
       ? `${definition.displayName} (${definition.type})`
       : definition.type;
-  return definition.description !== undefined ? `${label} - ${definition.description}` : label;
+  const label = `${badge}${name}`;
+  return definition.description !== undefined ? `${label} — ${definition.description}` : label;
 }
 
 export function formatProviderSetupHelpLinks(

package/src/provider/provider-startup.ts

@@ -10,6 +10,7 @@ import {
   applyProviderConfiguration,
 } from '@robota-sdk/agent-framework';
 
+import { runOnboardingBranch } from './provider-onboarding.js';
 import {
   formatProviderSetupSelectionPrompt,
   resolveProviderSetupSelection,
@@ -38,11 +39,22 @@ export async function runProviderStartupSetup(
   terminal: ITerminalOutput,
   providerDefinitions: readonly IProviderDefinition[],
 ): Promise<void> {
-  const providerChoice = await promptInput(formatProviderSetupSelectionPrompt(providerDefinitions));
-  const type = resolveProviderSetupSelection(providerChoice, providerDefinitions);
+  const onboarding = await runOnboardingBranch(promptInput, terminal);
+  const existingProfileNames = Object.keys(readMergedProviderSettings(cwd).providers ?? {});
   const settingsPath = resolveSettingsPathForScope(cwd, ctx.settingsScope);
+
+  let type: string;
+  if (onboarding.preselectedType !== undefined) {
+    type = onboarding.preselectedType;
+  } else {
+    const providerChoice = await promptInput(
+      formatProviderSetupSelectionPrompt(providerDefinitions),
+    );
+    type = resolveProviderSetupSelection(providerChoice, providerDefinitions);
+  }
+
   const input = await runProviderSetupPromptFlow(type, promptInput, providerDefinitions, {
-    existingProfileNames: Object.keys(readMergedProviderSettings(cwd).providers ?? {}),
+    existingProfileNames,
   });
   applyProviderConfiguration(settingsPath, input, { providerDefinitions });
   const language = await promptInput('  Response language (ko/en/ja/zh, default: en): ');

package/src/session/__tests__/model-pricing.test.ts

@@ -0,0 +1,64 @@
+import { describe, expect, it } from 'vitest';
+import { calculateCost, formatTokens, formatUsd } from '../model-pricing.js';
+
+describe('calculateCost', () => {
+  it('returns exact cost for known Anthropic Sonnet model', () => {
+    const cost = calculateCost('claude-sonnet-4-5', 1_000_000, 1_000_000);
+    expect(cost).toBeCloseTo(3 + 15, 5);
+  });
+
+  it('returns exact cost for known Anthropic Opus model', () => {
+    const cost = calculateCost('claude-opus-4-7', 1_000_000, 1_000_000);
+    expect(cost).toBeCloseTo(15 + 75, 5);
+  });
+
+  it('calculates fractional token costs correctly', () => {
+    const cost = calculateCost('claude-sonnet-4-5', 45_000, 12_000);
+    expect(cost).toBeDefined();
+    expect(cost!).toBeCloseTo((45_000 / 1_000_000) * 3 + (12_000 / 1_000_000) * 15, 8);
+  });
+
+  it('falls back to pattern matching for claude-sonnet variant', () => {
+    const cost = calculateCost('claude-sonnet-99-99', 1_000_000, 0);
+    expect(cost).toBeCloseTo(3, 5);
+  });
+
+  it('falls back to pattern matching for deepseek variant', () => {
+    const cost = calculateCost('deepseek-v3-turbo', 1_000_000, 0);
+    expect(cost).toBeDefined();
+  });
+
+  it('returns undefined for completely unknown model', () => {
+    const cost = calculateCost('unknown-model-xyz-123', 1_000, 1_000);
+    expect(cost).toBeUndefined();
+  });
+
+  it('returns zero cost when both token counts are zero', () => {
+    const cost = calculateCost('claude-sonnet-4-5', 0, 0);
+    expect(cost).toBe(0);
+  });
+});
+
+describe('formatUsd', () => {
+  it('formats small amounts with 4 decimal places', () => {
+    expect(formatUsd(0.0043)).toBe('$0.0043');
+  });
+
+  it('formats sub-dollar amounts with 3 decimal places', () => {
+    expect(formatUsd(0.187)).toBe('$0.187');
+  });
+
+  it('formats dollar+ amounts with 2 decimal places', () => {
+    expect(formatUsd(1.24)).toBe('$1.24');
+  });
+});
+
+describe('formatTokens', () => {
+  it('formats token counts with comma separators', () => {
+    expect(formatTokens(45000)).toBe('45,000');
+  });
+
+  it('formats small counts without separator', () => {
+    expect(formatTokens(999)).toBe('999');
+  });
+});

package/src/session/__tests__/session-command-module.test.ts

@@ -139,7 +139,7 @@ describe('createSessionCommandModule', () => {
     expect(entry).toEqual(
       expect.objectContaining({
         name: 'cost',
-        description: 'Show session info',
+        description: expect.stringContaining('token usage'),
         source: 'session',
         modelInvocable: false,
       }),
@@ -288,18 +288,62 @@ describe('createSessionCommandModule', () => {
     });
   });
 
-  it('shows session info through the session command API', async () => {
+  it('shows session info through the session command API (no token data yet)', async () => {
     const executor = new SystemCommandExecutor([
       ...(createSessionCommandModule().systemCommands ?? []),
     ]);
 
     const result = await executor.execute('cost', createCommandContext(), '');
 
-    expect(result).toEqual({
-      success: true,
-      message: 'Session: session_1\nMessages: 5',
-      data: { sessionId: 'session_1', messageCount: 5 },
-    });
+    expect(result?.success).toBe(true);
+    expect(result?.message).toContain('session_1');
+    expect(result?.message).toContain('Messages: 5');
+    expect(result?.message).toContain('not yet available');
+  });
+
+  it('shows token counts and estimated cost when session has usage data', async () => {
+    const runtime = {
+      ...createRuntime(),
+      getSessionTokenUsage: () => ({ inputTokens: 45_000, outputTokens: 12_000 }),
+      getModelId: () => 'claude-sonnet-4-5',
+    };
+    const context = { ...createCommandContext(), getSession: () => runtime };
+    const executor = new SystemCommandExecutor([
+      ...(createSessionCommandModule().systemCommands ?? []),
+    ]);
+
+    const result = await executor.execute('cost', context, '');
+
+    expect(result?.success).toBe(true);
+    expect(result?.message).toContain('45,000');
+    expect(result?.message).toContain('12,000');
+    expect(result?.message).toContain('$');
+    expect((result?.data as Record<string, unknown>)?.inputTokens).toBe(45_000);
+    expect((result?.data as Record<string, unknown>)?.outputTokens).toBe(12_000);
+    expect((result?.data as Record<string, unknown>)?.estimatedCostUsd).toBeDefined();
+  });
+
+  it('sets a monthly budget via /cost budget subcommand', async () => {
+    const executor = new SystemCommandExecutor([
+      ...(createSessionCommandModule().systemCommands ?? []),
+    ]);
+    const context = { ...createCommandContext(), getCwd: () => '/tmp/robota-test-budget' };
+
+    const result = await executor.execute('cost', context, 'budget 5.00');
+
+    expect(result?.success).toBe(true);
+    expect(result?.message).toContain('$5.00');
+  });
+
+  it('rejects invalid budget amount', async () => {
+    const executor = new SystemCommandExecutor([
+      ...(createSessionCommandModule().systemCommands ?? []),
+    ]);
+
+    const result = await executor.execute('cost', createCommandContext(), 'budget notanumber');
+
+    expect(result?.success).toBe(false);
+    expect(result?.message).toContain('Usage:');
   });
 
   it('validates the current session replay log through the SDK common API', async () => {

package/src/session/model-pricing.ts

@@ -0,0 +1,79 @@
+/** USD prices per 1,000,000 tokens (as of May 2026 — update when providers change rates). */
+interface IModelPrice {
+  inputPerMillion: number;
+  outputPerMillion: number;
+}
+
+const MODEL_PRICES: Record<string, IModelPrice> = {
+  // Anthropic Claude 4
+  'claude-opus-4-7': { inputPerMillion: 15, outputPerMillion: 75 },
+  'claude-opus-4-5': { inputPerMillion: 15, outputPerMillion: 75 },
+  'claude-sonnet-4-6': { inputPerMillion: 3, outputPerMillion: 15 },
+  'claude-sonnet-4-5': { inputPerMillion: 3, outputPerMillion: 15 },
+  'claude-haiku-4-5': { inputPerMillion: 0.8, outputPerMillion: 4 },
+  // Anthropic Claude 3
+  'claude-3-5-sonnet-20241022': { inputPerMillion: 3, outputPerMillion: 15 },
+  'claude-3-5-haiku-20241022': { inputPerMillion: 0.8, outputPerMillion: 4 },
+  'claude-3-opus-20240229': { inputPerMillion: 15, outputPerMillion: 75 },
+  // OpenAI
+  'gpt-4o': { inputPerMillion: 2.5, outputPerMillion: 10 },
+  'gpt-4o-mini': { inputPerMillion: 0.15, outputPerMillion: 0.6 },
+  o1: { inputPerMillion: 15, outputPerMillion: 60 },
+  'o1-mini': { inputPerMillion: 3, outputPerMillion: 12 },
+  o3: { inputPerMillion: 10, outputPerMillion: 40 },
+  'o3-mini': { inputPerMillion: 1.1, outputPerMillion: 4.4 },
+  // DeepSeek
+  'deepseek-chat': { inputPerMillion: 0.14, outputPerMillion: 0.28 },
+  'deepseek-reasoner': { inputPerMillion: 0.55, outputPerMillion: 2.19 },
+  // Google Gemini
+  'gemini-2.0-flash': { inputPerMillion: 0.1, outputPerMillion: 0.4 },
+  'gemini-2.0-flash-thinking': { inputPerMillion: 0.35, outputPerMillion: 3.5 },
+  'gemini-1.5-pro': { inputPerMillion: 1.25, outputPerMillion: 5 },
+  'gemini-1.5-flash': { inputPerMillion: 0.075, outputPerMillion: 0.3 },
+};
+
+const PATTERN_PRICES: Array<{ pattern: RegExp; price: IModelPrice }> = [
+  { pattern: /claude-opus/i, price: { inputPerMillion: 15, outputPerMillion: 75 } },
+  { pattern: /claude-sonnet/i, price: { inputPerMillion: 3, outputPerMillion: 15 } },
+  { pattern: /claude-haiku/i, price: { inputPerMillion: 0.8, outputPerMillion: 4 } },
+  { pattern: /gpt-4o-mini/i, price: { inputPerMillion: 0.15, outputPerMillion: 0.6 } },
+  { pattern: /gpt-4/i, price: { inputPerMillion: 2.5, outputPerMillion: 10 } },
+  { pattern: /deepseek/i, price: { inputPerMillion: 0.14, outputPerMillion: 0.28 } },
+  { pattern: /gemini-2/i, price: { inputPerMillion: 0.1, outputPerMillion: 0.4 } },
+  { pattern: /gemini-1/i, price: { inputPerMillion: 1.25, outputPerMillion: 5 } },
+];
+
+function lookupPrice(modelId: string): IModelPrice | undefined {
+  const exact = MODEL_PRICES[modelId];
+  if (exact) return exact;
+  for (const { pattern, price } of PATTERN_PRICES) {
+    if (pattern.test(modelId)) return price;
+  }
+  return undefined;
+}
+
+/** Returns USD cost, or undefined if the model is not in the pricing table. */
+export function calculateCost(
+  modelId: string,
+  inputTokens: number,
+  outputTokens: number,
+): number | undefined {
+  const price = lookupPrice(modelId);
+  if (!price) return undefined;
+  return (
+    (inputTokens / 1_000_000) * price.inputPerMillion +
+    (outputTokens / 1_000_000) * price.outputPerMillion
+  );
+}
+
+/** Format a USD amount for display (e.g. "$0.0043"). */
+export function formatUsd(amount: number): string {
+  if (amount < 0.01) return `$${amount.toFixed(4)}`;
+  if (amount < 1) return `$${amount.toFixed(3)}`;
+  return `$${amount.toFixed(2)}`;
+}
+
+/** Format a token count with comma separators (e.g. "45,000"). */
+export function formatTokens(count: number): string {
+  return count.toLocaleString('en-US');
+}

package/src/session/session-command.ts

@@ -1,3 +1,6 @@
+import { existsSync, mkdirSync, readFileSync, writeFileSync } from 'node:fs';
+import { join } from 'node:path';
+
 import {
   RENAME_COMMAND_USAGE,
   clearConversationHistory,
@@ -9,6 +12,8 @@ import {
   validateCommandSessionReplayLog,
 } from '@robota-sdk/agent-framework';
 
+import { calculateCost, formatTokens, formatUsd } from './model-pricing.js';
+
 import type { ICommandHostContext, ICommandResult } from '@robota-sdk/agent-framework';
 
 export const CLEAR_COMMAND_MESSAGE = 'Conversation cleared.';
@@ -45,16 +50,127 @@ export function executeResumeCommand(_context: ICommandHostContext, _args: strin
   };
 }
 
-export function executeCostCommand(context: ICommandHostContext, _args: string): ICommandResult {
+const BUDGET_FILE = '.robota/budget.json';
+
+interface IBudgetConfig {
+  monthly: number;
+}
+
+function readBudget(cwd: string): IBudgetConfig | undefined {
+  const file = join(cwd, BUDGET_FILE);
+  if (!existsSync(file)) return undefined;
+  let raw: string;
+  try {
+    raw = readFileSync(file, 'utf-8');
+  } catch {
+    // allow-fallback: budget file read failure disables feature gracefully
+    return undefined;
+  }
+  try {
+    return JSON.parse(raw) as IBudgetConfig;
+  } catch {
+    // allow-fallback: malformed budget JSON treated as no budget set
+    return undefined;
+  }
+}
+
+function writeBudget(cwd: string, config: IBudgetConfig): void {
+  const dir = join(cwd, '.robota');
+  mkdirSync(dir, { recursive: true });
+  writeFileSync(join(cwd, BUDGET_FILE), JSON.stringify(config, null, 2));
+}
+
+function clearBudget(cwd: string): void {
+  const file = join(cwd, BUDGET_FILE);
+  if (existsSync(file)) {
+    writeFileSync(file, '{}');
+  }
+}
+
+function buildCostOutput(context: ICommandHostContext): {
+  lines: string[];
+  data: Record<string, unknown>;
+} {
+  const session = context.getSession();
   const sessionInfo = readCommandSessionInfo(context);
-  return {
-    success: true,
-    message: `Session: ${sessionInfo.sessionId}\nMessages: ${sessionInfo.messageCount}`,
-    data: {
-      sessionId: sessionInfo.sessionId,
-      messageCount: sessionInfo.messageCount,
-    },
+  const tokenUsage = session.getSessionTokenUsage?.();
+  const modelId = session.getModelId?.();
+  const lines: string[] = [
+    `Session:  ${sessionInfo.sessionId}`,
+    `Messages: ${sessionInfo.messageCount}`,
+  ];
+  const data: Record<string, unknown> = {
+    sessionId: sessionInfo.sessionId,
+    messageCount: sessionInfo.messageCount,
   };
+
+  if (tokenUsage) {
+    lines.push(
+      `Tokens:   ${formatTokens(tokenUsage.inputTokens)} input  /  ${formatTokens(tokenUsage.outputTokens)} output`,
+    );
+    data.inputTokens = tokenUsage.inputTokens;
+    data.outputTokens = tokenUsage.outputTokens;
+
+    if (modelId) {
+      const cost = calculateCost(modelId, tokenUsage.inputTokens, tokenUsage.outputTokens);
+      if (cost !== undefined) {
+        lines.push(`Cost:     ${formatUsd(cost)}  (${modelId})`);
+        data.estimatedCostUsd = cost;
+
+        const budget = readBudget(context.getCwd());
+        if (budget?.monthly) {
+          const remaining = budget.monthly - cost;
+          const pct = Math.min(100, Math.round((cost / budget.monthly) * 100));
+          lines.push(
+            `Budget:   ${formatUsd(remaining)} remaining of ${formatUsd(budget.monthly)}/mo  (${pct}% used)`,
+          );
+          data.budgetMonthly = budget.monthly;
+          data.budgetRemainingUsd = remaining;
+        }
+      }
+    }
+  } else {
+    lines.push('Tokens:   not yet available (no turns completed)');
+  }
+
+  return { lines, data };
+}
+
+export function executeCostCommand(context: ICommandHostContext, args: string): ICommandResult {
+  const trimmed = args.trim();
+
+  if (trimmed.startsWith('budget')) {
+    const budgetArg = trimmed.slice('budget'.length).trim();
+
+    if (budgetArg === 'clear') {
+      clearBudget(context.getCwd());
+      return { success: true, message: 'Monthly budget cleared.' };
+    }
+
+    if (budgetArg === '') {
+      const current = readBudget(context.getCwd());
+      if (!current?.monthly) {
+        return {
+          success: true,
+          message: 'No budget set. Use: /cost budget <amount>',
+        };
+      }
+      return { success: true, message: `Monthly budget: ${formatUsd(current.monthly)}` };
+    }
+
+    const amount = parseFloat(budgetArg);
+    if (!Number.isFinite(amount) || amount <= 0) {
+      return {
+        success: false,
+        message: 'Usage: /cost budget <amount>  (e.g. /cost budget 5.00)',
+      };
+    }
+    writeBudget(context.getCwd(), { monthly: amount });
+    return { success: true, message: `Monthly budget set to ${formatUsd(amount)}.` };
+  }
+
+  const { lines, data } = buildCostOutput(context);
+  return { success: true, message: lines.join('\n'), data };
 }
 
 export function executeValidateSessionCommand(