@swarmclawai/swarmclaw 0.8.4 → 0.8.7

package/src/lib/server/wake-dispatcher.ts

@@ -0,0 +1,318 @@
+/**
+ * Wake Dispatcher — central routing for wake requests by WakeMode.
+ *
+ * Instead of callers directly choosing between `requestHeartbeatNow()` (immediate)
+ * and the heartbeat-service timer (deferred), they call `dispatchWake()` with an
+ * explicit WakeMode. The dispatcher routes to the correct execution path.
+ *
+ * This replaces the pattern where scheduler.ts both creates a task AND calls
+ * requestHeartbeatNow() — the dispatcher handles that routing centrally.
+ */
+
+import { genId } from '@/lib/id'
+import { log } from './logger'
+import type { WakeModeRequest, JobContext } from './wake-mode'
+import {
+  computeWakePriority,
+  resolveRunAt,
+  wakeModeToSource,
+  createJobContext,
+} from './wake-mode'
+import type { WakeRequestInput } from './heartbeat-wake'
+import { requestHeartbeatNow } from './heartbeat-wake'
+import { enqueueSystemEvent } from './system-events'
+import { errorMessage, hmrSingleton } from '@/lib/shared-utils'
+
+// ── Deferred queue for `next_heartbeat` mode ────────────────────────────
+
+interface DeferredWake {
+  request: WakeModeRequest
+  priority: number
+  enqueuedAt: number
+}
+
+const state = hmrSingleton('__swarmclaw_wake_dispatcher__', () => ({
+  deferredQueue: new Map<string, DeferredWake[]>(),
+  scheduledTimers: new Map<string, ReturnType<typeof setTimeout>>(),
+  activeJobs: new Map<string, JobContext>(),
+}))
+
+function deferredKey(request: WakeModeRequest): string {
+  return `${request.agentId || ''}::${request.sessionId || ''}`
+}
+
+// ── Public API ──────────────────────────────────────────────────────────
+
+/**
+ * Dispatch a wake request through the appropriate execution path.
+ *
+ * - `immediate`: forwards to requestHeartbeatNow with coalesce
+ * - `next_heartbeat`: queues for the next heartbeat-service tick to drain
+ * - `scheduled`: sets a timer to fire requestHeartbeatNow at the target time
+ */
+export function dispatchWake(request: WakeModeRequest): {
+  mode: string
+  priority: number
+  runAt: number | null
+  jobId: string
+} {
+  const priority = computeWakePriority(request)
+  const runAt = resolveRunAt(request)
+  const jobId = genId(8)
+
+  switch (request.mode) {
+    case 'immediate':
+      dispatchImmediate(request, priority, jobId)
+      break
+    case 'next_heartbeat':
+      dispatchDeferred(request, priority, jobId)
+      break
+    case 'scheduled':
+      dispatchScheduled(request, priority, runAt, jobId)
+      break
+  }
+
+  return { mode: request.mode, priority, runAt, jobId }
+}
+
+function dispatchImmediate(request: WakeModeRequest, priority: number, jobId: string): void {
+  const wakeInput: WakeRequestInput = {
+    eventId: request.eventId || jobId,
+    agentId: request.agentId,
+    sessionId: request.sessionId,
+    reason: request.reason || 'immediate-wake',
+    source: request.source,
+    resumeMessage: request.resumeMessage,
+    detail: request.detail,
+    priority,
+  }
+  requestHeartbeatNow(wakeInput)
+  log.info('wake-dispatcher', `Dispatched immediate wake ${jobId}`, {
+    agentId: request.agentId,
+    sessionId: request.sessionId,
+    reason: request.reason,
+    priority,
+  })
+}
+
+function dispatchDeferred(request: WakeModeRequest, priority: number, jobId: string): void {
+  const key = deferredKey(request)
+  const queue = state.deferredQueue.get(key) || []
+
+  // Deduplicate by eventId or reason+source
+  const existingIndex = queue.findIndex((entry: any) => {
+    if (request.eventId && entry.request.eventId) {
+      return request.eventId === entry.request.eventId
+    }
+    return entry.request.reason === request.reason
+      && entry.request.source === request.source
+  })
+
+  const entry: DeferredWake = { request, priority, enqueuedAt: Date.now() }
+
+  if (existingIndex >= 0) {
+    // Replace with higher-priority version
+    if (priority >= queue[existingIndex].priority) {
+      queue[existingIndex] = entry
+    }
+  } else {
+    queue.push(entry)
+  }
+
+  // Keep sorted by priority (highest first)
+  queue.sort((a: any, b: any) => b.priority - a.priority)
+  state.deferredQueue.set(key, queue)
+
+  log.info('wake-dispatcher', `Deferred wake ${jobId} queued for next heartbeat`, {
+    agentId: request.agentId,
+    sessionId: request.sessionId,
+    reason: request.reason,
+    queueDepth: queue.length,
+  })
+}
+
+function dispatchScheduled(
+  request: WakeModeRequest,
+  priority: number,
+  runAt: number | null,
+  jobId: string,
+): void {
+  const now = Date.now()
+  const targetTime = runAt ?? now
+  const delayMs = Math.max(0, targetTime - now)
+
+  // If the target time is now or in the past, dispatch immediately
+  if (delayMs <= 0) {
+    dispatchImmediate(request, priority, jobId)
+    return
+  }
+
+  // Set a timer to fire at the target time
+  const timerId = setTimeout(() => {
+    state.scheduledTimers.delete(jobId)
+    try {
+      dispatchImmediate(request, priority, jobId)
+    } catch (err: unknown) {
+      log.error('wake-dispatcher', `Scheduled wake ${jobId} failed, retrying once`, {
+        error: errorMessage(err),
+        agentId: request.agentId,
+        sessionId: request.sessionId,
+      })
+      // Single retry after 5s — if this also fails, the wake is lost (logged above)
+      setTimeout(() => {
+        try { dispatchImmediate(request, priority, jobId) } catch { /* give up */ }
+      }, 5000)
+    }
+    log.info('wake-dispatcher', `Scheduled wake ${jobId} fired after ${delayMs}ms delay`, {
+      agentId: request.agentId,
+      sessionId: request.sessionId,
+      reason: request.reason,
+    })
+  }, delayMs)
+
+  state.scheduledTimers.set(jobId, timerId)
+
+  log.info('wake-dispatcher', `Scheduled wake ${jobId} for ${new Date(targetTime).toISOString()}`, {
+    agentId: request.agentId,
+    sessionId: request.sessionId,
+    reason: request.reason,
+    delayMs,
+  })
+
+  // Also enqueue a system event so the agent knows something is scheduled
+  if (request.sessionId && request.resumeMessage) {
+    enqueueSystemEvent(
+      request.sessionId,
+      `[Scheduled] ${request.resumeMessage} (fires at ${new Date(targetTime).toISOString()})`,
+    )
+  }
+}
+
+// ── Deferred queue drain (called by heartbeat-service on each tick) ─────
+
+/**
+ * Drain all deferred wakes for a given session/agent. Called by heartbeat-service
+ * during its periodic tick to pick up `next_heartbeat` mode requests.
+ *
+ * Returns the deferred events so the heartbeat-service can include them in
+ * the prompt context alongside normal heartbeat content.
+ */
+export function drainDeferredWakes(agentId?: string, sessionId?: string): WakeModeRequest[] {
+  const key = `${agentId || ''}::${sessionId || ''}`
+  const queue = state.deferredQueue.get(key)
+  if (!queue || queue.length === 0) return []
+
+  const drained = queue.map((entry: any) => entry.request)
+  state.deferredQueue.delete(key)
+
+  log.info('wake-dispatcher', `Drained ${drained.length} deferred wakes`, {
+    agentId,
+    sessionId,
+    reasons: drained.map((r: any) => r.reason),
+  })
+
+  return drained
+}
+
+/**
+ * Check if there are pending deferred wakes for a target.
+ */
+export function hasDeferredWakes(agentId?: string, sessionId?: string): boolean {
+  const key = `${agentId || ''}::${sessionId || ''}`
+  const queue = state.deferredQueue.get(key)
+  return !!queue && queue.length > 0
+}
+
+// ── Job context management ──────────────────────────────────────────────
+
+/**
+ * Create and register an isolated job context for a wake execution.
+ * The job context provides per-job scratchpad, abort signal, and
+ * heartbeat snapshot isolation.
+ */
+export function startJobExecution(params: {
+  sessionId: string
+  agentId?: string
+  mode: WakeModeRequest['mode']
+  signal: AbortSignal
+  source?: string
+  reason?: string
+  heartbeatSnapshot?: string
+}): JobContext {
+  const jobId = genId(8)
+  const ctx = createJobContext({
+    jobId,
+    sessionId: params.sessionId,
+    agentId: params.agentId,
+    mode: params.mode,
+    signal: params.signal,
+    source: params.source,
+    reason: params.reason,
+    heartbeatSnapshot: params.heartbeatSnapshot,
+  })
+  ctx.startedAt = Date.now()
+  state.activeJobs.set(jobId, ctx)
+  return ctx
+}
+
+/**
+ * Mark a job as completed and remove from active tracking.
+ */
+export function endJobExecution(jobId: string): JobContext | null {
+  const ctx = state.activeJobs.get(jobId)
+  if (!ctx) return null
+  ctx.endedAt = Date.now()
+  state.activeJobs.delete(jobId)
+  return ctx
+}
+
+/**
+ * Get the active job context by ID.
+ */
+export function getActiveJob(jobId: string): JobContext | null {
+  return state.activeJobs.get(jobId) || null
+}
+
+/**
+ * List all active job contexts for a session.
+ */
+export function getActiveJobsForSession(sessionId: string): JobContext[] {
+  return [...state.activeJobs.values()].filter((ctx) => ctx.sessionId === sessionId)
+}
+
+// ── Cleanup ─────────────────────────────────────────────────────────────
+
+export function cancelScheduledWake(jobId: string): boolean {
+  const timer = state.scheduledTimers.get(jobId)
+  if (!timer) return false
+  clearTimeout(timer)
+  state.scheduledTimers.delete(jobId)
+  return true
+}
+
+export function resetWakeDispatcherForTests(): void {
+  for (const timer of state.scheduledTimers.values()) {
+    clearTimeout(timer)
+  }
+  state.deferredQueue.clear()
+  state.scheduledTimers.clear()
+  state.activeJobs.clear()
+}
+
+// ── Diagnostics ─────────────────────────────────────────────────────────
+
+export function getWakeDispatcherStatus(): {
+  deferredQueueCount: number
+  scheduledTimerCount: number
+  activeJobCount: number
+} {
+  let deferredCount = 0
+  for (const queue of state.deferredQueue.values()) {
+    deferredCount += queue.length
+  }
+  return {
+    deferredQueueCount: deferredCount,
+    scheduledTimerCount: state.scheduledTimers.size,
+    activeJobCount: state.activeJobs.size,
+  }
+}

