agent-orchestrator/ui-electrobun/src/bun/handlers/agent-handlers.ts

/**
 * Agent + Session persistence RPC handlers.
 */

import type { SidecarManager } from "../sidecar-manager.ts";
import type { SessionDb } from "../session-db.ts";
import type { SearchDb } from "../search-db.ts";

export function createAgentHandlers(
  sidecarManager: SidecarManager,
  sessionDb: SessionDb,
  sendToWebview: { send: Record<string, (...args: unknown[]) => void> },
  searchDb?: SearchDb,
) {
  return {
    "agent.start": ({ sessionId, provider, prompt, cwd, model, systemPrompt, maxTurns, permissionMode, claudeConfigDir, extraEnv, additionalDirectories, worktreeName }: Record<string, unknown>) => {
      try {
        const result = sidecarManager.startSession(
          sessionId as string,
          provider as string,
          prompt as string,
          { cwd, model, systemPrompt, maxTurns, permissionMode, claudeConfigDir, extraEnv, additionalDirectories, worktreeName } as Record<string, unknown>,
        );

        if (result.ok) {
          // Partial 3: Track last-sent cost to emit live updates only on change
          let lastCostUsd = 0;
          let lastInputTokens = 0;
          let lastOutputTokens = 0;

          /** Emit agent.cost if cost data changed since last emit. */
          function emitCostIfChanged(sid: string): void {
            const sessions = sidecarManager.listSessions();
            const session = sessions.find((s: Record<string, unknown>) => s.sessionId === sid);
            if (!session) return;
            const cost = (session.costUsd as number) ?? 0;
            const inTok = (session.inputTokens as number) ?? 0;
            const outTok = (session.outputTokens as number) ?? 0;
            if (cost === lastCostUsd && inTok === lastInputTokens && outTok === lastOutputTokens) return;
            lastCostUsd = cost;
            lastInputTokens = inTok;
            lastOutputTokens = outTok;
            try {
              (sendToWebview.send as Record<string, Function>)["agent.cost"]({
                sessionId: sid, costUsd: cost, inputTokens: inTok, outputTokens: outTok,
              });
            } catch { /* ignore */ }
          }

          sidecarManager.onMessage(sessionId as string, (sid: string, messages: unknown) => {
            try {
              (sendToWebview.send as Record<string, Function>)["agent.message"]({ sessionId: sid, messages });
            } catch (err) {
              console.error("[agent.message] forward error:", err);
            }
            // Partial 3: Stream cost on every message batch
            emitCostIfChanged(sid);
          });

          sidecarManager.onStatus(sessionId as string, (sid: string, status: string, error?: string) => {
            try {
              (sendToWebview.send as Record<string, Function>)["agent.status"]({ sessionId: sid, status, error });
            } catch (err) {
              console.error("[agent.status] forward error:", err);
            }
            emitCostIfChanged(sid);
          });
        }

        return result;
      } catch (err) {
        const error = err instanceof Error ? err.message : String(err);
        console.error("[agent.start]", err);
        return { ok: false, error };
      }
    },

    "agent.stop": ({ sessionId }: { sessionId: string }) => {
      try {
        return sidecarManager.stopSession(sessionId);
      } catch (err) {
        const error = err instanceof Error ? err.message : String(err);
        console.error("[agent.stop]", err);
        return { ok: false, error };
      }
    },

    "agent.prompt": ({ sessionId, prompt }: { sessionId: string; prompt: string }) => {
      try {
        return sidecarManager.writePrompt(sessionId, prompt);
      } catch (err) {
        const error = err instanceof Error ? err.message : String(err);
        console.error("[agent.prompt]", err);
        return { ok: false, error };
      }
    },

    "agent.list": () => {
      try {
        return { sessions: sidecarManager.listSessions() };
      } catch (err) {
        console.error("[agent.list]", err);
        return { sessions: [] };
      }
    },

    // Session persistence
    "session.save": ({ projectId, sessionId, provider, status, costUsd, inputTokens, outputTokens, model, error, createdAt, updatedAt }: Record<string, unknown>) => {
      try {
        sessionDb.saveSession({ projectId, sessionId, provider, status, costUsd, inputTokens, outputTokens, model, error, createdAt, updatedAt } as Record<string, unknown>);
        return { ok: true };
      } catch (err) {
        console.error("[session.save]", err);
        return { ok: false };
      }
    },

    "session.load": ({ projectId }: { projectId: string }) => {
      try {
        return { session: sessionDb.loadSession(projectId) };
      } catch (err) {
        console.error("[session.load]", err);
        return { session: null };
      }
    },

    "session.list": ({ projectId }: { projectId: string }) => {
      try {
        return { sessions: sessionDb.listSessionsByProject(projectId) };
      } catch (err) {
        console.error("[session.list]", err);
        return { sessions: [] };
      }
    },

    "session.messages.save": ({ messages }: { messages: Array<Record<string, unknown>> }) => {
      try {
        sessionDb.saveMessages(messages.map((m) => ({
          sessionId: m.sessionId, msgId: m.msgId, role: m.role,
          content: m.content, toolName: m.toolName, toolInput: m.toolInput,
          timestamp: m.timestamp, seqId: (m.seqId as number) ?? 0,
          costUsd: (m.costUsd as number) ?? 0,
          inputTokens: (m.inputTokens as number) ?? 0, outputTokens: (m.outputTokens as number) ?? 0,
        })));
        // Partial 1: Index each new message in FTS5 search
        if (searchDb) {
          for (const m of messages) {
            try {
              searchDb.indexMessage(
                String(m.sessionId ?? ""),
                String(m.role ?? ""),
                String(m.content ?? ""),
              );
            } catch { /* non-critical — don't fail the save */ }
          }
        }
        return { ok: true };
      } catch (err) {
        console.error("[session.messages.save]", err);
        return { ok: false };
      }
    },

    "session.messages.load": ({ sessionId }: { sessionId: string }) => {
      try {
        return { messages: sessionDb.loadMessages(sessionId) };
      } catch (err) {
        console.error("[session.messages.load]", err);
        return { messages: [] };
      }
    },
  };
}