From 86bb5cc82701f3f668389e31360bde98912250cc Mon Sep 17 00:00:00 2001 From: Mohamed Boudra Date: Thu, 16 Apr 2026 10:44:57 +0800 Subject: [PATCH] fix: make MCP tools work for archived agents, matching CLI code paths (#423) Extract shared functions from session.ts (toAgentPersistenceHandle, buildStoredAgentPayload, ensureAgentLoaded) so both CLI/WebSocket handlers and MCP tools use the same code paths for agent lookup. Fix get_agent_status, get_agent_activity, and list_agents MCP tools to fall back to persistent storage for archived agents. Add includeArchived param to list_agents. Fix setupFinishNotification to not wake archived callers. Delete dead agent-management-mcp.ts. --- .../server/src/server/agent/agent-loading.ts | 80 ++ .../src/server/agent/agent-management-mcp.ts | 1025 ----------------- .../src/server/agent/agent-projections.ts | 92 ++ .../src/server/agent/mcp-server.test.ts | 267 ++++- .../server/src/server/agent/mcp-server.ts | 64 +- .../src/server/agent/mcp-shared.test.ts | 78 ++ .../server/src/server/agent/mcp-shared.ts | 10 +- .../server/src/server/persistence-hooks.ts | 76 +- .../src/server/schedule/service.test.ts | 58 + .../server/src/server/schedule/service.ts | 79 +- packages/server/src/server/session.ts | 248 +--- 11 files changed, 751 insertions(+), 1326 deletions(-) create mode 100644 packages/server/src/server/agent/agent-loading.ts delete mode 100644 packages/server/src/server/agent/agent-management-mcp.ts create mode 100644 packages/server/src/server/agent/mcp-shared.test.ts diff --git a/packages/server/src/server/agent/agent-loading.ts b/packages/server/src/server/agent/agent-loading.ts new file mode 100644 index 000000000..01ae00630 --- /dev/null +++ b/packages/server/src/server/agent/agent-loading.ts @@ -0,0 +1,80 @@ +import type { Logger } from "pino"; + +import type { AgentProvider } from "./agent-sdk-types.js"; +import type { AgentManager, ManagedAgent } from "./agent-manager.js"; +import type { AgentStorage } from "./agent-storage.js"; +import { + buildConfigOverrides, + buildSessionConfig, + extractTimestamps, + toAgentPersistenceHandle, +} from "../persistence-hooks.js"; + +const pendingAgentInitializations = new Map>(); + +export interface EnsureAgentLoadedDeps { + agentManager: AgentManager; + agentStorage: AgentStorage; + validProviders?: Iterable; + logger: Logger; +} + +export async function ensureAgentLoaded( + agentId: string, + deps: EnsureAgentLoadedDeps, +): Promise { + const existing = deps.agentManager.getAgent(agentId); + if (existing) { + return existing; + } + + const inflight = pendingAgentInitializations.get(agentId); + if (inflight) { + return inflight; + } + + const initPromise = (async () => { + const record = await deps.agentStorage.get(agentId); + if (!record) { + throw new Error(`Agent not found: ${agentId}`); + } + + const validProviders = deps.validProviders ?? deps.agentManager.getRegisteredProviderIds(); + const handle = toAgentPersistenceHandle(deps.logger, validProviders, record.persistence); + + let snapshot: ManagedAgent; + if (handle) { + snapshot = await deps.agentManager.resumeAgentFromPersistence( + handle, + buildConfigOverrides(record), + agentId, + extractTimestamps(record), + ); + deps.logger.info({ agentId, provider: record.provider }, "Agent resumed from persistence"); + } else { + const config = buildSessionConfig(record, { + validProviders, + logger: deps.logger, + }); + if (!config) { + throw new Error(`Agent ${agentId} references unavailable provider '${record.provider}'`); + } + snapshot = await deps.agentManager.createAgent(config, agentId, { labels: record.labels }); + deps.logger.info({ agentId, provider: record.provider }, "Agent created from stored config"); + } + + await deps.agentManager.hydrateTimelineFromProvider(agentId); + return deps.agentManager.getAgent(agentId) ?? snapshot; + })(); + + pendingAgentInitializations.set(agentId, initPromise); + + try { + return await initPromise; + } finally { + const current = pendingAgentInitializations.get(agentId); + if (current === initPromise) { + pendingAgentInitializations.delete(agentId); + } + } +} diff --git a/packages/server/src/server/agent/agent-management-mcp.ts b/packages/server/src/server/agent/agent-management-mcp.ts deleted file mode 100644 index 6070fd5ec..000000000 --- a/packages/server/src/server/agent/agent-management-mcp.ts +++ /dev/null @@ -1,1025 +0,0 @@ -/** - * Agent Management MCP Server - * - * Purpose: Managing agents from the UI/voice assistant LLM - * Transport: In-memory (runs in-process with the voice assistant LLM) - * Server name: "paseo-agent-management" - * - * Tools: - * - create_agent - * - wait_for_agent - * - send_agent_prompt - * - get_agent_status - * - list_agents - * - cancel_agent - * - kill_agent - * - get_agent_activity - * - set_agent_mode - * - list_pending_permissions - * - respond_to_permission - * - * No callerAgentId needed - voice assistant is not an agent. - */ - -import { McpServer } from "@modelcontextprotocol/sdk/server/mcp.js"; -import { z } from "zod"; -import { ensureValidJson } from "../json-utils.js"; -import type { Logger } from "pino"; - -import type { AgentProvider } from "./agent-sdk-types.js"; -import type { AgentManager, WaitForAgentResult } from "./agent-manager.js"; -import { - AgentPermissionRequestPayloadSchema, - AgentPermissionResponseSchema, - AgentSnapshotPayloadSchema, -} from "../messages.js"; -import { toAgentPayload } from "./agent-projections.js"; -import { curateAgentActivity } from "./activity-curator.js"; -import { AgentStorage } from "./agent-storage.js"; -import { - appendTimelineItemIfAgentKnown, - emitLiveTimelineItemIfAgentKnown, -} from "./timeline-append.js"; -import { type WorktreeConfig } from "../../utils/worktree.js"; -import { WaitForAgentTracker } from "./wait-for-agent-tracker.js"; -import { scheduleAgentMetadataGeneration } from "./agent-metadata-generator.js"; -import { expandUserPath } from "../path-utils.js"; -import type { TerminalManager } from "../../terminal/terminal-manager.js"; -import { createAgentWorktree, runAsyncWorktreeBootstrap } from "../worktree-bootstrap.js"; -import type { ScheduleService } from "../schedule/service.js"; -import { ScheduleSummarySchema, StoredScheduleSchema } from "../schedule/types.js"; -import type { ProviderDefinition } from "./provider-registry.js"; -import { - AgentModelSchema, - AgentProviderEnum, - AgentStatusEnum, - ProviderSummarySchema, - parseDurationString, - resolveProviderAndModel, - sanitizePermissionRequest, - serializeSnapshotWithMetadata, - startAgentRun, - toScheduleSummary, - waitForAgentWithTimeout, -} from "./mcp-shared.js"; - -export interface AgentManagementMcpOptions { - agentManager: AgentManager; - agentStorage: AgentStorage; - terminalManager?: TerminalManager | null; - scheduleService?: ScheduleService | null; - providerRegistry?: Record | null; - paseoHome?: string; - logger: Logger; -} - -export async function createAgentManagementMcpServer( - options: AgentManagementMcpOptions, -): Promise { - const { agentManager, agentStorage, scheduleService, providerRegistry, logger } = options; - const childLogger = logger.child({ - module: "agent", - component: "agent-management-mcp", - }); - const waitTracker = new WaitForAgentTracker(logger); - const resolveNewAgentScheduleTarget = (params?: { provider?: string; cwd?: string }) => { - const resolvedProviderModel = resolveProviderAndModel({ - provider: params?.provider, - defaultProvider: "claude", - }); - return { - type: "new-agent" as const, - config: { - provider: resolvedProviderModel.provider, - cwd: params?.cwd?.trim() ? expandUserPath(params.cwd) : process.cwd(), - ...(resolvedProviderModel.model ? { model: resolvedProviderModel.model } : {}), - }, - }; - }; - - const server = new McpServer({ - name: "paseo-agent-management", - version: "1.0.0", - }); - - const inputSchema = { - cwd: z - .string() - .describe("Required working directory for the agent (absolute, relative, or ~)."), - title: z - .string() - .trim() - .min(1, "Title is required") - .max(60, "Title must be 60 characters or fewer") - .describe("Short descriptive title (<= 60 chars) summarizing the agent's focus."), - provider: AgentProviderEnum.optional().describe( - "Optional agent implementation to spawn. Defaults to 'claude'.", - ), - model: z.string().optional().describe("Model to use (e.g. claude-sonnet-4-20250514)"), - thinking: z.string().optional().describe("Thinking option ID"), - labels: z.record(z.string(), z.string()).optional().describe("Labels to set on the agent"), - initialPrompt: z - .string() - .optional() - .describe("Optional task to start immediately after creation (non-blocking)."), - mode: z - .string() - .optional() - .describe("Optional session mode to configure before the first run."), - worktreeName: z - .string() - .optional() - .describe("Optional git worktree branch name (lowercase alphanumerics + hyphen)."), - baseBranch: z - .string() - .optional() - .describe("Required when worktreeName is set: the base branch to diff/merge against."), - background: z - .boolean() - .optional() - .default(false) - .describe( - "Run agent in background. If false (default), waits for completion or permission request. If true, returns immediately.", - ), - }; - - server.registerTool( - "create_agent", - { - title: "Create agent", - description: - "Create a new Claude or Codex agent tied to a working directory. Optionally run an initial prompt immediately or create a git worktree for the agent.", - inputSchema, - outputSchema: { - agentId: z.string(), - type: AgentProviderEnum, - status: AgentStatusEnum, - cwd: z.string(), - currentModeId: z.string().nullable(), - availableModes: z.array( - z.object({ - id: z.string(), - label: z.string(), - description: z.string().nullable().optional(), - }), - ), - lastMessage: z.string().nullable().optional(), - permission: AgentPermissionRequestPayloadSchema.nullable().optional(), - }, - }, - async (args) => { - const { - cwd, - provider, - initialPrompt, - mode, - worktreeName, - baseBranch, - background = false, - title, - model, - thinking, - labels, - } = args as { - cwd: string; - provider?: AgentProvider; - initialPrompt?: string; - mode?: string; - worktreeName?: string; - baseBranch?: string; - background?: boolean; - title: string; - model?: string; - thinking?: string; - labels?: Record; - }; - - let resolvedCwd = expandUserPath(cwd); - let worktreeConfig: WorktreeConfig | undefined; - - if (worktreeName) { - if (!baseBranch) { - throw new Error("baseBranch is required when creating a worktree"); - } - const worktree = await createAgentWorktree({ - branchName: worktreeName, - cwd: resolvedCwd, - baseBranch, - worktreeSlug: worktreeName, - paseoHome: options.paseoHome, - }); - resolvedCwd = worktree.worktreePath; - worktreeConfig = worktree; - } - - const resolvedProvider: AgentProvider = provider ?? "claude"; - const normalizedTitle = title?.trim() ?? null; - const snapshot = await agentManager.createAgent( - { - provider: resolvedProvider, - cwd: resolvedCwd, - modeId: mode, - title: normalizedTitle ?? undefined, - model, - thinkingOptionId: thinking, - }, - undefined, - labels ? { labels } : undefined, - ); - - if (worktreeConfig) { - void runAsyncWorktreeBootstrap({ - agentId: snapshot.id, - worktree: worktreeConfig, - terminalManager: options.terminalManager ?? null, - appendTimelineItem: (item) => - appendTimelineItemIfAgentKnown({ - agentManager, - agentId: snapshot.id, - item, - }), - emitLiveTimelineItem: (item) => - emitLiveTimelineItemIfAgentKnown({ - agentManager, - agentId: snapshot.id, - item, - }), - logger: childLogger, - }); - } - - const trimmedPrompt = initialPrompt?.trim(); - if (trimmedPrompt) { - scheduleAgentMetadataGeneration({ - agentManager, - agentId: snapshot.id, - cwd: snapshot.cwd, - initialPrompt: trimmedPrompt, - explicitTitle: normalizedTitle ?? undefined, - paseoHome: options.paseoHome, - logger: childLogger, - }); - - try { - agentManager.recordUserMessage(snapshot.id, trimmedPrompt, { - emitState: false, - }); - } catch (error) { - childLogger.error( - { err: error, agentId: snapshot.id }, - "Failed to record initial prompt", - ); - } - - try { - startAgentRun(agentManager, snapshot.id, trimmedPrompt, childLogger); - - if (!background) { - const result = await waitForAgentWithTimeout(agentManager, snapshot.id, { - waitForActive: true, - }); - - const responseData = { - agentId: snapshot.id, - type: provider, - status: result.status, - cwd: snapshot.cwd, - currentModeId: snapshot.currentModeId, - availableModes: snapshot.availableModes, - lastMessage: result.lastMessage, - permission: sanitizePermissionRequest(result.permission), - }; - const validJson = ensureValidJson(responseData); - - return { - content: [], - structuredContent: validJson, - }; - } - } catch (error) { - childLogger.error({ err: error, agentId: snapshot.id }, "Failed to run initial prompt"); - } - } - - return { - content: [], - structuredContent: ensureValidJson({ - agentId: snapshot.id, - type: provider, - status: snapshot.lifecycle, - cwd: snapshot.cwd, - currentModeId: snapshot.currentModeId, - availableModes: snapshot.availableModes, - lastMessage: null, - permission: null, - }), - }; - }, - ); - - server.registerTool( - "wait_for_agent", - { - title: "Wait for agent", - description: - "Block until the agent requests permission or the current run completes. Returns the pending permission (if any) and recent activity summary.", - inputSchema: { - agentId: z.string().describe("Agent identifier returned by the create_agent tool"), - }, - outputSchema: { - agentId: z.string(), - status: AgentStatusEnum, - permission: AgentPermissionRequestPayloadSchema.nullable(), - lastMessage: z.string().nullable(), - }, - }, - async ({ agentId }, { signal }) => { - const abortController = new AbortController(); - const cleanupFns: Array<() => void> = []; - - const cleanup = () => { - while (cleanupFns.length) { - const fn = cleanupFns.pop(); - try { - fn?.(); - } catch { - // ignore cleanup errors - } - } - }; - - const forwardExternalAbort = () => { - if (!abortController.signal.aborted) { - const reason = signal?.reason ?? new Error("wait_for_agent aborted"); - abortController.abort(reason); - } - }; - - if (signal) { - if (signal.aborted) { - forwardExternalAbort(); - } else { - signal.addEventListener("abort", forwardExternalAbort, { - once: true, - }); - cleanupFns.push(() => signal.removeEventListener("abort", forwardExternalAbort)); - } - } - - const unregister = waitTracker.register(agentId, (reason) => { - if (!abortController.signal.aborted) { - abortController.abort(new Error(reason ?? "wait_for_agent cancelled")); - } - }); - cleanupFns.push(unregister); - - try { - const result: WaitForAgentResult = await waitForAgentWithTimeout(agentManager, agentId, { - signal: abortController.signal, - }); - - const validJson = ensureValidJson({ - agentId, - status: result.status, - permission: sanitizePermissionRequest(result.permission), - lastMessage: result.lastMessage, - }); - - return { - content: [], - structuredContent: validJson, - }; - } finally { - cleanup(); - } - }, - ); - - server.registerTool( - "send_agent_prompt", - { - title: "Send agent prompt", - description: - "Send a task to a running agent. Returns immediately after the agent begins processing.", - inputSchema: { - agentId: z.string(), - prompt: z.string(), - sessionMode: z - .string() - .optional() - .describe("Optional mode to set before running the prompt."), - background: z - .boolean() - .optional() - .default(false) - .describe( - "Run agent in background. If false (default), waits for completion or permission request. If true, returns immediately.", - ), - }, - outputSchema: { - success: z.boolean(), - status: AgentStatusEnum, - lastMessage: z.string().nullable().optional(), - permission: AgentPermissionRequestPayloadSchema.nullable().optional(), - }, - }, - async ({ agentId, prompt, sessionMode, background = false }) => { - const snapshot = agentManager.getAgent(agentId); - if (!snapshot) { - throw new Error(`Agent ${agentId} not found`); - } - - if (agentManager.hasInFlightRun(agentId)) { - waitTracker.cancel(agentId, "Agent run interrupted by new prompt"); - } - - if (sessionMode) { - await agentManager.setAgentMode(agentId, sessionMode); - } - - try { - agentManager.recordUserMessage(agentId, prompt, { - emitState: false, - }); - } catch (error) { - childLogger.error({ err: error, agentId }, "Failed to record user message"); - } - - startAgentRun(agentManager, agentId, prompt, childLogger, { - replaceRunning: true, - }); - - if (!background) { - const result = await waitForAgentWithTimeout(agentManager, agentId, { - waitForActive: true, - }); - - const responseData = { - success: true, - status: result.status, - lastMessage: result.lastMessage, - permission: sanitizePermissionRequest(result.permission), - }; - const validJson = ensureValidJson(responseData); - - return { - content: [], - structuredContent: validJson, - }; - } - - const currentSnapshot = agentManager.getAgent(agentId); - - const responseData = { - success: true, - status: currentSnapshot?.lifecycle ?? "idle", - lastMessage: null, - permission: null, - }; - const validJson = ensureValidJson(responseData); - - return { - content: [], - structuredContent: validJson, - }; - }, - ); - - server.registerTool( - "get_agent_status", - { - title: "Get agent status", - description: - "Return the latest snapshot for an agent, including lifecycle state, capabilities, and pending permissions.", - inputSchema: { - agentId: z.string(), - }, - outputSchema: { - status: AgentStatusEnum, - snapshot: AgentSnapshotPayloadSchema, - }, - }, - async ({ agentId }) => { - const snapshot = agentManager.getAgent(agentId); - if (!snapshot) { - throw new Error(`Agent ${agentId} not found`); - } - - const structuredSnapshot = await serializeSnapshotWithMetadata( - agentStorage, - snapshot, - childLogger, - ); - return { - content: [], - structuredContent: ensureValidJson({ - status: snapshot.lifecycle, - snapshot: structuredSnapshot, - }), - }; - }, - ); - - server.registerTool( - "list_agents", - { - title: "List agents", - description: "List all live agents managed by the server.", - inputSchema: {}, - outputSchema: { - agents: z.array(AgentSnapshotPayloadSchema), - }, - }, - async () => { - const snapshots = agentManager.listAgents(); - const agents = await Promise.all( - snapshots.map((snapshot) => - serializeSnapshotWithMetadata(agentStorage, snapshot, childLogger), - ), - ); - return { - content: [], - structuredContent: ensureValidJson({ agents }), - }; - }, - ); - - server.registerTool( - "cancel_agent", - { - title: "Cancel agent run", - description: "Abort the agent's current run but keep the agent alive for future tasks.", - inputSchema: { - agentId: z.string(), - }, - outputSchema: { - success: z.boolean(), - }, - }, - async ({ agentId }) => { - const success = await agentManager.cancelAgentRun(agentId); - if (success) { - waitTracker.cancel(agentId, "Agent run cancelled"); - } - return { - content: [], - structuredContent: ensureValidJson({ success }), - }; - }, - ); - - server.registerTool( - "archive_agent", - { - title: "Archive agent", - description: - "Archive an agent (soft-delete). The agent is interrupted if running and removed from the active list.", - inputSchema: { - agentId: z.string(), - }, - outputSchema: { - success: z.boolean(), - }, - }, - async ({ agentId }) => { - await agentManager.archiveAgent(agentId); - waitTracker.cancel(agentId, "Agent archived"); - return { - content: [], - structuredContent: ensureValidJson({ success: true }), - }; - }, - ); - - server.registerTool( - "kill_agent", - { - title: "Kill agent", - description: "Terminate an agent session permanently.", - inputSchema: { - agentId: z.string(), - }, - outputSchema: { - success: z.boolean(), - }, - }, - async ({ agentId }) => { - await agentManager.closeAgent(agentId); - waitTracker.cancel(agentId, "Agent terminated"); - return { - content: [], - structuredContent: ensureValidJson({ success: true }), - }; - }, - ); - - server.registerTool( - "update_agent", - { - title: "Update agent", - description: "Update an agent name and/or labels.", - inputSchema: { - agentId: z.string(), - name: z.string().optional(), - labels: z.record(z.string(), z.string()).optional().describe("Labels to set on the agent"), - }, - outputSchema: { - success: z.boolean(), - }, - }, - async ({ agentId, name, labels }) => { - const trimmedName = name?.trim(); - if (trimmedName) { - const record = await agentStorage.get(agentId); - if (!record) { - throw new Error(`Agent ${agentId} not found`); - } - await agentStorage.upsert({ - ...record, - title: trimmedName, - updatedAt: new Date().toISOString(), - }); - agentManager.notifyAgentState(agentId); - } - - if (labels) { - await agentManager.setLabels(agentId, labels); - } - - return { - content: [], - structuredContent: ensureValidJson({ success: true }), - }; - }, - ); - - server.registerTool( - "create_schedule", - { - title: "Create schedule", - description: "Create a recurring schedule that runs on an agent or a new agent.", - inputSchema: { - prompt: z.string().trim().min(1, "prompt is required"), - every: z.string().optional(), - cron: z.string().optional(), - name: z.string().optional(), - target: z.enum(["self", "new-agent"]).optional(), - provider: AgentProviderEnum.optional().describe( - "Provider, or provider/model (for example: codex or codex/gpt-5.4).", - ), - cwd: z.string().optional(), - maxRuns: z.number().int().positive().optional(), - expiresIn: z.string().optional(), - }, - outputSchema: ScheduleSummarySchema.shape, - }, - async ({ prompt, every, cron, name, target, provider, cwd, maxRuns, expiresIn }) => { - if (!scheduleService) { - throw new Error("Schedule service is not configured"); - } - - const cadenceCount = Number(every !== undefined) + Number(cron !== undefined); - if (cadenceCount !== 1) { - throw new Error("Specify exactly one of every or cron"); - } - if (target === "self") { - throw new Error("target=self requires a caller agent"); - } - - const schedule = await scheduleService.create({ - prompt: prompt.trim(), - cadence: every - ? { type: "every" as const, everyMs: parseDurationString(every) } - : { type: "cron" as const, expression: cron!.trim() }, - target: resolveNewAgentScheduleTarget({ provider, cwd }), - ...(name?.trim() ? { name: name.trim() } : {}), - ...(maxRuns === undefined ? {} : { maxRuns }), - ...(expiresIn === undefined - ? {} - : { expiresAt: new Date(Date.now() + parseDurationString(expiresIn)).toISOString() }), - }); - - return { - content: [], - structuredContent: ensureValidJson(toScheduleSummary(schedule)), - }; - }, - ); - - server.registerTool( - "list_schedules", - { - title: "List schedules", - description: "List all schedules managed by the daemon.", - inputSchema: {}, - outputSchema: { - schedules: z.array(ScheduleSummarySchema), - }, - }, - async () => { - if (!scheduleService) { - throw new Error("Schedule service is not configured"); - } - - const schedules = (await scheduleService.list()).map((schedule) => - toScheduleSummary(schedule), - ); - return { - content: [], - structuredContent: ensureValidJson({ schedules }), - }; - }, - ); - - server.registerTool( - "inspect_schedule", - { - title: "Inspect schedule", - description: "Inspect a schedule and its run history.", - inputSchema: { - id: z.string(), - }, - outputSchema: StoredScheduleSchema.shape, - }, - async ({ id }) => { - if (!scheduleService) { - throw new Error("Schedule service is not configured"); - } - - const schedule = await scheduleService.inspect(id); - return { - content: [], - structuredContent: ensureValidJson(schedule), - }; - }, - ); - - server.registerTool( - "pause_schedule", - { - title: "Pause schedule", - description: "Pause an active schedule.", - inputSchema: { - id: z.string(), - }, - outputSchema: { - success: z.boolean(), - }, - }, - async ({ id }) => { - if (!scheduleService) { - throw new Error("Schedule service is not configured"); - } - - await scheduleService.pause(id); - return { - content: [], - structuredContent: ensureValidJson({ success: true }), - }; - }, - ); - - server.registerTool( - "resume_schedule", - { - title: "Resume schedule", - description: "Resume a paused schedule.", - inputSchema: { - id: z.string(), - }, - outputSchema: { - success: z.boolean(), - }, - }, - async ({ id }) => { - if (!scheduleService) { - throw new Error("Schedule service is not configured"); - } - - await scheduleService.resume(id); - return { - content: [], - structuredContent: ensureValidJson({ success: true }), - }; - }, - ); - - server.registerTool( - "delete_schedule", - { - title: "Delete schedule", - description: "Delete a schedule permanently.", - inputSchema: { - id: z.string(), - }, - outputSchema: { - success: z.boolean(), - }, - }, - async ({ id }) => { - if (!scheduleService) { - throw new Error("Schedule service is not configured"); - } - - await scheduleService.delete(id); - return { - content: [], - structuredContent: ensureValidJson({ success: true }), - }; - }, - ); - - server.registerTool( - "list_providers", - { - title: "List providers", - description: "List available agent providers and their modes.", - inputSchema: {}, - outputSchema: { - providers: z.array(ProviderSummarySchema), - }, - }, - async () => ({ - content: [], - structuredContent: ensureValidJson({ - providers: Object.values(providerRegistry ?? {}).map((provider) => ({ - id: provider.id, - label: provider.label, - modes: provider.modes.map((mode) => ({ - id: mode.id, - label: mode.label, - ...(mode.description ? { description: mode.description } : {}), - })), - })), - }), - }), - ); - - server.registerTool( - "list_models", - { - title: "List models", - description: "List models for an agent provider.", - inputSchema: { - provider: AgentProviderEnum, - }, - outputSchema: { - provider: z.string(), - models: z.array(AgentModelSchema), - }, - }, - async ({ provider }) => { - if (!providerRegistry) { - throw new Error("Provider registry is not configured"); - } - - const definition = providerRegistry[provider]; - if (!definition) { - throw new Error(`Provider ${provider} is not configured`); - } - - const models = await definition.fetchModels(); - return { - content: [], - structuredContent: ensureValidJson({ - provider, - models, - }), - }; - }, - ); - - server.registerTool( - "get_agent_activity", - { - title: "Get agent activity", - description: "Return recent agent timeline entries as a curated summary.", - inputSchema: { - agentId: z.string(), - limit: z - .number() - .optional() - .describe("Optional limit for number of activities to include (most recent first)."), - }, - outputSchema: { - agentId: z.string(), - updateCount: z.number(), - currentModeId: z.string().nullable(), - content: z.string(), - }, - }, - async ({ agentId, limit }) => { - const timeline = agentManager.getTimeline(agentId); - const snapshot = agentManager.getAgent(agentId); - - const activitiesToCurate = limit ? timeline.slice(-limit) : timeline; - - const curatedContent = curateAgentActivity(activitiesToCurate); - const totalCount = timeline.length; - const shownCount = activitiesToCurate.length; - - let countHeader: string; - if (limit && shownCount < totalCount) { - countHeader = `Showing ${shownCount} of ${totalCount} ${totalCount === 1 ? "activity" : "activities"} (limited to ${limit})`; - } else { - countHeader = `Showing all ${totalCount} ${totalCount === 1 ? "activity" : "activities"}`; - } - - const contentWithCount = `${countHeader}\n\n${curatedContent}`; - - return { - content: [], - structuredContent: ensureValidJson({ - agentId, - updateCount: timeline.length, - currentModeId: snapshot?.currentModeId ?? null, - content: contentWithCount, - }), - }; - }, - ); - - server.registerTool( - "set_agent_mode", - { - title: "Set agent session mode", - description: - "Switch the agent's session mode (plan, bypassPermissions, read-only, auto, etc.).", - inputSchema: { - agentId: z.string(), - modeId: z.string(), - }, - outputSchema: { - success: z.boolean(), - newMode: z.string(), - }, - }, - async ({ agentId, modeId }) => { - await agentManager.setAgentMode(agentId, modeId); - return { - content: [], - structuredContent: ensureValidJson({ success: true, newMode: modeId }), - }; - }, - ); - - server.registerTool( - "list_pending_permissions", - { - title: "List pending permissions", - description: - "Return all pending permission requests across all agents with the normalized payloads.", - inputSchema: {}, - outputSchema: { - permissions: z.array( - z.object({ - agentId: z.string(), - status: AgentStatusEnum, - request: AgentPermissionRequestPayloadSchema, - }), - ), - }, - }, - async () => { - const permissions = agentManager.listAgents().flatMap((agent) => { - const payload = toAgentPayload(agent); - return payload.pendingPermissions.map((request) => ({ - agentId: agent.id, - status: payload.status, - request, - })); - }); - - return { - content: [], - structuredContent: ensureValidJson({ permissions }), - }; - }, - ); - - server.registerTool( - "respond_to_permission", - { - title: "Respond to permission", - description: - "Approve or deny a pending permission request with an AgentManager-compatible response payload.", - inputSchema: { - agentId: z.string(), - requestId: z.string(), - response: AgentPermissionResponseSchema, - }, - outputSchema: { - success: z.boolean(), - }, - }, - async ({ agentId, requestId, response }) => { - await agentManager.respondToPermission(agentId, requestId, response); - return { - content: [], - structuredContent: ensureValidJson({ success: true }), - }; - }, - ); - - return server; -} diff --git a/packages/server/src/server/agent/agent-projections.ts b/packages/server/src/server/agent/agent-projections.ts index 5d89a25f6..899df40f4 100644 --- a/packages/server/src/server/agent/agent-projections.ts +++ b/packages/server/src/server/agent/agent-projections.ts @@ -13,6 +13,9 @@ import type { } from "./agent-sdk-types.js"; import type { ManagedAgent } from "./agent-manager.js"; import type { JsonValue } from "../json-utils.js"; +import type { Logger } from "pino"; +import { buildProviderRegistry } from "./provider-registry.js"; +import { coerceAgentProvider, toAgentPersistenceHandle } from "../persistence-hooks.js"; export type { ManagedAgent }; @@ -128,6 +131,95 @@ export function toAgentPayload( return payload; } +export function buildStoredAgentPayload( + record: StoredAgentRecord, + providerRegistry: ReturnType, + logger: Logger, +): AgentSnapshotPayload { + const defaultCapabilities = { + supportsStreaming: false, + supportsSessionPersistence: true, + supportsDynamicModes: false, + supportsMcpServers: false, + supportsReasoningStream: false, + supportsToolInvocations: true, + } as const; + + const createdAt = new Date(record.createdAt); + const updatedAt = new Date(resolveStoredAgentPayloadUpdatedAt(record)); + const lastUserMessageAt = record.lastUserMessageAt ? new Date(record.lastUserMessageAt) : null; + + const provider = coerceAgentProvider(logger, providerRegistry, record.provider, record.id); + const runtimeInfo = record.runtimeInfo + ? { + provider: coerceAgentProvider( + logger, + providerRegistry, + record.runtimeInfo.provider, + record.id, + ), + sessionId: record.runtimeInfo.sessionId, + ...(Object.prototype.hasOwnProperty.call(record.runtimeInfo, "model") + ? { model: record.runtimeInfo.model ?? null } + : {}), + ...(Object.prototype.hasOwnProperty.call(record.runtimeInfo, "thinkingOptionId") + ? { thinkingOptionId: record.runtimeInfo.thinkingOptionId ?? null } + : {}), + ...(Object.prototype.hasOwnProperty.call(record.runtimeInfo, "modeId") + ? { modeId: record.runtimeInfo.modeId ?? null } + : {}), + ...(record.runtimeInfo.extra ? { extra: record.runtimeInfo.extra } : {}), + } + : undefined; + + return { + id: record.id, + provider, + cwd: record.cwd, + model: record.config?.model ?? null, + thinkingOptionId: record.config?.thinkingOptionId ?? null, + effectiveThinkingOptionId: resolveEffectiveThinkingOptionId({ + runtimeInfo, + configuredThinkingOptionId: record.config?.thinkingOptionId ?? null, + }), + ...(runtimeInfo ? { runtimeInfo } : {}), + createdAt: createdAt.toISOString(), + updatedAt: updatedAt.toISOString(), + lastUserMessageAt: lastUserMessageAt ? lastUserMessageAt.toISOString() : null, + status: record.lastStatus, + capabilities: defaultCapabilities, + currentModeId: record.lastModeId ?? null, + availableModes: [], + pendingPermissions: [], + persistence: toAgentPersistenceHandle(logger, providerRegistry, record.persistence), + lastUsage: undefined, + lastError: undefined, + title: record.title ?? record.config?.title ?? null, + requiresAttention: record.requiresAttention ?? false, + attentionReason: record.attentionReason ?? null, + attentionTimestamp: record.attentionTimestamp ?? null, + archivedAt: record.archivedAt ?? null, + labels: record.labels, + }; +} + +export function resolveStoredAgentPayloadUpdatedAt(record: StoredAgentRecord): string { + const timestamps = [record.updatedAt, record.lastActivityAt] + .filter((value): value is string => typeof value === "string" && value.length > 0) + .map((value) => ({ + raw: value, + parsed: Date.parse(value), + })) + .filter((value) => !Number.isNaN(value.parsed)); + + if (timestamps.length === 0) { + return record.updatedAt; + } + + timestamps.sort((a, b) => b.parsed - a.parsed); + return timestamps[0].raw; +} + function buildSerializableConfig(config: AgentSessionConfig): SerializableAgentConfig | null { const serializable: SerializableAgentConfig = {}; if (Object.prototype.hasOwnProperty.call(config, "title")) { diff --git a/packages/server/src/server/agent/mcp-server.test.ts b/packages/server/src/server/agent/mcp-server.test.ts index f2394be1b..007068afe 100644 --- a/packages/server/src/server/agent/mcp-server.test.ts +++ b/packages/server/src/server/agent/mcp-server.test.ts @@ -6,7 +6,7 @@ import { tmpdir } from "node:os"; import { createTestLogger } from "../../test-utils/test-logger.js"; import { createAgentMcpServer } from "./mcp-server.js"; import type { AgentManager, ManagedAgent } from "./agent-manager.js"; -import type { AgentStorage } from "./agent-storage.js"; +import type { AgentStorage, StoredAgentRecord } from "./agent-storage.js"; import type { ProviderDefinition } from "./provider-registry.js"; type TestDeps = { @@ -29,10 +29,17 @@ function createTestDeps(): TestDeps { archiveAgent: vi.fn().mockResolvedValue({ archivedAt: new Date().toISOString() }), notifyAgentState: vi.fn(), getAgent: vi.fn(), + listAgents: vi.fn().mockReturnValue([]), + getTimeline: vi.fn().mockReturnValue([]), + resumeAgentFromPersistence: vi.fn(), + hydrateTimelineFromProvider: vi.fn().mockResolvedValue(undefined), + hasInFlightRun: vi.fn().mockReturnValue(false), + subscribe: vi.fn().mockReturnValue(() => {}), streamAgent: vi.fn(() => (async function* noop() {})()), respondToPermission: vi.fn(), cancelAgentRun: vi.fn(), getPendingPermissions: vi.fn(), + getRegisteredProviderIds: vi.fn().mockReturnValue(["claude"]), }; const agentStorageSpies = { @@ -40,7 +47,7 @@ function createTestDeps(): TestDeps { setTitle: vi.fn().mockResolvedValue(undefined), upsert: vi.fn().mockResolvedValue(undefined), applySnapshot: vi.fn(), - list: vi.fn(), + list: vi.fn().mockResolvedValue([]), remove: vi.fn(), }; @@ -68,6 +75,43 @@ function createProviderDefinition(overrides: Partial): Provi }; } +function createStoredRecord(overrides: Partial = {}): StoredAgentRecord { + const now = "2026-04-11T00:00:00.000Z"; + return { + id: "stored-agent", + provider: "claude", + cwd: "/tmp/stored-project", + createdAt: now, + updatedAt: now, + lastActivityAt: now, + lastUserMessageAt: null, + title: "Stored agent", + labels: {}, + lastStatus: "closed", + lastModeId: "default", + config: { + modeId: "default", + model: "claude-sonnet-4-20250514", + }, + runtimeInfo: { + provider: "claude", + sessionId: "session-123", + model: "claude-sonnet-4-20250514", + }, + features: [], + persistence: { + provider: "claude", + sessionId: "session-123", + }, + requiresAttention: false, + attentionReason: null, + attentionTimestamp: null, + internal: false, + archivedAt: "2026-04-12T00:00:00.000Z", + ...overrides, + }; +} + describe("create_agent MCP tool", () => { const logger = createTestLogger(); const existingCwd = process.cwd(); @@ -483,4 +527,223 @@ describe("agent snapshot MCP serialization", () => { }); expect(Array.isArray(structured.agents[0].features)).toBe(true); }); + + it("returns archived agent snapshots from storage for get_agent_status", async () => { + const { agentManager, agentStorage, spies } = createTestDeps(); + const record = createStoredRecord({ + id: "archived-agent", + archivedAt: "2026-04-12T00:00:00.000Z", + }); + spies.agentManager.getAgent.mockReturnValue(null); + spies.agentStorage.get.mockResolvedValue(record); + + const server = await createAgentMcpServer({ + agentManager, + agentStorage, + logger, + providerRegistry: { + claude: createProviderDefinition({}), + } as any, + }); + const tool = (server as any)._registeredTools["get_agent_status"]; + const response = await tool.callback({ agentId: "archived-agent" }); + + expect(response.structuredContent).toEqual({ + status: "closed", + snapshot: expect.objectContaining({ + id: "archived-agent", + archivedAt: "2026-04-12T00:00:00.000Z", + title: "Stored agent", + status: "closed", + }), + }); + expect(spies.agentStorage.get).toHaveBeenCalledWith("archived-agent"); + }); + + it("does not expose internal stored agents from get_agent_status", async () => { + const { agentManager, agentStorage, spies } = createTestDeps(); + spies.agentManager.getAgent.mockReturnValue(null); + spies.agentStorage.get.mockResolvedValue( + createStoredRecord({ + id: "internal-agent", + internal: true, + }), + ); + + const server = await createAgentMcpServer({ + agentManager, + agentStorage, + logger, + providerRegistry: { + claude: createProviderDefinition({}), + } as any, + }); + const tool = (server as any)._registeredTools["get_agent_status"]; + + await expect(tool.callback({ agentId: "internal-agent" })).rejects.toThrow( + "Agent internal-agent not found", + ); + }); + + it("includes stored non-archived agents in list_agents by default", async () => { + const { agentManager, agentStorage, spies } = createTestDeps(); + const liveAgent = { + id: "live-agent", + provider: "claude", + cwd: "/tmp/live-project", + config: {}, + runtimeInfo: undefined, + createdAt: new Date("2026-04-11T00:00:00.000Z"), + updatedAt: new Date("2026-04-11T00:00:00.000Z"), + lastUserMessageAt: null, + lifecycle: "idle", + capabilities: { + supportsStreaming: false, + supportsSessionPersistence: false, + supportsDynamicModes: false, + supportsMcpServers: true, + supportsReasoningStream: false, + supportsToolInvocations: true, + }, + currentModeId: null, + availableModes: [], + features: [], + pendingPermissions: new Map(), + persistence: null, + labels: {}, + attention: { requiresAttention: false }, + } as unknown as ManagedAgent; + spies.agentManager.listAgents.mockReturnValue([liveAgent]); + spies.agentStorage.list.mockResolvedValue([ + createStoredRecord({ id: "closed-agent", archivedAt: null }), + createStoredRecord({ id: "archived-agent", archivedAt: "2026-04-12T00:00:00.000Z" }), + createStoredRecord({ id: "live-agent", archivedAt: null }), + createStoredRecord({ id: "internal-agent", archivedAt: null, internal: true }), + ]); + + const server = await createAgentMcpServer({ + agentManager, + agentStorage, + logger, + providerRegistry: { + claude: createProviderDefinition({}), + } as any, + }); + const tool = (server as any)._registeredTools["list_agents"]; + const response = await tool.callback({}); + + expect(response.structuredContent.agents).toEqual([ + expect.objectContaining({ id: "live-agent" }), + expect.objectContaining({ id: "closed-agent", archivedAt: null }), + ]); + }); + + it("includes archived stored agents in list_agents when requested", async () => { + const { agentManager, agentStorage, spies } = createTestDeps(); + const liveAgent = { + id: "live-agent", + provider: "claude", + cwd: "/tmp/live-project", + config: {}, + runtimeInfo: undefined, + createdAt: new Date("2026-04-11T00:00:00.000Z"), + updatedAt: new Date("2026-04-11T00:00:00.000Z"), + lastUserMessageAt: null, + lifecycle: "idle", + capabilities: { + supportsStreaming: false, + supportsSessionPersistence: false, + supportsDynamicModes: false, + supportsMcpServers: true, + supportsReasoningStream: false, + supportsToolInvocations: true, + }, + currentModeId: null, + availableModes: [], + features: [], + pendingPermissions: new Map(), + persistence: null, + labels: {}, + attention: { requiresAttention: false }, + } as unknown as ManagedAgent; + spies.agentManager.listAgents.mockReturnValue([liveAgent]); + spies.agentStorage.list.mockResolvedValue([ + createStoredRecord({ id: "archived-agent", archivedAt: "2026-04-12T00:00:00.000Z" }), + createStoredRecord({ id: "live-agent", archivedAt: "2026-04-12T00:00:00.000Z" }), + createStoredRecord({ + id: "internal-archived-agent", + archivedAt: "2026-04-12T00:00:00.000Z", + internal: true, + }), + createStoredRecord({ id: "not-archived-agent", archivedAt: null }), + ]); + + const server = await createAgentMcpServer({ + agentManager, + agentStorage, + logger, + providerRegistry: { + claude: createProviderDefinition({}), + } as any, + }); + const tool = (server as any)._registeredTools["list_agents"]; + const response = await tool.callback({ includeArchived: true }); + + expect(response.structuredContent.agents).toEqual([ + expect.objectContaining({ id: "live-agent" }), + expect.objectContaining({ + id: "archived-agent", + archivedAt: "2026-04-12T00:00:00.000Z", + }), + expect.objectContaining({ + id: "not-archived-agent", + archivedAt: null, + }), + ]); + }); + + it("loads archived agents before reading get_agent_activity", async () => { + const { agentManager, agentStorage, spies } = createTestDeps(); + const record = createStoredRecord({ id: "archived-activity-agent" }); + const snapshot = { + id: "archived-activity-agent", + currentModeId: "default", + } as ManagedAgent; + spies.agentManager.getAgent + .mockReturnValueOnce(null) + .mockReturnValue(snapshot) + .mockReturnValue(snapshot); + spies.agentStorage.get.mockResolvedValue(record); + spies.agentManager.resumeAgentFromPersistence.mockResolvedValue(snapshot); + spies.agentManager.getTimeline.mockReturnValue([ + { + kind: "status", + timestamp: "2026-04-11T00:00:00.000Z", + text: "Agent resumed", + }, + ]); + + const server = await createAgentMcpServer({ + agentManager, + agentStorage, + logger, + providerRegistry: { + claude: createProviderDefinition({}), + } as any, + }); + const tool = (server as any)._registeredTools["get_agent_activity"]; + const response = await tool.callback({ agentId: "archived-activity-agent" }); + + expect(response.structuredContent).toEqual( + expect.objectContaining({ + agentId: "archived-activity-agent", + updateCount: 1, + currentModeId: "default", + }), + ); + expect(spies.agentManager.resumeAgentFromPersistence).toHaveBeenCalled(); + expect(spies.agentManager.hydrateTimelineFromProvider).toHaveBeenCalledWith( + "archived-activity-agent", + ); + }); }); diff --git a/packages/server/src/server/agent/mcp-server.ts b/packages/server/src/server/agent/mcp-server.ts index d269bd301..a71134b25 100644 --- a/packages/server/src/server/agent/mcp-server.ts +++ b/packages/server/src/server/agent/mcp-server.ts @@ -12,9 +12,10 @@ import { AgentPermissionResponseSchema, AgentSnapshotPayloadSchema, } from "../messages.js"; -import { toAgentPayload } from "./agent-projections.js"; +import { buildStoredAgentPayload, toAgentPayload } from "./agent-projections.js"; import { curateAgentActivity } from "./activity-curator.js"; import { AgentStorage } from "./agent-storage.js"; +import { ensureAgentLoaded } from "./agent-loading.js"; import { appendTimelineItemIfAgentKnown, emitLiveTimelineItemIfAgentKnown, @@ -198,6 +199,13 @@ export async function createAgentMcpServer(options: AgentMcpServerOptions): Prom version: "2.0.0", }); + const requireProviderRegistry = (): Record => { + if (!providerRegistry) { + throw new Error("Provider registry is required to load stored agent records"); + } + return providerRegistry; + }; + const resolveCallerAgent = () => { if (!callerAgentId) { return null; @@ -594,6 +602,7 @@ export async function createAgentMcpServer(options: AgentMcpServerOptions): Prom if (notifyOnFinish && callerAgentId) { setupFinishNotification({ agentManager, + agentStorage, childAgentId: snapshot.id, callerAgentId, logger: childLogger, @@ -786,6 +795,7 @@ export async function createAgentMcpServer(options: AgentMcpServerOptions): Prom if (notifyOnFinish && callerAgentId) { setupFinishNotification({ agentManager, + agentStorage, childAgentId: agentId, callerAgentId, logger: childLogger, @@ -849,19 +859,35 @@ export async function createAgentMcpServer(options: AgentMcpServerOptions): Prom }, async ({ agentId }) => { const snapshot = agentManager.getAgent(agentId); - if (!snapshot) { + if (snapshot) { + const structuredSnapshot = await serializeSnapshotWithMetadata( + agentStorage, + snapshot, + childLogger, + ); + return { + content: [], + structuredContent: ensureValidJson({ + status: snapshot.lifecycle, + snapshot: structuredSnapshot, + }), + }; + } + + const record = await agentStorage.get(agentId); + if (!record || record.internal) { throw new Error(`Agent ${agentId} not found`); } - const structuredSnapshot = await serializeSnapshotWithMetadata( - agentStorage, - snapshot, + const structuredSnapshot = buildStoredAgentPayload( + record, + requireProviderRegistry(), childLogger, ); return { content: [], structuredContent: ensureValidJson({ - status: snapshot.lifecycle, + status: structuredSnapshot.status, snapshot: structuredSnapshot, }), }; @@ -873,21 +899,30 @@ export async function createAgentMcpServer(options: AgentMcpServerOptions): Prom { title: "List agents", description: "List all live agents managed by the server.", - inputSchema: {}, + inputSchema: { + includeArchived: z.boolean().optional().default(false), + }, outputSchema: { agents: z.array(AgentSnapshotPayloadSchema), }, }, - async () => { - const snapshots = agentManager.listAgents(); - const agents = await Promise.all( - snapshots.map((snapshot) => + async ({ includeArchived }) => { + const liveSnapshots = agentManager.listAgents(); + const liveAgents = await Promise.all( + liveSnapshots.map((snapshot) => serializeSnapshotWithMetadata(agentStorage, snapshot, childLogger), ), ); + const liveIds = new Set(liveSnapshots.map((snapshot) => snapshot.id)); + const storedRecords = await agentStorage.list(); + const storedAgents = storedRecords + .filter((record) => !record.internal && !liveIds.has(record.id)) + .filter((record) => includeArchived || !record.archivedAt) + .map((record) => buildStoredAgentPayload(record, requireProviderRegistry(), childLogger)); + return { content: [], - structuredContent: ensureValidJson({ agents }), + structuredContent: ensureValidJson({ agents: [...liveAgents, ...storedAgents] }), }; }, ); @@ -1562,6 +1597,11 @@ export async function createAgentMcpServer(options: AgentMcpServerOptions): Prom }, }, async ({ agentId, limit }) => { + await ensureAgentLoaded(agentId, { + agentManager, + agentStorage, + logger: childLogger, + }); const timeline = agentManager.getTimeline(agentId); const snapshot = agentManager.getAgent(agentId); diff --git a/packages/server/src/server/agent/mcp-shared.test.ts b/packages/server/src/server/agent/mcp-shared.test.ts new file mode 100644 index 000000000..2494e17a9 --- /dev/null +++ b/packages/server/src/server/agent/mcp-shared.test.ts @@ -0,0 +1,78 @@ +import { describe, expect, it, vi } from "vitest"; + +import { createTestLogger } from "../../test-utils/test-logger.js"; +import { setupFinishNotification } from "./mcp-shared.js"; +import type { AgentManager, AgentManagerEvent, ManagedAgent } from "./agent-manager.js"; +import type { AgentStorage } from "./agent-storage.js"; + +describe("setupFinishNotification", () => { + it("does not notify archived callers", async () => { + let subscriber: ((event: AgentManagerEvent) => void) | null = null; + + const childAgent = { + id: "child-agent", + lifecycle: "idle", + config: { title: "Child Agent" }, + } as ManagedAgent; + + const agentManager = { + getAgent: vi.fn((agentId: string) => { + if (agentId === "child-agent") { + return childAgent; + } + if (agentId === "caller-agent") { + return { + id: "caller-agent", + lifecycle: "idle", + config: { title: "Caller Agent" }, + } as ManagedAgent; + } + return null; + }), + subscribe: vi.fn((callback: (event: AgentManagerEvent) => void) => { + subscriber = callback; + return () => { + subscriber = null; + }; + }), + hasInFlightRun: vi.fn().mockReturnValue(false), + streamAgent: vi.fn(() => (async function* noop() {})()), + replaceAgentRun: vi.fn(() => (async function* noop() {})()), + } as unknown as AgentManager; + + const agentStorage = { + get: vi.fn(async (agentId: string) => + agentId === "caller-agent" ? { archivedAt: "2024-01-01" } : null, + ), + } as unknown as AgentStorage; + + setupFinishNotification({ + agentManager, + agentStorage, + childAgentId: "child-agent", + callerAgentId: "caller-agent", + logger: createTestLogger(), + }); + + expect(subscriber).not.toBeNull(); + + childAgent.lifecycle = "running"; + subscriber?.({ + type: "agent_state", + agent: childAgent, + }); + + childAgent.lifecycle = "idle"; + subscriber?.({ + type: "agent_state", + agent: childAgent, + }); + + await vi.waitFor(() => { + expect(agentStorage.get).toHaveBeenCalledWith("caller-agent"); + }); + + expect((agentManager as any).streamAgent).not.toHaveBeenCalled(); + expect((agentManager as any).replaceAgentRun).not.toHaveBeenCalled(); + }); +}); diff --git a/packages/server/src/server/agent/mcp-shared.ts b/packages/server/src/server/agent/mcp-shared.ts index 6065e5d6e..3edf8f6c7 100644 --- a/packages/server/src/server/agent/mcp-shared.ts +++ b/packages/server/src/server/agent/mcp-shared.ts @@ -187,18 +187,19 @@ export function startAgentRun( interface SetupFinishNotificationParams { agentManager: AgentManager; + agentStorage: AgentStorage; childAgentId: string; callerAgentId: string; logger: Logger; } export function setupFinishNotification(params: SetupFinishNotificationParams): void { - const { agentManager, childAgentId, callerAgentId, logger } = params; + const { agentManager, agentStorage, childAgentId, callerAgentId, logger } = params; let hasSeenRunning = false; let fired = false; let unsubscribe: (() => void) | null = null; - function notify(reason: "finished" | "errored" | "needs permission"): void { + async function notify(reason: "finished" | "errored" | "needs permission"): Promise { if (fired) { return; } @@ -209,6 +210,11 @@ export function setupFinishNotification(params: SetupFinishNotificationParams): return; } + const callerRecord = await agentStorage.get(callerAgentId); + if (callerRecord?.archivedAt) { + return; + } + const title = agentManager.getAgent(childAgentId)?.config?.title ?? childAgentId; const prompt = `\nAgent ${childAgentId} (${title}) ${reason}.\n`; diff --git a/packages/server/src/server/persistence-hooks.ts b/packages/server/src/server/persistence-hooks.ts index c14881b45..4ece48b1e 100644 --- a/packages/server/src/server/persistence-hooks.ts +++ b/packages/server/src/server/persistence-hooks.ts @@ -1,6 +1,11 @@ import type { AgentManager } from "./agent/agent-manager.js"; -import type { AgentProvider, AgentSessionConfig } from "./agent/agent-sdk-types.js"; +import type { + AgentPersistenceHandle, + AgentProvider, + AgentSessionConfig, +} from "./agent/agent-sdk-types.js"; import type { AgentStorage, StoredAgentRecord } from "./agent/agent-storage.js"; +import { buildProviderRegistry } from "./agent/provider-registry.js"; type LoggerLike = { child(bindings: Record): LoggerLike; @@ -8,6 +13,8 @@ type LoggerLike = { warn(...args: any[]): void; }; +const DEFAULT_AGENT_PROVIDER = "claude"; + function getLogger(logger: LoggerLike): LoggerLike { return logger.child({ module: "persistence" }); } @@ -20,6 +27,18 @@ type BuildSessionConfigOptions = { logger?: LoggerLike; }; +type RegisteredProviders = ReturnType | Iterable; + +function isProviderRegistry( + registeredProviders: RegisteredProviders, +): registeredProviders is ReturnType { + return ( + typeof registeredProviders === "object" && + registeredProviders !== null && + !(Symbol.iterator in registeredProviders) + ); +} + /** * Attach AgentStorage persistence to an AgentManager instance so every * agent_state snapshot is flushed to disk. @@ -97,3 +116,58 @@ export function extractTimestamps(record: StoredAgentRecord): { labels: record.labels, }; } + +function hasRegisteredProvider(registeredProviders: RegisteredProviders, value: string): boolean { + if (isProviderRegistry(registeredProviders)) { + return Object.prototype.hasOwnProperty.call(registeredProviders, value); + } + return new Set(registeredProviders).has(value as AgentProvider); +} + +export function isRegisteredProvider( + providerRegistry: ReturnType, + value: string, +): boolean { + return hasRegisteredProvider(providerRegistry, value); +} + +export function coerceAgentProvider( + logger: LoggerLike, + providerRegistry: ReturnType, + value: string, + agentId?: string, +): AgentProvider { + if (isRegisteredProvider(providerRegistry, value)) { + return value; + } + logger.warn( + { value, agentId, defaultProvider: DEFAULT_AGENT_PROVIDER }, + `Unknown provider '${value}' for agent ${agentId ?? "unknown"}; defaulting to '${DEFAULT_AGENT_PROVIDER}'`, + ); + return DEFAULT_AGENT_PROVIDER; +} + +export function toAgentPersistenceHandle( + logger: LoggerLike, + registeredProviders: RegisteredProviders, + handle: StoredAgentRecord["persistence"], +): AgentPersistenceHandle | null { + if (!handle) { + return null; + } + const provider = handle.provider; + if (!hasRegisteredProvider(registeredProviders, provider)) { + logger.warn({ provider }, `Ignoring persistence handle with unknown provider '${provider}'`); + return null; + } + if (!handle.sessionId) { + logger.warn("Ignoring persistence handle missing sessionId"); + return null; + } + return { + provider, + sessionId: handle.sessionId, + nativeHandle: handle.nativeHandle, + metadata: handle.metadata, + } satisfies AgentPersistenceHandle; +} diff --git a/packages/server/src/server/schedule/service.test.ts b/packages/server/src/server/schedule/service.test.ts index 9547d3e7e..0a539667c 100644 --- a/packages/server/src/server/schedule/service.test.ts +++ b/packages/server/src/server/schedule/service.test.ts @@ -265,4 +265,62 @@ describe("ScheduleService", () => { expect(inspected.runs).toHaveLength(1); expect(inspected.runs[0]?.status).toBe("succeeded"); }); + + test("rejects archived target agents before loading them", async () => { + const manager = new AgentManager({ logger: createTestLogger() }); + const service = new ScheduleService({ + paseoHome: tempDir, + logger: createTestLogger(), + agentManager: manager, + agentStorage, + now: () => now, + }); + + await agentStorage.upsert({ + id: "archived-agent", + provider: "claude", + cwd: tempDir, + createdAt: now.toISOString(), + updatedAt: now.toISOString(), + lastActivityAt: now.toISOString(), + lastUserMessageAt: null, + title: "Archived Agent", + labels: {}, + lastStatus: "closed", + lastModeId: "default", + config: { + modeId: "default", + }, + runtimeInfo: null, + features: [], + persistence: null, + requiresAttention: false, + attentionReason: null, + attentionTimestamp: null, + internal: false, + archivedAt: "2026-01-02T00:00:00.000Z", + }); + + await expect( + (service as any).executeSchedule({ + id: "schedule-1", + name: null, + prompt: "Check archived agent", + cadence: { type: "every", everyMs: 60_000 }, + target: { + type: "agent", + agentId: "archived-agent", + }, + status: "active", + createdAt: now.toISOString(), + updatedAt: now.toISOString(), + nextRunAt: now.toISOString(), + lastRunAt: null, + pausedAt: null, + expiresAt: null, + maxRuns: null, + runs: [], + }), + ).rejects.toThrow("Agent archived-agent is archived"); + }); }); diff --git a/packages/server/src/server/schedule/service.ts b/packages/server/src/server/schedule/service.ts index af2dcbe56..b6b1cc557 100644 --- a/packages/server/src/server/schedule/service.ts +++ b/packages/server/src/server/schedule/service.ts @@ -2,15 +2,10 @@ import { randomUUID } from "node:crypto"; import { join } from "node:path"; import type { Logger } from "pino"; import { AgentManager } from "../agent/agent-manager.js"; -import type { ManagedAgent } from "../agent/agent-manager.js"; import { AgentStorage } from "../agent/agent-storage.js"; import type { AgentPromptInput, AgentSessionConfig } from "../agent/agent-sdk-types.js"; import { curateAgentActivity } from "../agent/activity-curator.js"; -import { - buildConfigOverrides, - buildSessionConfig, - extractTimestamps, -} from "../persistence-hooks.js"; +import { ensureAgentLoaded } from "../agent/agent-loading.js"; import { ScheduleStore } from "./store.js"; import { computeNextRunAt, validateScheduleCadence } from "./cron.js"; import type { @@ -21,7 +16,6 @@ import type { } from "./types.js"; const SCHEDULE_TICK_INTERVAL_MS = 1000; -const pendingAgentInitializations = new Map>(); function trimOptionalName(value: string | null | undefined): string | null { if (typeof value !== "string") { @@ -386,7 +380,16 @@ export class ScheduleService { private async executeSchedule(schedule: StoredSchedule): Promise { if (schedule.target.type === "agent") { - const agent = await this.ensureAgentLoaded(schedule.target.agentId); + const record = await this.agentStorage.get(schedule.target.agentId); + if (record?.archivedAt) { + throw new Error(`Agent ${schedule.target.agentId} is archived`); + } + + const agent = await ensureAgentLoaded(schedule.target.agentId, { + agentManager: this.agentManager, + agentStorage: this.agentStorage, + logger: this.logger, + }); if (this.agentManager.hasInFlightRun(agent.id)) { throw new Error(`Agent ${agent.id} already has an active run`); } @@ -435,64 +438,4 @@ export class ScheduleService { }), }; } - - private async ensureAgentLoaded(agentId: string): Promise { - const existing = this.agentManager.getAgent(agentId); - if (existing) { - return existing; - } - - const inflight = pendingAgentInitializations.get(agentId); - if (inflight) { - return inflight; - } - - const initPromise = (async () => { - const record = await this.agentStorage.get(agentId); - if (!record) { - throw new Error(`Agent not found: ${agentId}`); - } - if (record.archivedAt) { - throw new Error(`Agent ${agentId} is archived`); - } - - let snapshot: ManagedAgent; - if (record.persistence?.provider && record.persistence?.sessionId) { - snapshot = await this.agentManager.resumeAgentFromPersistence( - { - provider: record.persistence.provider as AgentSessionConfig["provider"], - sessionId: record.persistence.sessionId, - nativeHandle: record.persistence.nativeHandle, - metadata: record.persistence.metadata, - }, - buildConfigOverrides(record), - agentId, - extractTimestamps(record), - ); - } else { - const config = buildSessionConfig(record, { - validProviders: this.agentManager.getRegisteredProviderIds(), - logger: this.logger, - }); - if (!config) { - throw new Error(`Agent ${agentId} references unavailable provider '${record.provider}'`); - } - snapshot = await this.agentManager.createAgent(config, agentId, { - labels: record.labels, - }); - } - - await this.agentManager.hydrateTimelineFromProvider(agentId); - return this.agentManager.getAgent(agentId) ?? snapshot; - })(); - - pendingAgentInitializations.set(agentId, initPromise); - try { - return await initPromise; - } finally { - if (pendingAgentInitializations.get(agentId) === initPromise) { - pendingAgentInitializations.delete(agentId); - } - } - } } diff --git a/packages/server/src/server/session.ts b/packages/server/src/server/session.ts index e141ceddf..08927cc92 100644 --- a/packages/server/src/server/session.ts +++ b/packages/server/src/server/session.ts @@ -61,9 +61,10 @@ import { } from "./voice/voice-turn-controller.js"; import { buildConfigOverrides, - buildSessionConfig, extractTimestamps, + toAgentPersistenceHandle, } from "./persistence-hooks.js"; +import { ensureAgentLoaded } from "./agent/agent-loading.js"; import { experimental_createMCPClient } from "ai"; import { StreamableHTTPClientTransport } from "@modelcontextprotocol/sdk/client/streamableHttp.js"; import type { VoiceCallerContext, VoiceSpeakHandler } from "./voice-types.js"; @@ -83,7 +84,12 @@ import type { ManagedAgent, } from "./agent/agent-manager.js"; import { scheduleAgentMetadataGeneration } from "./agent/agent-metadata-generator.js"; -import { resolveEffectiveThinkingOptionId, toAgentPayload } from "./agent/agent-projections.js"; +import { + buildStoredAgentPayload, + resolveEffectiveThinkingOptionId, + resolveStoredAgentPayloadUpdatedAt, + toAgentPayload, +} from "./agent/agent-projections.js"; import { MAX_EXPLICIT_AGENT_TITLE_CHARS } from "./agent/agent-title-limits.js"; import { appendTimelineItemIfAgentKnown, @@ -101,14 +107,14 @@ import { generateStructuredAgentResponseWithFallback, } from "./agent/agent-response-loop.js"; import type { + AgentPersistenceHandle, AgentPermissionResponse, + AgentProvider, AgentPromptContentBlock, AgentPromptInput, AgentRunOptions, AgentSessionConfig, AgentStreamEvent, - AgentProvider, - AgentPersistenceHandle, ProviderSnapshotEntry, } from "./agent/agent-sdk-types.js"; import { AgentStorage, type StoredAgentRecord } from "./agent/agent-storage.js"; @@ -187,8 +193,6 @@ import { const execAsync = promisify(exec); const MAX_INITIAL_AGENT_TITLE_CHARS = Math.min(60, MAX_EXPLICIT_AGENT_TITLE_CHARS); -const pendingAgentInitializations = new Map>(); -const DEFAULT_AGENT_PROVIDER = "claude"; // TODO: Remove once all app store clients are on >=0.1.45 and understand arbitrary provider strings. // Clients before 0.1.45 validate providers with z.enum(["claude", "codex", "opencode"]) and reject @@ -527,54 +531,6 @@ function convertPCMToWavBuffer( return wavBuffer; } -function isRegisteredProvider( - providerRegistry: ReturnType, - value: string, -): boolean { - return Object.prototype.hasOwnProperty.call(providerRegistry, value); -} - -function coerceAgentProvider( - logger: pino.Logger, - providerRegistry: ReturnType, - value: string, - agentId?: string, -): AgentProvider { - if (isRegisteredProvider(providerRegistry, value)) { - return value; - } - logger.warn( - { value, agentId, defaultProvider: DEFAULT_AGENT_PROVIDER }, - `Unknown provider '${value}' for agent ${agentId ?? "unknown"}; defaulting to '${DEFAULT_AGENT_PROVIDER}'`, - ); - return DEFAULT_AGENT_PROVIDER; -} - -function toAgentPersistenceHandle( - logger: pino.Logger, - providerRegistry: ReturnType, - handle: StoredAgentRecord["persistence"], -): AgentPersistenceHandle | null { - if (!handle) { - return null; - } - const provider = handle.provider; - if (!isRegisteredProvider(providerRegistry, provider)) { - logger.warn({ provider }, `Ignoring persistence handle with unknown provider '${provider}'`); - return null; - } - if (!handle.sessionId) { - logger.warn("Ignoring persistence handle missing sessionId"); - return null; - } - return { - provider, - sessionId: handle.sessionId, - nativeHandle: handle.nativeHandle, - metadata: handle.metadata, - } satisfies AgentPersistenceHandle; -} - /** * Session represents a single connected client session. * It owns all state management, orchestration logic, and message processing. @@ -1097,9 +1053,7 @@ export class Session { const storedRecord = await this.agentStorage.get(agent.id); const title = storedRecord?.title ?? storedRecord?.config?.title ?? null; const payload = toAgentPayload(agent, { title }); - const storedUpdatedAt = storedRecord - ? this.resolveStoredAgentPayloadUpdatedAt(storedRecord) - : null; + const storedUpdatedAt = storedRecord ? resolveStoredAgentPayloadUpdatedAt(storedRecord) : null; if (storedUpdatedAt) { const liveUpdatedAt = Date.parse(payload.updatedAt); const persistedUpdatedAt = Date.parse(storedUpdatedAt); @@ -1115,161 +1069,7 @@ export class Session { } private buildStoredAgentPayload(record: StoredAgentRecord): AgentSnapshotPayload { - const defaultCapabilities = { - supportsStreaming: false, - supportsSessionPersistence: true, - supportsDynamicModes: false, - supportsMcpServers: false, - supportsReasoningStream: false, - supportsToolInvocations: true, - } as const; - - const createdAt = new Date(record.createdAt); - const updatedAt = new Date(this.resolveStoredAgentPayloadUpdatedAt(record)); - const lastUserMessageAt = record.lastUserMessageAt ? new Date(record.lastUserMessageAt) : null; - - const provider = coerceAgentProvider( - this.sessionLogger, - this.providerRegistry, - record.provider, - record.id, - ); - const runtimeInfo = record.runtimeInfo - ? { - provider: coerceAgentProvider( - this.sessionLogger, - this.providerRegistry, - record.runtimeInfo.provider, - record.id, - ), - sessionId: record.runtimeInfo.sessionId, - ...(Object.prototype.hasOwnProperty.call(record.runtimeInfo, "model") - ? { model: record.runtimeInfo.model ?? null } - : {}), - ...(Object.prototype.hasOwnProperty.call(record.runtimeInfo, "thinkingOptionId") - ? { thinkingOptionId: record.runtimeInfo.thinkingOptionId ?? null } - : {}), - ...(Object.prototype.hasOwnProperty.call(record.runtimeInfo, "modeId") - ? { modeId: record.runtimeInfo.modeId ?? null } - : {}), - ...(record.runtimeInfo.extra ? { extra: record.runtimeInfo.extra } : {}), - } - : undefined; - return { - id: record.id, - provider, - cwd: record.cwd, - model: record.config?.model ?? null, - thinkingOptionId: record.config?.thinkingOptionId ?? null, - effectiveThinkingOptionId: resolveEffectiveThinkingOptionId({ - runtimeInfo, - configuredThinkingOptionId: record.config?.thinkingOptionId ?? null, - }), - ...(runtimeInfo ? { runtimeInfo } : {}), - createdAt: createdAt.toISOString(), - updatedAt: updatedAt.toISOString(), - lastUserMessageAt: lastUserMessageAt ? lastUserMessageAt.toISOString() : null, - status: record.lastStatus, - capabilities: defaultCapabilities, - currentModeId: record.lastModeId ?? null, - availableModes: [], - pendingPermissions: [], - persistence: toAgentPersistenceHandle( - this.sessionLogger, - this.providerRegistry, - record.persistence, - ), - lastUsage: undefined, - lastError: undefined, - title: record.title ?? record.config?.title ?? null, - requiresAttention: record.requiresAttention ?? false, - attentionReason: record.attentionReason ?? null, - attentionTimestamp: record.attentionTimestamp ?? null, - archivedAt: record.archivedAt ?? null, - labels: record.labels, - }; - } - - private resolveStoredAgentPayloadUpdatedAt(record: StoredAgentRecord): string { - const timestamps = [record.updatedAt, record.lastActivityAt] - .filter((value): value is string => typeof value === "string" && value.length > 0) - .map((value) => ({ - raw: value, - parsed: Date.parse(value), - })) - .filter((value) => !Number.isNaN(value.parsed)); - - if (timestamps.length === 0) { - return record.updatedAt; - } - - timestamps.sort((a, b) => b.parsed - a.parsed); - return timestamps[0].raw; - } - - private async ensureAgentLoaded(agentId: string): Promise { - const existing = this.agentManager.getAgent(agentId); - if (existing) { - return existing; - } - - const inflight = pendingAgentInitializations.get(agentId); - if (inflight) { - return inflight; - } - - const initPromise = (async () => { - const record = await this.agentStorage.get(agentId); - if (!record) { - throw new Error(`Agent not found: ${agentId}`); - } - - const handle = toAgentPersistenceHandle( - this.sessionLogger, - this.providerRegistry, - record.persistence, - ); - let snapshot: ManagedAgent; - if (handle) { - snapshot = await this.agentManager.resumeAgentFromPersistence( - handle, - buildConfigOverrides(record), - agentId, - extractTimestamps(record), - ); - this.sessionLogger.info( - { agentId, provider: record.provider }, - "Agent resumed from persistence", - ); - } else { - const config = buildSessionConfig(record, { - validProviders: Object.keys(this.providerRegistry), - logger: this.sessionLogger, - }); - if (!config) { - throw new Error(`Agent ${agentId} references unavailable provider '${record.provider}'`); - } - snapshot = await this.agentManager.createAgent(config, agentId, { labels: record.labels }); - this.sessionLogger.info( - { agentId, provider: record.provider }, - "Agent created from stored config", - ); - } - - await this.agentManager.hydrateTimelineFromProvider(agentId); - return this.agentManager.getAgent(agentId) ?? snapshot; - })(); - - pendingAgentInitializations.set(agentId, initPromise); - - try { - return await initPromise; - } finally { - const current = pendingAgentInitializations.get(agentId); - if (current === initPromise) { - pendingAgentInitializations.delete(agentId); - } - } + return buildStoredAgentPayload(record, this.providerRegistry, this.sessionLogger); } // TODO: Remove once all app store clients are on >=0.1.45. @@ -2698,7 +2498,11 @@ export class Session { private async enableVoiceModeForAgent(agentId: string): Promise { const startedAt = Date.now(); this.sessionLogger.info({ agentId }, "enableVoiceModeForAgent.ensureAgentLoaded.start"); - const existing = await this.ensureAgentLoaded(agentId); + const existing = await ensureAgentLoaded(agentId, { + agentManager: this.agentManager, + agentStorage: this.agentStorage, + logger: this.sessionLogger, + }); this.sessionLogger.info( { agentId, elapsedMs: Date.now() - startedAt }, "enableVoiceModeForAgent.ensureAgentLoaded.done", @@ -2910,7 +2714,11 @@ export class Session { await this.unarchiveAgentState(agentId); try { - await this.ensureAgentLoaded(agentId); + await ensureAgentLoaded(agentId, { + agentManager: this.agentManager, + agentStorage: this.agentStorage, + logger: this.sessionLogger, + }); } catch (error) { this.handleAgentRunError(agentId, error, "Failed to initialize agent before sending prompt"); return { @@ -6428,7 +6236,11 @@ export class Session { : undefined; try { - const snapshot = await this.ensureAgentLoaded(msg.agentId); + const snapshot = await ensureAgentLoaded(msg.agentId, { + agentManager: this.agentManager, + agentStorage: this.agentStorage, + logger: this.sessionLogger, + }); const agentPayload = await this.buildAgentPayload(snapshot); let timeline = this.agentManager.fetchTimeline(msg.agentId, { @@ -6583,7 +6395,11 @@ export class Session { const agentId = resolved.agentId; await this.unarchiveAgentState(agentId); - await this.ensureAgentLoaded(agentId); + await ensureAgentLoaded(agentId, { + agentManager: this.agentManager, + agentStorage: this.agentStorage, + logger: this.sessionLogger, + }); this.sessionLogger.trace( { agentId, messageId: msg.messageId, textPrefix: msg.text.slice(0, 80) },