package/src/lib/server/wake-mode.test.ts

@@ -0,0 +1,161 @@
+import assert from 'node:assert/strict'
+import { describe, it } from 'node:test'
+
+import {
+  computeWakePriority,
+  createJobContext,
+  resolveRunAt,
+  sourceToWakeMode,
+  wakeModeToSource,
+} from './wake-mode'
+import type { WakeModeRequest } from './wake-mode'
+
+describe('WakeMode', () => {
+  describe('computeWakePriority', () => {
+    it('returns mode-based default priority when none specified', () => {
+      assert.equal(computeWakePriority({ mode: 'immediate' }), 80)
+      assert.equal(computeWakePriority({ mode: 'next_heartbeat' }), 40)
+      assert.equal(computeWakePriority({ mode: 'scheduled' }), 60)
+    })
+
+    it('uses explicit priority when provided', () => {
+      assert.equal(computeWakePriority({ mode: 'immediate', priority: 95 }), 95)
+      assert.equal(computeWakePriority({ mode: 'next_heartbeat', priority: 10 }), 10)
+    })
+
+    it('clamps priority to [0, 100]', () => {
+      assert.equal(computeWakePriority({ mode: 'immediate', priority: 150 }), 100)
+      assert.equal(computeWakePriority({ mode: 'immediate', priority: -5 }), 0)
+    })
+
+    it('ignores non-finite priority values', () => {
+      assert.equal(computeWakePriority({ mode: 'immediate', priority: NaN }), 80)
+      assert.equal(computeWakePriority({ mode: 'immediate', priority: Infinity }), 80)
+    })
+  })
+
+  describe('resolveRunAt', () => {
+    const NOW = 1_700_000_000_000
+
+    it('returns now for immediate mode', () => {
+      assert.equal(resolveRunAt({ mode: 'immediate' }, NOW), NOW)
+    })
+
+    it('returns null for next_heartbeat mode (deferred)', () => {
+      assert.equal(resolveRunAt({ mode: 'next_heartbeat' }, NOW), null)
+    })
+
+    it('returns absolute runAt for scheduled mode', () => {
+      const target = NOW + 60_000
+      assert.equal(resolveRunAt({ mode: 'scheduled', runAt: target }, NOW), target)
+    })
+
+    it('computes runAt from delayMs for scheduled mode', () => {
+      assert.equal(resolveRunAt({ mode: 'scheduled', delayMs: 5_000 }, NOW), NOW + 5_000)
+    })
+
+    it('clamps scheduled runAt to at least now', () => {
+      const pastTime = NOW - 10_000
+      assert.equal(resolveRunAt({ mode: 'scheduled', runAt: pastTime }, NOW), NOW)
+    })
+
+    it('falls back to now for scheduled mode without runAt or delayMs', () => {
+      assert.equal(resolveRunAt({ mode: 'scheduled' }, NOW), NOW)
+    })
+  })
+
+  describe('wakeModeToSource (backward compat)', () => {
+    it('maps immediate to heartbeat-wake', () => {
+      assert.equal(wakeModeToSource('immediate'), 'heartbeat-wake')
+    })
+
+    it('maps next_heartbeat to heartbeat', () => {
+      assert.equal(wakeModeToSource('next_heartbeat'), 'heartbeat')
+    })
+
+    it('maps scheduled to heartbeat-wake', () => {
+      assert.equal(wakeModeToSource('scheduled'), 'heartbeat-wake')
+    })
+  })
+
+  describe('sourceToWakeMode (legacy migration)', () => {
+    it('infers next_heartbeat from heartbeat source', () => {
+      assert.equal(sourceToWakeMode('heartbeat'), 'next_heartbeat')
+    })
+
+    it('infers immediate from heartbeat-wake source', () => {
+      assert.equal(sourceToWakeMode('heartbeat-wake'), 'immediate')
+    })
+
+    it('infers scheduled from schedule-prefixed source', () => {
+      assert.equal(sourceToWakeMode('schedule:nightly'), 'scheduled')
+    })
+
+    it('defaults to immediate for unknown sources', () => {
+      assert.equal(sourceToWakeMode('connector:slack'), 'immediate')
+    })
+  })
+
+  describe('createJobContext', () => {
+    it('creates an isolated context with scratchpad', () => {
+      const controller = new AbortController()
+      const ctx = createJobContext({
+        jobId: 'job-1',
+        sessionId: 'sess-1',
+        agentId: 'agent-1',
+        mode: 'immediate',
+        signal: controller.signal,
+        source: 'connector:slack',
+        reason: 'New message arrived',
+      })
+
+      assert.equal(ctx.jobId, 'job-1')
+      assert.equal(ctx.sessionId, 'sess-1')
+      assert.equal(ctx.agentId, 'agent-1')
+      assert.equal(ctx.mode, 'immediate')
+      assert.equal(ctx.source, 'connector:slack')
+      assert.equal(ctx.reason, 'New message arrived')
+      assert.ok(ctx.createdAt > 0)
+      assert.equal(ctx.startedAt, undefined)
+      assert.equal(ctx.endedAt, undefined)
+      assert.ok(ctx.scratchpad instanceof Map)
+      assert.equal(ctx.scratchpad.size, 0)
+    })
+
+    it('scratchpad isolates state between jobs', () => {
+      const controller = new AbortController()
+      const ctx1 = createJobContext({
+        jobId: 'job-a',
+        sessionId: 'sess-1',
+        mode: 'immediate',
+        signal: controller.signal,
+      })
+      const ctx2 = createJobContext({
+        jobId: 'job-b',
+        sessionId: 'sess-1',
+        mode: 'next_heartbeat',
+        signal: controller.signal,
+      })
+
+      ctx1.scratchpad.set('key', 'value-a')
+      ctx2.scratchpad.set('key', 'value-b')
+
+      assert.equal(ctx1.scratchpad.get('key'), 'value-a')
+      assert.equal(ctx2.scratchpad.get('key'), 'value-b')
+    })
+
+    it('captures heartbeat snapshot for isolation', () => {
+      const controller = new AbortController()
+      const snapshot = '# Heartbeat Tasks\n## Active\n- [ ] Send report'
+      const ctx = createJobContext({
+        jobId: 'job-snap',
+        sessionId: 'sess-1',
+        mode: 'next_heartbeat',
+        signal: controller.signal,
+        heartbeatSnapshot: snapshot,
+      })
+
+      assert.equal(ctx.heartbeatSnapshot, snapshot)
+    })
+  })
+})

