polygala-ai · DorianZheng · Apr 24, 2026
diff --git a/src/electron/main.ts b/src/electron/main.ts
@@ -116,6 +116,12 @@ void app.whenReady().then(async () => {
     app,
     ...(agentLiteHomeDir ? { agentLiteHomeDir } : {}),
     onAgentIdle: workflowCoordinator.onAgentIdle,
+    onBudgetExceeded: (payload) => {
+      broadcast(ipcChannels.budgetExceeded, payload);
+    },
+    onBudgetWarning: (payload) => {
+      broadcast(ipcChannels.budgetWarning, payload);
+    },
     onItemActivityChanged: (payload) => {
       broadcast(ipcChannels.itemActivityUpdated, payload);
     },

diff --git a/src/electron/main/ipc/register-main-ipc-handlers.ts b/src/electron/main/ipc/register-main-ipc-handlers.ts
@@ -142,6 +142,23 @@ export function registerMainIpcHandlers(options: RegisterMainIpcHandlersOptions)
     async (_event, agentId: string) =>
       withRuntime((runtimeController) => runtimeController.selectAgent(agentId)),
   );
+  ipcMain.handle(ipcChannels.getBudget, async (_event, agentId: string) => {
+    try {
+      return await withRuntime((runtimeController) => runtimeController.getBudget(agentId));
+    } catch {
+      return null;
+    }
+  });
+  ipcMain.handle(
+    ipcChannels.setBudget,
+    async (_event, agentId: string, config: Record<string, unknown>) =>
+      withRuntime((runtimeController) => runtimeController.setBudget(agentId, config)),
+  );
+  ipcMain.handle(
+    ipcChannels.resumeBudget,
+    async (_event, agentId: string) =>
+      withRuntime((runtimeController) => runtimeController.resumeBudget(agentId)),
+  );
   ipcMain.handle(
     ipcChannels.updateAgentChannel,
     async (_event, input) =>

diff --git a/src/electron/main/runtime/agent-runtime/index.ts b/src/electron/main/runtime/agent-runtime/index.ts
@@ -434,6 +434,8 @@ export interface AgentRuntimeOptions {
   loadAgentLiteModule?: () => Promise<typeof import('@boxlite-ai/agentlite')>;
   now?: () => number;
   onAgentIdle?: (agentId: string) => void;
+  onBudgetExceeded?: (payload: unknown) => void;
+  onBudgetWarning?: (payload: unknown) => void;
   onItemActivityChanged?: (payload: { itemId: string; isWorking: boolean }) => void;
   resolveProjectName?: (projectId: string) => Promise<string | null>;
   resolveProjectRootPath?: (projectId: string) => Promise<string | null>;
@@ -497,6 +499,10 @@ export class AgentRuntime implements AgentRuntimeContract {
 
   private readonly onAgentIdle: AgentRuntimeOptions['onAgentIdle'];
 
+  private readonly onBudgetExceeded: AgentRuntimeOptions['onBudgetExceeded'];
+
+  private readonly onBudgetWarning: AgentRuntimeOptions['onBudgetWarning'];
+
   private readonly onItemActivityChanged: AgentRuntimeOptions['onItemActivityChanged'];
 
   /** Per-item ephemeral run state driven by AgentLite task.run.* events. */
@@ -539,6 +545,8 @@ export class AgentRuntime implements AgentRuntimeContract {
     this.actionServices = options.actionServices;
     this.homeDir = options.homeDir ?? os.homedir();
     this.onAgentIdle = options.onAgentIdle;
+    this.onBudgetExceeded = options.onBudgetExceeded;
+    this.onBudgetWarning = options.onBudgetWarning;
     this.onItemActivityChanged = options.onItemActivityChanged;
     this.runtimeRoot = resolveAgentLiteRuntimeRoot(options.homeDir);
     this.bundledAgentDir = options.bundledAgentDir ?? resolveBundledAgentDir();
@@ -617,6 +625,7 @@ export class AgentRuntime implements AgentRuntimeContract {
         this.telegram.getSetupSession(sessionId),
       getSnapshot: () => this.getSnapshot(),
       listAgents: () => this.getSnapshot().agents,
+      callAction: async (name, input) => this.callAgentLiteAction(name, input),
       selectAgent: (agentId) => {
         this.selectAgent(agentId);
       },
@@ -2354,6 +2363,37 @@ export class AgentRuntime implements AgentRuntimeContract {
     }
   }
 
+  private async callAgentLiteAction(
+    name: string,
+    input: Record<string, unknown>,
+  ): Promise<unknown> {
+    const groupJid = typeof input.group_jid === 'string' ? input.group_jid : null;
+    const agentId = groupJid
+      ? toAgentPathId(groupJid)
+      : this.snapshot.selectedAgentId ?? this.snapshot.agents[0]?.id ?? null;
+
+    if (!agentId) {
+      return null;
+    }
+
+    const record = this.records.get(agentId);
+
+    if (!record) {
+      return null;
+    }
+
+    const duneAgent = await this.ensureAgentRuntime(record);
+    const alAgent = duneAgent.agentLiteAgent as AgentLiteAgent & {
+      callAction?: (actionName: string, actionInput: Record<string, unknown>) => Promise<unknown>;
+    };
+
+    if (typeof alAgent.callAction !== 'function') {
+      return null;
+    }
+
+    return alAgent.callAction(name, input);
+  }
+
   private async ensureAgentRuntime(record: PersistedAgentRecord): Promise<DuneAgent> {
     const agentId = record.agent.id;
     const existing = this.lifecycle.getRuntime(agentId);
@@ -2457,6 +2497,22 @@ export class AgentRuntime implements AgentRuntimeContract {
         });
       });
 
+      alAgent.on('budget.exceeded', (event: Record<string, unknown>) => {
+        this.onBudgetExceeded?.({
+          agentId,
+          jid: toAgentChatJid(agentId),
+          ...event,
+        });
+      });
+
+      alAgent.on('budget.warning', (event: Record<string, unknown>) => {
+        this.onBudgetWarning?.({
+          agentId,
+          jid: toAgentChatJid(agentId),
+          ...event,
+        });
+      });
+
       alAgent.on('run.status', (event) => {
         this.pushActivityEvent(agentId, {
           id: `act-${Date.now()}-${Math.random().toString(36).slice(2, 6)}`,

diff --git a/src/electron/main/runtime/desktop-runtime-controller.ts b/src/electron/main/runtime/desktop-runtime-controller.ts
@@ -5,6 +5,7 @@ import type {
   AgentServiceListener,
   AgentServiceSnapshot,
 } from '@/shared/agents/agent-runtime';
+import { toAgentChatJid } from '@/shared/agents/agent-id';
 import { createMockAgentRuntime } from '@/renderer/features/agents/services/mock-agent-service';
 import type {
   AgentDefinition,
@@ -164,6 +165,38 @@ export class DesktopRuntimeController {
     await this.activeRuntime.service.cancelItemAssignment(agentId, taskId);
   }
 
+  /** Gets budget config and state for an agent. */
+  async getBudget(agentId: string): Promise<unknown> {
+    try {
+      if (typeof this.activeRuntime.service.callAction !== 'function') return null;
+      return await this.activeRuntime.service.callAction('budget_get', {
+        group_jid: toAgentChatJid(agentId),
+      });
+    } catch {
+      return null;
+    }
+  }
+
+  /** Sets budget config for an agent. */
+  async setBudget(
+    agentId: string,
+    config: { daily_limit_usd?: number | null; total_limit_usd?: number | null; reset_hour?: number },
+  ): Promise<void> {
+    if (typeof this.activeRuntime.service.callAction !== 'function') return;
+    await this.activeRuntime.service.callAction('budget_set', {
+      group_jid: toAgentChatJid(agentId),
+      ...config,
+    });
+  }
+
+  /** Resumes a paused agent. */
+  async resumeBudget(agentId: string): Promise<void> {
+    if (typeof this.activeRuntime.service.callAction !== 'function') return;
+    await this.activeRuntime.service.callAction('budget_resume', {
+      group_jid: toAgentChatJid(agentId),
+    });
+  }
+
   /** Returns true when the task still exists and remains active in agentlite. */
   isItemTaskKnown(agentId: string, taskId: string): boolean {
     return this.activeRuntime.service.isItemTaskKnown(agentId, taskId);

diff --git a/src/electron/main/runtime/runtime-bootstrap.ts b/src/electron/main/runtime/runtime-bootstrap.ts
@@ -12,6 +12,8 @@ interface RuntimeBootstrapOptions {
   agentStore: AppStorage;
   app: Pick<App, 'getAppPath'>;
   onAgentIdle: (agentId: string) => void;
+  onBudgetExceeded: (payload: unknown) => void;
+  onBudgetWarning: (payload: unknown) => void;
   onItemActivityChanged: (payload: { isWorking: boolean; itemId: string }) => void;
   onRuntimeSnapshot: (snapshot: AgentServiceSnapshot) => void;
   onStarted?: () => void;
@@ -67,6 +69,8 @@ export function createRuntimeBootstrap(options: RuntimeBootstrapOptions) {
         bundledAgentDir: path.join(options.app.getAppPath(), 'agent'),
         ...(options.agentLiteHomeDir ? { homeDir: options.agentLiteHomeDir } : {}),
         onAgentIdle: options.onAgentIdle,
+        onBudgetExceeded: options.onBudgetExceeded,
+        onBudgetWarning: options.onBudgetWarning,
         onItemActivityChanged: options.onItemActivityChanged,
         resolveModelCredentials: () => resolveDefaultModelCredentials({
           secretsStore: options.secretsStore,

diff --git a/src/electron/preload.ts b/src/electron/preload.ts
@@ -30,6 +30,7 @@ const bridge: DesktopBridge = {
     ipcRenderer.invoke(ipcChannels.getProjectActivityPage, projectId, options),
   getAgentTranscriptPage: (agentId, options) =>
     ipcRenderer.invoke(ipcChannels.getAgentTranscriptPage, agentId, options),
+  getBudget: (agentId) => ipcRenderer.invoke(ipcChannels.getBudget, agentId),
   getRuntimeSnapshot: () => ipcRenderer.invoke(ipcChannels.getRuntimeSnapshot),
   getTelegramSetupSession: (sessionId) =>
     ipcRenderer.invoke(ipcChannels.getTelegramSetupSession, sessionId),
@@ -41,10 +42,12 @@ const bridge: DesktopBridge = {
     ipcRenderer.invoke(ipcChannels.prepareProjectRootPath, rootPath, artifactFolderNames),
   reloadExternalChannels: () => ipcRenderer.invoke(ipcChannels.reloadExternalChannels),
   resetRuntime: () => ipcRenderer.invoke(ipcChannels.resetRuntime),
+  resumeBudget: (agentId) => ipcRenderer.invoke(ipcChannels.resumeBudget, agentId),
   restartApp: () => ipcRenderer.invoke(ipcChannels.restartApp),
   runIsolatedResearch: (agentId, input) =>
     ipcRenderer.invoke(ipcChannels.runIsolatedResearch, agentId, input),
   selectAgent: (agentId) => ipcRenderer.invoke(ipcChannels.selectAgent, agentId),
+  setBudget: (agentId, config) => ipcRenderer.invoke(ipcChannels.setBudget, agentId, config),
   updateAgentChannel: (input) => ipcRenderer.invoke(ipcChannels.updateAgentChannel, input),
   updateAgentDefinition: (agentId, definition) =>
     ipcRenderer.invoke(ipcChannels.updateAgentDefinition, agentId, definition),
@@ -93,6 +96,34 @@ const bridge: DesktopBridge = {
       ipcRenderer.removeListener(ipcChannels.itemActivityUpdated, handler);
     };
   },
+  subscribeBudgetExceeded: (listener) => {
+    const handler = (
+      _event: Electron.IpcRendererEvent,
+      payload: Parameters<typeof listener>[0],
+    ) => {
+      listener(payload);
+    };
+
+    ipcRenderer.on(ipcChannels.budgetExceeded, handler);
+
+    return () => {
+      ipcRenderer.removeListener(ipcChannels.budgetExceeded, handler);
+    };
+  },
+  subscribeBudgetWarning: (listener) => {
+    const handler = (
+      _event: Electron.IpcRendererEvent,
+      payload: Parameters<typeof listener>[0],
+    ) => {
+      listener(payload);
+    };
+
+    ipcRenderer.on(ipcChannels.budgetWarning, handler);
+
+    return () => {
+      ipcRenderer.removeListener(ipcChannels.budgetWarning, handler);
+    };
+  },
 };
 
 contextBridge.exposeInMainWorld('duneDesktop', Object.freeze(bridge));
diff --git a/src/renderer/features/agents/components/AgentPanel.tsx b/src/renderer/features/agents/components/AgentPanel.tsx
@@ -1,16 +1,23 @@
 // Agent panel UI.
 
-import { type KeyboardEvent, type RefObject, useState } from 'react';
+import { type KeyboardEvent, type RefObject, useEffect, useState } from 'react';
 import { ArrowUpRight } from 'lucide-react';
 
 import { Wrench, Bot, Info } from 'lucide-react';
 
 import { AgentMessageContent } from '@/renderer/features/agents/components/AgentMessageContent';
+import { BudgetExceededBanner } from '@/renderer/features/agents/components/BudgetExceededBanner';
+import { BudgetWarningBadge } from '@/renderer/features/agents/components/BudgetWarningBadge';
 import { CodingEngineCard, groupEngineRuns } from '@/renderer/features/agents/components/CodingEngineCard';
 import type { AgentActivityEvent, PresentedAgent } from '@/renderer/features/agents/types';
 import { useDesktopPlatform } from '@/renderer/shared/lib/use-desktop-platform';
 import { cn } from '@/renderer/shared/lib/utils';
 import { Button } from '@/renderer/shared/ui/button';
+import type {
+  BudgetExceededPayload,
+  BudgetResult,
+  BudgetWarningPayload,
+} from '@/shared/electron/desktop-bridge';
 
 /** Agent panel props. */
 interface AgentPanelProps {
@@ -193,6 +200,37 @@ function formatUsageCost(costUsd: number) {
   return `$${costUsd.toFixed(costUsd >= 0.01 ? 4 : 6).replace(/0+$/, '').replace(/\.$/, '')}`;
 }
 
+function deriveBudgetExceededPayload(
+  agentId: string,
+  budget: BudgetResult,
+): BudgetExceededPayload | null {
+  if (!budget.state.paused) {
+    return null;
+  }
+
+  const dailyPct = budget.usage.daily_pct ?? -1;
+  const totalPct = budget.usage.total_pct ?? -1;
+  const pausedReason = budget.state.paused_reason?.toLowerCase() ?? '';
+  const limitType = pausedReason.includes('total') || totalPct > dailyPct ? 'total' : 'daily';
+  const limitUsd = limitType === 'daily'
+    ? budget.config.daily_limit_usd
+    : budget.config.total_limit_usd;
+  const usedUsd = limitType === 'daily'
+    ? budget.usage.daily_cost_usd
+    : budget.usage.total_cost_usd;
+
+  return {
+    agentId,
+    jid: `dune:agent:${agentId}`,
+    limitType,
+    limitUsd: limitUsd ?? usedUsd,
+    timestamp: budget.state.paused_at
+      ? new Date(budget.state.paused_at).toISOString()
+      : new Date().toISOString(),
+    usedUsd,
+  };
+}
+
 /** Renders the agent panel UI. */
 export function AgentPanel({
   agent,
@@ -206,6 +244,47 @@ export function AgentPanel({
 }: AgentPanelProps) {
   const { modifierLabel } = useDesktopPlatform();
   const composerHint = `${modifierLabel} Enter to send · Shift Enter for a new line`;
+  const [budgetExceeded, setBudgetExceeded] = useState<BudgetExceededPayload | null>(null);
+  const [budgetWarning, setBudgetWarning] = useState<BudgetWarningPayload | null>(null);
+
+  useEffect(() => {
+    let isMounted = true;
+
+    const loadBudget = async () => {
+      const budget = await window.duneDesktop?.getBudget?.(agent.id) ?? null;
+
+      if (!isMounted || !budget) {
+        return;
+      }
+
+      setBudgetExceeded(deriveBudgetExceededPayload(agent.id, budget));
+    };
+
+    void loadBudget();
+
+    return () => {
+      isMounted = false;
+    };
+  }, [agent.id]);
+
+  useEffect(() => {
+    const unsubExceeded = window.duneDesktop?.subscribeBudgetExceeded?.((payload) => {
+      if (payload.jid === `dune:agent:${agent.id}`) {
+        setBudgetExceeded(payload);
+        setBudgetWarning(null);
+      }
+    });
+    const unsubWarning = window.duneDesktop?.subscribeBudgetWarning?.((payload) => {
+      if (payload.jid === `dune:agent:${agent.id}`) {
+        setBudgetWarning(payload);
+      }
+    });
+
+    return () => {
+      unsubExceeded?.();
+      unsubWarning?.();
+    };
+  }, [agent.id]);
 
   /** Handles key down composer. */
   const handleComposerKeyDown = async (
@@ -248,6 +327,25 @@ export function AgentPanel({
             </h2>
           </div>
 
+          {budgetExceeded ? (
+            <BudgetExceededBanner
+              agentId={agent.id}
+              limitType={budgetExceeded.limitType}
+              limitUsd={budgetExceeded.limitUsd}
+              onResume={() => {
+                setBudgetExceeded(null);
+              }}
+              usedUsd={budgetExceeded.usedUsd}
+            />
+          ) : null}
+
+          {budgetWarning && !budgetExceeded ? (
+            <BudgetWarningBadge
+              limitType={budgetWarning.limitType}
+              pctUsed={budgetWarning.pctUsed}
+            />
+          ) : null}
+
           {buildTimeline(agent).map((item) => {
             if (item.type === 'message') {
               const { message } = item;