package/src/lib/server/wake-mode.ts

@@ -0,0 +1,174 @@
+/**
+ * WakeMode — explicit scheduling semantics for heartbeat and task execution.
+ *
+ * Replaces the implicit `source: 'heartbeat' | 'heartbeat-wake'` convention
+ * with a formal enum that determines routing, priority, and isolation behavior.
+ *
+ * Inspired by OpenClaw's separation of "run now" vs "queue next heartbeat" vs
+ * scheduled execution with proper isolation.
+ */
+
+// ── WakeMode enum ───────────────────────────────────────────────────────
+
+export type WakeMode = 'immediate' | 'next_heartbeat' | 'scheduled'
+
+/**
+ * `immediate`       — Run now. Coalesced within a short window (250ms default),
+ *                      then dispatched. Used for connector events, watch-job
+ *                      triggers, approvals, webhooks.
+ *
+ * `next_heartbeat`  — Queue for the next periodic heartbeat tick. No coalesce
+ *                      window; the job waits until the heartbeat-service timer
+ *                      fires. Used for low-urgency background polling, system
+ *                      events that don't need instant reaction.
+ *
+ * `scheduled`        — Run at a specific future time (absolute or relative).
+ *                      Managed by the scheduler tick. Used for cron jobs,
+ *                      interval schedules, one-shot delayed wakes.
+ */
+
+// ── Job Context ─────────────────────────────────────────────────────────
+
+/**
+ * Isolated execution context for a wake job. Each job gets its own context
+ * so that failures, side effects, and state are contained.
+ */
+export interface JobContext {
+  /** Unique job execution ID. */
+  jobId: string
+  /** Which session this job targets. */
+  sessionId: string
+  /** Which agent (if any) owns this job. */
+  agentId?: string
+  /** The wake mode that created this job. */
+  mode: WakeMode
+  /** When the job was created/requested. */
+  createdAt: number
+  /** When the job actually started executing. */
+  startedAt?: number
+  /** When the job finished (success or failure). */
+  endedAt?: number
+  /** Abort controller for this specific job. */
+  signal: AbortSignal
+  /** Source identifier (e.g. 'connector:slack', 'schedule:nightly'). */
+  source?: string
+  /** Human-readable reason for this wake. */
+  reason?: string
+  /** Snapshot of HEARTBEAT.md at job start (for isolation). */
+  heartbeatSnapshot?: string
+  /** Job-scoped metadata accumulator — tools can stash results here without
+   *  polluting session state until the job completes successfully. */
+  scratchpad: Map<string, unknown>
+}
+
+export function createJobContext(params: {
+  jobId: string
+  sessionId: string
+  agentId?: string
+  mode: WakeMode
+  signal: AbortSignal
+  source?: string
+  reason?: string
+  heartbeatSnapshot?: string
+}): JobContext {
+  return {
+    jobId: params.jobId,
+    sessionId: params.sessionId,
+    agentId: params.agentId,
+    mode: params.mode,
+    createdAt: Date.now(),
+    signal: params.signal,
+    source: params.source,
+    reason: params.reason,
+    heartbeatSnapshot: params.heartbeatSnapshot,
+    scratchpad: new Map(),
+  }
+}
+
+// ── Wake Request with explicit mode ─────────────────────────────────────
+
+export interface WakeModeRequest {
+  mode: WakeMode
+  agentId?: string
+  sessionId?: string
+  reason?: string
+  source?: string
+  resumeMessage?: string
+  detail?: string
+  priority?: number
+  /** For `scheduled` mode: absolute timestamp (ms) to execute at. */
+  runAt?: number
+  /** For `scheduled` mode: relative delay (ms) from now. */
+  delayMs?: number
+  /** Event ID for deduplication. */
+  eventId?: string
+}
+
+// ── Priority mapping per mode ───────────────────────────────────────────
+
+const MODE_BASE_PRIORITY: Record<WakeMode, number> = {
+  immediate: 80,
+  scheduled: 60,
+  next_heartbeat: 40,
+}
+
+/**
+ * Compute effective priority for a wake request. Explicit priority overrides
+ * the mode-based default; otherwise the mode determines the base.
+ */
+export function computeWakePriority(request: WakeModeRequest): number {
+  if (typeof request.priority === 'number' && Number.isFinite(request.priority)) {
+    return Math.max(0, Math.min(100, Math.trunc(request.priority)))
+  }
+  return MODE_BASE_PRIORITY[request.mode]
+}
+
+/**
+ * Resolve the target execution time for a wake request.
+ * - `immediate`: now (or within coalesce window)
+ * - `next_heartbeat`: null (deferred to next tick)
+ * - `scheduled`: absolute time from runAt or now + delayMs
+ */
+export function resolveRunAt(request: WakeModeRequest, now = Date.now()): number | null {
+  switch (request.mode) {
+    case 'immediate':
+      return now
+    case 'next_heartbeat':
+      return null
+    case 'scheduled': {
+      if (typeof request.runAt === 'number' && Number.isFinite(request.runAt)) {
+        return Math.max(now, Math.trunc(request.runAt))
+      }
+      if (typeof request.delayMs === 'number' && Number.isFinite(request.delayMs)) {
+        return now + Math.max(0, Math.trunc(request.delayMs))
+      }
+      return now
+    }
+  }
+}
+
+/**
+ * Map a WakeMode to the session-run-manager source string.
+ * Maintains backward compatibility with existing heartbeat-source checks.
+ */
+export function wakeModeToSource(mode: WakeMode): string {
+  switch (mode) {
+    case 'immediate':
+      return 'heartbeat-wake'
+    case 'next_heartbeat':
+      return 'heartbeat'
+    case 'scheduled':
+      return 'heartbeat-wake'
+  }
+}
+
+/**
+ * Infer WakeMode from a legacy source string.
+ * Used during migration to preserve backward compat with existing callers.
+ */
+export function sourceToWakeMode(source: string): WakeMode {
+  if (source === 'heartbeat') return 'next_heartbeat'
+  if (source === 'heartbeat-wake') return 'immediate'
+  if (source.startsWith('schedule')) return 'scheduled'
+  return 'immediate'
+}

package/src/lib/server/wallet-service.ts

@@ -9,7 +9,8 @@ import {
   normalizeAtomicString,
 } from '@/lib/wallet'
 import type { Agent, AgentWallet, WalletChain, WalletTransaction } from '@/types'
-import { loadAgents, loadWalletTransactions, loadWallets, saveAgents, upsertWallet } from './storage'
+import { dedup } from '@/lib/shared-utils'
+import { loadAgent, loadAgents, loadWalletTransactions, loadWallets, upsertAgent, upsertWallet } from './storage'
 import { generateEthereumWallet, isValidEthereumAddress, sendEth } from './ethereum'
 import { generateSolanaKeypair, isValidSolanaAddress, sendSol } from './solana'
 import { notify } from './ws-hub'
@@ -31,7 +32,7 @@ export function getAgentWalletIds(agent: Pick<Agent, 'walletIds' | 'walletId'> |
   const legacy = typeof agent?.walletId === 'string' && agent.walletId.trim()
     ? [agent.walletId.trim()]
     : []
-  return [...new Set([...ids, ...legacy])]
+  return dedup([...ids, ...legacy])
 }
 
 export function getAgentActiveWalletId(
@@ -44,7 +45,7 @@ export function getAgentActiveWalletId(
 }
 
 function syncAgentWalletPointers(agent: Agent, walletIds: string[], activeWalletId?: string | null): Agent {
-  const normalizedIds = [...new Set(walletIds.filter(Boolean))]
+  const normalizedIds = dedup(walletIds.filter(Boolean))
   const normalizedActive = activeWalletId && normalizedIds.includes(activeWalletId)
     ? activeWalletId
     : normalizedIds[0] || null
@@ -102,8 +103,8 @@ export function createAgentWallet(input: {
   const agentId = String(input.agentId || '').trim()
   if (!agentId) throw new Error('agentId is required')
 
-  const agents = loadAgents()
-  if (!agents[agentId]) throw new Error('Agent not found')
+  const agent = loadAgent(agentId)
+  if (!agent) throw new Error('Agent not found')
 
   const chain = getWalletChainOrDefault(input.chain ?? input.provider, 'solana')
   const existing = getWalletByAgentId(agentId, chain)
@@ -128,11 +129,9 @@ export function createAgentWallet(input: {
   upsertWallet(id, wallet)
   clearWalletPortfolioCache(id)
 
-  const agent = agents[agentId]
-  linkWalletToAgent(agent, id, getAgentActiveWalletId(agent) == null)
+  linkWalletToAgent(agent as any, id, getAgentActiveWalletId(agent as any) == null)
   agent.updatedAt = now
-  agents[agentId] = agent
-  saveAgents(agents)
+  upsertAgent(agentId, agent)
 
   notify('wallets')
   notify('agents')