From f392d4e9b067204c81efea9e8875b413489a94d6 Mon Sep 17 00:00:00 2001 From: vmelikyan Date: Sat, 25 Apr 2026 23:57:25 -0700 Subject: [PATCH] add canonical agent-session control plane --- Tiltfile | 2 + .../[threadId]/conversation/route.test.ts | 117 ++ .../threads/[threadId]/conversation/route.ts | 4 +- .../canonical-api-acceptance.test.ts | 760 +++++++++++ src/app/api/v2/ai/agent/models/route.ts | 3 - .../[actionId]/respond/route.test.ts | 202 +++ .../[actionId]/respond/route.ts | 62 +- .../agent/runs/[runId]/cancel/route.test.ts | 104 ++ .../agent/runs/[runId]/events/route.test.ts | 196 +++ .../v2/ai/agent/runs/[runId]/events/route.ts | 190 +++ .../runs/[runId]/events/stream/route.test.ts | 147 +++ .../agent/runs/[runId]/events/stream/route.ts | 126 ++ .../agent/runs/[runId]/stream/route.test.ts | 297 ----- .../v2/ai/agent/runs/[runId]/stream/route.ts | 363 ----- .../api/v2/ai/agent/sandbox-sessions/route.ts | 54 +- .../v2/ai/agent/sessions/[sessionId]/route.ts | 153 +-- .../[sessionId]/sandbox/resume}/route.ts | 43 +- .../[sessionId]/sandbox/suspend}/route.ts | 44 +- src/app/api/v2/ai/agent/sessions/route.ts | 590 ++++----- .../threads/[threadId]/messages/route.test.ts | 268 +--- .../threads/[threadId]/messages/route.ts | 402 +----- .../[threadId]/pending-actions/route.test.ts | 131 ++ .../[threadId]/pending-actions/route.ts | 35 +- .../threads/[threadId]/runs/route.test.ts | 505 +++++++ .../ai/agent/threads/[threadId]/runs/route.ts | 385 ++++++ .../api/v2/ai/config/agent-session/route.ts | 12 + .../ai/config/agent-session/runtime/route.ts | 12 + src/server/database.ts | 2 +- .../db/migrations/019_remove_scale_to_zero.ts | 64 +- ...21_agent_session_control_plane_contract.ts | 565 ++++++++ src/server/db/seeds/.gitkeep | 1 + .../agentRunDispatchRecovery.test.ts | 100 ++ .../jobs/__tests__/agentRunExecute.test.ts | 187 +++ .../__tests__/agentSessionCleanup.test.ts | 178 ++- src/server/jobs/agentRunDispatchRecovery.ts | 62 + src/server/jobs/agentRunExecute.ts | 112 ++ src/server/jobs/agentSandboxSessionLaunch.ts | 10 +- src/server/jobs/agentSessionCleanup.ts | 44 +- src/server/jobs/index.ts | 57 +- src/server/lib/__mocks__/redisClientMock.ts | 1 + .../agentSession/__tests__/pvcFactory.test.ts | 18 +- .../__tests__/runtimeConfig.test.ts | 133 +- .../__tests__/systemPrompt.test.ts | 4 +- .../lib/agentSession/chatPreviewFactory.ts | 226 ++++ src/server/lib/agentSession/configSeeder.ts | 13 +- src/server/lib/agentSession/pvcFactory.ts | 40 +- src/server/lib/agentSession/runtimeConfig.ts | 185 ++- src/server/lib/createApiHandler.test.ts | 55 + src/server/lib/createApiHandler.ts | 8 + src/server/lib/createStreamHandler.test.ts | 42 + src/server/lib/createStreamHandler.ts | 8 + src/server/lib/kubernetes.ts | 4 +- src/server/lib/response.ts | 2 + .../validation/agentSessionConfigSchemas.ts | 43 + .../validation/agentSessionConfigValidator.ts | 59 + src/server/middlewares/cors.test.ts | 50 + src/server/middlewares/cors.ts | 3 +- src/server/models/AgentMessage.ts | 17 +- src/server/models/AgentRun.ts | 56 +- src/server/models/AgentRunEvent.ts | 64 + src/server/models/AgentSandbox.ts | 88 ++ src/server/models/AgentSandboxExposure.ts | 74 ++ src/server/models/AgentSession.ts | 67 +- src/server/models/AgentSource.ts | 74 ++ .../__tests__/AgentModelsValidation.test.ts | 8 +- src/server/models/index.ts | 12 + .../services/__tests__/agentSession.test.ts | 398 +++++- .../__tests__/agentSessionConfig.test.ts | 80 +- .../services/__tests__/globalConfig.test.ts | 27 + src/server/services/agent/AdminService.ts | 61 +- .../agent/AgentRunOwnershipLostError.ts | 47 + src/server/services/agent/ApprovalService.ts | 656 +++++++-- .../services/agent/CapabilityService.ts | 681 +++++++++- .../services/agent/ChatSessionService.ts | 96 ++ .../services/agent/LifecycleAiSdkHarness.ts | 795 +++++++++++ src/server/services/agent/MessageStore.ts | 488 ++++++- src/server/services/agent/ProviderRegistry.ts | 55 +- .../services/agent/RunAdmissionService.ts | 163 +++ src/server/services/agent/RunEventService.ts | 1171 +++++++++++++++++ src/server/services/agent/RunExecutor.ts | 375 ++++-- src/server/services/agent/RunQueueService.ts | 82 ++ src/server/services/agent/RunService.ts | 778 ++++++++++- src/server/services/agent/SandboxService.ts | 289 ++++ .../services/agent/SessionReadService.ts | 265 ++++ src/server/services/agent/SourceService.ts | 147 +++ src/server/services/agent/StreamBroker.ts | 205 --- .../agent/__tests__/AdminService.test.ts | 302 +++++ .../agent/__tests__/ApprovalService.test.ts | 848 +++++++++++- .../agent/__tests__/CapabilityService.test.ts | 543 +++++++- .../__tests__/LifecycleAiSdkHarness.test.ts | 325 +++++ .../agent/__tests__/MessageStore.test.ts | 354 ++++- .../agent/__tests__/PolicyService.test.ts | 7 + .../agent/__tests__/ProviderRegistry.test.ts | 73 +- .../__tests__/RunAdmissionService.test.ts | 221 ++++ .../agent/__tests__/RunEventService.test.ts | 874 ++++++++++++ .../agent/__tests__/RunExecutor.test.ts | 273 +++- .../agent/__tests__/RunQueueService.test.ts | 146 ++ .../agent/__tests__/RunService.test.ts | 508 +++++++ .../__tests__/SessionReadService.test.ts | 234 ++++ .../agent/__tests__/SourceService.test.ts | 140 ++ .../agent/__tests__/sandboxExecSafety.test.ts | 29 +- .../__tests__/sandboxToolCatalog.test.ts | 2 +- ...reamState.test.ts => streamChunks.test.ts} | 24 +- .../services/agent/canonicalMessages.ts | 231 ++++ src/server/services/agent/fileChanges.ts | 23 +- src/server/services/agent/payloadLimits.ts | 69 + src/server/services/agent/providerConfig.ts | 2 - .../services/agent/sandboxExecSafety.ts | 64 + .../services/agent/sandboxToolCatalog.ts | 2 +- .../services/agent/serializeSessionSummary.ts | 8 +- .../agent/{streamState.ts => streamChunks.ts} | 25 - src/server/services/agent/toolKeys.ts | 7 +- src/server/services/agent/types.ts | 2 + src/server/services/agentPrewarm.ts | 8 +- src/server/services/agentSandboxSession.ts | 9 +- src/server/services/agentSession.ts | 649 +++++++-- src/server/services/agentSessionConfig.ts | 162 ++- src/server/services/ai/mcp/config.ts | 14 +- src/server/services/globalConfig.ts | 6 + .../services/types/agentSessionConfig.ts | 28 + src/server/services/types/globalConfig.ts | 29 + src/shared/config.ts | 2 + src/shared/constants.ts | 21 + src/shared/openApiSpec.test.ts | 68 + src/shared/openApiSpec.ts | 1051 ++++++++++++++- sysops/workspace-gateway/index.mjs | 312 ++++- 126 files changed, 19640 insertions(+), 2857 deletions(-) create mode 100644 src/app/api/v2/ai/admin/agent/threads/[threadId]/conversation/route.test.ts create mode 100644 src/app/api/v2/ai/agent/__tests__/canonical-api-acceptance.test.ts create mode 100644 src/app/api/v2/ai/agent/pending-actions/[actionId]/respond/route.test.ts create mode 100644 src/app/api/v2/ai/agent/runs/[runId]/cancel/route.test.ts create mode 100644 src/app/api/v2/ai/agent/runs/[runId]/events/route.test.ts create mode 100644 src/app/api/v2/ai/agent/runs/[runId]/events/route.ts create mode 100644 src/app/api/v2/ai/agent/runs/[runId]/events/stream/route.test.ts create mode 100644 src/app/api/v2/ai/agent/runs/[runId]/events/stream/route.ts delete mode 100644 src/app/api/v2/ai/agent/runs/[runId]/stream/route.test.ts delete mode 100644 src/app/api/v2/ai/agent/runs/[runId]/stream/route.ts rename src/app/api/v2/ai/agent/{pending-actions/[actionId]/approve => sessions/[sessionId]/sandbox/resume}/route.ts (59%) rename src/app/api/v2/ai/agent/{pending-actions/[actionId]/deny => sessions/[sessionId]/sandbox/suspend}/route.ts (58%) create mode 100644 src/app/api/v2/ai/agent/threads/[threadId]/pending-actions/route.test.ts create mode 100644 src/app/api/v2/ai/agent/threads/[threadId]/runs/route.test.ts create mode 100644 src/app/api/v2/ai/agent/threads/[threadId]/runs/route.ts create mode 100644 src/server/db/migrations/021_agent_session_control_plane_contract.ts create mode 100644 src/server/db/seeds/.gitkeep create mode 100644 src/server/jobs/__tests__/agentRunDispatchRecovery.test.ts create mode 100644 src/server/jobs/__tests__/agentRunExecute.test.ts create mode 100644 src/server/jobs/agentRunDispatchRecovery.ts create mode 100644 src/server/jobs/agentRunExecute.ts create mode 100644 src/server/lib/agentSession/chatPreviewFactory.ts create mode 100644 src/server/lib/createApiHandler.test.ts create mode 100644 src/server/lib/createStreamHandler.test.ts create mode 100644 src/server/middlewares/cors.test.ts create mode 100644 src/server/models/AgentRunEvent.ts create mode 100644 src/server/models/AgentSandbox.ts create mode 100644 src/server/models/AgentSandboxExposure.ts create mode 100644 src/server/models/AgentSource.ts create mode 100644 src/server/services/agent/AgentRunOwnershipLostError.ts create mode 100644 src/server/services/agent/ChatSessionService.ts create mode 100644 src/server/services/agent/LifecycleAiSdkHarness.ts create mode 100644 src/server/services/agent/RunAdmissionService.ts create mode 100644 src/server/services/agent/RunEventService.ts create mode 100644 src/server/services/agent/RunQueueService.ts create mode 100644 src/server/services/agent/SandboxService.ts create mode 100644 src/server/services/agent/SessionReadService.ts create mode 100644 src/server/services/agent/SourceService.ts delete mode 100644 src/server/services/agent/StreamBroker.ts create mode 100644 src/server/services/agent/__tests__/LifecycleAiSdkHarness.test.ts create mode 100644 src/server/services/agent/__tests__/RunAdmissionService.test.ts create mode 100644 src/server/services/agent/__tests__/RunEventService.test.ts create mode 100644 src/server/services/agent/__tests__/RunQueueService.test.ts create mode 100644 src/server/services/agent/__tests__/SessionReadService.test.ts create mode 100644 src/server/services/agent/__tests__/SourceService.test.ts rename src/server/services/agent/__tests__/{streamState.test.ts => streamChunks.test.ts} (76%) create mode 100644 src/server/services/agent/canonicalMessages.ts create mode 100644 src/server/services/agent/payloadLimits.ts rename src/server/services/agent/{streamState.ts => streamChunks.ts} (82%) create mode 100644 src/shared/openApiSpec.test.ts diff --git a/Tiltfile b/Tiltfile index a189cfe8..eff9a358 100644 --- a/Tiltfile +++ b/Tiltfile @@ -20,6 +20,8 @@ load("ext://restart_process", "docker_build_with_restart") load("ext://secret", "secret_create_generic") load('ext://dotenv', 'dotenv') +update_settings(k8s_upsert_timeout_secs=180) + # Load .env file if it exists dotenv() diff --git a/src/app/api/v2/ai/admin/agent/threads/[threadId]/conversation/route.test.ts b/src/app/api/v2/ai/admin/agent/threads/[threadId]/conversation/route.test.ts new file mode 100644 index 00000000..699bba3f --- /dev/null +++ b/src/app/api/v2/ai/admin/agent/threads/[threadId]/conversation/route.test.ts @@ -0,0 +1,117 @@ +/** + * Copyright 2026 GoodRx, Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +import { NextRequest } from 'next/server'; + +jest.mock('server/lib/get-user', () => ({ + getRequestUserIdentity: jest.fn(), +})); + +jest.mock('server/services/agent/AdminService', () => ({ + __esModule: true, + default: { + getThreadConversation: jest.fn(), + }, +})); + +import { GET } from './route'; +import { getRequestUserIdentity } from 'server/lib/get-user'; +import AgentAdminService from 'server/services/agent/AdminService'; + +const mockGetRequestUserIdentity = getRequestUserIdentity as jest.Mock; +const mockGetThreadConversation = AgentAdminService.getThreadConversation as jest.Mock; + +function makeRequest(): NextRequest { + return { + headers: new Headers([['x-request-id', 'req-test']]), + nextUrl: new URL('http://localhost/api/v2/ai/admin/agent/threads/thread-1/conversation'), + } as unknown as NextRequest; +} + +describe('GET /api/v2/ai/admin/agent/threads/[threadId]/conversation', () => { + beforeEach(() => { + jest.clearAllMocks(); + }); + + it('returns 401 when the requester is not authenticated', async () => { + mockGetRequestUserIdentity.mockReturnValue(null); + + const response = await GET(makeRequest(), { params: { threadId: 'thread-1' } }); + + expect(response.status).toBe(401); + await expect(response.json()).resolves.toMatchObject({ + error: { message: 'Unauthorized' }, + }); + expect(mockGetThreadConversation).not.toHaveBeenCalled(); + }); + + it('returns the canonical admin replay payload from the service', async () => { + mockGetRequestUserIdentity.mockReturnValue({ + userId: 'sample-admin', + githubUsername: 'sample-admin', + }); + mockGetThreadConversation.mockResolvedValue({ + session: { id: 'session-1' }, + thread: { id: 'thread-1' }, + messages: [ + { + id: 'message-1', + role: 'user', + parts: [{ type: 'text', text: 'Hi' }], + }, + ], + runs: [], + events: [], + pendingActions: [], + toolExecutions: [ + { + id: 'tool-1', + toolCallId: 'tool-call-1', + }, + ], + }); + + const response = await GET(makeRequest(), { params: { threadId: 'thread-1' } }); + const body = await response.json(); + + expect(response.status).toBe(200); + expect(mockGetThreadConversation).toHaveBeenCalledWith('thread-1'); + expect(body.data.messages[0]).toEqual({ + id: 'message-1', + role: 'user', + parts: [{ type: 'text', text: 'Hi' }], + }); + expect(body.data.toolExecutions[0]).toEqual({ + id: 'tool-1', + toolCallId: 'tool-call-1', + }); + }); + + it.each(['Agent thread not found', 'Agent session not found'])('returns 404 for %s', async (message) => { + mockGetRequestUserIdentity.mockReturnValue({ + userId: 'sample-admin', + githubUsername: 'sample-admin', + }); + mockGetThreadConversation.mockRejectedValue(new Error(message)); + + const response = await GET(makeRequest(), { params: { threadId: 'missing-thread' } }); + + expect(response.status).toBe(404); + await expect(response.json()).resolves.toMatchObject({ + error: { message }, + }); + }); +}); diff --git a/src/app/api/v2/ai/admin/agent/threads/[threadId]/conversation/route.ts b/src/app/api/v2/ai/admin/agent/threads/[threadId]/conversation/route.ts index f42e9768..7c2edcda 100644 --- a/src/app/api/v2/ai/admin/agent/threads/[threadId]/conversation/route.ts +++ b/src/app/api/v2/ai/admin/agent/threads/[threadId]/conversation/route.ts @@ -26,8 +26,8 @@ import AgentAdminService from 'server/services/agent/AdminService'; * get: * summary: Get full agent thread conversation for admin review * description: > - * Returns the persisted UI messages for a thread together with runs, - * pending actions, and tool executions so admins can replay the session. + * Returns the canonical messages for a thread together with runs, run + * events, pending actions, and tool executions so admins can replay the session. * tags: * - Agent Admin * operationId: getAdminAgentThreadConversation diff --git a/src/app/api/v2/ai/agent/__tests__/canonical-api-acceptance.test.ts b/src/app/api/v2/ai/agent/__tests__/canonical-api-acceptance.test.ts new file mode 100644 index 00000000..5f24c4d8 --- /dev/null +++ b/src/app/api/v2/ai/agent/__tests__/canonical-api-acceptance.test.ts @@ -0,0 +1,760 @@ +/** + * Copyright 2026 GoodRx, Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +import { NextRequest } from 'next/server'; + +jest.mock('server/lib/get-user', () => ({ + getRequestUserIdentity: jest.fn(), +})); + +jest.mock('server/lib/agentSession/githubToken', () => ({ + resolveRequestGitHubToken: jest.fn(), +})); + +jest.mock('server/lib/agentSession/runtimeConfig', () => ({ + AgentSessionRuntimeConfigError: class AgentSessionRuntimeConfigError extends Error {}, + AgentSessionWorkspaceStorageConfigError: class AgentSessionWorkspaceStorageConfigError extends Error {}, + resolveAgentSessionRuntimeConfig: jest.fn().mockResolvedValue({ + workspaceStorage: { + defaultSize: '10Gi', + allowedSizes: ['10Gi'], + allowClientOverride: false, + accessMode: 'ReadWriteOnce', + }, + }), + resolveAgentSessionWorkspaceStorageIntent: jest.fn(() => ({ + requestedSize: null, + storageSize: '10Gi', + accessMode: 'ReadWriteOnce', + })), +})); + +jest.mock('server/services/agent/ChatSessionService', () => ({ + __esModule: true, + default: { + createChatSession: jest.fn(), + }, +})); + +jest.mock('server/services/agent/SessionReadService', () => ({ + __esModule: true, + DEFAULT_AGENT_SESSION_LIST_LIMIT: 25, + MAX_AGENT_SESSION_LIST_LIMIT: 100, + default: { + listOwnedSessionRecords: jest.fn(), + serializeSessionRecord: jest.fn(), + }, +})); + +jest.mock('server/services/agent/ThreadService', () => ({ + __esModule: true, + default: { + getOwnedThreadWithSession: jest.fn(), + }, +})); + +jest.mock('server/services/agentSession', () => ({ + __esModule: true, + default: { + canAcceptMessages: jest.fn(), + getMessageBlockReason: jest.fn(), + touchActivity: jest.fn(), + }, + ActiveEnvironmentSessionError: class ActiveEnvironmentSessionError extends Error {}, +})); + +jest.mock('server/services/agent/SourceService', () => ({ + __esModule: true, + default: { + getSessionSource: jest.fn(), + }, +})); + +jest.mock('server/services/agent/CapabilityService', () => ({ + __esModule: true, + default: { + resolveSessionContext: jest.fn(), + }, +})); + +jest.mock('server/services/agent/ProviderRegistry', () => ({ + __esModule: true, + MissingAgentProviderApiKeyError: class MissingAgentProviderApiKeyError extends Error {}, + default: { + resolveSelection: jest.fn(), + }, +})); + +jest.mock('server/services/agent/RunAdmissionService', () => ({ + __esModule: true, + default: { + createQueuedRunWithMessage: jest.fn(), + }, +})); + +jest.mock('server/services/agent/RunQueueService', () => ({ + __esModule: true, + default: { + enqueueRun: jest.fn(), + }, +})); + +jest.mock('server/services/agent/RunService', () => ({ + __esModule: true, + InvalidAgentRunDefaultsError: class InvalidAgentRunDefaultsError extends Error {}, + default: { + getOwnedRun: jest.fn(), + isActiveRunConflictError: jest.fn(), + isRunNotFoundError: jest.fn(), + markFailed: jest.fn(), + markQueuedRunDispatchFailed: jest.fn(), + serializeRun: jest.fn(), + }, +})); + +jest.mock('server/services/agent/MessageStore', () => ({ + __esModule: true, + DEFAULT_AGENT_MESSAGE_PAGE_LIMIT: 50, + MAX_AGENT_MESSAGE_PAGE_LIMIT: 100, + default: { + listCanonicalMessages: jest.fn(), + serializeCanonicalMessage: jest.fn(), + }, +})); + +jest.mock('server/services/agent/RunEventService', () => ({ + __esModule: true, + DEFAULT_RUN_EVENT_PAGE_LIMIT: 100, + MAX_RUN_EVENT_PAGE_LIMIT: 500, + default: { + createCanonicalRunEventStream: jest.fn(), + listRunEventsPageForRun: jest.fn(), + serializeRunEvent: jest.fn(), + }, +})); + +jest.mock('server/services/agent/ApprovalService', () => ({ + __esModule: true, + default: { + listPendingActions: jest.fn(), + normalizePendingActionResponseBody: jest.fn((body) => { + if (!body || typeof body !== 'object' || Array.isArray(body)) { + return new Error('Request body must be a JSON object'); + } + const response = body as Record; + if (typeof response.approved !== 'boolean') { + return new Error('approved must be a boolean'); + } + + return { + approved: response.approved, + reason: typeof response.reason === 'string' ? response.reason : null, + }; + }), + resolvePendingAction: jest.fn(), + serializePendingAction: jest.fn(), + }, +})); + +import { getRequestUserIdentity } from 'server/lib/get-user'; +import { resolveRequestGitHubToken } from 'server/lib/agentSession/githubToken'; +import AgentChatSessionService from 'server/services/agent/ChatSessionService'; +import AgentSessionReadService from 'server/services/agent/SessionReadService'; +import AgentThreadService from 'server/services/agent/ThreadService'; +import AgentSessionService from 'server/services/agentSession'; +import AgentSourceService from 'server/services/agent/SourceService'; +import AgentCapabilityService from 'server/services/agent/CapabilityService'; +import AgentProviderRegistry from 'server/services/agent/ProviderRegistry'; +import AgentRunAdmissionService from 'server/services/agent/RunAdmissionService'; +import AgentRunQueueService from 'server/services/agent/RunQueueService'; +import AgentRunService from 'server/services/agent/RunService'; +import AgentMessageStore from 'server/services/agent/MessageStore'; +import AgentRunEventService from 'server/services/agent/RunEventService'; +import ApprovalService from 'server/services/agent/ApprovalService'; +import { POST as createSession } from '../sessions/route'; +import { POST as createRun } from '../threads/[threadId]/runs/route'; +import { GET as getMessages } from '../threads/[threadId]/messages/route'; +import { GET as getRunEvents } from '../runs/[runId]/events/route'; +import { GET as streamRunEvents } from '../runs/[runId]/events/stream/route'; +import { GET as getPendingActions } from '../threads/[threadId]/pending-actions/route'; +import { POST as respondToPendingAction } from '../pending-actions/[actionId]/respond/route'; + +const mockGetRequestUserIdentity = getRequestUserIdentity as jest.Mock; +const mockResolveRequestGitHubToken = resolveRequestGitHubToken as jest.Mock; +const mockCreateChatSession = AgentChatSessionService.createChatSession as jest.Mock; +const mockSerializeSessionRecord = AgentSessionReadService.serializeSessionRecord as jest.Mock; +const mockGetOwnedThreadWithSession = AgentThreadService.getOwnedThreadWithSession as jest.Mock; +const mockCanAcceptMessages = AgentSessionService.canAcceptMessages as jest.Mock; +const mockTouchActivity = AgentSessionService.touchActivity as jest.Mock; +const mockGetSessionSource = AgentSourceService.getSessionSource as jest.Mock; +const mockResolveSessionContext = AgentCapabilityService.resolveSessionContext as jest.Mock; +const mockResolveSelection = AgentProviderRegistry.resolveSelection as jest.Mock; +const mockCreateQueuedRunWithMessage = AgentRunAdmissionService.createQueuedRunWithMessage as jest.Mock; +const mockEnqueueRun = AgentRunQueueService.enqueueRun as jest.Mock; +const mockGetOwnedRun = AgentRunService.getOwnedRun as jest.Mock; +const mockSerializeRun = AgentRunService.serializeRun as jest.Mock; +const mockListCanonicalMessages = AgentMessageStore.listCanonicalMessages as jest.Mock; +const mockSerializeCanonicalMessage = AgentMessageStore.serializeCanonicalMessage as jest.Mock; +const mockListRunEventsPageForRun = AgentRunEventService.listRunEventsPageForRun as jest.Mock; +const mockCreateCanonicalRunEventStream = AgentRunEventService.createCanonicalRunEventStream as jest.Mock; +const mockSerializeRunEvent = AgentRunEventService.serializeRunEvent as jest.Mock; +const mockListPendingActions = ApprovalService.listPendingActions as jest.Mock; +const mockResolvePendingAction = ApprovalService.resolvePendingAction as jest.Mock; +const mockSerializePendingAction = ApprovalService.serializePendingAction as jest.Mock; + +type CanonicalMessage = { + id: string; + clientMessageId: string | null; + threadId: string; + runId: string | null; + role: 'user' | 'assistant'; + parts: Array>; + createdAt: string; +}; + +type RunEvent = { + uuid: string; + runUuid: string; + threadUuid: string; + sessionUuid: string; + sequence: number; + eventType: string; + payload: Record; + createdAt: string; + updatedAt: string; +}; + +type PendingAction = { + id: string; + kind: string; + status: 'pending' | 'approved' | 'denied'; + threadId: string; + runId: string; + title: string; + description: string; + requestedAt: string; + expiresAt: string | null; + toolName: string; + argumentsSummary: Array<{ name: string; value: string }>; + commandPreview: string | null; + fileChangePreview: Array>; + riskLabels: string[]; +}; + +const sampleUser = { + userId: 'sample-user', + githubUsername: 'sample-user', +}; + +const state = { + session: { + id: 17, + uuid: 'session-1', + defaultHarness: 'lifecycle_ai_sdk', + defaultModel: 'gpt-5.4', + }, + thread: { + id: 7, + uuid: 'thread-1', + }, + run: { + id: 31, + uuid: 'run-1', + status: 'queued', + }, + messages: [] as CanonicalMessage[], + events: [] as RunEvent[], + pendingAction: null as PendingAction | null, +}; + +function makeRequest(url: string, body?: unknown, headers: [string, string][] = []): NextRequest { + return { + ...(body !== undefined ? { json: jest.fn().mockResolvedValue(body) } : {}), + headers: new Headers([['x-request-id', 'req-test'], ...headers]), + nextUrl: new URL(url), + } as unknown as NextRequest; +} + +function serializeEvent(event: RunEvent) { + return { + id: event.uuid, + runId: event.runUuid, + threadId: event.threadUuid, + sessionId: event.sessionUuid, + sequence: event.sequence, + eventType: event.eventType, + version: 1, + payload: event.payload, + createdAt: event.createdAt, + updatedAt: event.updatedAt, + }; +} + +function appendEvent(eventType: string, payload: Record) { + const sequence = state.events.length + 1; + state.events.push({ + uuid: `event-${sequence}`, + runUuid: state.run.uuid, + threadUuid: state.thread.uuid, + sessionUuid: state.session.uuid, + sequence, + eventType, + payload, + createdAt: '2026-04-25T00:00:00.000Z', + updatedAt: '2026-04-25T00:00:00.000Z', + }); +} + +function buildSseStream(runId: string, afterSequence: number): ReadableStream { + const events = state.events + .filter((event) => event.runUuid === runId && event.sequence > afterSequence) + .map(serializeEvent); + + const body = events + .map((event) => + [`id: ${event.sequence}`, `event: ${event.eventType}`, `data: ${JSON.stringify(event)}`, ''].join('\n') + ) + .join('\n'); + + return new ReadableStream({ + start(controller) { + controller.enqueue(new TextEncoder().encode(body ? `${body}\n` : '')); + controller.close(); + }, + }); +} + +function serializeMessage( + message: { + uuid: string; + clientMessageId?: string | null; + role: 'user' | 'assistant'; + parts: Array>; + createdAt?: string | null; + }, + runUuid?: string | null +): CanonicalMessage { + return { + id: message.uuid, + clientMessageId: message.clientMessageId || null, + threadId: state.thread.uuid, + runId: runUuid || null, + role: message.role, + parts: message.parts, + createdAt: message.createdAt || '2026-04-25T00:00:00.000Z', + }; +} + +function simulateApprovalRequest() { + state.run.status = 'waiting_for_approval'; + state.pendingAction = { + id: 'action-1', + kind: 'tool_approval', + status: 'pending', + threadId: state.thread.uuid, + runId: state.run.uuid, + title: 'Approve workspace edit', + description: 'A workspace edit requires approval.', + requestedAt: '2026-04-25T00:00:00.000Z', + expiresAt: null, + toolName: 'mcp__sandbox__workspace_edit_file', + argumentsSummary: [{ name: 'path', value: 'sample-file.txt' }], + commandPreview: null, + fileChangePreview: [ + { + path: 'sample-file.txt', + action: 'edited', + summary: 'Updated sample-file.txt', + additions: 1, + deletions: 1, + truncated: false, + }, + ], + riskLabels: ['Workspace write'], + }; + appendEvent('approval.requested', { + actionId: state.pendingAction.id, + approvalId: 'approval-1', + toolCallId: 'tool-call-1', + }); + appendEvent('run.waiting_for_approval', { + status: 'waiting_for_approval', + }); +} + +function simulateTerminalCompletion(approved: boolean, reason: string | null) { + if (!state.pendingAction) { + throw new Error('Expected a pending action before completion'); + } + + state.pendingAction = { + ...state.pendingAction, + status: approved ? 'approved' : 'denied', + }; + appendEvent('approval.resolved', { + actionId: state.pendingAction.id, + approvalId: 'approval-1', + toolCallId: 'tool-call-1', + approved, + reason, + }); + appendEvent('approval.responded', { + actionId: state.pendingAction.id, + approvalId: 'approval-1', + toolCallId: 'tool-call-1', + approved, + reason, + }); + state.run.status = 'completed'; + appendEvent('run.completed', { + status: 'completed', + }); + state.messages.push({ + id: 'message-2', + clientMessageId: null, + threadId: state.thread.uuid, + runId: state.run.uuid, + role: 'assistant', + parts: [{ type: 'text', text: 'The workspace edit is complete.' }], + createdAt: '2026-04-25T00:00:05.000Z', + }); +} + +describe('canonical agent session API acceptance flow', () => { + beforeEach(() => { + jest.clearAllMocks(); + state.run.status = 'queued'; + state.messages = []; + state.events = []; + state.pendingAction = null; + + mockGetRequestUserIdentity.mockReturnValue(sampleUser); + mockResolveRequestGitHubToken.mockResolvedValue('sample-gh-token'); + mockCreateChatSession.mockResolvedValue(state.session); + mockSerializeSessionRecord.mockResolvedValue({ + id: state.session.uuid, + status: 'ready', + model: 'gpt-5.4', + harness: 'lifecycle_ai_sdk', + userId: sampleUser.userId, + defaultThreadId: state.thread.uuid, + sessionKind: 'chat', + workspaceStatus: 'ready', + chatStatus: 'ready', + source: { adapter: 'blank_workspace', status: 'ready' }, + sandbox: { status: 'ready' }, + canonical: { id: state.session.uuid }, + }); + mockGetOwnedThreadWithSession.mockResolvedValue({ + thread: state.thread, + session: state.session, + }); + mockCanAcceptMessages.mockReturnValue(true); + mockGetSessionSource.mockResolvedValue({ + status: 'ready', + sandboxRequirements: { filesystem: 'persistent' }, + }); + mockResolveSessionContext.mockResolvedValue({ + approvalPolicy: { defaultMode: 'require_approval', rules: {} }, + repoFullName: 'example-org/example-repo', + }); + mockResolveSelection.mockResolvedValue({ + provider: 'openai', + modelId: 'gpt-5.4', + }); + mockCreateQueuedRunWithMessage.mockImplementation(async ({ message }) => { + const storedMessage = { + uuid: 'message-1', + clientMessageId: message.clientMessageId || null, + role: 'user' as const, + parts: message.parts, + createdAt: '2026-04-25T00:00:00.000Z', + }; + state.messages = [serializeMessage(storedMessage, state.run.uuid)]; + appendEvent('run.queued', { status: 'queued' }); + + return { + run: state.run, + message: storedMessage, + created: true, + }; + }); + mockTouchActivity.mockResolvedValue(undefined); + mockEnqueueRun.mockResolvedValue(undefined); + mockSerializeRun.mockImplementation((run) => ({ + id: run.uuid, + status: run.status, + })); + mockSerializeCanonicalMessage.mockImplementation((message, _threadUuid, runUuid) => + serializeMessage(message, runUuid) + ); + mockListCanonicalMessages.mockImplementation(async () => ({ + thread: { + id: state.thread.uuid, + sessionId: state.session.uuid, + title: null, + isDefault: true, + archivedAt: null, + lastRunAt: null, + metadata: {}, + createdAt: null, + updatedAt: null, + }, + messages: state.messages, + pagination: { + hasMore: false, + nextBeforeMessageId: null, + }, + })); + mockGetOwnedRun.mockImplementation(async (runId) => { + if (runId !== state.run.uuid) { + throw new Error('Agent run not found'); + } + + return state.run; + }); + (AgentRunService.isRunNotFoundError as jest.Mock).mockImplementation((error) => { + return error instanceof Error && error.message === 'Agent run not found'; + }); + (AgentRunService.isActiveRunConflictError as jest.Mock).mockReturnValue(false); + mockListRunEventsPageForRun.mockImplementation(async (_run, { afterSequence, limit }) => { + const events = state.events.filter((event) => event.sequence > afterSequence).slice(0, limit); + const nextSequence = events.length > 0 ? events[events.length - 1].sequence : afterSequence; + + return { + events, + nextSequence, + hasMore: state.events.some((event) => event.sequence > nextSequence), + run: { + id: state.run.uuid, + status: state.run.status, + }, + limit, + maxLimit: 500, + }; + }); + mockSerializeRunEvent.mockImplementation(serializeEvent); + mockCreateCanonicalRunEventStream.mockImplementation(buildSseStream); + mockListPendingActions.mockImplementation(async () => (state.pendingAction ? [state.pendingAction] : [])); + mockResolvePendingAction.mockImplementation(async (_actionId, _userId, _status, resolution) => { + simulateTerminalCompletion( + resolution.approved === true, + typeof resolution.reason === 'string' ? resolution.reason : null + ); + return state.pendingAction; + }); + mockSerializePendingAction.mockImplementation((action) => action); + }); + + it('exercises the canonical non-UI chat client contract end to end', async () => { + const sessionResponse = await createSession( + makeRequest('http://localhost/api/v2/ai/agent/sessions', { + source: { adapter: 'blank_workspace' }, + defaults: { model: 'gpt-5.4' }, + }) + ); + const sessionBody = await sessionResponse.json(); + const threadId = sessionBody.data.defaultThreadId; + + expect(sessionResponse.status).toBe(201); + expect(threadId).toBe('thread-1'); + expect(mockCreateChatSession).toHaveBeenCalledWith( + expect.objectContaining({ + userId: 'sample-user', + model: 'gpt-5.4', + }) + ); + + const runResponse = await createRun( + makeRequest(`http://localhost/api/v2/ai/agent/threads/${threadId}/runs`, { + message: { + clientMessageId: 'client-message-1', + parts: [{ type: 'text', text: 'Inspect the workspace and summarize the main entrypoints.' }], + }, + }), + { params: { threadId } } + ); + const runBody = await runResponse.json(); + const runId = runBody.data.run.id; + + expect(runResponse.status).toBe(201); + expect(runBody.data).toEqual( + expect.objectContaining({ + run: { + id: 'run-1', + status: 'queued', + threadId: 'thread-1', + sessionId: 'session-1', + }, + message: { + id: 'message-1', + clientMessageId: 'client-message-1', + threadId: 'thread-1', + runId: 'run-1', + role: 'user', + parts: [{ type: 'text', text: 'Inspect the workspace and summarize the main entrypoints.' }], + createdAt: '2026-04-25T00:00:00.000Z', + }, + links: { + events: '/api/v2/ai/agent/runs/run-1/events', + eventStream: '/api/v2/ai/agent/runs/run-1/events/stream', + pendingActions: '/api/v2/ai/agent/threads/thread-1/pending-actions', + }, + }) + ); + expect(mockEnqueueRun).toHaveBeenCalledWith('run-1', 'submit', { githubToken: 'sample-gh-token' }); + + const initialMessagesResponse = await getMessages( + makeRequest(`http://localhost/api/v2/ai/agent/threads/${threadId}/messages`), + { params: { threadId } } + ); + const initialMessagesBody = await initialMessagesResponse.json(); + + expect(initialMessagesResponse.status).toBe(200); + expect(initialMessagesBody.data.messages).toEqual([ + expect.objectContaining({ + id: 'message-1', + clientMessageId: 'client-message-1', + role: 'user', + parts: [{ type: 'text', text: 'Inspect the workspace and summarize the main entrypoints.' }], + }), + ]); + + simulateApprovalRequest(); + + const eventsResponse = await getRunEvents( + makeRequest(`http://localhost/api/v2/ai/agent/runs/${runId}/events?afterSequence=0&limit=100`), + { params: { runId } } + ); + const eventsBody = await eventsResponse.json(); + + expect(eventsResponse.status).toBe(200); + expect(eventsBody.data.events).toEqual( + expect.arrayContaining([ + expect.objectContaining({ sequence: 1, eventType: 'run.queued' }), + expect.objectContaining({ + sequence: 2, + eventType: 'approval.requested', + payload: expect.objectContaining({ actionId: 'action-1' }), + }), + expect.objectContaining({ sequence: 3, eventType: 'run.waiting_for_approval' }), + ]) + ); + + const streamResponse = await streamRunEvents( + makeRequest(`http://localhost/api/v2/ai/agent/runs/${runId}/events/stream?afterSequence=0`, undefined, [ + ['last-event-id', '1'], + ]), + { params: { runId } } + ); + const streamBody = await streamResponse.text(); + + expect(streamResponse.status).toBe(200); + expect(streamResponse.headers.get('content-type')).toBe('text/event-stream'); + expect(mockCreateCanonicalRunEventStream).toHaveBeenCalledWith('run-1', 1); + expect(streamBody).not.toContain('id: 1'); + expect(streamBody).toContain('id: 2'); + expect(streamBody).toContain('event: approval.requested'); + expect(streamBody).toContain('"actionId":"action-1"'); + + const pendingActionsResponse = await getPendingActions( + makeRequest(`http://localhost/api/v2/ai/agent/threads/${threadId}/pending-actions`), + { params: { threadId } } + ); + const pendingActionsBody = await pendingActionsResponse.json(); + + expect(pendingActionsResponse.status).toBe(200); + expect(pendingActionsBody.data.pendingActions).toEqual([ + expect.objectContaining({ + id: 'action-1', + kind: 'tool_approval', + status: 'pending', + threadId: 'thread-1', + runId: 'run-1', + title: 'Approve workspace edit', + argumentsSummary: [{ name: 'path', value: 'sample-file.txt' }], + riskLabels: ['Workspace write'], + }), + ]); + + const approvalResponse = await respondToPendingAction( + makeRequest('http://localhost/api/v2/ai/agent/pending-actions/action-1/respond', { + approved: true, + reason: 'approved for acceptance flow', + }), + { params: { actionId: 'action-1' } } + ); + const approvalBody = await approvalResponse.json(); + + expect(approvalResponse.status).toBe(200); + expect(mockResolvePendingAction).toHaveBeenCalledWith( + 'action-1', + 'sample-user', + 'approved', + { + approved: true, + reason: 'approved for acceptance flow', + source: 'endpoint', + }, + { githubToken: 'sample-gh-token' } + ); + expect(approvalBody.data).toEqual( + expect.objectContaining({ + id: 'action-1', + status: 'approved', + }) + ); + + const terminalEventsResponse = await getRunEvents( + makeRequest(`http://localhost/api/v2/ai/agent/runs/${runId}/events?afterSequence=0&limit=100`), + { params: { runId } } + ); + const terminalEventsBody = await terminalEventsResponse.json(); + + expect(terminalEventsResponse.status).toBe(200); + expect(terminalEventsBody.data.run).toEqual({ + id: 'run-1', + status: 'completed', + }); + expect(terminalEventsBody.data.events).toEqual( + expect.arrayContaining([ + expect.objectContaining({ eventType: 'approval.resolved' }), + expect.objectContaining({ eventType: 'approval.responded' }), + expect.objectContaining({ eventType: 'run.completed' }), + ]) + ); + + const finalMessagesResponse = await getMessages( + makeRequest(`http://localhost/api/v2/ai/agent/threads/${threadId}/messages`), + { params: { threadId } } + ); + const finalMessagesBody = await finalMessagesResponse.json(); + + expect(finalMessagesResponse.status).toBe(200); + expect(finalMessagesBody.data.messages).toEqual([ + expect.objectContaining({ id: 'message-1', role: 'user' }), + { + id: 'message-2', + clientMessageId: null, + threadId: 'thread-1', + runId: 'run-1', + role: 'assistant', + parts: [{ type: 'text', text: 'The workspace edit is complete.' }], + createdAt: '2026-04-25T00:00:05.000Z', + }, + ]); + }); +}); diff --git a/src/app/api/v2/ai/agent/models/route.ts b/src/app/api/v2/ai/agent/models/route.ts index 4fdb832d..b3f66712 100644 --- a/src/app/api/v2/ai/agent/models/route.ts +++ b/src/app/api/v2/ai/agent/models/route.ts @@ -19,7 +19,6 @@ import { createApiHandler } from 'server/lib/createApiHandler'; import { successResponse, errorResponse } from 'server/lib/response'; import { getRequestUserIdentity } from 'server/lib/get-user'; import AgentProviderRegistry from 'server/services/agent/ProviderRegistry'; -import { AGENT_API_KEY_HEADER, AGENT_API_KEY_PROVIDER_HEADER } from 'server/services/agent/providerConfig'; export const dynamic = 'force-dynamic'; @@ -73,8 +72,6 @@ const getHandler = async (req: NextRequest) => { const models = await AgentProviderRegistry.listAvailableModelsForUser({ repoFullName: repo, userIdentity, - requestApiKey: req.headers.get(AGENT_API_KEY_HEADER), - requestApiKeyProvider: req.headers.get(AGENT_API_KEY_PROVIDER_HEADER), }); return successResponse({ models }, { status: 200 }, req); diff --git a/src/app/api/v2/ai/agent/pending-actions/[actionId]/respond/route.test.ts b/src/app/api/v2/ai/agent/pending-actions/[actionId]/respond/route.test.ts new file mode 100644 index 00000000..b4608c28 --- /dev/null +++ b/src/app/api/v2/ai/agent/pending-actions/[actionId]/respond/route.test.ts @@ -0,0 +1,202 @@ +/** + * Copyright 2026 GoodRx, Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +import { NextRequest } from 'next/server'; + +jest.mock('server/lib/get-user', () => ({ + getRequestUserIdentity: jest.fn(), +})); + +jest.mock('server/lib/agentSession/githubToken', () => ({ + resolveRequestGitHubToken: jest.fn(), +})); + +jest.mock('server/services/agent/ApprovalService', () => ({ + __esModule: true, + default: { + normalizePendingActionResponseBody: jest.requireActual('server/services/agent/ApprovalService').default + .normalizePendingActionResponseBody, + resolvePendingAction: jest.fn(), + serializePendingAction: jest.fn((action) => action), + }, +})); + +import { POST } from './route'; +import { getRequestUserIdentity } from 'server/lib/get-user'; +import { resolveRequestGitHubToken } from 'server/lib/agentSession/githubToken'; +import ApprovalService from 'server/services/agent/ApprovalService'; + +const mockGetRequestUserIdentity = getRequestUserIdentity as jest.Mock; +const mockResolveRequestGitHubToken = resolveRequestGitHubToken as jest.Mock; +const mockResolvePendingAction = ApprovalService.resolvePendingAction as jest.Mock; +const mockSerializePendingAction = ApprovalService.serializePendingAction as jest.Mock; + +function makeRequest(body: unknown): NextRequest { + return { + json: jest.fn().mockResolvedValue(body), + headers: new Headers([['x-request-id', 'req-test']]), + nextUrl: new URL('http://localhost/api/v2/ai/agent/pending-actions/action-1/respond'), + } as unknown as NextRequest; +} + +function makeInvalidJsonRequest(): NextRequest { + return { + json: jest.fn().mockRejectedValue(new Error('invalid json')), + headers: new Headers([['x-request-id', 'req-test']]), + nextUrl: new URL('http://localhost/api/v2/ai/agent/pending-actions/action-1/respond'), + } as unknown as NextRequest; +} + +describe('POST /api/v2/ai/agent/pending-actions/[actionId]/respond', () => { + beforeEach(() => { + jest.clearAllMocks(); + mockGetRequestUserIdentity.mockReturnValue({ + userId: 'sample-user', + githubUsername: 'sample-user', + }); + mockResolveRequestGitHubToken.mockResolvedValue('sample-gh-token'); + mockResolvePendingAction.mockResolvedValue({ + id: 'action-1', + status: 'denied', + }); + mockSerializePendingAction.mockReturnValue({ + id: 'action-1', + kind: 'tool_approval', + status: 'denied', + threadId: 'thread-1', + runId: 'run-1', + title: 'Approve workspace edit', + description: 'A workspace edit requires approval.', + requestedAt: '2026-04-11T00:00:00.000Z', + expiresAt: null, + toolName: 'mcp__sandbox__workspace_edit_file', + argumentsSummary: [], + commandPreview: null, + fileChangePreview: [], + riskLabels: ['Workspace write'], + }); + }); + + it('returns 401 when the requester is not authenticated', async () => { + mockGetRequestUserIdentity.mockReturnValue(null); + + const response = await POST(makeRequest({ approved: true }), { params: { actionId: 'action-1' } }); + + expect(response.status).toBe(401); + await expect(response.json()).resolves.toMatchObject({ + error: { message: 'Unauthorized' }, + }); + expect(mockResolvePendingAction).not.toHaveBeenCalled(); + }); + + it('resolves the pending action through the canonical response API', async () => { + const response = await POST(makeRequest({ approved: false, reason: 'not needed' }), { + params: { actionId: 'action-1' }, + }); + const body = await response.json(); + + expect(response.status).toBe(200); + expect(mockResolvePendingAction).toHaveBeenCalledWith( + 'action-1', + 'sample-user', + 'denied', + { + approved: false, + reason: 'not needed', + source: 'endpoint', + }, + { githubToken: 'sample-gh-token' } + ); + expect(body.data).toEqual({ + id: 'action-1', + kind: 'tool_approval', + status: 'denied', + threadId: 'thread-1', + runId: 'run-1', + title: 'Approve workspace edit', + description: 'A workspace edit requires approval.', + requestedAt: '2026-04-11T00:00:00.000Z', + expiresAt: null, + toolName: 'mcp__sandbox__workspace_edit_file', + argumentsSummary: [], + commandPreview: null, + fileChangePreview: [], + riskLabels: ['Workspace write'], + }); + }); + + it('rejects malformed response bodies without resolving the action', async () => { + const cases = [ + { body: {}, message: 'approved must be a boolean' }, + { body: { approved: 'yes' }, message: 'approved must be a boolean' }, + { body: { approved: true, reason: 123 }, message: 'reason must be a string when provided' }, + { + body: { approved: true, rawApproval: true }, + message: 'Unsupported pending action response fields: rawApproval', + }, + ]; + + for (const testCase of cases) { + jest.clearAllMocks(); + mockGetRequestUserIdentity.mockReturnValue({ + userId: 'sample-user', + githubUsername: 'sample-user', + }); + + const response = await POST(makeRequest(testCase.body), { params: { actionId: 'action-1' } }); + + expect(response.status).toBe(400); + await expect(response.json()).resolves.toMatchObject({ + error: { message: testCase.message }, + }); + expect(mockResolveRequestGitHubToken).not.toHaveBeenCalled(); + expect(mockResolvePendingAction).not.toHaveBeenCalled(); + } + }); + + it('rejects invalid JSON instead of denying the action', async () => { + const response = await POST(makeInvalidJsonRequest(), { params: { actionId: 'action-1' } }); + + expect(response.status).toBe(400); + await expect(response.json()).resolves.toMatchObject({ + error: { message: 'Request body must be a JSON object' }, + }); + expect(mockResolveRequestGitHubToken).not.toHaveBeenCalled(); + expect(mockResolvePendingAction).not.toHaveBeenCalled(); + }); + + it('returns 404 when the pending action cannot be resolved for the requester', async () => { + mockResolvePendingAction.mockRejectedValue(new Error('Pending action not found')); + + const response = await POST(makeRequest({ approved: true }), { params: { actionId: 'missing-action' } }); + + expect(response.status).toBe(404); + await expect(response.json()).resolves.toMatchObject({ + error: { message: 'Pending action not found' }, + }); + expect(mockResolvePendingAction).toHaveBeenCalledWith( + 'missing-action', + 'sample-user', + 'approved', + { + approved: true, + reason: null, + source: 'endpoint', + }, + { githubToken: 'sample-gh-token' } + ); + }); +}); diff --git a/src/app/api/v2/ai/agent/pending-actions/[actionId]/respond/route.ts b/src/app/api/v2/ai/agent/pending-actions/[actionId]/respond/route.ts index 793e0f91..9f5d1c67 100644 --- a/src/app/api/v2/ai/agent/pending-actions/[actionId]/respond/route.ts +++ b/src/app/api/v2/ai/agent/pending-actions/[actionId]/respond/route.ts @@ -19,6 +19,7 @@ import 'server/lib/dependencies'; import { createApiHandler } from 'server/lib/createApiHandler'; import { errorResponse, successResponse } from 'server/lib/response'; import { getRequestUserIdentity } from 'server/lib/get-user'; +import { resolveRequestGitHubToken } from 'server/lib/agentSession/githubToken'; import ApprovalService from 'server/services/agent/ApprovalService'; /** @@ -41,11 +42,15 @@ import ApprovalService from 'server/services/agent/ApprovalService'; * application/json: * schema: * type: object + * required: + * - approved + * additionalProperties: false * properties: * approved: * type: boolean * reason: * type: string + * nullable: true * responses: * '200': * description: Pending action resolved @@ -59,6 +64,24 @@ import ApprovalService from 'server/services/agent/ApprovalService'; * properties: * data: * $ref: '#/components/schemas/AgentPendingAction' + * '400': + * description: Invalid pending action response + * content: + * application/json: + * schema: + * $ref: '#/components/schemas/ApiErrorResponse' + * '401': + * description: Unauthorized + * content: + * application/json: + * schema: + * $ref: '#/components/schemas/ApiErrorResponse' + * '404': + * description: Pending action not found + * content: + * application/json: + * schema: + * $ref: '#/components/schemas/ApiErrorResponse' */ const postHandler = async (req: NextRequest, { params }: { params: { actionId: string } }) => { const userIdentity = getRequestUserIdentity(req); @@ -66,20 +89,33 @@ const postHandler = async (req: NextRequest, { params }: { params: { actionId: s return errorResponse(new Error('Unauthorized'), { status: 401 }, req); } - const body = await req.json().catch(() => ({})); - const approved = body?.approved === true; - const action = await ApprovalService.resolvePendingAction( - params.actionId, - userIdentity.userId, - approved ? 'approved' : 'denied', - { - approved, - reason: typeof body?.reason === 'string' ? body.reason : null, - source: 'endpoint', - } - ); + const body = await req.json().catch(() => null); + const responseBody = ApprovalService.normalizePendingActionResponseBody(body); + if (responseBody instanceof Error) { + return errorResponse(responseBody, { status: 400 }, req); + } - return successResponse(ApprovalService.serializePendingAction(action), { status: 200 }, req); + const githubToken = await resolveRequestGitHubToken(req); + try { + const action = await ApprovalService.resolvePendingAction( + params.actionId, + userIdentity.userId, + responseBody.approved ? 'approved' : 'denied', + { + approved: responseBody.approved, + reason: responseBody.reason, + source: 'endpoint', + }, + { githubToken } + ); + + return successResponse(ApprovalService.serializePendingAction(action), { status: 200 }, req); + } catch (error) { + if (error instanceof Error && error.message === 'Pending action not found') { + return errorResponse(error, { status: 404 }, req); + } + throw error; + } }; export const POST = createApiHandler(postHandler); diff --git a/src/app/api/v2/ai/agent/runs/[runId]/cancel/route.test.ts b/src/app/api/v2/ai/agent/runs/[runId]/cancel/route.test.ts new file mode 100644 index 00000000..5bc0bd1d --- /dev/null +++ b/src/app/api/v2/ai/agent/runs/[runId]/cancel/route.test.ts @@ -0,0 +1,104 @@ +/** + * Copyright 2026 GoodRx, Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +import { NextRequest } from 'next/server'; + +jest.mock('server/lib/get-user', () => ({ + getRequestUserIdentity: jest.fn(), +})); + +jest.mock('server/services/agent/RunService', () => ({ + __esModule: true, + default: { + cancelRun: jest.fn(), + isRunNotFoundError: jest.fn(), + serializeRun: jest.fn((run) => ({ + id: run.uuid, + status: run.status, + })), + }, +})); + +import { POST } from './route'; +import { getRequestUserIdentity } from 'server/lib/get-user'; +import AgentRunService from 'server/services/agent/RunService'; + +const mockGetRequestUserIdentity = getRequestUserIdentity as jest.Mock; +const mockCancelRun = AgentRunService.cancelRun as jest.Mock; +const mockIsRunNotFoundError = AgentRunService.isRunNotFoundError as jest.Mock; + +function makeRequest(): NextRequest { + return { + headers: new Headers([['x-request-id', 'req-test']]), + nextUrl: new URL('http://localhost/api/v2/ai/agent/runs/run-1/cancel'), + } as unknown as NextRequest; +} + +describe('POST /api/v2/ai/agent/runs/[runId]/cancel', () => { + beforeEach(() => { + jest.clearAllMocks(); + mockGetRequestUserIdentity.mockReturnValue({ + userId: 'sample-user', + }); + mockIsRunNotFoundError.mockReturnValue(false); + }); + + it('cancels the owned run through the control-plane service path', async () => { + mockCancelRun.mockResolvedValue({ + uuid: 'run-1', + status: 'cancelled', + }); + + const response = await POST(makeRequest(), { + params: { runId: 'run-1' }, + }); + const body = await response.json(); + + expect(response.status).toBe(200); + expect(mockCancelRun).toHaveBeenCalledWith('run-1', 'sample-user'); + expect(body.data).toEqual({ + id: 'run-1', + status: 'cancelled', + }); + }); + + it('maps missing runs to 404', async () => { + const error = new Error('Agent run not found'); + mockCancelRun.mockRejectedValueOnce(error); + mockIsRunNotFoundError.mockReturnValueOnce(true); + + const response = await POST(makeRequest(), { + params: { runId: 'missing-run' }, + }); + const body = await response.json(); + + expect(response.status).toBe(404); + expect(body.error.message).toBe('Agent run not found'); + }); + + it('rejects unauthenticated requests', async () => { + mockGetRequestUserIdentity.mockReturnValueOnce(null); + + const response = await POST(makeRequest(), { + params: { runId: 'run-1' }, + }); + const body = await response.json(); + + expect(response.status).toBe(401); + expect(body.error.message).toBe('Unauthorized'); + expect(mockCancelRun).not.toHaveBeenCalled(); + }); +}); diff --git a/src/app/api/v2/ai/agent/runs/[runId]/events/route.test.ts b/src/app/api/v2/ai/agent/runs/[runId]/events/route.test.ts new file mode 100644 index 00000000..a19c2866 --- /dev/null +++ b/src/app/api/v2/ai/agent/runs/[runId]/events/route.test.ts @@ -0,0 +1,196 @@ +/** + * Copyright 2026 GoodRx, Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +import { NextRequest } from 'next/server'; + +jest.mock('server/lib/get-user', () => ({ + getRequestUserIdentity: jest.fn(), +})); + +jest.mock('server/services/agent/RunEventService', () => ({ + __esModule: true, + DEFAULT_RUN_EVENT_PAGE_LIMIT: 100, + MAX_RUN_EVENT_PAGE_LIMIT: 500, + default: { + listRunEventsPageForRun: jest.fn(), + serializeRunEvent: jest.fn((event) => ({ + id: event.uuid, + runId: event.runUuid, + sequence: event.sequence, + eventType: event.eventType, + payload: event.payload, + createdAt: event.createdAt || null, + updatedAt: event.updatedAt || null, + })), + }, +})); + +jest.mock('server/services/agent/RunService', () => ({ + __esModule: true, + default: { + getOwnedRun: jest.fn(), + isRunNotFoundError: jest.fn(), + }, +})); + +import { GET } from './route'; +import { getRequestUserIdentity } from 'server/lib/get-user'; +import AgentRunEventService from 'server/services/agent/RunEventService'; +import AgentRunService from 'server/services/agent/RunService'; + +const mockGetRequestUserIdentity = getRequestUserIdentity as jest.Mock; +const mockGetOwnedRun = AgentRunService.getOwnedRun as jest.Mock; +const mockIsRunNotFoundError = AgentRunService.isRunNotFoundError as jest.Mock; +const mockListRunEventsPageForRun = AgentRunEventService.listRunEventsPageForRun as jest.Mock; + +function makeRequest(url: string): NextRequest { + return { + headers: new Headers([['x-request-id', 'req-test']]), + nextUrl: new URL(url), + } as unknown as NextRequest; +} + +describe('GET /api/v2/ai/agent/runs/[runId]/events', () => { + beforeEach(() => { + jest.clearAllMocks(); + mockGetRequestUserIdentity.mockReturnValue({ + userId: 'sample-user', + }); + }); + + it('returns a cursor page of owned run events', async () => { + const run = { + id: 17, + uuid: 'run-1', + status: 'running', + }; + mockGetOwnedRun.mockResolvedValue(run); + mockListRunEventsPageForRun.mockResolvedValue({ + events: [ + { + uuid: 'event-1', + runUuid: 'run-1', + sequence: 6, + eventType: 'message.delta', + payload: { delta: 'Hello' }, + }, + ], + nextSequence: 6, + hasMore: false, + run: { + id: 'run-1', + status: 'running', + }, + limit: 2, + maxLimit: 500, + }); + + const response = await GET( + makeRequest('http://localhost/api/v2/ai/agent/runs/run-1/events?afterSequence=5&limit=2'), + { + params: { runId: 'run-1' }, + } + ); + const body = await response.json(); + + expect(response.status).toBe(200); + expect(mockGetOwnedRun).toHaveBeenCalledWith('run-1', 'sample-user'); + expect(mockListRunEventsPageForRun).toHaveBeenCalledWith(run, { + afterSequence: 5, + limit: 2, + }); + expect(body.data).toEqual({ + run: { + id: 'run-1', + status: 'running', + }, + events: [ + { + id: 'event-1', + runId: 'run-1', + sequence: 6, + eventType: 'message.delta', + payload: { delta: 'Hello' }, + createdAt: null, + updatedAt: null, + }, + ], + pagination: { + nextSequence: 6, + hasMore: false, + }, + }); + expect(body.metadata).toEqual({ + limit: 2, + maxLimit: 500, + }); + }); + + it('clamps oversized limits to the endpoint maximum', async () => { + const run = { + id: 17, + uuid: 'run-1', + status: 'completed', + }; + mockGetOwnedRun.mockResolvedValue(run); + mockListRunEventsPageForRun.mockResolvedValue({ + events: [], + nextSequence: 0, + hasMore: false, + run: { + id: 'run-1', + status: 'completed', + }, + limit: 500, + maxLimit: 500, + }); + + const response = await GET(makeRequest('http://localhost/api/v2/ai/agent/runs/run-1/events?limit=999'), { + params: { runId: 'run-1' }, + }); + + expect(response.status).toBe(200); + expect(mockListRunEventsPageForRun).toHaveBeenCalledWith(run, { + afterSequence: 0, + limit: 500, + }); + }); + + it('returns 400 for an invalid cursor', async () => { + const response = await GET(makeRequest('http://localhost/api/v2/ai/agent/runs/run-1/events?afterSequence=-1'), { + params: { runId: 'run-1' }, + }); + const body = await response.json(); + + expect(response.status).toBe(400); + expect(body.error.message).toBe('Expected a non-negative integer cursor.'); + expect(mockGetOwnedRun).not.toHaveBeenCalled(); + }); + + it('returns 404 when the run is not owned by the user', async () => { + const missingRunError = new Error('Agent run not found'); + mockGetOwnedRun.mockRejectedValue(missingRunError); + mockIsRunNotFoundError.mockReturnValue(true); + + const response = await GET(makeRequest('http://localhost/api/v2/ai/agent/runs/missing-run/events'), { + params: { runId: 'missing-run' }, + }); + const body = await response.json(); + + expect(response.status).toBe(404); + expect(body.error.message).toBe('Agent run not found'); + }); +}); diff --git a/src/app/api/v2/ai/agent/runs/[runId]/events/route.ts b/src/app/api/v2/ai/agent/runs/[runId]/events/route.ts new file mode 100644 index 00000000..81155859 --- /dev/null +++ b/src/app/api/v2/ai/agent/runs/[runId]/events/route.ts @@ -0,0 +1,190 @@ +/** + * Copyright 2026 GoodRx, Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +import { NextRequest } from 'next/server'; +import 'server/lib/dependencies'; +import { createApiHandler } from 'server/lib/createApiHandler'; +import { errorResponse, successResponse } from 'server/lib/response'; +import { getRequestUserIdentity } from 'server/lib/get-user'; +import AgentRunEventService, { + DEFAULT_RUN_EVENT_PAGE_LIMIT, + MAX_RUN_EVENT_PAGE_LIMIT, +} from 'server/services/agent/RunEventService'; +import AgentRunService from 'server/services/agent/RunService'; + +function parseNonNegativeInteger(value: string | null, fallback: number): number { + if (value == null || value.trim() === '') { + return fallback; + } + + const parsed = Number(value); + if (!Number.isInteger(parsed) || parsed < 0) { + throw new Error('Expected a non-negative integer cursor.'); + } + + return parsed; +} + +function parsePositiveInteger(value: string | null, fallback: number): number { + if (value == null || value.trim() === '') { + return fallback; + } + + const parsed = Number(value); + if (!Number.isInteger(parsed) || parsed <= 0) { + throw new Error('Expected a positive integer limit.'); + } + + return Math.min(parsed, MAX_RUN_EVENT_PAGE_LIMIT); +} + +/** + * @openapi + * /api/v2/ai/agent/runs/{runId}/events: + * get: + * summary: Replay canonical events for an agent run + * tags: + * - Agent Sessions + * operationId: getAgentRunEvents + * parameters: + * - in: path + * name: runId + * required: true + * schema: + * type: string + * - in: query + * name: afterSequence + * required: false + * schema: + * type: integer + * minimum: 0 + * default: 0 + * description: Return events with sequence greater than this cursor. + * - in: query + * name: limit + * required: false + * schema: + * type: integer + * minimum: 1 + * maximum: 500 + * default: 100 + * description: Maximum events to return. + * responses: + * '200': + * description: Canonical run events + * content: + * application/json: + * schema: + * allOf: + * - $ref: '#/components/schemas/SuccessApiResponse' + * - type: object + * required: [data] + * properties: + * data: + * type: object + * required: [run, events, pagination] + * properties: + * run: + * type: object + * required: [id, status] + * properties: + * id: + * type: string + * status: + * $ref: '#/components/schemas/AgentRunStatus' + * events: + * type: array + * items: + * $ref: '#/components/schemas/AgentRunEvent' + * pagination: + * type: object + * required: [nextSequence, hasMore] + * properties: + * nextSequence: + * type: integer + * hasMore: + * type: boolean + * '400': + * description: Invalid event cursor or page size. + * content: + * application/json: + * schema: + * $ref: '#/components/schemas/ApiErrorResponse' + * '401': + * description: Unauthorized + * content: + * application/json: + * schema: + * $ref: '#/components/schemas/ApiErrorResponse' + * '404': + * description: Run not found + * content: + * application/json: + * schema: + * $ref: '#/components/schemas/ApiErrorResponse' + */ +const getHandler = async (req: NextRequest, { params }: { params: { runId: string } }) => { + const userIdentity = getRequestUserIdentity(req); + if (!userIdentity) { + return errorResponse(new Error('Unauthorized'), { status: 401 }, req); + } + + let afterSequence: number; + let limit: number; + try { + afterSequence = parseNonNegativeInteger(req.nextUrl.searchParams.get('afterSequence'), 0); + limit = parsePositiveInteger(req.nextUrl.searchParams.get('limit'), DEFAULT_RUN_EVENT_PAGE_LIMIT); + } catch (error) { + return errorResponse(error, { status: 400 }, req); + } + + let run; + try { + run = await AgentRunService.getOwnedRun(params.runId, userIdentity.userId); + } catch (error) { + if (AgentRunService.isRunNotFoundError(error)) { + return errorResponse(new Error('Agent run not found'), { status: 404 }, req); + } + + throw error; + } + + const page = await AgentRunEventService.listRunEventsPageForRun(run, { + afterSequence, + limit, + }); + + return successResponse( + { + run: page.run, + events: page.events.map((event) => AgentRunEventService.serializeRunEvent(event)), + pagination: { + nextSequence: page.nextSequence, + hasMore: page.hasMore, + }, + }, + { + status: 200, + metadata: { + limit: page.limit, + maxLimit: page.maxLimit, + }, + }, + req + ); +}; + +export const GET = createApiHandler(getHandler); diff --git a/src/app/api/v2/ai/agent/runs/[runId]/events/stream/route.test.ts b/src/app/api/v2/ai/agent/runs/[runId]/events/stream/route.test.ts new file mode 100644 index 00000000..520db77e --- /dev/null +++ b/src/app/api/v2/ai/agent/runs/[runId]/events/stream/route.test.ts @@ -0,0 +1,147 @@ +/** + * Copyright 2026 GoodRx, Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +import { NextRequest } from 'next/server'; + +jest.mock('server/lib/get-user', () => ({ + getRequestUserIdentity: jest.fn(), +})); + +jest.mock('server/services/agent/RunEventService', () => ({ + __esModule: true, + default: { + createCanonicalRunEventStream: jest.fn(), + }, +})); + +jest.mock('server/services/agent/RunService', () => ({ + __esModule: true, + default: { + getOwnedRun: jest.fn(), + isRunNotFoundError: jest.fn(), + }, +})); + +import { GET } from './route'; +import { getRequestUserIdentity } from 'server/lib/get-user'; +import AgentRunEventService from 'server/services/agent/RunEventService'; +import AgentRunService from 'server/services/agent/RunService'; + +const mockGetRequestUserIdentity = getRequestUserIdentity as jest.Mock; +const mockCreateCanonicalRunEventStream = AgentRunEventService.createCanonicalRunEventStream as jest.Mock; +const mockGetOwnedRun = AgentRunService.getOwnedRun as jest.Mock; +const mockIsRunNotFoundError = AgentRunService.isRunNotFoundError as jest.Mock; + +function makeRequest(url: string, headers: [string, string][] = []): NextRequest { + return { + headers: new Headers([['x-request-id', 'req-test'], ...headers]), + nextUrl: new URL(url), + } as unknown as NextRequest; +} + +function makeStream(text: string): ReadableStream { + return new ReadableStream({ + start(controller) { + controller.enqueue(new TextEncoder().encode(text)); + controller.close(); + }, + }); +} + +describe('GET /api/v2/ai/agent/runs/[runId]/events/stream', () => { + beforeEach(() => { + jest.clearAllMocks(); + mockGetRequestUserIdentity.mockReturnValue({ + userId: 'sample-user', + }); + mockGetOwnedRun.mockResolvedValue({ + uuid: 'run-1', + status: 'running', + }); + mockCreateCanonicalRunEventStream.mockReturnValue(makeStream('id: 6\n\n')); + }); + + it('returns the canonical run event SSE stream', async () => { + const response = await GET( + makeRequest('http://localhost/api/v2/ai/agent/runs/run-1/events/stream?afterSequence=5'), + { + params: { runId: 'run-1' }, + } + ); + + expect(response.status).toBe(200); + expect(response.headers.get('content-type')).toBe('text/event-stream'); + expect(response.headers.get('cache-control')).toBe('no-cache, no-transform'); + expect(await response.text()).toBe('id: 6\n\n'); + expect(mockGetOwnedRun).toHaveBeenCalledWith('run-1', 'sample-user'); + expect(mockCreateCanonicalRunEventStream).toHaveBeenCalledWith('run-1', 5); + }); + + it('uses Last-Event-ID as the replay cursor when present', async () => { + const response = await GET( + makeRequest('http://localhost/api/v2/ai/agent/runs/run-1/events/stream?afterSequence=2', [ + ['last-event-id', '8'], + ]), + { + params: { runId: 'run-1' }, + } + ); + + expect(response.status).toBe(200); + expect(mockCreateCanonicalRunEventStream).toHaveBeenCalledWith('run-1', 8); + }); + + it('returns 400 for an invalid cursor', async () => { + const response = await GET( + makeRequest('http://localhost/api/v2/ai/agent/runs/run-1/events/stream?afterSequence=-1'), + { + params: { runId: 'run-1' }, + } + ); + const body = await response.json(); + + expect(response.status).toBe(400); + expect(body.error.message).toBe('Expected a non-negative integer cursor.'); + expect(mockGetOwnedRun).not.toHaveBeenCalled(); + }); + + it('returns 404 when the run is not owned by the user', async () => { + const missingRunError = new Error('Agent run not found'); + mockGetOwnedRun.mockRejectedValue(missingRunError); + mockIsRunNotFoundError.mockReturnValue(true); + + const response = await GET(makeRequest('http://localhost/api/v2/ai/agent/runs/missing-run/events/stream'), { + params: { runId: 'missing-run' }, + }); + const body = await response.json(); + + expect(response.status).toBe(404); + expect(body.error.message).toBe('Agent run not found'); + }); + + it('returns 401 when the request has no user identity', async () => { + mockGetRequestUserIdentity.mockReturnValue(null); + + const response = await GET(makeRequest('http://localhost/api/v2/ai/agent/runs/run-1/events/stream'), { + params: { runId: 'run-1' }, + }); + const body = await response.json(); + + expect(response.status).toBe(401); + expect(body.error.message).toBe('Unauthorized'); + expect(mockGetOwnedRun).not.toHaveBeenCalled(); + }); +}); diff --git a/src/app/api/v2/ai/agent/runs/[runId]/events/stream/route.ts b/src/app/api/v2/ai/agent/runs/[runId]/events/stream/route.ts new file mode 100644 index 00000000..543259cb --- /dev/null +++ b/src/app/api/v2/ai/agent/runs/[runId]/events/stream/route.ts @@ -0,0 +1,126 @@ +/** + * Copyright 2026 GoodRx, Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +import { NextRequest } from 'next/server'; +import 'server/lib/dependencies'; +import { errorResponse } from 'server/lib/response'; +import { getRequestUserIdentity } from 'server/lib/get-user'; +import AgentRunEventService from 'server/services/agent/RunEventService'; +import AgentRunService from 'server/services/agent/RunService'; + +function parseAfterSequence(req: NextRequest): number { + const rawValue = req.headers.get('last-event-id') || req.nextUrl.searchParams.get('afterSequence'); + if (rawValue == null || rawValue.trim() === '') { + return 0; + } + + const parsed = Number(rawValue); + if (!Number.isInteger(parsed) || parsed < 0) { + throw new Error('Expected a non-negative integer cursor.'); + } + + return parsed; +} + +/** + * @openapi + * /api/v2/ai/agent/runs/{runId}/events/stream: + * get: + * summary: Stream canonical events for an agent run + * tags: + * - Agent Sessions + * operationId: streamAgentRunEvents + * parameters: + * - in: path + * name: runId + * required: true + * schema: + * type: string + * - in: query + * name: afterSequence + * required: false + * schema: + * type: integer + * minimum: 0 + * default: 0 + * description: Replay events with sequence greater than this cursor before following live events. + * - in: header + * name: Last-Event-ID + * required: false + * schema: + * type: integer + * minimum: 0 + * description: Browser SSE resume cursor. Takes precedence over afterSequence. + * responses: + * '200': + * description: Canonical run event stream. + * content: + * text/event-stream: + * schema: + * type: string + * '400': + * description: Invalid event cursor. + * content: + * application/json: + * schema: + * $ref: '#/components/schemas/ApiErrorResponse' + * '401': + * description: Unauthorized + * content: + * application/json: + * schema: + * $ref: '#/components/schemas/ApiErrorResponse' + * '404': + * description: Run not found + * content: + * application/json: + * schema: + * $ref: '#/components/schemas/ApiErrorResponse' + */ +const getHandler = async (req: NextRequest, { params }: { params: { runId: string } }): Promise => { + const userIdentity = getRequestUserIdentity(req); + if (!userIdentity) { + return errorResponse(new Error('Unauthorized'), { status: 401 }, req); + } + + let afterSequence: number; + try { + afterSequence = parseAfterSequence(req); + } catch (error) { + return errorResponse(error, { status: 400 }, req); + } + + let run; + try { + run = await AgentRunService.getOwnedRun(params.runId, userIdentity.userId); + } catch (error) { + if (AgentRunService.isRunNotFoundError(error)) { + return errorResponse(new Error('Agent run not found'), { status: 404 }, req); + } + + throw error; + } + + return new Response(AgentRunEventService.createCanonicalRunEventStream(run.uuid, afterSequence), { + headers: { + 'Content-Type': 'text/event-stream', + 'Cache-Control': 'no-cache, no-transform', + Connection: 'keep-alive', + }, + }); +}; + +export const GET = getHandler; diff --git a/src/app/api/v2/ai/agent/runs/[runId]/stream/route.test.ts b/src/app/api/v2/ai/agent/runs/[runId]/stream/route.test.ts deleted file mode 100644 index 40758241..00000000 --- a/src/app/api/v2/ai/agent/runs/[runId]/stream/route.test.ts +++ /dev/null @@ -1,297 +0,0 @@ -/** - * Copyright 2026 GoodRx, Inc. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -import { NextRequest } from 'next/server'; - -const mockCreateUIMessageStream = jest.fn(); -const mockCreateUIMessageStreamResponse = jest.fn(); - -jest.mock('ai', () => ({ - createUIMessageStream: (...args: unknown[]) => mockCreateUIMessageStream(...args), - createUIMessageStreamResponse: (...args: unknown[]) => mockCreateUIMessageStreamResponse(...args), -})); - -jest.mock('server/lib/get-user', () => ({ - getRequestUserIdentity: jest.fn(), -})); - -jest.mock('server/services/agent/RunService', () => ({ - __esModule: true, - default: { - getOwnedRun: jest.fn(), - isRunNotFoundError: jest.fn(), - }, -})); - -jest.mock('server/services/agent/StreamBroker', () => ({ - __esModule: true, - default: { - open: jest.fn(), - }, -})); - -jest.mock('server/services/agent/MessageStore', () => ({ - __esModule: true, - default: { - listRunMessages: jest.fn(), - }, -})); - -import { GET } from './route'; -import { getRequestUserIdentity } from 'server/lib/get-user'; -import AgentRunService from 'server/services/agent/RunService'; -import AgentStreamBroker from 'server/services/agent/StreamBroker'; -import AgentMessageStore from 'server/services/agent/MessageStore'; - -const mockGetRequestUserIdentity = getRequestUserIdentity as jest.Mock; -const mockGetOwnedRun = AgentRunService.getOwnedRun as jest.Mock; -const mockIsRunNotFoundError = AgentRunService.isRunNotFoundError as jest.Mock; -const mockOpenStream = AgentStreamBroker.open as jest.Mock; -const mockListRunMessages = AgentMessageStore.listRunMessages as jest.Mock; - -function makeRequest(url: string): NextRequest { - return { - headers: new Headers([['x-request-id', 'req-test']]), - nextUrl: new URL(url), - } as unknown as NextRequest; -} - -describe('GET /api/v2/ai/agent/runs/[runId]/stream', () => { - beforeEach(() => { - jest.clearAllMocks(); - mockCreateUIMessageStream.mockImplementation(({ execute }) => { - const chunks: unknown[] = []; - void execute({ - writer: { - write: (chunk: unknown) => { - chunks.push(chunk); - }, - }, - }); - return chunks; - }); - mockCreateUIMessageStreamResponse.mockImplementation( - ({ stream }) => - new Response(JSON.stringify(stream), { - status: 200, - headers: { - 'content-type': 'application/json', - }, - }) - ); - }); - - it('returns 204 when the run cannot be found', async () => { - const missingRunError = new Error('Agent run not found'); - mockGetRequestUserIdentity.mockReturnValue({ - userId: 'sample-user', - }); - mockGetOwnedRun.mockRejectedValue(missingRunError); - mockIsRunNotFoundError.mockReturnValue(true); - - const response = await GET(makeRequest('http://localhost/api/v2/ai/agent/runs/unavailable/stream'), { - params: { runId: 'unavailable' }, - }); - - expect(response.status).toBe(204); - expect(mockOpenStream).not.toHaveBeenCalled(); - expect(mockListRunMessages).not.toHaveBeenCalled(); - }); - - it('replays only the final tool state for completed approval-backed tools', async () => { - mockGetRequestUserIdentity.mockReturnValue({ - userId: 'sample-user', - }); - mockGetOwnedRun.mockResolvedValue({ - uuid: 'run-123', - status: 'completed', - streamState: { - finishReason: 'stop', - }, - }); - mockOpenStream.mockReturnValue(null); - mockListRunMessages.mockResolvedValue([ - { - id: 'assistant-1', - role: 'assistant', - metadata: { - runId: 'run-123', - }, - parts: [ - { - type: 'tool-workspace_edit_file', - toolCallId: 'tool-1', - toolName: 'workspace_edit_file', - state: 'output-available', - input: { - path: '/workspace/sample.ts', - oldText: 'before', - newText: 'after', - }, - output: { - ok: true, - }, - approval: { - id: 'approval-1', - }, - }, - ], - }, - ]); - - const response = await GET(makeRequest('http://localhost/api/v2/ai/agent/runs/run-123/stream'), { - params: { runId: 'run-123' }, - }); - - const chunks = (await response.json()) as Array<{ type: string }>; - - expect(response.status).toBe(200); - expect(chunks.some((chunk) => chunk.type === 'tool-approval-request')).toBe(false); - expect(chunks.some((chunk) => chunk.type === 'tool-output-available')).toBe(true); - }); - - it('removes duplicate file-change payloads from replayed tool output when canonical file changes exist', async () => { - mockGetRequestUserIdentity.mockReturnValue({ - userId: 'sample-user', - }); - mockGetOwnedRun.mockResolvedValue({ - uuid: 'run-456', - status: 'completed', - streamState: { - finishReason: 'stop', - }, - }); - mockOpenStream.mockReturnValue(null); - mockListRunMessages.mockResolvedValue([ - { - id: 'assistant-2', - role: 'assistant', - metadata: { - runId: 'run-456', - }, - parts: [ - { - type: 'data-file-change', - id: 'tool-2:file.ts', - data: { - id: 'tool-2:file.ts', - toolCallId: 'tool-2', - path: 'file.ts', - displayPath: 'file.ts', - sourceTool: 'workspace.edit_file', - stage: 'applied', - kind: 'edited', - additions: 1, - deletions: 0, - truncated: false, - }, - }, - { - type: 'tool-workspace_edit_file', - toolCallId: 'tool-2', - toolName: 'workspace_edit_file', - state: 'output-available', - output: { - content: [ - { - type: 'text', - text: JSON.stringify( - { - ok: true, - path: 'file.ts', - fileChanges: [{ path: 'file.ts', additions: 1 }], - }, - null, - 2 - ), - }, - ], - }, - }, - ], - }, - ]); - - const response = await GET(makeRequest('http://localhost/api/v2/ai/agent/runs/run-456/stream'), { - params: { runId: 'run-456' }, - }); - - const chunks = (await response.json()) as Array<{ - type: string; - output?: { - content?: Array<{ text?: string }>; - }; - }>; - const toolOutput = chunks.find((chunk) => chunk.type === 'tool-output-available'); - - expect(response.status).toBe(200); - expect(toolOutput?.output?.content?.[0]?.text).not.toContain('fileChanges'); - }); - - it('replays stored stream chunks when terminal runs have no persisted assistant message', async () => { - mockGetRequestUserIdentity.mockReturnValue({ - userId: 'sample-user', - }); - mockGetOwnedRun.mockResolvedValue({ - uuid: 'run-234', - status: 'completed', - streamState: { - chunks: [ - { - type: 'start', - messageMetadata: { - runId: 'run-234', - }, - }, - { - type: 'text-start', - id: 'text-1', - }, - { - type: 'text-delta', - id: 'text-1', - delta: 'Recovered answer', - }, - { - type: 'text-end', - id: 'text-1', - }, - { - type: 'finish', - finishReason: 'stop', - }, - ], - }, - }); - mockOpenStream.mockReturnValue(null); - mockListRunMessages.mockResolvedValue([ - { - id: 'user-1', - role: 'user', - parts: [{ type: 'text', text: 'Hello' }], - }, - ]); - - const response = await GET(makeRequest('http://localhost/api/v2/ai/agent/runs/run-234/stream'), { - params: { runId: 'run-234' }, - }); - - const chunks = (await response.json()) as Array<{ type: string }>; - - expect(response.status).toBe(200); - expect(chunks.some((chunk) => chunk.type === 'text-delta')).toBe(true); - }); -}); diff --git a/src/app/api/v2/ai/agent/runs/[runId]/stream/route.ts b/src/app/api/v2/ai/agent/runs/[runId]/stream/route.ts deleted file mode 100644 index 6ac991b5..00000000 --- a/src/app/api/v2/ai/agent/runs/[runId]/stream/route.ts +++ /dev/null @@ -1,363 +0,0 @@ -/** - * Copyright 2026 GoodRx, Inc. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -import { createUIMessageStream, createUIMessageStreamResponse, type ProviderMetadata, type UIMessageChunk } from 'ai'; -import { NextRequest } from 'next/server'; -import 'server/lib/dependencies'; -import { getRequestUserIdentity } from 'server/lib/get-user'; -import AgentMessageStore from 'server/services/agent/MessageStore'; -import AgentRunService from 'server/services/agent/RunService'; -import AgentStreamBroker from 'server/services/agent/StreamBroker'; -import { sanitizeAgentRunStreamChunks } from 'server/services/agent/streamState'; -import type { AgentUIDataParts, AgentUIMessage } from 'server/services/agent/types'; - -type AgentUiMessageChunk = UIMessageChunk; -const STREAM_POLL_INTERVAL_MS = 300; - -function sleep(durationMs: number): Promise { - return new Promise((resolve) => setTimeout(resolve, durationMs)); -} - -function getStoredChunks(run: { streamState?: Record | null }): AgentUiMessageChunk[] { - const streamState = run.streamState || {}; - const rawChunks = Array.isArray(streamState.chunks) ? streamState.chunks : []; - - return sanitizeAgentRunStreamChunks(rawChunks as AgentUiMessageChunk[]); -} - -function buildDurableReplayStream(runId: string, originalMessages: AgentUIMessage[]): ReadableStream { - const stream = createUIMessageStream({ - execute: async ({ writer }) => { - let emittedCount = 0; - let shouldContinue = true; - - while (shouldContinue) { - const currentRun = await AgentRunService.getRunByUuid(runId); - if (!currentRun) { - return; - } - - const storedChunks = getStoredChunks(currentRun); - while (emittedCount < storedChunks.length) { - writer.write(storedChunks[emittedCount] as Parameters[0]); - emittedCount += 1; - } - - shouldContinue = currentRun.status === 'running'; - if (!shouldContinue) { - return; - } - - await sleep(STREAM_POLL_INTERVAL_MS); - } - }, - originalMessages, - }); - - return stream as ReadableStream; -} - -function buildStoredChunkReplayStream( - originalMessages: AgentUIMessage[], - chunks: AgentUiMessageChunk[] -): ReadableStream { - const stream = createUIMessageStream({ - originalMessages, - execute: ({ writer }) => { - for (const chunk of chunks) { - writer.write(chunk as Parameters[0]); - } - }, - }); - - return stream as ReadableStream; -} - -function replayTextPart(chunks: AgentUiMessageChunk[], part: Record, id: string): void { - chunks.push({ - type: 'text-start', - id, - providerMetadata: part.providerMetadata as ProviderMetadata | undefined, - }); - chunks.push({ - type: 'text-delta', - id, - delta: typeof part.text === 'string' ? part.text : '', - providerMetadata: part.providerMetadata as ProviderMetadata | undefined, - }); - chunks.push({ - type: 'text-end', - id, - providerMetadata: part.providerMetadata as ProviderMetadata | undefined, - }); -} - -function replayReasoningPart(chunks: AgentUiMessageChunk[], part: Record, id: string): void { - chunks.push({ - type: 'reasoning-start', - id, - providerMetadata: part.providerMetadata as ProviderMetadata | undefined, - }); - chunks.push({ - type: 'reasoning-delta', - id, - delta: typeof part.text === 'string' ? part.text : '', - providerMetadata: part.providerMetadata as ProviderMetadata | undefined, - }); - chunks.push({ - type: 'reasoning-end', - id, - providerMetadata: part.providerMetadata as ProviderMetadata | undefined, - }); -} - -function replayToolPart(chunks: AgentUiMessageChunk[], part: Record): void { - const toolCallId = typeof part.toolCallId === 'string' ? part.toolCallId : null; - if (!toolCallId) { - return; - } - - const toolName = - typeof part.toolName === 'string' - ? part.toolName - : typeof part.type === 'string' && part.type.startsWith('tool-') - ? part.type.replace(/^tool-/, '') - : 'unknown'; - const input = 'input' in part ? part.input : null; - const providerMetadata = part.callProviderMetadata as ProviderMetadata | undefined; - const resultProviderMetadata = part.resultProviderMetadata as ProviderMetadata | undefined; - const dynamic = part.type === 'dynamic-tool'; - const title = typeof part.title === 'string' ? part.title : undefined; - - if (part.state === 'output-error') { - chunks.push({ - type: 'tool-input-error', - toolCallId, - toolName, - input: input ?? part.rawInput ?? null, - errorText: typeof part.errorText === 'string' ? part.errorText : 'Tool execution failed.', - providerExecuted: part.providerExecuted === true, - providerMetadata, - dynamic, - title, - }); - } else if (part.state !== 'input-streaming') { - chunks.push({ - type: 'tool-input-available', - toolCallId, - toolName, - input, - providerExecuted: part.providerExecuted === true, - providerMetadata, - dynamic, - title, - }); - } - - const approval = part.approval as { id?: string } | undefined; - if (part.state === 'approval-requested' && approval?.id) { - chunks.push({ - type: 'tool-approval-request', - approvalId: approval.id, - toolCallId, - }); - } - - switch (part.state) { - case 'output-available': - chunks.push({ - type: 'tool-output-available', - toolCallId, - output: part.output, - providerExecuted: part.providerExecuted === true, - providerMetadata: resultProviderMetadata, - dynamic, - preliminary: part.preliminary === true, - }); - break; - case 'output-error': - chunks.push({ - type: 'tool-output-error', - toolCallId, - errorText: typeof part.errorText === 'string' ? part.errorText : 'Tool execution failed.', - providerExecuted: part.providerExecuted === true, - providerMetadata: resultProviderMetadata, - dynamic, - }); - break; - case 'output-denied': - chunks.push({ - type: 'tool-output-denied', - toolCallId, - }); - break; - default: - break; - } -} - -function buildReplayChunks(message: AgentUIMessage, finishReason?: string | null): AgentUiMessageChunk[] { - const chunks: AgentUiMessageChunk[] = [ - { - type: 'start', - messageId: message.id, - messageMetadata: message.metadata, - }, - ]; - - for (const [index, rawPart] of (message.parts || []).entries()) { - if (!rawPart || typeof rawPart !== 'object') { - continue; - } - - const part = rawPart as Record; - const partType = typeof part.type === 'string' ? part.type : ''; - - switch (partType) { - case 'step-start': - chunks.push({ type: 'start-step' }); - break; - case 'text': - replayTextPart(chunks, part, `${message.id}-text-${index}`); - break; - case 'reasoning': - replayReasoningPart(chunks, part, `${message.id}-reasoning-${index}`); - break; - case 'source-url': - case 'source-document': - case 'file': - chunks.push(part as unknown as AgentUiMessageChunk); - break; - case 'dynamic-tool': - replayToolPart(chunks, part); - break; - default: - if (partType.startsWith('tool-')) { - replayToolPart(chunks, part); - break; - } - - if (partType.startsWith('data-')) { - chunks.push(part as unknown as AgentUiMessageChunk); - } - break; - } - } - - chunks.push({ - type: 'finish', - finishReason: - (finishReason as 'stop' | 'length' | 'tool-calls' | 'content-filter' | 'error' | 'other' | undefined) || - undefined, - messageMetadata: message.metadata, - }); - - return sanitizeAgentRunStreamChunks(chunks); -} - -/** - * @openapi - * /api/v2/ai/agent/runs/{runId}/stream: - * get: - * summary: Reconnect to an agent run stream or replay the latest persisted assistant message - * tags: - * - Agent Sessions - * operationId: reconnectAgentRunStream - * parameters: - * - in: path - * name: runId - * required: true - * schema: - * type: string - * responses: - * '200': - * description: UI message stream - * content: - * text/event-stream: - * schema: - * type: string - * '204': - * description: No active or replayable stream is available for this run. - * '401': - * description: Unauthorized - * content: - * application/json: - * schema: - * $ref: '#/components/schemas/ApiErrorResponse' - */ -const getHandler = async (req: NextRequest, { params }: { params: { runId: string } }): Promise => { - const userIdentity = getRequestUserIdentity(req); - if (!userIdentity) { - return new Response('Unauthorized', { status: 401 }); - } - - let run; - try { - run = await AgentRunService.getOwnedRun(params.runId, userIdentity.userId); - } catch (error) { - if (AgentRunService.isRunNotFoundError(error)) { - return new Response(null, { status: 204 }); - } - - throw error; - } - - const activeStream = AgentStreamBroker.open(run.uuid); - if (activeStream) { - return createUIMessageStreamResponse({ - stream: activeStream, - }); - } - - const messages = await AgentMessageStore.listRunMessages(params.runId, userIdentity.userId); - if (run.status === 'running' && getStoredChunks(run).length > 0) { - return createUIMessageStreamResponse({ - stream: buildDurableReplayStream(run.uuid, messages), - }); - } - - const replayMessage = [...messages].reverse().find((message) => message.role === 'assistant'); - - if (!replayMessage) { - const storedChunks = getStoredChunks(run); - if (storedChunks.length === 0) { - return new Response(null, { status: 204 }); - } - - return createUIMessageStreamResponse({ - stream: buildStoredChunkReplayStream(messages, storedChunks), - }); - } - - const chunks = buildReplayChunks( - replayMessage, - typeof run.streamState?.finishReason === 'string' ? run.streamState.finishReason : null - ); - const stream = createUIMessageStream({ - execute: ({ writer }) => { - for (const chunk of chunks) { - writer.write(chunk as Parameters[0]); - } - }, - originalMessages: messages, - }); - - return createUIMessageStreamResponse({ - stream: stream as ReadableStream, - }); -}; - -export const GET = getHandler; diff --git a/src/app/api/v2/ai/agent/sandbox-sessions/route.ts b/src/app/api/v2/ai/agent/sandbox-sessions/route.ts index 7d7645ff..166d1461 100644 --- a/src/app/api/v2/ai/agent/sandbox-sessions/route.ts +++ b/src/app/api/v2/ai/agent/sandbox-sessions/route.ts @@ -22,14 +22,15 @@ import { getRequestUserIdentity } from 'server/lib/get-user'; import { resolveRequestGitHubToken } from 'server/lib/agentSession/githubToken'; import { AgentSessionRuntimeConfigError, + AgentSessionWorkspaceStorageConfigError, resolveAgentSessionRuntimeConfig, + resolveAgentSessionWorkspaceStorageIntent, } from 'server/lib/agentSession/runtimeConfig'; import { encrypt } from 'server/lib/encryption'; import { redisClient } from 'server/lib/dependencies'; import QueueManager from 'server/lib/queueManager'; import { QUEUE_NAMES } from 'shared/config'; import { setSandboxLaunchState, toPublicSandboxLaunchState } from 'server/lib/agentSession/sandboxLaunchState'; -import { AGENT_API_KEY_HEADER, AGENT_API_KEY_PROVIDER_HEADER } from 'server/services/agent/providerConfig'; import AgentSandboxSessionService, { formatRequestedSandboxServicesLabel, summarizeRequestedSandboxServices, @@ -44,6 +45,9 @@ interface CreateSandboxSessionBody { service?: unknown; services?: unknown; model?: string; + workspace?: { + storageSize?: string; + }; } function isRequestedSandboxServiceRef(value: unknown): value is RequestedAgentSessionServiceRef { @@ -98,6 +102,27 @@ function parseRequestedSandboxServicesFromBody(body: CreateSandboxSessionBody): throw new Error('service or services is required'); } +function parseRequestedWorkspaceStorageSize(body: CreateSandboxSessionBody): string | undefined { + const workspace = body.workspace; + if (workspace === undefined) { + return undefined; + } + + if (!workspace || typeof workspace !== 'object' || Array.isArray(workspace)) { + throw new Error('workspace must be an object'); + } + + if (workspace.storageSize === undefined) { + return undefined; + } + + if (typeof workspace.storageSize !== 'string' || !workspace.storageSize.trim()) { + throw new Error('workspace.storageSize must be a non-empty string'); + } + + return workspace.storageSize.trim(); +} + const sandboxLaunchQueue = QueueManager.getInstance().registerQueue(QUEUE_NAMES.AGENT_SANDBOX_SESSION_LAUNCH, { connection: redisClient.getConnection(), defaultJobOptions: { @@ -198,6 +223,12 @@ const sandboxLaunchQueue = QueueManager.getInstance().registerQueue(QUEUE_NAMES. * type: string * model: * type: string + * workspace: + * type: object + * properties: + * storageSize: + * type: string + * description: Optional workspace PVC size. Accepted only when admin runtime settings allow client overrides. * - type: object * required: * - baseBuildUuid @@ -223,6 +254,12 @@ const sandboxLaunchQueue = QueueManager.getInstance().registerQueue(QUEUE_NAMES. * type: string * model: * type: string + * workspace: + * type: object + * properties: + * storageSize: + * type: string + * description: Optional workspace PVC size. Accepted only when admin runtime settings allow client overrides. * responses: * '200': * description: Sandbox session launch queued @@ -348,12 +385,15 @@ const postHandler = async (req: NextRequest) => { try { const requestedServices = parseRequestedSandboxServicesFromBody(body); + const requestedWorkspaceStorageSize = parseRequestedWorkspaceStorageSize(body); const requestedServiceSummary = summarizeRequestedSandboxServices(requestedServices); const requestedServiceLabel = formatRequestedSandboxServicesLabel(requestedServices); const runtimeConfig = await resolveAgentSessionRuntimeConfig(); + const workspaceStorage = resolveAgentSessionWorkspaceStorageIntent({ + requestedSize: requestedWorkspaceStorageSize, + storage: runtimeConfig.workspaceStorage, + }); const githubToken = await resolveRequestGitHubToken(req); - const requestApiKey = req.headers.get(AGENT_API_KEY_HEADER); - const requestApiKeyProvider = req.headers.get(AGENT_API_KEY_PROVIDER_HEADER); const launchId = uuid(); const now = new Date().toISOString(); await setSandboxLaunchState(redisClient.getRedis(), { @@ -380,8 +420,6 @@ const postHandler = async (req: NextRequest) => { userId: userIdentity.userId, userIdentity, encryptedGithubToken: githubToken ? encrypt(githubToken) : null, - encryptedRequestApiKey: requestApiKey ? encrypt(requestApiKey) : null, - requestApiKeyProvider, baseBuildUuid: body.baseBuildUuid, services: requestedServices, model: body.model, @@ -392,6 +430,8 @@ const postHandler = async (req: NextRequest) => { keepAttachedServicesOnSessionNode: runtimeConfig.keepAttachedServicesOnSessionNode, readiness: runtimeConfig.readiness, resources: runtimeConfig.resources, + workspaceStorage, + redisTtlSeconds: runtimeConfig.cleanup.redisTtlSeconds, } as SandboxSessionLaunchJob, { jobId: launchId, @@ -427,6 +467,10 @@ const postHandler = async (req: NextRequest) => { return errorResponse(err, { status: 503 }, req); } + if (err instanceof AgentSessionWorkspaceStorageConfigError) { + return errorResponse(err, { status: 400 }, req); + } + return errorResponse(err, { status: 400 }, req); } }; diff --git a/src/app/api/v2/ai/agent/sessions/[sessionId]/route.ts b/src/app/api/v2/ai/agent/sessions/[sessionId]/route.ts index b37e33bc..91128ebf 100644 --- a/src/app/api/v2/ai/agent/sessions/[sessionId]/route.ts +++ b/src/app/api/v2/ai/agent/sessions/[sessionId]/route.ts @@ -19,7 +19,7 @@ import 'server/lib/dependencies'; import { createApiHandler } from 'server/lib/createApiHandler'; import { successResponse, errorResponse } from 'server/lib/response'; import { getRequestUserIdentity } from 'server/lib/get-user'; -import { serializeAgentSessionSummary } from 'server/services/agent/serializeSessionSummary'; +import AgentSessionReadService from 'server/services/agent/SessionReadService'; import AgentSessionService from 'server/services/agentSession'; /** @@ -48,146 +48,7 @@ import AgentSessionService from 'server/services/agentSession'; * request_id: * type: string * data: - * type: object - * required: - * - id - * - buildUuid - * - baseBuildUuid - * - buildKind - * - userId - * - ownerGithubUsername - * - podName - * - namespace - * - model - * - status - * - repo - * - branch - * - services - * - lastActivity - * - createdAt - * - updatedAt - * - endedAt - * - startupFailure - * - editorUrl - * properties: - * id: - * type: string - * buildUuid: - * type: string - * nullable: true - * baseBuildUuid: - * type: string - * nullable: true - * buildKind: - * $ref: '#/components/schemas/BuildKind' - * userId: - * type: string - * ownerGithubUsername: - * type: string - * nullable: true - * podName: - * type: string - * namespace: - * type: string - * model: - * type: string - * status: - * type: string - * enum: [starting, active, ended, error] - * repo: - * type: string - * nullable: true - * branch: - * type: string - * nullable: true - * primaryRepo: - * type: string - * nullable: true - * primaryBranch: - * type: string - * nullable: true - * workspaceRepos: - * type: array - * items: - * type: object - * required: [repo, repoUrl, branch, mountPath] - * properties: - * repo: - * type: string - * repoUrl: - * type: string - * branch: - * type: string - * revision: - * type: string - * nullable: true - * mountPath: - * type: string - * primary: - * type: boolean - * selectedServices: - * type: array - * items: - * type: object - * required: [name, deployId, repo, branch, workspacePath] - * properties: - * name: - * type: string - * deployId: - * type: integer - * repo: - * type: string - * branch: - * type: string - * revision: - * type: string - * nullable: true - * resourceName: - * type: string - * nullable: true - * workspacePath: - * type: string - * workDir: - * type: string - * nullable: true - * services: - * type: array - * items: - * type: string - * lastActivity: - * type: string - * format: date-time - * createdAt: - * type: string - * format: date-time - * updatedAt: - * type: string - * format: date-time - * endedAt: - * type: string - * nullable: true - * format: date-time - * startupFailure: - * type: object - * nullable: true - * required: - * - stage - * - title - * - message - * - recordedAt - * properties: - * stage: - * type: string - * enum: [create_session, connect_runtime, attach_services] - * title: - * type: string - * message: - * type: string - * recordedAt: - * type: string - * format: date-time - * editorUrl: - * type: string + * $ref: '#/components/schemas/AgentSessionSummary' * error: * nullable: true * '401': @@ -251,16 +112,12 @@ const getHandler = async (req: NextRequest, { params }: { params: Promise<{ sess if (!userIdentity) return errorResponse(new Error('Unauthorized'), { status: 401 }, req); const { sessionId } = await params; - const session = await AgentSessionService.getSession(sessionId); - if (!session) { + const sessionRecord = await AgentSessionReadService.getOwnedSessionRecord(sessionId, userIdentity.userId); + if (!sessionRecord) { return errorResponse(new Error('Session not found'), { status: 404 }, req); } - if (session.userId !== userIdentity.userId) { - return errorResponse(new Error('Forbidden: you do not own this session'), { status: 401 }, req); - } - - return successResponse(serializeAgentSessionSummary(session), { status: 200 }, req); + return successResponse(sessionRecord, { status: 200 }, req); }; const deleteHandler = async (req: NextRequest, { params }: { params: Promise<{ sessionId: string }> }) => { diff --git a/src/app/api/v2/ai/agent/pending-actions/[actionId]/approve/route.ts b/src/app/api/v2/ai/agent/sessions/[sessionId]/sandbox/resume/route.ts similarity index 59% rename from src/app/api/v2/ai/agent/pending-actions/[actionId]/approve/route.ts rename to src/app/api/v2/ai/agent/sessions/[sessionId]/sandbox/resume/route.ts index d24036f8..f56c4ce2 100644 --- a/src/app/api/v2/ai/agent/pending-actions/[actionId]/approve/route.ts +++ b/src/app/api/v2/ai/agent/sessions/[sessionId]/sandbox/resume/route.ts @@ -17,27 +17,29 @@ import { NextRequest } from 'next/server'; import 'server/lib/dependencies'; import { createApiHandler } from 'server/lib/createApiHandler'; -import { errorResponse, successResponse } from 'server/lib/response'; import { getRequestUserIdentity } from 'server/lib/get-user'; -import ApprovalService from 'server/services/agent/ApprovalService'; +import { errorResponse, successResponse } from 'server/lib/response'; +import { resolveRequestGitHubToken } from 'server/lib/agentSession/githubToken'; +import AgentSessionService from 'server/services/agentSession'; +import AgentSessionReadService from 'server/services/agent/SessionReadService'; /** * @openapi - * /api/v2/ai/agent/pending-actions/{actionId}/approve: + * /api/v2/ai/agent/sessions/{sessionId}/sandbox/resume: * post: - * summary: Approve a pending action + * summary: Resume a chat session sandbox runtime * tags: * - Agent Sessions - * operationId: approveAgentPendingAction + * operationId: resumeAgentSessionSandbox * parameters: * - in: path - * name: actionId + * name: sessionId * required: true * schema: * type: string * responses: * '200': - * description: Pending action approved + * description: Resumed session * content: * application/json: * schema: @@ -47,22 +49,31 @@ import ApprovalService from 'server/services/agent/ApprovalService'; * required: [data] * properties: * data: - * $ref: '#/components/schemas/AgentPendingAction' + * $ref: '#/components/schemas/AgentSessionSummary' + * '400': + * description: Session cannot be resumed + * '401': + * description: Unauthorized */ -const postHandler = async (req: NextRequest, { params }: { params: { actionId: string } }) => { +const postHandler = async (req: NextRequest, { params }: { params: { sessionId: string } }) => { const userIdentity = getRequestUserIdentity(req); if (!userIdentity) { return errorResponse(new Error('Unauthorized'), { status: 401 }, req); } - const body = await req.json().catch(() => ({})); - const action = await ApprovalService.resolvePendingAction(params.actionId, userIdentity.userId, 'approved', { - approved: true, - reason: typeof body?.reason === 'string' ? body.reason : null, - source: 'endpoint', - }); + try { + const githubToken = await resolveRequestGitHubToken(req); + const session = await AgentSessionService.resumeChatRuntime({ + sessionId: params.sessionId, + userId: userIdentity.userId, + userIdentity, + githubToken, + }); - return successResponse(ApprovalService.serializePendingAction(action), { status: 200 }, req); + return successResponse(await AgentSessionReadService.serializeSessionRecord(session), { status: 200 }, req); + } catch (error) { + return errorResponse(error, { status: 400 }, req); + } }; export const POST = createApiHandler(postHandler); diff --git a/src/app/api/v2/ai/agent/pending-actions/[actionId]/deny/route.ts b/src/app/api/v2/ai/agent/sessions/[sessionId]/sandbox/suspend/route.ts similarity index 58% rename from src/app/api/v2/ai/agent/pending-actions/[actionId]/deny/route.ts rename to src/app/api/v2/ai/agent/sessions/[sessionId]/sandbox/suspend/route.ts index a5b46e40..56223e80 100644 --- a/src/app/api/v2/ai/agent/pending-actions/[actionId]/deny/route.ts +++ b/src/app/api/v2/ai/agent/sessions/[sessionId]/sandbox/suspend/route.ts @@ -17,27 +17,28 @@ import { NextRequest } from 'next/server'; import 'server/lib/dependencies'; import { createApiHandler } from 'server/lib/createApiHandler'; -import { errorResponse, successResponse } from 'server/lib/response'; import { getRequestUserIdentity } from 'server/lib/get-user'; -import ApprovalService from 'server/services/agent/ApprovalService'; +import { errorResponse, successResponse } from 'server/lib/response'; +import AgentSessionService, { ActiveAgentRunSuspensionError } from 'server/services/agentSession'; +import AgentSessionReadService from 'server/services/agent/SessionReadService'; /** * @openapi - * /api/v2/ai/agent/pending-actions/{actionId}/deny: + * /api/v2/ai/agent/sessions/{sessionId}/sandbox/suspend: * post: - * summary: Deny a pending action + * summary: Suspend a chat session sandbox runtime * tags: * - Agent Sessions - * operationId: denyAgentPendingAction + * operationId: suspendAgentSessionSandbox * parameters: * - in: path - * name: actionId + * name: sessionId * required: true * schema: * type: string * responses: * '200': - * description: Pending action denied + * description: Suspended session * content: * application/json: * schema: @@ -47,22 +48,33 @@ import ApprovalService from 'server/services/agent/ApprovalService'; * required: [data] * properties: * data: - * $ref: '#/components/schemas/AgentPendingAction' + * $ref: '#/components/schemas/AgentSessionSummary' + * '400': + * description: Session cannot be suspended + * '401': + * description: Unauthorized + * '409': + * description: Session has an active agent run */ -const postHandler = async (req: NextRequest, { params }: { params: { actionId: string } }) => { +const postHandler = async (req: NextRequest, { params }: { params: { sessionId: string } }) => { const userIdentity = getRequestUserIdentity(req); if (!userIdentity) { return errorResponse(new Error('Unauthorized'), { status: 401 }, req); } - const body = await req.json().catch(() => ({})); - const action = await ApprovalService.resolvePendingAction(params.actionId, userIdentity.userId, 'denied', { - approved: false, - reason: typeof body?.reason === 'string' ? body.reason : null, - source: 'endpoint', - }); + try { + const session = await AgentSessionService.suspendChatRuntime({ + sessionId: params.sessionId, + userId: userIdentity.userId, + }); - return successResponse(ApprovalService.serializePendingAction(action), { status: 200 }, req); + return successResponse(await AgentSessionReadService.serializeSessionRecord(session), { status: 200 }, req); + } catch (error) { + if (error instanceof ActiveAgentRunSuspensionError) { + return errorResponse(error, { status: 409 }, req); + } + return errorResponse(error, { status: 400 }, req); + } }; export const POST = createApiHandler(postHandler); diff --git a/src/app/api/v2/ai/agent/sessions/route.ts b/src/app/api/v2/ai/agent/sessions/route.ts index c0f1f7b1..36f2bb29 100644 --- a/src/app/api/v2/ai/agent/sessions/route.ts +++ b/src/app/api/v2/ai/agent/sessions/route.ts @@ -19,26 +19,22 @@ import 'server/lib/dependencies'; import { createApiHandler } from 'server/lib/createApiHandler'; import { successResponse, errorResponse } from 'server/lib/response'; import { getRequestUserIdentity } from 'server/lib/get-user'; -import { resolveRequestGitHubToken } from 'server/lib/agentSession/githubToken'; -import { - AgentSessionRuntimeConfigError, - mergeAgentSessionReadinessForServices, - mergeAgentSessionResources, - resolveAgentSessionRuntimeConfig, -} from 'server/lib/agentSession/runtimeConfig'; +import type { DevConfig } from 'server/models/yaml/YamlService'; +import type { LifecycleConfig } from 'server/models/yaml'; +import AgentChatSessionService from 'server/services/agent/ChatSessionService'; import { MissingAgentProviderApiKeyError } from 'server/services/agent/ProviderRegistry'; -import { AGENT_API_KEY_HEADER, AGENT_API_KEY_PROVIDER_HEADER } from 'server/services/agent/providerConfig'; -import AgentSessionService, { ActiveEnvironmentSessionError } from 'server/services/agentSession'; +import AgentSessionReadService from 'server/services/agent/SessionReadService'; import { - resolveAgentSessionServiceCandidatesForBuild, - resolveRequestedAgentSessionServices, - type RequestedAgentSessionServiceRef, -} from 'server/services/agentSessionCandidates'; -import { serializeAgentSessionSummary } from 'server/services/agent/serializeSessionSummary'; -import Build from 'server/models/Build'; -import { fetchLifecycleConfig, type LifecycleConfig } from 'server/models/yaml'; -import type { DevConfig } from 'server/models/yaml/YamlService'; -import { BuildKind } from 'shared/constants'; + DEFAULT_AGENT_SESSION_LIST_LIMIT, + MAX_AGENT_SESSION_LIST_LIMIT, +} from 'server/services/agent/SessionReadService'; +import { AgentSessionKind, BuildKind } from 'shared/constants'; + +interface RequestedAgentSessionServiceRef { + name: string; + repo?: string | null; + branch?: string | null; +} interface ResolvedSessionService { name: string; @@ -53,13 +49,17 @@ interface ResolvedSessionService { } interface CreateSessionBody { - buildUuid?: string; - services?: unknown[]; - model?: string; - repoUrl?: string; - branch?: string; - prNumber?: number; - namespace?: string; + defaults?: { + model?: string; + harness?: string; + }; + source?: { + adapter?: string; + input?: Record; + }; + workspace?: { + storageSize?: string; + }; } function repoNameFromRepoUrl(repoUrl?: string | null) { @@ -71,6 +71,27 @@ function repoNameFromRepoUrl(repoUrl?: string | null) { return normalized || null; } +function parseRequestedWorkspaceStorageSize(body: CreateSessionBody): string | undefined { + const workspace = body.workspace; + if (workspace === undefined) { + return undefined; + } + + if (!workspace || typeof workspace !== 'object' || Array.isArray(workspace)) { + throw new Error('workspace must be an object'); + } + + if (workspace.storageSize === undefined) { + return undefined; + } + + if (typeof workspace.storageSize !== 'string' || !workspace.storageSize.trim()) { + throw new Error('workspace.storageSize must be a non-empty string'); + } + + return workspace.storageSize.trim(); +} + async function resolveLifecycleConfigForSession({ buildContext, repoUrl, @@ -80,6 +101,8 @@ async function resolveLifecycleConfigForSession({ repoUrl?: string | null; branch?: string | null; }): Promise { + const { fetchLifecycleConfig } = await import('server/models/yaml'); + if (buildContext?.pullRequest?.fullName && buildContext.pullRequest.branchName) { return fetchLifecycleConfig(buildContext.pullRequest.fullName, buildContext.pullRequest.branchName); } @@ -115,6 +138,7 @@ function isRequestedSessionServiceRef(value: unknown): value is RequestedAgentSe } async function resolveBuildContext(buildUuid: string) { + const { default: Build } = await import('server/models/Build'); return Build.query() .findOne({ uuid: buildUuid }) .withGraphFetched('[pullRequest, deploys.[deployable, repository, service]]'); @@ -141,6 +165,10 @@ async function resolveRequestedServices( throw new Error('Build not found'); } + const { resolveAgentSessionServiceCandidatesForBuild, resolveRequestedAgentSessionServices } = await import( + 'server/services/agentSessionCandidates' + ); + const requestedRefs = requestedServices.map((service) => { if (typeof service === 'string') { return service; @@ -181,6 +209,21 @@ async function resolveRequestedServices( * schema: * type: boolean * description: When true, include ended and errored sessions in the response. + * - in: query + * name: page + * schema: + * type: integer + * default: 1 + * minimum: 1 + * description: Page number for pagination. + * - in: query + * name: limit + * schema: + * type: integer + * default: 25 + * minimum: 1 + * maximum: 100 + * description: Number of sessions per page. * responses: * '200': * description: Agent sessions @@ -195,146 +238,9 @@ async function resolveRequestedServices( * data: * type: array * items: - * type: object - * required: - * - id - * - buildUuid - * - baseBuildUuid - * - buildKind - * - userId - * - ownerGithubUsername - * - podName - * - namespace - * - model - * - status - * - repo - * - branch - * - services - * - lastActivity - * - createdAt - * - updatedAt - * - endedAt - * - startupFailure - * - editorUrl - * properties: - * id: - * type: string - * buildUuid: - * type: string - * nullable: true - * baseBuildUuid: - * type: string - * nullable: true - * buildKind: - * $ref: '#/components/schemas/BuildKind' - * userId: - * type: string - * ownerGithubUsername: - * type: string - * nullable: true - * podName: - * type: string - * namespace: - * type: string - * model: - * type: string - * status: - * type: string - * enum: [starting, active, ended, error] - * repo: - * type: string - * nullable: true - * branch: - * type: string - * nullable: true - * primaryRepo: - * type: string - * nullable: true - * primaryBranch: - * type: string - * nullable: true - * workspaceRepos: - * type: array - * items: - * type: object - * required: [repo, repoUrl, branch, mountPath] - * properties: - * repo: - * type: string - * repoUrl: - * type: string - * branch: - * type: string - * revision: - * type: string - * nullable: true - * mountPath: - * type: string - * primary: - * type: boolean - * selectedServices: - * type: array - * items: - * type: object - * required: [name, deployId, repo, branch, workspacePath] - * properties: - * name: - * type: string - * deployId: - * type: integer - * repo: - * type: string - * branch: - * type: string - * revision: - * type: string - * nullable: true - * resourceName: - * type: string - * nullable: true - * workspacePath: - * type: string - * workDir: - * type: string - * nullable: true - * services: - * type: array - * items: - * type: string - * lastActivity: - * type: string - * format: date-time - * createdAt: - * type: string - * format: date-time - * updatedAt: - * type: string - * format: date-time - * endedAt: - * type: string - * nullable: true - * format: date-time - * startupFailure: - * type: object - * nullable: true - * required: - * - stage - * - title - * - message - * - recordedAt - * properties: - * stage: - * type: string - * enum: [create_session, connect_runtime, attach_services] - * title: - * type: string - * message: - * type: string - * recordedAt: - * type: string - * format: date-time - * editorUrl: - * type: string + * $ref: '#/components/schemas/AgentSessionSummary' + * metadata: + * $ref: '#/components/schemas/ResponseMetadata' * error: * nullable: true * '401': @@ -344,7 +250,7 @@ async function resolveRequestedServices( * schema: * $ref: '#/components/schemas/ApiErrorResponse' * post: - * summary: Create a new interactive agent session + * summary: Create a new agent session * tags: * - Agent Sessions * operationId: createAgentSession @@ -354,28 +260,43 @@ async function resolveRequestedServices( * application/json: * schema: * type: object - * required: - * - buildUuid + * required: [source] * properties: - * buildUuid: - * type: string - * services: - * type: array - * description: Optional service names or repo-qualified service references to enable dev mode for. - * items: - * oneOf: - * - type: string - * - type: object - * required: [name] - * properties: - * name: - * type: string - * repo: - * type: string - * branch: - * type: string - * model: - * type: string + * defaults: + * type: object + * properties: + * model: + * type: string + * harness: + * type: string + * source: + * type: object + * required: [adapter] + * properties: + * adapter: + * type: string + * input: + * type: object + * additionalProperties: true + * workspace: + * type: object + * properties: + * storageSize: + * type: string + * description: Optional workspace PVC size. Accepted only when admin runtime settings allow client overrides. + * sandbox: + * type: object + * properties: + * providerHint: + * type: string + * requirements: + * type: object + * additionalProperties: true + * thread: + * type: object + * properties: + * createDefault: + * type: boolean * responses: * '201': * description: Agent session created @@ -388,146 +309,7 @@ async function resolveRequestedServices( * request_id: * type: string * data: - * type: object - * required: - * - id - * - buildUuid - * - baseBuildUuid - * - buildKind - * - userId - * - ownerGithubUsername - * - podName - * - namespace - * - model - * - status - * - repo - * - branch - * - services - * - editorUrl - * - lastActivity - * - createdAt - * - updatedAt - * - endedAt - * - startupFailure - * properties: - * id: - * type: string - * buildUuid: - * type: string - * nullable: true - * baseBuildUuid: - * type: string - * nullable: true - * buildKind: - * $ref: '#/components/schemas/BuildKind' - * userId: - * type: string - * ownerGithubUsername: - * type: string - * nullable: true - * podName: - * type: string - * namespace: - * type: string - * model: - * type: string - * status: - * type: string - * enum: [starting, active, ended, error] - * repo: - * type: string - * nullable: true - * branch: - * type: string - * nullable: true - * primaryRepo: - * type: string - * nullable: true - * primaryBranch: - * type: string - * nullable: true - * workspaceRepos: - * type: array - * items: - * type: object - * required: [repo, repoUrl, branch, mountPath] - * properties: - * repo: - * type: string - * repoUrl: - * type: string - * branch: - * type: string - * revision: - * type: string - * nullable: true - * mountPath: - * type: string - * primary: - * type: boolean - * selectedServices: - * type: array - * items: - * type: object - * required: [name, deployId, repo, branch, workspacePath] - * properties: - * name: - * type: string - * deployId: - * type: integer - * repo: - * type: string - * branch: - * type: string - * revision: - * type: string - * nullable: true - * resourceName: - * type: string - * nullable: true - * workspacePath: - * type: string - * workDir: - * type: string - * nullable: true - * services: - * type: array - * items: - * type: string - * editorUrl: - * type: string - * lastActivity: - * type: string - * format: date-time - * createdAt: - * type: string - * format: date-time - * updatedAt: - * type: string - * format: date-time - * endedAt: - * type: string - * nullable: true - * format: date-time - * startupFailure: - * type: object - * nullable: true - * required: - * - stage - * - title - * - message - * - recordedAt - * properties: - * stage: - * type: string - * enum: [create_session, connect_runtime, attach_services] - * title: - * type: string - * message: - * type: string - * recordedAt: - * type: string - * format: date-time + * $ref: '#/components/schemas/AgentSessionSummary' * error: * nullable: true * '400': @@ -560,10 +342,26 @@ const getHandler = async (req: NextRequest) => { if (!userIdentity) return errorResponse(new Error('Unauthorized'), { status: 401 }, req); const includeEnded = req.nextUrl.searchParams.get('includeEnded') === 'true'; - const sessions = await AgentSessionService.getSessions(userIdentity.userId, { includeEnded }); + const page = parseInt(req.nextUrl.searchParams.get('page') || '1', 10); + const requestedLimit = parseInt( + req.nextUrl.searchParams.get('limit') || String(DEFAULT_AGENT_SESSION_LIST_LIMIT), + 10 + ); + const limit = Number.isFinite(requestedLimit) + ? Math.min(requestedLimit, MAX_AGENT_SESSION_LIST_LIMIT) + : DEFAULT_AGENT_SESSION_LIST_LIMIT; + const result = await AgentSessionReadService.listOwnedSessionRecords(userIdentity.userId, { + includeEnded, + page, + limit, + }); + return successResponse( - sessions.map((session) => serializeAgentSessionSummary(session)), - { status: 200 }, + result.records, + { + status: 200, + metadata: result.metadata, + }, req ); }; @@ -573,12 +371,89 @@ const postHandler = async (req: NextRequest) => { if (!userIdentity) return errorResponse(new Error('Unauthorized'), { status: 401 }, req); const body = (await req.json()) as CreateSessionBody; - const { buildUuid, services, model } = body; + let requestedWorkspaceStorageSize: string | undefined; + try { + requestedWorkspaceStorageSize = parseRequestedWorkspaceStorageSize(body); + } catch (err) { + return errorResponse(err, { status: 400 }, req); + } + const sourceInput = + body.source?.input && typeof body.source.input === 'object' && !Array.isArray(body.source.input) + ? body.source.input + : {}; + const buildUuid = + typeof (sourceInput as { buildUuid?: unknown }).buildUuid === 'string' + ? (sourceInput as { buildUuid: string }).buildUuid + : undefined; + const services = Array.isArray((sourceInput as { services?: unknown[] }).services) + ? (sourceInput as { services: unknown[] }).services + : undefined; + const requestedModel = body.defaults?.model; + const sessionKind = + body.source?.adapter === 'blank_workspace' + ? AgentSessionKind.CHAT + : body.source?.adapter === 'lifecycle_fork' + ? AgentSessionKind.SANDBOX + : AgentSessionKind.ENVIRONMENT; + + if (sessionKind === AgentSessionKind.CHAT) { + if (buildUuid || (Array.isArray(services) && services.length > 0)) { + return errorResponse( + new Error('Chat sessions cannot be created with buildUuid or attached services'), + { status: 400 }, + req + ); + } + + try { + const { resolveAgentSessionRuntimeConfig, resolveAgentSessionWorkspaceStorageIntent } = await import( + 'server/lib/agentSession/runtimeConfig' + ); + const workspaceStorage = requestedWorkspaceStorageSize + ? resolveAgentSessionWorkspaceStorageIntent({ + requestedSize: requestedWorkspaceStorageSize, + storage: (await resolveAgentSessionRuntimeConfig()).workspaceStorage, + }) + : undefined; + const session = await AgentChatSessionService.createChatSession({ + userId: userIdentity.userId, + userIdentity, + model: requestedModel, + workspaceStorage, + }); + + return successResponse(await AgentSessionReadService.serializeSessionRecord(session), { status: 201 }, req); + } catch (err) { + const { AgentSessionRuntimeConfigError, AgentSessionWorkspaceStorageConfigError } = await import( + 'server/lib/agentSession/runtimeConfig' + ); + if (err instanceof MissingAgentProviderApiKeyError) { + return errorResponse(err, { status: 400 }, req); + } + if (err instanceof AgentSessionRuntimeConfigError || err instanceof AgentSessionWorkspaceStorageConfigError) { + return errorResponse(err, { status: 400 }, req); + } + + return errorResponse(err, { status: 500 }, req); + } + } - let repoUrl = body.repoUrl; - let branch = body.branch; - let prNumber = body.prNumber; - let namespace = body.namespace; + let repoUrl = + typeof (sourceInput as { repoUrl?: unknown }).repoUrl === 'string' + ? (sourceInput as { repoUrl: string }).repoUrl + : undefined; + let branch = + typeof (sourceInput as { branch?: unknown }).branch === 'string' + ? (sourceInput as { branch: string }).branch + : undefined; + let prNumber = + typeof (sourceInput as { prNumber?: unknown }).prNumber === 'number' + ? (sourceInput as { prNumber: number }).prNumber + : undefined; + let namespace = + typeof (sourceInput as { namespace?: unknown }).namespace === 'string' + ? (sourceInput as { namespace: string }).namespace + : undefined; let buildKind = BuildKind.ENVIRONMENT; let buildContext: Awaited> | null = null; let lifecycleConfig: LifecycleConfig | null = null; @@ -618,18 +493,34 @@ const postHandler = async (req: NextRequest) => { } try { + const [ + { resolveRequestGitHubToken }, + { + mergeAgentSessionReadinessForServices, + mergeAgentSessionResources, + resolveAgentSessionRuntimeConfig, + resolveAgentSessionWorkspaceStorageIntent, + }, + { default: AgentSessionService }, + ] = await Promise.all([ + import('server/lib/agentSession/githubToken'), + import('server/lib/agentSession/runtimeConfig'), + import('server/services/agentSession'), + ]); const runtimeConfig = await resolveAgentSessionRuntimeConfig(); + const workspaceStorage = resolveAgentSessionWorkspaceStorageIntent({ + requestedSize: requestedWorkspaceStorageSize, + storage: runtimeConfig.workspaceStorage, + }); const githubToken = await resolveRequestGitHubToken(req); const session = await AgentSessionService.createSession({ userId: userIdentity.userId, userIdentity, githubToken, - requestApiKey: req.headers.get(AGENT_API_KEY_HEADER), - requestApiKeyProvider: req.headers.get(AGENT_API_KEY_PROVIDER_HEADER), buildUuid, buildKind, services: resolvedServices, - model, + model: requestedModel, environmentSkillRefs: lifecycleConfig?.environment?.agentSession?.skills, repoUrl, branch, @@ -648,21 +539,17 @@ const postHandler = async (req: NextRequest) => { runtimeConfig.resources, lifecycleConfig?.environment?.agentSession?.resources ), + workspaceStorage, + redisTtlSeconds: runtimeConfig.cleanup.redisTtlSeconds, }); - return successResponse( - serializeAgentSessionSummary({ - ...session, - baseBuildUuid: null, - repo: repoNameFromRepoUrl(repoUrl), - branch, - services: resolvedServices.map((service) => service.name), - startupFailure: null, - }), - { status: 201 }, - req - ); + return successResponse(await AgentSessionReadService.serializeSessionRecord(session), { status: 201 }, req); } catch (err) { + const { ActiveEnvironmentSessionError } = await import('server/services/agentSession'); + const { AgentSessionRuntimeConfigError, AgentSessionWorkspaceStorageConfigError } = await import( + 'server/lib/agentSession/runtimeConfig' + ); + if (err instanceof ActiveEnvironmentSessionError) { return errorResponse(err, { status: 409 }, req); } @@ -672,6 +559,9 @@ const postHandler = async (req: NextRequest) => { if (err instanceof AgentSessionRuntimeConfigError) { return errorResponse(err, { status: 503 }, req); } + if (err instanceof AgentSessionWorkspaceStorageConfigError) { + return errorResponse(err, { status: 400 }, req); + } throw err; } }; diff --git a/src/app/api/v2/ai/agent/threads/[threadId]/messages/route.test.ts b/src/app/api/v2/ai/agent/threads/[threadId]/messages/route.test.ts index 2b8d0d7d..3f3448fa 100644 --- a/src/app/api/v2/ai/agent/threads/[threadId]/messages/route.test.ts +++ b/src/app/api/v2/ai/agent/threads/[threadId]/messages/route.test.ts @@ -16,254 +16,110 @@ import { NextRequest } from 'next/server'; -const mockValidateUIMessages = jest.fn(); -const mockCreateAgentUIStream = jest.fn(); -const mockCreateUIMessageStream = jest.fn(); -const mockCreateUIMessageStreamResponse = jest.fn(); - -jest.mock('ai', () => ({ - validateUIMessages: (...args: unknown[]) => mockValidateUIMessages(...args), - createAgentUIStream: (...args: unknown[]) => mockCreateAgentUIStream(...args), - createUIMessageStream: (...args: unknown[]) => mockCreateUIMessageStream(...args), - createUIMessageStreamResponse: (...args: unknown[]) => mockCreateUIMessageStreamResponse(...args), -})); - jest.mock('server/lib/get-user', () => ({ getRequestUserIdentity: jest.fn(), })); -jest.mock('server/services/agent/ThreadService', () => ({ - __esModule: true, - default: { - getOwnedThreadWithSession: jest.fn(), - }, -})); - -jest.mock('server/services/agent/RunService', () => ({ - __esModule: true, - default: { - getLatestOwnedThreadRun: jest.fn(), - isTerminalStatus: jest.fn(), - }, -})); - -jest.mock('server/services/agent/ApprovalService', () => ({ - __esModule: true, - default: { - syncApprovalResponsesFromMessages: jest.fn(), - }, -})); - jest.mock('server/services/agent/MessageStore', () => ({ __esModule: true, + DEFAULT_AGENT_MESSAGE_PAGE_LIMIT: 50, + MAX_AGENT_MESSAGE_PAGE_LIMIT: 100, default: { - syncMessages: jest.fn(), - }, -})); - -jest.mock('server/services/agent/RunExecutor', () => ({ - __esModule: true, - default: { - execute: jest.fn(), - }, -})); - -jest.mock('server/services/agent/StreamBroker', () => ({ - __esModule: true, - default: { - attach: jest.fn(), - }, -})); - -jest.mock('server/services/agentSession', () => ({ - __esModule: true, - default: { - touchActivity: jest.fn(), + listCanonicalMessages: jest.fn(), }, })); -import { POST } from './route'; +import { GET } from './route'; import { getRequestUserIdentity } from 'server/lib/get-user'; -import AgentThreadService from 'server/services/agent/ThreadService'; -import AgentRunService from 'server/services/agent/RunService'; -import ApprovalService from 'server/services/agent/ApprovalService'; import AgentMessageStore from 'server/services/agent/MessageStore'; -import AgentRunExecutor from 'server/services/agent/RunExecutor'; -import AgentStreamBroker from 'server/services/agent/StreamBroker'; -import AgentSessionService from 'server/services/agentSession'; const mockGetRequestUserIdentity = getRequestUserIdentity as jest.Mock; -const mockGetOwnedThreadWithSession = AgentThreadService.getOwnedThreadWithSession as jest.Mock; -const mockGetLatestOwnedThreadRun = AgentRunService.getLatestOwnedThreadRun as jest.Mock; -const mockIsTerminalStatus = AgentRunService.isTerminalStatus as jest.Mock; -const mockSyncApprovalResponses = ApprovalService.syncApprovalResponsesFromMessages as jest.Mock; -const mockSyncMessages = AgentMessageStore.syncMessages as jest.Mock; -const mockExecute = AgentRunExecutor.execute as jest.Mock; -const mockAttachStream = AgentStreamBroker.attach as jest.Mock; -const mockTouchActivity = AgentSessionService.touchActivity as jest.Mock; +const mockListCanonicalMessages = AgentMessageStore.listCanonicalMessages as jest.Mock; -function makeRequest(body: unknown): NextRequest { +function makeRequest(url = 'http://localhost/api/v2/ai/agent/threads/thread-1/messages'): NextRequest { return { headers: new Headers([['x-request-id', 'req-test']]), - json: async () => body, + nextUrl: new URL(url), } as unknown as NextRequest; } -describe('POST /api/v2/ai/agent/threads/[threadId]/messages', () => { +describe('GET /api/v2/ai/agent/threads/[threadId]/messages', () => { beforeEach(() => { jest.clearAllMocks(); - mockGetRequestUserIdentity.mockReturnValue({ userId: 'sample-user', + githubUsername: 'sample-user', }); - mockGetOwnedThreadWithSession.mockResolvedValue({ + mockListCanonicalMessages.mockResolvedValue({ thread: { - id: 10, - uuid: '7a972d88-3c05-4b80-93b9-7c420fb4d67d', - }, - session: { - id: 20, - uuid: '8fcaebb5-9392-4a81-a2ea-6dd367afd9f7', - status: 'active', + id: 'thread-1', + sessionId: 'session-1', + title: null, + isDefault: true, + archivedAt: null, + lastRunAt: null, + metadata: {}, + createdAt: null, + updatedAt: null, }, - }); - mockValidateUIMessages.mockImplementation(async ({ messages }) => messages); - mockIsTerminalStatus.mockImplementation( - (status) => status === 'completed' || status === 'failed' || status === 'cancelled' - ); - mockSyncApprovalResponses.mockResolvedValue(undefined); - mockSyncMessages.mockImplementation(async (_threadId, _userId, messages) => messages); - mockCreateAgentUIStream.mockResolvedValue(new ReadableStream()); - mockCreateUIMessageStream.mockReturnValue(new ReadableStream()); - mockCreateUIMessageStreamResponse.mockReturnValue(new Response('ok', { status: 200 })); - mockExecute.mockResolvedValue({ - agent: { tools: {} }, - abortSignal: new AbortController().signal, - onStreamFinish: jest.fn(), - selection: { - provider: 'openai', - modelId: 'gpt-5.4', - }, - run: { - uuid: 'b2e0f5e1-3342-4c83-a2f2-37d3104fc8e4', + messages: [ + { + id: 'message-1', + clientMessageId: 'client-message-1', + threadId: 'thread-1', + runId: 'run-1', + role: 'user', + parts: [{ type: 'text', text: 'Hi' }], + createdAt: '2026-04-25T00:00:00.000Z', + }, + ], + pagination: { + hasMore: false, + nextBeforeMessageId: null, }, }); }); - it('rejects a new user turn when the thread already has an active run', async () => { - mockGetLatestOwnedThreadRun.mockResolvedValue({ - status: 'running', - }); - - const response = await POST( - makeRequest({ - messages: [ - { - id: 'user-1', - role: 'user', - parts: [{ type: 'text', text: 'Inspect the workspace.' }], - }, - ], - }), - { - params: { threadId: '7a972d88-3c05-4b80-93b9-7c420fb4d67d' }, - } + it('returns canonical messages with cursor options', async () => { + const response = await GET( + makeRequest('http://localhost/api/v2/ai/agent/threads/thread-1/messages?limit=25&beforeMessageId=message-2'), + { params: { threadId: 'thread-1' } } ); - const body = await response.json(); - expect(response.status).toBe(409); - expect(body.error.message).toBe('Wait for the current agent run to finish before sending another message.'); - expect(mockExecute).not.toHaveBeenCalled(); - expect(mockCreateAgentUIStream).not.toHaveBeenCalled(); - }); - - it('rejects new messages while the session is still starting', async () => { - mockGetOwnedThreadWithSession.mockResolvedValue({ - thread: { - id: 10, - uuid: '7a972d88-3c05-4b80-93b9-7c420fb4d67d', - }, - session: { - id: 20, - uuid: '8fcaebb5-9392-4a81-a2ea-6dd367afd9f7', - status: 'starting', - }, + expect(response.status).toBe(200); + expect(mockListCanonicalMessages).toHaveBeenCalledWith('thread-1', 'sample-user', { + limit: 25, + beforeMessageId: 'message-2', }); - - const response = await POST( - makeRequest({ - messages: [ - { - id: 'user-1', - role: 'user', - parts: [{ type: 'text', text: 'Inspect the workspace.' }], - }, - ], + expect(body.data.messages).toEqual([ + expect.objectContaining({ + id: 'message-1', + clientMessageId: 'client-message-1', + role: 'user', + parts: [{ type: 'text', text: 'Hi' }], }), - { - params: { threadId: '7a972d88-3c05-4b80-93b9-7c420fb4d67d' }, - } - ); + ]); + }); + it('rejects invalid limits', async () => { + const response = await GET(makeRequest('http://localhost/api/v2/ai/agent/threads/thread-1/messages?limit=0'), { + params: { threadId: 'thread-1' }, + }); const body = await response.json(); - expect(response.status).toBe(409); - expect(body.error.message).toBe('Wait for the session to finish starting before sending a message.'); - expect(mockSyncApprovalResponses).not.toHaveBeenCalled(); - expect(mockSyncMessages).not.toHaveBeenCalled(); - expect(mockExecute).not.toHaveBeenCalled(); - expect(mockCreateAgentUIStream).not.toHaveBeenCalled(); + expect(response.status).toBe(400); + expect(body.error.message).toBe('Expected a positive integer limit.'); + expect(mockListCanonicalMessages).not.toHaveBeenCalled(); }); - it('allows assistant approval continuations to reuse the waiting run', async () => { - mockGetLatestOwnedThreadRun.mockResolvedValue({ - status: 'waiting_for_approval', - }); + it('maps missing threads to 404', async () => { + mockListCanonicalMessages.mockRejectedValueOnce(new Error('Agent thread not found')); - const response = await POST( - makeRequest({ - messages: [ - { - id: 'assistant-1', - role: 'assistant', - metadata: { - runId: 'b2e0f5e1-3342-4c83-a2f2-37d3104fc8e4', - }, - parts: [ - { - type: 'tool-workspace_exec_mutation', - toolCallId: 'tool-1', - state: 'approval-responded', - approval: { - id: 'approval-1', - approved: true, - }, - }, - ], - }, - ], - }), - { - params: { threadId: '7a972d88-3c05-4b80-93b9-7c420fb4d67d' }, - } - ); + const response = await GET(makeRequest(), { params: { threadId: 'missing-thread' } }); + const body = await response.json(); - expect(response.status).toBe(200); - expect(mockSyncApprovalResponses).toHaveBeenCalled(); - expect(mockSyncMessages).toHaveBeenCalled(); - expect(mockTouchActivity).toHaveBeenCalledWith('8fcaebb5-9392-4a81-a2ea-6dd367afd9f7'); - expect(mockExecute).toHaveBeenCalled(); - expect(mockCreateAgentUIStream).toHaveBeenCalledWith( - expect.objectContaining({ - generateMessageId: expect.any(Function), - }) - ); - expect(mockCreateUIMessageStream).toHaveBeenCalledWith( - expect.objectContaining({ - execute: expect.any(Function), - }) - ); - expect(mockAttachStream).toHaveBeenCalledWith('b2e0f5e1-3342-4c83-a2f2-37d3104fc8e4', expect.any(ReadableStream)); + expect(response.status).toBe(404); + expect(body.error.message).toBe('Agent thread not found'); }); }); diff --git a/src/app/api/v2/ai/agent/threads/[threadId]/messages/route.ts b/src/app/api/v2/ai/agent/threads/[threadId]/messages/route.ts index eb64db2c..69eb1d19 100644 --- a/src/app/api/v2/ai/agent/threads/[threadId]/messages/route.ts +++ b/src/app/api/v2/ai/agent/threads/[threadId]/messages/route.ts @@ -14,77 +14,34 @@ * limitations under the License. */ -import { - createAgentUIStream, - createUIMessageStream, - createUIMessageStreamResponse, - type UIMessageChunk, - validateUIMessages, -} from 'ai'; import { NextRequest } from 'next/server'; import 'server/lib/dependencies'; import { createApiHandler } from 'server/lib/createApiHandler'; import { errorResponse, successResponse } from 'server/lib/response'; import { getRequestUserIdentity } from 'server/lib/get-user'; -import AgentMessageStore from 'server/services/agent/MessageStore'; -import { buildMessageObservabilityMetadataPatch, normalizeSdkUsageSummary } from 'server/services/agent/observability'; -import AgentRunService from 'server/services/agent/RunService'; -import AgentThreadService from 'server/services/agent/ThreadService'; -import ApprovalService from 'server/services/agent/ApprovalService'; -import AgentRunExecutor from 'server/services/agent/RunExecutor'; -import AgentStreamBroker from 'server/services/agent/StreamBroker'; -import AgentSessionService from 'server/services/agentSession'; -import { applyApprovalResponsesToFileChangeParts } from 'server/services/agent/fileChanges'; -import type { AgentUIDataParts, AgentUIMessage, AgentUIMessageMetadata } from 'server/services/agent/types'; -import { MissingAgentProviderApiKeyError } from 'server/services/agent/ProviderRegistry'; -import { AGENT_API_KEY_HEADER, AGENT_API_KEY_PROVIDER_HEADER } from 'server/services/agent/providerConfig'; - -type AgentUiMessageChunk = UIMessageChunk; - -function createChunkStream() { - let controller: ReadableStreamDefaultController | null = null; - let closed = false; - - return { - stream: new ReadableStream({ - start(nextController) { - controller = nextController; - if (closed) { - nextController.close(); - controller = null; - } - }, - cancel() { - closed = true; - controller = null; - }, - }), - write(chunk: AgentUiMessageChunk) { - if (closed || !controller) { - return; - } +import AgentMessageStore, { + DEFAULT_AGENT_MESSAGE_PAGE_LIMIT, + MAX_AGENT_MESSAGE_PAGE_LIMIT, +} from 'server/services/agent/MessageStore'; + +function parseLimit(value: string | null): number { + if (value == null || value.trim() === '') { + return DEFAULT_AGENT_MESSAGE_PAGE_LIMIT; + } - controller.enqueue(chunk); - }, - close() { - if (closed) { - return; - } + const parsed = Number(value); + if (!Number.isInteger(parsed) || parsed <= 0) { + throw new Error('Expected a positive integer limit.'); + } - closed = true; - if (controller) { - controller.close(); - controller = null; - } - }, - }; + return Math.min(parsed, MAX_AGENT_MESSAGE_PAGE_LIMIT); } /** * @openapi * /api/v2/ai/agent/threads/{threadId}/messages: * get: - * summary: List persisted messages for an agent thread + * summary: List canonical messages for an agent thread * tags: * - Agent Sessions * operationId: getAgentThreadMessages @@ -94,9 +51,23 @@ function createChunkStream() { * required: true * schema: * type: string + * - in: query + * name: limit + * required: false + * schema: + * type: integer + * minimum: 1 + * maximum: 100 + * default: 50 + * - in: query + * name: beforeMessageId + * required: false + * schema: + * type: string + * description: Return messages older than this message id. * responses: * '200': - * description: Thread and persisted messages + * description: Canonical thread messages * content: * application/json: * schema: @@ -106,61 +77,25 @@ function createChunkStream() { * required: [data] * properties: * data: - * type: object - * required: [thread, messages] - * properties: - * thread: - * $ref: '#/components/schemas/AgentThread' - * messages: - * type: array - * items: - * $ref: '#/components/schemas/AgentUIMessage' - * post: - * summary: Send messages to an agent thread and stream the response - * tags: - * - Agent Sessions - * operationId: postAgentThreadMessages - * parameters: - * - in: path - * name: threadId - * required: true - * schema: - * type: string - * requestBody: - * required: true - * content: - * application/json: - * schema: - * type: object - * required: [messages] - * properties: - * provider: - * type: string - * modelId: - * type: string - * messages: - * type: array - * items: - * $ref: '#/components/schemas/AgentUIMessage' - * responses: + * $ref: '#/components/schemas/AgentThreadMessagesResponse' * '400': - * description: Invalid request or missing stored provider API key + * description: Invalid message cursor or page size * content: * application/json: * schema: * $ref: '#/components/schemas/ApiErrorResponse' - * '409': - * description: Session is still starting or another run is already active for this thread + * '401': + * description: Unauthorized * content: * application/json: * schema: * $ref: '#/components/schemas/ApiErrorResponse' - * '200': - * description: UI message stream + * '404': + * description: Thread not found * content: - * text/event-stream: + * application/json: * schema: - * type: string + * $ref: '#/components/schemas/ApiErrorResponse' */ const getHandler = async (req: NextRequest, { params }: { params: { threadId: string } }) => { const userIdentity = getRequestUserIdentity(req); @@ -168,263 +103,32 @@ const getHandler = async (req: NextRequest, { params }: { params: { threadId: st return errorResponse(new Error('Unauthorized'), { status: 401 }, req); } - const { thread, session } = await AgentThreadService.getOwnedThreadWithSession(params.threadId, userIdentity.userId); - const messages = await AgentMessageStore.listMessages(params.threadId, userIdentity.userId); - - return successResponse( - { - thread: AgentThreadService.serializeThread(thread, session.uuid), - messages, - }, - { status: 200 }, - req - ); -}; - -const postHandler = async (req: NextRequest, { params }: { params: { threadId: string } }): Promise => { - const userIdentity = getRequestUserIdentity(req); - if (!userIdentity) { - return errorResponse(new Error('Unauthorized'), { status: 401 }, req); - } - - const body = await req.json().catch(() => ({})); - if (!Array.isArray(body?.messages)) { - return errorResponse(new Error('messages array is required'), { status: 400 }, req); - } - let submittedMessages: AgentUIMessage[]; + let limit; try { - submittedMessages = await validateUIMessages({ - messages: body.messages, - }); + limit = parseLimit(req.nextUrl.searchParams.get('limit')); } catch (error) { - return errorResponse(error instanceof Error ? error : new Error('Invalid UI messages'), { status: 400 }, req); + return errorResponse(error, { status: 400 }, req); } - const { thread, session } = await AgentThreadService.getOwnedThreadWithSession(params.threadId, userIdentity.userId); - if (session.status !== 'active') { - return errorResponse( - new Error( - session.status === 'starting' - ? 'Wait for the session to finish starting before sending a message.' - : 'This session is no longer available for new messages.' - ), - { status: 409 }, - req - ); - } - const latestSubmittedMessage = submittedMessages[submittedMessages.length - 1]; - const isNewUserTurn = latestSubmittedMessage?.role === 'user'; - if (isNewUserTurn) { - const latestRun = await AgentRunService.getLatestOwnedThreadRun(params.threadId, userIdentity.userId); - if (latestRun && !AgentRunService.isTerminalStatus(latestRun.status)) { - return errorResponse( - new Error( - latestRun.status === 'waiting_for_approval' - ? 'Respond to the approval request before sending another message.' - : 'Wait for the current agent run to finish before sending another message.' - ), - { status: 409 }, - req - ); - } - } - - const normalizedSubmittedMessages = applyApprovalResponsesToFileChangeParts(submittedMessages); - await ApprovalService.syncApprovalResponsesFromMessages( - params.threadId, - userIdentity.userId, - normalizedSubmittedMessages - ); - const syncedMessages = await AgentMessageStore.syncMessages( - params.threadId, - userIdentity.userId, - normalizedSubmittedMessages - ); - await AgentSessionService.touchActivity(session.uuid); - const fileChangeStream = createChunkStream(); - let execution: Awaited>; try { - execution = await AgentRunExecutor.execute({ - session, - thread, - userIdentity, - messages: syncedMessages, - requestedProvider: typeof body?.provider === 'string' ? body.provider : undefined, - requestedModelId: typeof body?.modelId === 'string' ? body.modelId : undefined, - requestApiKey: req.headers.get(AGENT_API_KEY_HEADER), - requestApiKeyProvider: req.headers.get(AGENT_API_KEY_PROVIDER_HEADER), - onFileChange: async (change) => { - fileChangeStream.write({ - type: 'data-file-change', - id: change.id, - data: change, - }); - }, + const result = await AgentMessageStore.listCanonicalMessages(params.threadId, userIdentity.userId, { + limit, + beforeMessageId: req.nextUrl.searchParams.get('beforeMessageId'), }); + + return successResponse(result, { status: 200 }, req); } catch (error) { - if (error instanceof MissingAgentProviderApiKeyError) { + if ( + error instanceof Error && + (error.message === 'Agent thread not found' || error.message === 'Agent session not found') + ) { + return errorResponse(error, { status: 404 }, req); + } + if (error instanceof Error && error.message === 'Agent message cursor not found') { return errorResponse(error, { status: 400 }, req); } - throw error; } - - let finishContext: { - finishReason?: string; - isAborted: boolean; - } = { - finishReason: undefined, - isAborted: false, - }; - - const agentUiMessageStream = await createAgentUIStream< - never, - typeof execution.agent.tools, - never, - AgentUIMessageMetadata - >({ - agent: execution.agent, - uiMessages: syncedMessages, - generateMessageId: () => crypto.randomUUID(), - abortSignal: execution.abortSignal, - onFinish: async ({ finishReason, isAborted }) => { - finishContext = { - finishReason, - isAborted, - }; - fileChangeStream.close(); - }, - messageMetadata: ({ part }) => { - const eventType = (part as { type?: string }).type; - if (eventType === 'start') { - return { - sessionId: session.uuid, - threadId: thread.uuid, - runId: execution.run.uuid, - provider: execution.selection.provider, - model: execution.selection.modelId, - createdAt: new Date().toISOString(), - }; - } - - if (eventType === 'finish-step') { - const response = (part as { response?: unknown }).response as - | { - id?: string; - modelId?: string; - timestamp?: string | Date | number; - } - | undefined; - const providerMetadata = (part as { providerMetadata?: unknown }).providerMetadata as - | Record - | undefined; - const warningCount = Array.isArray((part as { warnings?: unknown[] }).warnings) - ? (part as { warnings: unknown[] }).warnings.length - : undefined; - - return { - ...(response?.id ? { responseId: response.id } : {}), - ...(response?.modelId - ? { - responseModelId: response.modelId, - model: response.modelId, - } - : {}), - ...(response?.timestamp - ? { - responseTimestamp: - response.timestamp instanceof Date - ? response.timestamp.toISOString() - : typeof response.timestamp === 'number' - ? new Date(response.timestamp).toISOString() - : response.timestamp, - } - : {}), - ...(providerMetadata ? { providerMetadata } : {}), - ...(warningCount != null ? { warningCount } : {}), - }; - } - - if (eventType === 'finish') { - const totalUsage = - ( - part as { - totalUsage?: { - inputTokens?: number; - outputTokens?: number; - totalTokens?: number; - reasoningTokens?: number; - cachedInputTokens?: number; - inputTokenDetails?: { - cacheReadTokens?: number; - cacheWriteTokens?: number; - noCacheTokens?: number; - }; - outputTokenDetails?: { - reasoningTokens?: number; - textTokens?: number; - }; - raw?: unknown; - }; - finishReason?: string; - rawFinishReason?: string; - } - ).totalUsage ?? undefined; - const usageSummary = totalUsage - ? normalizeSdkUsageSummary({ - usage: totalUsage, - finishReason: - typeof (part as { finishReason?: unknown }).finishReason === 'string' - ? (part as { finishReason: string }).finishReason - : undefined, - rawFinishReason: - typeof (part as { rawFinishReason?: unknown }).rawFinishReason === 'string' - ? (part as { rawFinishReason: string }).rawFinishReason - : undefined, - }) - : undefined; - - return { - sessionId: session.uuid, - threadId: thread.uuid, - runId: execution.run.uuid, - provider: execution.selection.provider, - model: execution.selection.modelId, - completedAt: new Date().toISOString(), - ...(usageSummary ? buildMessageObservabilityMetadataPatch(usageSummary) : {}), - }; - } - - return undefined; - }, - }); - - const uiMessageStream = createUIMessageStream({ - originalMessages: syncedMessages, - generateId: () => crypto.randomUUID(), - execute: ({ writer }) => { - writer.merge(agentUiMessageStream as ReadableStream); - writer.merge(fileChangeStream.stream); - }, - onFinish: async ({ messages }) => { - await execution.onStreamFinish({ - messages, - finishReason: finishContext.finishReason, - isAborted: finishContext.isAborted, - }); - }, - }); - - const [responseStream, replayStream] = uiMessageStream.tee() as [ - ReadableStream, - ReadableStream - ]; - AgentStreamBroker.attach(execution.run.uuid, replayStream); - - return createUIMessageStreamResponse({ - stream: responseStream, - }); }; export const GET = createApiHandler(getHandler); -export const POST = postHandler; diff --git a/src/app/api/v2/ai/agent/threads/[threadId]/pending-actions/route.test.ts b/src/app/api/v2/ai/agent/threads/[threadId]/pending-actions/route.test.ts new file mode 100644 index 00000000..4bc379ee --- /dev/null +++ b/src/app/api/v2/ai/agent/threads/[threadId]/pending-actions/route.test.ts @@ -0,0 +1,131 @@ +/** + * Copyright 2026 GoodRx, Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +import { NextRequest } from 'next/server'; + +jest.mock('server/lib/get-user', () => ({ + getRequestUserIdentity: jest.fn(), +})); + +jest.mock('server/services/agent/ApprovalService', () => ({ + __esModule: true, + default: { + listPendingActions: jest.fn(), + serializePendingAction: jest.fn((action) => action), + }, +})); + +import { GET } from './route'; +import { getRequestUserIdentity } from 'server/lib/get-user'; +import ApprovalService from 'server/services/agent/ApprovalService'; + +const mockGetRequestUserIdentity = getRequestUserIdentity as jest.Mock; +const mockListPendingActions = ApprovalService.listPendingActions as jest.Mock; +const mockSerializePendingAction = ApprovalService.serializePendingAction as jest.Mock; + +function makeRequest(): NextRequest { + return { + headers: new Headers([['x-request-id', 'req-test']]), + nextUrl: new URL('http://localhost/api/v2/ai/agent/threads/thread-1/pending-actions'), + } as unknown as NextRequest; +} + +describe('GET /api/v2/ai/agent/threads/[threadId]/pending-actions', () => { + beforeEach(() => { + jest.clearAllMocks(); + }); + + it('returns 401 when the requester is not authenticated', async () => { + mockGetRequestUserIdentity.mockReturnValue(null); + + const response = await GET(makeRequest(), { params: { threadId: 'thread-1' } }); + + expect(response.status).toBe(401); + await expect(response.json()).resolves.toMatchObject({ + error: { message: 'Unauthorized' }, + }); + expect(mockListPendingActions).not.toHaveBeenCalled(); + }); + + it('returns canonical pending action display payloads for the owned thread', async () => { + mockGetRequestUserIdentity.mockReturnValue({ + userId: 'sample-user', + githubUsername: 'sample-user', + }); + mockListPendingActions.mockResolvedValue([ + { + id: 'action-db-id', + publicPayload: true, + }, + ]); + mockSerializePendingAction.mockReturnValue({ + id: 'action-1', + kind: 'tool_approval', + status: 'pending', + threadId: 'thread-1', + runId: 'run-1', + title: 'Approve workspace edit', + description: 'A workspace edit requires approval.', + requestedAt: '2026-04-11T00:00:00.000Z', + expiresAt: null, + toolName: 'mcp__sandbox__workspace_edit_file', + argumentsSummary: [{ name: 'path', value: 'sample-file.txt' }], + commandPreview: null, + fileChangePreview: [{ path: 'sample-file.txt', action: 'edited', summary: 'Updated sample-file.txt' }], + riskLabels: ['Workspace write'], + }); + + const response = await GET(makeRequest(), { params: { threadId: 'thread-1' } }); + const body = await response.json(); + + expect(response.status).toBe(200); + expect(mockListPendingActions).toHaveBeenCalledWith('thread-1', 'sample-user'); + expect(body.data.pendingActions).toEqual([ + { + id: 'action-1', + kind: 'tool_approval', + status: 'pending', + threadId: 'thread-1', + runId: 'run-1', + title: 'Approve workspace edit', + description: 'A workspace edit requires approval.', + requestedAt: '2026-04-11T00:00:00.000Z', + expiresAt: null, + toolName: 'mcp__sandbox__workspace_edit_file', + argumentsSummary: [{ name: 'path', value: 'sample-file.txt' }], + commandPreview: null, + fileChangePreview: [{ path: 'sample-file.txt', action: 'edited', summary: 'Updated sample-file.txt' }], + riskLabels: ['Workspace write'], + }, + ]); + }); + + it('returns 404 when the thread is not owned by the requester', async () => { + mockGetRequestUserIdentity.mockReturnValue({ + userId: 'sample-user', + githubUsername: 'sample-user', + }); + mockListPendingActions.mockRejectedValue(new Error('Agent thread not found')); + + const response = await GET(makeRequest(), { params: { threadId: 'missing-thread' } }); + + expect(response.status).toBe(404); + await expect(response.json()).resolves.toMatchObject({ + error: { message: 'Agent thread not found' }, + }); + expect(mockListPendingActions).toHaveBeenCalledWith('missing-thread', 'sample-user'); + }); +}); diff --git a/src/app/api/v2/ai/agent/threads/[threadId]/pending-actions/route.ts b/src/app/api/v2/ai/agent/threads/[threadId]/pending-actions/route.ts index 9edc6ad4..b1fc2b54 100644 --- a/src/app/api/v2/ai/agent/threads/[threadId]/pending-actions/route.ts +++ b/src/app/api/v2/ai/agent/threads/[threadId]/pending-actions/route.ts @@ -54,6 +54,18 @@ import ApprovalService from 'server/services/agent/ApprovalService'; * type: array * items: * $ref: '#/components/schemas/AgentPendingAction' + * '401': + * description: Unauthorized + * content: + * application/json: + * schema: + * $ref: '#/components/schemas/ApiErrorResponse' + * '404': + * description: Thread not found + * content: + * application/json: + * schema: + * $ref: '#/components/schemas/ApiErrorResponse' */ const getHandler = async (req: NextRequest, { params }: { params: { threadId: string } }) => { const userIdentity = getRequestUserIdentity(req); @@ -61,14 +73,21 @@ const getHandler = async (req: NextRequest, { params }: { params: { threadId: st return errorResponse(new Error('Unauthorized'), { status: 401 }, req); } - const pendingActions = await ApprovalService.listPendingActions(params.threadId, userIdentity.userId); - return successResponse( - { - pendingActions: pendingActions.map((action) => ApprovalService.serializePendingAction(action)), - }, - { status: 200 }, - req - ); + try { + const pendingActions = await ApprovalService.listPendingActions(params.threadId, userIdentity.userId); + return successResponse( + { + pendingActions: pendingActions.map((action) => ApprovalService.serializePendingAction(action)), + }, + { status: 200 }, + req + ); + } catch (error) { + if (error instanceof Error && error.message === 'Agent thread not found') { + return errorResponse(error, { status: 404 }, req); + } + throw error; + } }; export const GET = createApiHandler(getHandler); diff --git a/src/app/api/v2/ai/agent/threads/[threadId]/runs/route.test.ts b/src/app/api/v2/ai/agent/threads/[threadId]/runs/route.test.ts new file mode 100644 index 00000000..d806a3ad --- /dev/null +++ b/src/app/api/v2/ai/agent/threads/[threadId]/runs/route.test.ts @@ -0,0 +1,505 @@ +/** + * Copyright 2026 GoodRx, Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +import { NextRequest } from 'next/server'; + +jest.mock('server/lib/get-user', () => ({ + getRequestUserIdentity: jest.fn(), +})); + +jest.mock('server/lib/agentSession/githubToken', () => ({ + resolveRequestGitHubToken: jest.fn(), +})); + +jest.mock('server/services/agent/CapabilityService', () => ({ + __esModule: true, + default: { + resolveSessionContext: jest.fn(), + }, +})); + +jest.mock('server/services/agent/RunAdmissionService', () => ({ + __esModule: true, + default: { + createQueuedRunWithMessage: jest.fn(), + }, +})); + +jest.mock('server/services/agent/MessageStore', () => ({ + __esModule: true, + default: { + serializeCanonicalMessage: jest.fn((message) => ({ + id: message.uuid, + clientMessageId: message.clientMessageId || null, + threadId: 'thread-1', + runId: 'run-1', + role: message.role, + parts: message.parts, + createdAt: message.createdAt || null, + })), + }, +})); + +jest.mock('server/services/agent/ProviderRegistry', () => ({ + __esModule: true, + default: { + resolveSelection: jest.fn(), + }, +})); + +jest.mock('server/services/agent/RunQueueService', () => ({ + __esModule: true, + default: { + enqueueRun: jest.fn(), + }, +})); + +jest.mock('server/services/agent/RunService', () => ({ + __esModule: true, + default: { + isActiveRunConflictError: jest.fn(), + markFailed: jest.fn(), + markQueuedRunDispatchFailed: jest.fn(), + serializeRun: jest.fn((run) => ({ id: run.uuid, status: run.status })), + }, + InvalidAgentRunDefaultsError: class InvalidAgentRunDefaultsError extends Error {}, +})); + +jest.mock('server/services/agent/SourceService', () => ({ + __esModule: true, + default: { + getSessionSource: jest.fn(), + }, +})); + +jest.mock('server/services/agent/ThreadService', () => ({ + __esModule: true, + default: { + getOwnedThreadWithSession: jest.fn(), + }, +})); + +jest.mock('server/services/agentSession', () => ({ + __esModule: true, + default: { + canAcceptMessages: jest.fn(), + getMessageBlockReason: jest.fn(), + touchActivity: jest.fn(), + }, +})); + +import { POST } from './route'; +import { getRequestUserIdentity } from 'server/lib/get-user'; +import { resolveRequestGitHubToken } from 'server/lib/agentSession/githubToken'; +import AgentCapabilityService from 'server/services/agent/CapabilityService'; +import AgentProviderRegistry from 'server/services/agent/ProviderRegistry'; +import AgentRunAdmissionService from 'server/services/agent/RunAdmissionService'; +import AgentRunQueueService from 'server/services/agent/RunQueueService'; +import AgentRunService from 'server/services/agent/RunService'; +import AgentSourceService from 'server/services/agent/SourceService'; +import AgentThreadService from 'server/services/agent/ThreadService'; +import AgentSessionService from 'server/services/agentSession'; + +const mockGetRequestUserIdentity = getRequestUserIdentity as jest.Mock; +const mockResolveRequestGitHubToken = resolveRequestGitHubToken as jest.Mock; +const mockResolveSessionContext = AgentCapabilityService.resolveSessionContext as jest.Mock; +const mockResolveSelection = AgentProviderRegistry.resolveSelection as jest.Mock; +const mockCreateQueuedRunWithMessage = AgentRunAdmissionService.createQueuedRunWithMessage as jest.Mock; +const mockEnqueueRun = AgentRunQueueService.enqueueRun as jest.Mock; +const mockMarkQueuedRunDispatchFailed = AgentRunService.markQueuedRunDispatchFailed as jest.Mock; +const mockGetSessionSource = AgentSourceService.getSessionSource as jest.Mock; +const mockGetOwnedThreadWithSession = AgentThreadService.getOwnedThreadWithSession as jest.Mock; +const mockCanAcceptMessages = AgentSessionService.canAcceptMessages as jest.Mock; +const mockTouchActivity = AgentSessionService.touchActivity as jest.Mock; + +function makeRequest(body: Record): NextRequest { + return { + json: jest.fn().mockResolvedValue(body), + headers: new Headers([['x-request-id', 'req-test']]), + nextUrl: new URL('http://localhost/api/v2/ai/agent/threads/thread-1/runs'), + } as unknown as NextRequest; +} + +describe('POST /api/v2/ai/agent/threads/[threadId]/runs', () => { + beforeEach(() => { + jest.clearAllMocks(); + mockGetRequestUserIdentity.mockReturnValue({ + userId: 'sample-user', + githubUsername: 'sample-user', + }); + mockResolveRequestGitHubToken.mockResolvedValue('sample-gh-token'); + mockGetOwnedThreadWithSession.mockResolvedValue({ + thread: { id: 7, uuid: 'thread-1' }, + session: { + id: 17, + uuid: 'session-1', + defaultHarness: 'lifecycle_ai_sdk', + defaultModel: 'gpt-5.4', + }, + }); + mockCanAcceptMessages.mockReturnValue(true); + mockGetSessionSource.mockResolvedValue({ + status: 'ready', + sandboxRequirements: { filesystem: 'persistent' }, + }); + mockResolveSessionContext.mockResolvedValue({ + repoFullName: 'example-org/example-repo', + approvalPolicy: 'on-request', + }); + mockResolveSelection.mockResolvedValue({ + provider: 'openai', + modelId: 'gpt-5.4', + }); + mockCreateQueuedRunWithMessage.mockResolvedValue({ + run: { + uuid: 'run-1', + status: 'queued', + }, + message: { + uuid: 'message-1', + clientMessageId: 'client-message-1', + role: 'user', + parts: [{ type: 'text', text: 'Hi' }], + }, + created: true, + }); + mockTouchActivity.mockResolvedValue(undefined); + mockEnqueueRun.mockResolvedValue(undefined); + }); + + it('rejects run admission when no explicit or session model exists', async () => { + mockGetOwnedThreadWithSession.mockResolvedValueOnce({ + thread: { id: 7, uuid: 'thread-1' }, + session: { + id: 17, + uuid: 'session-1', + defaultHarness: 'lifecycle_ai_sdk', + defaultModel: null, + }, + }); + + const response = await POST( + makeRequest({ + message: { + clientMessageId: 'client-message-1', + parts: [{ type: 'text', text: 'Hi' }], + }, + }), + { params: { threadId: 'thread-1' } } + ); + const body = await response.json(); + + expect(response.status).toBe(400); + expect(body.error.message).toBe('Agent run model is required'); + expect(mockCreateQueuedRunWithMessage).not.toHaveBeenCalled(); + }); + + it('resolves explicit-or-default values before queueing', async () => { + const response = await POST( + makeRequest({ + message: { + clientMessageId: 'client-message-1', + parts: [{ type: 'text', text: 'Hi' }], + }, + runtimeOptions: { maxIterations: 12 }, + }), + { params: { threadId: 'thread-1' } } + ); + + expect(response.status).toBe(201); + expect(mockResolveSelection).toHaveBeenCalledWith({ + repoFullName: 'example-org/example-repo', + requestedProvider: undefined, + requestedModelId: 'gpt-5.4', + }); + expect(mockCreateQueuedRunWithMessage).toHaveBeenCalledWith( + expect.objectContaining({ + message: { + clientMessageId: 'client-message-1', + parts: [{ type: 'text', text: 'Hi' }], + }, + requestedHarness: null, + requestedProvider: null, + requestedModel: null, + resolvedHarness: 'lifecycle_ai_sdk', + resolvedProvider: 'openai', + resolvedModel: 'gpt-5.4', + runtimeOptions: { maxIterations: 12 }, + }) + ); + expect(mockEnqueueRun).toHaveBeenCalledWith('run-1', 'submit', { githubToken: 'sample-gh-token' }); + const body = await response.json(); + expect(body.data).toEqual( + expect.objectContaining({ + run: expect.objectContaining({ id: 'run-1', threadId: 'thread-1', sessionId: 'session-1' }), + message: expect.objectContaining({ id: 'message-1', clientMessageId: 'client-message-1' }), + links: { + events: '/api/v2/ai/agent/runs/run-1/events', + eventStream: '/api/v2/ai/agent/runs/run-1/events/stream', + pendingActions: '/api/v2/ai/agent/threads/thread-1/pending-actions', + }, + }) + ); + }); + + it('rejects tool or UI payload parts in canonical input messages', async () => { + const response = await POST( + makeRequest({ + message: { + clientMessageId: 'client-message-1', + parts: [ + { + type: 'dynamic-tool', + toolCallId: 'tool-call-1', + state: 'output-available', + }, + ], + }, + }), + { params: { threadId: 'thread-1' } } + ); + const body = await response.json(); + + expect(response.status).toBe(400); + expect(body.error.message).toBe('message must contain supported canonical parts and no role or metadata fields'); + expect(mockCreateQueuedRunWithMessage).not.toHaveBeenCalled(); + }); + + it('rejects extra canonical part fields instead of stripping them', async () => { + const response = await POST( + makeRequest({ + message: { + clientMessageId: 'client-message-1', + parts: [ + { + type: 'text', + text: 'Hi', + providerMetadata: { traceId: 'trace-1' }, + }, + ], + }, + }), + { params: { threadId: 'thread-1' } } + ); + const body = await response.json(); + + expect(response.status).toBe(400); + expect(body.error.message).toBe('message must contain supported canonical parts and no role or metadata fields'); + expect(mockCreateQueuedRunWithMessage).not.toHaveBeenCalled(); + }); + + it('rejects public message roles', async () => { + const response = await POST( + makeRequest({ + message: { + role: 'assistant', + parts: [{ type: 'text', text: 'Nope' }], + }, + }), + { params: { threadId: 'thread-1' } } + ); + const body = await response.json(); + + expect(response.status).toBe(400); + expect(body.error.message).toBe('message must contain supported canonical parts and no role or metadata fields'); + expect(mockCreateQueuedRunWithMessage).not.toHaveBeenCalled(); + }); + + it('rejects public message metadata', async () => { + const response = await POST( + makeRequest({ + message: { + metadata: { runId: 'run-1' }, + parts: [{ type: 'text', text: 'Nope' }], + }, + }), + { params: { threadId: 'thread-1' } } + ); + const body = await response.json(); + + expect(response.status).toBe(400); + expect(body.error.message).toBe('message must contain supported canonical parts and no role or metadata fields'); + expect(mockCreateQueuedRunWithMessage).not.toHaveBeenCalled(); + }); + + it('rejects unsupported runtime options', async () => { + const response = await POST( + makeRequest({ + message: { + parts: [{ type: 'text', text: 'Hi' }], + }, + runtimeOptions: { temperature: 0.7 }, + }), + { params: { threadId: 'thread-1' } } + ); + const body = await response.json(); + + expect(response.status).toBe(400); + expect(body.error.message).toBe('runtimeOptions contains unsupported or invalid fields'); + expect(mockCreateQueuedRunWithMessage).not.toHaveBeenCalled(); + }); + + it('rejects invalid model field types', async () => { + const response = await POST( + makeRequest({ + message: { + parts: [{ type: 'text', text: 'Hi' }], + }, + model: { id: 123 }, + }), + { params: { threadId: 'thread-1' } } + ); + const body = await response.json(); + + expect(response.status).toBe(400); + expect(body.error.message).toBe('model must contain only provider and id fields'); + expect(mockResolveSelection).not.toHaveBeenCalled(); + expect(mockCreateQueuedRunWithMessage).not.toHaveBeenCalled(); + }); + + it('rejects public harness selection', async () => { + const response = await POST( + makeRequest({ + message: { + parts: [{ type: 'text', text: 'Hi' }], + }, + harness: { kind: 'lifecycle_ai_sdk' }, + }), + { params: { threadId: 'thread-1' } } + ); + const body = await response.json(); + + expect(response.status).toBe(400); + expect(body.error.message).toBe('Unsupported run request fields: harness'); + expect(mockCreateQueuedRunWithMessage).not.toHaveBeenCalled(); + }); + + it('returns an idempotent response and emits a fresh dispatch signal for a queued run', async () => { + mockCreateQueuedRunWithMessage.mockResolvedValueOnce({ + run: { + uuid: 'run-1', + status: 'queued', + }, + message: { + uuid: 'message-1', + clientMessageId: 'client-message-1', + role: 'user', + parts: [{ type: 'text', text: 'Hi' }], + }, + created: false, + }); + + const response = await POST( + makeRequest({ + message: { + clientMessageId: 'client-message-1', + parts: [{ type: 'text', text: 'Hi' }], + }, + }), + { params: { threadId: 'thread-1' } } + ); + + expect(response.status).toBe(200); + expect(mockEnqueueRun).toHaveBeenCalledWith('run-1', 'submit', { githubToken: 'sample-gh-token' }); + }); + + it('marks a newly admitted queued run failed when activity touch fails before dispatch', async () => { + mockTouchActivity.mockRejectedValueOnce(new Error('touch failed')); + + const response = await POST( + makeRequest({ + message: { + clientMessageId: 'client-message-1', + parts: [{ type: 'text', text: 'Hi' }], + }, + }), + { params: { threadId: 'thread-1' } } + ); + + expect(response.status).toBe(500); + expect(mockMarkQueuedRunDispatchFailed).toHaveBeenCalledWith('run-1', expect.any(Error)); + expect(mockEnqueueRun).not.toHaveBeenCalled(); + }); + + it('does not mark an existing queued run failed when idempotent retry activity touch fails', async () => { + mockCreateQueuedRunWithMessage.mockResolvedValueOnce({ + run: { + uuid: 'run-1', + status: 'queued', + }, + message: { + uuid: 'message-1', + clientMessageId: 'client-message-1', + role: 'user', + parts: [{ type: 'text', text: 'Hi' }], + }, + created: false, + }); + mockTouchActivity.mockRejectedValueOnce(new Error('touch failed')); + + const response = await POST( + makeRequest({ + message: { + clientMessageId: 'client-message-1', + parts: [{ type: 'text', text: 'Hi' }], + }, + }), + { params: { threadId: 'thread-1' } } + ); + + expect(response.status).toBe(500); + expect(mockMarkQueuedRunDispatchFailed).not.toHaveBeenCalled(); + expect(mockEnqueueRun).not.toHaveBeenCalled(); + }); + + it('maps missing threads to 404', async () => { + mockGetOwnedThreadWithSession.mockRejectedValueOnce(new Error('Agent thread not found')); + + const response = await POST( + makeRequest({ + message: { + clientMessageId: 'client-message-1', + parts: [{ type: 'text', text: 'Hi' }], + }, + }), + { params: { threadId: 'missing-thread' } } + ); + const body = await response.json(); + + expect(response.status).toBe(404); + expect(body.error.message).toBe('Agent thread not found'); + expect(mockCreateQueuedRunWithMessage).not.toHaveBeenCalled(); + }); + + it('maps missing thread sessions to 404', async () => { + mockGetOwnedThreadWithSession.mockRejectedValueOnce(new Error('Agent session not found')); + + const response = await POST( + makeRequest({ + message: { + clientMessageId: 'client-message-1', + parts: [{ type: 'text', text: 'Hi' }], + }, + }), + { params: { threadId: 'thread-1' } } + ); + const body = await response.json(); + + expect(response.status).toBe(404); + expect(body.error.message).toBe('Agent session not found'); + expect(mockCreateQueuedRunWithMessage).not.toHaveBeenCalled(); + }); +}); diff --git a/src/app/api/v2/ai/agent/threads/[threadId]/runs/route.ts b/src/app/api/v2/ai/agent/threads/[threadId]/runs/route.ts new file mode 100644 index 00000000..17469b7b --- /dev/null +++ b/src/app/api/v2/ai/agent/threads/[threadId]/runs/route.ts @@ -0,0 +1,385 @@ +/** + * Copyright 2026 GoodRx, Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +import { NextRequest } from 'next/server'; +import 'server/lib/dependencies'; +import { createApiHandler } from 'server/lib/createApiHandler'; +import { errorResponse, successResponse } from 'server/lib/response'; +import { getRequestUserIdentity } from 'server/lib/get-user'; +import { resolveRequestGitHubToken } from 'server/lib/agentSession/githubToken'; +import AgentCapabilityService from 'server/services/agent/CapabilityService'; +import AgentProviderRegistry from 'server/services/agent/ProviderRegistry'; +import AgentRunAdmissionService from 'server/services/agent/RunAdmissionService'; +import AgentRunQueueService from 'server/services/agent/RunQueueService'; +import AgentRunService, { InvalidAgentRunDefaultsError } from 'server/services/agent/RunService'; +import AgentThreadService from 'server/services/agent/ThreadService'; +import { + normalizeCanonicalAgentMessagePart, + type AgentRunRuntimeOptions, + type CanonicalAgentRunMessageInput, +} from 'server/services/agent/canonicalMessages'; +import AgentSourceService from 'server/services/agent/SourceService'; +import AgentSessionService from 'server/services/agentSession'; +import AgentMessageStore from 'server/services/agent/MessageStore'; + +const MAX_RUN_MAX_ITERATIONS = 100; + +function getUnknownKeys(value: Record, allowedKeys: string[]): string[] { + return Object.keys(value).filter((key) => !allowedKeys.includes(key)); +} + +function hasOnlyAllowedCanonicalPartKeys(value: unknown): boolean { + if (!value || typeof value !== 'object' || Array.isArray(value)) { + return false; + } + + const part = value as Record; + switch (part.type) { + case 'text': + case 'reasoning': + return getUnknownKeys(part, ['type', 'text']).length === 0; + case 'file_ref': + return getUnknownKeys(part, ['type', 'path', 'url', 'mediaType', 'title']).length === 0; + case 'source_ref': + return getUnknownKeys(part, ['type', 'url', 'title', 'sourceType']).length === 0; + default: + return false; + } +} + +function normalizeCanonicalRunMessage(value: unknown): CanonicalAgentRunMessageInput | null { + if (!value || typeof value !== 'object') { + return null; + } + + const message = value as Record; + if (getUnknownKeys(message, ['clientMessageId', 'parts']).length > 0) { + return null; + } + + if (message.clientMessageId !== undefined && typeof message.clientMessageId !== 'string') { + return null; + } + + if (!Array.isArray(message.parts)) { + return null; + } + + if (!message.parts.every(hasOnlyAllowedCanonicalPartKeys)) { + return null; + } + + const parts = message.parts.map(normalizeCanonicalAgentMessagePart); + if (parts.some((part) => !part) || parts.length === 0) { + return null; + } + + return { + ...(typeof message.clientMessageId === 'string' && message.clientMessageId.trim() + ? { clientMessageId: message.clientMessageId.trim() } + : {}), + parts: parts as CanonicalAgentRunMessageInput['parts'], + }; +} + +function readString(value: unknown): string | null { + return typeof value === 'string' && value.trim() ? value.trim() : null; +} + +function normalizeModelRequest( + value: unknown +): { requestedProvider: string | null; requestedModel: string | null } | null { + if (value === undefined) { + return { + requestedProvider: null, + requestedModel: null, + }; + } + if (!value || typeof value !== 'object' || Array.isArray(value)) { + return null; + } + + const model = value as Record; + if (getUnknownKeys(model, ['provider', 'id']).length > 0) { + return null; + } + if (model.provider !== undefined && typeof model.provider !== 'string') { + return null; + } + if (model.id !== undefined && typeof model.id !== 'string') { + return null; + } + + return { + requestedProvider: readString(model.provider), + requestedModel: readString(model.id), + }; +} + +function normalizeRuntimeOptions(value: unknown): AgentRunRuntimeOptions | null { + if (value === undefined) { + return {}; + } + if (!value || typeof value !== 'object' || Array.isArray(value)) { + return null; + } + + const options = value as Record; + if (getUnknownKeys(options, ['maxIterations']).length > 0) { + return null; + } + + const normalized: AgentRunRuntimeOptions = {}; + if (options.maxIterations !== undefined) { + if ( + typeof options.maxIterations !== 'number' || + !Number.isInteger(options.maxIterations) || + options.maxIterations < 1 || + options.maxIterations > MAX_RUN_MAX_ITERATIONS + ) { + return null; + } + normalized.maxIterations = options.maxIterations; + } + + return normalized; +} + +/** + * @openapi + * /api/v2/ai/agent/threads/{threadId}/runs: + * post: + * summary: Create and enqueue a managed run for an agent thread + * tags: + * - Agent Sessions + * operationId: createAgentThreadRun + * parameters: + * - in: path + * name: threadId + * required: true + * schema: + * type: string + * requestBody: + * required: true + * content: + * application/json: + * schema: + * $ref: '#/components/schemas/CreateAgentThreadRunRequest' + * responses: + * '200': + * description: Existing managed run returned for an idempotent client message retry + * content: + * application/json: + * schema: + * allOf: + * - $ref: '#/components/schemas/SuccessApiResponse' + * - type: object + * required: [data] + * properties: + * data: + * $ref: '#/components/schemas/CreateAgentThreadRunResponse' + * '201': + * description: Managed run created + * content: + * application/json: + * schema: + * allOf: + * - $ref: '#/components/schemas/SuccessApiResponse' + * - type: object + * required: [data] + * properties: + * data: + * $ref: '#/components/schemas/CreateAgentThreadRunResponse' + * '400': + * description: Invalid run request, model selection, harness, or runtime options. + * content: + * application/json: + * schema: + * $ref: '#/components/schemas/ApiErrorResponse' + * '401': + * description: Unauthorized + * content: + * application/json: + * schema: + * $ref: '#/components/schemas/ApiErrorResponse' + * '404': + * description: Thread or session not found + * content: + * application/json: + * schema: + * $ref: '#/components/schemas/ApiErrorResponse' + * '409': + * description: Session source is not ready or another run is already active for the session + * content: + * application/json: + * schema: + * $ref: '#/components/schemas/ApiErrorResponse' + */ +const postHandler = async (req: NextRequest, { params }: { params: { threadId: string } }) => { + const userIdentity = getRequestUserIdentity(req); + if (!userIdentity) { + return errorResponse(new Error('Unauthorized'), { status: 401 }, req); + } + + const body = await req.json().catch(() => ({})); + if (!body || typeof body !== 'object' || Array.isArray(body)) { + return errorResponse(new Error('Request body must be an object'), { status: 400 }, req); + } + + const requestBody = body as Record; + const unknownTopLevelKeys = getUnknownKeys(requestBody, ['message', 'model', 'runtimeOptions']); + if (unknownTopLevelKeys.length > 0) { + return errorResponse( + new Error(`Unsupported run request fields: ${unknownTopLevelKeys.join(', ')}`), + { status: 400 }, + req + ); + } + + const message = normalizeCanonicalRunMessage(requestBody.message); + if (!message) { + return errorResponse( + new Error('message must contain supported canonical parts and no role or metadata fields'), + { status: 400 }, + req + ); + } + + const modelRequest = normalizeModelRequest(requestBody.model); + if (!modelRequest) { + return errorResponse(new Error('model must contain only provider and id fields'), { status: 400 }, req); + } + + const runtimeOptions = normalizeRuntimeOptions(requestBody.runtimeOptions); + if (!runtimeOptions) { + return errorResponse(new Error('runtimeOptions contains unsupported or invalid fields'), { status: 400 }, req); + } + + let threadWithSession; + try { + threadWithSession = await AgentThreadService.getOwnedThreadWithSession(params.threadId, userIdentity.userId); + } catch (error) { + if ( + error instanceof Error && + (error.message === 'Agent thread not found' || error.message === 'Agent session not found') + ) { + return errorResponse(error, { status: 404 }, req); + } + throw error; + } + + const { thread, session } = threadWithSession; + if (!AgentSessionService.canAcceptMessages(session)) { + return errorResponse(new Error(AgentSessionService.getMessageBlockReason(session)), { status: 409 }, req); + } + + const source = await AgentSourceService.getSessionSource(session.id); + if (!source || source.status !== 'ready') { + return errorResponse(new Error('Session source is not ready yet.'), { status: 409 }, req); + } + + const { approvalPolicy, repoFullName } = await AgentCapabilityService.resolveSessionContext( + session.uuid, + userIdentity + ); + const requestedHarness = null; + const resolvedHarness = readString(session.defaultHarness); + if (!resolvedHarness) { + return errorResponse(new Error('Agent run harness is required'), { status: 400 }, req); + } + if (resolvedHarness !== 'lifecycle_ai_sdk') { + return errorResponse(new Error(`Unsupported agent run harness: ${resolvedHarness}`), { status: 400 }, req); + } + + const requestedProvider = modelRequest.requestedProvider; + const requestedModel = modelRequest.requestedModel; + const resolvedModelRequest = requestedModel || readString(session.defaultModel); + if (!resolvedModelRequest) { + return errorResponse(new Error('Agent run model is required'), { status: 400 }, req); + } + + let selection; + try { + selection = await AgentProviderRegistry.resolveSelection({ + repoFullName, + requestedProvider: requestedProvider || undefined, + requestedModelId: resolvedModelRequest, + }); + } catch (error) { + return errorResponse(error instanceof Error ? error : new Error('Invalid agent run model'), { status: 400 }, req); + } + + let admission: Awaited>; + try { + admission = await AgentRunAdmissionService.createQueuedRunWithMessage({ + thread, + session, + policy: approvalPolicy, + message, + requestedHarness, + requestedProvider, + requestedModel, + resolvedHarness, + resolvedProvider: selection.provider, + resolvedModel: selection.modelId, + sandboxRequirement: source.sandboxRequirements || {}, + runtimeOptions, + }); + } catch (error) { + if (AgentRunService.isActiveRunConflictError(error)) { + return errorResponse(error, { status: 409 }, req); + } + if (error instanceof InvalidAgentRunDefaultsError) { + return errorResponse(error, { status: 400 }, req); + } + + throw error; + } + + try { + await AgentSessionService.touchActivity(session.uuid); + } catch (error) { + if (admission.created) { + await AgentRunService.markQueuedRunDispatchFailed(admission.run.uuid, error).catch(() => {}); + } + throw error; + } + + if (admission.created || admission.run.status === 'queued') { + const githubToken = await resolveRequestGitHubToken(req); + await AgentRunQueueService.enqueueRun(admission.run.uuid, 'submit', { githubToken }); + } + + return successResponse( + { + run: { + ...AgentRunService.serializeRun(admission.run), + threadId: thread.uuid, + sessionId: session.uuid, + }, + message: AgentMessageStore.serializeCanonicalMessage(admission.message, thread.uuid, admission.run.uuid), + links: { + events: `/api/v2/ai/agent/runs/${admission.run.uuid}/events`, + eventStream: `/api/v2/ai/agent/runs/${admission.run.uuid}/events/stream`, + pendingActions: `/api/v2/ai/agent/threads/${thread.uuid}/pending-actions`, + }, + }, + { status: admission.created ? 201 : 200 }, + req + ); +}; + +export const POST = createApiHandler(postHandler); diff --git a/src/app/api/v2/ai/config/agent-session/route.ts b/src/app/api/v2/ai/config/agent-session/route.ts index 95511b33..c30e5dd8 100644 --- a/src/app/api/v2/ai/config/agent-session/route.ts +++ b/src/app/api/v2/ai/config/agent-session/route.ts @@ -37,6 +37,12 @@ import AgentSessionConfigService from 'server/services/agentSessionConfig'; * application/json: * schema: * $ref: '#/components/schemas/GetGlobalAgentSessionConfigSuccessResponse' + * '401': + * description: Unauthorized + * content: + * application/json: + * schema: + * $ref: '#/components/schemas/ApiErrorResponse' * put: * summary: Update global Agent Session configuration * tags: @@ -61,6 +67,12 @@ import AgentSessionConfigService from 'server/services/agentSessionConfig'; * application/json: * schema: * $ref: '#/components/schemas/ApiErrorResponse' + * '401': + * description: Unauthorized + * content: + * application/json: + * schema: + * $ref: '#/components/schemas/ApiErrorResponse' */ const getHandler = async (req: NextRequest) => { const config = await AgentSessionConfigService.getInstance().getGlobalConfig(); diff --git a/src/app/api/v2/ai/config/agent-session/runtime/route.ts b/src/app/api/v2/ai/config/agent-session/runtime/route.ts index 84c7cbae..edd84cc7 100644 --- a/src/app/api/v2/ai/config/agent-session/runtime/route.ts +++ b/src/app/api/v2/ai/config/agent-session/runtime/route.ts @@ -37,6 +37,12 @@ import AgentSessionConfigService from 'server/services/agentSessionConfig'; * application/json: * schema: * $ref: '#/components/schemas/GetGlobalAgentSessionRuntimeConfigSuccessResponse' + * '401': + * description: Unauthorized + * content: + * application/json: + * schema: + * $ref: '#/components/schemas/ApiErrorResponse' * put: * summary: Update global Agent Session runtime configuration * tags: @@ -61,6 +67,12 @@ import AgentSessionConfigService from 'server/services/agentSessionConfig'; * application/json: * schema: * $ref: '#/components/schemas/ApiErrorResponse' + * '401': + * description: Unauthorized + * content: + * application/json: + * schema: + * $ref: '#/components/schemas/ApiErrorResponse' */ const getHandler = async (req: NextRequest) => { const config = await AgentSessionConfigService.getInstance().getGlobalRuntimeConfig(); diff --git a/src/server/database.ts b/src/server/database.ts index 7764a881..e2393f22 100644 --- a/src/server/database.ts +++ b/src/server/database.ts @@ -17,7 +17,7 @@ import knex, { Knex } from 'knex'; import { merge } from 'lodash'; import * as models from './models'; -import { IServices } from 'server/services/types'; +import type { IServices } from 'server/services/types'; import Model from 'server/models/_Model'; import knexfile from '../../knexfile'; diff --git a/src/server/db/migrations/019_remove_scale_to_zero.ts b/src/server/db/migrations/019_remove_scale_to_zero.ts index 021cbd38..38862bb5 100644 --- a/src/server/db/migrations/019_remove_scale_to_zero.ts +++ b/src/server/db/migrations/019_remove_scale_to_zero.ts @@ -16,21 +16,40 @@ import { Knex } from 'knex'; -export async function up(knex: Knex): Promise { - await knex.schema.alterTable('services', (table) => { - table.dropColumn('scaleToZero'); - table.dropColumn('scaleToZeroMetricsCheckInterval'); - }); +async function dropColumnIfExists(knex: Knex, tableName: string, columnName: string): Promise { + const exists = await knex.schema.hasColumn(tableName, columnName); - await knex.schema.alterTable('deployables', (table) => { - table.dropColumn('scaleToZero'); - table.dropColumn('scaleToZeroMetricsCheckInterval'); - table.dropColumn('kedaScaleToZero'); - }); + if (!exists) { + return; + } - await knex.schema.alterTable('deploys', (table) => { - table.dropColumn('kedaScaleToZero'); + await knex.schema.alterTable(tableName, (table) => { + table.dropColumn(columnName); }); +} + +async function addColumnIfMissing( + knex: Knex, + tableName: string, + columnName: string, + addColumn: (table: Knex.AlterTableBuilder) => void +): Promise { + const exists = await knex.schema.hasColumn(tableName, columnName); + + if (exists) { + return; + } + + await knex.schema.alterTable(tableName, addColumn); +} + +export async function up(knex: Knex): Promise { + await dropColumnIfExists(knex, 'services', 'scaleToZero'); + await dropColumnIfExists(knex, 'services', 'scaleToZeroMetricsCheckInterval'); + await dropColumnIfExists(knex, 'deployables', 'scaleToZero'); + await dropColumnIfExists(knex, 'deployables', 'scaleToZeroMetricsCheckInterval'); + await dropColumnIfExists(knex, 'deployables', 'kedaScaleToZero'); + await dropColumnIfExists(knex, 'deploys', 'kedaScaleToZero'); await knex.raw(` DELETE FROM global_config @@ -46,18 +65,22 @@ export async function up(knex: Knex): Promise { } export async function down(knex: Knex): Promise { - await knex.schema.alterTable('services', (table) => { + await addColumnIfMissing(knex, 'services', 'scaleToZero', (table) => { table.boolean('scaleToZero').defaultTo(false); + }); + await addColumnIfMissing(knex, 'services', 'scaleToZeroMetricsCheckInterval', (table) => { table.integer('scaleToZeroMetricsCheckInterval').defaultTo(1800); }); - - await knex.schema.alterTable('deployables', (table) => { + await addColumnIfMissing(knex, 'deployables', 'scaleToZero', (table) => { table.boolean('scaleToZero').defaultTo(false); + }); + await addColumnIfMissing(knex, 'deployables', 'scaleToZeroMetricsCheckInterval', (table) => { table.integer('scaleToZeroMetricsCheckInterval').defaultTo(1800); + }); + await addColumnIfMissing(knex, 'deployables', 'kedaScaleToZero', (table) => { table.json('kedaScaleToZero').defaultTo('{}'); }); - - await knex.schema.alterTable('deploys', (table) => { + await addColumnIfMissing(knex, 'deploys', 'kedaScaleToZero', (table) => { table.json('kedaScaleToZero').defaultTo('{}'); }); @@ -70,7 +93,12 @@ export async function down(knex: Knex): Promise { now(), null, 'This is the default configuration for Keda Scale To Zero' - ); + ) + ON CONFLICT (key) DO UPDATE + SET config = EXCLUDED.config, + "updatedAt" = now(), + "deletedAt" = null, + description = EXCLUDED.description; `); await knex.raw(` diff --git a/src/server/db/migrations/021_agent_session_control_plane_contract.ts b/src/server/db/migrations/021_agent_session_control_plane_contract.ts new file mode 100644 index 00000000..86334f45 --- /dev/null +++ b/src/server/db/migrations/021_agent_session_control_plane_contract.ts @@ -0,0 +1,565 @@ +/** + * Copyright 2026 GoodRx, Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +import { Knex } from 'knex'; + +export const config = { + transaction: true, +}; + +const EMPTY_OBJECT = '{}'; +const EMPTY_ARRAY = '[]'; +const AGENT_MESSAGES_CLIENT_MESSAGE_ID_INDEX_NAME = 'agent_messages_thread_client_message_id_unique'; +const AGENT_SESSION_DEFAULTS_KEY = 'agentSessionDefaults'; + +type JsonObject = Record; + +function isObject(value: unknown): value is JsonObject { + return Boolean(value) && typeof value === 'object' && !Array.isArray(value); +} + +function getPvcAccessMode(): 'ReadWriteOnce' | 'ReadWriteMany' { + return process.env.AGENT_SESSION_PVC_ACCESS_MODE === 'ReadWriteMany' ? 'ReadWriteMany' : 'ReadWriteOnce'; +} + +function mergeAgentSessionDefaults(currentConfig: JsonObject): JsonObject { + const currentWorkspaceStorage = isObject(currentConfig.workspaceStorage) ? currentConfig.workspaceStorage : {}; + const currentCleanup = isObject(currentConfig.cleanup) ? currentConfig.cleanup : {}; + const currentDurability = isObject(currentConfig.durability) ? currentConfig.durability : {}; + + return { + ...currentConfig, + workspaceStorage: { + defaultSize: '10Gi', + allowedSizes: ['10Gi'], + allowClientOverride: false, + accessMode: getPvcAccessMode(), + ...currentWorkspaceStorage, + }, + cleanup: { + activeIdleSuspendMs: 30 * 60 * 1000, + startingTimeoutMs: 15 * 60 * 1000, + hibernatedRetentionMs: 24 * 60 * 60 * 1000, + intervalMs: 5 * 60 * 1000, + redisTtlSeconds: 7200, + ...currentCleanup, + }, + durability: { + runExecutionLeaseMs: 30 * 60 * 1000, + queuedRunDispatchStaleMs: 30 * 1000, + dispatchRecoveryLimit: 50, + maxDurablePayloadBytes: 64 * 1024, + payloadPreviewBytes: 16 * 1024, + fileChangePreviewChars: 4000, + ...currentDurability, + }, + }; +} + +async function upsertAgentSessionDefaults(knex: Knex): Promise { + const row = await knex('global_config').where('key', AGENT_SESSION_DEFAULTS_KEY).first(); + const currentConfig = isObject(row?.config) ? row.config : {}; + const nextConfig = mergeAgentSessionDefaults(currentConfig); + + if (!row) { + await knex('global_config').insert({ + key: AGENT_SESSION_DEFAULTS_KEY, + config: nextConfig, + createdAt: knex.fn.now(), + updatedAt: knex.fn.now(), + deletedAt: null, + description: 'Default configuration for agent session workspace runtime.', + }); + return; + } + + await knex('global_config').where('key', AGENT_SESSION_DEFAULTS_KEY).update({ + config: nextConfig, + updatedAt: knex.fn.now(), + }); +} + +async function removeAgentSessionOperationalDefaults(knex: Knex): Promise { + const row = await knex('global_config').where('key', AGENT_SESSION_DEFAULTS_KEY).first(); + if (!row || !isObject(row.config)) { + return; + } + + const { workspaceStorage: _workspaceStorage, cleanup: _cleanup, durability: _durability, ...nextConfig } = row.config; + await knex('global_config').where('key', AGENT_SESSION_DEFAULTS_KEY).update({ + config: nextConfig, + updatedAt: knex.fn.now(), + }); +} + +export async function up(knex: Knex): Promise { + await knex.schema.alterTable('agent_sessions', (table) => { + table.string('sessionKind').notNullable().defaultTo('environment'); + table.string('chatStatus').notNullable().defaultTo('ready'); + table.string('workspaceStatus').notNullable().defaultTo('ready'); + table.string('defaultModel').nullable(); + table.string('defaultHarness').nullable(); + table.integer('defaultThreadId').nullable(); + }); + + await knex('agent_sessions').update({ + sessionKind: knex.raw(`case when "buildKind" = 'sandbox' then 'sandbox' else 'environment' end`), + chatStatus: knex.raw( + `case + when "status" = 'ended' then 'ended' + when "status" = 'error' then 'error' + else 'ready' + end` + ), + workspaceStatus: knex.raw( + `case + when "status" = 'starting' then 'provisioning' + when "status" = 'active' then 'ready' + when "status" = 'error' then 'failed' + when "status" = 'ended' then 'ended' + else 'ready' + end` + ), + defaultModel: knex.raw('coalesce("defaultModel", "model")'), + defaultHarness: knex.raw(`coalesce("defaultHarness", 'lifecycle_ai_sdk')`), + }); + + await knex.schema.alterTable('agent_sessions', (table) => { + table.string('buildKind').nullable().alter(); + table.string('podName').nullable().alter(); + table.string('namespace').nullable().alter(); + table.string('pvcName').nullable().alter(); + table.string('defaultModel').notNullable().alter(); + table.foreign(['defaultThreadId']).references(['id']).inTable('agent_threads').onDelete('SET NULL'); + table.index(['defaultThreadId']); + }); + + await knex.raw(` + update agent_sessions as session + set "defaultThreadId" = thread.id + from agent_threads as thread + where thread."sessionId" = session.id + and thread."isDefault" = true + and thread."archivedAt" is null + and session."defaultThreadId" is null + `); + + await knex.schema.createTable('agent_sources', (table) => { + table.increments('id').primary(); + table.uuid('uuid').notNullable().unique().defaultTo(knex.raw('gen_random_uuid()')); + table.integer('sessionId').notNullable().references('id').inTable('agent_sessions').onDelete('CASCADE'); + table.string('adapter').notNullable(); + table.string('status').notNullable().defaultTo('requested'); + table.jsonb('input').notNullable().defaultTo(EMPTY_OBJECT); + table.jsonb('preparedSource').notNullable().defaultTo(EMPTY_OBJECT); + table.jsonb('sandboxRequirements').notNullable().defaultTo(EMPTY_OBJECT); + table.jsonb('error').nullable(); + table.timestamp('preparedAt').nullable(); + table.timestamp('cleanedUpAt').nullable(); + table.timestamp('createdAt').notNullable().defaultTo(knex.fn.now()); + table.timestamp('updatedAt').notNullable().defaultTo(knex.fn.now()); + + table.unique(['sessionId']); + table.index(['status', 'createdAt']); + }); + + await knex.schema.createTable('agent_sandboxes', (table) => { + table.increments('id').primary(); + table.uuid('uuid').notNullable().unique().defaultTo(knex.raw('gen_random_uuid()')); + table.integer('sessionId').notNullable().references('id').inTable('agent_sessions').onDelete('CASCADE'); + table.integer('generation').notNullable(); + table.string('provider').notNullable(); + table.string('status').notNullable().defaultTo('provisioning'); + table.jsonb('capabilitySnapshot').notNullable().defaultTo(EMPTY_OBJECT); + table.jsonb('providerState').notNullable().defaultTo(EMPTY_OBJECT); + table.jsonb('metadata').notNullable().defaultTo(EMPTY_OBJECT); + table.jsonb('error').nullable(); + table.timestamp('suspendedAt').nullable(); + table.timestamp('endedAt').nullable(); + table.timestamp('createdAt').notNullable().defaultTo(knex.fn.now()); + table.timestamp('updatedAt').notNullable().defaultTo(knex.fn.now()); + + table.unique(['sessionId', 'generation']); + table.index(['sessionId', 'createdAt']); + table.index(['status', 'createdAt']); + }); + + await knex.schema.createTable('agent_sandbox_exposures', (table) => { + table.increments('id').primary(); + table.uuid('uuid').notNullable().unique().defaultTo(knex.raw('gen_random_uuid()')); + table.integer('sandboxId').notNullable().references('id').inTable('agent_sandboxes').onDelete('CASCADE'); + table.string('kind').notNullable(); + table.string('status').notNullable().defaultTo('ready'); + table.integer('targetPort').nullable(); + table.text('url').nullable(); + table.jsonb('metadata').notNullable().defaultTo(EMPTY_OBJECT); + table.jsonb('providerState').notNullable().defaultTo(EMPTY_OBJECT); + table.timestamp('lastVerifiedAt').nullable(); + table.timestamp('endedAt').nullable(); + table.timestamp('createdAt').notNullable().defaultTo(knex.fn.now()); + table.timestamp('updatedAt').notNullable().defaultTo(knex.fn.now()); + + table.index(['sandboxId', 'kind']); + table.index(['status', 'createdAt']); + }); + + await knex.schema.createTable('agent_run_events', (table) => { + table.increments('id').primary(); + table.uuid('uuid').notNullable().unique().defaultTo(knex.raw('gen_random_uuid()')); + table.integer('runId').notNullable().references('id').inTable('agent_runs').onDelete('CASCADE'); + table.integer('sequence').notNullable(); + table.string('eventType').notNullable(); + table.jsonb('payload').notNullable().defaultTo(EMPTY_OBJECT); + table.timestamp('createdAt').notNullable().defaultTo(knex.fn.now()); + table.timestamp('updatedAt').notNullable().defaultTo(knex.fn.now()); + + table.unique(['runId', 'sequence']); + table.index(['runId', 'sequence']); + table.index(['eventType', 'createdAt']); + }); + + await knex.schema.alterTable('agent_runs', (table) => { + table.string('requestedHarness').nullable(); + table.string('resolvedHarness').nullable(); + table.string('requestedProvider').nullable(); + table.string('requestedModel').nullable(); + table.string('resolvedProvider').nullable(); + table.string('resolvedModel').nullable(); + table.jsonb('sandboxRequirement').notNullable().defaultTo(EMPTY_OBJECT); + table.integer('sandboxGeneration').nullable(); + table.string('executionOwner').nullable(); + table.timestamp('leaseExpiresAt').nullable(); + table.timestamp('heartbeatAt').nullable(); + table.dropColumn('streamState'); + table.index(['status', 'leaseExpiresAt']); + }); + + await knex('agent_runs').update({ + requestedHarness: knex.raw(`coalesce("requestedHarness", 'lifecycle_ai_sdk')`), + resolvedHarness: knex.raw(`coalesce("resolvedHarness", 'lifecycle_ai_sdk')`), + requestedProvider: knex.raw('coalesce("requestedProvider", "provider")'), + requestedModel: knex.raw('coalesce("requestedModel", "model")'), + resolvedProvider: knex.raw('coalesce("resolvedProvider", "provider")'), + resolvedModel: knex.raw('coalesce("resolvedModel", "model")'), + }); + + await knex.schema.alterTable('agent_messages', (table) => { + table.uuid('uuid').nullable(); + table.string('clientMessageId').nullable(); + table.jsonb('parts').notNullable().defaultTo(EMPTY_ARRAY); + table.jsonb('uiMessage').nullable().alter(); + }); + + await knex('agent_messages').update({ + uuid: knex.raw('gen_random_uuid()'), + }); + + await knex.schema.alterTable('agent_messages', (table) => { + table.uuid('uuid').notNullable().alter(); + table.unique(['uuid']); + }); + + await knex.raw(` + with ranked as ( + select + id, + nullif(btrim(metadata->>'clientMessageId'), '') as "clientMessageId", + row_number() over ( + partition by "threadId", nullif(btrim(metadata->>'clientMessageId'), '') + order by id + ) as rank + from agent_messages + where nullif(btrim(metadata->>'clientMessageId'), '') is not null + ) + update agent_messages as message + set "clientMessageId" = ranked."clientMessageId" + from ranked + where message.id = ranked.id + and ranked.rank = 1 + `); + + await knex.raw(` + create unique index ${AGENT_MESSAGES_CLIENT_MESSAGE_ID_INDEX_NAME} + on agent_messages ("threadId", "clientMessageId") + where "clientMessageId" is not null + `); + + await knex.raw(` + insert into agent_sources ( + "sessionId", + "adapter", + "status", + "input", + "preparedSource", + "sandboxRequirements", + "error", + "preparedAt", + "cleanedUpAt", + "createdAt", + "updatedAt" + ) + select + session.id, + case + when session."sessionKind" = 'chat' then 'blank_workspace' + when session."buildKind" = 'sandbox' then 'lifecycle_fork' + else 'lifecycle_environment' + end, + case + when session.status = 'error' or session."workspaceStatus" = 'failed' then 'failed' + when session.status = 'ended' then 'cleaned_up' + else 'ready' + end, + jsonb_strip_nulls( + jsonb_build_object( + 'buildUuid', session."buildUuid", + 'buildKind', session."buildKind", + 'sessionKind', session."sessionKind", + 'ownerGithubUsername', session."ownerGithubUsername" + ) + ), + jsonb_build_object( + 'kind', 'workspace_snapshot', + 'metadata', + jsonb_strip_nulls( + jsonb_build_object( + 'buildUuid', session."buildUuid", + 'buildKind', session."buildKind", + 'sessionKind', session."sessionKind" + ) + ) + ), + case + when session."sessionKind" = 'chat' then '{"filesystem":"persistent","suspendMode":"filesystem","editorAccess":true,"previewPorts":true}'::jsonb + else '{"filesystem":"persistent","suspendMode":"none","editorAccess":true,"previewPorts":true}'::jsonb + end, + case + when session.status = 'error' or session."workspaceStatus" = 'failed' then jsonb_build_object('message', 'Source is not ready') + else null + end, + case + when session.status = 'ended' then null + else session."updatedAt" + end, + case + when session.status = 'ended' then coalesce(session."endedAt", session."updatedAt") + else null + end, + session."createdAt", + session."updatedAt" + from agent_sessions as session + where not exists ( + select 1 + from agent_sources as source + where source."sessionId" = session.id + ) + `); + + await knex.raw(` + insert into agent_sandboxes ( + "sessionId", + "generation", + "provider", + "status", + "capabilitySnapshot", + "providerState", + "metadata", + "error", + "suspendedAt", + "endedAt", + "createdAt", + "updatedAt" + ) + select + session.id, + 1, + 'lifecycle_kubernetes', + case + when session."workspaceStatus" = 'provisioning' then 'provisioning' + when session."workspaceStatus" = 'hibernated' then 'suspended' + when session.status = 'ended' or session."workspaceStatus" = 'ended' then 'ended' + when session.status = 'error' or session."workspaceStatus" = 'failed' then 'failed' + else 'ready' + end, + jsonb_build_object( + 'toolTransport', 'mcp', + 'persistentFilesystem', coalesce(session."pvcName", '') <> '', + 'portExposure', true, + 'editorAccess', true + ), + jsonb_strip_nulls( + jsonb_build_object( + 'namespace', session.namespace, + 'podName', session."podName", + 'pvcName', session."pvcName" + ) + ), + jsonb_strip_nulls( + jsonb_build_object( + 'sessionKind', session."sessionKind", + 'buildUuid', session."buildUuid", + 'buildKind', session."buildKind" + ) + ), + case + when session.status = 'error' or session."workspaceStatus" = 'failed' then jsonb_build_object('message', 'Sandbox is not available') + else null + end, + case + when session."workspaceStatus" = 'hibernated' then session."updatedAt" + else null + end, + case + when session.status = 'ended' or session."workspaceStatus" = 'ended' then coalesce(session."endedAt", session."updatedAt") + else null + end, + session."createdAt", + session."updatedAt" + from agent_sessions as session + where (coalesce(session.namespace, '') <> '' or coalesce(session."podName", '') <> '' or coalesce(session."pvcName", '') <> '') + and not exists ( + select 1 + from agent_sandboxes as sandbox + where sandbox."sessionId" = session.id + and sandbox.generation = 1 + ) + `); + + await knex.raw(` + insert into agent_sandbox_exposures ( + "sandboxId", + "kind", + "status", + "targetPort", + "url", + "metadata", + "providerState", + "lastVerifiedAt", + "endedAt", + "createdAt", + "updatedAt" + ) + select + sandbox.id, + 'editor', + case + when sandbox.status = 'ended' then 'ended' + when sandbox.status = 'failed' then 'failed' + when sandbox.status = 'provisioning' then 'provisioning' + else 'ready' + end, + null, + '/api/agent-session/workspace-editor/' || session.uuid || '/', + '{"attachmentKind":"mcp_gateway"}'::jsonb, + '{}'::jsonb, + case + when sandbox.status = 'ready' then sandbox."updatedAt" + else null + end, + sandbox."endedAt", + sandbox."createdAt", + sandbox."updatedAt" + from agent_sandboxes as sandbox + join agent_sessions as session + on session.id = sandbox."sessionId" + where not exists ( + select 1 + from agent_sandbox_exposures as exposure + where exposure."sandboxId" = sandbox.id + and exposure.kind = 'editor' + and exposure."endedAt" is null + ) + `); + + await upsertAgentSessionDefaults(knex); +} + +export async function down(knex: Knex): Promise { + const chatSession = await knex('agent_sessions').where('sessionKind', 'chat').first(); + if (chatSession) { + throw new Error('Cannot roll back agent session control-plane contract while chat sessions exist'); + } + + await knex.schema.dropTableIfExists('agent_run_events'); + await knex.schema.dropTableIfExists('agent_sandbox_exposures'); + await knex.schema.dropTableIfExists('agent_sandboxes'); + await knex.schema.dropTableIfExists('agent_sources'); + await removeAgentSessionOperationalDefaults(knex); + await knex.raw(`drop index if exists ${AGENT_MESSAGES_CLIENT_MESSAGE_ID_INDEX_NAME}`); + + await knex('agent_messages') + .whereNull('uiMessage') + .update({ + uiMessage: knex.raw(` + jsonb_build_object( + 'id', "uuid", + 'role', "role", + 'parts', "parts", + 'metadata', "metadata" + ) + `), + }); + + await knex.schema.alterTable('agent_messages', (table) => { + table.jsonb('uiMessage').notNullable().alter(); + }); + + await knex.schema.alterTable('agent_messages', (table) => { + table.dropUnique(['uuid']); + table.dropColumn('clientMessageId'); + table.dropColumn('parts'); + table.dropColumn('uuid'); + }); + + await knex.schema.alterTable('agent_runs', (table) => { + table.dropColumn('sandboxGeneration'); + table.dropColumn('sandboxRequirement'); + table.dropColumn('resolvedModel'); + table.dropColumn('resolvedProvider'); + table.dropColumn('requestedModel'); + table.dropColumn('requestedProvider'); + table.dropColumn('resolvedHarness'); + table.dropColumn('requestedHarness'); + table.dropColumn('heartbeatAt'); + table.dropColumn('leaseExpiresAt'); + table.dropColumn('executionOwner'); + table.jsonb('streamState').notNullable().defaultTo(EMPTY_OBJECT); + }); + + await knex.schema.alterTable('agent_sessions', (table) => { + table.dropForeign(['defaultThreadId']); + table.dropColumn('defaultThreadId'); + table.dropColumn('defaultHarness'); + table.dropColumn('defaultModel'); + }); + + await knex('agent_sessions').update({ + buildKind: knex.raw(`coalesce("buildKind", 'environment')`), + podName: knex.raw(`coalesce("podName", '')`), + namespace: knex.raw(`coalesce("namespace", '')`), + pvcName: knex.raw(`coalesce("pvcName", '')`), + }); + + await knex.schema.alterTable('agent_sessions', (table) => { + table.string('buildKind').notNullable().defaultTo('environment').alter(); + table.string('podName').notNullable().alter(); + table.string('namespace').notNullable().alter(); + table.string('pvcName').notNullable().alter(); + table.dropColumn('workspaceStatus'); + table.dropColumn('chatStatus'); + table.dropColumn('sessionKind'); + }); +} diff --git a/src/server/db/seeds/.gitkeep b/src/server/db/seeds/.gitkeep new file mode 100644 index 00000000..3b1281a6 --- /dev/null +++ b/src/server/db/seeds/.gitkeep @@ -0,0 +1 @@ +# Keeps the Knex seeds directory present for local dev startup. diff --git a/src/server/jobs/__tests__/agentRunDispatchRecovery.test.ts b/src/server/jobs/__tests__/agentRunDispatchRecovery.test.ts new file mode 100644 index 00000000..b23fd4cd --- /dev/null +++ b/src/server/jobs/__tests__/agentRunDispatchRecovery.test.ts @@ -0,0 +1,100 @@ +/** + * Copyright 2026 GoodRx, Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +jest.mock('server/services/agent/RunService', () => ({ + __esModule: true, + default: { + listRunsNeedingDispatch: jest.fn(), + }, +})); + +jest.mock('server/services/agent/RunQueueService', () => ({ + __esModule: true, + default: { + enqueueRun: jest.fn(), + }, +})); + +jest.mock('server/lib/logger', () => { + const logger = { + info: jest.fn(), + warn: jest.fn(), + }; + + return { + getLogger: jest.fn(() => logger), + }; +}); + +import AgentRunService from 'server/services/agent/RunService'; +import AgentRunQueueService from 'server/services/agent/RunQueueService'; +import { getLogger } from 'server/lib/logger'; +import { processAgentRunDispatchRecovery } from '../agentRunDispatchRecovery'; + +const mockListRunsNeedingDispatch = AgentRunService.listRunsNeedingDispatch as jest.Mock; +const mockEnqueueRun = AgentRunQueueService.enqueueRun as jest.Mock; +const mockLogger = getLogger() as { + info: jest.Mock; + warn: jest.Mock; +}; + +describe('agentRunDispatchRecovery', () => { + beforeEach(() => { + jest.clearAllMocks(); + }); + + it('re-enqueues stale queued and expired-lease runs', async () => { + mockListRunsNeedingDispatch.mockResolvedValue([{ uuid: 'run-1' }, { uuid: 'run-2' }]); + mockEnqueueRun + .mockResolvedValueOnce({ dispatchAttemptId: 'attempt-1' }) + .mockResolvedValueOnce({ dispatchAttemptId: 'attempt-2' }); + + const result = await processAgentRunDispatchRecovery(); + + expect(mockListRunsNeedingDispatch).toHaveBeenCalledTimes(1); + expect(mockEnqueueRun).toHaveBeenNthCalledWith(1, 'run-1', 'resume'); + expect(mockEnqueueRun).toHaveBeenNthCalledWith(2, 'run-2', 'resume'); + expect(result).toEqual({ + runs: 2, + enqueued: [ + { runId: 'run-1', dispatchAttemptId: 'attempt-1' }, + { runId: 'run-2', dispatchAttemptId: 'attempt-2' }, + ], + failed: [], + }); + expect(mockLogger.info).toHaveBeenCalledWith( + expect.objectContaining({ runId: 'run-1', dispatchAttemptId: 'attempt-1' }), + 'AgentExec: recovery enqueued runId=run-1 reason=resume dispatchAttemptId=attempt-1' + ); + }); + + it('continues re-enqueueing remaining runs after one enqueue fails', async () => { + mockListRunsNeedingDispatch.mockResolvedValue([{ uuid: 'run-1' }, { uuid: 'run-2' }]); + mockEnqueueRun + .mockRejectedValueOnce(new Error('redis unavailable')) + .mockResolvedValueOnce({ dispatchAttemptId: 'attempt-2' }); + + const result = await processAgentRunDispatchRecovery(); + + expect(mockEnqueueRun).toHaveBeenNthCalledWith(1, 'run-1', 'resume'); + expect(mockEnqueueRun).toHaveBeenNthCalledWith(2, 'run-2', 'resume'); + expect(result).toEqual({ + runs: 2, + enqueued: [{ runId: 'run-2', dispatchAttemptId: 'attempt-2' }], + failed: [{ runId: 'run-1' }], + }); + }); +}); diff --git a/src/server/jobs/__tests__/agentRunExecute.test.ts b/src/server/jobs/__tests__/agentRunExecute.test.ts new file mode 100644 index 00000000..5eaea4e7 --- /dev/null +++ b/src/server/jobs/__tests__/agentRunExecute.test.ts @@ -0,0 +1,187 @@ +/** + * Copyright 2026 GoodRx, Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +jest.mock('server/services/agent/LifecycleAiSdkHarness', () => ({ + __esModule: true, + default: { + executeRun: jest.fn(), + }, +})); + +jest.mock('server/services/agent/RunService', () => ({ + __esModule: true, + default: { + claimQueuedRunForExecution: jest.fn(), + getRunByUuid: jest.fn(), + isTerminalStatus: jest.fn(), + markFailedForExecutionOwner: jest.fn(), + }, +})); + +jest.mock('server/lib/logger', () => ({ + getLogger: jest.fn(() => ({ + info: jest.fn(), + warn: jest.fn(), + })), +})); + +jest.mock('server/lib/logger/context', () => ({ + withLogContext: jest.fn((_context, fn) => fn()), +})); + +jest.mock('server/lib/encryption', () => ({ + decrypt: jest.fn((value: string) => `decrypted:${value}`), +})); + +import LifecycleAiSdkHarness from 'server/services/agent/LifecycleAiSdkHarness'; +import AgentRunService from 'server/services/agent/RunService'; +import { AgentRunOwnershipLostError } from 'server/services/agent/AgentRunOwnershipLostError'; +import { processAgentRunExecute } from '../agentRunExecute'; + +const mockClaimQueuedRunForExecution = AgentRunService.claimQueuedRunForExecution as jest.Mock; +const mockGetRunByUuid = AgentRunService.getRunByUuid as jest.Mock; +const mockExecuteRun = LifecycleAiSdkHarness.executeRun as jest.Mock; +const mockIsTerminalStatus = AgentRunService.isTerminalStatus as jest.Mock; +const mockMarkFailedForExecutionOwner = AgentRunService.markFailedForExecutionOwner as jest.Mock; + +describe('agentRunExecute', () => { + beforeEach(() => { + jest.clearAllMocks(); + }); + + it('skips duplicate jobs once the run is no longer queued', async () => { + mockClaimQueuedRunForExecution.mockResolvedValue(null); + + await processAgentRunExecute({ + data: { + runId: 'run-1', + dispatchAttemptId: 'attempt-1', + reason: 'submit', + }, + } as any); + + expect(mockClaimQueuedRunForExecution).toHaveBeenCalledWith('run-1', expect.stringMatching(/^bull:unknown:/)); + expect(mockExecuteRun).not.toHaveBeenCalled(); + expect(mockMarkFailedForExecutionOwner).not.toHaveBeenCalled(); + }); + + it('rejects jobs missing a dispatch attempt id before claiming the run', async () => { + await expect( + processAgentRunExecute({ + data: { + runId: 'run-1', + reason: 'submit', + }, + } as any) + ).rejects.toThrow('Invalid agent run execute job payload: dispatchAttemptId is required'); + + expect(mockClaimQueuedRunForExecution).not.toHaveBeenCalled(); + expect(mockExecuteRun).not.toHaveBeenCalled(); + }); + + it('decrypts the queued GitHub token before executing the run', async () => { + const run = { uuid: 'run-1', status: 'starting' }; + mockClaimQueuedRunForExecution.mockResolvedValue(run); + mockExecuteRun.mockResolvedValue({ run }); + + await processAgentRunExecute({ + data: { + runId: 'run-1', + dispatchAttemptId: 'attempt-1', + reason: 'submit', + encryptedGithubToken: 'encrypted-token', + }, + } as any); + + expect(mockExecuteRun).toHaveBeenCalledWith(run, { + requestGitHubToken: 'decrypted:encrypted-token', + dispatchAttemptId: 'attempt-1', + }); + }); + + it('marks a queued run failed when harness setup throws before terminal status is recorded', async () => { + const run = { uuid: 'run-1', status: 'starting' }; + mockClaimQueuedRunForExecution.mockResolvedValue(run); + mockGetRunByUuid.mockResolvedValue(run); + mockExecuteRun.mockRejectedValue(new Error('setup failed')); + mockIsTerminalStatus.mockReturnValue(false); + mockMarkFailedForExecutionOwner.mockResolvedValue(undefined); + + await expect( + processAgentRunExecute({ + data: { + runId: 'run-1', + dispatchAttemptId: 'attempt-1', + reason: 'submit', + }, + } as any) + ).rejects.toThrow('setup failed'); + + expect(mockMarkFailedForExecutionOwner).toHaveBeenCalledWith( + 'run-1', + expect.stringMatching(/^bull:unknown:/), + expect.objectContaining({ message: 'setup failed' }), + undefined, + { dispatchAttemptId: 'attempt-1' } + ); + }); + + it('does not overwrite a run failure already recorded by the executor', async () => { + const run = { uuid: 'run-1', status: 'starting' }; + const failedRun = { uuid: 'run-1', status: 'failed' }; + mockClaimQueuedRunForExecution.mockResolvedValue(run); + mockGetRunByUuid.mockResolvedValue(failedRun); + mockExecuteRun.mockRejectedValue(new Error('setup failed')); + mockIsTerminalStatus.mockReturnValue(true); + + await expect( + processAgentRunExecute({ + data: { + runId: 'run-1', + dispatchAttemptId: 'attempt-1', + reason: 'submit', + }, + } as any) + ).rejects.toThrow('setup failed'); + + expect(mockMarkFailedForExecutionOwner).not.toHaveBeenCalled(); + }); + + it('treats ownership loss from execution as a clean stale-worker exit', async () => { + const run = { uuid: 'run-1', status: 'starting' }; + mockClaimQueuedRunForExecution.mockResolvedValue(run); + mockExecuteRun.mockRejectedValue( + new AgentRunOwnershipLostError({ + runUuid: 'run-1', + expectedExecutionOwner: 'worker-1', + currentStatus: 'running', + currentExecutionOwner: 'worker-2', + }) + ); + + await expect( + processAgentRunExecute({ + data: { + runId: 'run-1', + dispatchAttemptId: 'attempt-1', + reason: 'resume', + }, + } as any) + ).resolves.toBeUndefined(); + + expect(mockMarkFailedForExecutionOwner).not.toHaveBeenCalled(); + }); +}); diff --git a/src/server/jobs/__tests__/agentSessionCleanup.test.ts b/src/server/jobs/__tests__/agentSessionCleanup.test.ts index b9cc4e8d..f7eb9577 100644 --- a/src/server/jobs/__tests__/agentSessionCleanup.test.ts +++ b/src/server/jobs/__tests__/agentSessionCleanup.test.ts @@ -15,21 +15,58 @@ */ jest.mock('server/models/AgentSession'); -jest.mock('server/services/agentSession'); +jest.mock('server/services/agentSession', () => { + class ActiveAgentRunSuspensionError extends Error { + constructor() { + super('Cannot suspend a chat runtime while an agent run is active'); + this.name = 'ActiveAgentRunSuspensionError'; + } + } + + return { + __esModule: true, + ActiveAgentRunSuspensionError, + default: { + endSession: jest.fn(), + suspendChatRuntime: jest.fn(), + }, + }; +}); jest.mock('server/lib/logger', () => ({ getLogger: jest.fn(() => ({ info: jest.fn(), error: jest.fn(), })), })); +jest.mock('server/lib/agentSession/runtimeConfig', () => { + const actual = jest.requireActual('server/lib/agentSession/runtimeConfig'); + return { + __esModule: true, + ...actual, + resolveAgentSessionCleanupConfig: jest.fn().mockResolvedValue({ + activeIdleSuspendMs: 30 * 60 * 1000, + startingTimeoutMs: 15 * 60 * 1000, + hibernatedRetentionMs: 24 * 60 * 60 * 1000, + intervalMs: 5 * 60 * 1000, + redisTtlSeconds: 7200, + }), + }; +}); import AgentSession from 'server/models/AgentSession'; -import AgentSessionService from 'server/services/agentSession'; +import AgentSessionService, { ActiveAgentRunSuspensionError } from 'server/services/agentSession'; +import { getLogger } from 'server/lib/logger'; import { processAgentSessionCleanup } from '../agentSessionCleanup'; describe('agentSessionCleanup', () => { + const mockLogger = { + info: jest.fn(), + error: jest.fn(), + }; + beforeEach(() => { jest.clearAllMocks(); + (getLogger as jest.Mock).mockReturnValue(mockLogger); jest.useFakeTimers().setSystemTime(new Date('2026-03-23T12:00:00.000Z')); }); @@ -60,21 +97,152 @@ describe('agentSessionCleanup', () => { const activeQuery = { where: jest.fn() }; activeQuery.where .mockImplementationOnce(() => activeQuery) - .mockImplementationOnce(() => Promise.resolve(activeSessions)); + .mockImplementationOnce(() => activeQuery) + .mockImplementationOnce((callback) => { + callback({ + whereNot: jest.fn().mockReturnValue({ + orWhereNot: jest.fn(), + }), + }); + return Promise.resolve(activeSessions); + }); const startingQuery = { where: jest.fn() }; startingQuery.where .mockImplementationOnce(() => startingQuery) .mockImplementationOnce(() => Promise.resolve(startingSessions)); - (AgentSession.query as jest.Mock) = jest.fn().mockReturnValueOnce(activeQuery).mockReturnValueOnce(startingQuery); + const suspendedQuery = { where: jest.fn() }; + suspendedQuery.where + .mockImplementationOnce(() => suspendedQuery) + .mockImplementationOnce(() => suspendedQuery) + .mockImplementationOnce(() => suspendedQuery) + .mockImplementationOnce(() => Promise.resolve([])); + + (AgentSession.query as jest.Mock) = jest + .fn() + .mockReturnValueOnce(activeQuery) + .mockReturnValueOnce(startingQuery) + .mockReturnValueOnce(suspendedQuery); (AgentSessionService.endSession as jest.Mock).mockResolvedValue(undefined); await processAgentSessionCleanup(); - expect(AgentSession.query).toHaveBeenCalledTimes(2); + expect(AgentSession.query).toHaveBeenCalledTimes(3); expect(AgentSessionService.endSession).toHaveBeenCalledTimes(2); expect(AgentSessionService.endSession).toHaveBeenNthCalledWith(1, 'active-session'); expect(AgentSessionService.endSession).toHaveBeenNthCalledWith(2, 'starting-session'); }); + + it('suspends idle chat runtimes before terminal cleanup', async () => { + const activeSessions = [ + { + id: 1, + uuid: 'chat-session', + userId: 'sample-user', + sessionKind: 'chat', + workspaceStatus: 'ready', + status: 'active', + lastActivity: '2026-03-23T11:00:00.000Z', + updatedAt: '2026-03-23T11:00:00.000Z', + }, + ]; + + const activeQuery = { where: jest.fn() }; + activeQuery.where + .mockImplementationOnce(() => activeQuery) + .mockImplementationOnce(() => activeQuery) + .mockImplementationOnce((callback) => { + callback({ + whereNot: jest.fn().mockReturnValue({ + orWhereNot: jest.fn(), + }), + }); + return Promise.resolve(activeSessions); + }); + + const emptyTwoWhereQuery = { where: jest.fn() }; + emptyTwoWhereQuery.where + .mockImplementationOnce(() => emptyTwoWhereQuery) + .mockImplementationOnce(() => Promise.resolve([])); + + const emptyFourWhereQuery = { where: jest.fn() }; + emptyFourWhereQuery.where + .mockImplementationOnce(() => emptyFourWhereQuery) + .mockImplementationOnce(() => emptyFourWhereQuery) + .mockImplementationOnce(() => emptyFourWhereQuery) + .mockImplementationOnce(() => Promise.resolve([])); + + (AgentSession.query as jest.Mock) = jest + .fn() + .mockReturnValueOnce(activeQuery) + .mockReturnValueOnce(emptyTwoWhereQuery) + .mockReturnValueOnce(emptyFourWhereQuery); + (AgentSessionService.suspendChatRuntime as jest.Mock).mockResolvedValue(undefined); + + await processAgentSessionCleanup(); + + expect(AgentSessionService.suspendChatRuntime).toHaveBeenCalledWith({ + sessionId: 'chat-session', + userId: 'sample-user', + }); + expect(AgentSessionService.endSession).not.toHaveBeenCalled(); + }); + + it('skips idle chat suspension when a run is still active', async () => { + const activeSessions = [ + { + id: 1, + uuid: 'chat-session', + userId: 'sample-user', + sessionKind: 'chat', + workspaceStatus: 'ready', + status: 'active', + lastActivity: '2026-03-23T11:00:00.000Z', + updatedAt: '2026-03-23T11:00:00.000Z', + }, + ]; + + const activeQuery = { where: jest.fn() }; + activeQuery.where + .mockImplementationOnce(() => activeQuery) + .mockImplementationOnce(() => activeQuery) + .mockImplementationOnce((callback) => { + callback({ + whereNot: jest.fn().mockReturnValue({ + orWhereNot: jest.fn(), + }), + }); + return Promise.resolve(activeSessions); + }); + + const emptyTwoWhereQuery = { where: jest.fn() }; + emptyTwoWhereQuery.where + .mockImplementationOnce(() => emptyTwoWhereQuery) + .mockImplementationOnce(() => Promise.resolve([])); + + const emptyFourWhereQuery = { where: jest.fn() }; + emptyFourWhereQuery.where + .mockImplementationOnce(() => emptyFourWhereQuery) + .mockImplementationOnce(() => emptyFourWhereQuery) + .mockImplementationOnce(() => emptyFourWhereQuery) + .mockImplementationOnce(() => Promise.resolve([])); + + (AgentSession.query as jest.Mock) = jest + .fn() + .mockReturnValueOnce(activeQuery) + .mockReturnValueOnce(emptyTwoWhereQuery) + .mockReturnValueOnce(emptyFourWhereQuery); + (AgentSessionService.suspendChatRuntime as jest.Mock).mockRejectedValue(new ActiveAgentRunSuspensionError()); + + await processAgentSessionCleanup(); + + expect(AgentSessionService.suspendChatRuntime).toHaveBeenCalledWith({ + sessionId: 'chat-session', + userId: 'sample-user', + }); + expect(AgentSessionService.endSession).not.toHaveBeenCalled(); + expect(mockLogger.error).not.toHaveBeenCalled(); + expect(mockLogger.info).toHaveBeenCalledWith('Session: cleanup skipped sessionId=chat-session reason=active_run'); + }); }); diff --git a/src/server/jobs/agentRunDispatchRecovery.ts b/src/server/jobs/agentRunDispatchRecovery.ts new file mode 100644 index 00000000..5a6f25a6 --- /dev/null +++ b/src/server/jobs/agentRunDispatchRecovery.ts @@ -0,0 +1,62 @@ +/** + * Copyright 2026 GoodRx, Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +import { getLogger } from 'server/lib/logger'; +import AgentRunQueueService from 'server/services/agent/RunQueueService'; +import AgentRunService from 'server/services/agent/RunService'; + +const logger = () => getLogger(); + +export async function processAgentRunDispatchRecovery(): Promise<{ + runs: number; + enqueued: Array<{ runId: string; dispatchAttemptId: string }>; + failed: Array<{ runId: string }>; +}> { + const runs = await AgentRunService.listRunsNeedingDispatch(); + if (runs.length === 0) { + return { + runs: 0, + enqueued: [], + failed: [], + }; + } + + logger().info(`AgentExec: recovery enqueue runs=${runs.length}`); + const enqueued: Array<{ runId: string; dispatchAttemptId: string }> = []; + const failed: Array<{ runId: string }> = []; + for (const run of runs) { + try { + const dispatch = await AgentRunQueueService.enqueueRun(run.uuid, 'resume'); + enqueued.push({ + runId: run.uuid, + dispatchAttemptId: dispatch.dispatchAttemptId, + }); + logger().info( + { runId: run.uuid, reason: 'resume', dispatchAttemptId: dispatch.dispatchAttemptId }, + `AgentExec: recovery enqueued runId=${run.uuid} reason=resume dispatchAttemptId=${dispatch.dispatchAttemptId}` + ); + } catch (error) { + failed.push({ runId: run.uuid }); + logger().warn({ error, runId: run.uuid }, `AgentExec: recovery enqueue failed runId=${run.uuid}`); + } + } + + return { + runs: runs.length, + enqueued, + failed, + }; +} diff --git a/src/server/jobs/agentRunExecute.ts b/src/server/jobs/agentRunExecute.ts new file mode 100644 index 00000000..b9ffcd67 --- /dev/null +++ b/src/server/jobs/agentRunExecute.ts @@ -0,0 +1,112 @@ +/** + * Copyright 2026 GoodRx, Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +import { Job } from 'bullmq'; +import { randomBytes } from 'crypto'; +import os from 'os'; +import { getLogger } from 'server/lib/logger'; +import { withLogContext } from 'server/lib/logger/context'; +import { decrypt } from 'server/lib/encryption'; +import LifecycleAiSdkHarness from 'server/services/agent/LifecycleAiSdkHarness'; +import AgentRunService from 'server/services/agent/RunService'; +import { AgentRunOwnershipLostError } from 'server/services/agent/AgentRunOwnershipLostError'; +import type { AgentRunExecuteJob } from 'server/services/agent/RunQueueService'; + +const logger = () => getLogger(); + +function requireJobString(value: unknown, field: keyof AgentRunExecuteJob): string { + if (typeof value !== 'string' || !value.trim()) { + throw new Error(`Invalid agent run execute job payload: ${field} is required`); + } + + return value.trim(); +} + +function buildExecutionOwner(jobId: string): string { + return `bull:${jobId}:${os.hostname()}:${process.pid}:${randomBytes(6).toString('hex')}`; +} + +export async function processAgentRunExecute(job: Job): Promise { + await withLogContext(job.data, async () => { + const runId = requireJobString(job.data.runId, 'runId'); + const dispatchAttemptId = requireJobString(job.data.dispatchAttemptId, 'dispatchAttemptId'); + const executionOwner = buildExecutionOwner(String(job.id || 'unknown')); + const run = await AgentRunService.claimQueuedRunForExecution(runId, executionOwner); + if (!run) { + logger().info( + `AgentExec: queued run skip runId=${runId} reason=${ + job.data.reason || 'submit' + } dispatchAttemptId=${dispatchAttemptId} owner=${executionOwner}` + ); + return; + } + + try { + logger().info( + `AgentExec: queued run start runId=${run.uuid} reason=${ + job.data.reason || 'submit' + } dispatchAttemptId=${dispatchAttemptId} owner=${executionOwner}` + ); + await LifecycleAiSdkHarness.executeRun(run, { + requestGitHubToken: job.data.encryptedGithubToken ? decrypt(job.data.encryptedGithubToken) : null, + dispatchAttemptId, + }); + logger().info( + `AgentExec: queued run finish runId=${run.uuid} dispatchAttemptId=${dispatchAttemptId} owner=${executionOwner}` + ); + } catch (error) { + if (error instanceof AgentRunOwnershipLostError) { + logger().info( + { + runId: run.uuid, + owner: executionOwner, + currentStatus: error.currentStatus || null, + currentOwner: error.currentExecutionOwner || null, + }, + `AgentExec: ownership lost runId=${run.uuid} owner=${executionOwner}` + ); + return; + } + + const latestRun = await AgentRunService.getRunByUuid(run.uuid); + if (!latestRun || !AgentRunService.isTerminalStatus(latestRun.status)) { + await AgentRunService.markFailedForExecutionOwner(run.uuid, executionOwner, error, undefined, { + dispatchAttemptId, + }).catch((failureRecordError) => { + if (failureRecordError instanceof AgentRunOwnershipLostError) { + logger().info( + { + runId: run.uuid, + owner: executionOwner, + currentStatus: failureRecordError.currentStatus || null, + currentOwner: failureRecordError.currentExecutionOwner || null, + }, + `AgentExec: ownership lost runId=${run.uuid} owner=${executionOwner}` + ); + return; + } + + logger().warn( + { error: failureRecordError, runId: run.uuid }, + `AgentExec: queued run failure record failed runId=${run.uuid}` + ); + }); + } + + throw error; + } + }); +} diff --git a/src/server/jobs/agentSandboxSessionLaunch.ts b/src/server/jobs/agentSandboxSessionLaunch.ts index e8e1a8b2..1b3c67f7 100644 --- a/src/server/jobs/agentSandboxSessionLaunch.ts +++ b/src/server/jobs/agentSandboxSessionLaunch.ts @@ -36,8 +36,6 @@ const logger = () => getLogger(); export interface SandboxSessionLaunchJob extends Omit { launchId: string; encryptedGithubToken?: string | null; - encryptedRequestApiKey?: string | null; - requestApiKeyProvider?: string | null; } export async function processAgentSandboxSessionLaunch(job: Job): Promise { @@ -47,8 +45,6 @@ export async function processAgentSandboxSessionLaunch(job: Job => { @@ -75,8 +73,6 @@ export async function processAgentSandboxSessionLaunch(job: Job getLogger(); -const IDLE_TIMEOUT_MS = 30 * 60 * 1000; -const STARTING_TIMEOUT_MS = 15 * 60 * 1000; export async function processAgentSessionCleanup(): Promise { - const activeCutoff = new Date(Date.now() - IDLE_TIMEOUT_MS); - const startingCutoff = new Date(Date.now() - STARTING_TIMEOUT_MS); + const cleanupConfig = await resolveAgentSessionCleanupConfig(); + const activeCutoff = new Date(Date.now() - cleanupConfig.activeIdleSuspendMs); + const startingCutoff = new Date(Date.now() - cleanupConfig.startingTimeoutMs); + const suspendedExpiryCutoff = new Date(Date.now() - cleanupConfig.hibernatedRetentionMs); + const idleActiveSessions = await AgentSession.query() + .where('status', 'active') + .where('lastActivity', '<', activeCutoff) + .where((builder) => { + builder + .whereNot('sessionKind', AgentSessionKind.CHAT) + .orWhereNot('workspaceStatus', AgentWorkspaceStatus.HIBERNATED); + }); const staleSessions = [ - ...(await AgentSession.query().where('status', 'active').where('lastActivity', '<', activeCutoff)), + ...idleActiveSessions, ...(await AgentSession.query().where('status', 'starting').where('updatedAt', '<', startingCutoff)), + ...(await AgentSession.query() + .where('status', 'active') + .where('sessionKind', AgentSessionKind.CHAT) + .where('workspaceStatus', AgentWorkspaceStatus.HIBERNATED) + .where('updatedAt', '<', suspendedExpiryCutoff)), ]; for (const session of staleSessions) { const sessionId = session.uuid || String(session.id); try { + if ( + session.status === 'active' && + session.sessionKind === AgentSessionKind.CHAT && + session.workspaceStatus !== AgentWorkspaceStatus.HIBERNATED + ) { + logger().info(`Session: cleanup suspending sessionId=${sessionId} lastActivity=${session.lastActivity}`); + await AgentSessionService.suspendChatRuntime({ + sessionId, + userId: session.userId, + }); + continue; + } + logger().info( `Session: cleanup starting sessionId=${sessionId} status=${session.status} lastActivity=${session.lastActivity}` ); await AgentSessionService.endSession(sessionId); } catch (err) { + if (err instanceof ActiveAgentRunSuspensionError) { + logger().info(`Session: cleanup skipped sessionId=${sessionId} reason=active_run`); + continue; + } logger().error({ error: err, sessionId }, `Session: cleanup failed sessionId=${sessionId}`); } } diff --git a/src/server/jobs/index.ts b/src/server/jobs/index.ts index 1cfabb42..e5b65d3d 100644 --- a/src/server/jobs/index.ts +++ b/src/server/jobs/index.ts @@ -23,6 +23,12 @@ import { MAX_GITHUB_API_REQUEST, GITHUB_API_REQUEST_INTERVAL, QUEUE_NAMES } from import { processAgentSessionCleanup } from './agentSessionCleanup'; import { processAgentSessionPrewarm } from './agentSessionPrewarm'; import { processAgentSandboxSessionLaunch } from './agentSandboxSessionLaunch'; +import { processAgentRunExecute } from './agentRunExecute'; +import { processAgentRunDispatchRecovery } from './agentRunDispatchRecovery'; +import { + DEFAULT_AGENT_SESSION_CLEANUP_INTERVAL_MS, + resolveAgentSessionCleanupConfig, +} from 'server/lib/agentSession/runtimeConfig'; let isBootstrapped = false; @@ -121,6 +127,16 @@ export default function bootstrapJobs(services: IServices) { concurrency: 5, }); + queueManager.registerWorker(QUEUE_NAMES.AGENT_RUN_EXECUTE, processAgentRunExecute, { + connection: redisClient.getConnection(), + concurrency: 5, + }); + + queueManager.registerWorker(QUEUE_NAMES.AGENT_RUN_RECOVERY, processAgentRunDispatchRecovery, { + connection: redisClient.getConnection(), + concurrency: 1, + }); + const agentCleanupQueue = queueManager.registerQueue(QUEUE_NAMES.AGENT_SESSION_CLEANUP, { connection: redisClient.getConnection(), defaultJobOptions: { @@ -130,12 +146,47 @@ export default function bootstrapJobs(services: IServices) { }, }); - agentCleanupQueue.add( - 'agent-session-cleanup', + void resolveAgentSessionCleanupConfig() + .then((cleanupConfig) => + agentCleanupQueue.add( + 'agent-session-cleanup', + {}, + { + repeat: { + every: cleanupConfig.intervalMs, + }, + } + ) + ) + .catch((error) => { + getLogger().warn({ error }, 'Jobs: cleanup schedule config failed queue=agent_session_cleanup'); + return agentCleanupQueue.add( + 'agent-session-cleanup', + {}, + { + repeat: { + every: DEFAULT_AGENT_SESSION_CLEANUP_INTERVAL_MS, + }, + } + ); + }); + + const agentRunRecoveryQueue = queueManager.registerQueue(QUEUE_NAMES.AGENT_RUN_RECOVERY, { + connection: redisClient.getConnection(), + defaultJobOptions: { + attempts: 1, + removeOnComplete: true, + removeOnFail: false, + }, + }); + + agentRunRecoveryQueue.add( + 'agent-run-recovery', {}, { + jobId: 'agent-run-recovery', repeat: { - every: 5 * 60 * 1000, + every: 60 * 1000, }, } ); diff --git a/src/server/lib/__mocks__/redisClientMock.ts b/src/server/lib/__mocks__/redisClientMock.ts index bcb1fff1..fa1f4a77 100644 --- a/src/server/lib/__mocks__/redisClientMock.ts +++ b/src/server/lib/__mocks__/redisClientMock.ts @@ -29,6 +29,7 @@ jest.mock('ioredis', () => { info = jest.fn().mockResolvedValue('redis_version:6.0.5'); hgetall = jest.fn().mockResolvedValue(null); hset = jest.fn().mockResolvedValue(1); + del = jest.fn().mockResolvedValue(1); expire = jest.fn().mockResolvedValue(1); hmset = jest.fn().mockResolvedValue('OK'); disconnect = jest.fn(); diff --git a/src/server/lib/agentSession/__tests__/pvcFactory.test.ts b/src/server/lib/agentSession/__tests__/pvcFactory.test.ts index a4bf64cd..2415670f 100644 --- a/src/server/lib/agentSession/__tests__/pvcFactory.test.ts +++ b/src/server/lib/agentSession/__tests__/pvcFactory.test.ts @@ -17,6 +17,7 @@ import * as k8s from '@kubernetes/client-node'; const mockCreatePvc = jest.fn(); +const mockReadPvc = jest.fn(); const mockDeletePvc = jest.fn(); jest.mock('@kubernetes/client-node', () => { @@ -27,6 +28,7 @@ jest.mock('@kubernetes/client-node', () => { loadFromDefault: jest.fn(), makeApiClient: jest.fn().mockReturnValue({ createNamespacedPersistentVolumeClaim: mockCreatePvc, + readNamespacedPersistentVolumeClaim: mockReadPvc, deleteNamespacedPersistentVolumeClaim: mockDeletePvc, }), })), @@ -94,15 +96,25 @@ describe('pvcFactory', () => { expect(pvcBody.spec.accessModes).toEqual(['ReadWriteOnce']); }); - it('honors AGENT_SESSION_PVC_ACCESS_MODE when configured', async () => { - process.env.AGENT_SESSION_PVC_ACCESS_MODE = 'ReadWriteMany'; + it('honors the configured access mode when provided', async () => { mockCreatePvc.mockResolvedValue({ body: { metadata: { name: 'test-pvc' } } }); - await createAgentPvc('test-ns', 'test-pvc'); + await createAgentPvc('test-ns', 'test-pvc', '10Gi', undefined, 'ReadWriteMany'); const [, pvcBody] = mockCreatePvc.mock.calls[0]; expect(pvcBody.spec.accessModes).toEqual(['ReadWriteMany']); }); + + it('reuses an existing PVC on resume', async () => { + mockCreatePvc.mockRejectedValue(new k8s.HttpError({ statusCode: 409 } as any, 'already exists', 409)); + mockReadPvc.mockResolvedValue({ body: { metadata: { name: 'test-pvc' } } }); + + await expect(createAgentPvc('test-ns', 'test-pvc')).resolves.toEqual({ + metadata: { name: 'test-pvc' }, + }); + + expect(mockReadPvc).toHaveBeenCalledWith('test-pvc', 'test-ns'); + }); }); describe('deleteAgentPvc', () => { diff --git a/src/server/lib/agentSession/__tests__/runtimeConfig.test.ts b/src/server/lib/agentSession/__tests__/runtimeConfig.test.ts index 85910268..5c9698a3 100644 --- a/src/server/lib/agentSession/__tests__/runtimeConfig.test.ts +++ b/src/server/lib/agentSession/__tests__/runtimeConfig.test.ts @@ -26,9 +26,23 @@ jest.mock('server/services/globalConfig', () => ({ })); import { + AgentSessionWorkspaceStorageConfigError, AgentSessionRuntimeConfigError, + DEFAULT_AGENT_SESSION_ACTIVE_IDLE_SUSPEND_MS, + DEFAULT_AGENT_SESSION_CLEANUP_INTERVAL_MS, + DEFAULT_AGENT_SESSION_DISPATCH_RECOVERY_LIMIT, + DEFAULT_AGENT_SESSION_FILE_CHANGE_PREVIEW_CHARS, + DEFAULT_AGENT_SESSION_HIBERNATED_RETENTION_MS, DEFAULT_AGENT_SESSION_KEEP_ATTACHED_SERVICES_ON_SESSION_NODE, DEFAULT_AGENT_SESSION_MAX_ITERATIONS, + DEFAULT_AGENT_SESSION_MAX_DURABLE_PAYLOAD_BYTES, + DEFAULT_AGENT_SESSION_PAYLOAD_PREVIEW_BYTES, + DEFAULT_AGENT_SESSION_QUEUED_RUN_DISPATCH_STALE_MS, + DEFAULT_AGENT_SESSION_REDIS_TTL_SECONDS, + DEFAULT_AGENT_SESSION_RUN_EXECUTION_LEASE_MS, + DEFAULT_AGENT_SESSION_STARTING_TIMEOUT_MS, + DEFAULT_AGENT_SESSION_WORKSPACE_STORAGE_ACCESS_MODE, + DEFAULT_AGENT_SESSION_WORKSPACE_STORAGE_SIZE, DEFAULT_AGENT_SESSION_WORKSPACE_TOOL_DISCOVERY_TIMEOUT_MS, DEFAULT_AGENT_SESSION_WORKSPACE_TOOL_EXECUTION_TIMEOUT_MS, mergeAgentSessionReadiness, @@ -36,9 +50,13 @@ import { mergeAgentSessionResources, resolveAgentSessionControlPlaneConfig, resolveAgentSessionControlPlaneConfigFromDefaults, + resolveAgentSessionDurabilityFromDefaults, + resolveAgentSessionCleanupFromDefaults, resolveAgentSessionReadinessFromDefaults, resolveAgentSessionResourcesFromDefaults, resolveAgentSessionRuntimeConfig, + resolveAgentSessionWorkspaceStorageFromDefaults, + resolveAgentSessionWorkspaceStorageIntent, } from '../runtimeConfig'; const DEFAULT_READINESS = { @@ -79,11 +97,38 @@ const DEFAULT_RESOURCES = { }, }; +const DEFAULT_WORKSPACE_STORAGE = { + defaultSize: DEFAULT_AGENT_SESSION_WORKSPACE_STORAGE_SIZE, + allowedSizes: [DEFAULT_AGENT_SESSION_WORKSPACE_STORAGE_SIZE], + allowClientOverride: false, + accessMode: DEFAULT_AGENT_SESSION_WORKSPACE_STORAGE_ACCESS_MODE, +}; + +const DEFAULT_CLEANUP = { + activeIdleSuspendMs: DEFAULT_AGENT_SESSION_ACTIVE_IDLE_SUSPEND_MS, + startingTimeoutMs: DEFAULT_AGENT_SESSION_STARTING_TIMEOUT_MS, + hibernatedRetentionMs: DEFAULT_AGENT_SESSION_HIBERNATED_RETENTION_MS, + intervalMs: DEFAULT_AGENT_SESSION_CLEANUP_INTERVAL_MS, + redisTtlSeconds: DEFAULT_AGENT_SESSION_REDIS_TTL_SECONDS, +}; + +const DEFAULT_DURABILITY = { + runExecutionLeaseMs: DEFAULT_AGENT_SESSION_RUN_EXECUTION_LEASE_MS, + queuedRunDispatchStaleMs: DEFAULT_AGENT_SESSION_QUEUED_RUN_DISPATCH_STALE_MS, + dispatchRecoveryLimit: DEFAULT_AGENT_SESSION_DISPATCH_RECOVERY_LIMIT, + maxDurablePayloadBytes: DEFAULT_AGENT_SESSION_MAX_DURABLE_PAYLOAD_BYTES, + payloadPreviewBytes: DEFAULT_AGENT_SESSION_PAYLOAD_PREVIEW_BYTES, + fileChangePreviewChars: DEFAULT_AGENT_SESSION_FILE_CHANGE_PREVIEW_CHARS, +}; + function buildExpectedRuntimeConfig(overrides?: { nodeSelector?: Record; keepAttachedServicesOnSessionNode?: boolean; readiness?: typeof DEFAULT_READINESS; resources?: typeof DEFAULT_RESOURCES; + workspaceStorage?: typeof DEFAULT_WORKSPACE_STORAGE; + cleanup?: typeof DEFAULT_CLEANUP; + durability?: typeof DEFAULT_DURABILITY; }) { return { workspaceImage: 'lifecycle-workspace:sha-123', @@ -93,6 +138,9 @@ function buildExpectedRuntimeConfig(overrides?: { keepAttachedServicesOnSessionNode: DEFAULT_AGENT_SESSION_KEEP_ATTACHED_SERVICES_ON_SESSION_NODE, readiness: DEFAULT_READINESS, resources: DEFAULT_RESOURCES, + workspaceStorage: DEFAULT_WORKSPACE_STORAGE, + cleanup: DEFAULT_CLEANUP, + durability: DEFAULT_DURABILITY, ...overrides, }; } @@ -233,6 +281,88 @@ describe('runtimeConfig', () => { ); }); + it('returns configured workspace storage, cleanup, and durability settings when present', async () => { + getAllConfigs.mockResolvedValue({ + agentSessionDefaults: { + workspaceImage: 'lifecycle-workspace:sha-123', + workspaceEditorImage: 'codercom/code-server:4.98.2', + workspaceStorage: { + defaultSize: '20Gi', + allowedSizes: ['10Gi', '20Gi'], + allowClientOverride: true, + accessMode: 'ReadWriteMany', + }, + cleanup: { + activeIdleSuspendMs: 60_000, + startingTimeoutMs: 120_000, + hibernatedRetentionMs: 180_000, + intervalMs: 30_000, + redisTtlSeconds: 900, + }, + durability: { + runExecutionLeaseMs: 45_000, + queuedRunDispatchStaleMs: 5_000, + dispatchRecoveryLimit: 12, + maxDurablePayloadBytes: 4096, + payloadPreviewBytes: 512, + fileChangePreviewChars: 600, + }, + }, + }); + + await expect(resolveAgentSessionRuntimeConfig()).resolves.toEqual( + buildExpectedRuntimeConfig({ + workspaceStorage: { + defaultSize: '20Gi', + allowedSizes: ['10Gi', '20Gi'], + allowClientOverride: true, + accessMode: 'ReadWriteMany', + }, + cleanup: { + activeIdleSuspendMs: 60_000, + startingTimeoutMs: 120_000, + hibernatedRetentionMs: 180_000, + intervalMs: 30_000, + redisTtlSeconds: 900, + }, + durability: { + runExecutionLeaseMs: 45_000, + queuedRunDispatchStaleMs: 5_000, + dispatchRecoveryLimit: 12, + maxDurablePayloadBytes: 4096, + payloadPreviewBytes: 512, + fileChangePreviewChars: 600, + }, + }) + ); + }); + + it('resolves client workspace storage intent only when overrides are enabled and allowed', () => { + const storage = resolveAgentSessionWorkspaceStorageFromDefaults({ + defaultSize: '10Gi', + allowedSizes: ['10Gi', '20Gi'], + allowClientOverride: true, + accessMode: 'ReadWriteOnce', + }); + + expect(resolveAgentSessionWorkspaceStorageIntent({ requestedSize: '20Gi', storage })).toEqual({ + requestedSize: '20Gi', + storageSize: '20Gi', + accessMode: 'ReadWriteOnce', + }); + expect(() => + resolveAgentSessionWorkspaceStorageIntent({ + requestedSize: '30Gi', + storage, + }) + ).toThrow(AgentSessionWorkspaceStorageConfigError); + }); + + it('resolves cleanup and durability defaults independently', () => { + expect(resolveAgentSessionCleanupFromDefaults()).toEqual(DEFAULT_CLEANUP); + expect(resolveAgentSessionDurabilityFromDefaults()).toEqual(DEFAULT_DURABILITY); + }); + it('returns the configured control-plane append prompt from the neutral path', async () => { getAllConfigs.mockResolvedValue({ agentSessionDefaults: { @@ -264,7 +394,8 @@ describe('runtimeConfig', () => { 'Do not claim that a file was read, a command was run, or a change was made unless that happened through an actual tool call in this conversation.\n' + 'If a tool call fails or a capability is unavailable, say that plainly and explain what failed.', appendSystemPrompt: - 'When a tool execution is not approved, do not retry the denied action. Use the denial reason as updated guidance and continue from there.', + 'When a tool execution is not approved, do not retry the denied action. Use the denial reason as updated guidance and continue from there.\n' + + 'When showing multi-line exact text such as file contents, command output, diffs, or JSON, use a fenced code block instead of inline code.', maxIterations: DEFAULT_AGENT_SESSION_MAX_ITERATIONS, workspaceToolDiscoveryTimeoutMs: DEFAULT_AGENT_SESSION_WORKSPACE_TOOL_DISCOVERY_TIMEOUT_MS, workspaceToolExecutionTimeoutMs: DEFAULT_AGENT_SESSION_WORKSPACE_TOOL_EXECUTION_TIMEOUT_MS, diff --git a/src/server/lib/agentSession/__tests__/systemPrompt.test.ts b/src/server/lib/agentSession/__tests__/systemPrompt.test.ts index b0264d89..2b840154 100644 --- a/src/server/lib/agentSession/__tests__/systemPrompt.test.ts +++ b/src/server/lib/agentSession/__tests__/systemPrompt.test.ts @@ -48,7 +48,7 @@ describe('agent session system prompt', () => { skillsAvailable: true, toolLines: [ '- inspect files, services, and git state: workspace.read_file, workspace.exec', - '- run mutating or networked shell commands: workspace.exec_mutation', + '- run mutating or networked shell commands that are not direct file edits: workspace.exec_mutation', ], services: [ { @@ -68,7 +68,7 @@ describe('agent session system prompt', () => { '- equipped skills: use skills.list to discover them and skills.learn to load a skill before using it', '- equipped tools:', ' - inspect files, services, and git state: workspace.read_file, workspace.exec', - ' - run mutating or networked shell commands: workspace.exec_mutation', + ' - run mutating or networked shell commands that are not direct file edits: workspace.exec_mutation', ].join('\n') ); }); diff --git a/src/server/lib/agentSession/chatPreviewFactory.ts b/src/server/lib/agentSession/chatPreviewFactory.ts new file mode 100644 index 00000000..f0d6b962 --- /dev/null +++ b/src/server/lib/agentSession/chatPreviewFactory.ts @@ -0,0 +1,226 @@ +/** + * Copyright 2026 GoodRx, Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +import * as k8s from '@kubernetes/client-node'; +import { APP_HOST } from 'shared/config'; +import { buildLifecycleLabels } from 'server/lib/kubernetes/labels'; +import { normalizeKubernetesLabelValue } from 'server/lib/kubernetes/utils'; +import GlobalConfigService from 'server/services/globalConfig'; + +export interface ChatPreviewPublication { + url: string; + host: string | null; + path: string; + serviceName: string; + ingressName: string; + port: number; +} + +function getClients() { + const kc = new k8s.KubeConfig(); + kc.loadFromDefault(); + + return { + coreApi: kc.makeApiClient(k8s.CoreV1Api), + networkingApi: kc.makeApiClient(k8s.NetworkingV1Api), + }; +} + +function buildResourceName(prefix: string, sessionUuid: string, port: number): string { + return normalizeKubernetesLabelValue(`${prefix}-${sessionUuid.slice(0, 8)}-${port}`).replace(/[_.]/g, '-'); +} + +function buildPreviewPath(sessionUuid: string, port: number): string { + return `/_chat/${sessionUuid}/${port}`; +} + +function resolvePreviewUrl({ + sessionUuid, + port, + httpDomain, +}: { + sessionUuid: string; + port: number; + httpDomain?: string | null; +}): Pick { + const previewPath = buildPreviewPath(sessionUuid, port); + const appUrl = new URL(APP_HOST); + + if (httpDomain?.trim()) { + const host = `${buildResourceName('chat', sessionUuid, port)}.${httpDomain.trim()}`; + return { + url: `${appUrl.protocol}//${host}`, + host, + path: '/', + }; + } + + return { + url: new URL(previewPath, APP_HOST).toString(), + host: appUrl.hostname, + path: previewPath, + }; +} + +async function upsertService(coreApi: k8s.CoreV1Api, namespace: string, service: k8s.V1Service): Promise { + try { + const existing = await coreApi.readNamespacedService(service.metadata!.name!, namespace); + service.metadata = { + ...(service.metadata || {}), + resourceVersion: existing.body.metadata?.resourceVersion, + }; + await coreApi.replaceNamespacedService(service.metadata!.name!, namespace, service); + } catch (error) { + if (error instanceof k8s.HttpError && error.response?.statusCode === 404) { + await coreApi.createNamespacedService(namespace, service); + return; + } + + throw error; + } +} + +async function upsertIngress( + networkingApi: k8s.NetworkingV1Api, + namespace: string, + ingress: k8s.V1Ingress +): Promise { + try { + const existing = await networkingApi.readNamespacedIngress(ingress.metadata!.name!, namespace); + ingress.metadata = { + ...(ingress.metadata || {}), + resourceVersion: existing.body.metadata?.resourceVersion, + }; + await networkingApi.replaceNamespacedIngress(ingress.metadata!.name!, namespace, ingress); + } catch (error) { + if (error instanceof k8s.HttpError && error.response?.statusCode === 404) { + await networkingApi.createNamespacedIngress(namespace, ingress); + return; + } + + throw error; + } +} + +export async function createOrUpdateChatPreview({ + sessionUuid, + namespace, + podName, + port, +}: { + sessionUuid: string; + namespace: string; + podName: string; + port: number; +}): Promise { + const { coreApi, networkingApi } = getClients(); + const { lifecycleDefaults, domainDefaults } = await GlobalConfigService.getInstance().getAllConfigs(); + const publication = resolvePreviewUrl({ + sessionUuid, + port, + httpDomain: domainDefaults?.http, + }); + const serviceName = buildResourceName('agent-preview', sessionUuid, port); + const ingressName = buildResourceName('agent-preview-ingress', sessionUuid, port); + const labels = { + ...buildLifecycleLabels(), + 'app.kubernetes.io/component': 'agent-session-preview', + 'lfc/agent-session': sessionUuid, + }; + + const service: k8s.V1Service = { + apiVersion: 'v1', + kind: 'Service', + metadata: { + name: serviceName, + namespace, + labels, + }, + spec: { + selector: { + 'app.kubernetes.io/name': podName, + }, + ports: [ + { + name: 'http', + port: 80, + targetPort: port, + }, + ], + }, + }; + + const ingressAnnotations: Record = {}; + const pathRule = + publication.path === '/' + ? { + path: '/', + pathType: 'Prefix' as const, + } + : { + path: `${publication.path}(/|$)(.*)`, + pathType: 'ImplementationSpecific' as const, + }; + + if (publication.path !== '/') { + ingressAnnotations['nginx.ingress.kubernetes.io/use-regex'] = 'true'; + ingressAnnotations['nginx.ingress.kubernetes.io/rewrite-target'] = '/$2'; + } + + const ingress: k8s.V1Ingress = { + apiVersion: 'networking.k8s.io/v1', + kind: 'Ingress', + metadata: { + name: ingressName, + namespace, + labels, + ...(Object.keys(ingressAnnotations).length > 0 ? { annotations: ingressAnnotations } : {}), + }, + spec: { + ingressClassName: lifecycleDefaults?.ingressClassName || 'nginx', + rules: [ + { + ...(publication.host ? { host: publication.host } : {}), + http: { + paths: [ + { + ...pathRule, + backend: { + service: { + name: serviceName, + port: { + number: 80, + }, + }, + }, + }, + ], + }, + }, + ], + }, + }; + + await upsertService(coreApi, namespace, service); + await upsertIngress(networkingApi, namespace, ingress); + + return { + ...publication, + serviceName, + ingressName, + port, + }; +} diff --git a/src/server/lib/agentSession/configSeeder.ts b/src/server/lib/agentSession/configSeeder.ts index 573528c3..6e8d1510 100644 --- a/src/server/lib/agentSession/configSeeder.ts +++ b/src/server/lib/agentSession/configSeeder.ts @@ -15,7 +15,7 @@ */ import { posix as pathPosix } from 'path'; -import type { AgentSessionWorkspaceRepo } from './workspace'; +import { SESSION_WORKSPACE_ROOT, type AgentSessionWorkspaceRepo } from './workspace'; export const SESSION_WORKSPACE_SHARED_HOME_DIR = '/home/agent/.lifecycle-session'; export const SESSION_WORKSPACE_HOME_VOLUME_NAME = 'session-home'; @@ -69,8 +69,8 @@ function resolveInitWorkspaceRepos(opts: InitScriptOpts): AgentSessionWorkspaceR return opts.workspaceRepos; } - if (!opts.repoUrl || !opts.branch || !opts.workspacePath) { - throw new Error('repoUrl, branch, and workspacePath are required when workspaceRepos is not provided'); + if (!opts.repoUrl || !opts.branch) { + return []; } return [ @@ -161,12 +161,15 @@ export function generateInitScript(opts: InitScriptOpts): string { const { installCommand } = opts; const workspaceRepos = resolveInitWorkspaceRepos(opts); const primaryRepo = workspaceRepos.find((repo) => repo.primary) || workspaceRepos[0]; + const workspaceRoot = opts.workspacePath || SESSION_WORKSPACE_ROOT; const lines = ['#!/bin/sh', 'set -e']; appendGitIdentityAndAuthLines(lines, opts); - if (workspaceRepos.length === 1) { + if (workspaceRepos.length === 0) { + lines.push('', `mkdir -p "${escapeDoubleQuotedShell(workspaceRoot)}"`); + } else if (workspaceRepos.length === 1) { lines.push('', ...buildRepoCloneLines(workspaceRepos[0])); } else { lines.push('', 'clone_pids=""'); @@ -182,7 +185,7 @@ export function generateInitScript(opts: InitScriptOpts): string { lines.push('for clone_pid in $clone_pids; do', ' wait "$clone_pid"', 'done'); } - if (installCommand) { + if (installCommand && primaryRepo) { lines.push('', `cd "${escapeDoubleQuotedShell(primaryRepo.mountPath)}"`, installCommand); } diff --git a/src/server/lib/agentSession/pvcFactory.ts b/src/server/lib/agentSession/pvcFactory.ts index b9a7add4..8d0968ed 100644 --- a/src/server/lib/agentSession/pvcFactory.ts +++ b/src/server/lib/agentSession/pvcFactory.ts @@ -17,6 +17,11 @@ import * as k8s from '@kubernetes/client-node'; import { getLogger } from 'server/lib/logger'; import { buildLifecycleLabels } from 'server/lib/kubernetes/labels'; +import { + DEFAULT_AGENT_SESSION_WORKSPACE_STORAGE_ACCESS_MODE, + DEFAULT_AGENT_SESSION_WORKSPACE_STORAGE_SIZE, +} from './runtimeConfig'; +import type { AgentSessionWorkspaceStorageAccessMode } from 'server/services/types/globalConfig'; function getCoreApi(): k8s.CoreV1Api { const kc = new k8s.KubeConfig(); @@ -24,24 +29,15 @@ function getCoreApi(): k8s.CoreV1Api { return kc.makeApiClient(k8s.CoreV1Api); } -function getAccessMode(): 'ReadWriteMany' | 'ReadWriteOnce' { - const configured = process.env.AGENT_SESSION_PVC_ACCESS_MODE; - if (configured === 'ReadWriteMany' || configured === 'ReadWriteOnce') { - return configured; - } - - return 'ReadWriteOnce'; -} - export async function createAgentPvc( namespace: string, pvcName: string, - storageSize = '10Gi', - buildUuid?: string + storageSize = DEFAULT_AGENT_SESSION_WORKSPACE_STORAGE_SIZE, + buildUuid?: string, + accessMode: AgentSessionWorkspaceStorageAccessMode = DEFAULT_AGENT_SESSION_WORKSPACE_STORAGE_ACCESS_MODE ): Promise { const logger = getLogger(); const coreApi = getCoreApi(); - const accessMode = getAccessMode(); const pvc: k8s.V1PersistentVolumeClaim = { apiVersion: 'v1', @@ -64,11 +60,21 @@ export async function createAgentPvc( }, }; - const { body: result } = await coreApi.createNamespacedPersistentVolumeClaim(namespace, pvc); - logger.info( - `AgentRuntime: workspace prepared pvcName=${pvcName} namespace=${namespace} size=${storageSize} accessMode=${accessMode}` - ); - return result; + try { + const { body: result } = await coreApi.createNamespacedPersistentVolumeClaim(namespace, pvc); + logger.info( + `AgentRuntime: workspace prepared pvcName=${pvcName} namespace=${namespace} size=${storageSize} accessMode=${accessMode}` + ); + return result; + } catch (error: any) { + if (error instanceof k8s.HttpError && error.response?.statusCode === 409) { + const { body: existing } = await coreApi.readNamespacedPersistentVolumeClaim(pvcName, namespace); + logger.info(`AgentRuntime: workspace prepared reason=exists pvcName=${pvcName} namespace=${namespace}`); + return existing; + } + + throw error; + } } export async function deleteAgentPvc(namespace: string, pvcName: string): Promise { diff --git a/src/server/lib/agentSession/runtimeConfig.ts b/src/server/lib/agentSession/runtimeConfig.ts index b34b4624..06404980 100644 --- a/src/server/lib/agentSession/runtimeConfig.ts +++ b/src/server/lib/agentSession/runtimeConfig.ts @@ -17,10 +17,14 @@ import GlobalConfigService from 'server/services/globalConfig'; import type { AgentSessionControlPlaneConfig, + AgentSessionCleanupConfig, AgentSessionDefaults, + AgentSessionDurabilityConfig, AgentSessionReadinessConfig, AgentSessionResourcesConfig, AgentSessionSchedulingConfig, + AgentSessionWorkspaceStorageAccessMode, + AgentSessionWorkspaceStorageConfig, ResourceRequirements, } from 'server/services/types/globalConfig'; @@ -32,6 +36,9 @@ export interface AgentSessionRuntimeConfig { keepAttachedServicesOnSessionNode: boolean; readiness: ResolvedAgentSessionReadinessConfig; resources: ResolvedAgentSessionResources; + workspaceStorage: ResolvedAgentSessionWorkspaceStorageConfig; + cleanup: ResolvedAgentSessionCleanupConfig; + durability: ResolvedAgentSessionDurabilityConfig; } export interface ResolvedAgentSessionReadinessConfig { @@ -50,6 +57,36 @@ export interface ResolvedAgentSessionResources { workspaceGateway: ResolvedAgentSessionResourceRequirements; } +export interface ResolvedAgentSessionWorkspaceStorageConfig { + defaultSize: string; + allowedSizes: string[]; + allowClientOverride: boolean; + accessMode: AgentSessionWorkspaceStorageAccessMode; +} + +export interface ResolvedAgentSessionWorkspaceStorageIntent { + requestedSize: string | null; + storageSize: string; + accessMode: AgentSessionWorkspaceStorageAccessMode; +} + +export interface ResolvedAgentSessionCleanupConfig { + activeIdleSuspendMs: number; + startingTimeoutMs: number; + hibernatedRetentionMs: number; + intervalMs: number; + redisTtlSeconds: number; +} + +export interface ResolvedAgentSessionDurabilityConfig { + runExecutionLeaseMs: number; + queuedRunDispatchStaleMs: number; + dispatchRecoveryLimit: number; + maxDurablePayloadBytes: number; + payloadPreviewBytes: number; + fileChangePreviewChars: number; +} + export interface ResolvedAgentSessionControlPlaneConfig { systemPrompt?: string; appendSystemPrompt?: string; @@ -66,8 +103,10 @@ export const DEFAULT_AGENT_SESSION_CONTROL_PLANE_SYSTEM_PROMPT = [ 'If a tool call fails or a capability is unavailable, say that plainly and explain what failed.', ].join('\n'); -export const DEFAULT_AGENT_SESSION_CONTROL_PLANE_APPEND_SYSTEM_PROMPT = - 'When a tool execution is not approved, do not retry the denied action. Use the denial reason as updated guidance and continue from there.'; +export const DEFAULT_AGENT_SESSION_CONTROL_PLANE_APPEND_SYSTEM_PROMPT = [ + 'When a tool execution is not approved, do not retry the denied action. Use the denial reason as updated guidance and continue from there.', + 'When showing multi-line exact text such as file contents, command output, diffs, or JSON, use a fenced code block instead of inline code.', +].join('\n'); export const DEFAULT_AGENT_SESSION_MAX_ITERATIONS = 8; export const DEFAULT_AGENT_SESSION_WORKSPACE_TOOL_DISCOVERY_TIMEOUT_MS = 3000; export const DEFAULT_AGENT_SESSION_WORKSPACE_TOOL_EXECUTION_TIMEOUT_MS = 15000; @@ -75,6 +114,20 @@ export const DEFAULT_AGENT_SESSION_KEEP_ATTACHED_SERVICES_ON_SESSION_NODE = true const DEFAULT_AGENT_READY_TIMEOUT_MS = 60000; const DEFAULT_AGENT_READY_POLL_MS = 1000; +export const DEFAULT_AGENT_SESSION_WORKSPACE_STORAGE_SIZE = '10Gi'; +export const DEFAULT_AGENT_SESSION_WORKSPACE_STORAGE_ACCESS_MODE: AgentSessionWorkspaceStorageAccessMode = + 'ReadWriteOnce'; +export const DEFAULT_AGENT_SESSION_ACTIVE_IDLE_SUSPEND_MS = 30 * 60 * 1000; +export const DEFAULT_AGENT_SESSION_STARTING_TIMEOUT_MS = 15 * 60 * 1000; +export const DEFAULT_AGENT_SESSION_HIBERNATED_RETENTION_MS = 24 * 60 * 60 * 1000; +export const DEFAULT_AGENT_SESSION_CLEANUP_INTERVAL_MS = 5 * 60 * 1000; +export const DEFAULT_AGENT_SESSION_REDIS_TTL_SECONDS = 7200; +export const DEFAULT_AGENT_SESSION_RUN_EXECUTION_LEASE_MS = 30 * 60 * 1000; +export const DEFAULT_AGENT_SESSION_QUEUED_RUN_DISPATCH_STALE_MS = 30 * 1000; +export const DEFAULT_AGENT_SESSION_DISPATCH_RECOVERY_LIMIT = 50; +export const DEFAULT_AGENT_SESSION_MAX_DURABLE_PAYLOAD_BYTES = 64 * 1024; +export const DEFAULT_AGENT_SESSION_PAYLOAD_PREVIEW_BYTES = 16 * 1024; +export const DEFAULT_AGENT_SESSION_FILE_CHANGE_PREVIEW_CHARS = 4000; const DEFAULT_WORKSPACE_RESOURCES: ResolvedAgentSessionResourceRequirements = { requests: { cpu: '500m', @@ -153,6 +206,20 @@ function normalizeBoolean(value: unknown): boolean | undefined { return undefined; } +function normalizeStringArray(value: unknown): string[] { + if (!Array.isArray(value)) { + return []; + } + + return Array.from( + new Set(value.filter((entry): entry is string => typeof entry === 'string').map((entry) => entry.trim())) + ).filter(Boolean); +} + +function normalizeAccessMode(value: unknown): AgentSessionWorkspaceStorageAccessMode | undefined { + return value === 'ReadWriteMany' || value === 'ReadWriteOnce' ? value : undefined; +} + function normalizeResourceQuantityMap(values: unknown): Record { if (!values || typeof values !== 'object' || Array.isArray(values)) { return {}; @@ -276,6 +343,107 @@ export function mergeAgentSessionResources( }; } +export function resolveAgentSessionWorkspaceStorageFromDefaults( + storageDefaults?: AgentSessionWorkspaceStorageConfig | null +): ResolvedAgentSessionWorkspaceStorageConfig { + const defaultSize = + normalizeOptionalString(storageDefaults?.defaultSize) || DEFAULT_AGENT_SESSION_WORKSPACE_STORAGE_SIZE; + const configuredAllowedSizes = normalizeStringArray(storageDefaults?.allowedSizes); + const allowedSizes = configuredAllowedSizes.includes(defaultSize) + ? configuredAllowedSizes + : [defaultSize, ...configuredAllowedSizes]; + + return { + defaultSize, + allowedSizes, + allowClientOverride: normalizeBoolean(storageDefaults?.allowClientOverride) ?? false, + accessMode: + normalizeAccessMode(storageDefaults?.accessMode) || + normalizeAccessMode(process.env.AGENT_SESSION_PVC_ACCESS_MODE) || + DEFAULT_AGENT_SESSION_WORKSPACE_STORAGE_ACCESS_MODE, + }; +} + +export class AgentSessionWorkspaceStorageConfigError extends Error { + constructor(message: string) { + super(message); + this.name = 'AgentSessionWorkspaceStorageConfigError'; + } +} + +export function resolveAgentSessionWorkspaceStorageIntent({ + requestedSize, + storage, +}: { + requestedSize?: string | null; + storage: ResolvedAgentSessionWorkspaceStorageConfig; +}): ResolvedAgentSessionWorkspaceStorageIntent { + const normalizedRequestedSize = normalizeOptionalString(requestedSize) || null; + + if (!normalizedRequestedSize) { + return { + requestedSize: null, + storageSize: storage.defaultSize, + accessMode: storage.accessMode, + }; + } + + if (!storage.allowClientOverride) { + throw new AgentSessionWorkspaceStorageConfigError('workspace.storageSize overrides are not enabled.'); + } + + if (!storage.allowedSizes.includes(normalizedRequestedSize)) { + throw new AgentSessionWorkspaceStorageConfigError( + `workspace.storageSize must be one of: ${storage.allowedSizes.join(', ')}.` + ); + } + + return { + requestedSize: normalizedRequestedSize, + storageSize: normalizedRequestedSize, + accessMode: storage.accessMode, + }; +} + +export function resolveAgentSessionCleanupFromDefaults( + cleanupDefaults?: AgentSessionCleanupConfig | null +): ResolvedAgentSessionCleanupConfig { + return { + activeIdleSuspendMs: + normalizePositiveInteger(cleanupDefaults?.activeIdleSuspendMs) ?? DEFAULT_AGENT_SESSION_ACTIVE_IDLE_SUSPEND_MS, + startingTimeoutMs: + normalizePositiveInteger(cleanupDefaults?.startingTimeoutMs) ?? DEFAULT_AGENT_SESSION_STARTING_TIMEOUT_MS, + hibernatedRetentionMs: + normalizePositiveInteger(cleanupDefaults?.hibernatedRetentionMs) ?? DEFAULT_AGENT_SESSION_HIBERNATED_RETENTION_MS, + intervalMs: normalizePositiveInteger(cleanupDefaults?.intervalMs) ?? DEFAULT_AGENT_SESSION_CLEANUP_INTERVAL_MS, + redisTtlSeconds: + normalizePositiveInteger(cleanupDefaults?.redisTtlSeconds) ?? DEFAULT_AGENT_SESSION_REDIS_TTL_SECONDS, + }; +} + +export function resolveAgentSessionDurabilityFromDefaults( + durabilityDefaults?: AgentSessionDurabilityConfig | null +): ResolvedAgentSessionDurabilityConfig { + return { + runExecutionLeaseMs: + normalizePositiveInteger(durabilityDefaults?.runExecutionLeaseMs) ?? DEFAULT_AGENT_SESSION_RUN_EXECUTION_LEASE_MS, + queuedRunDispatchStaleMs: + normalizePositiveInteger(durabilityDefaults?.queuedRunDispatchStaleMs) ?? + DEFAULT_AGENT_SESSION_QUEUED_RUN_DISPATCH_STALE_MS, + dispatchRecoveryLimit: + normalizePositiveInteger(durabilityDefaults?.dispatchRecoveryLimit) ?? + DEFAULT_AGENT_SESSION_DISPATCH_RECOVERY_LIMIT, + maxDurablePayloadBytes: + normalizePositiveInteger(durabilityDefaults?.maxDurablePayloadBytes) ?? + DEFAULT_AGENT_SESSION_MAX_DURABLE_PAYLOAD_BYTES, + payloadPreviewBytes: + normalizePositiveInteger(durabilityDefaults?.payloadPreviewBytes) ?? DEFAULT_AGENT_SESSION_PAYLOAD_PREVIEW_BYTES, + fileChangePreviewChars: + normalizePositiveInteger(durabilityDefaults?.fileChangePreviewChars) ?? + DEFAULT_AGENT_SESSION_FILE_CHANGE_PREVIEW_CHARS, + }; +} + export function resolveAgentSessionControlPlaneConfigFromDefaults( agentSessionDefaults?: AgentSessionDefaults | null ): ResolvedAgentSessionControlPlaneConfig { @@ -308,6 +476,16 @@ export async function resolveAgentSessionControlPlaneConfig(): Promise { + const { agentSessionDefaults } = await GlobalConfigService.getInstance().getAllConfigs(); + return resolveAgentSessionCleanupFromDefaults(agentSessionDefaults?.cleanup); +} + +export async function resolveAgentSessionDurabilityConfig(): Promise { + const { agentSessionDefaults } = await GlobalConfigService.getInstance().getAllConfigs(); + return resolveAgentSessionDurabilityFromDefaults(agentSessionDefaults?.durability); +} + export class AgentSessionRuntimeConfigError extends Error { readonly missingFields: Array<'workspaceImage' | 'workspaceEditorImage'>; @@ -345,5 +523,8 @@ export async function resolveAgentSessionRuntimeConfig(): Promise { + it('returns a standard error response for application errors', async () => { + const handler = createApiHandler(async () => { + throw new Error('sample failure'); + }); + + const response = await handler(new NextRequest('http://localhost/api/v2/sample')); + + await expect(response.json()).resolves.toMatchObject({ + data: null, + error: { + message: 'sample failure', + }, + }); + expect(response.status).toBe(500); + }); + + it('rethrows Next dynamic usage errors so Next can mark the route dynamic', async () => { + const dynamicError = new DynamicServerError('Route /api/v2/sample used request headers.'); + const handler = createApiHandler(async () => { + throw dynamicError; + }); + + await expect(handler(new NextRequest('http://localhost/api/v2/sample'))).rejects.toBe(dynamicError); + }); + + it('returns successful responses unchanged', async () => { + const handler = createApiHandler(async () => NextResponse.json({ ok: true }, { status: 201 })); + + const response = await handler(new NextRequest('http://localhost/api/v2/sample')); + + await expect(response.json()).resolves.toEqual({ ok: true }); + expect(response.status).toBe(201); + }); +}); diff --git a/src/server/lib/createApiHandler.ts b/src/server/lib/createApiHandler.ts index 8b70268a..3f9ab79f 100644 --- a/src/server/lib/createApiHandler.ts +++ b/src/server/lib/createApiHandler.ts @@ -15,6 +15,8 @@ */ import { NextRequest, NextResponse } from 'next/server'; +import { unstable_noStore as noStore } from 'next/cache'; +import { isDynamicUsageError } from 'next/dist/export/helpers/is-dynamic-usage-error'; import { errorResponse } from './response'; import { requireRole, type LifecycleRole } from './roles'; @@ -48,9 +50,15 @@ export function createApiHandler(handler: RouteHandler, options?: ApiHandlerOpti } return async (req: NextRequest, ...args: any[]) => { + noStore(); + try { return await wrapped(req, ...args); } catch (error) { + if (isDynamicUsageError(error)) { + throw error; + } + return errorResponse(error, { status: 500 }, req); } }; diff --git a/src/server/lib/createStreamHandler.test.ts b/src/server/lib/createStreamHandler.test.ts new file mode 100644 index 00000000..a6b34f97 --- /dev/null +++ b/src/server/lib/createStreamHandler.test.ts @@ -0,0 +1,42 @@ +/** + * Copyright 2026 GoodRx, Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +import { DynamicServerError } from 'next/dist/client/components/hooks-server-context'; +import { NextRequest } from 'next/server'; +import { createStreamHandler } from './createStreamHandler'; + +describe('createStreamHandler', () => { + it('returns an SSE error frame for application errors', async () => { + const handler = createStreamHandler(async () => { + throw new Error('sample stream failure'); + }); + + const response = await handler(new NextRequest('http://localhost/api/v2/sample/stream')); + + await expect(response.text()).resolves.toContain('"error":"sample stream failure"'); + expect(response.status).toBe(200); + expect(response.headers.get('content-type')).toBe('text/event-stream; charset=utf-8'); + }); + + it('rethrows Next dynamic usage errors so Next can mark the stream route dynamic', async () => { + const dynamicError = new DynamicServerError('Route /api/v2/sample/stream used request headers.'); + const handler = createStreamHandler(async () => { + throw dynamicError; + }); + + await expect(handler(new NextRequest('http://localhost/api/v2/sample/stream'))).rejects.toBe(dynamicError); + }); +}); diff --git a/src/server/lib/createStreamHandler.ts b/src/server/lib/createStreamHandler.ts index 19780c30..2fc821e2 100644 --- a/src/server/lib/createStreamHandler.ts +++ b/src/server/lib/createStreamHandler.ts @@ -15,6 +15,8 @@ */ import { NextRequest } from 'next/server'; +import { unstable_noStore as noStore } from 'next/cache'; +import { isDynamicUsageError } from 'next/dist/export/helpers/is-dynamic-usage-error'; import { getLogger } from 'server/lib/logger'; // eslint-disable-next-line no-unused-vars @@ -22,9 +24,15 @@ type StreamRouteHandler = (req: NextRequest, ...args: any[]) => Promise { + noStore(); + try { return await handler(req, ...args); } catch (error) { + if (isDynamicUsageError(error)) { + throw error; + } + let errorMessage = 'An unexpected error occurred.'; let errorStack = ''; diff --git a/src/server/lib/kubernetes.ts b/src/server/lib/kubernetes.ts index d2dc0700..805f0dc1 100644 --- a/src/server/lib/kubernetes.ts +++ b/src/server/lib/kubernetes.ts @@ -592,7 +592,9 @@ export async function deleteBuild(build: Build) { * @param name namespace to delete */ export async function deleteNamespace(name: string) { - if (!name.startsWith('env-') && !name.startsWith('sbx-')) return; + if (!name.startsWith('env-') && !name.startsWith('sbx-') && !name.startsWith('prj-') && !name.startsWith('chat-')) { + return; + } try { await shellPromise(`kubectl delete ns ${name} --grace-period 120`); diff --git a/src/server/lib/response.ts b/src/server/lib/response.ts index 7a22441e..a8c7ed73 100644 --- a/src/server/lib/response.ts +++ b/src/server/lib/response.ts @@ -20,6 +20,8 @@ import { getLogger } from 'server/lib/logger'; interface Metadata { pagination?: PaginationMetadata; + limit?: number; + maxLimit?: number; } type SuccessStatusCode = 200 | 201; diff --git a/src/server/lib/validation/agentSessionConfigSchemas.ts b/src/server/lib/validation/agentSessionConfigSchemas.ts index 6ae885c0..ea6a12bf 100644 --- a/src/server/lib/validation/agentSessionConfigSchemas.ts +++ b/src/server/lib/validation/agentSessionConfigSchemas.ts @@ -54,6 +54,12 @@ const resourceRequirementsSchema = { additionalProperties: false, }; +const workspaceStorageSizeSchema = { + type: 'string', + minLength: 1, + maxLength: 64, +}; + export const agentSessionControlPlaneConfigSchema = { type: 'object', properties: { @@ -101,6 +107,43 @@ export const agentSessionRuntimeSettingsSchema = { }, additionalProperties: false, }, + workspaceStorage: { + type: 'object', + properties: { + defaultSize: workspaceStorageSizeSchema, + allowedSizes: { + type: 'array', + items: workspaceStorageSizeSchema, + uniqueItems: true, + }, + allowClientOverride: { type: 'boolean' }, + accessMode: { type: 'string', enum: ['ReadWriteOnce', 'ReadWriteMany'] }, + }, + additionalProperties: false, + }, + cleanup: { + type: 'object', + properties: { + activeIdleSuspendMs: positiveIntegerSchema, + startingTimeoutMs: positiveIntegerSchema, + hibernatedRetentionMs: positiveIntegerSchema, + intervalMs: positiveIntegerSchema, + redisTtlSeconds: positiveIntegerSchema, + }, + additionalProperties: false, + }, + durability: { + type: 'object', + properties: { + runExecutionLeaseMs: positiveIntegerSchema, + queuedRunDispatchStaleMs: positiveIntegerSchema, + dispatchRecoveryLimit: positiveIntegerSchema, + maxDurablePayloadBytes: positiveIntegerSchema, + payloadPreviewBytes: positiveIntegerSchema, + fileChangePreviewChars: positiveIntegerSchema, + }, + additionalProperties: false, + }, }, additionalProperties: false, }; diff --git a/src/server/lib/validation/agentSessionConfigValidator.ts b/src/server/lib/validation/agentSessionConfigValidator.ts index 7ba4b6fb..3c349e38 100644 --- a/src/server/lib/validation/agentSessionConfigValidator.ts +++ b/src/server/lib/validation/agentSessionConfigValidator.ts @@ -84,6 +84,50 @@ function validateBooleanField(value: unknown, fieldName: string): void { } } +function validateOptionalStringField(value: unknown, fieldName: string, maxLength = 2048): void { + if (value === undefined) { + return; + } + + if (typeof value !== 'string' || !value.trim()) { + throw new AgentSessionConfigValidationError(`${fieldName} must be a non-empty string.`); + } + + if (value.length > maxLength) { + throw new AgentSessionConfigValidationError(`${fieldName} exceeds maximum length of ${maxLength} characters.`); + } +} + +function validateStringArrayField(value: unknown, fieldName: string): void { + if (value === undefined) { + return; + } + + if (!Array.isArray(value)) { + throw new AgentSessionConfigValidationError(`${fieldName} must be an array.`); + } + + const seen = new Set(); + for (const item of value) { + validateOptionalStringField(item, `${fieldName} entry`, 64); + const normalized = item.trim(); + if (seen.has(normalized)) { + throw new AgentSessionConfigValidationError(`${fieldName} contains duplicate value "${normalized}".`); + } + seen.add(normalized); + } +} + +function validateAccessModeField(value: unknown, fieldName: string): void { + if (value === undefined) { + return; + } + + if (value !== 'ReadWriteOnce' && value !== 'ReadWriteMany') { + throw new AgentSessionConfigValidationError(`${fieldName} must be ReadWriteOnce or ReadWriteMany.`); + } +} + export function validateAgentSessionControlPlaneConfig(config: Partial): void { validatePromptField(config.systemPrompt, 'systemPrompt'); validatePromptField(config.appendSystemPrompt, 'appendSystemPrompt'); @@ -132,4 +176,19 @@ export function validateAgentSessionRuntimeSettings(config: AgentSessionRuntimeS validateStringRecord(config.resources?.editor?.limits, 'resources.editor.limits'); validateStringRecord(config.resources?.workspaceGateway?.requests, 'resources.workspaceGateway.requests'); validateStringRecord(config.resources?.workspaceGateway?.limits, 'resources.workspaceGateway.limits'); + validateOptionalStringField(config.workspaceStorage?.defaultSize, 'workspaceStorage.defaultSize', 64); + validateStringArrayField(config.workspaceStorage?.allowedSizes, 'workspaceStorage.allowedSizes'); + validateBooleanField(config.workspaceStorage?.allowClientOverride, 'workspaceStorage.allowClientOverride'); + validateAccessModeField(config.workspaceStorage?.accessMode, 'workspaceStorage.accessMode'); + validatePositiveIntegerField(config.cleanup?.activeIdleSuspendMs, 'cleanup.activeIdleSuspendMs'); + validatePositiveIntegerField(config.cleanup?.startingTimeoutMs, 'cleanup.startingTimeoutMs'); + validatePositiveIntegerField(config.cleanup?.hibernatedRetentionMs, 'cleanup.hibernatedRetentionMs'); + validatePositiveIntegerField(config.cleanup?.intervalMs, 'cleanup.intervalMs'); + validatePositiveIntegerField(config.cleanup?.redisTtlSeconds, 'cleanup.redisTtlSeconds'); + validatePositiveIntegerField(config.durability?.runExecutionLeaseMs, 'durability.runExecutionLeaseMs'); + validatePositiveIntegerField(config.durability?.queuedRunDispatchStaleMs, 'durability.queuedRunDispatchStaleMs'); + validatePositiveIntegerField(config.durability?.dispatchRecoveryLimit, 'durability.dispatchRecoveryLimit'); + validatePositiveIntegerField(config.durability?.maxDurablePayloadBytes, 'durability.maxDurablePayloadBytes'); + validatePositiveIntegerField(config.durability?.payloadPreviewBytes, 'durability.payloadPreviewBytes'); + validatePositiveIntegerField(config.durability?.fileChangePreviewChars, 'durability.fileChangePreviewChars'); } diff --git a/src/server/middlewares/cors.test.ts b/src/server/middlewares/cors.test.ts new file mode 100644 index 00000000..d8c1be54 --- /dev/null +++ b/src/server/middlewares/cors.test.ts @@ -0,0 +1,50 @@ +/** + * Copyright 2026 GoodRx, Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +import { NextRequest, NextResponse } from 'next/server'; + +describe('corsMiddleware', () => { + const originalAllowedOrigins = process.env.ALLOWED_ORIGINS; + let corsMiddleware: typeof import('./cors').corsMiddleware; + + beforeEach(async () => { + jest.resetModules(); + process.env.ALLOWED_ORIGINS = 'http://localhost:3000'; + ({ corsMiddleware } = await import('./cors')); + }); + + afterAll(() => { + process.env.ALLOWED_ORIGINS = originalAllowedOrigins; + }); + + it('allows the SSE resume header during preflight', async () => { + const next = jest.fn().mockResolvedValue(NextResponse.next()); + const request = new NextRequest('http://localhost/api/v2/ai/agent/runs/run-1/events/stream', { + headers: { + origin: 'http://localhost:3000', + 'access-control-request-headers': 'last-event-id', + }, + method: 'OPTIONS', + }); + + const response = await corsMiddleware(request, next); + + expect(response.status).toBe(204); + expect(next).not.toHaveBeenCalled(); + expect(response.headers.get('access-control-allow-origin')).toBe('http://localhost:3000'); + expect(response.headers.get('access-control-allow-headers')).toContain('Last-Event-ID'); + }); +}); diff --git a/src/server/middlewares/cors.ts b/src/server/middlewares/cors.ts index 7bf06918..937a6594 100644 --- a/src/server/middlewares/cors.ts +++ b/src/server/middlewares/cors.ts @@ -18,6 +18,7 @@ import { NextResponse } from 'next/server'; import type { Middleware } from './chain'; const allowedOrigins = process.env.ALLOWED_ORIGINS?.split(',').map((o) => o.trim()) ?? []; +const allowedHeaders = ['Content-Type', 'Authorization', 'Last-Event-ID']; export const corsMiddleware: Middleware = async (request, next) => { const origin = request.headers.get('origin'); @@ -27,7 +28,7 @@ export const corsMiddleware: Middleware = async (request, next) => { 'Access-Control-Allow-Origin': isAllowedOrigin ? origin : allowedOrigins[0], 'Access-Control-Allow-Credentials': 'true', 'Access-Control-Allow-Methods': 'GET, POST, PUT, PATCH, DELETE, OPTIONS', - 'Access-Control-Allow-Headers': 'Content-Type, Authorization', + 'Access-Control-Allow-Headers': allowedHeaders.join(', '), }; if (request.method === 'OPTIONS') { diff --git a/src/server/models/AgentMessage.ts b/src/server/models/AgentMessage.ts index 74deb850..5dc33dcd 100644 --- a/src/server/models/AgentMessage.ts +++ b/src/server/models/AgentMessage.ts @@ -17,10 +17,13 @@ import Model from './_Model'; export default class AgentMessage extends Model { + uuid!: string; + clientMessageId!: string | null; threadId!: number; runId!: number | null; role!: 'user' | 'assistant' | 'system' | 'tool'; - uiMessage!: Record; + parts!: Array>; + uiMessage!: Record | null; metadata!: Record; static tableName = 'agent_messages'; @@ -29,19 +32,25 @@ export default class AgentMessage extends Model { static jsonSchema = { type: 'object', - required: ['threadId', 'role', 'uiMessage'], + required: ['threadId', 'role'], properties: { id: { type: 'integer' }, + uuid: { + type: 'string', + pattern: '^[0-9a-fA-F-]{36}$', + }, threadId: { type: 'integer' }, + clientMessageId: { type: ['string', 'null'] }, runId: { type: ['integer', 'null'] }, role: { type: 'string', enum: ['user', 'assistant', 'system', 'tool'] }, - uiMessage: { type: 'object' }, + parts: { type: 'array', items: { type: 'object' }, default: [] }, + uiMessage: { type: ['object', 'null'] }, metadata: { type: 'object', default: {} }, }, }; static get jsonAttributes() { - return ['uiMessage', 'metadata']; + return ['parts', 'uiMessage', 'metadata']; } static get relationMappings() { diff --git a/src/server/models/AgentRun.ts b/src/server/models/AgentRun.ts index c153ecca..2d417be1 100644 --- a/src/server/models/AgentRun.ts +++ b/src/server/models/AgentRun.ts @@ -20,7 +20,26 @@ export default class AgentRun extends Model { uuid!: string; threadId!: number; sessionId!: number; - status!: 'queued' | 'running' | 'waiting_for_approval' | 'waiting_for_input' | 'completed' | 'failed' | 'cancelled'; + status!: + | 'queued' + | 'starting' + | 'running' + | 'waiting_for_approval' + | 'waiting_for_input' + | 'completed' + | 'failed' + | 'cancelled'; + requestedHarness!: string | null; + resolvedHarness!: string | null; + requestedProvider!: string | null; + requestedModel!: string | null; + resolvedProvider!: string | null; + resolvedModel!: string | null; + sandboxRequirement!: Record; + sandboxGeneration!: number | null; + executionOwner!: string | null; + leaseExpiresAt!: string | null; + heartbeatAt!: string | null; provider!: string; model!: string; queuedAt!: string; @@ -29,7 +48,6 @@ export default class AgentRun extends Model { cancelledAt!: string | null; usageSummary!: Record; policySnapshot!: Record; - streamState!: Record; error!: Record | null; static tableName = 'agent_runs'; @@ -49,8 +67,28 @@ export default class AgentRun extends Model { sessionId: { type: 'integer' }, status: { type: 'string', - enum: ['queued', 'running', 'waiting_for_approval', 'waiting_for_input', 'completed', 'failed', 'cancelled'], + enum: [ + 'queued', + 'starting', + 'running', + 'waiting_for_approval', + 'waiting_for_input', + 'completed', + 'failed', + 'cancelled', + ], }, + requestedHarness: { type: ['string', 'null'] }, + resolvedHarness: { type: ['string', 'null'] }, + requestedProvider: { type: ['string', 'null'] }, + requestedModel: { type: ['string', 'null'] }, + resolvedProvider: { type: ['string', 'null'] }, + resolvedModel: { type: ['string', 'null'] }, + sandboxRequirement: { type: 'object', default: {} }, + sandboxGeneration: { type: ['integer', 'null'] }, + executionOwner: { type: ['string', 'null'] }, + leaseExpiresAt: { type: ['string', 'null'] }, + heartbeatAt: { type: ['string', 'null'] }, provider: { type: 'string' }, model: { type: 'string' }, queuedAt: { type: 'string' }, @@ -59,13 +97,12 @@ export default class AgentRun extends Model { cancelledAt: { type: ['string', 'null'] }, usageSummary: { type: 'object', default: {} }, policySnapshot: { type: 'object', default: {} }, - streamState: { type: 'object', default: {} }, error: { type: ['object', 'null'], default: null }, }, }; static get jsonAttributes() { - return ['usageSummary', 'policySnapshot', 'streamState', 'error']; + return ['sandboxRequirement', 'usageSummary', 'policySnapshot', 'error']; } static get relationMappings() { @@ -74,6 +111,7 @@ export default class AgentRun extends Model { const AgentMessage = require('./AgentMessage').default; const AgentPendingAction = require('./AgentPendingAction').default; const AgentToolExecution = require('./AgentToolExecution').default; + const AgentRunEvent = require('./AgentRunEvent').default; return { thread: { @@ -116,6 +154,14 @@ export default class AgentRun extends Model { to: 'agent_tool_executions.runId', }, }, + events: { + relation: Model.HasManyRelation, + modelClass: AgentRunEvent, + join: { + from: 'agent_runs.id', + to: 'agent_run_events.runId', + }, + }, }; } } diff --git a/src/server/models/AgentRunEvent.ts b/src/server/models/AgentRunEvent.ts new file mode 100644 index 00000000..baa7d2dc --- /dev/null +++ b/src/server/models/AgentRunEvent.ts @@ -0,0 +1,64 @@ +/** + * Copyright 2026 GoodRx, Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +import Model from './_Model'; + +export default class AgentRunEvent extends Model { + uuid!: string; + runId!: number; + sequence!: number; + eventType!: string; + payload!: Record; + + static tableName = 'agent_run_events'; + static timestamps = true; + static idColumn = 'id'; + + static jsonSchema = { + type: 'object', + required: ['runId', 'sequence', 'eventType'], + properties: { + id: { type: 'integer' }, + uuid: { + type: 'string', + pattern: '^[0-9a-fA-F-]{36}$', + }, + runId: { type: 'integer' }, + sequence: { type: 'integer' }, + eventType: { type: 'string' }, + payload: { type: 'object', default: {} }, + }, + }; + + static get jsonAttributes() { + return ['payload']; + } + + static get relationMappings() { + const AgentRun = require('./AgentRun').default; + + return { + run: { + relation: Model.BelongsToOneRelation, + modelClass: AgentRun, + join: { + from: 'agent_run_events.runId', + to: 'agent_runs.id', + }, + }, + }; + } +} diff --git a/src/server/models/AgentSandbox.ts b/src/server/models/AgentSandbox.ts new file mode 100644 index 00000000..12cb9c1e --- /dev/null +++ b/src/server/models/AgentSandbox.ts @@ -0,0 +1,88 @@ +/** + * Copyright 2026 GoodRx, Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +import Model from './_Model'; + +export default class AgentSandbox extends Model { + uuid!: string; + sessionId!: number; + generation!: number; + provider!: string; + status!: 'provisioning' | 'ready' | 'suspending' | 'suspended' | 'resuming' | 'failed' | 'ended'; + capabilitySnapshot!: Record; + providerState!: Record; + metadata!: Record; + error!: Record | null; + suspendedAt!: string | null; + endedAt!: string | null; + + static tableName = 'agent_sandboxes'; + static timestamps = true; + static idColumn = 'id'; + + static jsonSchema = { + type: 'object', + required: ['sessionId', 'generation', 'provider', 'status'], + properties: { + id: { type: 'integer' }, + uuid: { + type: 'string', + pattern: '^[0-9a-fA-F-]{36}$', + }, + sessionId: { type: 'integer' }, + generation: { type: 'integer' }, + provider: { type: 'string' }, + status: { + type: 'string', + enum: ['provisioning', 'ready', 'suspending', 'suspended', 'resuming', 'failed', 'ended'], + }, + capabilitySnapshot: { type: 'object', default: {} }, + providerState: { type: 'object', default: {} }, + metadata: { type: 'object', default: {} }, + error: { type: ['object', 'null'], default: null }, + suspendedAt: { type: ['string', 'null'] }, + endedAt: { type: ['string', 'null'] }, + }, + }; + + static get jsonAttributes() { + return ['capabilitySnapshot', 'providerState', 'metadata', 'error']; + } + + static get relationMappings() { + const AgentSession = require('./AgentSession').default; + const AgentSandboxExposure = require('./AgentSandboxExposure').default; + + return { + session: { + relation: Model.BelongsToOneRelation, + modelClass: AgentSession, + join: { + from: 'agent_sandboxes.sessionId', + to: 'agent_sessions.id', + }, + }, + exposures: { + relation: Model.HasManyRelation, + modelClass: AgentSandboxExposure, + join: { + from: 'agent_sandboxes.id', + to: 'agent_sandbox_exposures.sandboxId', + }, + }, + }; + } +} diff --git a/src/server/models/AgentSandboxExposure.ts b/src/server/models/AgentSandboxExposure.ts new file mode 100644 index 00000000..de7a9737 --- /dev/null +++ b/src/server/models/AgentSandboxExposure.ts @@ -0,0 +1,74 @@ +/** + * Copyright 2026 GoodRx, Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +import Model from './_Model'; + +export default class AgentSandboxExposure extends Model { + uuid!: string; + sandboxId!: number; + kind!: string; + status!: 'provisioning' | 'ready' | 'failed' | 'ended'; + targetPort!: number | null; + url!: string | null; + metadata!: Record; + providerState!: Record; + lastVerifiedAt!: string | null; + endedAt!: string | null; + + static tableName = 'agent_sandbox_exposures'; + static timestamps = true; + static idColumn = 'id'; + + static jsonSchema = { + type: 'object', + required: ['sandboxId', 'kind', 'status'], + properties: { + id: { type: 'integer' }, + uuid: { + type: 'string', + pattern: '^[0-9a-fA-F-]{36}$', + }, + sandboxId: { type: 'integer' }, + kind: { type: 'string' }, + status: { type: 'string', enum: ['provisioning', 'ready', 'failed', 'ended'] }, + targetPort: { type: ['integer', 'null'] }, + url: { type: ['string', 'null'] }, + metadata: { type: 'object', default: {} }, + providerState: { type: 'object', default: {} }, + lastVerifiedAt: { type: ['string', 'null'] }, + endedAt: { type: ['string', 'null'] }, + }, + }; + + static get jsonAttributes() { + return ['metadata', 'providerState']; + } + + static get relationMappings() { + const AgentSandbox = require('./AgentSandbox').default; + + return { + sandbox: { + relation: Model.BelongsToOneRelation, + modelClass: AgentSandbox, + join: { + from: 'agent_sandbox_exposures.sandboxId', + to: 'agent_sandboxes.id', + }, + }, + }; + } +} diff --git a/src/server/models/AgentSession.ts b/src/server/models/AgentSession.ts index dc127954..ce962dad 100644 --- a/src/server/models/AgentSession.ts +++ b/src/server/models/AgentSession.ts @@ -19,19 +19,25 @@ import type { DevModeResourceSnapshot } from 'server/lib/agentSession/devModeMan import type { AgentSessionSkillPlan } from 'server/lib/agentSession/skillPlan'; import { EMPTY_AGENT_SESSION_SKILL_PLAN } from 'server/lib/agentSession/skillPlan'; import type { AgentSessionSelectedService, AgentSessionWorkspaceRepo } from 'server/lib/agentSession/workspace'; -import { BuildKind } from 'shared/constants'; +import { AgentChatStatus, AgentSessionKind, AgentWorkspaceStatus, BuildKind } from 'shared/constants'; export default class AgentSession extends Model { uuid!: string; + defaultThreadId!: number | null; + defaultModel!: string; + defaultHarness!: string | null; buildUuid!: string | null; - buildKind!: BuildKind; + buildKind!: BuildKind | null; + sessionKind!: AgentSessionKind; userId!: string; ownerGithubUsername!: string | null; - podName!: string; - namespace!: string; - pvcName!: string; + podName!: string | null; + namespace!: string | null; + pvcName!: string | null; model!: string; status!: 'starting' | 'active' | 'ended' | 'error'; + chatStatus!: AgentChatStatus; + workspaceStatus!: AgentWorkspaceStatus; keepAttachedServicesOnSessionNode!: boolean | null; lastActivity!: string; endedAt!: string | null; @@ -47,22 +53,36 @@ export default class AgentSession extends Model { static jsonSchema = { type: 'object', - required: ['userId', 'podName', 'namespace', 'pvcName', 'model'], + required: ['userId', 'model', 'sessionKind', 'chatStatus', 'workspaceStatus'], properties: { id: { type: 'integer' }, uuid: { type: 'string', pattern: '^[0-9a-fA-F]{8}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{12}$', }, + defaultThreadId: { type: ['integer', 'null'] }, + defaultModel: { type: 'string' }, + defaultHarness: { type: ['string', 'null'] }, buildUuid: { type: ['string', 'null'] }, - buildKind: { type: 'string', enum: Object.values(BuildKind), default: BuildKind.ENVIRONMENT }, + buildKind: { + type: ['string', 'null'], + enum: [...Object.values(BuildKind), null], + default: BuildKind.ENVIRONMENT, + }, + sessionKind: { type: 'string', enum: Object.values(AgentSessionKind), default: AgentSessionKind.ENVIRONMENT }, userId: { type: 'string' }, ownerGithubUsername: { type: ['string', 'null'] }, - podName: { type: 'string' }, - namespace: { type: 'string' }, - pvcName: { type: 'string' }, + podName: { type: ['string', 'null'] }, + namespace: { type: ['string', 'null'] }, + pvcName: { type: ['string', 'null'] }, model: { type: 'string' }, status: { type: 'string', enum: ['starting', 'active', 'ended', 'error'], default: 'starting' }, + chatStatus: { type: 'string', enum: Object.values(AgentChatStatus), default: AgentChatStatus.READY }, + workspaceStatus: { + type: 'string', + enum: Object.values(AgentWorkspaceStatus), + default: AgentWorkspaceStatus.READY, + }, keepAttachedServicesOnSessionNode: { type: ['boolean', 'null'] }, lastActivity: { type: 'string' }, endedAt: { type: ['string', 'null'] }, @@ -87,6 +107,9 @@ export default class AgentSession extends Model { static get relationMappings() { const Deploy = require('./Deploy').default; + const AgentThread = require('./AgentThread').default; + const AgentSource = require('./AgentSource').default; + const AgentSandbox = require('./AgentSandbox').default; return { deploys: { relation: Model.HasManyRelation, @@ -96,6 +119,30 @@ export default class AgentSession extends Model { to: 'deploys.devModeSessionId', }, }, + defaultThread: { + relation: Model.BelongsToOneRelation, + modelClass: AgentThread, + join: { + from: 'agent_sessions.defaultThreadId', + to: 'agent_threads.id', + }, + }, + source: { + relation: Model.HasOneRelation, + modelClass: AgentSource, + join: { + from: 'agent_sessions.id', + to: 'agent_sources.sessionId', + }, + }, + sandboxes: { + relation: Model.HasManyRelation, + modelClass: AgentSandbox, + join: { + from: 'agent_sessions.id', + to: 'agent_sandboxes.sessionId', + }, + }, }; } } diff --git a/src/server/models/AgentSource.ts b/src/server/models/AgentSource.ts new file mode 100644 index 00000000..4b67ad59 --- /dev/null +++ b/src/server/models/AgentSource.ts @@ -0,0 +1,74 @@ +/** + * Copyright 2026 GoodRx, Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +import Model from './_Model'; + +export default class AgentSource extends Model { + uuid!: string; + sessionId!: number; + adapter!: string; + status!: 'requested' | 'preparing' | 'ready' | 'failed' | 'cleaned_up'; + input!: Record; + preparedSource!: Record; + sandboxRequirements!: Record; + error!: Record | null; + preparedAt!: string | null; + cleanedUpAt!: string | null; + + static tableName = 'agent_sources'; + static timestamps = true; + static idColumn = 'id'; + + static jsonSchema = { + type: 'object', + required: ['sessionId', 'adapter', 'status'], + properties: { + id: { type: 'integer' }, + uuid: { + type: 'string', + pattern: '^[0-9a-fA-F-]{36}$', + }, + sessionId: { type: 'integer' }, + adapter: { type: 'string' }, + status: { type: 'string', enum: ['requested', 'preparing', 'ready', 'failed', 'cleaned_up'] }, + input: { type: 'object', default: {} }, + preparedSource: { type: 'object', default: {} }, + sandboxRequirements: { type: 'object', default: {} }, + error: { type: ['object', 'null'], default: null }, + preparedAt: { type: ['string', 'null'] }, + cleanedUpAt: { type: ['string', 'null'] }, + }, + }; + + static get jsonAttributes() { + return ['input', 'preparedSource', 'sandboxRequirements', 'error']; + } + + static get relationMappings() { + const AgentSession = require('./AgentSession').default; + + return { + session: { + relation: Model.BelongsToOneRelation, + modelClass: AgentSession, + join: { + from: 'agent_sources.sessionId', + to: 'agent_sessions.id', + }, + }, + }; + } +} diff --git a/src/server/models/__tests__/AgentModelsValidation.test.ts b/src/server/models/__tests__/AgentModelsValidation.test.ts index d15c34f6..b2322190 100644 --- a/src/server/models/__tests__/AgentModelsValidation.test.ts +++ b/src/server/models/__tests__/AgentModelsValidation.test.ts @@ -27,16 +27,12 @@ describe('Agent model validation', () => { ).not.toThrow(); }); - test('allows ui messages with non-uuid ids because row identity comes from uiMessage.id', () => { + test('allows canonical messages without a uiMessage projection', () => { expect(() => AgentMessage.fromJson({ threadId: 42, role: 'user', - uiMessage: { - id: 'msg_local_123', - role: 'user', - parts: [{ type: 'text', text: 'hello' }], - }, + parts: [{ type: 'text', text: 'hello' }], metadata: {}, }) ).not.toThrow(); diff --git a/src/server/models/index.ts b/src/server/models/index.ts index 60549629..cea619af 100644 --- a/src/server/models/index.ts +++ b/src/server/models/index.ts @@ -33,10 +33,14 @@ import ConversationMessage from './ConversationMessage'; import MessageFeedback from './MessageFeedback'; import ConversationFeedback from './ConversationFeedback'; import AgentSession from './AgentSession'; +import AgentSource from './AgentSource'; +import AgentSandbox from './AgentSandbox'; +import AgentSandboxExposure from './AgentSandboxExposure'; import AgentPrewarm from './AgentPrewarm'; import UserApiKey from './UserApiKey'; import AgentThread from './AgentThread'; import AgentRun from './AgentRun'; +import AgentRunEvent from './AgentRunEvent'; import AgentMessage from './AgentMessage'; import AgentPendingAction from './AgentPendingAction'; import AgentToolExecution from './AgentToolExecution'; @@ -62,10 +66,14 @@ export interface IModels { MessageFeedback: typeof MessageFeedback; ConversationFeedback: typeof ConversationFeedback; AgentSession: typeof AgentSession; + AgentSource: typeof AgentSource; + AgentSandbox: typeof AgentSandbox; + AgentSandboxExposure: typeof AgentSandboxExposure; AgentPrewarm: typeof AgentPrewarm; UserApiKey: typeof UserApiKey; AgentThread: typeof AgentThread; AgentRun: typeof AgentRun; + AgentRunEvent: typeof AgentRunEvent; AgentMessage: typeof AgentMessage; AgentPendingAction: typeof AgentPendingAction; AgentToolExecution: typeof AgentToolExecution; @@ -92,10 +100,14 @@ export { MessageFeedback, ConversationFeedback, AgentSession, + AgentSource, + AgentSandbox, + AgentSandboxExposure, AgentPrewarm, UserApiKey, AgentThread, AgentRun, + AgentRunEvent, AgentMessage, AgentPendingAction, AgentToolExecution, diff --git a/src/server/services/__tests__/agentSession.test.ts b/src/server/services/__tests__/agentSession.test.ts index 1064a23d..8d15c325 100644 --- a/src/server/services/__tests__/agentSession.test.ts +++ b/src/server/services/__tests__/agentSession.test.ts @@ -23,8 +23,22 @@ const mockGetReadyPrewarmByPvc = jest.fn(); const mockExecInPod = jest.fn(); const mockResolveSessionPodServersForRepo = jest.fn().mockResolvedValue([]); const mockGetDefaultThreadForSession = jest.fn().mockResolvedValue({ uuid: 'default-thread-1' }); +const mockCreateOrUpdateNamespace = jest.fn().mockResolvedValue(undefined); +const mockDeleteNamespace = jest.fn().mockResolvedValue(undefined); +const mockCreateOrUpdateChatPreview = jest.fn().mockResolvedValue({ + url: 'https://chat-aaaaaaaa-3000.example.test', + host: 'chat-aaaaaaaa-3000.example.test', + path: '/', + port: 3000, + serviceName: 'agent-preview-aaaaaaaa-3000', + ingressName: 'agent-preview-ingress-aaaaaaaa-3000', +}); jest.mock('server/models/AgentSession'); +jest.mock('server/models/AgentThread'); +jest.mock('server/models/AgentSource'); +jest.mock('server/models/AgentSandbox'); +jest.mock('server/models/AgentSandboxExposure'); jest.mock('server/models/Build'); jest.mock('server/models/Deploy'); jest.mock('server/lib/dependencies', () => ({})); @@ -37,6 +51,13 @@ jest.mock('server/lib/agentSession/gvisorCheck'); jest.mock('server/lib/agentSession/configSeeder'); jest.mock('server/lib/agentSession/devModeManager'); jest.mock('server/lib/agentSession/forwardedEnv'); +jest.mock('server/lib/agentSession/chatPreviewFactory', () => ({ + createOrUpdateChatPreview: (...args: unknown[]) => mockCreateOrUpdateChatPreview(...args), +})); +jest.mock('server/lib/kubernetes', () => ({ + createOrUpdateNamespace: (...args: unknown[]) => mockCreateOrUpdateNamespace(...args), + deleteNamespace: (...args: unknown[]) => mockDeleteNamespace(...args), +})); jest.mock('server/lib/kubernetes/networkPolicyFactory'); jest.mock('server/services/ai/mcp/config', () => ({ __esModule: true, @@ -50,6 +71,7 @@ jest.mock('server/lib/agentSession/runtimeConfig', () => { __esModule: true, ...actual, resolveAgentSessionControlPlaneConfig: jest.fn(actual.resolveAgentSessionControlPlaneConfig), + resolveAgentSessionRuntimeConfig: jest.fn(actual.resolveAgentSessionRuntimeConfig), }; }); jest.mock('server/lib/agentSession/systemPrompt', () => { @@ -204,6 +226,10 @@ jest.mock('server/services/globalConfig', () => ({ import AgentSessionService, { CreateSessionOptions, buildAgentSessionPodName } from 'server/services/agentSession'; import AgentSession from 'server/models/AgentSession'; +import AgentThread from 'server/models/AgentThread'; +import AgentSource from 'server/models/AgentSource'; +import AgentSandbox from 'server/models/AgentSandbox'; +import AgentSandboxExposure from 'server/models/AgentSandboxExposure'; import Build from 'server/models/Build'; import Deploy from 'server/models/Deploy'; import { createAgentPvc, deleteAgentPvc } from 'server/lib/agentSession/pvcFactory'; @@ -231,6 +257,7 @@ import { deployHelm } from 'server/lib/nativeHelm/helm'; import { DeploymentManager } from 'server/lib/deploymentManager/deploymentManager'; import BuildServiceModule from 'server/services/build'; import { loadAgentSessionServiceCandidates } from 'server/services/agentSessionCandidates'; +import { AgentChatStatus, AgentSessionKind, AgentWorkspaceStatus } from 'shared/constants'; const mockRedis = { setex: jest.fn().mockResolvedValue('OK'), @@ -314,10 +341,43 @@ const mockSessionQuery = { select: jest.fn(), findById: jest.fn().mockReturnThis(), patch: jest.fn().mockResolvedValue(1), + patchAndFetchById: jest.fn(), insert: jest.fn().mockResolvedValue({}), insertAndFetch: jest.fn(), }; (AgentSession.query as jest.Mock) = jest.fn().mockReturnValue(mockSessionQuery); +(AgentSession.transaction as jest.Mock) = jest.fn(); + +const mockThreadQuery = { + insertAndFetch: jest.fn(), +}; +(AgentThread.query as jest.Mock) = jest.fn().mockReturnValue(mockThreadQuery); + +const mockSourceQuery = { + findOne: jest.fn(), + insert: jest.fn().mockResolvedValue({}), + insertAndFetch: jest.fn(), + patchAndFetchById: jest.fn(), +}; +(AgentSource.query as jest.Mock) = jest.fn().mockReturnValue(mockSourceQuery); + +const mockSandboxQuery = { + where: jest.fn().mockReturnThis(), + orderBy: jest.fn().mockReturnThis(), + first: jest.fn(), + insertAndFetch: jest.fn(), + patchAndFetchById: jest.fn(), +}; +(AgentSandbox.query as jest.Mock) = jest.fn().mockReturnValue(mockSandboxQuery); + +const mockSandboxExposureQuery = { + where: jest.fn().mockReturnThis(), + whereNull: jest.fn().mockReturnThis(), + first: jest.fn(), + insert: jest.fn().mockResolvedValue({}), + patchAndFetchById: jest.fn(), +}; +(AgentSandboxExposure.query as jest.Mock) = jest.fn().mockReturnValue(mockSandboxExposureQuery); const mockDeployQuery = { where: jest.fn().mockReturnThis(), @@ -363,6 +423,11 @@ describe('AgentSessionService', () => { beforeEach(() => { jest.clearAllMocks(); (AgentSession.query as jest.Mock) = jest.fn().mockReturnValue(mockSessionQuery); + (AgentSession.transaction as jest.Mock) = jest.fn(async (callback) => callback({ trx: true })); + (AgentThread.query as jest.Mock) = jest.fn().mockReturnValue(mockThreadQuery); + (AgentSource.query as jest.Mock) = jest.fn().mockReturnValue(mockSourceQuery); + (AgentSandbox.query as jest.Mock) = jest.fn().mockReturnValue(mockSandboxQuery); + (AgentSandboxExposure.query as jest.Mock) = jest.fn().mockReturnValue(mockSandboxExposureQuery); (Deploy.query as jest.Mock) = jest.fn().mockReturnValue(mockDeployQuery); mockSessionQuery.where.mockReturnThis(); mockSessionQuery.whereIn.mockReturnThis(); @@ -372,21 +437,87 @@ describe('AgentSessionService', () => { mockSessionQuery.select.mockResolvedValue({ id: 123 }); mockSessionQuery.findById.mockReturnThis(); mockSessionQuery.patch.mockResolvedValue(1); + mockSessionQuery.patchAndFetchById.mockResolvedValue({ + id: 123, + uuid: 'aaaaaaaa-bbbb-cccc-dddd-eeeeeeeeeeee', + userId: 'user-123', + ownerGithubUsername: null, + sessionKind: 'environment', + podName: 'agent-aaaaaaaa', + namespace: 'test-ns', + pvcName: 'agent-pvc-aaaaaaaa', + model: 'claude-sonnet-4-6', + buildKind: 'environment', + status: 'starting', + chatStatus: 'ready', + workspaceStatus: 'provisioning', + defaultThreadId: 456, + devModeSnapshots: {}, + forwardedAgentSecretProviders: [], + }); mockSessionQuery.insert.mockResolvedValue({}); mockSessionQuery.insertAndFetch.mockResolvedValue({ id: 123, uuid: 'aaaaaaaa-bbbb-cccc-dddd-eeeeeeeeeeee', userId: 'user-123', ownerGithubUsername: null, + sessionKind: 'environment', podName: 'agent-aaaaaaaa', namespace: 'test-ns', pvcName: 'agent-pvc-aaaaaaaa', model: 'claude-sonnet-4-6', buildKind: 'environment', status: 'starting', + chatStatus: 'ready', + workspaceStatus: 'provisioning', devModeSnapshots: {}, forwardedAgentSecretProviders: [], }); + mockThreadQuery.insertAndFetch.mockResolvedValue({ + id: 456, + uuid: 'default-thread-1', + sessionId: 123, + title: 'Default thread', + isDefault: true, + metadata: { + sessionUuid: 'aaaaaaaa-bbbb-cccc-dddd-eeeeeeeeeeee', + }, + }); + mockSourceQuery.findOne.mockResolvedValue(null); + mockSourceQuery.insert.mockResolvedValue({}); + mockSourceQuery.insertAndFetch.mockResolvedValue({ + id: 321, + uuid: 'source-1', + sessionId: 123, + status: 'ready', + }); + mockSourceQuery.patchAndFetchById.mockResolvedValue({}); + mockSandboxQuery.where.mockReturnThis(); + mockSandboxQuery.orderBy.mockReturnThis(); + mockSandboxQuery.first.mockResolvedValue(null); + mockSandboxQuery.insertAndFetch.mockResolvedValue({ + id: 654, + uuid: 'sandbox-1', + sessionId: 123, + generation: 1, + provider: 'lifecycle_kubernetes', + status: 'provisioning', + endedAt: null, + }); + mockSandboxQuery.patchAndFetchById.mockResolvedValue({ + id: 654, + uuid: 'sandbox-1', + sessionId: 123, + generation: 1, + provider: 'lifecycle_kubernetes', + status: 'ready', + endedAt: null, + }); + mockSandboxExposureQuery.where.mockReturnThis(); + mockSandboxExposureQuery.whereNull.mockReturnThis(); + mockSandboxExposureQuery.first.mockResolvedValue(null); + mockSandboxExposureQuery.insert.mockResolvedValue({}); + mockSandboxExposureQuery.patchAndFetchById.mockResolvedValue({}); mockDeployQuery.where.mockReturnThis(); mockDeployQuery.whereIn.mockReturnThis(); mockDeployQuery.findById.mockReturnThis(); @@ -428,6 +559,16 @@ describe('AgentSessionService', () => { mockGetCompatibleReadyPrewarm.mockResolvedValue(null); mockGetReadyPrewarmByPvc.mockResolvedValue(null); mockGetDefaultThreadForSession.mockResolvedValue({ uuid: 'default-thread-1' }); + mockCreateOrUpdateNamespace.mockResolvedValue(undefined); + mockDeleteNamespace.mockResolvedValue(undefined); + mockCreateOrUpdateChatPreview.mockResolvedValue({ + url: 'https://chat-aaaaaaaa-3000.example.test', + host: 'chat-aaaaaaaa-3000.example.test', + path: '/', + port: 3000, + serviceName: 'agent-preview-aaaaaaaa-3000', + ingressName: 'agent-preview-ingress-aaaaaaaa-3000', + }); mockExecInPod.mockImplementation( async ( _namespace: string, @@ -458,6 +599,36 @@ describe('AgentSessionService', () => { (runtimeConfig.resolveAgentSessionControlPlaneConfig as jest.Mock).mockResolvedValue({ appendSystemPrompt: undefined, }); + (runtimeConfig.resolveAgentSessionRuntimeConfig as jest.Mock).mockResolvedValue({ + workspaceImage: 'lifecycle-agent:latest', + workspaceEditorImage: 'codercom/code-server:4.98.2', + workspaceGatewayImage: 'lifecycle-agent:latest', + nodeSelector: undefined, + keepAttachedServicesOnSessionNode: true, + readiness: undefined, + resources: undefined, + workspaceStorage: { + defaultSize: '10Gi', + allowedSizes: ['10Gi'], + allowClientOverride: false, + accessMode: 'ReadWriteOnce', + }, + cleanup: { + activeIdleSuspendMs: 30 * 60 * 1000, + startingTimeoutMs: 15 * 60 * 1000, + hibernatedRetentionMs: 24 * 60 * 60 * 1000, + intervalMs: 5 * 60 * 1000, + redisTtlSeconds: 7200, + }, + durability: { + runExecutionLeaseMs: 30 * 60 * 1000, + queuedRunDispatchStaleMs: 30 * 1000, + dispatchRecoveryLimit: 50, + maxDurablePayloadBytes: 64 * 1024, + payloadPreviewBytes: 16 * 1024, + fileChangePreviewChars: 4000, + }, + }); (systemPrompt.buildAgentSessionDynamicSystemPrompt as jest.Mock).mockImplementation( jest.requireActual('server/lib/agentSession/systemPrompt').buildAgentSessionDynamicSystemPrompt ); @@ -469,6 +640,221 @@ describe('AgentSessionService', () => { ); }); + it('creates a chat session without provisioning runtime resources', async () => { + mockSessionQuery.insertAndFetch.mockResolvedValueOnce({ + id: 321, + uuid: 'aaaaaaaa-bbbb-cccc-dddd-eeeeeeeeeeee', + userId: 'user-123', + ownerGithubUsername: null, + sessionKind: 'chat', + podName: null, + namespace: null, + pvcName: null, + model: 'claude-sonnet-4-6', + buildKind: null, + status: 'active', + chatStatus: 'ready', + workspaceStatus: 'none', + devModeSnapshots: {}, + forwardedAgentSecretProviders: [], + workspaceRepos: [], + selectedServices: [], + skillPlan: { version: 1, skills: [] }, + }); + mockThreadQuery.insertAndFetch.mockResolvedValueOnce({ + id: 654, + uuid: 'default-thread-1', + sessionId: 321, + title: 'Default thread', + isDefault: true, + metadata: { + sessionUuid: 'aaaaaaaa-bbbb-cccc-dddd-eeeeeeeeeeee', + }, + }); + mockSessionQuery.patchAndFetchById.mockResolvedValueOnce({ + id: 321, + uuid: 'aaaaaaaa-bbbb-cccc-dddd-eeeeeeeeeeee', + userId: 'user-123', + ownerGithubUsername: null, + sessionKind: 'chat', + podName: null, + namespace: null, + pvcName: null, + model: 'claude-sonnet-4-6', + buildKind: null, + status: 'active', + chatStatus: 'ready', + workspaceStatus: 'none', + defaultThreadId: 654, + devModeSnapshots: {}, + forwardedAgentSecretProviders: [], + workspaceRepos: [], + selectedServices: [], + skillPlan: { version: 1, skills: [] }, + }); + + const session = await AgentSessionService.createChatSession({ + userId: 'user-123', + model: 'claude-sonnet-4-6', + }); + await new Promise((resolve) => setImmediate(resolve)); + + expect(mockSessionQuery.insertAndFetch).toHaveBeenCalledWith( + expect.objectContaining({ + sessionKind: AgentSessionKind.CHAT, + buildKind: null, + podName: null, + namespace: null, + pvcName: null, + status: 'active', + chatStatus: AgentChatStatus.READY, + workspaceStatus: AgentWorkspaceStatus.NONE, + }) + ); + expect(createAgentPvc).not.toHaveBeenCalled(); + expect(createSessionWorkspacePod).not.toHaveBeenCalled(); + expect(createSessionWorkspaceService).not.toHaveBeenCalled(); + expect(mockThreadQuery.insertAndFetch).toHaveBeenCalledWith( + expect.objectContaining({ + sessionId: 321, + title: 'Default thread', + isDefault: true, + }) + ); + expect(mockSourceQuery.insertAndFetch).toHaveBeenCalledWith( + expect.objectContaining({ + sessionId: 321, + adapter: 'blank_workspace', + status: 'ready', + }) + ); + expect(AgentSession.transaction).toHaveBeenCalledTimes(1); + expect(mockSessionQuery.patchAndFetchById).toHaveBeenCalledWith( + 321, + expect.objectContaining({ + defaultThreadId: 654, + }) + ); + expect(mockGetDefaultThreadForSession).not.toHaveBeenCalled(); + expect(session.sessionKind).toBe('chat'); + expect(session.workspaceStatus).toBe('none'); + expect(session.defaultThreadId).toBe(654); + }); + + it('provisions a blank workspace runtime for a chat session on demand', async () => { + const chatSession = { + id: 321, + uuid: 'aaaaaaaa-bbbb-cccc-dddd-eeeeeeeeeeee', + userId: 'user-123', + ownerGithubUsername: 'sample-user', + sessionKind: 'chat', + podName: null, + namespace: null, + pvcName: null, + model: 'claude-sonnet-4-6', + buildKind: null, + status: 'active', + chatStatus: 'ready', + workspaceStatus: 'none', + devModeSnapshots: {}, + forwardedAgentSecretProviders: [], + workspaceRepos: [], + selectedServices: [], + skillPlan: { version: 1, skills: [] }, + }; + const readyChatSession = { + ...chatSession, + namespace: 'chat-aaaaaaaa', + podName: 'agent-aaaaaaaa', + pvcName: 'agent-pvc-aaaaaaaa', + workspaceStatus: 'ready', + }; + + mockSessionQuery.findOne.mockResolvedValueOnce(chatSession).mockResolvedValueOnce(readyChatSession); + + const session = await AgentSessionService.provisionChatRuntime({ + sessionId: 'aaaaaaaa-bbbb-cccc-dddd-eeeeeeeeeeee', + userId: 'user-123', + userIdentity: { + userId: 'user-123', + githubUsername: 'sample-user', + } as any, + githubToken: 'sample-gh-token', + }); + + expect(mockCreateOrUpdateNamespace).toHaveBeenCalledWith( + expect.objectContaining({ + name: 'chat-aaaaaaaa', + buildUUID: 'aaaaaaaa-bbbb-cccc-dddd-eeeeeeeeeeee', + author: 'sample-user', + }) + ); + expect(createAgentPvc).toHaveBeenCalledWith( + 'chat-aaaaaaaa', + 'agent-pvc-aaaaaaaa', + '10Gi', + undefined, + 'ReadWriteOnce' + ); + expect(createAgentApiKeySecret).toHaveBeenCalledWith( + 'chat-aaaaaaaa', + 'agent-secret-aaaaaaaa', + {}, + 'sample-gh-token', + undefined, + {}, + { + LIFECYCLE_SESSION_MCP_CONFIG_JSON: '[]', + } + ); + expect(createSessionWorkspacePod).toHaveBeenCalledWith( + expect.objectContaining({ + namespace: 'chat-aaaaaaaa', + podName: 'agent-aaaaaaaa', + pvcName: 'agent-pvc-aaaaaaaa', + workspaceRepos: [], + }) + ); + expect(createSessionWorkspaceService).toHaveBeenCalledWith('chat-aaaaaaaa', 'agent-aaaaaaaa'); + expect(mockSessionQuery.patch).toHaveBeenCalledWith( + expect.objectContaining({ + workspaceStatus: AgentWorkspaceStatus.READY, + namespace: 'chat-aaaaaaaa', + podName: 'agent-aaaaaaaa', + pvcName: 'agent-pvc-aaaaaaaa', + }) + ); + expect(session.workspaceStatus).toBe('ready'); + expect(session.namespace).toBe('chat-aaaaaaaa'); + }); + + it('publishes a chat session HTTP port through ingress', async () => { + mockSessionQuery.findOne.mockResolvedValue({ + id: 321, + uuid: 'aaaaaaaa-bbbb-cccc-dddd-eeeeeeeeeeee', + userId: 'user-123', + sessionKind: 'chat', + namespace: 'chat-aaaaaaaa', + podName: 'agent-aaaaaaaa', + workspaceStatus: 'ready', + status: 'active', + }); + + const publication = await AgentSessionService.publishChatHttpPort({ + sessionId: 'aaaaaaaa-bbbb-cccc-dddd-eeeeeeeeeeee', + userId: 'user-123', + port: 3000, + }); + + expect(mockCreateOrUpdateChatPreview).toHaveBeenCalledWith({ + sessionUuid: 'aaaaaaaa-bbbb-cccc-dddd-eeeeeeeeeeee', + namespace: 'chat-aaaaaaaa', + podName: 'agent-aaaaaaaa', + port: 3000, + }); + expect(publication.url).toBe('https://chat-aaaaaaaa-3000.example.test'); + }); + describe('createSession', () => { it('throws an active environment session error when another user already owns the environment session', async () => { const conflictingOpts: CreateSessionOptions = { @@ -518,7 +904,7 @@ describe('AgentSessionService', () => { it('creates PVC, pod, network policy, and session record', async () => { const session = await AgentSessionService.createSession(baseOpts); - expect(createAgentPvc).toHaveBeenCalledWith('test-ns', 'agent-pvc-aaaaaaaa', '10Gi', undefined); + expect(createAgentPvc).toHaveBeenCalledWith('test-ns', 'agent-pvc-aaaaaaaa', '10Gi', undefined, 'ReadWriteOnce'); expect(ensureAgentSessionServiceAccount).toHaveBeenCalledWith('test-ns'); expect(createAgentApiKeySecret).toHaveBeenCalledWith( 'test-ns', @@ -546,6 +932,7 @@ describe('AgentSessionService', () => { }) ); expect(createSessionWorkspaceService).toHaveBeenCalledWith('test-ns', 'agent-aaaaaaaa', undefined); + expect(AgentSession.transaction).toHaveBeenCalledTimes(1); expect(mockSessionQuery.insertAndFetch).toHaveBeenCalledWith( expect.objectContaining({ uuid: 'aaaaaaaa-bbbb-cccc-dddd-eeeeeeeeeeee', @@ -676,7 +1063,13 @@ describe('AgentSessionService', () => { ], }) ); - expect(createAgentPvc).toHaveBeenCalledWith('test-ns', 'agent-pvc-aaaaaaaa', '10Gi', 'build-123'); + expect(createAgentPvc).toHaveBeenCalledWith( + 'test-ns', + 'agent-pvc-aaaaaaaa', + '10Gi', + 'build-123', + 'ReadWriteOnce' + ); expect(createSessionWorkspacePod).toHaveBeenCalledWith( expect.objectContaining({ podName: 'agent-build-123', @@ -2218,7 +2611,6 @@ describe('AgentSessionService', () => { expect(DeploymentManager).toHaveBeenCalledWith(devModeDeploys); expect(deployManagerDeploy).toHaveBeenCalled(); - expect(mockDisableDevMode).toHaveBeenCalledTimes(1); expect(mockDisableDevMode).toHaveBeenCalledWith( 'test-ns', 'deploy-10', diff --git a/src/server/services/__tests__/agentSessionConfig.test.ts b/src/server/services/__tests__/agentSessionConfig.test.ts index 404f87b8..11ba4466 100644 --- a/src/server/services/__tests__/agentSessionConfig.test.ts +++ b/src/server/services/__tests__/agentSessionConfig.test.ts @@ -54,7 +54,7 @@ describe('AgentSessionConfigService', () => { mockGlobalConfigSetConfig.mockResolvedValue(undefined); }); - it('lists only admin-visible sandbox tools in tool inventory', async () => { + it('lists admin-visible built-in tools in tool inventory', async () => { const service = makeService(); jest.spyOn(service, 'getGlobalConfig').mockResolvedValue({}); @@ -80,9 +80,24 @@ describe('AgentSessionConfigService', () => { 'git.add', 'git.commit', 'git.branch', + 'publish_http', ]); expect(entries.find((entry) => entry.toolName === 'skills.list')).toBeUndefined(); expect(entries.find((entry) => entry.toolName === 'session.get_workspace_state')).toBeUndefined(); + expect(entries.find((entry) => entry.toolName === 'publish_http')).toEqual( + expect.objectContaining({ + toolKey: 'mcp__lifecycle__publish_http', + serverSlug: 'lifecycle', + serverName: 'Lifecycle', + sourceType: 'builtin', + sourceScope: 'session', + capabilityKey: 'deploy_k8s_mutation', + approvalMode: 'require_approval', + scopeRuleMode: 'inherit', + effectiveRuleMode: 'inherit', + availability: 'available', + }) + ); }); it('merges repo control-plane numeric overrides over global defaults', async () => { @@ -218,6 +233,27 @@ describe('AgentSessionConfigService', () => { }, }, }, + workspaceStorage: { + defaultSize: '20Gi', + allowedSizes: ['10Gi', '20Gi'], + allowClientOverride: true, + accessMode: 'ReadWriteMany', + }, + cleanup: { + activeIdleSuspendMs: 60000, + startingTimeoutMs: 120000, + hibernatedRetentionMs: 180000, + intervalMs: 30000, + redisTtlSeconds: 900, + }, + durability: { + runExecutionLeaseMs: 45000, + queuedRunDispatchStaleMs: 5000, + dispatchRecoveryLimit: 12, + maxDurablePayloadBytes: 4096, + payloadPreviewBytes: 512, + fileChangePreviewChars: 600, + }, }) ).resolves.toEqual({ workspaceImage: 'workspace-image:v2', @@ -240,6 +276,27 @@ describe('AgentSessionConfigService', () => { }, }, }, + workspaceStorage: { + defaultSize: '20Gi', + allowedSizes: ['10Gi', '20Gi'], + allowClientOverride: true, + accessMode: 'ReadWriteMany', + }, + cleanup: { + activeIdleSuspendMs: 60000, + startingTimeoutMs: 120000, + hibernatedRetentionMs: 180000, + intervalMs: 30000, + redisTtlSeconds: 900, + }, + durability: { + runExecutionLeaseMs: 45000, + queuedRunDispatchStaleMs: 5000, + dispatchRecoveryLimit: 12, + maxDurablePayloadBytes: 4096, + payloadPreviewBytes: 512, + fileChangePreviewChars: 600, + }, }); expect(mockGlobalConfigSetConfig).toHaveBeenCalledWith('agentSessionDefaults', { @@ -266,6 +323,27 @@ describe('AgentSessionConfigService', () => { }, }, }, + workspaceStorage: { + defaultSize: '20Gi', + allowedSizes: ['10Gi', '20Gi'], + allowClientOverride: true, + accessMode: 'ReadWriteMany', + }, + cleanup: { + activeIdleSuspendMs: 60000, + startingTimeoutMs: 120000, + hibernatedRetentionMs: 180000, + intervalMs: 30000, + redisTtlSeconds: 900, + }, + durability: { + runExecutionLeaseMs: 45000, + queuedRunDispatchStaleMs: 5000, + dispatchRecoveryLimit: 12, + maxDurablePayloadBytes: 4096, + payloadPreviewBytes: 512, + fileChangePreviewChars: 600, + }, }); }); diff --git a/src/server/services/__tests__/globalConfig.test.ts b/src/server/services/__tests__/globalConfig.test.ts index afaa3a27..ef23d0c0 100644 --- a/src/server/services/__tests__/globalConfig.test.ts +++ b/src/server/services/__tests__/globalConfig.test.ts @@ -30,6 +30,7 @@ jest.mock('ioredis', () => { return jest.fn().mockImplementation(() => ({ hgetall: jest.fn(), hmset: jest.fn(), + del: jest.fn(), })); }); jest.mock('@octokit/auth-app', () => ({ @@ -89,6 +90,32 @@ describe('GlobalConfigService', () => { }); }); + describe('setConfig', () => { + it('updates shared cache and clears the in-memory config cache after writing', async () => { + const upsertQuery = { + insert: jest.fn().mockReturnThis(), + onConflict: jest.fn().mockReturnThis(), + merge: jest.fn().mockResolvedValue(undefined), + }; + service.db = { + knex: jest.fn().mockReturnValue(upsertQuery), + }; + service.memoryCache = { agentSessionDefaults: { workspaceImage: 'stale-image' } }; + service.memoryCacheExpiry = Date.now() + 10000; + + const config = { workspaceImage: 'workspace-image:v2' }; + await service.setConfig('agentSessionDefaults', config); + + expect(service.db.knex).toHaveBeenCalledWith('global_config'); + expect(upsertQuery.insert).toHaveBeenCalledWith({ key: 'agentSessionDefaults', config }); + expect(upsertQuery.onConflict).toHaveBeenCalledWith('key'); + expect(upsertQuery.merge).toHaveBeenCalledWith(); + expect(service.redis.del).toHaveBeenCalledWith('global_config'); + expect(service.memoryCache).toBeNull(); + expect(service.memoryCacheExpiry).toBe(0); + }); + }); + describe('setupCacheRefreshJob', () => { it('should set up a cache refresh job', async () => { await service.setupCacheRefreshJob(); diff --git a/src/server/services/agent/AdminService.ts b/src/server/services/agent/AdminService.ts index e93f9416..fbf0f595 100644 --- a/src/server/services/agent/AdminService.ts +++ b/src/server/services/agent/AdminService.ts @@ -17,6 +17,7 @@ import AgentMessage from 'server/models/AgentMessage'; import AgentPendingAction from 'server/models/AgentPendingAction'; import AgentRun from 'server/models/AgentRun'; +import AgentRunEvent from 'server/models/AgentRunEvent'; import AgentSession from 'server/models/AgentSession'; import AgentThread from 'server/models/AgentThread'; import AgentToolExecution from 'server/models/AgentToolExecution'; @@ -25,9 +26,11 @@ import UserMcpConnection from 'server/models/UserMcpConnection'; import AgentSessionService from 'server/services/agentSession'; import UserMcpConnectionService from 'server/services/userMcpConnection'; import AgentRunService from './RunService'; +import AgentRunEventService from './RunEventService'; import ApprovalService from './ApprovalService'; import AgentThreadService from './ThreadService'; -import type { AgentUIMessage } from './types'; +import AgentMessageStore from './MessageStore'; +import type { CanonicalAgentMessage } from './canonicalMessages'; import type { McpAuthConfig, McpDiscoveredTool, @@ -69,6 +72,7 @@ type AdminToolExecutionRecord = { source: string; serverSlug: string | null; toolName: string; + toolCallId: string | null; args: Record; result: Record | null; status: AgentToolExecution['status']; @@ -117,10 +121,6 @@ function normalizeSearchTerm(value?: string | null): string | null { return trimmed ? trimmed : null; } -function toAgentUiMessage(message: AgentMessage): AgentUIMessage { - return message.uiMessage as unknown as AgentUIMessage; -} - function toToolExecutionRecord(row: AgentToolExecution): AdminToolExecutionRecord { const enrichedRow = row as AgentToolExecution & { threadUuid?: string; @@ -136,6 +136,7 @@ function toToolExecutionRecord(row: AgentToolExecution): AdminToolExecutionRecor source: row.source, serverSlug: row.serverSlug, toolName: row.toolName, + toolCallId: row.toolCallId || null, args: (row.args || {}) as Record, result: (row.result as Record | null) || null, status: row.status, @@ -149,6 +150,18 @@ function toToolExecutionRecord(row: AgentToolExecution): AdminToolExecutionRecor }; } +function toCanonicalMessageRecord(message: AgentMessage, threadUuid: string): CanonicalAgentMessage | null { + try { + return AgentMessageStore.serializeCanonicalMessage( + message, + threadUuid, + (message as AgentMessage & { runUuid?: string | null }).runUuid || null + ); + } catch { + return null; + } +} + function paginateArray(items: T[], page = 1, limit = 25) { const safePage = Number.isFinite(page) && page > 0 ? page : 1; const safeLimit = Number.isFinite(limit) && limit > 0 ? limit : 25; @@ -177,6 +190,7 @@ function serializeSessionSummary( ) { return { id: session.uuid, + sessionKind: session.sessionKind, buildUuid: session.buildUuid, baseBuildUuid: session.baseBuildUuid, buildKind: session.buildKind, @@ -187,6 +201,8 @@ function serializeSessionSummary( pvcName: session.pvcName, model: session.model, status: session.status, + chatStatus: session.chatStatus, + workspaceStatus: session.workspaceStatus, repo: session.repo, branch: session.branch, primaryRepo: session.primaryRepo, @@ -202,7 +218,7 @@ function serializeSessionSummary( lastRunAt: counts?.lastRunAt ?? null, createdAt: session.createdAt || null, updatedAt: session.updatedAt || null, - editorUrl: `/api/agent-session/workspace-editor/${session.uuid}/`, + editorUrl: session.podName && session.namespace ? `/api/agent-session/workspace-editor/${session.uuid}/` : null, }; } @@ -372,9 +388,14 @@ export default class AgentAdminService { throw new Error('Agent session not found'); } - const [sessionDetail, messageRows, runRows, pendingRows, toolRows] = await Promise.all([ + const [sessionDetail, messageRows, runRows, pendingRows, toolRows, eventRows] = await Promise.all([ this.getSession(session.uuid), - AgentMessage.query().where({ threadId: thread.id }).orderBy('createdAt', 'asc'), + AgentMessage.query() + .alias('message') + .leftJoinRelated('run') + .where('message.threadId', thread.id) + .select('message.*', 'run.uuid as runUuid') + .orderBy('message.createdAt', 'asc'), AgentRun.query().where({ threadId: thread.id }).orderBy('createdAt', 'asc'), AgentPendingAction.query() .alias('action') @@ -384,11 +405,18 @@ export default class AgentAdminService { .orderBy('action.createdAt', 'asc'), AgentToolExecution.query() .alias('tool') - .joinRelated('run') + .joinRelated('[run, thread]') .leftJoinRelated('pendingAction') .where('tool.threadId', thread.id) - .select('tool.*', 'run.uuid as runUuid', 'pendingAction.uuid as pendingActionUuid') + .select('tool.*', 'thread.uuid as threadUuid', 'run.uuid as runUuid', 'pendingAction.uuid as pendingActionUuid') .orderBy('tool.createdAt', 'asc'), + AgentRunEvent.query() + .alias('event') + .joinRelated('run') + .where('run.threadId', thread.id) + .select('event.*', 'run.uuid as runUuid') + .orderBy('event.runId', 'asc') + .orderBy('event.sequence', 'asc'), ]); const runs = runRows.map((run) => @@ -415,8 +443,19 @@ export default class AgentAdminService { return { session: sessionDetail.session, thread: threadSummary, - messages: messageRows.map(toAgentUiMessage), + messages: messageRows.flatMap((message) => { + const serialized = toCanonicalMessageRecord(message, thread.uuid); + return serialized ? [serialized] : []; + }), runs, + events: eventRows.map((event) => + AgentRunEventService.serializeRunEvent({ + ...event, + runUuid: (event as AgentRunEvent & { runUuid?: string }).runUuid, + threadUuid: thread.uuid, + sessionUuid: session.uuid, + } as AgentRunEvent) + ), pendingActions, toolExecutions: toolRows.map(toToolExecutionRecord), }; diff --git a/src/server/services/agent/AgentRunOwnershipLostError.ts b/src/server/services/agent/AgentRunOwnershipLostError.ts new file mode 100644 index 00000000..01b7e2b8 --- /dev/null +++ b/src/server/services/agent/AgentRunOwnershipLostError.ts @@ -0,0 +1,47 @@ +/** + * Copyright 2026 GoodRx, Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +import type { AgentRunStatus } from './types'; + +export class AgentRunOwnershipLostError extends Error { + readonly runUuid: string; + readonly expectedExecutionOwner: string; + readonly currentStatus?: AgentRunStatus | null; + readonly currentExecutionOwner?: string | null; + + constructor({ + runUuid, + expectedExecutionOwner, + currentStatus, + currentExecutionOwner, + }: { + runUuid: string; + expectedExecutionOwner: string; + currentStatus?: AgentRunStatus | null; + currentExecutionOwner?: string | null; + }) { + super( + `Agent run execution ownership lost runUuid=${runUuid} expectedOwner=${expectedExecutionOwner} currentStatus=${ + currentStatus || 'unknown' + } currentOwner=${currentExecutionOwner || 'none'}` + ); + this.name = 'AgentRunOwnershipLostError'; + this.runUuid = runUuid; + this.expectedExecutionOwner = expectedExecutionOwner; + this.currentStatus = currentStatus; + this.currentExecutionOwner = currentExecutionOwner; + } +} diff --git a/src/server/services/agent/ApprovalService.ts b/src/server/services/agent/ApprovalService.ts index 3195a422..169ce934 100644 --- a/src/server/services/agent/ApprovalService.ts +++ b/src/server/services/agent/ApprovalService.ts @@ -16,25 +16,304 @@ import { getToolName, isToolUIPart, type DynamicToolUIPart, type ToolUIPart, type UITools } from 'ai'; import AgentPendingAction from 'server/models/AgentPendingAction'; -import type AgentRun from 'server/models/AgentRun'; +import AgentRun from 'server/models/AgentRun'; import type AgentThread from 'server/models/AgentThread'; -import type { AgentCapabilityKey, AgentPendingActionStatus, AgentUIMessage } from './types'; -import { addFileChangesToApprovalPayload } from './fileChanges'; +import type { Transaction } from 'objection'; +import type { + AgentApprovalPolicy, + AgentCapabilityKey, + AgentFileChangeData, + AgentPendingActionStatus, + AgentUIMessage, +} from './types'; +import type { AgentSessionToolRule } from 'server/services/types/agentSessionConfig'; +import { listMessageFileChanges } from './fileChanges'; import AgentThreadService from './ThreadService'; +import AgentRunQueueService from './RunQueueService'; +import AgentRunEventService from './RunEventService'; +import AgentPolicyService from './PolicyService'; +import { + buildAgentToolKey, + CHAT_PUBLISH_HTTP_TOOL_NAME, + LIFECYCLE_BUILTIN_SERVER_SLUG, + SESSION_WORKSPACE_SERVER_SLUG, +} from './toolKeys'; type ToolLikePart = ToolUIPart | DynamicToolUIPart; +const SESSION_WORKSPACE_TOOL_KEY_PREFIX = `mcp__${SESSION_WORKSPACE_SERVER_SLUG}__`; +const ARGUMENT_PREVIEW_MAX_LENGTH = 160; +const PENDING_ACTION_RESPONSE_FIELDS = new Set(['approved', 'reason']); + +type PendingActionResponseBody = { + approved: boolean; + reason: string | null; +}; + +type ApprovalRequestSyncResult = { + pendingActions: AgentPendingAction[]; + resolvedActionCount: number; +}; + +type ApprovalRequestSyncOptions = { + thread: AgentThread; + run: AgentRun; + messages: AgentUIMessage[]; + capabilityKey?: AgentCapabilityKey; + approvalPolicy?: AgentApprovalPolicy; + toolRules?: AgentSessionToolRule[]; + trx?: Transaction; +}; function isToolLikePart(part: unknown): part is ToolLikePart { return !!part && typeof part === 'object' && isToolUIPart(part as ToolLikePart); } +function isRecord(value: unknown): value is Record { + return !!value && typeof value === 'object' && !Array.isArray(value); +} + +function readString(value: unknown): string | null { + return typeof value === 'string' && value.trim() ? value : null; +} + +function truncatePreview(value: string): string { + return value.length > ARGUMENT_PREVIEW_MAX_LENGTH ? `${value.slice(0, ARGUMENT_PREVIEW_MAX_LENGTH - 3)}...` : value; +} + +function formatArgumentValue(value: unknown): string { + if (typeof value === 'string') { + return truncatePreview(value); + } + + if (typeof value === 'number' || typeof value === 'boolean') { + return String(value); + } + + if (value == null) { + return 'null'; + } + + try { + return truncatePreview(JSON.stringify(value)); + } catch { + return '[unserializable]'; + } +} + +function summarizeArguments(input: unknown): Array<{ name: string; value: string }> { + if (!isRecord(input)) { + return []; + } + + return Object.entries(input) + .filter(([name]) => !['content', 'oldText', 'newText', 'command', 'cmd'].includes(name)) + .slice(0, 6) + .map(([name, value]) => ({ + name, + value: formatArgumentValue(value), + })); +} + +function getCommandPreview(input: unknown): string | null { + if (!isRecord(input)) { + return null; + } + + const command = readString(input.command) || readString(input.cmd); + if (command) { + return truncatePreview(command); + } + + if (Array.isArray(input.command)) { + const commandParts = input.command.filter( + (part): part is string => typeof part === 'string' && part.trim().length > 0 + ); + return commandParts.length > 0 ? truncatePreview(commandParts.join(' ')) : null; + } + + return null; +} + +function summarizeFileChanges(value: unknown): Array<{ + path: string; + action: string; + summary: string; + additions: number | null; + deletions: number | null; + truncated: boolean; +}> { + if (!Array.isArray(value)) { + return []; + } + + return value + .filter(isRecord) + .slice(0, 10) + .map((change) => { + const path = readString(change.displayPath) || readString(change.path) || 'unknown'; + const action = readString(change.kind) || readString(change.stage) || 'change'; + return { + path, + action, + summary: readString(change.summary) || `${action} ${path}`, + additions: typeof change.additions === 'number' ? change.additions : null, + deletions: typeof change.deletions === 'number' ? change.deletions : null, + truncated: change.truncated === true, + }; + }); +} + +function getRiskLabels(capabilityKey: string | null | undefined): string[] { + switch (capabilityKey) { + case 'workspace_write': + return ['Workspace write']; + case 'shell_exec': + return ['Shell command']; + case 'git_write': + return ['Git write']; + case 'network_access': + return ['Network access']; + case 'deploy_k8s_mutation': + return ['Deployment change']; + case 'external_mcp_write': + return ['MCP write']; + case 'external_mcp_read': + return ['MCP read']; + case 'read': + return ['Read-only']; + default: + return []; + } +} + +function resolveApprovalCapabilityKey(toolName: string, fallback: AgentCapabilityKey): AgentCapabilityKey { + if (toolName === buildAgentToolKey(LIFECYCLE_BUILTIN_SERVER_SLUG, CHAT_PUBLISH_HTTP_TOOL_NAME)) { + return 'deploy_k8s_mutation'; + } + + if (!toolName.startsWith(SESSION_WORKSPACE_TOOL_KEY_PREFIX)) { + return fallback; + } + + const sessionWorkspaceToolName = toolName.slice(SESSION_WORKSPACE_TOOL_KEY_PREFIX.length).replace(/_/g, '.'); + + return AgentPolicyService.capabilityForSessionWorkspaceTool(sessionWorkspaceToolName); +} + +function shouldPersistApprovalRequest({ + toolName, + fallbackCapabilityKey, + approvalPolicy, + toolRules, +}: { + toolName: string; + fallbackCapabilityKey: AgentCapabilityKey; + approvalPolicy?: AgentApprovalPolicy; + toolRules?: AgentSessionToolRule[]; +}): boolean { + if (!approvalPolicy) { + return true; + } + + const capabilityKey = resolveApprovalCapabilityKey(toolName, fallbackCapabilityKey); + const toolRule = toolRules?.find((rule) => rule.toolKey === toolName); + const mode = toolRule?.mode || AgentPolicyService.modeForCapability(approvalPolicy, capabilityKey); + + return mode === 'require_approval'; +} + +async function upsertApprovalRequestRecord({ + thread, + run, + approvalId, + toolCallId, + toolName, + input, + fileChanges, + capabilityKey, + trx, +}: { + thread: AgentThread; + run: AgentRun; + approvalId: string; + toolCallId: string | null; + toolName: string; + input: unknown; + fileChanges?: AgentFileChangeData[]; + capabilityKey: AgentCapabilityKey; + trx?: Transaction; +}): Promise { + const existing = await AgentPendingAction.query(trx) + .where({ runId: run.id, threadId: thread.id }) + .whereRaw(`payload->>'approvalId' = ?`, [approvalId]) + .first(); + + const payload = { + approvalId, + toolCallId, + toolName, + input: input ?? null, + ...(fileChanges?.length ? { fileChanges } : {}), + }; + const resolvedCapabilityKey = resolveApprovalCapabilityKey(toolName, capabilityKey); + + if (existing) { + if (existing.status !== 'pending') { + return null; + } + + return AgentPendingAction.query(trx).patchAndFetchById(existing.id, { + capabilityKey: resolvedCapabilityKey, + payload, + } as Partial); + } + + return AgentPendingAction.query(trx).insertAndFetch({ + threadId: thread.id, + runId: run.id, + kind: 'tool_approval', + status: 'pending', + capabilityKey: resolvedCapabilityKey, + title: `Approve ${toolName}`, + description: `${toolName} requires approval before it can run.`, + payload, + resolution: null, + resolvedAt: null, + } as Partial); +} + export default class ApprovalService { + static normalizePendingActionResponseBody(body: unknown): PendingActionResponseBody | Error { + if (!isRecord(body)) { + return new Error('Request body must be a JSON object'); + } + + const unsupportedFields = Object.keys(body).filter((field) => !PENDING_ACTION_RESPONSE_FIELDS.has(field)); + if (unsupportedFields.length > 0) { + return new Error(`Unsupported pending action response fields: ${unsupportedFields.join(', ')}`); + } + + if (typeof body.approved !== 'boolean') { + return new Error('approved must be a boolean'); + } + + if (body.reason != null && typeof body.reason !== 'string') { + return new Error('reason must be a string when provided'); + } + + return { + approved: body.approved, + reason: typeof body.reason === 'string' ? body.reason : null, + }; + } + static async listPendingActions(threadUuid: string, userId: string): Promise { const thread = await AgentThreadService.getOwnedThread(threadUuid, userId); return AgentPendingAction.query() .alias('action') .leftJoinRelated('[thread, run]') .where('action.threadId', thread.id) + .where('action.status', 'pending') .select('action.*', 'thread.uuid as threadUuid', 'run.uuid as runUuid') .orderBy('action.createdAt', 'asc'); } @@ -45,117 +324,100 @@ export default class ApprovalService { message, toolPart, capabilityKey, + trx, }: { thread: AgentThread; run: AgentRun; message: AgentUIMessage; toolPart: ToolLikePart; capabilityKey: AgentCapabilityKey; - }): Promise { + trx?: Transaction; + }): Promise { const approvalId = toolPart.approval?.id; if (!approvalId) { throw new Error('Missing approval id'); } - const existing = await AgentPendingAction.query() - .where({ runId: run.id, threadId: thread.id }) - .whereRaw(`payload->>'approvalId' = ?`, [approvalId]) - .first(); - - const payload = addFileChangesToApprovalPayload({ - payload: { - approvalId, - toolCallId: toolPart.toolCallId || null, - toolName: getToolName(toolPart) || 'tool', - input: toolPart.input || null, - }, - message, - toolCallId: toolPart.toolCallId || null, - }); - - if (existing) { - return AgentPendingAction.query().patchAndFetchById(existing.id, { - status: 'pending', - payload, - } as Partial); - } + const toolCallId = toolPart.toolCallId || null; + const fileChanges = toolCallId + ? listMessageFileChanges(message).filter((change) => change.toolCallId === toolCallId) + : []; - return AgentPendingAction.query().insertAndFetch({ - threadId: thread.id, - runId: run.id, - kind: 'tool_approval', - status: 'pending', + return upsertApprovalRequestRecord({ + thread, + run, + approvalId, + toolCallId, + toolName: getToolName(toolPart) || 'tool', + input: toolPart.input, + fileChanges, capabilityKey, - title: `Approve ${payload.toolName}`, - description: `${payload.toolName} requires approval before it can run.`, - payload, - resolution: null, - resolvedAt: null, - } as Partial); + trx, + }); } - static async syncApprovalResponsesFromMessages( - threadUuid: string, - userId: string, - messages: AgentUIMessage[] - ): Promise { - const thread = await AgentThreadService.getOwnedThread(threadUuid, userId); - - for (const message of messages) { - for (const part of message.parts || []) { - if (!isToolLikePart(part) || part.state !== 'approval-responded' || !part.approval?.id) { - continue; - } - - const pending = await AgentPendingAction.query() - .alias('action') - .joinRelated('run') - .where('action.threadId', thread.id) - .where('action.status', 'pending') - .whereRaw(`action.payload->>'approvalId' = ?`, [part.approval.id]) - .modify((queryBuilder) => { - if (part.toolCallId) { - queryBuilder.whereRaw(`action.payload->>'toolCallId' = ?`, [part.toolCallId]); - } - - if (typeof message.metadata?.runId === 'string' && message.metadata.runId.trim()) { - queryBuilder.where('run.uuid', message.metadata.runId); - } - }) - .select('action.*') - .orderBy('action.createdAt', 'desc') - .first(); - - if (!pending) { - continue; - } + static async upsertApprovalRequestFromStream({ + thread, + run, + approvalId, + toolCallId, + toolName, + input, + fileChanges, + capabilityKey = 'external_mcp_write', + approvalPolicy, + toolRules, + trx, + }: { + thread: AgentThread; + run: AgentRun; + approvalId: string; + toolCallId: string; + toolName?: string | null; + input?: unknown; + fileChanges?: AgentFileChangeData[]; + capabilityKey?: AgentCapabilityKey; + approvalPolicy?: AgentApprovalPolicy; + toolRules?: AgentSessionToolRule[]; + trx?: Transaction; + }): Promise { + const resolvedToolName = toolName?.trim() || 'tool'; - const approved = part.approval.approved === true; - await AgentPendingAction.query().patchAndFetchById(pending.id, { - status: approved ? 'approved' : 'denied', - resolvedAt: new Date().toISOString(), - resolution: { - approved, - reason: part.approval.reason || null, - source: 'message', - }, - } as Partial); - } + if ( + !shouldPersistApprovalRequest({ + toolName: resolvedToolName, + fallbackCapabilityKey: capabilityKey, + approvalPolicy, + toolRules, + }) + ) { + return null; } + + return upsertApprovalRequestRecord({ + thread, + run, + approvalId, + toolCallId, + toolName: resolvedToolName, + input, + fileChanges, + capabilityKey, + trx, + }); } - static async syncApprovalRequestsFromMessages({ + static async syncApprovalRequestStateFromMessages({ thread, run, messages, capabilityKey = 'external_mcp_write', - }: { - thread: AgentThread; - run: AgentRun; - messages: AgentUIMessage[]; - capabilityKey?: AgentCapabilityKey; - }): Promise { - const created: AgentPendingAction[] = []; + approvalPolicy, + toolRules, + trx, + }: ApprovalRequestSyncOptions): Promise { + const pendingActions: AgentPendingAction[] = []; + let resolvedActionCount = 0; for (const message of messages) { if (message.role !== 'assistant') { @@ -167,27 +429,106 @@ export default class ApprovalService { continue; } + const toolName = getToolName(part) || 'tool'; + if ( + !shouldPersistApprovalRequest({ + toolName, + fallbackCapabilityKey: capabilityKey, + approvalPolicy, + toolRules, + }) + ) { + continue; + } + const action = await this.upsertApprovalRequest({ thread, run, message, toolPart: part, capabilityKey, + trx, }); - created.push(action); + if (action?.status === 'pending') { + pendingActions.push(action); + } else if (!action) { + resolvedActionCount += 1; + } } } - return created; + return { + pendingActions, + resolvedActionCount, + }; + } + + static async syncApprovalRequestsFromMessages(options: ApprovalRequestSyncOptions): Promise { + const result = await this.syncApprovalRequestStateFromMessages(options); + return result.pendingActions; } static async resolvePendingAction( actionUuid: string, userId: string, status: Extract, - resolution?: Record + resolution?: Record, + options: { + githubToken?: string | null; + } = {} ): Promise { - const action = await AgentPendingAction.query() + const resolvedAt = new Date().toISOString(); + const resolvedActionPatch = { + status, + resolvedAt, + resolution: resolution || { + approved: status === 'approved', + }, + } as Partial; + const approved = status === 'approved'; + const eventNotifications: Array<{ runUuid: string; sequence: number }> = []; + let runToEnqueue: string | null = null; + + const resumeRunIfApprovalBlocked = async (actionRun: AgentRun, runId: number, trx: Transaction) => { + const remainingPendingAction = await AgentPendingAction.query(trx).where({ runId, status: 'pending' }).first(); + + if (remainingPendingAction) { + return; + } + + if (actionRun.status === 'queued') { + runToEnqueue = actionRun.uuid; + return; + } + + if (actionRun.status !== 'waiting_for_approval') { + return; + } + + const queuedRun = await AgentRun.query(trx).patchAndFetchById(actionRun.id, { + status: 'queued', + queuedAt: resolvedAt, + executionOwner: null, + leaseExpiresAt: null, + heartbeatAt: null, + } as Partial); + const queuedSequence = await AgentRunEventService.appendStatusEventForRunInTransaction( + queuedRun, + 'run.queued', + { + status: 'queued', + error: queuedRun.error || null, + usageSummary: queuedRun.usageSummary || {}, + }, + trx + ); + if (queuedSequence) { + eventNotifications.push({ runUuid: queuedRun.uuid, sequence: queuedSequence }); + } + runToEnqueue = queuedRun.uuid; + }; + + const actionSeed = await AgentPendingAction.query() .alias('action') .joinRelated('[thread.session, run]') .where('action.uuid', actionUuid) @@ -195,27 +536,99 @@ export default class ApprovalService { .select('action.*', 'thread.uuid as threadUuid', 'run.uuid as runUuid') .first(); - if (!action) { + if (!actionSeed) { throw new Error('Pending action not found'); } - await AgentPendingAction.query().patchAndFetchById(action.id, { - status, - resolvedAt: new Date().toISOString(), - resolution: resolution || { - approved: status === 'approved', - }, - } as Partial); + const updatedAction = await AgentPendingAction.transaction(async (trx) => { + const actionRun = await AgentRun.query(trx).findById(actionSeed.runId).forUpdate(); + if (!actionRun) { + throw new Error('Agent run not found'); + } - const updatedAction = await AgentPendingAction.query() - .alias('action') - .joinRelated('[thread, run]') - .where('action.id', action.id) - .select('action.*', 'thread.uuid as threadUuid', 'run.uuid as runUuid') - .first(); + const action = await AgentPendingAction.query(trx) + .alias('action') + .joinRelated('[thread, run]') + .where('action.id', actionSeed.id) + .select('action.*', 'thread.uuid as threadUuid', 'run.uuid as runUuid') + .forUpdate() + .first(); - if (!updatedAction) { - throw new Error('Pending action not found'); + if (!action) { + throw new Error('Pending action not found'); + } + + if (action.status !== 'pending') { + await resumeRunIfApprovalBlocked(actionRun, action.runId, trx); + return action; + } + + await AgentPendingAction.query(trx).patchAndFetchById(action.id, resolvedActionPatch); + + const approvalId = + typeof action.payload?.approvalId === 'string' && action.payload.approvalId.trim() + ? action.payload.approvalId + : null; + const toolCallId = + typeof action.payload?.toolCallId === 'string' && action.payload.toolCallId.trim() + ? action.payload.toolCallId + : null; + + if (approvalId) { + const approvalEventPayload = { + actionId: action.uuid, + approvalId, + toolCallId, + approved, + reason: + resolution && typeof resolution.reason === 'string' && resolution.reason.trim() + ? String(resolution.reason) + : null, + }; + const resolvedSequence = await AgentRunEventService.appendStatusEventForRunInTransaction( + actionRun, + 'approval.resolved', + approvalEventPayload, + trx + ); + if (resolvedSequence) { + eventNotifications.push({ runUuid: actionRun.uuid, sequence: resolvedSequence }); + } + const respondedSequence = await AgentRunEventService.appendStatusEventForRunInTransaction( + actionRun, + 'approval.responded', + approvalEventPayload, + trx + ); + if (respondedSequence) { + eventNotifications.push({ runUuid: actionRun.uuid, sequence: respondedSequence }); + } + } + + await resumeRunIfApprovalBlocked(actionRun, action.runId, trx); + + const currentAction = await AgentPendingAction.query(trx) + .alias('action') + .joinRelated('[thread, run]') + .where('action.id', action.id) + .select('action.*', 'thread.uuid as threadUuid', 'run.uuid as runUuid') + .first(); + + if (!currentAction) { + throw new Error('Pending action not found'); + } + + return currentAction; + }); + + for (const notification of eventNotifications) { + await AgentRunEventService.notifyRunEventsInserted(notification.runUuid, notification.sequence); + } + + if (runToEnqueue) { + await AgentRunQueueService.enqueueRun(runToEnqueue, 'approval_resolved', { + githubToken: options.githubToken, + }); } return updatedAction; @@ -225,22 +638,27 @@ export default class ApprovalService { const enrichedAction = action as AgentPendingAction & { threadUuid?: string; runUuid?: string; + expiresAt?: string | null; }; + const payload = isRecord(action.payload) ? action.payload : {}; + const toolName = readString(payload.toolName); + const input = payload.input; return { id: action.uuid, - threadId: enrichedAction.threadUuid || String(action.threadId), - runId: enrichedAction.runUuid || String(action.runId), kind: action.kind, status: action.status, - capabilityKey: action.capabilityKey, + threadId: enrichedAction.threadUuid || String(action.threadId), + runId: enrichedAction.runUuid || String(action.runId), title: action.title, description: action.description, - payload: action.payload || {}, - resolution: action.resolution, - resolvedAt: action.resolvedAt, - createdAt: action.createdAt || null, - updatedAt: action.updatedAt || null, + requestedAt: action.createdAt || null, + expiresAt: enrichedAction.expiresAt || null, + toolName, + argumentsSummary: summarizeArguments(input), + commandPreview: getCommandPreview(input), + fileChangePreview: summarizeFileChanges(payload.fileChanges), + riskLabels: getRiskLabels(action.capabilityKey), }; } } diff --git a/src/server/services/agent/CapabilityService.ts b/src/server/services/agent/CapabilityService.ts index 11db8ff0..87e99df9 100644 --- a/src/server/services/agent/CapabilityService.ts +++ b/src/server/services/agent/CapabilityService.ts @@ -16,6 +16,7 @@ import { dynamicTool, jsonSchema, type ToolSet } from 'ai'; import AgentSession from 'server/models/AgentSession'; +import AgentSessionService from 'server/services/agentSession'; import { SESSION_WORKSPACE_GATEWAY_PORT } from 'server/lib/agentSession/podFactory'; import { McpConfigService } from 'server/services/ai/mcp/config'; import { McpClientManager } from 'server/services/ai/mcp/client'; @@ -27,16 +28,24 @@ import type { AgentSessionToolRule } from 'server/services/types/agentSessionCon import AgentPolicyService from './PolicyService'; import type { AgentApprovalMode, AgentApprovalPolicy, AgentCapabilityKey, AgentToolAuditRecord } from './types'; import type { ResolvedMcpServer } from 'server/services/ai/mcp/types'; -import { isReadOnlyWorkspaceCommand } from './sandboxExecSafety'; +import { assertSafeWorkspaceMutationCommand, isReadOnlyWorkspaceCommand } from './sandboxExecSafety'; import { buildProposedFileChanges, buildResultFileChanges, didToolResultFail } from './fileChanges'; import type { AgentFileChangeData } from './types'; +import { resolveAgentSessionDurabilityConfig } from 'server/lib/agentSession/runtimeConfig'; import { buildAgentToolKey, + CHAT_PUBLISH_HTTP_TOOL_NAME, + LIFECYCLE_BUILTIN_SERVER_SLUG, SESSION_WORKSPACE_MUTATION_TOOL_NAME, SESSION_WORKSPACE_READONLY_TOOL_NAME, + SESSION_WORKSPACE_SERVER_NAME, + SESSION_WORKSPACE_SERVER_SLUG, + buildWorkspaceMutationExecDescription, + buildWorkspaceReadonlyExecDescription, } from './toolKeys'; import { getSessionWorkspaceCatalogEntriesForRuntimeTool } from './sandboxToolCatalog'; import { SessionWorkspaceGatewayUnavailableError } from './errors'; +import AgentSandboxService from './SandboxService'; type ToolExecutionHooks = { onToolStarted?: (audit: AgentToolAuditRecord) => Promise; @@ -49,6 +58,81 @@ type SessionWorkspaceGatewayTimeouts = { executionTimeoutMs: number; }; +const WORKSPACE_EXEC_RUNTIME_TOOL_NAME = 'workspace.exec'; +const WORKSPACE_WRITE_FILE_RUNTIME_TOOL_NAME = 'workspace.write_file'; +const WORKSPACE_EDIT_FILE_RUNTIME_TOOL_NAME = 'workspace.edit_file'; +const WORKSPACE_EXEC_INPUT_SCHEMA = { + type: 'object', + required: ['command'], + additionalProperties: false, + properties: { + command: { + type: 'string', + minLength: 1, + description: 'Command to run with bash -lc', + }, + cwd: { + type: 'string', + description: 'Working directory relative to the workspace', + }, + timeoutMs: { + type: 'integer', + minimum: 1, + maximum: 120000, + description: 'Command timeout in milliseconds', + }, + }, +} as const; +const WORKSPACE_WRITE_FILE_INPUT_SCHEMA = { + type: 'object', + required: ['path', 'content'], + additionalProperties: false, + properties: { + path: { + type: 'string', + minLength: 1, + description: 'Workspace-relative file path to write', + }, + content: { + type: 'string', + description: 'Complete file content to write', + }, + }, +} as const; +const WORKSPACE_EDIT_FILE_INPUT_SCHEMA = { + type: 'object', + required: ['path', 'oldText', 'newText'], + additionalProperties: false, + properties: { + path: { + type: 'string', + minLength: 1, + description: 'Workspace-relative file path to edit', + }, + oldText: { + type: 'string', + description: 'Exact existing text to replace', + }, + newText: { + type: 'string', + description: 'Replacement text', + }, + }, +} as const; +const PUBLISH_HTTP_INPUT_SCHEMA = { + type: 'object', + required: ['port'], + additionalProperties: false, + properties: { + port: { + type: 'integer', + minimum: 1, + maximum: 65535, + description: 'Workspace HTTP port to expose through ingress', + }, + }, +} as const; + function resolvePrimaryRepo(session: AgentSession): string | undefined { const primaryRepo = (session.workspaceRepos || []).find((repo) => repo.primary)?.repo; if (primaryRepo) { @@ -79,11 +163,23 @@ function resolveSessionWorkspaceGatewayBaseUrl(session: AgentSession): string | return `http://${session.podName}.${session.namespace}.svc.cluster.local:${SESSION_WORKSPACE_GATEWAY_PORT}`; } +function isChatWorkspaceRuntimeReady(session: AgentSession): boolean { + return ( + session.sessionKind === 'chat' && + session.status === 'active' && + session.workspaceStatus === 'ready' && + Boolean(session.namespace) && + Boolean(session.podName) + ); +} + async function resolveSessionWorkspaceGatewayServer( session: AgentSession, timeouts: SessionWorkspaceGatewayTimeouts ): Promise { - const baseUrl = resolveSessionWorkspaceGatewayBaseUrl(session); + const baseUrl = + (await AgentSandboxService.resolveWorkspaceGatewayBaseUrl(session.uuid)) || + resolveSessionWorkspaceGatewayBaseUrl(session); if (!baseUrl) { return null; } @@ -137,6 +233,526 @@ function resolveSessionExecutionServer(session: AgentSession, server: ResolvedMc }; } +async function loadLatestSession(sessionUuid: string): Promise { + const session = await AgentSession.query().findOne({ uuid: sessionUuid }); + if (!session) { + throw new Error('Agent session not found'); + } + + return session; +} + +async function getFileChangePreviewChars(): Promise { + return (await resolveAgentSessionDurabilityConfig()).fileChangePreviewChars; +} + +async function ensureChatWorkspaceRuntime({ + session, + userIdentity, + requestGitHubToken, +}: { + session: AgentSession; + userIdentity: RequestUserIdentity; + requestGitHubToken?: string | null; +}): Promise { + const latestSession = await loadLatestSession(session.uuid); + if (latestSession.sessionKind !== 'chat') { + return latestSession; + } + + const ensured = await AgentSandboxService.ensureChatSandbox({ + sessionId: latestSession.uuid, + userId: userIdentity.userId, + userIdentity, + githubToken: requestGitHubToken, + }); + + return ensured.session; +} + +async function executeWorkspaceRuntimeTool({ + session, + runtimeToolName, + input, + timeoutMs, + userIdentity, + requestGitHubToken, +}: { + session: AgentSession; + runtimeToolName: string; + input: Record; + timeoutMs: number; + userIdentity: RequestUserIdentity; + requestGitHubToken?: string | null; +}) { + const runtimeSession = await ensureChatWorkspaceRuntime({ + session, + userIdentity, + requestGitHubToken, + }); + const baseUrl = + (await AgentSandboxService.resolveWorkspaceGatewayBaseUrl(runtimeSession.uuid)) || + resolveSessionWorkspaceGatewayBaseUrl(runtimeSession); + if (!baseUrl) { + throw new SessionWorkspaceGatewayUnavailableError({ + sessionId: runtimeSession.uuid, + cause: new Error('Session workspace gateway URL is not available'), + }); + } + + const client = new McpClientManager(); + try { + await client.connect({ type: 'http', url: `${baseUrl}/mcp` }, timeoutMs); + return await client.callTool(runtimeToolName, input, timeoutMs); + } catch (error) { + throw new SessionWorkspaceGatewayUnavailableError({ + sessionId: runtimeSession.uuid, + cause: error, + }); + } finally { + await client.close(); + } +} + +async function emitResultFileChanges({ + hooks, + toolCallId, + sourceTool, + input, + result, + failed, +}: { + hooks?: ToolExecutionHooks; + toolCallId?: string; + sourceTool: string; + input: Record; + result: unknown; + failed: boolean; +}) { + if (!toolCallId) { + return; + } + + const changes = buildResultFileChanges({ + toolCallId, + sourceTool, + input, + result, + failed, + previewChars: await getFileChangePreviewChars(), + }); + + for (const change of changes) { + await hooks?.onFileChange?.(change); + } +} + +function registerChatWorkspaceExecTool({ + tools, + session, + userIdentity, + approvalPolicy, + workspaceToolExecutionTimeoutMs, + requestGitHubToken, + hooks, + toolRules, + toolName, + capabilityKey, + description, + readOnly, +}: { + tools: ToolSet; + session: AgentSession; + userIdentity: RequestUserIdentity; + approvalPolicy: AgentApprovalPolicy; + workspaceToolExecutionTimeoutMs: number; + requestGitHubToken?: string | null; + hooks?: ToolExecutionHooks; + toolRules?: AgentSessionToolRule[]; + toolName: string; + capabilityKey: AgentCapabilityKey; + description: string; + readOnly: boolean; +}) { + const toolKey = buildAgentToolKey(SESSION_WORKSPACE_SERVER_SLUG, toolName); + const mode = resolveToolApprovalMode({ + toolRules, + toolKey, + capabilityMode: AgentPolicyService.modeForCapability(approvalPolicy, capabilityKey), + }); + + if (mode === 'deny') { + return; + } + + tools[toolKey] = dynamicTool({ + description, + inputSchema: jsonSchema(WORKSPACE_EXEC_INPUT_SCHEMA), + needsApproval: mode === 'require_approval', + execute: async (input, context) => { + const args = (input as Record) || {}; + const command = typeof args.command === 'string' ? args.command : ''; + if (readOnly && !isReadOnlyWorkspaceCommand(command)) { + throw new Error( + 'This command is not a safe read-only inspection command. Use the workspace exec mutation tool for state-changing, networked, or process-managing commands.' + ); + } + if (!readOnly) { + assertSafeWorkspaceMutationCommand(command); + } + + const toolCallId = context?.toolCallId; + const audit: AgentToolAuditRecord = { + source: 'mcp', + serverSlug: SESSION_WORKSPACE_SERVER_SLUG, + toolName, + toolCallId, + args, + capabilityKey, + }; + + await hooks?.onToolStarted?.(audit); + + try { + const runtimeArgs = readOnly ? args : { ...args, captureFileChanges: true }; + const result = await executeWorkspaceRuntimeTool({ + session, + runtimeToolName: WORKSPACE_EXEC_RUNTIME_TOOL_NAME, + input: runtimeArgs, + timeoutMs: workspaceToolExecutionTimeoutMs, + userIdentity, + requestGitHubToken, + }); + const failed = result.isError || didToolResultFail(result); + if (!readOnly) { + await emitResultFileChanges({ + hooks, + toolCallId, + sourceTool: toolName, + input: args, + result, + failed, + }); + } + await hooks?.onToolFinished?.({ + ...audit, + result, + status: failed ? 'failed' : 'completed', + }); + return result; + } catch (error) { + getLogger().warn({ error }, `AgentExec: chat workspace tool failed sessionId=${session.uuid} tool=${toolName}`); + await hooks?.onToolFinished?.({ + ...audit, + result: { + error: error instanceof Error ? error.message : String(error), + }, + status: 'failed', + }); + throw error; + } + }, + }); +} + +function registerChatWorkspaceFileTool({ + tools, + session, + userIdentity, + approvalPolicy, + workspaceToolExecutionTimeoutMs, + requestGitHubToken, + hooks, + toolRules, + toolName, + inputSchema, + description, +}: { + tools: ToolSet; + session: AgentSession; + userIdentity: RequestUserIdentity; + approvalPolicy: AgentApprovalPolicy; + workspaceToolExecutionTimeoutMs: number; + requestGitHubToken?: string | null; + hooks?: ToolExecutionHooks; + toolRules?: AgentSessionToolRule[]; + toolName: string; + inputSchema: Record; + description: string; +}) { + const toolKey = buildAgentToolKey(SESSION_WORKSPACE_SERVER_SLUG, toolName); + const capabilityKey: AgentCapabilityKey = 'workspace_write'; + const mode = resolveToolApprovalMode({ + toolRules, + toolKey, + capabilityMode: AgentPolicyService.modeForCapability(approvalPolicy, capabilityKey), + }); + + if (mode === 'deny') { + return; + } + + tools[toolKey] = dynamicTool({ + description, + inputSchema: jsonSchema(inputSchema), + needsApproval: mode === 'require_approval', + onInputAvailable: async ({ input, toolCallId }) => { + if (!toolCallId) { + return; + } + + const args = (input as Record) || {}; + const changes = buildProposedFileChanges({ + toolCallId, + sourceTool: toolName, + input: args, + previewChars: await getFileChangePreviewChars(), + }); + + for (const change of changes) { + await hooks?.onFileChange?.(change); + } + }, + execute: async (input, context) => { + const args = (input as Record) || {}; + const toolCallId = context?.toolCallId; + const audit: AgentToolAuditRecord = { + source: 'mcp', + serverSlug: SESSION_WORKSPACE_SERVER_SLUG, + toolName, + toolCallId, + args, + capabilityKey, + }; + + await hooks?.onToolStarted?.(audit); + + try { + const result = await executeWorkspaceRuntimeTool({ + session, + runtimeToolName: toolName, + input: args, + timeoutMs: workspaceToolExecutionTimeoutMs, + userIdentity, + requestGitHubToken, + }); + const failed = result.isError || didToolResultFail(result); + if (toolCallId) { + const changes = buildResultFileChanges({ + toolCallId, + sourceTool: toolName, + input: args, + result, + failed, + previewChars: await getFileChangePreviewChars(), + }); + + for (const change of changes) { + await hooks?.onFileChange?.(change); + } + } + await hooks?.onToolFinished?.({ + ...audit, + result, + status: failed ? 'failed' : 'completed', + }); + return result; + } catch (error) { + getLogger().warn( + { error }, + `AgentExec: chat workspace file tool failed sessionId=${session.uuid} tool=${toolName}` + ); + if (toolCallId) { + const changes = buildResultFileChanges({ + toolCallId, + sourceTool: toolName, + input: args, + result: { + error: error instanceof Error ? error.message : String(error), + }, + failed: true, + previewChars: await getFileChangePreviewChars(), + }); + + for (const change of changes) { + await hooks?.onFileChange?.(change); + } + } + await hooks?.onToolFinished?.({ + ...audit, + result: { + error: error instanceof Error ? error.message : String(error), + }, + status: 'failed', + }); + throw error; + } + }, + }); +} + +function registerChatPublishHttpTool({ + tools, + session, + approvalPolicy, + userIdentity, + requestGitHubToken, + hooks, + toolRules, +}: { + tools: ToolSet; + session: AgentSession; + approvalPolicy: AgentApprovalPolicy; + userIdentity: RequestUserIdentity; + requestGitHubToken?: string | null; + hooks?: ToolExecutionHooks; + toolRules?: AgentSessionToolRule[]; +}) { + const toolKey = buildAgentToolKey(LIFECYCLE_BUILTIN_SERVER_SLUG, CHAT_PUBLISH_HTTP_TOOL_NAME); + const capabilityKey: AgentCapabilityKey = 'deploy_k8s_mutation'; + const mode = resolveToolApprovalMode({ + toolRules, + toolKey, + capabilityMode: AgentPolicyService.modeForCapability(approvalPolicy, capabilityKey), + }); + + if (mode === 'deny') { + return; + } + + tools[toolKey] = dynamicTool({ + description: + 'Expose a running HTTP app from the chat workspace through lifecycle-managed ingress and return the reachable URL.', + inputSchema: jsonSchema(PUBLISH_HTTP_INPUT_SCHEMA), + needsApproval: mode === 'require_approval', + execute: async (input, context) => { + const args = (input as Record) || {}; + const toolCallId = context?.toolCallId; + const audit: AgentToolAuditRecord = { + source: 'mcp', + serverSlug: LIFECYCLE_BUILTIN_SERVER_SLUG, + toolName: CHAT_PUBLISH_HTTP_TOOL_NAME, + toolCallId, + args, + capabilityKey, + }; + + await hooks?.onToolStarted?.(audit); + + try { + const runtimeSession = await ensureChatWorkspaceRuntime({ + session, + userIdentity, + requestGitHubToken, + }); + const port = Number(args.port); + if (!Number.isInteger(port) || port < 1 || port > 65535) { + throw new Error('port must be an integer between 1 and 65535'); + } + + const result = await AgentSessionService.publishChatHttpPort({ + sessionId: runtimeSession.uuid, + userId: userIdentity.userId, + port, + }); + await hooks?.onToolFinished?.({ + ...audit, + result, + status: 'completed', + }); + return result; + } catch (error) { + getLogger().warn({ error }, `AgentExec: chat publish failed sessionId=${session.uuid}`); + await hooks?.onToolFinished?.({ + ...audit, + result: { + error: error instanceof Error ? error.message : String(error), + }, + status: 'failed', + }); + throw error; + } + }, + }); +} + +function registerChatWorkspaceTools({ + tools, + session, + userIdentity, + approvalPolicy, + workspaceToolExecutionTimeoutMs, + requestGitHubToken, + hooks, + toolRules, +}: { + tools: ToolSet; + session: AgentSession; + userIdentity: RequestUserIdentity; + approvalPolicy: AgentApprovalPolicy; + workspaceToolExecutionTimeoutMs: number; + requestGitHubToken?: string | null; + hooks?: ToolExecutionHooks; + toolRules?: AgentSessionToolRule[]; +}) { + registerChatWorkspaceExecTool({ + tools, + session, + userIdentity, + approvalPolicy, + workspaceToolExecutionTimeoutMs, + requestGitHubToken, + hooks, + toolRules, + toolName: SESSION_WORKSPACE_READONLY_TOOL_NAME, + capabilityKey: 'read', + description: buildWorkspaceReadonlyExecDescription(SESSION_WORKSPACE_SERVER_NAME), + readOnly: true, + }); + registerChatWorkspaceExecTool({ + tools, + session, + userIdentity, + approvalPolicy, + workspaceToolExecutionTimeoutMs, + requestGitHubToken, + hooks, + toolRules, + toolName: SESSION_WORKSPACE_MUTATION_TOOL_NAME, + capabilityKey: 'shell_exec', + description: buildWorkspaceMutationExecDescription(SESSION_WORKSPACE_SERVER_NAME), + readOnly: false, + }); + registerChatWorkspaceFileTool({ + tools, + session, + userIdentity, + approvalPolicy, + workspaceToolExecutionTimeoutMs, + requestGitHubToken, + hooks, + toolRules, + toolName: WORKSPACE_WRITE_FILE_RUNTIME_TOOL_NAME, + inputSchema: WORKSPACE_WRITE_FILE_INPUT_SCHEMA, + description: + 'Write a file in the chat workspace. Use this when the user asks to create or replace file contents. This provisions the workspace only when the tool runs.', + }); + registerChatWorkspaceFileTool({ + tools, + session, + userIdentity, + approvalPolicy, + workspaceToolExecutionTimeoutMs, + requestGitHubToken, + hooks, + toolRules, + toolName: WORKSPACE_EDIT_FILE_RUNTIME_TOOL_NAME, + inputSchema: WORKSPACE_EDIT_FILE_INPUT_SCHEMA, + description: + 'Edit a file in the chat workspace by replacing exact text. Use this for targeted file modifications. This provisions the workspace only when the tool runs.', + }); +} + function registerGenericMcpTool({ tools, session, @@ -178,6 +794,7 @@ function registerGenericMcpTool({ toolCallId, sourceTool: exposedToolName, input: args, + previewChars: await getFileChangePreviewChars(), }); for (const change of changes) { @@ -214,6 +831,7 @@ function registerGenericMcpTool({ input: args, result, failed, + previewChars: await getFileChangePreviewChars(), }); for (const change of changes) { @@ -240,6 +858,7 @@ function registerGenericMcpTool({ error: error instanceof Error ? error.message : String(error), }, failed: true, + previewChars: await getFileChangePreviewChars(), }); for (const change of changes) { @@ -297,6 +916,7 @@ export default class AgentCapabilityService { approvalPolicy, workspaceToolDiscoveryTimeoutMs, workspaceToolExecutionTimeoutMs, + requestGitHubToken, hooks, toolRules, }: { @@ -306,21 +926,45 @@ export default class AgentCapabilityService { approvalPolicy: AgentApprovalPolicy; workspaceToolDiscoveryTimeoutMs: number; workspaceToolExecutionTimeoutMs: number; + requestGitHubToken?: string | null; hooks?: ToolExecutionHooks; toolRules?: AgentSessionToolRule[]; }): Promise { const tools: ToolSet = {}; - if (!repoFullName) { - return tools; + const chatWorkspaceRuntimeReady = isChatWorkspaceRuntimeReady(session); + + if (session.sessionKind === 'chat') { + registerChatWorkspaceTools({ + tools, + session, + userIdentity, + approvalPolicy, + workspaceToolExecutionTimeoutMs, + requestGitHubToken, + hooks, + toolRules, + }); + + registerChatPublishHttpTool({ + tools, + session, + approvalPolicy, + userIdentity, + requestGitHubToken, + hooks, + toolRules, + }); } const mcpConfigService = new McpConfigService(); const [repoServers, workspaceGatewayServer] = await Promise.all([ - mcpConfigService.resolveServersForRepo(repoFullName, undefined, userIdentity), - resolveSessionWorkspaceGatewayServer(session, { - discoveryTimeoutMs: workspaceToolDiscoveryTimeoutMs, - executionTimeoutMs: workspaceToolExecutionTimeoutMs, - }), + mcpConfigService.resolveServers(repoFullName, undefined, userIdentity), + session.sessionKind === 'chat' && !chatWorkspaceRuntimeReady + ? Promise.resolve(null) + : resolveSessionWorkspaceGatewayServer(session, { + discoveryTimeoutMs: workspaceToolDiscoveryTimeoutMs, + executionTimeoutMs: workspaceToolExecutionTimeoutMs, + }), ]); const resolvedRepoServers = repoServers.flatMap((server) => { if (!usesSessionWorkspaceGatewayExecution(server.transport)) { @@ -367,7 +1011,7 @@ export default class AgentCapabilityService { } if (entry.toolName === SESSION_WORKSPACE_READONLY_TOOL_NAME) { - const inputSchema = jsonSchema(discoveredTool.inputSchema as Record); + const inputSchema = jsonSchema(WORKSPACE_EXEC_INPUT_SCHEMA); tools[entry.toolKey] = dynamicTool({ description: entry.description, @@ -428,20 +1072,23 @@ export default class AgentCapabilityService { } if (entry.toolName === SESSION_WORKSPACE_MUTATION_TOOL_NAME) { - const inputSchema = jsonSchema(discoveredTool.inputSchema as Record); + const inputSchema = jsonSchema(WORKSPACE_EXEC_INPUT_SCHEMA); tools[entry.toolKey] = dynamicTool({ description: entry.description, inputSchema, needsApproval: mode === 'require_approval', execute: async (input, context) => { + const args = (input as Record) || {}; + const command = typeof args.command === 'string' ? args.command : ''; + assertSafeWorkspaceMutationCommand(command); const toolCallId = context?.toolCallId; const audit: AgentToolAuditRecord = { source: 'mcp', serverSlug: server.slug, toolName: entry.toolName, toolCallId, - args: (input as Record) || {}, + args, capabilityKey, }; @@ -452,10 +1099,18 @@ export default class AgentCapabilityService { await client.connect(server.transport, server.timeout); const result = await client.callTool( discoveredTool.name, - (input as Record) || {}, + { ...args, captureFileChanges: true }, server.timeout ); const failed = result.isError || didToolResultFail(result); + await emitResultFileChanges({ + hooks, + toolCallId, + sourceTool: entry.toolName, + input: args, + result, + failed, + }); await hooks?.onToolFinished?.({ ...audit, result, diff --git a/src/server/services/agent/ChatSessionService.ts b/src/server/services/agent/ChatSessionService.ts new file mode 100644 index 00000000..ee575736 --- /dev/null +++ b/src/server/services/agent/ChatSessionService.ts @@ -0,0 +1,96 @@ +/** + * Copyright 2026 GoodRx, Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +import { v4 as uuid } from 'uuid'; +import type { RequestUserIdentity } from 'server/lib/get-user'; +import { getLogger } from 'server/lib/logger'; +import { EMPTY_AGENT_SESSION_SKILL_PLAN } from 'server/lib/agentSession/skillPlan'; +import AgentSession from 'server/models/AgentSession'; +import AgentThread from 'server/models/AgentThread'; +import { AgentChatStatus, AgentSessionKind, AgentWorkspaceStatus } from 'shared/constants'; +import AgentProviderRegistry from './ProviderRegistry'; +import AgentSourceService from './SourceService'; +import type { ResolvedAgentSessionWorkspaceStorageIntent } from 'server/lib/agentSession/runtimeConfig'; + +export interface CreateChatSessionOptions { + userId: string; + userIdentity?: RequestUserIdentity; + model?: string; + workspaceStorage?: ResolvedAgentSessionWorkspaceStorageIntent; +} + +export default class AgentChatSessionService { + static async createChatSession(opts: CreateChatSessionOptions): Promise { + const sessionUuid = uuid(); + const requestedModelId = opts.model?.trim() || undefined; + const providerUserIdentity = { + userId: opts.userId, + githubUsername: opts.userIdentity?.githubUsername || null, + }; + const selection = await AgentProviderRegistry.resolveSelection({ + requestedModelId, + }); + await AgentProviderRegistry.getRequiredStoredApiKey({ + provider: selection.provider, + userIdentity: providerUserIdentity, + }); + + const finalizedSession = await AgentSession.transaction(async (trx) => { + const session = await AgentSession.query(trx).insertAndFetch({ + uuid: sessionUuid, + defaultThreadId: null, + defaultModel: selection.modelId, + defaultHarness: 'lifecycle_ai_sdk', + buildUuid: null, + buildKind: null, + sessionKind: AgentSessionKind.CHAT, + userId: opts.userId, + ownerGithubUsername: opts.userIdentity?.githubUsername || null, + podName: null, + namespace: null, + pvcName: null, + model: selection.modelId, + status: 'active', + chatStatus: AgentChatStatus.READY, + workspaceStatus: AgentWorkspaceStatus.NONE, + keepAttachedServicesOnSessionNode: null, + devModeSnapshots: {}, + forwardedAgentSecretProviders: [], + workspaceRepos: [], + selectedServices: [], + skillPlan: EMPTY_AGENT_SESSION_SKILL_PLAN, + } as unknown as Partial); + + const defaultThread = await AgentThread.query(trx).insertAndFetch({ + sessionId: session.id, + title: 'Default thread', + isDefault: true, + metadata: { + sessionUuid: session.uuid, + }, + } as Partial); + + await AgentSourceService.createSessionSource(session, { trx, workspaceStorage: opts.workspaceStorage }); + + return AgentSession.query(trx).patchAndFetchById(session.id, { + defaultThreadId: defaultThread.id, + } as Partial); + }); + + getLogger().info(`Session: created chat sessionId=${sessionUuid} workspaceStatus=none`); + return finalizedSession; + } +} diff --git a/src/server/services/agent/LifecycleAiSdkHarness.ts b/src/server/services/agent/LifecycleAiSdkHarness.ts new file mode 100644 index 00000000..2f651591 --- /dev/null +++ b/src/server/services/agent/LifecycleAiSdkHarness.ts @@ -0,0 +1,795 @@ +/** + * Copyright 2026 GoodRx, Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +import { + createAgentUIStream, + createUIMessageStream, + readUIMessageStream, + safeValidateUIMessages, + type ToolSet, + type UIMessageChunk, +} from 'ai'; +import type AgentRunEvent from 'server/models/AgentRunEvent'; +import type AgentRun from 'server/models/AgentRun'; +import AgentSession from 'server/models/AgentSession'; +import AgentThread from 'server/models/AgentThread'; +import { getLogger } from 'server/lib/logger'; +import type { RequestUserIdentity } from 'server/lib/get-user'; +import AgentMessageStore from './MessageStore'; +import { buildMessageObservabilityMetadataPatch, normalizeSdkUsageSummary } from './observability'; +import ApprovalService from './ApprovalService'; +import AgentRunExecutor from './RunExecutor'; +import AgentRunService from './RunService'; +import AgentRunEventService from './RunEventService'; +import type { AgentFileChangeData, AgentUIDataParts, AgentUIMessage, AgentUIMessageMetadata } from './types'; +import { applyApprovalResponsesToFileChangeParts } from './fileChanges'; +import { AgentRunTerminalFailure } from './errors'; +import type { Transaction } from 'objection'; +import { AgentRunOwnershipLostError } from './AgentRunOwnershipLostError'; + +type AgentUiMessageChunk = UIMessageChunk; +const CONTINUATION_EVENT_PAGE_LIMIT = 500; +const CONTINUATION_EVENT_MAX_PAGES = 20; + +type ApprovalResponse = { + approved: boolean; + reason?: string | null; +}; + +type StreamFinishPayload = { + messages: AgentUIMessage[]; + finishReason?: string; + isAborted: boolean; +}; + +function createChunkReplayStream(chunks: AgentUiMessageChunk[]): ReadableStream { + return new ReadableStream({ + start(controller) { + for (const chunk of chunks) { + controller.enqueue(chunk); + } + + controller.close(); + }, + }); +} + +function extractApprovalResponses(events: AgentRunEvent[]): Map { + const responses = new Map(); + + for (const event of events) { + if (event.eventType !== 'approval.responded') { + continue; + } + + const payload = event.payload || {}; + const approvalId = typeof payload.approvalId === 'string' && payload.approvalId.trim() ? payload.approvalId : null; + if (!approvalId || typeof payload.approved !== 'boolean') { + continue; + } + + responses.set(approvalId, { + approved: payload.approved, + reason: typeof payload.reason === 'string' && payload.reason.trim() ? payload.reason : null, + }); + } + + return responses; +} + +function isToolMessagePart(value: unknown): value is Record { + if (!value || typeof value !== 'object') { + return false; + } + + const type = (value as { type?: unknown }).type; + return type === 'dynamic-tool' || (typeof type === 'string' && type.startsWith('tool-')); +} + +function readApprovalId(part: Record): string | null { + const approval = + part.approval && typeof part.approval === 'object' ? (part.approval as Record) : null; + const approvalId = typeof approval?.id === 'string' && approval.id.trim() ? approval.id : null; + + return approvalId; +} + +function buildResolvedApproval( + approvalId: string, + existingApproval: Record, + response: ApprovalResponse +): Record { + return { + ...existingApproval, + id: approvalId, + approved: response.approved, + ...(response.reason ? { reason: response.reason } : {}), + }; +} + +export function applyApprovalResponsesToToolParts( + message: AgentUIMessage, + responses: Map +): AgentUIMessage { + if (responses.size === 0) { + return message; + } + + return { + ...message, + parts: message.parts.map((rawPart) => { + if (!isToolMessagePart(rawPart)) { + return rawPart; + } + + const part = rawPart as Record; + const approval = + part.approval && typeof part.approval === 'object' ? (part.approval as Record) : {}; + const approvalId = readApprovalId(part); + const response = approvalId ? responses.get(approvalId) : null; + if (!approvalId || !response) { + return rawPart; + } + + const resolvedApproval = buildResolvedApproval(approvalId, approval, response); + if (part.state === 'approval-requested' || part.state === 'approval-responded') { + return { + ...part, + state: 'approval-responded', + approval: resolvedApproval, + } as AgentUIMessage['parts'][number]; + } + + if (response.approved && (part.state === 'output-available' || part.state === 'output-error')) { + return { + ...part, + approval: resolvedApproval, + } as AgentUIMessage['parts'][number]; + } + + if (!response.approved && part.state === 'output-denied') { + return { + ...part, + approval: resolvedApproval, + } as AgentUIMessage['parts'][number]; + } + + return rawPart; + }), + }; +} + +export function normalizeUnavailableToolPartsForAgentInput( + messages: AgentUIMessage[], + tools: ToolSet +): AgentUIMessage[] { + const availableToolNames = new Set(Object.keys(tools)); + let messagesChanged = false; + + const normalizedMessages = messages.map((message) => { + let messageChanged = false; + const parts = message.parts.map((rawPart) => { + if (!isToolMessagePart(rawPart)) { + return rawPart; + } + + const part = rawPart as Record; + const partType = typeof part.type === 'string' ? part.type : ''; + const staticToolName = partType.startsWith('tool-') ? partType.slice('tool-'.length) : null; + let nextPart = part; + let partChanged = false; + + if (staticToolName && !availableToolNames.has(staticToolName)) { + nextPart = { + ...nextPart, + type: 'dynamic-tool', + toolName: staticToolName, + }; + partChanged = true; + } + + if ( + (nextPart.state === 'output-available' || + nextPart.state === 'output-error' || + nextPart.state === 'output-denied') && + !Object.prototype.hasOwnProperty.call(nextPart, 'input') + ) { + nextPart = { + ...nextPart, + input: nextPart.rawInput, + }; + partChanged = true; + } + + if (!partChanged) { + return rawPart; + } + + messageChanged = true; + return nextPart as AgentUIMessage['parts'][number]; + }); + + if (!messageChanged) { + return message; + } + + messagesChanged = true; + return { + ...message, + parts, + }; + }); + + return messagesChanged ? normalizedMessages : messages; +} + +async function validateMessagesForAgentInput({ + runUuid, + messages, + tools, +}: { + runUuid: string; + messages: AgentUIMessage[]; + tools: ToolSet; +}): Promise { + const normalizedMessages = normalizeUnavailableToolPartsForAgentInput(messages, tools); + const validation = await safeValidateUIMessages({ + messages: normalizedMessages, + tools, + }); + + if (validation.success) { + return validation.data as AgentUIMessage[]; + } + + getLogger().warn( + { error: validation.error, runId: runUuid }, + `AgentExec: saved message validation failed runId=${runUuid}` + ); + + throw new AgentRunTerminalFailure({ + code: 'run_resume_state_invalid', + message: + 'Lifecycle could not resume this response because the saved run state is invalid. Send a new message to continue from the last saved chat state.', + details: { + reason: 'ui_message_validation', + }, + }); +} + +function sanitizeAgentStreamError(runUuid: string, error: unknown): never { + if ( + error instanceof Error && + (error.name === 'AI_TypeValidationError' || error.message.startsWith('Type validation failed')) + ) { + getLogger().warn({ error, runId: runUuid }, `AgentExec: stream validation failed runId=${runUuid}`); + + throw new AgentRunTerminalFailure({ + code: 'run_resume_state_invalid', + message: + 'Lifecycle could not resume this response because the saved run state is invalid. Send a new message to continue from the last saved chat state.', + details: { + reason: 'ui_message_validation', + }, + }); + } + + throw error; +} + +async function listRunEventsForContinuation(runUuid: string): Promise { + const events: AgentRunEvent[] = []; + let afterSequence = 0; + + for (let pageIndex = 0; pageIndex < CONTINUATION_EVENT_MAX_PAGES; pageIndex += 1) { + const page = await AgentRunEventService.listRunEventsPage(runUuid, { + afterSequence, + limit: CONTINUATION_EVENT_PAGE_LIMIT, + }); + if (!page) { + return []; + } + + events.push(...page.events); + afterSequence = page.nextSequence; + if (!page.hasMore) { + return events; + } + } + + throw new Error('Agent run event history is too large to rebuild approval continuation.'); +} + +async function rebuildAssistantMessageFromEvents(runUuid: string): Promise { + const events = await listRunEventsForContinuation(runUuid); + const approvalResponses = extractApprovalResponses(events); + if (approvalResponses.size === 0) { + return null; + } + + const chunks = AgentRunEventService.projectUiChunksFromEvents(events) as AgentUiMessageChunk[]; + let latestMessage: AgentUIMessage | null = null; + + for await (const message of readUIMessageStream({ + stream: createChunkReplayStream(chunks), + terminateOnError: false, + onError: (error) => { + getLogger().warn({ error, runId: runUuid }, `AgentExec: continuation replay skipped chunk runId=${runUuid}`); + }, + })) { + if (message.parts.length > 0) { + latestMessage = message; + } + } + + return latestMessage ? applyApprovalResponsesToToolParts(latestMessage, approvalResponses) : null; +} + +async function loadMessagesForRun( + run: AgentRun, + thread: AgentThread, + session: AgentSession +): Promise { + const storedMessages = await AgentMessageStore.listMessages(thread.uuid, session.userId); + const continuationMessage = run.startedAt ? await rebuildAssistantMessageFromEvents(run.uuid) : null; + if (!continuationMessage) { + return applyApprovalResponsesToFileChangeParts(storedMessages); + } + + return applyApprovalResponsesToFileChangeParts([ + ...storedMessages.filter((message) => message.metadata?.runId !== run.uuid), + continuationMessage, + ]); +} + +function getSessionUserIdentity(session: AgentSession): RequestUserIdentity { + const githubUsername = session.ownerGithubUsername || null; + const displayName = githubUsername || session.userId; + + return { + userId: session.userId, + githubUsername, + preferredUsername: githubUsername, + email: null, + firstName: null, + lastName: null, + displayName, + gitUserName: displayName, + gitUserEmail: githubUsername ? `${githubUsername}@users.noreply.github.com` : `${session.userId}@local.lifecycle`, + }; +} + +function createChunkStream() { + let controller: ReadableStreamDefaultController | null = null; + let closed = false; + + return { + stream: new ReadableStream({ + start(nextController) { + controller = nextController; + if (closed) { + nextController.close(); + controller = null; + } + }, + cancel() { + closed = true; + controller = null; + }, + }), + write(chunk: AgentUiMessageChunk) { + if (!closed && controller) { + controller.enqueue(chunk); + } + }, + close() { + if (closed) { + return; + } + + closed = true; + if (controller) { + controller.close(); + controller = null; + } + }, + }; +} + +function isRecord(value: unknown): value is Record { + return !!value && typeof value === 'object' && !Array.isArray(value); +} + +function readStringField(record: Record, key: string): string | null { + const value = record[key]; + return typeof value === 'string' && value.trim() ? value : null; +} + +function readFileChangeData(value: unknown): AgentFileChangeData | null { + if (!isRecord(value)) { + return null; + } + + if ( + typeof value.id !== 'string' || + typeof value.toolCallId !== 'string' || + typeof value.sourceTool !== 'string' || + typeof value.displayPath !== 'string' || + typeof value.stage !== 'string' + ) { + return null; + } + + return value as unknown as AgentFileChangeData; +} + +function createEagerApprovalRequestSync({ + thread, + run, + approvalPolicy, + toolRules, +}: { + thread: AgentThread; + run: AgentRun; + approvalPolicy: Awaited>['approvalPolicy']; + toolRules: Awaited>['toolRules']; +}) { + const draftsByToolCallId = new Map< + string, + { + toolName?: string; + input?: unknown; + fileChangesById: Map; + } + >(); + const handledApprovals = new Set(); + + const getDraft = (toolCallId: string) => { + const existing = draftsByToolCallId.get(toolCallId); + if (existing) { + return existing; + } + + const draft = { + fileChangesById: new Map(), + }; + draftsByToolCallId.set(toolCallId, draft); + return draft; + }; + + return async (chunks: AgentUiMessageChunk[], options: { trx?: Transaction; run?: AgentRun } = {}) => { + for (const chunk of chunks) { + if (!isRecord(chunk)) { + continue; + } + + const type = readStringField(chunk, 'type'); + if (!type) { + continue; + } + + if (type === 'data-file-change') { + const fileChange = readFileChangeData(chunk.data); + if (fileChange) { + getDraft(fileChange.toolCallId).fileChangesById.set(fileChange.id, fileChange); + } + continue; + } + + const toolCallId = readStringField(chunk, 'toolCallId'); + if (!toolCallId) { + continue; + } + + if (type === 'tool-input-start' || type === 'tool-input-available' || type === 'tool-input-error') { + const draft = getDraft(toolCallId); + const toolName = readStringField(chunk, 'toolName'); + if (toolName) { + draft.toolName = toolName; + } + + if (type === 'tool-input-available' && Object.prototype.hasOwnProperty.call(chunk, 'input')) { + draft.input = chunk.input; + } + continue; + } + + if (type !== 'tool-approval-request') { + continue; + } + + const approvalId = readStringField(chunk, 'approvalId'); + if (!approvalId) { + continue; + } + + const approvalKey = `${approvalId}:${toolCallId}`; + if (handledApprovals.has(approvalKey)) { + continue; + } + handledApprovals.add(approvalKey); + + const draft = getDraft(toolCallId); + try { + const action = await ApprovalService.upsertApprovalRequestFromStream({ + thread, + run: options.run || run, + approvalId, + toolCallId, + toolName: draft.toolName, + input: draft.input, + fileChanges: [...draft.fileChangesById.values()], + approvalPolicy, + toolRules, + trx: options.trx, + }); + if (action) { + chunk.actionId = action.uuid; + } + } catch (error) { + getLogger().warn( + { error, runId: run.uuid, approvalId, toolCallId }, + `AgentExec: approval request persistence failed runId=${run.uuid} approvalId=${approvalId}` + ); + throw error; + } + } + }; +} + +async function consumeStream( + runUuid: string, + executionOwner: string, + stream: ReadableStream, + beforeAppendChunks?: (chunks: AgentUiMessageChunk[], options: { trx?: Transaction; run?: AgentRun }) => Promise +): Promise { + const reader = stream.getReader(); + const batch: AgentUiMessageChunk[] = []; + + try { + let streamDone = false; + while (!streamDone) { + const { value, done } = await reader.read(); + if (done) { + streamDone = true; + continue; + } + + if (!value) { + continue; + } + + batch.push(value); + if (batch.length >= 10) { + const chunks = batch.splice(0, batch.length); + await AgentRunService.appendStreamChunksForExecutionOwner(runUuid, executionOwner, chunks, { + beforeAppendChunks: ({ trx, run }) => beforeAppendChunks?.(chunks, { trx, run }), + }); + } + } + + if (batch.length > 0) { + const chunks = batch.splice(0, batch.length); + await AgentRunService.appendStreamChunksForExecutionOwner(runUuid, executionOwner, chunks, { + beforeAppendChunks: ({ trx, run }) => beforeAppendChunks?.(chunks, { trx, run }), + }); + } + } finally { + reader.releaseLock(); + } +} + +export default class LifecycleAiSdkHarness { + static async executeRun( + run: AgentRun, + options: { + requestGitHubToken?: string | null; + dispatchAttemptId?: string; + } = {} + ): Promise { + const thread = await AgentThread.query().findById(run.threadId); + const session = await AgentSession.query().findById(run.sessionId); + if (!thread || !session) { + throw new Error('Agent run context not found'); + } + + const userIdentity = getSessionUserIdentity(session); + const normalizedMessages = await loadMessagesForRun(run, thread, session); + const fileChangeStream = createChunkStream(); + const execution = await AgentRunExecutor.execute({ + session, + thread, + userIdentity, + messages: normalizedMessages, + requestedProvider: run.resolvedProvider || run.requestedProvider || undefined, + requestedModelId: run.resolvedModel || run.requestedModel || undefined, + requestGitHubToken: options.requestGitHubToken, + existingRun: run, + dispatchAttemptId: options.dispatchAttemptId, + onFileChange: async (change) => { + fileChangeStream.write({ + type: 'data-file-change', + id: change.id, + data: change, + }); + }, + }); + const syncApprovalRequestsBeforeAppend = createEagerApprovalRequestSync({ + thread, + run: execution.run, + approvalPolicy: execution.approvalPolicy, + toolRules: execution.toolRules, + }); + const executionOwner = execution.run.executionOwner; + if (!executionOwner) { + throw new Error('Agent run execution owner is required.'); + } + + let finishContext: { + finishReason?: string; + isAborted: boolean; + } = { + finishReason: undefined, + isAborted: false, + }; + let streamFinishPayload: StreamFinishPayload | null = null; + + const agentInputMessages = await validateMessagesForAgentInput({ + runUuid: run.uuid, + messages: normalizedMessages, + tools: execution.agent.tools, + }); + + let agentUiMessageStream: ReadableStream; + try { + agentUiMessageStream = (await createAgentUIStream< + never, + typeof execution.agent.tools, + never, + AgentUIMessageMetadata + >({ + agent: execution.agent, + uiMessages: agentInputMessages, + generateMessageId: () => crypto.randomUUID(), + abortSignal: execution.abortSignal, + onFinish: async ({ finishReason, isAborted }) => { + finishContext = { + finishReason, + isAborted, + }; + fileChangeStream.close(); + }, + messageMetadata: ({ part }) => { + const eventType = (part as { type?: string }).type; + if (eventType === 'start') { + return { + sessionId: session.uuid, + threadId: thread.uuid, + runId: execution.run.uuid, + provider: execution.selection.provider, + model: execution.selection.modelId, + createdAt: new Date().toISOString(), + }; + } + + if (eventType === 'finish') { + const totalUsage = + ( + part as { + totalUsage?: { + inputTokens?: number; + outputTokens?: number; + totalTokens?: number; + reasoningTokens?: number; + cachedInputTokens?: number; + inputTokenDetails?: { + cacheReadTokens?: number; + cacheWriteTokens?: number; + noCacheTokens?: number; + }; + outputTokenDetails?: { + reasoningTokens?: number; + textTokens?: number; + }; + raw?: unknown; + }; + finishReason?: string; + rawFinishReason?: string; + } + ).totalUsage ?? undefined; + const usageSummary = totalUsage + ? normalizeSdkUsageSummary({ + usage: totalUsage, + finishReason: + typeof (part as { finishReason?: unknown }).finishReason === 'string' + ? (part as { finishReason: string }).finishReason + : undefined, + rawFinishReason: + typeof (part as { rawFinishReason?: unknown }).rawFinishReason === 'string' + ? (part as { rawFinishReason: string }).rawFinishReason + : undefined, + }) + : undefined; + + return { + sessionId: session.uuid, + threadId: thread.uuid, + runId: execution.run.uuid, + provider: execution.selection.provider, + model: execution.selection.modelId, + completedAt: new Date().toISOString(), + ...(usageSummary ? buildMessageObservabilityMetadataPatch(usageSummary) : {}), + }; + } + + return undefined; + }, + })) as ReadableStream; + } catch (error) { + sanitizeAgentStreamError(run.uuid, error); + } + + const uiMessageStream = createUIMessageStream({ + originalMessages: agentInputMessages, + generateId: () => crypto.randomUUID(), + execute: ({ writer }) => { + writer.merge(agentUiMessageStream as ReadableStream); + writer.merge(fileChangeStream.stream); + }, + onFinish: async ({ messages }) => { + streamFinishPayload = { + messages, + finishReason: finishContext.finishReason, + isAborted: finishContext.isAborted, + }; + }, + }); + + try { + await consumeStream( + run.uuid, + executionOwner, + uiMessageStream as ReadableStream, + syncApprovalRequestsBeforeAppend + ); + if (!streamFinishPayload) { + throw new Error('Agent run stream finished without final message state.'); + } + await execution.onStreamFinish(streamFinishPayload); + } catch (error) { + if (error instanceof AgentRunOwnershipLostError) { + getLogger().info( + { + runId: run.uuid, + owner: executionOwner, + currentStatus: error.currentStatus || null, + currentOwner: error.currentExecutionOwner || null, + }, + `AgentExec: ownership lost runId=${run.uuid} owner=${executionOwner}` + ); + throw error; + } + + getLogger().warn({ error, runId: run.uuid }, `AgentExec: stream consumption failed runId=${run.uuid}`); + await AgentRunService.markFailedForExecutionOwner(run.uuid, executionOwner, error, undefined, { + dispatchAttemptId: options.dispatchAttemptId, + }); + throw error; + } finally { + execution.dispose?.(); + } + } +} diff --git a/src/server/services/agent/MessageStore.ts b/src/server/services/agent/MessageStore.ts index bdf3e453..4c1f152b 100644 --- a/src/server/services/agent/MessageStore.ts +++ b/src/server/services/agent/MessageStore.ts @@ -14,85 +14,368 @@ * limitations under the License. */ -import type { PartialModelObject } from 'objection'; +import type { PartialModelObject, Transaction } from 'objection'; +import { v4 as uuid } from 'uuid'; import AgentMessage from 'server/models/AgentMessage'; import AgentRun from 'server/models/AgentRun'; +import type AgentThread from 'server/models/AgentThread'; import type { AgentUIMessage } from './types'; import AgentThreadService from './ThreadService'; +import { + getCanonicalPartsFromUiMessage, + normalizeCanonicalAgentMessageParts, + toUiMessageFromCanonicalInput, + type CanonicalAgentMessage, + type CanonicalAgentInputMessage, + type CanonicalAgentMessagePart, + type CanonicalAgentRunMessageInput, +} from './canonicalMessages'; + +const AGENT_MESSAGE_UUID_PATTERN = /^[0-9a-fA-F]{8}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{4}-[0-9a-fA-F]{12}$/; +const CLIENT_MESSAGE_ID_METADATA_KEY = 'clientMessageId'; +export const DEFAULT_AGENT_MESSAGE_PAGE_LIMIT = 50; +export const MAX_AGENT_MESSAGE_PAGE_LIMIT = 100; function toAgentUiMessage(message: AgentMessage): AgentUIMessage { - return message.uiMessage as unknown as AgentUIMessage; + return toUiMessageFromCanonicalInput( + { + id: message.uuid, + role: message.role as CanonicalAgentInputMessage['role'], + parts: normalizeCanonicalAgentMessageParts(message.parts || []), + }, + message.metadata || {} + ); +} + +function toNonEmptyAgentUiMessage(message: AgentMessage): AgentUIMessage | null { + const uiMessage = toAgentUiMessage(message); + return uiMessage.parts.length > 0 ? uiMessage : null; } function toJsonRecord(value: unknown): Record { return value as Record; } -function getUiMessageId(value: Pick | AgentMessage): string | null { - if ('uiMessage' in value) { - const messageId = (value.uiMessage as { id?: unknown } | null | undefined)?.id; - return typeof messageId === 'string' && messageId.trim() ? messageId : null; +function normalizeMessageId(value: unknown): string | null { + return typeof value === 'string' && value.trim() ? value.trim() : null; +} + +function getIncomingMessageId(message: Pick): string | null { + return normalizeMessageId(message.id); +} + +function getIncomingClientMessageId( + message: Pick +): string | null { + const explicitClientMessageId = normalizeMessageId(message.clientMessageId); + if (explicitClientMessageId) { + return explicitClientMessageId; + } + + const incomingMessageId = getIncomingMessageId(message); + return incomingMessageId && !AGENT_MESSAGE_UUID_PATTERN.test(incomingMessageId) ? incomingMessageId : null; +} + +function getStoredMessageIds(message: AgentMessage): string[] { + const ids = new Set(); + const rowUuid = normalizeMessageId(message.uuid); + const rowClientMessageId = normalizeMessageId(message.clientMessageId); + const metadataClientMessageId = normalizeMessageId(message.metadata?.[CLIENT_MESSAGE_ID_METADATA_KEY]); + + for (const id of [rowUuid, rowClientMessageId, metadataClientMessageId]) { + if (id) { + ids.add(id); + } + } + + return [...ids]; +} + +function buildStoredCanonicalMessage( + message: CanonicalAgentInputMessage, + metadata?: Record, + row?: AgentMessage +): { + uuid: string; + clientMessageId: string | null; + metadata: Record; + parts: CanonicalAgentMessagePart[]; +} { + const incomingMessageId = getIncomingMessageId(message); + const clientMessageId = + getIncomingClientMessageId(message) || + normalizeMessageId(metadata?.[CLIENT_MESSAGE_ID_METADATA_KEY]) || + normalizeMessageId(row?.clientMessageId); + const messageUuid = + row?.uuid || (incomingMessageId && AGENT_MESSAGE_UUID_PATTERN.test(incomingMessageId) ? incomingMessageId : uuid()); + const storedMetadata = { + ...(metadata || {}), + ...(clientMessageId ? { [CLIENT_MESSAGE_ID_METADATA_KEY]: clientMessageId } : {}), + }; + + return { + uuid: messageUuid, + clientMessageId, + metadata: storedMetadata, + parts: normalizeCanonicalAgentMessageParts(message.parts || []), + }; +} + +async function loadExistingMessagesForIncomingIds( + threadId: number, + messages: CanonicalAgentInputMessage[], + trx?: Transaction +): Promise { + const incomingIds = [ + ...new Set( + messages + .flatMap((message) => [getIncomingMessageId(message), getIncomingClientMessageId(message)]) + .filter((id): id is string => Boolean(id)) + ), + ]; + const uuidIds = incomingIds.filter((id) => AGENT_MESSAGE_UUID_PATTERN.test(id)); + + if (incomingIds.length === 0) { + return []; } - return typeof value.id === 'string' && value.id.trim() ? value.id : null; + return AgentMessage.query(trx) + .where({ threadId }) + .where((builder) => { + if (uuidIds.length > 0) { + builder + .whereIn('uuid', uuidIds) + .orWhereIn('clientMessageId', incomingIds) + .orWhereRaw('"metadata"->>? = ANY(?::text[])', [CLIENT_MESSAGE_ID_METADATA_KEY, incomingIds]); + return; + } + + builder + .whereIn('clientMessageId', incomingIds) + .orWhereRaw('"metadata"->>? = ANY(?::text[])', [CLIENT_MESSAGE_ID_METADATA_KEY, incomingIds]); + }); +} + +function serializeCanonicalAgentMessage( + message: AgentMessage, + threadUuid: string, + runUuid?: string | null +): CanonicalAgentMessage | null { + if (message.role !== 'user' && message.role !== 'assistant') { + return null; + } + + const parts = normalizeCanonicalAgentMessageParts(message.parts || []); + if (parts.length === 0) { + return null; + } + + const enrichedMessage = message as AgentMessage & { + runUuid?: string | null; + createdAt?: string | null; + }; + + return { + id: message.uuid, + clientMessageId: + normalizeMessageId(message.clientMessageId) || + normalizeMessageId(message.metadata?.[CLIENT_MESSAGE_ID_METADATA_KEY]), + threadId: threadUuid, + runId: runUuid || normalizeMessageId(enrichedMessage.runUuid), + role: message.role, + parts, + createdAt: enrichedMessage.createdAt || null, + }; +} + +function normalizeLimit(limit?: number): number { + if (!Number.isFinite(limit) || !limit) { + return DEFAULT_AGENT_MESSAGE_PAGE_LIMIT; + } + + return Math.min(Math.max(1, Math.trunc(limit)), MAX_AGENT_MESSAGE_PAGE_LIMIT); } export default class AgentMessageStore { + static serializeCanonicalMessage( + message: AgentMessage, + threadUuid: string, + runUuid?: string | null + ): CanonicalAgentMessage { + const serialized = serializeCanonicalAgentMessage(message, threadUuid, runUuid); + if (!serialized) { + throw new Error('Agent message is not a public canonical message'); + } + + return serialized; + } + static async listMessages(threadUuid: string, userId: string): Promise { const thread = await AgentThreadService.getOwnedThread(threadUuid, userId); const rows = await AgentMessage.query().where({ threadId: thread.id }).orderBy('createdAt', 'asc'); - return rows.map(toAgentUiMessage); + return rows.flatMap((row) => { + const message = toNonEmptyAgentUiMessage(row); + return message ? [message] : []; + }); } - static async listRunMessages(runUuid: string, userId: string): Promise { - const rows = await AgentMessage.query() + static async listCanonicalMessages( + threadUuid: string, + userId: string, + options: { + limit?: number; + beforeMessageId?: string | null; + } = {} + ): Promise<{ + thread: ReturnType; + messages: CanonicalAgentMessage[]; + pagination: { + hasMore: boolean; + nextBeforeMessageId: string | null; + }; + }> { + const { thread, session } = await AgentThreadService.getOwnedThreadWithSession(threadUuid, userId); + const limit = normalizeLimit(options.limit); + let cursor: AgentMessage | undefined; + + if (options.beforeMessageId) { + cursor = await AgentMessage.query().findOne({ + threadId: thread.id, + uuid: options.beforeMessageId, + }); + if (!cursor) { + throw new Error('Agent message cursor not found'); + } + } + + const query = AgentMessage.query() .alias('message') - .joinRelated('run.thread.session') - .where('run.uuid', runUuid) - .where('run:thread:session.userId', userId) - .select('message.*') - .orderBy('message.createdAt', 'asc'); + .leftJoin('agent_runs as run', 'message.runId', 'run.id') + .where('message.threadId', thread.id) + .whereIn('message.role', ['user', 'assistant']) + .select('message.*', 'run.uuid as runUuid') + .orderBy('message.createdAt', 'desc') + .orderBy('message.id', 'desc') + .limit(limit + 1); - return rows.map(toAgentUiMessage); + if (cursor) { + const cursorCreatedAt = (cursor as AgentMessage & { createdAt?: string | null }).createdAt; + if (!cursorCreatedAt) { + throw new Error('Agent message cursor not found'); + } + const cursorId = cursor.id; + query.where((builder) => { + builder.where('message.createdAt', '<', cursorCreatedAt).orWhere((sameTimestampBuilder) => { + sameTimestampBuilder.where('message.createdAt', '=', cursorCreatedAt).where('message.id', '<', cursorId); + }); + }); + } + + const rows = await query; + const hasMore = rows.length > limit; + const pageRows = rows.slice(0, limit).reverse(); + const messages = pageRows.flatMap((row) => { + const serialized = serializeCanonicalAgentMessage(row, thread.uuid); + return serialized ? [serialized] : []; + }); + + return { + thread: AgentThreadService.serializeThread(thread, session.uuid), + messages, + pagination: { + hasMore, + nextBeforeMessageId: hasMore && pageRows[0] ? pageRows[0].uuid : null, + }, + }; + } + + static async findCanonicalMessageByClientMessageId( + thread: Pick, + clientMessageId: string, + trx?: Transaction + ): Promise { + const normalizedClientMessageId = normalizeMessageId(clientMessageId); + if (!normalizedClientMessageId) { + return undefined; + } + + const message = await AgentMessage.query(trx) + .where({ threadId: thread.id, clientMessageId: normalizedClientMessageId }) + .orWhere((builder) => { + builder + .where({ threadId: thread.id }) + .whereRaw('"metadata"->>? = ?', [CLIENT_MESSAGE_ID_METADATA_KEY, normalizedClientMessageId]); + }) + .first(); + + return message || undefined; } - static async syncMessages( + static async insertUserMessageForRun( + thread: Pick, + run: Pick, + message: CanonicalAgentRunMessageInput, + trx?: Transaction + ): Promise { + const clientMessageId = normalizeMessageId(message.clientMessageId); + const metadata = clientMessageId ? { [CLIENT_MESSAGE_ID_METADATA_KEY]: clientMessageId } : {}; + + return AgentMessage.query(trx).insertAndFetch({ + uuid: uuid(), + threadId: thread.id, + runId: run.id, + role: 'user', + parts: normalizeCanonicalAgentMessageParts(message.parts) as unknown as Record[], + uiMessage: null, + clientMessageId, + metadata, + }); + } + + static async syncCanonicalMessages( threadUuid: string, userId: string, - messages: AgentUIMessage[], + messages: CanonicalAgentInputMessage[], runUuid?: string ): Promise { const thread = await AgentThreadService.getOwnedThread(threadUuid, userId); const run = runUuid ? await AgentRun.query().findOne({ uuid: runUuid, threadId: thread.id }) : null; const runId = run?.id ?? null; const existing = await AgentMessage.query().where({ threadId: thread.id }); - const existingByMessageId = new Map( - existing.map((message) => [getUiMessageId(message), message]).filter(([messageId]) => !!messageId) as Array< - [string, AgentMessage] - > - ); - - for (const message of messages) { - const uiMessageId = getUiMessageId(message); - if (!uiMessageId) { - continue; + const existingByMessageId = new Map(); + for (const message of existing) { + for (const messageId of getStoredMessageIds(message)) { + existingByMessageId.set(messageId, message); } + } - const row = existingByMessageId.get(uiMessageId); - const metadata = toJsonRecord(message.metadata || {}); + const nonEmptyMessages = messages.filter( + (message) => normalizeCanonicalAgentMessageParts(message.parts).length > 0 + ); + + for (const message of nonEmptyMessages) { + const incomingMessageId = getIncomingMessageId(message); + const row = incomingMessageId ? existingByMessageId.get(incomingMessageId) : undefined; + const stored = buildStoredCanonicalMessage(message, undefined, row); + const metadata = toJsonRecord(stored.metadata); const patch: PartialModelObject = { role: message.role, - uiMessage: toJsonRecord(message), + parts: stored.parts as unknown as Record[], + uiMessage: null, + clientMessageId: stored.clientMessageId, metadata, runId: message.role === 'assistant' && runId ? runId : row?.runId ?? null, }; if (!row) { const inserted = await AgentMessage.query().insert({ + uuid: stored.uuid, threadId: thread.id, ...patch, }); - existingByMessageId.set(uiMessageId, inserted); + for (const messageId of getStoredMessageIds(inserted)) { + existingByMessageId.set(messageId, inserted); + } continue; } @@ -100,6 +383,145 @@ export default class AgentMessageStore { } const reloaded = await AgentMessage.query().where({ threadId: thread.id }).orderBy('createdAt', 'asc'); - return reloaded.map(toAgentUiMessage); + return reloaded.flatMap((row) => { + const message = toNonEmptyAgentUiMessage(row); + return message ? [message] : []; + }); + } + + static async upsertCanonicalMessagesForThread( + thread: Pick, + messages: CanonicalAgentInputMessage[], + options?: { + trx?: Transaction; + runId?: number | null; + } + ): Promise { + const nonEmptyMessages = messages.filter( + (message) => normalizeCanonicalAgentMessageParts(message.parts).length > 0 + ); + if (nonEmptyMessages.length === 0) { + return; + } + + const existing = await loadExistingMessagesForIncomingIds(thread.id, nonEmptyMessages, options?.trx); + const existingByMessageId = new Map(); + for (const message of existing) { + for (const messageId of getStoredMessageIds(message)) { + existingByMessageId.set(messageId, message); + } + } + + for (const message of nonEmptyMessages) { + const incomingMessageId = getIncomingMessageId(message); + const row = incomingMessageId ? existingByMessageId.get(incomingMessageId) : undefined; + const stored = buildStoredCanonicalMessage(message, undefined, row); + const patch: PartialModelObject = { + role: message.role, + parts: stored.parts as unknown as Record[], + uiMessage: null, + clientMessageId: stored.clientMessageId, + metadata: toJsonRecord(stored.metadata), + runId: message.role === 'assistant' && options?.runId ? options.runId : row?.runId ?? null, + }; + + if (!row) { + const inserted = await AgentMessage.query(options?.trx).insert({ + uuid: stored.uuid, + threadId: thread.id, + ...patch, + }); + for (const messageId of getStoredMessageIds(inserted)) { + existingByMessageId.set(messageId, inserted); + } + continue; + } + + const updated = await AgentMessage.query(options?.trx).patchAndFetchById(row.id, patch); + for (const messageId of getStoredMessageIds(updated)) { + existingByMessageId.set(messageId, updated); + } + } + } + + static async syncCanonicalMessagesFromUiMessages( + threadUuid: string, + userId: string, + messages: AgentUIMessage[], + runUuid?: string + ): Promise { + const thread = await AgentThreadService.getOwnedThread(threadUuid, userId); + const run = runUuid ? await AgentRun.query().findOne({ uuid: runUuid, threadId: thread.id }) : null; + await this.upsertCanonicalUiMessagesForThread(thread, messages, { + runId: run?.id ?? null, + }); + + const reloaded = await AgentMessage.query().where({ threadId: thread.id }).orderBy('createdAt', 'asc'); + return reloaded.flatMap((row) => { + const message = toNonEmptyAgentUiMessage(row); + return message ? [message] : []; + }); + } + + static async upsertCanonicalUiMessagesForThread( + thread: Pick, + messages: AgentUIMessage[], + options?: { + trx?: Transaction; + runId?: number | null; + } + ): Promise { + const metadataById = new Map>(); + const canonicalMessages = messages + .filter((message) => ['user', 'assistant', 'system'].includes(message.role)) + .map((message) => { + metadataById.set(message.id, toJsonRecord(message.metadata || {})); + return { + id: message.id, + role: message.role as CanonicalAgentInputMessage['role'], + parts: getCanonicalPartsFromUiMessage(message), + }; + }) + .filter((message) => message.parts.length > 0); + + const existing = await AgentMessage.query(options?.trx).where({ threadId: thread.id }); + const existingByMessageId = new Map(); + for (const message of existing) { + for (const messageId of getStoredMessageIds(message)) { + existingByMessageId.set(messageId, message); + } + } + + for (const message of canonicalMessages) { + const incomingMessageId = getIncomingMessageId(message); + const row = incomingMessageId ? existingByMessageId.get(incomingMessageId) : undefined; + const stored = buildStoredCanonicalMessage( + message, + incomingMessageId ? metadataById.get(incomingMessageId) : undefined, + row + ); + const patch: PartialModelObject = { + role: message.role, + parts: stored.parts as unknown as Record[], + uiMessage: null, + clientMessageId: stored.clientMessageId, + metadata: toJsonRecord(stored.metadata), + runId: message.role === 'assistant' && options?.runId ? options.runId : row?.runId ?? null, + }; + + if (!row) { + const inserted = await AgentMessage.query(options?.trx).insert({ + uuid: stored.uuid, + threadId: thread.id, + ...patch, + }); + for (const messageId of getStoredMessageIds(inserted)) { + existingByMessageId.set(messageId, inserted); + } + continue; + } + + await AgentMessage.query(options?.trx).patchAndFetchById(row.id, patch); + } } } diff --git a/src/server/services/agent/ProviderRegistry.ts b/src/server/services/agent/ProviderRegistry.ts index 1165c13a..18ea471c 100644 --- a/src/server/services/agent/ProviderRegistry.ts +++ b/src/server/services/agent/ProviderRegistry.ts @@ -48,26 +48,6 @@ export class MissingAgentProviderApiKeyError extends Error { } } -function resolveRequestApiKeyForProvider( - provider: string, - requestApiKey?: string | null, - requestApiKeyProvider?: string | null -): string | null { - if (process.env.ENABLE_AUTH === 'true') { - return null; - } - - const normalizedProvider = normalizeStoredAgentProviderName(provider); - const normalizedRequestProvider = normalizeStoredAgentProviderName(requestApiKeyProvider); - - if (!normalizedProvider || normalizedRequestProvider !== normalizedProvider) { - return null; - } - - const normalized = requestApiKey?.trim(); - return normalized ? normalized : null; -} - function getProviderInstance(provider: AgentResolvedModelSelection['provider'], apiKey: string) { switch (provider) { case 'anthropic': @@ -170,13 +150,9 @@ export default class AgentProviderRegistry { static async listAvailableModelsForUser({ repoFullName, userIdentity, - requestApiKey, - requestApiKeyProvider, }: { repoFullName?: string; userIdentity: Pick; - requestApiKey?: string | null; - requestApiKeyProvider?: string | null; }): Promise { const models = await this.listAvailableModels(repoFullName); const uniqueProviders = [...new Set(models.map((model) => model.provider))]; @@ -196,24 +172,15 @@ export default class AgentProviderRegistry { }) ); - const localRequestProvider = normalizeStoredAgentProviderName(requestApiKeyProvider); - if (process.env.ENABLE_AUTH !== 'true' && requestApiKey?.trim() && localRequestProvider) { - configuredProviders.add(localRequestProvider); - } - return models.filter((model) => configuredProviders.has(model.provider)); } static async resolveCredentialEnvMap({ repoFullName, userIdentity, - requestApiKey, - requestApiKeyProvider, }: { repoFullName?: string; userIdentity: Pick; - requestApiKey?: string | null; - requestApiKeyProvider?: string | null; }): Promise> { let config; try { @@ -234,9 +201,11 @@ export default class AgentProviderRegistry { .filter((provider) => provider?.enabled !== false && typeof provider.name === 'string') .map(async (provider) => { const envVarCandidates = this.getProviderEnvVarCandidates(provider.name, provider.apiKeyEnvVar); - const apiKey = - resolveRequestApiKeyForProvider(provider.name, requestApiKey, requestApiKeyProvider) || - (await UserApiKeyService.getDecryptedKey(userIdentity.userId, provider.name, userIdentity.githubUsername)); + const apiKey = await UserApiKeyService.getDecryptedKey( + userIdentity.userId, + provider.name, + userIdentity.githubUsername + ); if (!apiKey) { return; @@ -267,17 +236,11 @@ export default class AgentProviderRegistry { static async getRequiredStoredApiKey({ provider, userIdentity, - requestApiKey, - requestApiKeyProvider, }: { provider: string; userIdentity: Pick; - requestApiKey?: string | null; - requestApiKeyProvider?: string | null; }): Promise { - const apiKey = - resolveRequestApiKeyForProvider(provider, requestApiKey, requestApiKeyProvider) || - (await UserApiKeyService.getDecryptedKey(userIdentity.userId, provider, userIdentity.githubUsername)); + const apiKey = await UserApiKeyService.getDecryptedKey(userIdentity.userId, provider, userIdentity.githubUsername); if (!apiKey) { throw new MissingAgentProviderApiKeyError(provider); @@ -289,20 +252,14 @@ export default class AgentProviderRegistry { static async createLanguageModel({ selection, userIdentity, - requestApiKey, - requestApiKeyProvider, }: { repoFullName?: string; selection: AgentResolvedModelSelection; userIdentity: RequestUserIdentity; - requestApiKey?: string | null; - requestApiKeyProvider?: string | null; }): Promise { const apiKey = await this.getRequiredStoredApiKey({ provider: selection.provider, userIdentity, - requestApiKey, - requestApiKeyProvider, }); const provider = getProviderInstance(selection.provider, apiKey); diff --git a/src/server/services/agent/RunAdmissionService.ts b/src/server/services/agent/RunAdmissionService.ts new file mode 100644 index 00000000..5d15e3ec --- /dev/null +++ b/src/server/services/agent/RunAdmissionService.ts @@ -0,0 +1,163 @@ +/** + * Copyright 2026 GoodRx, Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +import type { PartialModelObject } from 'objection'; +import AgentRun from 'server/models/AgentRun'; +import AgentSession from 'server/models/AgentSession'; +import AgentThread from 'server/models/AgentThread'; +import AgentMessage from 'server/models/AgentMessage'; +import type { AgentApprovalPolicy } from './types'; +import type { AgentRunRuntimeOptions, CanonicalAgentRunMessageInput } from './canonicalMessages'; +import AgentMessageStore from './MessageStore'; +import AgentRunEventService from './RunEventService'; +import { ActiveAgentRunError, InvalidAgentRunDefaultsError, TERMINAL_RUN_STATUSES } from './RunService'; + +function buildPolicySnapshot( + policy: AgentApprovalPolicy, + runtimeOptions?: AgentRunRuntimeOptions +): Record { + if (!runtimeOptions || Object.keys(runtimeOptions).length === 0) { + return policy as unknown as Record; + } + + return { + ...(policy as unknown as Record), + runtimeOptions, + }; +} + +export default class AgentRunAdmissionService { + static async createQueuedRunWithMessage({ + thread, + session, + policy, + message, + requestedHarness, + requestedProvider, + requestedModel, + resolvedHarness, + resolvedProvider, + resolvedModel, + sandboxRequirement, + runtimeOptions, + }: { + thread: AgentThread; + session: AgentSession; + policy: AgentApprovalPolicy; + message: CanonicalAgentRunMessageInput; + requestedHarness?: string | null; + requestedProvider?: string | null; + requestedModel?: string | null; + resolvedHarness: string; + resolvedProvider: string; + resolvedModel: string; + sandboxRequirement?: Record; + runtimeOptions?: AgentRunRuntimeOptions; + }): Promise<{ run: AgentRun; message: AgentMessage; created: boolean }> { + if (!resolvedHarness?.trim()) { + throw new InvalidAgentRunDefaultsError('Agent run harness is required.'); + } + if (!resolvedProvider?.trim()) { + throw new InvalidAgentRunDefaultsError('Agent run provider is required.'); + } + if (!resolvedModel?.trim()) { + throw new InvalidAgentRunDefaultsError('Agent run model is required.'); + } + + const now = new Date().toISOString(); + const record: PartialModelObject = { + threadId: thread.id, + sessionId: session.id, + status: 'queued', + provider: resolvedProvider, + model: resolvedModel, + requestedHarness: requestedHarness || null, + resolvedHarness, + requestedProvider: requestedProvider || null, + requestedModel: requestedModel || null, + resolvedProvider, + resolvedModel, + sandboxRequirement: sandboxRequirement || {}, + sandboxGeneration: null, + queuedAt: now, + startedAt: null, + usageSummary: {}, + policySnapshot: buildPolicySnapshot(policy, runtimeOptions), + error: null, + }; + + const admitted = await AgentRun.transaction(async (trx) => { + await AgentSession.query(trx).findById(session.id).forUpdate(); + + if (message.clientMessageId) { + const existingMessage = await AgentMessageStore.findCanonicalMessageByClientMessageId( + thread, + message.clientMessageId, + trx + ); + if (existingMessage) { + if (!existingMessage.runId) { + throw new ActiveAgentRunError(); + } + + const existingRun = await AgentRun.query(trx).findById(existingMessage.runId); + if (existingRun) { + return { + run: existingRun, + message: existingMessage, + created: false, + }; + } + } + } + + const activeRun = await AgentRun.query(trx) + .where({ sessionId: session.id }) + .whereNotIn('status', TERMINAL_RUN_STATUSES) + .orderBy('createdAt', 'desc') + .orderBy('id', 'desc') + .first(); + if (activeRun) { + throw new ActiveAgentRunError(); + } + + const queuedRun = await AgentRun.query(trx).insertAndFetch(record); + const storedMessage = await AgentMessageStore.insertUserMessageForRun(thread, queuedRun, message, trx); + await AgentThread.query(trx).patchAndFetchById(thread.id, { + lastRunAt: now, + metadata: { + ...(thread.metadata || {}), + latestRunId: queuedRun.uuid, + }, + } as Partial); + + return { + run: queuedRun, + message: storedMessage, + created: true, + }; + }); + + if (admitted.created) { + await AgentRunEventService.appendStatusEvent(admitted.run.uuid, 'run.queued', { + threadId: thread.uuid, + sessionId: session.uuid, + }); + } + + return admitted; + } +} diff --git a/src/server/services/agent/RunEventService.ts b/src/server/services/agent/RunEventService.ts new file mode 100644 index 00000000..65570e58 --- /dev/null +++ b/src/server/services/agent/RunEventService.ts @@ -0,0 +1,1171 @@ +/** + * Copyright 2026 GoodRx, Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +import AgentRun from 'server/models/AgentRun'; +import AgentRunEvent from 'server/models/AgentRunEvent'; +import { getLogger } from 'server/lib/logger'; +import { sanitizeAgentRunStreamChunks, type AgentUiMessageChunk } from './streamChunks'; +import { limitDurablePayloadRecord } from './payloadLimits'; +import { resolveAgentSessionDurabilityConfig } from 'server/lib/agentSession/runtimeConfig'; +import { AgentRunOwnershipLostError } from './AgentRunOwnershipLostError'; +import type { Transaction } from 'objection'; + +type ChunkEvent = { + eventType: string; + payload: Record; +}; + +type RunEventAppendTarget = Pick & Partial>; + +type RunEventAppendOptions = { + executionOwner?: string; + trx?: Transaction; + lockRun?: boolean; +}; + +export const DEFAULT_RUN_EVENT_PAGE_LIMIT = 100; +export const MAX_RUN_EVENT_PAGE_LIMIT = 500; +export const RUN_EVENT_STREAM_PAGE_LIMIT = 100; +export const RUN_EVENT_STREAM_POLL_INTERVAL_MS = 2000; +const AGENT_RUN_EVENT_VERSION = 1; +const RUN_EVENT_NOTIFY_CHANNEL = 'agent_run_events'; +const RUN_EVENT_TERMINAL_STATUSES = new Set(['completed', 'failed', 'cancelled']); +const RUN_EVENT_TERMINAL_EVENT_TYPES = new Set(['run.completed', 'run.failed', 'run.cancelled']); +const textEncoder = new TextEncoder(); + +type RunEventPageOptions = { + afterSequence?: number; + limit?: number; +}; + +type RunEventPageRun = Pick & { + threadUuid?: string; + sessionUuid?: string; +}; + +type RunEventPage = { + events: AgentRunEvent[]; + nextSequence: number; + hasMore: boolean; + run: { + id: string; + status: AgentRun['status']; + }; + limit: number; + maxLimit: number; +}; + +type RunEventNotification = { + runId: string; + latestSequence: number; +}; + +type SerializedRunEvent = { + id: string; + runId: string; + threadId: string; + sessionId: string; + sequence: number; + eventType: string; + version: number; + payload: Record; + createdAt: string | null; + updatedAt: string | null; +}; + +type PgListenConnection = { + on(event: 'notification', listener: (notification: { channel?: string; payload?: string }) => void): void; + on(event: 'error', listener: (error: unknown) => void): void; + query(sql: string): Promise; +}; + +type RunEventNotificationSubscriber = (notification: RunEventNotification) => void; + +const notificationSubscribers = new Map>(); +let notificationConnection: PgListenConnection | null = null; +let notificationListenPromise: Promise | null = null; + +function cloneValue(value: T): T { + return JSON.parse(JSON.stringify(value)) as T; +} + +function isRecord(value: unknown): value is Record { + return !!value && typeof value === 'object' && !Array.isArray(value); +} + +function asRecord(value: unknown): Record { + return isRecord(value) ? value : {}; +} + +function readString(value: unknown): string | undefined { + return typeof value === 'string' ? value : undefined; +} + +function readBoolean(value: unknown): boolean | undefined { + return typeof value === 'boolean' ? value : undefined; +} + +function isRunEventStreamOpen(run: Pick): boolean { + return !RUN_EVENT_TERMINAL_STATUSES.has(run.status); +} + +function sleep(durationMs: number): Promise { + return new Promise((resolve) => setTimeout(resolve, durationMs)); +} + +function encodeCanonicalSseEvent(event: SerializedRunEvent): Uint8Array { + return textEncoder.encode( + [`id: ${event.sequence}`, `event: ${event.eventType}`, `data: ${JSON.stringify(event)}`, ''].join('\n') + '\n' + ); +} + +function pickDefined(source: Record, keys: string[]): Record { + const picked: Record = {}; + + for (const key of keys) { + if (source[key] !== undefined) { + picked[key] = cloneValue(source[key]); + } + } + + return picked; +} + +function compactChunk(fields: Record): AgentUiMessageChunk { + const chunk: Record = {}; + + for (const [key, value] of Object.entries(fields)) { + if (value !== undefined) { + chunk[key] = value; + } + } + + return chunk as AgentUiMessageChunk; +} + +function parseRunEventNotification(payload: string | undefined): RunEventNotification | null { + if (!payload) { + return null; + } + + try { + const parsed = JSON.parse(payload) as Record; + const runId = readString(parsed.runId); + const latestSequence = typeof parsed.latestSequence === 'number' ? parsed.latestSequence : null; + if (!runId || latestSequence == null) { + return null; + } + + return { + runId, + latestSequence, + }; + } catch (error) { + getLogger().warn({ error }, 'AgentExec: ignored invalid run-event notification'); + return null; + } +} + +function notifySubscribers(notification: RunEventNotification): void { + const subscribers = notificationSubscribers.get(notification.runId); + if (!subscribers) { + return; + } + + for (const subscriber of [...subscribers]) { + subscriber(notification); + } +} + +function clearNotificationConnection(): void { + notificationConnection = null; + notificationListenPromise = null; +} + +function handleNotification(notification: { channel?: string; payload?: string }): void { + if (notification.channel !== RUN_EVENT_NOTIFY_CHANNEL) { + return; + } + + const parsed = parseRunEventNotification(notification.payload); + if (parsed) { + notifySubscribers(parsed); + } +} + +function handleNotificationError(error: unknown): void { + getLogger().warn({ error }, 'AgentExec: run-event notification listener failed'); + clearNotificationConnection(); +} + +function toChunkEvents(chunk: AgentUiMessageChunk): ChunkEvent[] { + const chunkRecord = chunk as unknown as Record; + + switch (chunk.type) { + case 'start': + return [ + { + eventType: 'message.created', + payload: { + messageId: chunk.messageId, + metadata: chunk.messageMetadata || {}, + }, + }, + ]; + case 'message-metadata': + return [ + { + eventType: 'message.metadata', + payload: { + metadata: cloneValue(chunk.messageMetadata || {}), + }, + }, + ]; + case 'text-start': + return [ + { + eventType: 'message.part.started', + payload: { + partType: 'text', + partId: chunk.id, + ...pickDefined(chunkRecord, ['providerMetadata']), + }, + }, + ]; + case 'text-delta': + return [ + { + eventType: 'message.delta', + payload: { + partType: 'text', + partId: chunk.id, + delta: chunk.delta, + ...pickDefined(chunkRecord, ['providerMetadata']), + }, + }, + ]; + case 'text-end': + return [ + { + eventType: 'message.part.completed', + payload: { + partType: 'text', + partId: chunk.id, + ...pickDefined(chunkRecord, ['providerMetadata']), + }, + }, + ]; + case 'reasoning-start': + return [ + { + eventType: 'message.part.started', + payload: { + partType: 'reasoning', + partId: chunk.id, + ...pickDefined(chunkRecord, ['providerMetadata']), + }, + }, + ]; + case 'reasoning-delta': + return [ + { + eventType: 'message.delta', + payload: { + partType: 'reasoning', + partId: chunk.id, + delta: chunk.delta, + ...pickDefined(chunkRecord, ['providerMetadata']), + }, + }, + ]; + case 'reasoning-end': + return [ + { + eventType: 'message.part.completed', + payload: { + partType: 'reasoning', + partId: chunk.id, + ...pickDefined(chunkRecord, ['providerMetadata']), + }, + }, + ]; + case 'tool-input-start': + return [ + { + eventType: 'tool.call.input.started', + payload: { + toolCallId: chunk.toolCallId, + toolName: chunk.toolName, + ...pickDefined(chunkRecord, ['providerExecuted', 'providerMetadata', 'dynamic', 'title']), + }, + }, + ]; + case 'tool-input-delta': + return [ + { + eventType: 'tool.call.input.delta', + payload: { + toolCallId: chunk.toolCallId, + inputTextDelta: chunk.inputTextDelta, + }, + }, + ]; + case 'tool-input-available': + case 'tool-input-error': + return [ + { + eventType: 'tool.call.started', + payload: { + toolCallId: chunk.toolCallId, + toolName: chunk.toolName, + inputStatus: chunk.type === 'tool-input-error' ? 'error' : 'available', + input: 'input' in chunk ? chunk.input : null, + errorText: 'errorText' in chunk ? chunk.errorText : null, + ...pickDefined(chunkRecord, ['providerExecuted', 'providerMetadata', 'dynamic', 'title']), + }, + }, + ]; + case 'tool-output-available': + case 'tool-output-error': + case 'tool-output-denied': + return [ + { + eventType: 'tool.call.completed', + payload: { + toolCallId: chunk.toolCallId, + output: 'output' in chunk ? chunk.output : null, + errorText: 'errorText' in chunk ? chunk.errorText : null, + status: + chunk.type === 'tool-output-available' + ? 'completed' + : chunk.type === 'tool-output-denied' + ? 'denied' + : 'failed', + ...pickDefined(chunkRecord, ['providerExecuted', 'providerMetadata', 'dynamic', 'preliminary']), + }, + }, + ]; + case 'tool-approval-request': + return [ + { + eventType: 'approval.requested', + payload: { + ...pickDefined(chunkRecord, ['actionId']), + approvalId: chunk.approvalId, + toolCallId: chunk.toolCallId, + }, + }, + ]; + case 'data-file-change': + return [ + { + eventType: 'tool.file_change', + payload: { + id: chunk.id, + data: cloneValue(chunk.data), + transient: chunk.transient, + }, + }, + ]; + case 'source-url': + return [ + { + eventType: 'message.source', + payload: { + sourceType: 'url', + sourceId: chunk.sourceId, + url: chunk.url, + ...pickDefined(chunkRecord, ['title', 'providerMetadata']), + }, + }, + ]; + case 'source-document': + return [ + { + eventType: 'message.source', + payload: { + sourceType: 'document', + sourceId: chunk.sourceId, + mediaType: chunk.mediaType, + title: chunk.title, + ...pickDefined(chunkRecord, ['filename', 'providerMetadata']), + }, + }, + ]; + case 'file': + return [ + { + eventType: 'message.file', + payload: { + url: chunk.url, + mediaType: chunk.mediaType, + ...pickDefined(chunkRecord, ['providerMetadata']), + }, + }, + ]; + case 'start-step': + return [ + { + eventType: 'run.step.started', + payload: {}, + }, + ]; + case 'finish-step': + return [ + { + eventType: 'run.step.completed', + payload: {}, + }, + ]; + case 'finish': + return [ + { + eventType: 'run.finished', + payload: { + finishReason: chunk.finishReason, + metadata: chunk.messageMetadata || {}, + }, + }, + ]; + case 'abort': + return [ + { + eventType: 'run.aborted', + payload: { + reason: chunk.reason, + }, + }, + ]; + case 'error': + return [ + { + eventType: 'run.error', + payload: { + errorText: chunk.errorText, + }, + }, + ]; + } + + return []; +} + +function chunkFromMessagePartEvent(eventType: string, payload: Record): AgentUiMessageChunk | null { + const partType = readString(payload.partType); + const partId = readString(payload.partId) || readString(payload.messageId); + if ((partType !== 'text' && partType !== 'reasoning') || !partId) { + return null; + } + + const providerMetadata = payload.providerMetadata; + + if (eventType === 'message.part.started') { + return compactChunk({ + type: partType === 'text' ? 'text-start' : 'reasoning-start', + id: partId, + providerMetadata, + }); + } + + if (eventType === 'message.delta') { + return compactChunk({ + type: partType === 'text' ? 'text-delta' : 'reasoning-delta', + id: partId, + delta: readString(payload.delta) || '', + providerMetadata, + }); + } + + if (eventType === 'message.part.completed') { + return compactChunk({ + type: partType === 'text' ? 'text-end' : 'reasoning-end', + id: partId, + providerMetadata, + }); + } + + return null; +} + +function chunkFromToolStartedEvent(payload: Record): AgentUiMessageChunk | null { + const toolCallId = readString(payload.toolCallId); + const toolName = readString(payload.toolName); + if (!toolCallId || !toolName) { + return null; + } + + const inputStatus = readString(payload.inputStatus); + return compactChunk({ + type: inputStatus === 'error' ? 'tool-input-error' : 'tool-input-available', + toolCallId, + toolName, + input: payload.input, + errorText: inputStatus === 'error' ? readString(payload.errorText) || 'Tool input failed.' : undefined, + providerExecuted: readBoolean(payload.providerExecuted), + providerMetadata: payload.providerMetadata, + dynamic: readBoolean(payload.dynamic), + title: readString(payload.title), + }); +} + +function chunkFromToolCompletedEvent(payload: Record): AgentUiMessageChunk | null { + const toolCallId = readString(payload.toolCallId); + if (!toolCallId) { + return null; + } + + const status = readString(payload.status); + if (status === 'denied') { + return compactChunk({ + type: 'tool-output-denied', + toolCallId, + }); + } + + if (status === 'failed') { + return compactChunk({ + type: 'tool-output-error', + toolCallId, + errorText: readString(payload.errorText) || 'Tool execution failed.', + providerExecuted: readBoolean(payload.providerExecuted), + providerMetadata: payload.providerMetadata, + dynamic: readBoolean(payload.dynamic), + }); + } + + return compactChunk({ + type: 'tool-output-available', + toolCallId, + output: payload.output, + providerExecuted: readBoolean(payload.providerExecuted), + providerMetadata: payload.providerMetadata, + dynamic: readBoolean(payload.dynamic), + preliminary: readBoolean(payload.preliminary), + }); +} + +function chunkFromEvent(event: AgentRunEvent): AgentUiMessageChunk | null { + const payload = asRecord(event.payload); + + switch (event.eventType) { + case 'message.created': + return compactChunk({ + type: 'start', + messageId: readString(payload.messageId), + messageMetadata: payload.metadata, + }); + case 'message.metadata': + return compactChunk({ + type: 'message-metadata', + messageMetadata: payload.metadata || {}, + }); + case 'message.part.started': + case 'message.delta': + case 'message.part.completed': + return chunkFromMessagePartEvent(event.eventType, payload); + case 'tool.call.input.started': { + const toolCallId = readString(payload.toolCallId); + const toolName = readString(payload.toolName); + if (!toolCallId || !toolName) { + return null; + } + + return compactChunk({ + type: 'tool-input-start', + toolCallId, + toolName, + providerExecuted: readBoolean(payload.providerExecuted), + providerMetadata: payload.providerMetadata, + dynamic: readBoolean(payload.dynamic), + title: readString(payload.title), + }); + } + case 'tool.call.input.delta': { + const toolCallId = readString(payload.toolCallId); + if (!toolCallId) { + return null; + } + + return compactChunk({ + type: 'tool-input-delta', + toolCallId, + inputTextDelta: readString(payload.inputTextDelta) || '', + }); + } + case 'tool.call.started': + return chunkFromToolStartedEvent(payload); + case 'tool.call.completed': + return chunkFromToolCompletedEvent(payload); + case 'approval.requested': { + const approvalId = readString(payload.approvalId); + const toolCallId = readString(payload.toolCallId); + if (!approvalId || !toolCallId) { + return null; + } + + return compactChunk({ + type: 'tool-approval-request', + actionId: readString(payload.actionId), + approvalId, + toolCallId, + }); + } + case 'tool.file_change': + if (!payload.data) { + return null; + } + + return compactChunk({ + type: 'data-file-change', + id: readString(payload.id), + data: payload.data, + transient: readBoolean(payload.transient), + }); + case 'message.source': + if (payload.sourceType === 'url') { + const sourceId = readString(payload.sourceId); + const url = readString(payload.url); + if (!sourceId || !url) { + return null; + } + + return compactChunk({ + type: 'source-url', + sourceId, + url, + title: readString(payload.title), + providerMetadata: payload.providerMetadata, + }); + } + + if (payload.sourceType === 'document') { + const sourceId = readString(payload.sourceId); + const mediaType = readString(payload.mediaType); + const title = readString(payload.title); + if (!sourceId || !mediaType || !title) { + return null; + } + + return compactChunk({ + type: 'source-document', + sourceId, + mediaType, + title, + filename: readString(payload.filename), + providerMetadata: payload.providerMetadata, + }); + } + + return null; + case 'message.file': { + const url = readString(payload.url); + const mediaType = readString(payload.mediaType); + if (!url || !mediaType) { + return null; + } + + return compactChunk({ + type: 'file', + url, + mediaType, + providerMetadata: payload.providerMetadata, + }); + } + case 'run.step.started': + return compactChunk({ type: 'start-step' }); + case 'run.step.completed': + return compactChunk({ type: 'finish-step' }); + case 'run.finished': + return compactChunk({ + type: 'finish', + finishReason: readString(payload.finishReason), + messageMetadata: payload.metadata, + }); + case 'run.aborted': + return compactChunk({ + type: 'abort', + reason: readString(payload.reason), + }); + case 'run.error': + return compactChunk({ + type: 'error', + errorText: readString(payload.errorText) || 'Agent run failed.', + }); + case 'run.failed': { + const error = asRecord(payload.error); + return compactChunk({ + type: 'error', + errorText: readString(error.message) || readString(payload.errorText) || 'Agent run failed.', + }); + } + default: + return null; + } +} + +export function normalizeRunEventPageLimit(limit?: number | null): number { + if (!Number.isFinite(limit)) { + return DEFAULT_RUN_EVENT_PAGE_LIMIT; + } + + return Math.min(Math.max(Math.floor(limit as number), 1), MAX_RUN_EVENT_PAGE_LIMIT); +} + +function normalizeRunEventAfterSequence(afterSequence?: number | null): number { + if (!Number.isFinite(afterSequence)) { + return 0; + } + + return Math.max(Math.floor(afterSequence as number), 0); +} + +export default class AgentRunEventService { + private static async ensureNotificationListener(): Promise { + if (notificationConnection) { + return; + } + + if (notificationListenPromise) { + return notificationListenPromise; + } + + notificationListenPromise = (async () => { + const knex = AgentRunEvent.knex() as unknown as { + client: { + acquireConnection(): Promise; + releaseConnection(connection: PgListenConnection): Promise; + }; + }; + const connection = await knex.client.acquireConnection(); + + try { + connection.on('notification', handleNotification); + connection.on('error', handleNotificationError); + await connection.query(`LISTEN ${RUN_EVENT_NOTIFY_CHANNEL}`); + notificationConnection = connection; + } catch (error) { + await knex.client.releaseConnection(connection); + throw error; + } + })() + .catch((error) => { + clearNotificationConnection(); + getLogger().warn({ error }, 'AgentExec: run-event notification listener unavailable'); + throw error; + }) + .finally(() => { + notificationListenPromise = null; + }); + + return notificationListenPromise; + } + + static async waitForRunEventNotification( + runUuid: string, + afterSequence: number, + timeoutMs: number + ): Promise { + if (timeoutMs <= 0) { + return false; + } + + try { + await this.ensureNotificationListener(); + } catch { + await sleep(timeoutMs); + return false; + } + + return new Promise((resolve) => { + let timeout: NodeJS.Timeout | null = null; + const subscribers = notificationSubscribers.get(runUuid) || new Set(); + + const cleanup = (notified: boolean) => { + if (timeout) { + clearTimeout(timeout); + timeout = null; + } + + subscribers.delete(subscriber); + if (subscribers.size === 0) { + notificationSubscribers.delete(runUuid); + } + + resolve(notified); + }; + + const subscriber: RunEventNotificationSubscriber = (notification) => { + if (notification.latestSequence > afterSequence) { + cleanup(true); + } + }; + + subscribers.add(subscriber); + notificationSubscribers.set(runUuid, subscribers); + timeout = setTimeout(() => cleanup(false), timeoutMs); + }); + } + + static async notifyRunEventsInserted(runUuid: string, latestSequence: number): Promise { + try { + await AgentRunEvent.knex().raw('select pg_notify(?, ?)', [ + RUN_EVENT_NOTIFY_CHANNEL, + JSON.stringify({ + runId: runUuid, + latestSequence, + }), + ]); + } catch (error) { + getLogger().warn({ error, runUuid, latestSequence }, `AgentExec: run-event notify failed runId=${runUuid}`); + } + } + + static async listRunEventsPageForRun(run: RunEventPageRun, options: RunEventPageOptions = {}): Promise { + const limit = normalizeRunEventPageLimit(options.limit); + const afterSequence = normalizeRunEventAfterSequence(options.afterSequence); + const threadUuid = run.threadUuid || String(run.threadId); + const sessionUuid = run.sessionUuid || String(run.sessionId); + const rows = await AgentRunEvent.query() + .where({ runId: run.id }) + .where('sequence', '>', afterSequence) + .orderBy('sequence', 'asc') + .orderBy('id', 'asc') + .limit(limit + 1); + const events = rows + .slice(0, limit) + .map((event) => Object.assign(event, { runUuid: run.uuid, threadUuid, sessionUuid })); + const lastEvent = events[events.length - 1]; + + return { + events, + nextSequence: lastEvent?.sequence || afterSequence, + hasMore: rows.length > limit, + run: { + id: run.uuid, + status: run.status, + }, + limit, + maxLimit: MAX_RUN_EVENT_PAGE_LIMIT, + }; + } + + static async listRunEventsPage(runUuid: string, options: RunEventPageOptions = {}): Promise { + const run = await AgentRun.query() + .alias('run') + .joinRelated('[thread, session]') + .where('run.uuid', runUuid) + .select('run.*', 'thread.uuid as threadUuid', 'session.uuid as sessionUuid') + .first(); + if (!run) { + return null; + } + + return this.listRunEventsPageForRun(run, options); + } + + static createCanonicalRunEventStream( + runUuid: string, + afterSequence: number, + options: { + pageLimit?: number; + pollIntervalMs?: number; + } = {} + ): ReadableStream { + const pageLimit = normalizeRunEventPageLimit(options.pageLimit ?? RUN_EVENT_STREAM_PAGE_LIMIT); + const pollIntervalMs = options.pollIntervalMs ?? RUN_EVENT_STREAM_POLL_INTERVAL_MS; + + return new ReadableStream({ + start: async (controller) => { + let cursor = normalizeRunEventAfterSequence(afterSequence); + let sawTerminalEvent = false; + + const drainAvailableEvents = async (): Promise => { + let hasMoreEvents = true; + + while (hasMoreEvents) { + const page = await this.listRunEventsPage(runUuid, { + afterSequence: cursor, + limit: pageLimit, + }); + if (!page) { + return false; + } + + for (const event of page.events) { + if (RUN_EVENT_TERMINAL_EVENT_TYPES.has(event.eventType)) { + sawTerminalEvent = true; + } + controller.enqueue(encodeCanonicalSseEvent(this.serializeRunEvent(event))); + } + + cursor = page.nextSequence; + hasMoreEvents = page.hasMore; + } + + return true; + }; + + let streamOpen = true; + while (streamOpen) { + if (!(await drainAvailableEvents())) { + controller.close(); + return; + } + + const currentRun = await AgentRun.query().findOne({ uuid: runUuid }); + if (!currentRun) { + controller.close(); + return; + } + + if (!isRunEventStreamOpen(currentRun)) { + if (sawTerminalEvent) { + streamOpen = false; + controller.close(); + return; + } + + if (!(await drainAvailableEvents())) { + controller.close(); + return; + } + + if (sawTerminalEvent) { + streamOpen = false; + controller.close(); + return; + } + + if (!sawTerminalEvent) { + await this.waitForRunEventNotification(runUuid, cursor, pollIntervalMs); + continue; + } + } + + await this.waitForRunEventNotification(runUuid, cursor, pollIntervalMs); + } + }, + }); + } + + private static requireExecutionOwner( + runUuid: string, + expectedExecutionOwner: string, + run: Pick + ): void { + if (run.executionOwner !== expectedExecutionOwner || RUN_EVENT_TERMINAL_STATUSES.has(run.status)) { + throw new AgentRunOwnershipLostError({ + runUuid, + expectedExecutionOwner, + currentStatus: run.status, + currentExecutionOwner: run.executionOwner, + }); + } + } + + private static async appendEventsForRun( + run: RunEventAppendTarget, + events: ChunkEvent[], + options: RunEventAppendOptions = {} + ): Promise { + if (events.length === 0) { + return null; + } + + const durability = await resolveAgentSessionDurabilityConfig(); + const append = async (trx: Transaction): Promise => { + const lockedRun = options.lockRun === false ? run : await AgentRun.query(trx).findById(run.id).forUpdate(); + if (!lockedRun) { + return null; + } + if (options.executionOwner) { + this.requireExecutionOwner( + lockedRun.uuid || String(run.id), + options.executionOwner, + lockedRun as Pick + ); + } + + const latest = await AgentRunEvent.query(trx).where({ runId: run.id }).orderBy('sequence', 'desc').first(); + let sequence = latest?.sequence || 0; + const rows = events.map((event) => { + sequence += 1; + return { + runId: run.id, + sequence, + eventType: event.eventType, + payload: limitDurablePayloadRecord(event.payload, durability), + } as Partial; + }); + + await AgentRunEvent.query(trx).insert(rows); + return sequence; + }; + + return options.trx ? append(options.trx) : AgentRun.transaction(append); + } + + static async appendEvent(runId: number, eventType: string, payload: Record): Promise { + const sequence = await this.appendEventsForRun( + { + id: runId, + }, + [ + { + eventType, + payload, + }, + ] + ); + + if (!sequence) { + throw new Error('Agent run not found'); + } + + return sequence; + } + + static async appendStatusEvent(runUuid: string, eventType: string, payload: Record): Promise { + const run = await AgentRun.query().findOne({ uuid: runUuid }); + if (!run) { + return; + } + + const sequence = await this.appendEvent(run.id, eventType, payload); + await this.notifyRunEventsInserted(run.uuid, sequence); + } + + static async appendStatusEventForRunInTransaction( + run: RunEventAppendTarget, + eventType: string, + payload: Record, + trx: Transaction + ): Promise { + return this.appendEventsForRun( + run, + [ + { + eventType, + payload, + }, + ], + { trx, lockRun: false } + ); + } + + static async appendEventsForChunks(runUuid: string, chunks: AgentUiMessageChunk[]): Promise { + if (chunks.length === 0) { + return; + } + + const run = await AgentRun.query().findOne({ uuid: runUuid }); + if (!run) { + return; + } + + const events: ChunkEvent[] = []; + + for (const chunk of chunks) { + for (const event of toChunkEvents(chunk)) { + events.push(event); + } + } + + const sequence = await this.appendEventsForRun(run, events); + if (sequence) { + await this.notifyRunEventsInserted(run.uuid, sequence); + } + } + + static async appendEventsForChunksForExecutionOwner( + runUuid: string, + executionOwner: string, + chunks: AgentUiMessageChunk[] + ): Promise { + if (chunks.length === 0) { + return; + } + + const run = await AgentRun.query().findOne({ uuid: runUuid }); + if (!run) { + return; + } + + const events: ChunkEvent[] = []; + + for (const chunk of chunks) { + for (const event of toChunkEvents(chunk)) { + events.push(event); + } + } + + const sequence = await this.appendEventsForRun(run, events, { executionOwner }); + if (sequence) { + await this.notifyRunEventsInserted(run.uuid, sequence); + } + } + + static async appendChunkEventsForRunInTransaction( + run: RunEventAppendTarget, + chunks: AgentUiMessageChunk[], + trx: Transaction + ): Promise { + if (chunks.length === 0) { + return null; + } + + const events: ChunkEvent[] = []; + for (const chunk of chunks) { + for (const event of toChunkEvents(chunk)) { + events.push(event); + } + } + + return this.appendEventsForRun(run, events, { trx, lockRun: false }); + } + + static projectUiChunksFromEvents(events: AgentRunEvent[]): AgentUiMessageChunk[] { + const chunks = events.flatMap((event) => { + const chunk = chunkFromEvent(event); + return chunk ? [chunk] : []; + }); + + return sanitizeAgentRunStreamChunks(chunks); + } + + static serializeRunEvent(event: AgentRunEvent): SerializedRunEvent { + const enrichedEvent = event as AgentRunEvent & { + runUuid?: string; + threadUuid?: string; + sessionUuid?: string; + threadId?: number; + sessionId?: number; + }; + + return { + id: event.uuid, + runId: enrichedEvent.runUuid || String(event.runId), + threadId: enrichedEvent.threadUuid || String(enrichedEvent.threadId), + sessionId: enrichedEvent.sessionUuid || String(enrichedEvent.sessionId), + sequence: event.sequence, + eventType: event.eventType, + version: AGENT_RUN_EVENT_VERSION, + payload: event.payload || {}, + createdAt: event.createdAt || null, + updatedAt: event.updatedAt || null, + }; + } +} diff --git a/src/server/services/agent/RunExecutor.ts b/src/server/services/agent/RunExecutor.ts index 72d25788..ff09add2 100644 --- a/src/server/services/agent/RunExecutor.ts +++ b/src/server/services/agent/RunExecutor.ts @@ -15,6 +15,9 @@ */ import { stepCountIs, ToolLoopAgent } from 'ai'; +import { randomBytes } from 'crypto'; +import os from 'os'; +import type AgentRun from 'server/models/AgentRun'; import type AgentSession from 'server/models/AgentSession'; import type AgentThread from 'server/models/AgentThread'; import { getLogger } from 'server/lib/logger'; @@ -28,10 +31,14 @@ import AgentCapabilityService from './CapabilityService'; import AgentMessageStore from './MessageStore'; import { AgentRunObservabilityTracker, buildMessageObservabilityMetadataPatch } from './observability'; import AgentProviderRegistry from './ProviderRegistry'; +import AgentRunQueueService from './RunQueueService'; import AgentRunService from './RunService'; import type { AgentFileChangeData, AgentUIMessage } from './types'; import { applyApprovalResponsesToFileChangeParts, buildResultFileChanges } from './fileChanges'; import { AgentRunTerminalFailure, SessionWorkspaceGatewayUnavailableError } from './errors'; +import { limitDurablePayloadValue } from './payloadLimits'; +import { resolveAgentSessionDurabilityConfig } from 'server/lib/agentSession/runtimeConfig'; +import { AgentRunOwnershipLostError } from './AgentRunOwnershipLostError'; function buildSystemPrompt(parts: Array): string | undefined { const normalized = parts.map((part) => part?.trim()).filter(Boolean) as string[]; @@ -154,16 +161,36 @@ function classifyTerminalRunFailure({ } } +function readRunMaxIterations(run?: AgentRun): number | null { + const runtimeOptions = run?.policySnapshot?.runtimeOptions; + if (!runtimeOptions || typeof runtimeOptions !== 'object' || Array.isArray(runtimeOptions)) { + return null; + } + + const maxIterations = (runtimeOptions as Record).maxIterations; + return typeof maxIterations === 'number' && Number.isInteger(maxIterations) && maxIterations > 0 + ? maxIterations + : null; +} + +function resolveHeartbeatIntervalMs(runExecutionLeaseMs: number): number { + return Math.min(Math.max(Math.floor(runExecutionLeaseMs / 3), 10_000), 60_000); +} + +function buildDirectExecutionOwner(): string { + return `direct:${os.hostname()}:${process.pid}:${randomBytes(6).toString('hex')}`; +} + export default class AgentRunExecutor { static async execute({ session, thread, userIdentity, - messages: _messages, requestedProvider, requestedModelId, - requestApiKey, - requestApiKeyProvider, + requestGitHubToken, + existingRun, + dispatchAttemptId, onFileChange, }: { session: AgentSession; @@ -172,8 +199,9 @@ export default class AgentRunExecutor { messages: AgentUIMessage[]; requestedProvider?: string; requestedModelId?: string; - requestApiKey?: string | null; - requestApiKeyProvider?: string | null; + requestGitHubToken?: string | null; + existingRun?: AgentRun; + dispatchAttemptId?: string; onFileChange?: (change: AgentFileChangeData) => Promise | void; }) { const { repoFullName, approvalPolicy } = await AgentCapabilityService.resolveSessionContext( @@ -189,8 +217,6 @@ export default class AgentRunExecutor { repoFullName, selection, userIdentity, - requestApiKey, - requestApiKeyProvider, }); const observabilityTracker = new AgentRunObservabilityTracker(); const touchSessionActivity = async () => { @@ -204,12 +230,18 @@ export default class AgentRunExecutor { } }; const effectiveSessionConfig = await AgentSessionConfigService.getInstance().getEffectiveConfig(repoFullName); + const runMaxIterations = readRunMaxIterations(existingRun); + const runControlPlaneConfig = { + ...effectiveSessionConfig, + ...(runMaxIterations ? { maxIterations: runMaxIterations } : {}), + }; const sessionPrompt = await AgentSessionService.getSessionAppendSystemPrompt( session.uuid, repoFullName, - effectiveSessionConfig.appendSystemPrompt + runControlPlaneConfig.appendSystemPrompt ); let run: Awaited> | null = null; + let heartbeatTimer: NodeJS.Timeout | null = null; const requireRun = () => { if (!run) { @@ -218,6 +250,12 @@ export default class AgentRunExecutor { return run; }; + const clearHeartbeatTimer = () => { + if (heartbeatTimer) { + clearInterval(heartbeatTimer); + heartbeatTimer = null; + } + }; try { const tools = await AgentCapabilityService.buildToolSet({ @@ -225,9 +263,10 @@ export default class AgentRunExecutor { repoFullName, userIdentity, approvalPolicy, - workspaceToolDiscoveryTimeoutMs: effectiveSessionConfig.workspaceToolDiscoveryTimeoutMs, - workspaceToolExecutionTimeoutMs: effectiveSessionConfig.workspaceToolExecutionTimeoutMs, - toolRules: effectiveSessionConfig.toolRules, + workspaceToolDiscoveryTimeoutMs: runControlPlaneConfig.workspaceToolDiscoveryTimeoutMs, + workspaceToolExecutionTimeoutMs: runControlPlaneConfig.workspaceToolExecutionTimeoutMs, + requestGitHubToken, + toolRules: runControlPlaneConfig.toolRules, hooks: { onToolStarted: async (audit) => { const activeRun = requireRun(); @@ -271,6 +310,7 @@ export default class AgentRunExecutor { } const completedAt = new Date().toISOString(); + const durability = await resolveAgentSessionDurabilityConfig(); const fileChanges = audit.toolCallId ? buildResultFileChanges({ toolCallId: audit.toolCallId, @@ -278,12 +318,13 @@ export default class AgentRunExecutor { input: audit.args, result: audit.result, failed: audit.status === 'failed', + previewChars: durability.fileChangePreviewChars, }) : []; await AgentToolExecution.query().patchAndFetchById(execution.id, { status: audit.status, result: { - value: audit.result, + value: limitDurablePayloadValue(audit.result, durability), ...(fileChanges.length > 0 ? { fileChanges } : {}), }, completedAt, @@ -296,20 +337,84 @@ export default class AgentRunExecutor { }, }); - run = await AgentRunService.createRun({ - thread, - session, - provider: selection.provider, - model: selection.modelId, - policy: approvalPolicy, - }); + if (existingRun) { + if (!existingRun.executionOwner) { + throw new Error('Agent run execution owner is required.'); + } + + run = await AgentRunService.startRunForExecutionOwner( + existingRun.uuid, + existingRun.executionOwner, + { + resolvedHarness: existingRun.requestedHarness || session.defaultHarness || 'lifecycle_ai_sdk', + provider: selection.provider, + model: selection.modelId, + sandboxGeneration: existingRun.sandboxGeneration, + }, + { dispatchAttemptId } + ); + } else { + const queuedRun = await AgentRunService.createQueuedRun({ + thread, + session, + policy: approvalPolicy, + requestedHarness: session.defaultHarness, + requestedProvider, + requestedModel: requestedModelId, + resolvedHarness: session.defaultHarness || 'lifecycle_ai_sdk', + resolvedProvider: selection.provider, + resolvedModel: selection.modelId, + }); + const executionOwner = buildDirectExecutionOwner(); + const claimedRun = await AgentRunService.claimQueuedRunForExecution(queuedRun.uuid, executionOwner); + if (!claimedRun) { + throw new Error('Agent run could not be claimed for execution.'); + } + run = await AgentRunService.startRunForExecutionOwner( + claimedRun.uuid, + executionOwner, + { + resolvedHarness: session.defaultHarness || 'lifecycle_ai_sdk', + provider: selection.provider, + model: selection.modelId, + sandboxGeneration: claimedRun.sandboxGeneration, + }, + { dispatchAttemptId } + ); + } + const activeRun = run; const controller = new AbortController(); - AgentRunService.registerAbortController(run.uuid, controller); + const activeExecutionOwner = activeRun.executionOwner || null; + AgentRunService.registerAbortController(activeRun.uuid, controller); + if (activeExecutionOwner) { + const { runExecutionLeaseMs } = await resolveAgentSessionDurabilityConfig(); + heartbeatTimer = setInterval(() => { + void AgentRunService.heartbeatRunExecution(activeRun.uuid, activeExecutionOwner).catch((error) => { + if (error instanceof AgentRunOwnershipLostError) { + getLogger().info( + { + runId: activeRun.uuid, + owner: activeExecutionOwner, + currentStatus: error.currentStatus || null, + currentOwner: error.currentExecutionOwner || null, + }, + `AgentExec: ownership lost runId=${activeRun.uuid} owner=${activeExecutionOwner}` + ); + controller.abort(error); + clearHeartbeatTimer(); + return; + } + + getLogger().warn({ error, runId: activeRun.uuid }, `AgentExec: heartbeat failed runId=${activeRun.uuid}`); + }); + }, resolveHeartbeatIntervalMs(runExecutionLeaseMs)); + heartbeatTimer.unref?.(); + } const agent = new ToolLoopAgent({ model, - instructions: buildSystemPrompt([effectiveSessionConfig.systemPrompt, sessionPrompt]), + instructions: buildSystemPrompt([runControlPlaneConfig.systemPrompt, sessionPrompt]), tools, - stopWhen: stepCountIs(effectiveSessionConfig.maxIterations), + stopWhen: stepCountIs(runControlPlaneConfig.maxIterations), onStepFinish: async (step) => { try { const usageSummary = observabilityTracker.updateFromStep({ @@ -325,14 +430,22 @@ export default class AgentRunExecutor { : undefined, }); - await AgentRunService.patchRun(run.uuid, { - usageSummary: usageSummary as Record, - }); + if (activeExecutionOwner) { + await AgentRunService.patchProgressForExecutionOwner(activeRun.uuid, activeExecutionOwner, { + usageSummary: usageSummary as Record, + }); + } await touchSessionActivity(); } catch (error) { + if (error instanceof AgentRunOwnershipLostError) { + controller.abort(error); + clearHeartbeatTimer(); + throw error; + } + getLogger().warn( - { error, runId: run.uuid }, - `AgentExec: step observability patch failed runId=${run.uuid}` + { error, runId: activeRun.uuid }, + `AgentExec: step observability patch failed runId=${activeRun.uuid}` ); } }, @@ -366,10 +479,12 @@ export default class AgentRunExecutor { }); return { - run, + run: activeRun, agent, abortSignal: controller.signal, selection, + approvalPolicy, + toolRules: runControlPlaneConfig.toolRules, onStreamFinish: async ({ messages: updatedMessages, finishReason, @@ -381,90 +496,170 @@ export default class AgentRunExecutor { }) => { const observabilitySummary = observabilityTracker.getSummary(); try { + if (!activeExecutionOwner) { + throw new Error('Agent run execution owner is required.'); + } + const messagesWithApprovalStages = applyApprovalResponsesToFileChangeParts(updatedMessages); const messagesWithObservability = applyFinalObservabilityToMessages( messagesWithApprovalStages, - run.uuid, + activeRun.uuid, buildMessageObservabilityMetadataPatch(observabilitySummary) ); - const persistedMessages = await AgentMessageStore.syncMessages( - thread.uuid, - userIdentity.userId, - messagesWithObservability, - run.uuid - ); - const pendingApprovals = await ApprovalService.syncApprovalRequestsFromMessages({ - thread, - run, - messages: persistedMessages, + const terminalFailure = classifyTerminalRunFailure({ + finishReason, + maxIterations: runControlPlaneConfig.maxIterations, }); - await touchSessionActivity(); - - const currentRun = await AgentRunService.getRunByUuid(run.uuid); - if (currentRun?.status === 'cancelled') { - return; - } + const completedAt = new Date().toISOString(); - if (pendingApprovals.length > 0) { - await AgentRunService.patchStatus(run.uuid, 'waiting_for_approval', { - usageSummary: observabilitySummary as Record, - streamState: { - finishReason: finishReason || null, - }, + const finalizedRun = await AgentRunService.finalizeRunForExecutionOwner( + activeRun.uuid, + activeExecutionOwner, + async ({ run, trx }) => { + await AgentMessageStore.upsertCanonicalUiMessagesForThread(thread, messagesWithObservability, { + trx, + runId: run.id, + }); + const approvalSync = await ApprovalService.syncApprovalRequestStateFromMessages({ + thread, + run, + messages: messagesWithObservability, + approvalPolicy, + toolRules: runControlPlaneConfig.toolRules, + trx, + }); + + if (approvalSync.pendingActions.length > 0) { + return { + status: 'waiting_for_approval', + patch: { + usageSummary: observabilitySummary as Record, + }, + }; + } + + if (approvalSync.resolvedActionCount > 0) { + return { + status: 'queued', + patch: { + queuedAt: completedAt, + usageSummary: observabilitySummary as Record, + }, + }; + } + + if (terminalFailure) { + return { + status: 'failed', + error: terminalFailure, + patch: { + completedAt, + usageSummary: observabilitySummary as Record, + }, + }; + } + + return { + status: 'completed', + patch: { + completedAt, + usageSummary: observabilitySummary as Record, + }, + }; + }, + { dispatchAttemptId } + ); + if (finalizedRun.status === 'queued') { + await AgentRunQueueService.enqueueRun(finalizedRun.uuid, 'approval_resolved', { + githubToken: requestGitHubToken, + }).catch((error) => { + getLogger().warn( + { error, runId: finalizedRun.uuid }, + `AgentExec: approval resume enqueue failed runId=${finalizedRun.uuid}` + ); }); - return; + } + await touchSessionActivity(); + } catch (error) { + if (error instanceof AgentRunOwnershipLostError) { + controller.abort(error); + throw error; } - const terminalFailure = classifyTerminalRunFailure({ - finishReason, - maxIterations: effectiveSessionConfig.maxIterations, - }); - if (terminalFailure) { - await AgentRunService.markFailed(run.uuid, terminalFailure, observabilitySummary, { - finishReason: finishReason || null, + if (activeExecutionOwner) { + await AgentRunService.markFailedForExecutionOwner( + activeRun.uuid, + activeExecutionOwner, + error, + observabilitySummary, + { dispatchAttemptId } + ).catch((runFailureError) => { + if (runFailureError instanceof AgentRunOwnershipLostError) { + throw runFailureError; + } + + getLogger().warn( + { error: runFailureError, runId: activeRun.uuid }, + `AgentExec: stream finalization failure record failed runId=${activeRun.uuid}` + ); }); - return; } - await AgentRunService.markCompleted(run.uuid, observabilitySummary, { - finishReason: finishReason || null, - }); - } catch (error) { - await AgentRunService.markFailed(run.uuid, error, observabilitySummary, { - finishReason: finishReason || null, - }).catch((runFailureError) => { - getLogger().warn( - { error: runFailureError, runId: run.uuid }, - `AgentExec: stream finalization failure record failed runId=${run.uuid}` - ); - }); - throw error; + } finally { + clearHeartbeatTimer(); + AgentRunService.clearAbortController(activeRun.uuid); } }, + dispose: () => { + clearHeartbeatTimer(); + AgentRunService.clearAbortController(activeRun.uuid); + }, }; } catch (error) { - if (error instanceof SessionWorkspaceGatewayUnavailableError) { - await AgentSessionService.markSessionRuntimeFailure(session.uuid, error).catch((runtimeFailureError) => { - getLogger().warn( - { error: runtimeFailureError, sessionId: session.uuid }, - `Session: runtime failure record failed sessionId=${session.uuid}` - ); - }); - } - - if (run) { - await AgentRunService.markFailed(run.uuid, error, observabilityTracker.getSummary()).catch( - (runFailureError) => { + try { + if (error instanceof SessionWorkspaceGatewayUnavailableError) { + await AgentSessionService.markSessionRuntimeFailure(session.uuid, error).catch((runtimeFailureError) => { getLogger().warn( - { error: runFailureError, runId: run.uuid }, - `AgentExec: run failure record failed runId=${run.uuid}` + { error: runtimeFailureError, sessionId: session.uuid }, + `Session: runtime failure record failed sessionId=${session.uuid}` ); + }); + } + + const failedRun = run || existingRun; + if (error instanceof AgentRunOwnershipLostError) { + throw error; + } + + if (failedRun) { + const failureOwner = failedRun.executionOwner || existingRun?.executionOwner || null; + if (!failureOwner) { + throw error; } - ); - } - throw error; + await AgentRunService.markFailedForExecutionOwner( + failedRun.uuid, + failureOwner, + error, + observabilityTracker.getSummary(), + { dispatchAttemptId } + ).catch((runFailureError) => { + if (runFailureError instanceof AgentRunOwnershipLostError) { + throw runFailureError; + } + + getLogger().warn( + { error: runFailureError, runId: failedRun.uuid }, + `AgentExec: run failure record failed runId=${failedRun.uuid}` + ); + }); + } + + throw error; + } finally { + clearHeartbeatTimer(); + } } } } diff --git a/src/server/services/agent/RunQueueService.ts b/src/server/services/agent/RunQueueService.ts new file mode 100644 index 00000000..f87cc5a3 --- /dev/null +++ b/src/server/services/agent/RunQueueService.ts @@ -0,0 +1,82 @@ +/** + * Copyright 2026 GoodRx, Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +import QueueManager from 'server/lib/queueManager'; +import RedisClient from 'server/lib/redisClient'; +import { encrypt } from 'server/lib/encryption'; +import { extractContextForQueue } from 'server/lib/logger'; +import { QUEUE_NAMES } from 'shared/config'; +import { randomUUID } from 'crypto'; + +export type AgentRunExecuteJob = { + runId: string; + dispatchAttemptId: string; + reason?: 'submit' | 'approval_resolved' | 'resume'; + encryptedGithubToken?: string | null; + correlationId?: string; + buildUuid?: string; + deployUuid?: string; + serviceName?: string; + sender?: string; + repo?: string; + pr?: number; + branch?: string; + sha?: string; + _ddTraceContext?: Record; +}; + +type EnqueueRunOptions = { + githubToken?: string | null; +}; + +type EnqueueRunResult = { + dispatchAttemptId: string; +}; + +export default class AgentRunQueueService { + private static queue = QueueManager.getInstance().registerQueue(QUEUE_NAMES.AGENT_RUN_EXECUTE, { + connection: RedisClient.getInstance().getConnection(), + defaultJobOptions: { + attempts: 1, + removeOnComplete: true, + removeOnFail: 100, + }, + }); + + static async enqueueRun( + runId: string, + reason: AgentRunExecuteJob['reason'] = 'submit', + options: EnqueueRunOptions = {} + ): Promise { + const githubToken = options.githubToken?.trim(); + const dispatchAttemptId = randomUUID(); + await this.queue.add( + 'execute-run', + { + runId, + dispatchAttemptId, + reason, + encryptedGithubToken: githubToken ? encrypt(githubToken) : null, + ...extractContextForQueue(), + }, + { + jobId: `agent-run:${runId}:${dispatchAttemptId}`, + } + ); + + return { dispatchAttemptId }; + } +} diff --git a/src/server/services/agent/RunService.ts b/src/server/services/agent/RunService.ts index 96001dc0..e7864f67 100644 --- a/src/server/services/agent/RunService.ts +++ b/src/server/services/agent/RunService.ts @@ -14,22 +14,42 @@ * limitations under the License. */ -import type { PartialModelObject } from 'objection'; -import type { UIMessageChunk } from 'ai'; +import type { PartialModelObject, Transaction } from 'objection'; import 'server/lib/dependencies'; import AgentRun from 'server/models/AgentRun'; import AgentThread from 'server/models/AgentThread'; -import type AgentSession from 'server/models/AgentSession'; +import AgentSession from 'server/models/AgentSession'; import type { AgentApprovalPolicy, AgentRunStatus, AgentRunUsageSummary } from './types'; -import { sanitizeAgentRunStreamChunks, sanitizeAgentRunStreamState } from './streamState'; +import type { AgentUiMessageChunk } from './streamChunks'; +import AgentRunEventService from './RunEventService'; +import { AgentRunOwnershipLostError } from './AgentRunOwnershipLostError'; +import { + DEFAULT_AGENT_SESSION_DISPATCH_RECOVERY_LIMIT, + DEFAULT_AGENT_SESSION_QUEUED_RUN_DISPATCH_STALE_MS, + DEFAULT_AGENT_SESSION_RUN_EXECUTION_LEASE_MS, + resolveAgentSessionDurabilityConfig, +} from 'server/lib/agentSession/runtimeConfig'; const activeRunControllers = new Map(); const RUN_NOT_FOUND_ERROR = 'Agent run not found'; -const TERMINAL_RUN_STATUSES: AgentRunStatus[] = ['completed', 'failed', 'cancelled']; +export const TERMINAL_RUN_STATUSES: AgentRunStatus[] = ['completed', 'failed', 'cancelled']; +export const DEFAULT_RUN_EXECUTION_LEASE_MS = DEFAULT_AGENT_SESSION_RUN_EXECUTION_LEASE_MS; +export const DEFAULT_RUN_DISPATCH_RECOVERY_LIMIT = DEFAULT_AGENT_SESSION_DISPATCH_RECOVERY_LIMIT; +export const DEFAULT_QUEUED_RUN_DISPATCH_STALE_MS = DEFAULT_AGENT_SESSION_QUEUED_RUN_DISPATCH_STALE_MS; const UUID_PATTERN = /^[0-9a-f]{8}-[0-9a-f]{4}-[1-8][0-9a-f]{3}-[89ab][0-9a-f]{3}-[0-9a-f]{12}$/i; -function cloneChunk(chunk: T): T { - return JSON.parse(JSON.stringify(chunk)) as T; +export class ActiveAgentRunError extends Error { + constructor() { + super('Wait for the current agent run to finish before starting another run.'); + this.name = 'ActiveAgentRunError'; + } +} + +export class InvalidAgentRunDefaultsError extends Error { + constructor(message: string) { + super(message); + this.name = 'InvalidAgentRunDefaultsError'; + } } function serializeRunError(error: unknown): Record { @@ -38,6 +58,18 @@ function serializeRunError(error: unknown): Record { code?: unknown; details?: unknown; }; + if (error.name === 'AI_TypeValidationError' || error.message.startsWith('Type validation failed')) { + return { + name: error.name || 'Error', + code: 'run_resume_state_invalid', + message: + 'Lifecycle could not resume this response because the saved run state is invalid. Send a new message to continue from the last saved chat state.', + details: { + reason: 'ui_message_validation', + }, + }; + } + const serialized: Record = { message: error.message, stack: error.stack || null, @@ -77,47 +109,184 @@ function isUuid(value: string): boolean { return UUID_PATTERN.test(value); } +function isLeaseExpired(leaseExpiresAt: string | null | undefined, now: Date): boolean { + if (!leaseExpiresAt) { + return false; + } + + return new Date(leaseExpiresAt).getTime() <= now.getTime(); +} + +function shouldReleaseExecution(status: AgentRunStatus): boolean { + return ( + status === 'queued' || + status === 'waiting_for_approval' || + status === 'waiting_for_input' || + TERMINAL_RUN_STATUSES.includes(status) + ); +} + +function statusEventType(status: AgentRunStatus): string { + return status === 'waiting_for_approval' + ? 'run.waiting_for_approval' + : status === 'queued' + ? 'run.queued' + : status === 'completed' + ? 'run.completed' + : status === 'failed' + ? 'run.failed' + : status === 'cancelled' + ? 'run.cancelled' + : status === 'running' || status === 'starting' + ? 'run.started' + : 'run.updated'; +} + +function runOwnershipLost(runUuid: string, expectedExecutionOwner: string, run?: AgentRun | null) { + return new AgentRunOwnershipLostError({ + runUuid, + expectedExecutionOwner, + currentStatus: run?.status, + currentExecutionOwner: run?.executionOwner, + }); +} + +type OwnerTransactionContext = { + run: AgentRun; + trx: Transaction; +}; + +type FinalizeRunForOwnerResult = { + status: AgentRunStatus; + patch?: Partial; + error?: unknown; +}; + +type OwnerStatusEventContext = { + dispatchAttemptId?: string; +}; + export default class AgentRunService { - static async createRun({ + static async createQueuedRun({ thread, session, - provider, - model, policy, + requestedHarness, + requestedProvider, + requestedModel, + resolvedHarness, + resolvedProvider, + resolvedModel, + sandboxRequirement, }: { thread: AgentThread; session: AgentSession; - provider: string; - model: string; policy: AgentApprovalPolicy; + requestedHarness?: string | null; + requestedProvider?: string | null; + requestedModel?: string | null; + resolvedHarness: string; + resolvedProvider: string; + resolvedModel: string; + sandboxRequirement?: Record; }): Promise { + if (!resolvedHarness?.trim()) { + throw new InvalidAgentRunDefaultsError('Agent run harness is required.'); + } + if (!resolvedProvider?.trim()) { + throw new InvalidAgentRunDefaultsError('Agent run provider is required.'); + } + if (!resolvedModel?.trim()) { + throw new InvalidAgentRunDefaultsError('Agent run model is required.'); + } + const now = new Date().toISOString(); const record: PartialModelObject = { threadId: thread.id, sessionId: session.id, - status: 'running', - provider, - model, + status: 'queued', + provider: resolvedProvider, + model: resolvedModel, + requestedHarness: requestedHarness || null, + resolvedHarness, + requestedProvider: requestedProvider || null, + requestedModel: requestedModel || null, + resolvedProvider, + resolvedModel, + sandboxRequirement: sandboxRequirement || {}, + sandboxGeneration: null, queuedAt: now, - startedAt: now, + startedAt: null, usageSummary: {}, policySnapshot: policy as unknown as Record, - streamState: {}, error: null, }; - const run = await AgentRun.query().insertAndFetch(record); - await AgentThread.query().patchAndFetchById(thread.id, { - lastRunAt: now, - metadata: { - ...(thread.metadata || {}), - latestRunId: run.uuid, - }, - } as Partial); + const run = await AgentRun.transaction(async (trx) => { + await AgentSession.query(trx).findById(session.id).forUpdate(); + + const activeRun = await AgentRun.query(trx) + .where({ sessionId: session.id }) + .whereNotIn('status', TERMINAL_RUN_STATUSES) + .orderBy('createdAt', 'desc') + .orderBy('id', 'desc') + .first(); + if (activeRun) { + throw new ActiveAgentRunError(); + } + + const queuedRun = await AgentRun.query(trx).insertAndFetch(record); + await AgentThread.query(trx).patchAndFetchById(thread.id, { + lastRunAt: now, + metadata: { + ...(thread.metadata || {}), + latestRunId: queuedRun.uuid, + }, + } as Partial); + + return queuedRun; + }); + + await AgentRunEventService.appendStatusEvent(run.uuid, 'run.queued', { + threadId: thread.uuid, + sessionId: session.uuid, + }); return run; } + static async createRun({ + thread, + session, + provider, + model, + policy, + }: { + thread: AgentThread; + session: AgentSession; + provider: string; + model: string; + policy: AgentApprovalPolicy; + }): Promise { + const run = await this.createQueuedRun({ + thread, + session, + policy, + requestedHarness: session.defaultHarness, + requestedProvider: provider, + requestedModel: model, + resolvedHarness: session.defaultHarness || 'lifecycle_ai_sdk', + resolvedProvider: provider, + resolvedModel: model, + }); + + return this.startRun(run.uuid, { + resolvedHarness: session.defaultHarness || 'lifecycle_ai_sdk', + provider, + model, + }); + } + static registerAbortController(runUuid: string, controller: AbortController): void { activeRunControllers.set(runUuid, controller); } @@ -135,6 +304,96 @@ export default class AgentRunService { return run || undefined; } + static async listRunsNeedingDispatch({ + limit, + now = new Date(), + queuedStaleMs, + }: { + limit?: number; + now?: Date; + queuedStaleMs?: number; + } = {}): Promise { + const durability = + limit === undefined || queuedStaleMs === undefined ? await resolveAgentSessionDurabilityConfig() : null; + const effectiveLimit = limit ?? durability!.dispatchRecoveryLimit; + const effectiveQueuedStaleMs = queuedStaleMs ?? durability!.queuedRunDispatchStaleMs; + const nowIso = now.toISOString(); + const queuedCutoff = new Date(now.getTime() - effectiveQueuedStaleMs).toISOString(); + + return AgentRun.query() + .where((builder) => { + builder.where('status', 'queued').where('queuedAt', '<', queuedCutoff); + }) + .orWhere((builder) => { + builder + .whereIn('status', ['starting', 'running']) + .whereNotNull('leaseExpiresAt') + .where('leaseExpiresAt', '<=', nowIso); + }) + .orderBy('updatedAt', 'asc') + .limit(Math.max(1, Math.floor(effectiveLimit))); + } + + static async claimQueuedRunForExecution( + runUuid: string, + executionOwner: string, + leaseMs?: number + ): Promise { + if (!isUuid(runUuid)) { + throw new Error(RUN_NOT_FOUND_ERROR); + } + + const now = new Date(); + const nowIso = now.toISOString(); + const effectiveLeaseMs = leaseMs ?? (await resolveAgentSessionDurabilityConfig()).runExecutionLeaseMs; + const leaseExpiresAt = new Date(now.getTime() + effectiveLeaseMs).toISOString(); + + return AgentRun.transaction(async (trx) => { + const run = await AgentRun.query(trx).findOne({ uuid: runUuid }).forUpdate(); + if (!run) { + throw new Error(RUN_NOT_FOUND_ERROR); + } + + const staleClaim = + (run.status === 'starting' || run.status === 'running') && isLeaseExpired(run.leaseExpiresAt, now); + if (run.status !== 'queued' && !staleClaim) { + return null; + } + + await AgentSession.query(trx).findById(run.sessionId).forUpdate(); + + return AgentRun.query(trx).patchAndFetchById(run.id, { + status: 'starting', + executionOwner, + leaseExpiresAt, + heartbeatAt: nowIso, + } as Partial); + }); + } + + static async heartbeatRunExecution(runUuid: string, executionOwner: string): Promise { + const now = new Date(); + const { runExecutionLeaseMs } = await resolveAgentSessionDurabilityConfig(); + const updatedCount = await AgentRun.query() + .where({ + uuid: runUuid, + executionOwner, + }) + .whereNotIn('status', TERMINAL_RUN_STATUSES) + .patch({ + heartbeatAt: now.toISOString(), + leaseExpiresAt: new Date(now.getTime() + runExecutionLeaseMs).toISOString(), + } as Partial); + + if (updatedCount === 0) { + const currentRun = await this.getRunByUuid(runUuid); + if (!currentRun) { + throw new Error(RUN_NOT_FOUND_ERROR); + } + throw runOwnershipLost(runUuid, executionOwner, currentRun); + } + } + static async getLatestOwnedThreadRun(threadUuid: string, userId: string): Promise { if (!isUuid(threadUuid)) { return undefined; @@ -153,6 +412,24 @@ export default class AgentRunService { return run || undefined; } + static async getLatestOwnedSessionRun(sessionUuid: string, userId: string): Promise { + if (!isUuid(sessionUuid)) { + return undefined; + } + + const run = await AgentRun.query() + .alias('run') + .joinRelated('session') + .where('session.uuid', sessionUuid) + .where('session.userId', userId) + .select('run.*', 'session.uuid as sessionUuid') + .orderBy('run.createdAt', 'desc') + .orderBy('run.id', 'desc') + .first(); + + return run || undefined; + } + static async getOwnedRun(runUuid: string, userId: string): Promise { if (!isUuid(runUuid)) { throw new Error(RUN_NOT_FOUND_ERROR); @@ -175,79 +452,451 @@ export default class AgentRunService { static async cancelRun(runUuid: string, userId: string): Promise { const run = await this.getOwnedRun(runUuid, userId); - activeRunControllers.get(runUuid)?.abort(); + activeRunControllers.get(run.uuid)?.abort(); - await AgentRun.query().patchAndFetchById(run.id, { - status: 'cancelled', + await this.patchStatus(run.uuid, 'cancelled', { cancelledAt: new Date().toISOString(), completedAt: new Date().toISOString(), } as Partial); - this.clearAbortController(runUuid); - return this.getOwnedRun(runUuid, userId); + this.clearAbortController(run.uuid); + return this.getOwnedRun(run.uuid, userId); } static isTerminalStatus(status: AgentRunStatus): boolean { return TERMINAL_RUN_STATUSES.includes(status); } + // Compatibility path for owner-independent control actions. Executor writes should use owner-aware helpers. static async patchRun(runUuid: string, patch: Partial): Promise { const run = await AgentRun.query().findOne({ uuid: runUuid }); if (!run) { throw new Error(RUN_NOT_FOUND_ERROR); } - const nextPatch = { ...patch } as Partial; - if (patch.streamState) { - nextPatch.streamState = { - ...(run.streamState || {}), - ...(patch.streamState as Record), - }; - } - - return AgentRun.query().patchAndFetchById(run.id, nextPatch); + return AgentRun.query().patchAndFetchById(run.id, patch); } static async patchStatus(runUuid: string, status: AgentRunStatus, patch?: Partial): Promise { - return this.patchRun(runUuid, { + const releaseExecution = shouldReleaseExecution(status); + const updatedRun = await this.patchRun(runUuid, { status, ...patch, + ...(releaseExecution + ? { + executionOwner: null, + leaseExpiresAt: null, + heartbeatAt: null, + } + : {}), } as Partial); + + await this.appendRunStatusEvent(runUuid, statusEventType(status), status, updatedRun); + + return updatedRun; } - static async markWaitingForApproval(runUuid: string): Promise { - return this.patchStatus(runUuid, 'waiting_for_approval'); + static async assertRunExecutionOwner(runUuid: string, executionOwner: string): Promise { + if (!isUuid(runUuid)) { + throw new Error(RUN_NOT_FOUND_ERROR); + } + + const run = await AgentRun.query().findOne({ uuid: runUuid }); + if (!run) { + throw new Error(RUN_NOT_FOUND_ERROR); + } + + this.requireRunExecutionOwner(runUuid, executionOwner, run); + return run; + } + + static async patchRunForExecutionOwner( + runUuid: string, + executionOwner: string, + patch: Partial + ): Promise { + if (!isUuid(runUuid)) { + throw new Error(RUN_NOT_FOUND_ERROR); + } + + return AgentRun.transaction((trx) => + this.patchRunForExecutionOwnerInTransaction(runUuid, executionOwner, patch, trx) + ); + } + + static async patchStatusForExecutionOwner( + runUuid: string, + executionOwner: string, + status: AgentRunStatus, + patch?: Partial, + eventContext: OwnerStatusEventContext = {} + ): Promise { + const releaseExecution = shouldReleaseExecution(status); + let latestSequence: number | null = null; + const updatedRun = await AgentRun.transaction(async (trx) => { + const run = await AgentRun.query(trx).findOne({ uuid: runUuid }).forUpdate(); + if (!run) { + throw new Error(RUN_NOT_FOUND_ERROR); + } + + this.requireRunExecutionOwner(runUuid, executionOwner, run); + + const nextRun = await AgentRun.query(trx).patchAndFetchById(run.id, { + status, + ...patch, + ...(releaseExecution + ? { + executionOwner: null, + leaseExpiresAt: null, + heartbeatAt: null, + } + : {}), + } as Partial); + + latestSequence = await AgentRunEventService.appendStatusEventForRunInTransaction( + nextRun, + statusEventType(status), + this.buildStatusEventPayload(status, nextRun, executionOwner, eventContext), + trx + ); + + return nextRun; + }); + + if (latestSequence) { + await AgentRunEventService.notifyRunEventsInserted(updatedRun.uuid, latestSequence); + } + return updatedRun; + } + + static async startRunForExecutionOwner( + runUuid: string, + executionOwner: string, + resolved: { + resolvedHarness: string; + provider: string; + model: string; + sandboxGeneration?: number | null; + }, + eventContext: OwnerStatusEventContext = {} + ): Promise { + const now = new Date().toISOString(); + return this.patchStatusForExecutionOwner( + runUuid, + executionOwner, + 'running', + { + startedAt: now, + completedAt: null, + cancelledAt: null, + error: null, + resolvedHarness: resolved.resolvedHarness, + resolvedProvider: resolved.provider, + resolvedModel: resolved.model, + provider: resolved.provider, + model: resolved.model, + sandboxGeneration: resolved.sandboxGeneration ?? null, + } as Partial, + eventContext + ); + } + + static async markWaitingForApprovalForExecutionOwner( + runUuid: string, + executionOwner: string, + usageSummary?: AgentRunUsageSummary, + eventContext: OwnerStatusEventContext = {} + ): Promise { + return this.patchStatusForExecutionOwner( + runUuid, + executionOwner, + 'waiting_for_approval', + usageSummary + ? { + usageSummary: usageSummary as Record, + } + : undefined, + eventContext + ); } - static async markCompleted( + static async markCompletedForExecutionOwner( runUuid: string, + executionOwner: string, usageSummary?: AgentRunUsageSummary, - streamState?: Record + eventContext: OwnerStatusEventContext = {} ): Promise { + const completedRun = await this.patchStatusForExecutionOwner( + runUuid, + executionOwner, + 'completed', + { + completedAt: new Date().toISOString(), + usageSummary: (usageSummary || {}) as Record, + }, + eventContext + ); this.clearAbortController(runUuid); - return this.patchStatus(runUuid, 'completed', { - completedAt: new Date().toISOString(), - usageSummary: (usageSummary || {}) as Record, - streamState: streamState || {}, - }); + return completedRun; } - static async markFailed( + static async markFailedForExecutionOwner( runUuid: string, + executionOwner: string, error: unknown, usageSummary?: AgentRunUsageSummary, - streamState?: Record + eventContext: OwnerStatusEventContext = {} ): Promise { + const failedRun = await this.patchStatusForExecutionOwner( + runUuid, + executionOwner, + 'failed', + { + completedAt: new Date().toISOString(), + usageSummary: (usageSummary || {}) as Record, + error: serializeRunError(error), + }, + eventContext + ); + this.clearAbortController(runUuid); + return failedRun; + } + + static async patchProgressForExecutionOwner( + runUuid: string, + executionOwner: string, + patch: Partial + ): Promise { + const now = new Date(); + const { runExecutionLeaseMs } = await resolveAgentSessionDurabilityConfig(); + return this.patchRunForExecutionOwner(runUuid, executionOwner, { + ...patch, + heartbeatAt: now.toISOString(), + leaseExpiresAt: new Date(now.getTime() + runExecutionLeaseMs).toISOString(), + } as Partial); + } + + static async appendStreamChunksForExecutionOwner( + runUuid: string, + executionOwner: string, + chunks: AgentUiMessageChunk[], + options: { + beforeAppendChunks?: (context: OwnerTransactionContext) => Promise; + } = {} + ): Promise { + if (!isUuid(runUuid)) { + throw new Error(RUN_NOT_FOUND_ERROR); + } + + let latestSequence: number | null = null; + const run = await AgentRun.transaction(async (trx) => { + const lockedRun = await AgentRun.query(trx).findOne({ uuid: runUuid }).forUpdate(); + if (!lockedRun) { + throw new Error(RUN_NOT_FOUND_ERROR); + } + + this.requireRunExecutionOwner(runUuid, executionOwner, lockedRun); + + if (chunks.length > 0) { + await options.beforeAppendChunks?.({ run: lockedRun, trx }); + latestSequence = await AgentRunEventService.appendChunkEventsForRunInTransaction(lockedRun, chunks, trx); + } + + return lockedRun; + }); + + if (latestSequence) { + await AgentRunEventService.notifyRunEventsInserted(run.uuid, latestSequence); + } + + return run; + } + + static async finalizeRunForExecutionOwner( + runUuid: string, + executionOwner: string, + finalize: (context: OwnerTransactionContext) => Promise, + eventContext: OwnerStatusEventContext = {} + ): Promise { + if (!isUuid(runUuid)) { + throw new Error(RUN_NOT_FOUND_ERROR); + } + + let latestSequence: number | null = null; + const updatedRun = await AgentRun.transaction(async (trx) => { + const run = await AgentRun.query(trx).findOne({ uuid: runUuid }).forUpdate(); + if (!run) { + throw new Error(RUN_NOT_FOUND_ERROR); + } + + this.requireRunExecutionOwner(runUuid, executionOwner, run); + + const result = await finalize({ run, trx }); + const releaseExecution = shouldReleaseExecution(result.status); + const nextRun = await AgentRun.query(trx).patchAndFetchById(run.id, { + status: result.status, + ...(result.patch || {}), + ...(result.status === 'failed' && result.error !== undefined ? { error: serializeRunError(result.error) } : {}), + ...(releaseExecution + ? { + executionOwner: null, + leaseExpiresAt: null, + heartbeatAt: null, + } + : {}), + } as Partial); + + latestSequence = await AgentRunEventService.appendStatusEventForRunInTransaction( + nextRun, + statusEventType(result.status), + this.buildStatusEventPayload(result.status, nextRun, executionOwner, eventContext), + trx + ); + + return nextRun; + }); + + if (latestSequence) { + await AgentRunEventService.notifyRunEventsInserted(updatedRun.uuid, latestSequence); + } + + if (TERMINAL_RUN_STATUSES.includes(updatedRun.status)) { + this.clearAbortController(runUuid); + } + + return updatedRun; + } + + private static async appendRunStatusEvent( + runUuid: string, + eventType: string, + status: AgentRunStatus, + updatedRun: AgentRun + ): Promise { + await AgentRunEventService.appendStatusEvent(runUuid, eventType, this.buildStatusEventPayload(status, updatedRun)); + } + + private static buildStatusEventPayload( + status: AgentRunStatus, + updatedRun: AgentRun, + executionOwner?: string, + eventContext: OwnerStatusEventContext = {} + ): Record { + return { + status, + error: updatedRun.error || null, + usageSummary: updatedRun.usageSummary || {}, + ...(executionOwner ? { executionOwner } : {}), + ...(eventContext.dispatchAttemptId ? { dispatchAttemptId: eventContext.dispatchAttemptId } : {}), + }; + } + + private static async patchRunForExecutionOwnerInTransaction( + runUuid: string, + executionOwner: string, + patch: Partial, + trx: Transaction + ): Promise { + const run = await AgentRun.query(trx).findOne({ uuid: runUuid }).forUpdate(); + if (!run) { + throw new Error(RUN_NOT_FOUND_ERROR); + } + + this.requireRunExecutionOwner(runUuid, executionOwner, run); + + return AgentRun.query(trx).patchAndFetchById(run.id, patch); + } + + private static requireRunExecutionOwner(runUuid: string, executionOwner: string, run: AgentRun): void { + if (run.executionOwner !== executionOwner || TERMINAL_RUN_STATUSES.includes(run.status)) { + throw runOwnershipLost(runUuid, executionOwner, run); + } + } + + static async startRun( + runUuid: string, + resolved: { + resolvedHarness: string; + provider: string; + model: string; + sandboxGeneration?: number | null; + } + ): Promise { + const now = new Date().toISOString(); + return this.patchStatus(runUuid, 'running', { + startedAt: now, + completedAt: null, + cancelledAt: null, + error: null, + resolvedHarness: resolved.resolvedHarness, + resolvedProvider: resolved.provider, + resolvedModel: resolved.model, + provider: resolved.provider, + model: resolved.model, + sandboxGeneration: resolved.sandboxGeneration ?? null, + } as Partial); + } + + static async markWaitingForApproval(runUuid: string): Promise { + return this.patchStatus(runUuid, 'waiting_for_approval'); + } + + static async markCompleted(runUuid: string, usageSummary?: AgentRunUsageSummary): Promise { + this.clearAbortController(runUuid); + return this.patchStatus(runUuid, 'completed', { + completedAt: new Date().toISOString(), + usageSummary: (usageSummary || {}) as Record, + }); + } + + static async markFailed(runUuid: string, error: unknown, usageSummary?: AgentRunUsageSummary): Promise { this.clearAbortController(runUuid); return this.patchStatus(runUuid, 'failed', { completedAt: new Date().toISOString(), usageSummary: (usageSummary || {}) as Record, - streamState: streamState || {}, error: serializeRunError(error), }); } - static async appendStreamChunks(runUuid: string, chunks: UIMessageChunk[]): Promise { + static async markQueuedRunDispatchFailed(runUuid: string, error: unknown): Promise { + let latestSequence: number | null = null; + const failedRun = await AgentRun.transaction(async (trx) => { + const run = await AgentRun.query(trx).findOne({ uuid: runUuid }).forUpdate(); + if (!run) { + throw new Error(RUN_NOT_FOUND_ERROR); + } + + if (run.status !== 'queued' || run.executionOwner) { + return run; + } + + const nextRun = await AgentRun.query(trx).patchAndFetchById(run.id, { + status: 'failed', + completedAt: new Date().toISOString(), + executionOwner: null, + leaseExpiresAt: null, + heartbeatAt: null, + error: serializeRunError(error), + } as Partial); + + latestSequence = await AgentRunEventService.appendStatusEventForRunInTransaction( + nextRun, + 'run.failed', + this.buildStatusEventPayload('failed', nextRun), + trx + ); + + return nextRun; + }); + + if (latestSequence) { + await AgentRunEventService.notifyRunEventsInserted(failedRun.uuid, latestSequence); + } + + return failedRun; + } + + static async appendStreamChunks(runUuid: string, chunks: AgentUiMessageChunk[]): Promise { const run = await AgentRun.query().findOne({ uuid: runUuid }); if (!run) { throw new Error(RUN_NOT_FOUND_ERROR); @@ -257,14 +906,8 @@ export default class AgentRunService { return run; } - const existingChunks = Array.isArray(run.streamState?.chunks) ? (run.streamState.chunks as UIMessageChunk[]) : []; - const nextChunks = sanitizeAgentRunStreamChunks([...existingChunks, ...chunks.map((chunk) => cloneChunk(chunk))]); - - return this.patchRun(runUuid, { - streamState: { - chunks: nextChunks, - }, - } as Partial); + await AgentRunEventService.appendEventsForChunks(runUuid, chunks); + return run; } static serializeRun(run: AgentRun) { @@ -278,15 +921,22 @@ export default class AgentRunService { threadId: enrichedRun.threadUuid || String(run.threadId), sessionId: enrichedRun.sessionUuid || String(run.sessionId), status: run.status, + requestedHarness: run.requestedHarness || null, + resolvedHarness: run.resolvedHarness || null, + requestedProvider: run.requestedProvider || null, + requestedModel: run.requestedModel || null, + resolvedProvider: run.resolvedProvider || run.provider, + resolvedModel: run.resolvedModel || run.model, provider: run.provider, model: run.model, + sandboxRequirement: run.sandboxRequirement || {}, + sandboxGeneration: run.sandboxGeneration, queuedAt: run.queuedAt, startedAt: run.startedAt, completedAt: run.completedAt, cancelledAt: run.cancelledAt, usageSummary: run.usageSummary || {}, policySnapshot: run.policySnapshot || {}, - streamState: sanitizeAgentRunStreamState(run.streamState || {}), error: run.error, createdAt: run.createdAt || null, updatedAt: run.updatedAt || null, @@ -296,4 +946,8 @@ export default class AgentRunService { static isRunNotFoundError(error: unknown): boolean { return error instanceof Error && error.message === RUN_NOT_FOUND_ERROR; } + + static isActiveRunConflictError(error: unknown): error is ActiveAgentRunError { + return error instanceof ActiveAgentRunError; + } } diff --git a/src/server/services/agent/SandboxService.ts b/src/server/services/agent/SandboxService.ts new file mode 100644 index 00000000..c3846c53 --- /dev/null +++ b/src/server/services/agent/SandboxService.ts @@ -0,0 +1,289 @@ +/** + * Copyright 2026 GoodRx, Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +import AgentSandbox from 'server/models/AgentSandbox'; +import AgentSandboxExposure from 'server/models/AgentSandboxExposure'; +import AgentSession from 'server/models/AgentSession'; +import type { RequestUserIdentity } from 'server/lib/get-user'; +import type { Transaction } from 'objection'; +import type { ResolvedAgentSessionWorkspaceStorageIntent } from 'server/lib/agentSession/runtimeConfig'; + +const SESSION_WORKSPACE_GATEWAY_PORT = parseInt(process.env.AGENT_SESSION_WORKSPACE_GATEWAY_PORT || '13338', 10); + +function mapSessionToSandboxStatus(session: AgentSession): AgentSandbox['status'] { + if (session.status === 'ended' || session.workspaceStatus === 'ended') { + return 'ended'; + } + + if (session.workspaceStatus === 'failed' || session.status === 'error') { + return 'failed'; + } + + if (session.workspaceStatus === 'hibernated') { + return 'suspended'; + } + + if (session.workspaceStatus === 'provisioning') { + return 'provisioning'; + } + + return 'ready'; +} + +function isRecord(value: unknown): value is Record { + return Boolean(value) && typeof value === 'object' && !Array.isArray(value); +} + +function buildProviderState( + session: AgentSession, + workspaceStorage?: ResolvedAgentSessionWorkspaceStorageIntent, + existingProviderState?: Record +): Record { + const existingWorkspaceStorage = isRecord(existingProviderState?.workspaceStorage) + ? existingProviderState.workspaceStorage + : undefined; + + return { + ...(session.namespace ? { namespace: session.namespace } : {}), + ...(session.podName ? { podName: session.podName } : {}), + ...(session.pvcName ? { pvcName: session.pvcName } : {}), + ...(workspaceStorage + ? { + workspaceStorage: { + size: workspaceStorage.storageSize, + accessMode: workspaceStorage.accessMode, + ...(session.pvcName ? { pvcName: session.pvcName } : {}), + }, + } + : existingWorkspaceStorage + ? { workspaceStorage: existingWorkspaceStorage } + : {}), + }; +} + +function toTimestampString(value: unknown): string | null { + if (value instanceof Date) { + return value.toISOString(); + } + + return typeof value === 'string' ? value : null; +} + +export default class AgentSandboxService { + static async getLatestSandboxForSession( + sessionId: number, + options: { trx?: Transaction } = {} + ): Promise { + return AgentSandbox.query(options.trx) + .where({ sessionId }) + .orderBy('generation', 'desc') + .orderBy('createdAt', 'desc') + .first(); + } + + static async recordSessionSandboxState( + session: AgentSession, + options: { trx?: Transaction; workspaceStorage?: ResolvedAgentSessionWorkspaceStorageIntent } = {} + ): Promise { + if (!session.namespace && !session.podName && !session.pvcName) { + return this.getLatestSandboxForSession(session.id, options); + } + + const existing = await this.getLatestSandboxForSession(session.id, options); + const sandbox = existing + ? await AgentSandbox.query(options.trx).patchAndFetchById(existing.id, { + provider: 'lifecycle_kubernetes', + status: mapSessionToSandboxStatus(session), + capabilitySnapshot: { + toolTransport: 'mcp', + persistentFilesystem: Boolean(session.pvcName), + portExposure: true, + editorAccess: true, + }, + providerState: buildProviderState(session, options.workspaceStorage, existing.providerState), + metadata: { + sessionKind: session.sessionKind, + buildUuid: session.buildUuid, + buildKind: session.buildKind, + }, + error: + session.workspaceStatus === 'failed' || session.status === 'error' ? { message: 'Sandbox failed' } : null, + suspendedAt: + session.workspaceStatus === 'hibernated' + ? toTimestampString(session.updatedAt) || new Date().toISOString() + : null, + endedAt: + session.status === 'ended' + ? toTimestampString(session.endedAt) || toTimestampString(session.updatedAt) || new Date().toISOString() + : null, + } as Partial) + : await AgentSandbox.query(options.trx).insertAndFetch({ + sessionId: session.id, + generation: 1, + provider: 'lifecycle_kubernetes', + status: mapSessionToSandboxStatus(session), + capabilitySnapshot: { + toolTransport: 'mcp', + persistentFilesystem: Boolean(session.pvcName), + portExposure: true, + editorAccess: true, + }, + providerState: buildProviderState(session, options.workspaceStorage), + metadata: { + sessionKind: session.sessionKind, + buildUuid: session.buildUuid, + buildKind: session.buildKind, + }, + error: + session.workspaceStatus === 'failed' || session.status === 'error' ? { message: 'Sandbox failed' } : null, + suspendedAt: + session.workspaceStatus === 'hibernated' + ? toTimestampString(session.updatedAt) || new Date().toISOString() + : null, + endedAt: + session.status === 'ended' + ? toTimestampString(session.endedAt) || toTimestampString(session.updatedAt) || new Date().toISOString() + : null, + } as Partial); + + if (sandbox.status === 'suspended' || sandbox.status === 'ended') { + await AgentSandboxExposure.query(options.trx) + .where({ sandboxId: sandbox.id }) + .whereNull('endedAt') + .patch({ + status: 'ended', + endedAt: + toTimestampString(sandbox.suspendedAt) || toTimestampString(sandbox.endedAt) || new Date().toISOString(), + } as Partial); + } + + if (session.podName && session.namespace) { + const editorUrl = `/api/agent-session/workspace-editor/${session.uuid}/`; + const existingEditorExposure = await AgentSandboxExposure.query(options.trx) + .where({ sandboxId: sandbox.id, kind: 'editor' }) + .whereNull('endedAt') + .first(); + + if (existingEditorExposure) { + await AgentSandboxExposure.query(options.trx).patchAndFetchById(existingEditorExposure.id, { + status: + sandbox.status === 'provisioning' + ? 'provisioning' + : sandbox.status === 'failed' + ? 'failed' + : sandbox.status === 'ended' + ? 'ended' + : 'ready', + url: editorUrl, + metadata: { + attachmentKind: 'mcp_gateway', + }, + providerState: {}, + lastVerifiedAt: sandbox.status === 'ready' ? new Date().toISOString() : null, + endedAt: toTimestampString(sandbox.endedAt), + } as Partial); + } else { + await AgentSandboxExposure.query(options.trx).insert({ + sandboxId: sandbox.id, + kind: 'editor', + status: + sandbox.status === 'provisioning' + ? 'provisioning' + : sandbox.status === 'failed' + ? 'failed' + : sandbox.status === 'ended' + ? 'ended' + : 'ready', + url: editorUrl, + metadata: { + attachmentKind: 'mcp_gateway', + }, + providerState: {}, + lastVerifiedAt: sandbox.status === 'ready' ? new Date().toISOString() : null, + endedAt: toTimestampString(sandbox.endedAt), + } as Partial); + } + } + + return sandbox; + } + + static async ensureChatSandbox({ + sessionId, + userId, + userIdentity, + githubToken, + }: { + sessionId: string; + userId: string; + userIdentity: RequestUserIdentity; + githubToken?: string | null; + }): Promise<{ session: AgentSession; sandbox: AgentSandbox | null }> { + let session = await AgentSession.query().findOne({ uuid: sessionId, userId }); + if (!session) { + throw new Error('Agent session not found'); + } + + if ( + session.sessionKind === 'chat' && + (session.workspaceStatus !== 'ready' || !session.namespace || !session.podName) + ) { + const AgentSessionService = (await import('server/services/agentSession')).default; + session = await AgentSessionService.provisionChatRuntime({ + sessionId, + userId, + userIdentity, + githubToken, + }); + } + + const sandbox = await this.recordSessionSandboxState(session); + return { session, sandbox }; + } + + static async resolveWorkspaceGatewayBaseUrl(sessionUuid: string): Promise { + const session = await AgentSession.query().findOne({ uuid: sessionUuid }); + if (!session) { + return null; + } + + const sandbox = await this.recordSessionSandboxState(session); + const providerState = sandbox?.providerState || {}; + const podName = typeof providerState.podName === 'string' ? providerState.podName : session.podName; + const namespace = typeof providerState.namespace === 'string' ? providerState.namespace : session.namespace; + + if (!podName || !namespace || session.status !== 'active') { + return null; + } + + return `http://${podName}.${namespace}.svc.cluster.local:${SESSION_WORKSPACE_GATEWAY_PORT}`; + } + + static serializeSandboxExposure(exposure: AgentSandboxExposure) { + return { + id: exposure.uuid, + kind: exposure.kind, + status: exposure.status, + targetPort: exposure.targetPort, + url: exposure.url, + metadata: exposure.metadata || {}, + lastVerifiedAt: exposure.lastVerifiedAt, + endedAt: exposure.endedAt, + createdAt: exposure.createdAt || null, + updatedAt: exposure.updatedAt || null, + }; + } +} diff --git a/src/server/services/agent/SessionReadService.ts b/src/server/services/agent/SessionReadService.ts new file mode 100644 index 00000000..a21f27ba --- /dev/null +++ b/src/server/services/agent/SessionReadService.ts @@ -0,0 +1,265 @@ +/** + * Copyright 2026 GoodRx, Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +import AgentSession from 'server/models/AgentSession'; +import AgentSandbox from 'server/models/AgentSandbox'; +import AgentSandboxExposure from 'server/models/AgentSandboxExposure'; +import AgentSource from 'server/models/AgentSource'; +import AgentThread from 'server/models/AgentThread'; +import type { PaginationMetadata } from 'server/lib/paginate'; +import AgentThreadService from './ThreadService'; +import AgentSandboxService from './SandboxService'; + +export const DEFAULT_AGENT_SESSION_LIST_LIMIT = 25; +export const MAX_AGENT_SESSION_LIST_LIMIT = 100; + +interface ListOwnedSessionRecordOptions { + includeEnded?: boolean; + page?: number; + limit?: number; +} + +interface SessionRecordRelations { + source: AgentSource; + sandbox: AgentSandbox | null; + exposures: AgentSandboxExposure[]; + defaultThread: AgentThread | null; +} + +function mapSessionStatus(session: AgentSession): 'ready' | 'ended' | 'error' { + if (session.status === 'ended') { + return 'ended'; + } + + if (session.status === 'error' || session.workspaceStatus === 'failed') { + return 'error'; + } + + return 'ready'; +} + +function buildDerivedSourceInput(session: AgentSession, source: AgentSource) { + const workspaceRepos = session.workspaceRepos ?? []; + const primaryRepo = workspaceRepos.find((repo) => repo.primary) ?? workspaceRepos[0] ?? null; + + return { + ...(source.input || {}), + sessionKind: session.sessionKind, + buildUuid: session.buildUuid, + buildKind: session.buildKind, + repo: primaryRepo?.repo ?? null, + branch: primaryRepo?.branch ?? null, + primaryRepo: primaryRepo?.repo ?? null, + primaryBranch: primaryRepo?.branch ?? null, + workspaceRepos, + selectedServices: session.selectedServices ?? [], + services: (session.selectedServices ?? []).map((service) => service.name), + }; +} + +export default class AgentSessionReadService { + static async getOwnedSessionRecord(sessionId: string, userId: string) { + const session = await AgentSession.query().findOne({ uuid: sessionId, userId }); + if (!session) { + return null; + } + + return this.serializeSessionRecord(session); + } + + static async listOwnedSessionRecords(userId: string, options?: ListOwnedSessionRecordOptions) { + const page = + typeof options?.page === 'number' && Number.isFinite(options.page) && options.page > 0 + ? Math.floor(options.page) + : 1; + const limit = + typeof options?.limit === 'number' && Number.isFinite(options.limit) && options.limit > 0 + ? Math.min(Math.floor(options.limit), MAX_AGENT_SESSION_LIST_LIMIT) + : DEFAULT_AGENT_SESSION_LIST_LIMIT; + const query = AgentSession.query().where({ userId }); + + if (!options?.includeEnded) { + query.whereIn('status', ['starting', 'active']); + } + + const result = await query + .orderBy('updatedAt', 'desc') + .orderBy('createdAt', 'desc') + .page(page - 1, limit); + + return { + records: await this.listSessionRecords(result.results), + metadata: { + pagination: { + current: page, + total: Math.max(1, Math.ceil(result.total / limit)), + items: result.total, + limit, + }, + } as { pagination: PaginationMetadata }, + }; + } + + static async serializeSessionRecord(session: AgentSession) { + const [record] = await this.listSessionRecords([session]); + return record; + } + + private static serializeSessionRecordWithRelations(session: AgentSession, relations: SessionRecordRelations) { + const { source, sandbox, defaultThread } = relations; + return { + session: { + id: session.uuid, + status: mapSessionStatus(session), + userId: session.userId, + ownerGithubUsername: session.ownerGithubUsername, + defaults: { + model: session.defaultModel || session.model, + harness: session.defaultHarness, + }, + defaultThreadId: defaultThread?.uuid || null, + lastActivity: session.lastActivity || null, + endedAt: session.endedAt || null, + createdAt: session.createdAt || null, + updatedAt: session.updatedAt || null, + }, + source: { + id: source.uuid, + adapter: source.adapter, + status: source.status, + input: buildDerivedSourceInput(session, source), + sandboxRequirements: source.sandboxRequirements || {}, + error: source.error, + preparedAt: source.preparedAt, + cleanedUpAt: source.cleanedUpAt, + createdAt: source.createdAt || null, + updatedAt: source.updatedAt || null, + }, + sandbox: sandbox + ? { + id: sandbox.uuid, + generation: sandbox.generation, + provider: sandbox.provider, + status: sandbox.status, + capabilitySnapshot: sandbox.capabilitySnapshot || {}, + exposures: relations.exposures.map((exposure) => AgentSandboxService.serializeSandboxExposure(exposure)), + suspendedAt: sandbox.suspendedAt, + endedAt: sandbox.endedAt, + error: sandbox.error, + createdAt: sandbox.createdAt || null, + updatedAt: sandbox.updatedAt || null, + } + : { + id: null, + generation: null, + provider: null, + status: 'none', + capabilitySnapshot: {}, + exposures: [], + suspendedAt: null, + endedAt: null, + error: null, + createdAt: null, + updatedAt: null, + }, + }; + } + + static async listSessionRecords(sessions: AgentSession[]) { + if (sessions.length === 0) { + return []; + } + + const sessionIds = sessions.map((session) => session.id); + const defaultThreadIds = sessions + .map((session) => session.defaultThreadId) + .filter((threadId): threadId is number => Number.isInteger(threadId)); + const [sources, sandboxes, defaultThreads, fallbackThreads] = await Promise.all([ + AgentSource.query().whereIn('sessionId', sessionIds), + AgentSandbox.query().whereIn('sessionId', sessionIds).orderBy('generation', 'desc').orderBy('createdAt', 'desc'), + defaultThreadIds.length ? AgentThread.query().whereIn('id', defaultThreadIds) : Promise.resolve([]), + AgentThread.query() + .whereIn('sessionId', sessionIds) + .where({ isDefault: true }) + .whereNull('archivedAt') + .orderBy('createdAt', 'asc'), + ]); + const sourceBySessionId = new Map(); + for (const source of sources) { + sourceBySessionId.set(source.sessionId, source); + } + + const sandboxBySessionId = new Map(); + for (const sandbox of sandboxes) { + if (!sandboxBySessionId.has(sandbox.sessionId)) { + sandboxBySessionId.set(sandbox.sessionId, sandbox); + } + } + + const latestSandboxIds = [...sandboxBySessionId.values()].map((sandbox) => sandbox.id); + const exposures = latestSandboxIds.length + ? await AgentSandboxExposure.query().whereIn('sandboxId', latestSandboxIds).orderBy('createdAt', 'asc') + : []; + const exposuresBySandboxId = new Map(); + for (const exposure of exposures) { + const existing = exposuresBySandboxId.get(exposure.sandboxId) || []; + existing.push(exposure); + exposuresBySandboxId.set(exposure.sandboxId, existing); + } + + const defaultThreadById = new Map(); + for (const thread of defaultThreads) { + defaultThreadById.set(thread.id, thread); + } + + const fallbackThreadBySessionId = new Map(); + for (const thread of fallbackThreads) { + if (!fallbackThreadBySessionId.has(thread.sessionId)) { + fallbackThreadBySessionId.set(thread.sessionId, thread); + } + } + + return sessions.map((session) => { + const source = sourceBySessionId.get(session.id); + if (!source) { + throw new Error(`Agent session source missing for session ${session.uuid}`); + } + + const sandbox = sandboxBySessionId.get(session.id) || null; + const defaultThread = + (session.defaultThreadId ? defaultThreadById.get(session.defaultThreadId) : null) || + fallbackThreadBySessionId.get(session.id) || + null; + + return this.serializeSessionRecordWithRelations(session, { + source, + sandbox, + defaultThread, + exposures: sandbox ? exposuresBySandboxId.get(sandbox.id) || [] : [], + }); + }); + } + + static async serializeThread(thread: AgentThread, session: AgentSession) { + const serialized = AgentThreadService.serializeThread(thread, session.uuid); + return { + ...serialized, + session: { + id: session.uuid, + }, + }; + } +} diff --git a/src/server/services/agent/SourceService.ts b/src/server/services/agent/SourceService.ts new file mode 100644 index 00000000..087766df --- /dev/null +++ b/src/server/services/agent/SourceService.ts @@ -0,0 +1,147 @@ +/** + * Copyright 2026 GoodRx, Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +import AgentSession from 'server/models/AgentSession'; +import AgentSource from 'server/models/AgentSource'; +import { SESSION_WORKSPACE_ROOT } from 'server/lib/agentSession/workspace'; +import { AgentSessionKind } from 'shared/constants'; +import type { Transaction } from 'objection'; +import type { ResolvedAgentSessionWorkspaceStorageIntent } from 'server/lib/agentSession/runtimeConfig'; + +function deriveAdapter(session: AgentSession): string { + if (session.sessionKind === 'chat') { + return 'blank_workspace'; + } + + return session.buildKind === 'sandbox' ? 'lifecycle_fork' : 'lifecycle_environment'; +} + +function deriveStatus(session: AgentSession): AgentSource['status'] { + if (session.status === 'ended') { + return 'cleaned_up'; + } + + if (session.status === 'error') { + return 'failed'; + } + + return 'ready'; +} + +function toTimestampString(value: unknown): string | null { + if (value instanceof Date) { + return value.toISOString(); + } + + return typeof value === 'string' ? value : null; +} + +export default class AgentSourceService { + static async getSessionSource(sessionId: number, options: { trx?: Transaction } = {}): Promise { + return AgentSource.query(options.trx).findOne({ sessionId }); + } + + static async createSessionSource( + session: AgentSession, + options: { trx?: Transaction; workspaceStorage?: ResolvedAgentSessionWorkspaceStorageIntent } = {} + ): Promise { + const status = deriveStatus(session); + const workspaceRepos = session.workspaceRepos ?? []; + const primaryRepo = workspaceRepos.find((repo) => repo.primary) ?? workspaceRepos[0] ?? null; + const workspaceLayout = + session.sessionKind === AgentSessionKind.CHAT + ? { + repos: [], + primaryPath: SESSION_WORKSPACE_ROOT, + } + : { + repos: workspaceRepos, + primaryPath: primaryRepo?.mountPath || SESSION_WORKSPACE_ROOT, + }; + + return AgentSource.query(options.trx).insertAndFetch({ + sessionId: session.id, + adapter: deriveAdapter(session), + status, + input: { + buildUuid: session.buildUuid, + buildKind: session.buildKind, + sessionKind: session.sessionKind, + ...(options.workspaceStorage?.requestedSize + ? { + workspace: { + storageSize: options.workspaceStorage.requestedSize, + }, + } + : {}), + }, + preparedSource: { + kind: 'workspace_snapshot', + workspaceLayout, + artifactRefs: [], + metadata: { + buildUuid: session.buildUuid, + buildKind: session.buildKind, + sessionKind: session.sessionKind, + }, + }, + sandboxRequirements: { + filesystem: 'persistent', + suspendMode: session.sessionKind === AgentSessionKind.CHAT ? 'filesystem' : 'none', + editorAccess: true, + previewPorts: true, + }, + error: status === 'failed' ? { message: 'Source failed' } : null, + preparedAt: status === 'cleaned_up' ? null : toTimestampString(session.updatedAt) || new Date().toISOString(), + cleanedUpAt: + status === 'cleaned_up' + ? toTimestampString(session.endedAt) || toTimestampString(session.updatedAt) || new Date().toISOString() + : null, + } as Partial); + } + + static async recordSessionState( + session: AgentSession, + options: { trx?: Transaction } = {} + ): Promise { + const existing = await this.getSessionSource(session.id, options); + if (!existing) { + return null; + } + + const status = deriveStatus(session); + const patch: Partial = {}; + + if (existing.status !== status) { + patch.status = status; + } + + if (status === 'failed' && !existing.error) { + patch.error = { message: 'Source failed' }; + } + + if (status === 'cleaned_up' && !existing.cleanedUpAt) { + patch.cleanedUpAt = + toTimestampString(session.endedAt) || toTimestampString(session.updatedAt) || new Date().toISOString(); + } + + if (Object.keys(patch).length === 0) { + return existing; + } + + return AgentSource.query(options.trx).patchAndFetchById(existing.id, patch); + } +} diff --git a/src/server/services/agent/StreamBroker.ts b/src/server/services/agent/StreamBroker.ts deleted file mode 100644 index 78401fcb..00000000 --- a/src/server/services/agent/StreamBroker.ts +++ /dev/null @@ -1,205 +0,0 @@ -/** - * Copyright 2026 GoodRx, Inc. - * - * Licensed under the Apache License, Version 2.0 (the "License"); - * you may not use this file except in compliance with the License. - * You may obtain a copy of the License at - * - * http://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ - -import type { UIMessageChunk } from 'ai'; -import { getLogger } from 'server/lib/logger'; -import AgentRunService from './RunService'; - -type BrokerEntry = { - history: UIMessageChunk[]; - subscribers: Set>; - active: boolean; - cleanupTimer?: NodeJS.Timeout; - persistTimer?: NodeJS.Timeout; - pendingPersistChunks: UIMessageChunk[]; - persistPromise: Promise; -}; - -const STREAM_TTL_MS = 5 * 60 * 1000; -const PERSIST_DEBOUNCE_MS = 200; - -function cloneChunk(chunk: T): T { - return JSON.parse(JSON.stringify(chunk)) as T; -} - -export default class AgentStreamBroker { - private static entries = new Map(); - - static attach(runUuid: string, stream: ReadableStream): void { - const existing = this.entries.get(runUuid); - if (existing?.active) { - return; - } - - if (existing?.cleanupTimer) { - clearTimeout(existing.cleanupTimer); - } - - const entry: BrokerEntry = existing || { - history: [], - subscribers: new Set(), - active: true, - pendingPersistChunks: [], - persistPromise: Promise.resolve(), - }; - - entry.active = true; - entry.cleanupTimer = undefined; - this.entries.set(runUuid, entry); - - void this.consume(runUuid, stream, entry); - } - - static open(runUuid: string): ReadableStream | null { - const entry = this.entries.get(runUuid); - if (!entry) { - return null; - } - - let controllerRef: ReadableStreamDefaultController | null = null; - - return new ReadableStream({ - start: (controller) => { - controllerRef = controller; - - for (const chunk of entry.history) { - controller.enqueue(cloneChunk(chunk)); - } - - if (entry.active) { - entry.subscribers.add(controller); - } else { - controller.close(); - } - }, - cancel: () => { - if (controllerRef) { - entry.subscribers.delete(controllerRef); - } - }, - }); - } - - private static async consume( - runUuid: string, - stream: ReadableStream, - entry: BrokerEntry - ): Promise { - const reader = stream.getReader(); - - try { - let streamDone = false; - - while (!streamDone) { - const { value, done } = await reader.read(); - if (done) { - streamDone = true; - continue; - } - - if (value === undefined) { - continue; - } - - const clonedChunk = cloneChunk(value); - entry.history.push(clonedChunk); - this.queuePersist(runUuid, entry, clonedChunk); - - for (const subscriber of [...entry.subscribers]) { - try { - subscriber.enqueue(cloneChunk(clonedChunk)); - } catch (error) { - entry.subscribers.delete(subscriber); - getLogger().debug({ error, runUuid }, `AgentExec: stream subscriber dropped runId=${runUuid}`); - } - } - } - - await this.flushPersistedChunks(runUuid, entry); - this.finish(runUuid, entry); - } catch (error) { - getLogger().warn({ error, runUuid }, `AgentExec: stream broker failed runId=${runUuid}`); - - await this.flushPersistedChunks(runUuid, entry); - for (const subscriber of [...entry.subscribers]) { - try { - subscriber.error(error); - } catch { - // Ignore subscriber shutdown failures. - } - } - - this.finish(runUuid, entry); - } finally { - reader.releaseLock(); - } - } - - private static queuePersist(runUuid: string, entry: BrokerEntry, chunk: UIMessageChunk): void { - entry.pendingPersistChunks.push(cloneChunk(chunk)); - if (entry.persistTimer) { - return; - } - - entry.persistTimer = setTimeout(() => { - entry.persistTimer = undefined; - void this.flushPersistedChunks(runUuid, entry); - }, PERSIST_DEBOUNCE_MS); - } - - private static async flushPersistedChunks(runUuid: string, entry: BrokerEntry): Promise { - if (entry.persistTimer) { - clearTimeout(entry.persistTimer); - entry.persistTimer = undefined; - } - - if (entry.pendingPersistChunks.length === 0) { - return; - } - - const chunks = entry.pendingPersistChunks.splice(0, entry.pendingPersistChunks.length); - entry.persistPromise = entry.persistPromise - .catch(() => undefined) - .then(() => AgentRunService.appendStreamChunks(runUuid, chunks)) - .then(() => undefined) - .catch((error) => { - getLogger().warn({ error, runUuid }, `AgentExec: stream persist failed runId=${runUuid}`); - }); - - await entry.persistPromise; - } - - private static finish(runUuid: string, entry: BrokerEntry): void { - entry.active = false; - if (entry.persistTimer) { - clearTimeout(entry.persistTimer); - entry.persistTimer = undefined; - } - - for (const subscriber of [...entry.subscribers]) { - try { - subscriber.close(); - } catch { - // Ignore subscriber shutdown failures. - } - } - - entry.subscribers.clear(); - entry.cleanupTimer = setTimeout(() => { - this.entries.delete(runUuid); - }, STREAM_TTL_MS); - } -} diff --git a/src/server/services/agent/__tests__/AdminService.test.ts b/src/server/services/agent/__tests__/AdminService.test.ts index 2f56f9c4..286e63f3 100644 --- a/src/server/services/agent/__tests__/AdminService.test.ts +++ b/src/server/services/agent/__tests__/AdminService.test.ts @@ -18,6 +18,14 @@ const mockEnrichSessions = jest.fn(); const mockSessionQuery = jest.fn(); const mockThreadQuery = jest.fn(); const mockPendingActionQuery = jest.fn(); +const mockMessageQuery = jest.fn(); +const mockRunQuery = jest.fn(); +const mockRunEventQuery = jest.fn(); +const mockToolExecutionQuery = jest.fn(); +const mockSerializeRun = jest.fn(); +const mockSerializeRunEvent = jest.fn(); +const mockSerializeThread = jest.fn(); +const mockSerializeCanonicalMessage = jest.fn(); jest.mock('server/services/agentSession', () => ({ __esModule: true, @@ -47,6 +55,62 @@ jest.mock('server/models/AgentPendingAction', () => ({ }, })); +jest.mock('server/models/AgentMessage', () => ({ + __esModule: true, + default: { + query: (...args: unknown[]) => mockMessageQuery(...args), + }, +})); + +jest.mock('server/models/AgentRun', () => ({ + __esModule: true, + default: { + query: (...args: unknown[]) => mockRunQuery(...args), + }, +})); + +jest.mock('server/models/AgentRunEvent', () => ({ + __esModule: true, + default: { + query: (...args: unknown[]) => mockRunEventQuery(...args), + }, +})); + +jest.mock('server/models/AgentToolExecution', () => ({ + __esModule: true, + default: { + query: (...args: unknown[]) => mockToolExecutionQuery(...args), + }, +})); + +jest.mock('../RunService', () => ({ + __esModule: true, + default: { + serializeRun: (...args: unknown[]) => mockSerializeRun(...args), + }, +})); + +jest.mock('../RunEventService', () => ({ + __esModule: true, + default: { + serializeRunEvent: (...args: unknown[]) => mockSerializeRunEvent(...args), + }, +})); + +jest.mock('../ThreadService', () => ({ + __esModule: true, + default: { + serializeThread: (...args: unknown[]) => mockSerializeThread(...args), + }, +})); + +jest.mock('../MessageStore', () => ({ + __esModule: true, + default: { + serializeCanonicalMessage: (...args: unknown[]) => mockSerializeCanonicalMessage(...args), + }, +})); + import AgentAdminService from '../AdminService'; describe('AgentAdminService.listSessions', () => { @@ -157,3 +221,241 @@ describe('AgentAdminService.listSessions', () => { ]); }); }); + +describe('AgentAdminService.getThreadConversation', () => { + beforeEach(() => { + jest.clearAllMocks(); + mockSerializeThread.mockImplementation((thread, sessionId) => ({ + id: thread.uuid, + sessionId, + title: thread.title || null, + lastRunAt: thread.lastRunAt || null, + })); + mockSerializeRun.mockImplementation((run) => ({ + id: run.uuid, + threadId: run.threadUuid, + sessionId: run.sessionUuid, + status: run.status, + })); + mockSerializeRunEvent.mockImplementation((event) => ({ + id: event.uuid, + runId: event.runUuid, + threadId: event.threadUuid, + sessionId: event.sessionUuid, + sequence: event.sequence, + eventType: event.eventType, + version: 1, + payload: event.payload, + })); + mockSerializeCanonicalMessage.mockImplementation((message, threadUuid, runUuid) => ({ + id: message.uuid, + clientMessageId: message.clientMessageId || null, + threadId: threadUuid, + runId: runUuid, + role: message.role, + parts: message.parts, + createdAt: message.createdAt || null, + })); + }); + + it('returns canonical messages, runs, events, pending actions, and tool executions for admin replay', async () => { + jest.spyOn(AgentAdminService, 'getSession').mockResolvedValueOnce({ + session: { + id: 'session-1', + status: 'active', + }, + threads: [ + { + id: 'thread-1', + sessionId: 'session-1', + messageCount: 1, + runCount: 1, + pendingActionsCount: 1, + latestRun: null, + }, + ], + } as any); + + mockThreadQuery.mockReturnValueOnce({ + findOne: jest.fn().mockResolvedValue({ + id: 7, + uuid: 'thread-1', + sessionId: 17, + }), + }); + mockSessionQuery.mockReturnValueOnce({ + findById: jest.fn().mockResolvedValue({ + id: 17, + uuid: 'session-1', + }), + }); + + mockMessageQuery.mockReturnValueOnce({ + alias: jest.fn().mockReturnThis(), + leftJoinRelated: jest.fn().mockReturnThis(), + where: jest.fn().mockReturnThis(), + select: jest.fn().mockReturnThis(), + orderBy: jest.fn().mockResolvedValue([ + { + uuid: 'message-1', + clientMessageId: 'client-message-1', + role: 'user', + parts: [{ type: 'text', text: 'Hi' }], + runUuid: 'run-1', + createdAt: '2026-04-11T00:00:00.000Z', + }, + ]), + }); + mockRunQuery.mockReturnValueOnce({ + where: jest.fn().mockReturnThis(), + orderBy: jest.fn().mockResolvedValue([ + { + uuid: 'run-1', + status: 'completed', + }, + ]), + }); + mockPendingActionQuery.mockReturnValueOnce({ + alias: jest.fn().mockReturnThis(), + joinRelated: jest.fn().mockReturnThis(), + where: jest.fn().mockReturnThis(), + select: jest.fn().mockReturnThis(), + orderBy: jest.fn().mockResolvedValue([ + { + uuid: 'action-1', + threadId: 7, + runId: 11, + runUuid: 'run-1', + kind: 'tool_approval', + status: 'approved', + capabilityKey: 'workspace_write', + title: 'Approve workspace edit', + description: 'A workspace edit requires approval.', + payload: { + toolName: 'mcp__sandbox__workspace_edit_file', + input: { + path: 'sample-file.txt', + }, + }, + resolution: { + approved: true, + }, + resolvedAt: '2026-04-11T00:01:00.000Z', + createdAt: '2026-04-11T00:00:00.000Z', + }, + ]), + }); + mockToolExecutionQuery.mockReturnValueOnce({ + alias: jest.fn().mockReturnThis(), + joinRelated: jest.fn().mockReturnThis(), + leftJoinRelated: jest.fn().mockReturnThis(), + where: jest.fn().mockReturnThis(), + select: jest.fn().mockReturnThis(), + orderBy: jest.fn().mockResolvedValue([ + { + uuid: 'tool-1', + source: 'mcp', + serverSlug: 'sandbox', + toolName: 'workspace.edit_file', + toolCallId: 'tool-call-1', + args: { path: 'sample-file.txt' }, + result: null, + status: 'completed', + safetyLevel: null, + approved: true, + startedAt: null, + completedAt: null, + durationMs: null, + createdAt: '2026-04-11T00:00:00.000Z', + updatedAt: '2026-04-11T00:00:00.000Z', + threadUuid: 'thread-1', + runUuid: 'run-1', + pendingActionUuid: 'action-1', + }, + ]), + }); + const eventQuery: any = { + alias: jest.fn().mockReturnThis(), + joinRelated: jest.fn().mockReturnThis(), + where: jest.fn().mockReturnThis(), + select: jest.fn().mockReturnThis(), + orderBy: jest.fn(), + }; + eventQuery.orderBy + .mockImplementationOnce(() => eventQuery) + .mockResolvedValueOnce([ + { + uuid: 'event-1', + runUuid: 'run-1', + sequence: 1, + eventType: 'approval.resolved', + payload: { + actionId: 'action-1', + approved: true, + }, + }, + ]); + mockRunEventQuery.mockReturnValueOnce(eventQuery); + + const result = await AgentAdminService.getThreadConversation('thread-1'); + + expect(mockSerializeCanonicalMessage).toHaveBeenCalledWith( + expect.objectContaining({ uuid: 'message-1' }), + 'thread-1', + 'run-1' + ); + expect(result.messages).toEqual([ + { + id: 'message-1', + clientMessageId: 'client-message-1', + threadId: 'thread-1', + runId: 'run-1', + role: 'user', + parts: [{ type: 'text', text: 'Hi' }], + createdAt: '2026-04-11T00:00:00.000Z', + }, + ]); + expect(result.pendingActions).toEqual([ + expect.objectContaining({ + id: 'action-1', + threadId: 'thread-1', + runId: 'run-1', + requestedAt: '2026-04-11T00:00:00.000Z', + toolName: 'mcp__sandbox__workspace_edit_file', + }), + ]); + expect(result.events).toEqual([ + { + id: 'event-1', + runId: 'run-1', + threadId: 'thread-1', + sessionId: 'session-1', + sequence: 1, + eventType: 'approval.resolved', + version: 1, + payload: { + actionId: 'action-1', + approved: true, + }, + }, + ]); + expect(mockSerializeRunEvent).toHaveBeenCalledWith( + expect.objectContaining({ + uuid: 'event-1', + runUuid: 'run-1', + threadUuid: 'thread-1', + sessionUuid: 'session-1', + }) + ); + expect(result.toolExecutions).toEqual([ + expect.objectContaining({ + id: 'tool-1', + threadId: 'thread-1', + runId: 'run-1', + pendingActionId: 'action-1', + toolCallId: 'tool-call-1', + }), + ]); + expect(result.messages[0]).not.toHaveProperty('metadata'); + }); +}); diff --git a/src/server/services/agent/__tests__/ApprovalService.test.ts b/src/server/services/agent/__tests__/ApprovalService.test.ts index 522c2127..0b414e42 100644 --- a/src/server/services/agent/__tests__/ApprovalService.test.ts +++ b/src/server/services/agent/__tests__/ApprovalService.test.ts @@ -21,6 +21,14 @@ jest.mock('ai', () => ({ })); jest.mock('server/models/AgentPendingAction', () => ({ + __esModule: true, + default: { + query: jest.fn(), + transaction: jest.fn((callback) => callback({ trx: true })), + }, +})); + +jest.mock('server/models/AgentRun', () => ({ __esModule: true, default: { query: jest.fn(), @@ -34,29 +42,139 @@ jest.mock('../ThreadService', () => ({ }, })); +const mockGetRunByUuid = jest.fn(); +const mockPatchStatus = jest.fn(); + +jest.mock('../RunService', () => ({ + __esModule: true, + default: { + getRunByUuid: (...args: unknown[]) => mockGetRunByUuid(...args), + patchStatus: (...args: unknown[]) => mockPatchStatus(...args), + }, +})); + +const mockAppendStatusEvent = jest.fn(); +const mockAppendStatusEventForRunInTransaction = jest.fn(); +const mockNotifyRunEventsInserted = jest.fn(); + +jest.mock('../RunEventService', () => ({ + __esModule: true, + default: { + appendStatusEvent: (...args: unknown[]) => mockAppendStatusEvent(...args), + appendStatusEventForRunInTransaction: (...args: unknown[]) => mockAppendStatusEventForRunInTransaction(...args), + notifyRunEventsInserted: (...args: unknown[]) => mockNotifyRunEventsInserted(...args), + }, +})); + +const mockEnqueueRun = jest.fn(); + +jest.mock('../RunQueueService', () => ({ + __esModule: true, + default: { + enqueueRun: (...args: unknown[]) => mockEnqueueRun(...args), + }, +})); + import AgentPendingAction from 'server/models/AgentPendingAction'; +import AgentRun from 'server/models/AgentRun'; import ApprovalService from '../ApprovalService'; import AgentThreadService from '../ThreadService'; +import { getToolName } from 'ai'; const mockPendingActionQuery = AgentPendingAction.query as jest.Mock; +const mockPendingActionTransaction = AgentPendingAction.transaction as jest.Mock; +const mockRunQuery = AgentRun.query as jest.Mock; const mockGetOwnedThread = AgentThreadService.getOwnedThread as jest.Mock; +const mockGetToolName = getToolName as jest.Mock; -describe('ApprovalService.serializePendingAction', () => { +function makeTransactionalPendingActionQuery(...firstResults: unknown[]) { + const query: any = {}; + query.alias = jest.fn().mockReturnValue(query); + query.joinRelated = jest.fn().mockReturnValue(query); + query.where = jest.fn().mockReturnValue(query); + query.select = jest.fn().mockReturnValue(query); + query.forUpdate = jest.fn().mockReturnValue(query); + query.first = jest.fn(); + for (const result of firstResults) { + query.first.mockResolvedValueOnce(result); + } + query.patchAndFetchById = jest + .fn() + .mockImplementation((_id, patch) => Promise.resolve({ ...firstResults[0], ...patch })); + return query; +} + +function makeTransactionalRunQuery(run: unknown, queuedRun?: unknown) { + const query: any = {}; + query.findById = jest.fn().mockReturnValue({ + forUpdate: jest.fn().mockResolvedValue(run), + }); + query.patchAndFetchById = jest.fn().mockResolvedValue(queuedRun || run); + return query; +} + +describe('ApprovalService', () => { beforeEach(() => { jest.clearAllMocks(); + mockPendingActionTransaction.mockImplementation((callback) => callback({ trx: true })); + mockAppendStatusEventForRunInTransaction.mockResolvedValue(7); + }); + + it('normalizes canonical pending action response bodies', () => { + expect( + ApprovalService.normalizePendingActionResponseBody({ + approved: true, + reason: 'looks fine', + }) + ).toEqual({ + approved: true, + reason: 'looks fine', + }); + + expect(ApprovalService.normalizePendingActionResponseBody({ approved: false })).toEqual({ + approved: false, + reason: null, + }); + expect(ApprovalService.normalizePendingActionResponseBody({})).toEqual(new Error('approved must be a boolean')); + expect(ApprovalService.normalizePendingActionResponseBody(null)).toEqual( + new Error('Request body must be a JSON object') + ); + expect(ApprovalService.normalizePendingActionResponseBody({ approved: true, rawApproval: true })).toEqual( + new Error('Unsupported pending action response fields: rawApproval') + ); }); - it('omits the unused expiresAt field from API output', () => { + it('serializes display-ready pending action fields without exposing raw payload state', () => { const serialized = ApprovalService.serializePendingAction({ uuid: 'action-1', threadId: 3, runId: 4, kind: 'tool_approval', status: 'pending', - capabilityKey: 'external_mcp_write', + capabilityKey: 'workspace_write', title: 'Approve tool', description: 'Tool requires approval', - payload: { approvalId: 'approval-1' }, + payload: { + approvalId: 'approval-1', + toolCallId: 'tool-call-1', + toolName: 'mcp__sandbox__workspace_edit_file', + input: { + path: 'sample-file.txt', + oldText: 'before', + newText: 'after', + }, + fileChanges: [ + { + path: '/workspace/sample-file.txt', + displayPath: 'sample-file.txt', + kind: 'edited', + summary: 'Updated sample-file.txt', + additions: 1, + deletions: 1, + truncated: false, + }, + ], + }, resolution: null, expiresAt: '2026-04-12T00:00:00.000Z', resolvedAt: null, @@ -66,76 +184,702 @@ describe('ApprovalService.serializePendingAction', () => { expect(serialized).toEqual({ id: 'action-1', - threadId: '3', - runId: '4', kind: 'tool_approval', status: 'pending', - capabilityKey: 'external_mcp_write', + threadId: '3', + runId: '4', title: 'Approve tool', description: 'Tool requires approval', - payload: { approvalId: 'approval-1' }, - resolution: null, - resolvedAt: null, - createdAt: '2026-04-11T00:00:00.000Z', - updatedAt: '2026-04-11T00:00:00.000Z', + requestedAt: '2026-04-11T00:00:00.000Z', + expiresAt: '2026-04-12T00:00:00.000Z', + toolName: 'mcp__sandbox__workspace_edit_file', + argumentsSummary: [ + { + name: 'path', + value: 'sample-file.txt', + }, + ], + commandPreview: null, + fileChangePreview: [ + { + path: 'sample-file.txt', + action: 'edited', + summary: 'Updated sample-file.txt', + additions: 1, + deletions: 1, + truncated: false, + }, + ], + riskLabels: ['Workspace write'], }); }); - it('queries approval responses through the action alias before resolving the pending action', async () => { + it('lists only pending actions for the owned thread', async () => { + const query: any = {}; + query.alias = jest.fn().mockReturnValue(query); + query.leftJoinRelated = jest.fn().mockReturnValue(query); + query.where = jest.fn().mockReturnValue(query); + query.select = jest.fn().mockReturnValue(query); + query.orderBy = jest.fn().mockResolvedValue([{ uuid: 'action-1' }]); + mockGetOwnedThread.mockResolvedValue({ id: 7 }); + mockPendingActionQuery.mockReturnValue(query); - const pendingLookupQuery: any = {}; - pendingLookupQuery.alias = jest.fn().mockReturnValue(pendingLookupQuery); - pendingLookupQuery.joinRelated = jest.fn().mockReturnValue(pendingLookupQuery); - pendingLookupQuery.where = jest.fn().mockReturnValue(pendingLookupQuery); - pendingLookupQuery.whereRaw = jest.fn().mockReturnValue(pendingLookupQuery); - pendingLookupQuery.modify = jest.fn((callback) => { - callback(pendingLookupQuery); - return pendingLookupQuery; + await expect(ApprovalService.listPendingActions('thread-1', 'sample-user')).resolves.toEqual([ + { uuid: 'action-1' }, + ]); + + expect(mockGetOwnedThread).toHaveBeenCalledWith('thread-1', 'sample-user'); + expect(query.where).toHaveBeenCalledWith('action.threadId', 7); + expect(query.where).toHaveBeenCalledWith('action.status', 'pending'); + expect(query.orderBy).toHaveBeenCalledWith('action.createdAt', 'asc'); + }); + + it('classifies session workspace approval requests by their workspace capability', async () => { + mockGetToolName.mockReturnValue('mcp__sandbox__workspace_edit_file'); + + const existingLookupQuery: any = {}; + existingLookupQuery.where = jest.fn().mockReturnValue(existingLookupQuery); + existingLookupQuery.whereRaw = jest.fn().mockReturnValue(existingLookupQuery); + existingLookupQuery.first = jest.fn().mockResolvedValue(null); + + const insertQuery = { + insertAndFetch: jest.fn().mockResolvedValue({ id: 1 }), + }; + + mockPendingActionQuery.mockImplementationOnce(() => existingLookupQuery).mockImplementationOnce(() => insertQuery); + + await ApprovalService.upsertApprovalRequest({ + thread: { id: 7 } as any, + run: { id: 11 } as any, + message: { parts: [] } as any, + toolPart: { + approval: { id: 'approval-1' }, + input: { + path: 'approval-check.txt', + oldText: 'original', + newText: 'updated', + }, + state: 'approval-requested', + toolCallId: 'tool-call-1', + } as any, + capabilityKey: 'external_mcp_write', }); - pendingLookupQuery.select = jest.fn().mockReturnValue(pendingLookupQuery); - pendingLookupQuery.orderBy = jest.fn().mockReturnValue(pendingLookupQuery); - pendingLookupQuery.first = jest.fn().mockResolvedValue({ id: 99 }); - const patchQuery = { - patchAndFetchById: jest.fn().mockResolvedValue(undefined), + expect(insertQuery.insertAndFetch).toHaveBeenCalledWith( + expect.objectContaining({ + capabilityKey: 'workspace_write', + payload: expect.objectContaining({ + toolName: 'mcp__sandbox__workspace_edit_file', + }), + }) + ); + }); + + it('does not persist approval requests when the current policy allows the tool', async () => { + mockGetToolName.mockReturnValue('mcp__sandbox__workspace_write_file'); + + await ApprovalService.syncApprovalRequestsFromMessages({ + thread: { id: 7 } as any, + run: { id: 11 } as any, + messages: [ + { + role: 'assistant', + parts: [ + { + approval: { id: 'approval-1' }, + input: { + path: 'approval-check.txt', + content: 'hello', + }, + state: 'approval-requested', + toolCallId: 'tool-call-1', + }, + ], + } as any, + ], + approvalPolicy: { + defaultMode: 'allow', + rules: { + workspace_write: 'allow', + }, + } as any, + toolRules: [], + }); + + expect(mockPendingActionQuery).not.toHaveBeenCalled(); + }); + + it('persists approval requests when a tool rule requires approval', async () => { + mockGetToolName.mockReturnValue('mcp__sandbox__workspace_write_file'); + + const existingLookupQuery: any = {}; + existingLookupQuery.where = jest.fn().mockReturnValue(existingLookupQuery); + existingLookupQuery.whereRaw = jest.fn().mockReturnValue(existingLookupQuery); + existingLookupQuery.first = jest.fn().mockResolvedValue(null); + + const insertQuery = { + insertAndFetch: jest.fn().mockResolvedValue({ id: 1 }), }; - mockPendingActionQuery.mockImplementationOnce(() => pendingLookupQuery).mockImplementationOnce(() => patchQuery); + mockPendingActionQuery.mockImplementationOnce(() => existingLookupQuery).mockImplementationOnce(() => insertQuery); - await ApprovalService.syncApprovalResponsesFromMessages('thread-uuid', 'sample-user', [ - { - metadata: { runId: 'run-uuid' }, - parts: [ - { - state: 'approval-responded', - toolCallId: 'tool-call-1', - approval: { - id: 'approval-1', - approved: true, - reason: 'approved in UI', + await ApprovalService.syncApprovalRequestsFromMessages({ + thread: { id: 7 } as any, + run: { id: 11 } as any, + messages: [ + { + role: 'assistant', + parts: [ + { + approval: { id: 'approval-1' }, + input: { + path: 'approval-check.txt', + content: 'hello', + }, + state: 'approval-requested', + toolCallId: 'tool-call-1', }, - }, + ], + } as any, + ], + approvalPolicy: { + defaultMode: 'allow', + rules: { + workspace_write: 'allow', + }, + } as any, + toolRules: [ + { + toolKey: 'mcp__sandbox__workspace_write_file', + mode: 'require_approval', + }, + ], + }); + + expect(insertQuery.insertAndFetch).toHaveBeenCalledWith( + expect.objectContaining({ + capabilityKey: 'workspace_write', + payload: expect.objectContaining({ + toolName: 'mcp__sandbox__workspace_write_file', + }), + }) + ); + }); + + it('persists stream approval requests before message finalization', async () => { + const existingLookupQuery: any = {}; + existingLookupQuery.where = jest.fn().mockReturnValue(existingLookupQuery); + existingLookupQuery.whereRaw = jest.fn().mockReturnValue(existingLookupQuery); + existingLookupQuery.first = jest.fn().mockResolvedValue(null); + + const insertQuery = { + insertAndFetch: jest.fn().mockResolvedValue({ id: 1 }), + }; + + mockPendingActionQuery.mockImplementationOnce(() => existingLookupQuery).mockImplementationOnce(() => insertQuery); + + await ApprovalService.upsertApprovalRequestFromStream({ + thread: { id: 7 } as any, + run: { id: 11 } as any, + approvalId: 'approval-1', + toolCallId: 'tool-call-1', + toolName: 'mcp__sandbox__workspace_write_file', + input: { + path: 'sample-file.txt', + content: 'hello', + }, + fileChanges: [ + { + id: 'change-1', + toolCallId: 'tool-call-1', + sourceTool: 'workspace.write_file', + path: 'sample-file.txt', + displayPath: 'sample-file.txt', + kind: 'write', + stage: 'awaiting-approval', + } as any, + ], + approvalPolicy: { + defaultMode: 'allow', + rules: { + workspace_write: 'require_approval', + }, + } as any, + toolRules: [], + }); + + expect(insertQuery.insertAndFetch).toHaveBeenCalledWith( + expect.objectContaining({ + capabilityKey: 'workspace_write', + payload: expect.objectContaining({ + approvalId: 'approval-1', + toolCallId: 'tool-call-1', + toolName: 'mcp__sandbox__workspace_write_file', + fileChanges: expect.arrayContaining([ + expect.objectContaining({ + id: 'change-1', + path: 'sample-file.txt', + }), + ]), + }), + }) + ); + }); + + it('does not reset resolved approval requests during final message sync', async () => { + mockGetToolName.mockReturnValue('mcp__sandbox__workspace_write_file'); + + const existingLookupQuery: any = {}; + existingLookupQuery.where = jest.fn().mockReturnValue(existingLookupQuery); + existingLookupQuery.whereRaw = jest.fn().mockReturnValue(existingLookupQuery); + existingLookupQuery.first = jest.fn().mockResolvedValue({ + id: 1, + status: 'approved', + resolution: { approved: true }, + resolvedAt: '2026-04-12T00:00:00.000Z', + }); + existingLookupQuery.patchAndFetchById = jest.fn(); + + mockPendingActionQuery.mockReturnValueOnce(existingLookupQuery); + + await expect( + ApprovalService.syncApprovalRequestStateFromMessages({ + thread: { id: 7 } as any, + run: { id: 11 } as any, + messages: [ + { + role: 'assistant', + parts: [ + { + approval: { id: 'approval-1' }, + input: { + path: 'approval-check.txt', + content: 'hello', + }, + state: 'approval-requested', + toolCallId: 'tool-call-1', + }, + ], + } as any, ], + approvalPolicy: { + defaultMode: 'allow', + rules: { + workspace_write: 'require_approval', + }, + } as any, + toolRules: [], + }) + ).resolves.toEqual({ + pendingActions: [], + resolvedActionCount: 1, + }); + + expect(existingLookupQuery.patchAndFetchById).not.toHaveBeenCalled(); + }); + + it('classifies chat HTTP publish approvals as deploy mutations', async () => { + mockGetToolName.mockReturnValue('mcp__lifecycle__publish_http'); + + const existingLookupQuery: any = {}; + existingLookupQuery.where = jest.fn().mockReturnValue(existingLookupQuery); + existingLookupQuery.whereRaw = jest.fn().mockReturnValue(existingLookupQuery); + existingLookupQuery.first = jest.fn().mockResolvedValue(null); + + const insertQuery = { + insertAndFetch: jest.fn().mockResolvedValue({ id: 1 }), + }; + + mockPendingActionQuery.mockImplementationOnce(() => existingLookupQuery).mockImplementationOnce(() => insertQuery); + + await ApprovalService.upsertApprovalRequest({ + thread: { id: 7 } as any, + run: { id: 11 } as any, + message: { parts: [] } as any, + toolPart: { + approval: { id: 'approval-1' }, + input: { + port: 8000, + }, + state: 'approval-requested', + toolCallId: 'tool-call-1', } as any, - ]); + capabilityKey: 'external_mcp_write', + }); - expect(pendingLookupQuery.alias).toHaveBeenCalledWith('action'); - expect(pendingLookupQuery.whereRaw).toHaveBeenNthCalledWith(1, `action.payload->>'approvalId' = ?`, ['approval-1']); - expect(pendingLookupQuery.whereRaw).toHaveBeenNthCalledWith(2, `action.payload->>'toolCallId' = ?`, [ - 'tool-call-1', - ]); - expect(pendingLookupQuery.where).toHaveBeenCalledWith('run.uuid', 'run-uuid'); - expect(patchQuery.patchAndFetchById).toHaveBeenCalledWith( + expect(insertQuery.insertAndFetch).toHaveBeenCalledWith( + expect.objectContaining({ + capabilityKey: 'deploy_k8s_mutation', + payload: expect.objectContaining({ + toolName: 'mcp__lifecycle__publish_http', + }), + }) + ); + }); + + it('does not resume a waiting run while another pending action remains unresolved', async () => { + const action = { + id: 99, + uuid: 'action-1', + threadId: 7, + runId: 11, + status: 'pending', + payload: { approvalId: 'approval-1' }, + runUuid: 'run-uuid', + }; + const updatedAction = { + ...action, + status: 'approved', + }; + const pendingQuery = makeTransactionalPendingActionQuery(action, action, { id: 100 }, updatedAction); + const runQuery = makeTransactionalRunQuery({ + id: 11, + uuid: 'run-uuid', + status: 'waiting_for_approval', + usageSummary: {}, + error: null, + }); + + mockPendingActionQuery.mockReturnValue(pendingQuery); + mockRunQuery.mockReturnValue(runQuery); + + await ApprovalService.resolvePendingAction('action-1', 'sample-user', 'approved', { + approved: true, + }); + + expect(pendingQuery.where).toHaveBeenCalledWith({ runId: 11, status: 'pending' }); + expect(pendingQuery.patchAndFetchById).toHaveBeenCalledWith( 99, expect.objectContaining({ status: 'approved', - resolution: { - approved: true, - reason: 'approved in UI', - source: 'message', - }, + resolvedAt: expect.any(String), }) ); + expect(mockAppendStatusEventForRunInTransaction).toHaveBeenCalledWith( + expect.objectContaining({ uuid: 'run-uuid' }), + 'approval.resolved', + { + actionId: 'action-1', + approvalId: 'approval-1', + toolCallId: null, + approved: true, + reason: null, + }, + { trx: true } + ); + expect(mockAppendStatusEventForRunInTransaction).toHaveBeenCalledWith( + expect.objectContaining({ uuid: 'run-uuid' }), + 'approval.responded', + { + actionId: 'action-1', + approvalId: 'approval-1', + toolCallId: null, + approved: true, + reason: null, + }, + { trx: true } + ); + expect(mockPatchStatus).not.toHaveBeenCalled(); + expect(mockEnqueueRun).not.toHaveBeenCalled(); + }); + + it('resumes a waiting run after the last pending action is resolved', async () => { + const action = { + id: 99, + uuid: 'action-1', + threadId: 7, + runId: 11, + status: 'pending', + payload: { approvalId: 'approval-1' }, + runUuid: 'run-uuid', + }; + const updatedAction = { + ...action, + status: 'approved', + }; + const pendingQuery = makeTransactionalPendingActionQuery(action, action, null, updatedAction); + const queuedRun = { + id: 11, + uuid: 'run-uuid', + status: 'queued', + usageSummary: {}, + error: null, + }; + const runQuery = makeTransactionalRunQuery( + { + id: 11, + uuid: 'run-uuid', + status: 'waiting_for_approval', + usageSummary: {}, + error: null, + }, + queuedRun + ); + + mockEnqueueRun.mockResolvedValue(undefined); + mockPendingActionQuery.mockReturnValue(pendingQuery); + mockRunQuery.mockReturnValue(runQuery); + + await ApprovalService.resolvePendingAction( + 'action-1', + 'sample-user', + 'approved', + { + approved: true, + }, + { githubToken: 'sample-gh-token' } + ); + + expect(pendingQuery.where).toHaveBeenCalledWith({ runId: 11, status: 'pending' }); + expect(runQuery.patchAndFetchById).toHaveBeenCalledWith( + 11, + expect.objectContaining({ + status: 'queued', + queuedAt: expect.any(String), + executionOwner: null, + leaseExpiresAt: null, + heartbeatAt: null, + }) + ); + expect(mockEnqueueRun).toHaveBeenCalledWith('run-uuid', 'approval_resolved', { + githubToken: 'sample-gh-token', + }); + expect(mockAppendStatusEventForRunInTransaction).toHaveBeenCalledWith( + queuedRun, + 'run.queued', + expect.objectContaining({ + status: 'queued', + }), + { trx: true } + ); + }); + + it('emits the denial reason before requeueing the waiting run', async () => { + const action = { + id: 99, + uuid: 'action-1', + threadId: 7, + runId: 11, + status: 'pending', + payload: { approvalId: 'approval-1' }, + runUuid: 'run-uuid', + }; + const updatedAction = { + ...action, + status: 'denied', + resolution: { + approved: false, + reason: 'not needed', + }, + }; + const queuedRun = { + id: 11, + uuid: 'run-uuid', + status: 'queued', + usageSummary: {}, + error: null, + }; + const pendingQuery = makeTransactionalPendingActionQuery(action, action, null, updatedAction); + const runQuery = makeTransactionalRunQuery( + { + id: 11, + uuid: 'run-uuid', + status: 'waiting_for_approval', + usageSummary: {}, + error: null, + }, + queuedRun + ); + + mockEnqueueRun.mockResolvedValue(undefined); + mockPendingActionQuery.mockReturnValue(pendingQuery); + mockRunQuery.mockReturnValue(runQuery); + + await ApprovalService.resolvePendingAction('action-1', 'sample-user', 'denied', { + approved: false, + reason: 'not needed', + }); + + expect(mockAppendStatusEventForRunInTransaction).toHaveBeenCalledWith( + expect.objectContaining({ uuid: 'run-uuid' }), + 'approval.resolved', + { + actionId: 'action-1', + approvalId: 'approval-1', + toolCallId: null, + approved: false, + reason: 'not needed', + }, + { trx: true } + ); + expect(mockAppendStatusEventForRunInTransaction).toHaveBeenCalledWith( + expect.objectContaining({ uuid: 'run-uuid' }), + 'approval.responded', + { + actionId: 'action-1', + approvalId: 'approval-1', + toolCallId: null, + approved: false, + reason: 'not needed', + }, + { trx: true } + ); + expect(runQuery.patchAndFetchById).toHaveBeenCalledWith( + 11, + expect.objectContaining({ + status: 'queued', + }) + ); + expect(mockEnqueueRun).toHaveBeenCalledWith('run-uuid', 'approval_resolved', { + githubToken: undefined, + }); + }); + + it('resumes a waiting run from an already resolved action without duplicate approval side effects', async () => { + const action = { + id: 99, + uuid: 'action-1', + threadId: 7, + runId: 11, + status: 'approved', + payload: { approvalId: 'approval-1' }, + runUuid: 'run-uuid', + resolution: { + approved: true, + }, + }; + const queuedRun = { + id: 11, + uuid: 'run-uuid', + status: 'queued', + usageSummary: {}, + error: null, + }; + const pendingQuery = makeTransactionalPendingActionQuery(action, action, null); + const runQuery = makeTransactionalRunQuery( + { + id: 11, + uuid: 'run-uuid', + status: 'waiting_for_approval', + usageSummary: {}, + error: null, + }, + queuedRun + ); + mockPendingActionQuery.mockReturnValue(pendingQuery); + mockRunQuery.mockReturnValue(runQuery); + + await expect( + ApprovalService.resolvePendingAction('action-1', 'sample-user', 'denied', { + approved: false, + }) + ).resolves.toBe(action); + + expect(pendingQuery.patchAndFetchById).not.toHaveBeenCalled(); + expect(runQuery.patchAndFetchById).toHaveBeenCalledWith( + 11, + expect.objectContaining({ + status: 'queued', + queuedAt: expect.any(String), + }) + ); + expect(mockAppendStatusEventForRunInTransaction).toHaveBeenCalledWith( + queuedRun, + 'run.queued', + expect.objectContaining({ + status: 'queued', + }), + { trx: true } + ); + expect(mockAppendStatusEventForRunInTransaction).not.toHaveBeenCalledWith( + expect.anything(), + 'approval.resolved', + expect.anything(), + expect.anything() + ); + expect(mockAppendStatusEventForRunInTransaction).not.toHaveBeenCalledWith( + expect.anything(), + 'approval.responded', + expect.anything(), + expect.anything() + ); + expect(mockPatchStatus).not.toHaveBeenCalled(); + expect(mockEnqueueRun).toHaveBeenCalledWith('run-uuid', 'approval_resolved', { + githubToken: undefined, + }); + }); + + it('requeues an already queued run from an already resolved action', async () => { + const resolvedAction = { + id: 99, + uuid: 'action-1', + threadId: 7, + runId: 11, + status: 'approved', + payload: { approvalId: 'approval-1' }, + runUuid: 'run-uuid', + resolution: { + approved: true, + }, + }; + const pendingQuery = makeTransactionalPendingActionQuery(resolvedAction, resolvedAction, null); + const runQuery = makeTransactionalRunQuery({ + id: 11, + uuid: 'run-uuid', + status: 'queued', + usageSummary: {}, + error: null, + }); + mockPendingActionQuery.mockReturnValue(pendingQuery); + mockRunQuery.mockReturnValue(runQuery); + + await expect( + ApprovalService.resolvePendingAction('action-1', 'sample-user', 'approved', { + approved: true, + }) + ).resolves.toBe(resolvedAction); + + expect(pendingQuery.patchAndFetchById).not.toHaveBeenCalled(); + expect(runQuery.patchAndFetchById).not.toHaveBeenCalled(); + expect(mockAppendStatusEventForRunInTransaction).not.toHaveBeenCalled(); + expect(mockEnqueueRun).toHaveBeenCalledWith('run-uuid', 'approval_resolved', { + githubToken: undefined, + }); + }); + + it('does not emit approval side effects when the locked action is already resolved', async () => { + const resolvedAction = { + id: 99, + uuid: 'action-1', + threadId: 7, + runId: 11, + status: 'denied', + payload: { approvalId: 'approval-1' }, + runUuid: 'run-uuid', + resolution: { + approved: false, + }, + }; + const pendingQuery = makeTransactionalPendingActionQuery(resolvedAction, resolvedAction); + const runQuery = makeTransactionalRunQuery({ + id: 11, + uuid: 'run-uuid', + status: 'completed', + usageSummary: {}, + error: null, + }); + mockPendingActionQuery.mockReturnValue(pendingQuery); + mockRunQuery.mockReturnValue(runQuery); + + await expect( + ApprovalService.resolvePendingAction('action-1', 'sample-user', 'approved', { + approved: true, + }) + ).resolves.toBe(resolvedAction); + + expect(mockAppendStatusEventForRunInTransaction).not.toHaveBeenCalled(); + expect(mockPatchStatus).not.toHaveBeenCalled(); + expect(mockEnqueueRun).not.toHaveBeenCalled(); }); }); diff --git a/src/server/services/agent/__tests__/CapabilityService.test.ts b/src/server/services/agent/__tests__/CapabilityService.test.ts index 4f79e353..ac0a1940 100644 --- a/src/server/services/agent/__tests__/CapabilityService.test.ts +++ b/src/server/services/agent/__tests__/CapabilityService.test.ts @@ -16,13 +16,17 @@ const mockDynamicTool = jest.fn((config) => config); const mockJsonSchema = jest.fn((schema) => schema); -const mockResolveServersForRepo = jest.fn(); +const mockResolveServers = jest.fn(); const mockConnect = jest.fn(); const mockListTools = jest.fn(); const mockCallTool = jest.fn(); const mockClose = jest.fn(); const mockLoggerWarn = jest.fn(); const mockModeForCapability = jest.fn(() => 'allow'); +const mockPublishChatHttpPort = jest.fn(); +const mockFindSession = jest.fn(); +const mockResolveWorkspaceGatewayBaseUrl = jest.fn(); +const mockEnsureChatSandbox = jest.fn(); let currentTransport: Record | null = null; @@ -33,10 +37,34 @@ jest.mock('ai', () => ({ jest.mock('server/services/ai/mcp/config', () => ({ McpConfigService: jest.fn().mockImplementation(() => ({ - resolveServersForRepo: (...args: unknown[]) => mockResolveServersForRepo(...args), + resolveServers: (...args: unknown[]) => mockResolveServers(...args), })), })); +jest.mock('server/models/AgentSession', () => ({ + __esModule: true, + default: { + query: jest.fn(() => ({ + findOne: (...args: unknown[]) => mockFindSession(...args), + })), + }, +})); + +jest.mock('server/services/agentSession', () => ({ + __esModule: true, + default: { + publishChatHttpPort: (...args: unknown[]) => mockPublishChatHttpPort(...args), + }, +})); + +jest.mock('../SandboxService', () => ({ + __esModule: true, + default: { + resolveWorkspaceGatewayBaseUrl: (...args: unknown[]) => mockResolveWorkspaceGatewayBaseUrl(...args), + ensureChatSandbox: (...args: unknown[]) => mockEnsureChatSandbox(...args), + }, +})); + jest.mock('server/services/ai/mcp/client', () => ({ McpClientManager: jest.fn().mockImplementation(() => ({ connect: (...args: unknown[]) => mockConnect(...args), @@ -54,6 +82,22 @@ jest.mock('server/lib/logger', () => ({ }), })); +jest.mock('server/lib/agentSession/runtimeConfig', () => { + const actual = jest.requireActual('server/lib/agentSession/runtimeConfig'); + return { + __esModule: true, + ...actual, + resolveAgentSessionDurabilityConfig: jest.fn().mockResolvedValue({ + runExecutionLeaseMs: 30 * 60 * 1000, + queuedRunDispatchStaleMs: 30 * 1000, + dispatchRecoveryLimit: 50, + maxDurablePayloadBytes: 64 * 1024, + payloadPreviewBytes: 16 * 1024, + fileChangePreviewChars: 4000, + }), + }; +}); + jest.mock('../PolicyService', () => ({ __esModule: true, default: { @@ -72,6 +116,8 @@ describe('AgentCapabilityService.buildToolSet', () => { podName: 'agent-123', namespace: 'env-sample', status: 'active', + sessionKind: 'environment', + workspaceStatus: 'ready', } as any; const userIdentity = { userId: 'sample-user', @@ -112,7 +158,35 @@ describe('AgentCapabilityService.buildToolSet', () => { jest.clearAllMocks(); mockModeForCapability.mockReturnValue('allow'); currentTransport = null; - mockResolveServersForRepo.mockResolvedValue([stdioServer]); + mockResolveServers.mockResolvedValue([stdioServer]); + mockFindSession.mockResolvedValue(session); + mockPublishChatHttpPort.mockResolvedValue({ + url: 'https://chat-session.example.test', + host: 'chat-session.example.test', + path: '/', + port: 3000, + serviceName: 'agent-preview-sample', + ingressName: 'agent-preview-ingress-sample', + }); + mockResolveWorkspaceGatewayBaseUrl.mockImplementation(async (sessionUuid: string) => { + if (sessionUuid === 'session-chat') { + return 'http://agent-chat.chat-sample.svc.cluster.local:13338'; + } + + return 'http://agent-123.env-sample.svc.cluster.local:13338'; + }); + mockEnsureChatSandbox.mockResolvedValue({ + session: { + uuid: 'session-chat', + sessionKind: 'chat', + workspaceStatus: 'ready', + status: 'active', + podName: 'agent-chat', + namespace: 'chat-sample', + pvcName: 'agent-pvc-sample', + }, + sandbox: null, + }); mockConnect.mockImplementation(async (transport) => { currentTransport = transport as Record; }); @@ -258,4 +332,467 @@ describe('AgentCapabilityService.buildToolSet', () => { }) ); }); + + it('keeps global MCP tools available even when the session has no primary repo', async () => { + mockModeForCapability.mockImplementation((_policy, capability) => + capability === 'deploy_k8s_mutation' ? 'require_approval' : 'allow' + ); + mockResolveServers.mockResolvedValue([ + { + slug: 'docs', + name: 'Docs', + transport: { + type: 'http', + url: 'https://mcp.example.test', + }, + timeout: 30000, + defaultArgs: {}, + env: {}, + discoveredTools: [ + { + name: 'search_docs', + description: 'Search docs', + inputSchema: { + type: 'object', + properties: {}, + }, + annotations: { + readOnlyHint: true, + }, + }, + ], + }, + ]); + + const tools = await AgentCapabilityService.buildToolSet({ + session: { + ...session, + sessionKind: 'chat', + podName: null, + namespace: null, + workspaceStatus: 'none', + } as any, + repoFullName: undefined, + userIdentity, + approvalPolicy: {} as any, + workspaceToolDiscoveryTimeoutMs: 4500, + workspaceToolExecutionTimeoutMs: 22000, + }); + + expect(mockResolveServers).toHaveBeenCalledWith(undefined, undefined, userIdentity); + expect(tools.mcp__docs__search_docs).toBeDefined(); + expect(tools.mcp__sandbox__workspace_exec).toEqual( + expect.objectContaining({ + needsApproval: false, + }) + ); + expect(tools.mcp__sandbox__workspace_exec_mutation).toEqual( + expect.objectContaining({ + needsApproval: false, + }) + ); + expect(tools.mcp__lifecycle__publish_http).toEqual( + expect.objectContaining({ + needsApproval: true, + }) + ); + expect(Object.keys(tools).some((key) => key.includes('__source_'))).toBe(false); + expect(tools.mcp__lifecycle__workspace_provision).toBeUndefined(); + }); + + it('lets tool rules require approval for chat HTTP publishing', async () => { + mockResolveServers.mockResolvedValue([]); + mockModeForCapability.mockReturnValue('allow'); + + const tools = await AgentCapabilityService.buildToolSet({ + session: { + ...session, + sessionKind: 'chat', + } as any, + repoFullName: undefined, + userIdentity, + approvalPolicy: {} as any, + toolRules: [ + { + toolKey: 'mcp__lifecycle__publish_http', + mode: 'require_approval', + }, + ], + workspaceToolDiscoveryTimeoutMs: 4500, + workspaceToolExecutionTimeoutMs: 22000, + }); + + expect(tools.mcp__lifecycle__publish_http).toEqual( + expect.objectContaining({ + needsApproval: true, + }) + ); + }); + + it('exposes lazy chat workspace tools before runtime without provisioning during setup', async () => { + mockResolveServers.mockResolvedValue([]); + mockFindSession.mockResolvedValue({ + uuid: 'session-chat', + sessionKind: 'chat', + workspaceStatus: 'none', + status: 'active', + podName: null, + namespace: null, + }); + + const tools = await AgentCapabilityService.buildToolSet({ + session: { + uuid: 'session-chat', + sessionKind: 'chat', + workspaceStatus: 'none', + status: 'active', + podName: null, + namespace: null, + } as any, + repoFullName: undefined, + userIdentity, + approvalPolicy: {} as any, + workspaceToolDiscoveryTimeoutMs: 4500, + workspaceToolExecutionTimeoutMs: 22000, + requestGitHubToken: 'sample-gh-token', + }); + + expect(tools.mcp__lifecycle__workspace_provision).toBeUndefined(); + expect(mockEnsureChatSandbox).not.toHaveBeenCalled(); + expect(tools.mcp__sandbox__workspace_exec).toEqual( + expect.objectContaining({ + needsApproval: false, + }) + ); + expect(tools.mcp__sandbox__workspace_exec_mutation).toEqual( + expect.objectContaining({ + needsApproval: false, + }) + ); + expect(tools.mcp__sandbox__workspace_write_file).toEqual( + expect.objectContaining({ + needsApproval: false, + }) + ); + expect(tools.mcp__sandbox__workspace_edit_file).toEqual( + expect.objectContaining({ + needsApproval: false, + }) + ); + expect(tools.mcp__lifecycle__publish_http).toEqual( + expect.objectContaining({ + needsApproval: false, + }) + ); + expect(Object.keys(tools).some((key) => key.includes('__source_'))).toBe(false); + + const tool = tools.mcp__sandbox__workspace_write_file as { + execute: (input: Record, context?: { toolCallId?: string }) => Promise; + }; + + await expect( + tool.execute( + { + path: 'sample.txt', + content: 'hello', + }, + { toolCallId: 'tool-write' } + ) + ).resolves.toEqual({ + content: [{ type: 'text', text: 'ok' }], + isError: false, + }); + + expect(mockEnsureChatSandbox).toHaveBeenCalledWith({ + sessionId: 'session-chat', + userId: 'sample-user', + userIdentity, + githubToken: 'sample-gh-token', + }); + expect(mockCallTool).toHaveBeenCalledWith( + 'workspace.write_file', + { + path: 'sample.txt', + content: 'hello', + }, + 22000 + ); + }); + + it('lets tool rules require approval for lazy chat workspace tools before runtime exists', async () => { + mockResolveServers.mockResolvedValue([]); + + const tools = await AgentCapabilityService.buildToolSet({ + session: { + uuid: 'session-chat', + sessionKind: 'chat', + workspaceStatus: 'none', + status: 'active', + podName: null, + namespace: null, + } as any, + repoFullName: undefined, + userIdentity, + approvalPolicy: {} as any, + toolRules: [ + { + toolKey: 'mcp__sandbox__workspace_write_file', + mode: 'require_approval', + }, + ], + workspaceToolDiscoveryTimeoutMs: 4500, + workspaceToolExecutionTimeoutMs: 22000, + }); + + expect(tools.mcp__sandbox__workspace_write_file).toEqual( + expect.objectContaining({ + needsApproval: true, + }) + ); + }); + + it('runs GitHub CLI commands through the generic workspace mutation tool with request GitHub auth', async () => { + mockResolveServers.mockResolvedValue([]); + + const tools = await AgentCapabilityService.buildToolSet({ + session: { + uuid: 'session-chat', + sessionKind: 'chat', + workspaceStatus: 'none', + status: 'active', + podName: null, + namespace: null, + } as any, + repoFullName: undefined, + userIdentity, + approvalPolicy: {} as any, + workspaceToolDiscoveryTimeoutMs: 4500, + workspaceToolExecutionTimeoutMs: 22000, + requestGitHubToken: 'sample-gh-token', + }); + + const tool = tools.mcp__sandbox__workspace_exec_mutation as { + execute: (input: Record) => Promise; + }; + mockFindSession.mockResolvedValueOnce({ + uuid: 'session-chat', + sessionKind: 'chat', + workspaceStatus: 'none', + status: 'active', + podName: null, + namespace: null, + }); + + await tool.execute({ + command: 'gh repo clone example-org/private-repo private-repo', + cwd: '.', + }); + + expect(mockEnsureChatSandbox).toHaveBeenCalledWith({ + sessionId: 'session-chat', + userId: 'sample-user', + userIdentity, + githubToken: 'sample-gh-token', + }); + expect(mockCallTool).toHaveBeenCalledWith( + 'workspace.exec', + { + command: 'gh repo clone example-org/private-repo private-repo', + cwd: '.', + captureFileChanges: true, + }, + 22000 + ); + }); + + it('emits file changes returned by lazy chat workspace mutation commands', async () => { + mockResolveServers.mockResolvedValue([]); + mockCallTool.mockResolvedValueOnce({ + content: [ + { + type: 'text', + text: JSON.stringify({ + ok: true, + command: "printf 'number 1\\n' > fresh-e2e-artifacts/numbers.txt", + success: true, + fileChanges: [ + { + path: 'fresh-e2e-artifacts/numbers.txt', + kind: 'created', + additions: 1, + deletions: 0, + beforeTextPreview: '', + afterTextPreview: 'number 1\n', + summary: 'Created fresh-e2e-artifacts/numbers.txt', + }, + ], + }), + }, + ], + isError: false, + }); + const onFileChange = jest.fn(); + + const tools = await AgentCapabilityService.buildToolSet({ + session: { + uuid: 'session-chat', + sessionKind: 'chat', + workspaceStatus: 'ready', + status: 'active', + podName: 'agent-chat', + namespace: 'chat-sample', + } as any, + repoFullName: undefined, + userIdentity, + approvalPolicy: {} as any, + workspaceToolDiscoveryTimeoutMs: 4500, + workspaceToolExecutionTimeoutMs: 22000, + hooks: { + onFileChange, + }, + }); + + const tool = tools.mcp__sandbox__workspace_exec_mutation as { + execute: (input: Record, context?: { toolCallId?: string }) => Promise; + }; + + await tool.execute( + { + command: "printf 'number 1\\n' > fresh-e2e-artifacts/numbers.txt", + }, + { toolCallId: 'tool-call-1' } + ); + + expect(mockCallTool).toHaveBeenCalledWith( + 'workspace.exec', + { + command: "printf 'number 1\\n' > fresh-e2e-artifacts/numbers.txt", + captureFileChanges: true, + }, + 22000 + ); + expect(onFileChange).toHaveBeenCalledWith( + expect.objectContaining({ + id: 'tool-call-1:fresh-e2e-artifacts/numbers.txt', + toolCallId: 'tool-call-1', + sourceTool: 'workspace.exec_mutation', + path: 'fresh-e2e-artifacts/numbers.txt', + kind: 'created', + additions: 1, + stage: 'applied', + }) + ); + }); + + it('emits file changes returned by discovered workspace mutation commands', async () => { + mockResolveServers.mockResolvedValue([]); + mockListTools.mockResolvedValueOnce([ + { + name: 'workspace.exec', + inputSchema: { + type: 'object', + properties: { + command: { type: 'string' }, + }, + }, + annotations: { + destructiveHint: true, + }, + }, + ]); + mockCallTool.mockResolvedValueOnce({ + content: [ + { + type: 'text', + text: JSON.stringify({ + ok: true, + command: "printf 'number 1\\n' > fresh-e2e-artifacts/numbers.txt", + success: true, + fileChanges: [ + { + path: 'fresh-e2e-artifacts/numbers.txt', + kind: 'created', + additions: 1, + deletions: 0, + }, + ], + }), + }, + ], + isError: false, + }); + const onFileChange = jest.fn(); + + const tools = await AgentCapabilityService.buildToolSet({ + session, + repoFullName: undefined, + userIdentity, + approvalPolicy: {} as any, + workspaceToolDiscoveryTimeoutMs: 4500, + workspaceToolExecutionTimeoutMs: 22000, + hooks: { + onFileChange, + }, + }); + + const tool = tools.mcp__sandbox__workspace_exec_mutation as { + execute: (input: Record, context?: { toolCallId?: string }) => Promise; + }; + + await tool.execute( + { + command: "printf 'number 1\\n' > fresh-e2e-artifacts/numbers.txt", + }, + { toolCallId: 'tool-call-2' } + ); + + expect(mockCallTool).toHaveBeenCalledWith( + 'workspace.exec', + { + command: "printf 'number 1\\n' > fresh-e2e-artifacts/numbers.txt", + captureFileChanges: true, + }, + 22000 + ); + expect(onFileChange).toHaveBeenCalledWith( + expect.objectContaining({ + id: 'tool-call-2:fresh-e2e-artifacts/numbers.txt', + toolCallId: 'tool-call-2', + sourceTool: 'workspace.exec_mutation', + path: 'fresh-e2e-artifacts/numbers.txt', + kind: 'created', + stage: 'applied', + }) + ); + }); + + it('blocks unsafe broad process kill commands before they reach the workspace gateway', async () => { + mockResolveServers.mockResolvedValue([]); + + const tools = await AgentCapabilityService.buildToolSet({ + session: { + uuid: 'session-chat', + sessionKind: 'chat', + workspaceStatus: 'ready', + status: 'active', + podName: 'agent-chat', + namespace: 'chat-sample', + } as any, + repoFullName: undefined, + userIdentity, + approvalPolicy: {} as any, + workspaceToolDiscoveryTimeoutMs: 4500, + workspaceToolExecutionTimeoutMs: 22000, + }); + + const tool = tools.mcp__sandbox__workspace_exec_mutation as { + execute: (input: Record) => Promise; + }; + + mockConnect.mockClear(); + mockCallTool.mockClear(); + await expect(tool.execute({ command: 'kill -9 $(pidof node)' })).rejects.toThrow('workspace gateway'); + expect(mockConnect).not.toHaveBeenCalled(); + expect(mockCallTool).not.toHaveBeenCalled(); + }); }); diff --git a/src/server/services/agent/__tests__/LifecycleAiSdkHarness.test.ts b/src/server/services/agent/__tests__/LifecycleAiSdkHarness.test.ts new file mode 100644 index 00000000..9a4ab87a --- /dev/null +++ b/src/server/services/agent/__tests__/LifecycleAiSdkHarness.test.ts @@ -0,0 +1,325 @@ +/** + * Copyright 2026 GoodRx, Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +var mockCreateAgentUIStream: jest.Mock; +var mockCreateUIMessageStream: jest.Mock; +var mockSafeValidateUIMessages: jest.Mock; +var mockReadUIMessageStream: jest.Mock; + +jest.mock('ai', () => ({ + __esModule: true, + createAgentUIStream: (mockCreateAgentUIStream = jest.fn()), + createUIMessageStream: (mockCreateUIMessageStream = jest.fn()), + readUIMessageStream: (mockReadUIMessageStream = jest.fn()), + safeValidateUIMessages: (mockSafeValidateUIMessages = jest.fn()), +})); + +jest.mock('server/models/AgentSession', () => ({ + __esModule: true, + default: { + query: jest.fn(), + }, +})); + +jest.mock('server/models/AgentThread', () => ({ + __esModule: true, + default: { + query: jest.fn(), + }, +})); + +jest.mock('../MessageStore', () => ({ + __esModule: true, + default: { + listMessages: jest.fn(), + }, +})); + +jest.mock('../RunExecutor', () => ({ + __esModule: true, + default: { + execute: jest.fn(), + }, +})); + +jest.mock('../RunService', () => ({ + __esModule: true, + default: { + appendStreamChunksForExecutionOwner: jest.fn(), + }, +})); + +jest.mock('../RunEventService', () => ({ + __esModule: true, + default: { + listRunEventsPage: jest.fn(), + projectUiChunksFromEvents: jest.fn(), + }, +})); + +jest.mock('../ApprovalService', () => ({ + __esModule: true, + default: { + upsertApprovalRequestFromStream: jest.fn(), + }, +})); + +import AgentSession from 'server/models/AgentSession'; +import AgentThread from 'server/models/AgentThread'; +import AgentMessageStore from '../MessageStore'; +import AgentRunExecutor from '../RunExecutor'; +import AgentRunService from '../RunService'; +import type { AgentUIMessage } from '../types'; +import LifecycleAiSdkHarness from '../LifecycleAiSdkHarness'; +import { + applyApprovalResponsesToToolParts, + normalizeUnavailableToolPartsForAgentInput, +} from '../LifecycleAiSdkHarness'; + +const mockSessionQuery = AgentSession.query as jest.Mock; +const mockThreadQuery = AgentThread.query as jest.Mock; +const mockListMessages = AgentMessageStore.listMessages as jest.Mock; +const mockExecuteRun = AgentRunExecutor.execute as jest.Mock; +const mockAppendStreamChunksForExecutionOwner = AgentRunService.appendStreamChunksForExecutionOwner as jest.Mock; + +beforeEach(() => { + jest.clearAllMocks(); + mockReadUIMessageStream.mockImplementation(async function* () {}); +}); + +describe('LifecycleAiSdkHarness.executeRun', () => { + it('flushes stream chunks before finalizing a waiting approval run', async () => { + const operations: string[] = []; + const userMessage = { + id: 'user-1', + role: 'user', + parts: [{ type: 'text', text: 'Create a simple web app.' }], + } as AgentUIMessage; + const finalMessages = [ + userMessage, + { + id: 'assistant-1', + role: 'assistant', + parts: [{ type: 'text', text: 'Done.' }], + }, + ] as AgentUIMessage[]; + const onStreamFinish = jest.fn(async () => { + operations.push('finalize'); + }); + + mockSessionQuery.mockReturnValue({ + findById: jest.fn().mockResolvedValue({ + id: 13, + uuid: 'session-1', + userId: 'sample-user', + ownerGithubUsername: null, + }), + }); + mockThreadQuery.mockReturnValue({ + findById: jest.fn().mockResolvedValue({ + id: 17, + uuid: 'thread-1', + }), + }); + mockListMessages.mockResolvedValue([userMessage]); + mockSafeValidateUIMessages.mockResolvedValue({ + success: true, + data: [userMessage], + }); + mockCreateAgentUIStream.mockResolvedValue( + new ReadableStream({ + start(controller) { + controller.close(); + }, + }) + ); + mockCreateUIMessageStream.mockImplementation(({ onFinish }) => { + return new ReadableStream({ + async start(controller) { + controller.enqueue({ type: 'text-delta', id: 'text-1', delta: 'Done.' }); + await onFinish({ messages: finalMessages }); + controller.close(); + }, + }); + }); + mockExecuteRun.mockResolvedValue({ + run: { + id: 19, + uuid: 'run-1', + executionOwner: 'owner-1', + }, + agent: { + tools: {}, + }, + abortSignal: new AbortController().signal, + selection: { + provider: 'openai', + modelId: 'gpt-5.4', + }, + approvalPolicy: { + rules: {}, + defaultMode: 'require_approval', + }, + toolRules: [], + onStreamFinish, + dispose: jest.fn(), + }); + mockAppendStreamChunksForExecutionOwner.mockImplementation(async () => { + operations.push('append'); + return { id: 19, uuid: 'run-1' }; + }); + + await LifecycleAiSdkHarness.executeRun({ + id: 19, + uuid: 'run-1', + threadId: 17, + sessionId: 13, + startedAt: null, + } as any); + + expect(operations).toEqual(['append', 'finalize']); + expect(onStreamFinish).toHaveBeenCalledWith({ + messages: finalMessages, + finishReason: undefined, + isAborted: false, + }); + }); +}); + +describe('applyApprovalResponsesToToolParts', () => { + it('hydrates approved output tool parts so continuation messages validate', () => { + const message = { + id: 'assistant-1', + role: 'assistant', + parts: [ + { + type: 'dynamic-tool', + toolName: 'mcp__sandbox__workspace_write_file', + toolCallId: 'call-1', + state: 'output-error', + input: { + path: 'sample.txt', + content: 'hello', + }, + errorText: 'Session workspace gateway unavailable.', + approval: { + id: 'approval-1', + }, + }, + ], + } as AgentUIMessage; + + const result = applyApprovalResponsesToToolParts( + message, + new Map([ + [ + 'approval-1', + { + approved: true, + reason: 'Looks fine', + }, + ], + ]) + ); + + expect(result.parts[0]).toEqual( + expect.objectContaining({ + state: 'output-error', + approval: { + id: 'approval-1', + approved: true, + reason: 'Looks fine', + }, + }) + ); + }); + + it('marks pending approval parts as responded for resumed runs', () => { + const message = { + id: 'assistant-1', + role: 'assistant', + parts: [ + { + type: 'tool-mcp__sandbox__workspace_write_file', + toolCallId: 'call-1', + state: 'approval-requested', + input: { + path: 'sample.txt', + content: 'hello', + }, + approval: { + id: 'approval-1', + }, + }, + ], + } as AgentUIMessage; + + const result = applyApprovalResponsesToToolParts( + message, + new Map([ + [ + 'approval-1', + { + approved: false, + reason: 'Not needed', + }, + ], + ]) + ); + + expect(result.parts[0]).toEqual( + expect.objectContaining({ + state: 'approval-responded', + approval: { + id: 'approval-1', + approved: false, + reason: 'Not needed', + }, + }) + ); + }); +}); + +describe('normalizeUnavailableToolPartsForAgentInput', () => { + it('converts unavailable static tool parts to dynamic tool parts for continuation', () => { + const message = { + id: 'assistant-1', + role: 'assistant', + parts: [ + { + type: 'tool-mcp__sandbox__lifecycle__publish_http', + toolCallId: 'call-1', + state: 'output-error', + errorText: 'Model tried to call unavailable tool.', + }, + ], + } as unknown as AgentUIMessage; + + const [result] = normalizeUnavailableToolPartsForAgentInput([message], { + mcp__lifecycle__publish_http: {} as never, + }); + + expect(result.parts[0]).toEqual( + expect.objectContaining({ + type: 'dynamic-tool', + toolName: 'mcp__sandbox__lifecycle__publish_http', + toolCallId: 'call-1', + state: 'output-error', + input: undefined, + }) + ); + }); +}); diff --git a/src/server/services/agent/__tests__/MessageStore.test.ts b/src/server/services/agent/__tests__/MessageStore.test.ts index d4062288..83601d08 100644 --- a/src/server/services/agent/__tests__/MessageStore.test.ts +++ b/src/server/services/agent/__tests__/MessageStore.test.ts @@ -35,32 +35,356 @@ jest.mock('../ThreadService', () => ({ }, })); +jest.mock('uuid', () => ({ + v4: jest.fn(() => '11111111-1111-4111-8111-111111111111'), +})); + import AgentMessage from 'server/models/AgentMessage'; import AgentMessageStore from '../MessageStore'; +import AgentThreadService from '../ThreadService'; const mockMessageQuery = AgentMessage.query as jest.Mock; +const mockGetOwnedThread = AgentThreadService.getOwnedThread as jest.Mock; describe('AgentMessageStore', () => { beforeEach(() => { jest.clearAllMocks(); }); - describe('listRunMessages', () => { - it('filters through the joined run thread session alias', async () => { + describe('serializeCanonicalMessage', () => { + it('returns the public canonical message shape', () => { + expect( + AgentMessageStore.serializeCanonicalMessage( + { + uuid: '22222222-2222-4222-8222-222222222222', + clientMessageId: 'client-message-1', + role: 'user', + parts: [{ type: 'text', text: 'Hello' }], + metadata: { ignored: true }, + createdAt: '2026-04-25T00:00:00.000Z', + } as any, + 'thread-uuid', + 'run-uuid' + ) + ).toEqual({ + id: '22222222-2222-4222-8222-222222222222', + clientMessageId: 'client-message-1', + threadId: 'thread-uuid', + runId: 'run-uuid', + role: 'user', + parts: [{ type: 'text', text: 'Hello' }], + createdAt: '2026-04-25T00:00:00.000Z', + }); + }); + }); + + describe('listMessages', () => { + it('omits stored messages with no canonical parts', async () => { + const orderBy = jest.fn().mockResolvedValue([ + { + id: 11, + uuid: '22222222-2222-4222-8222-222222222222', + threadId: 17, + role: 'assistant', + parts: [], + metadata: { runId: 'run-empty' }, + }, + { + id: 12, + uuid: '33333333-3333-4333-8333-333333333333', + threadId: 17, + role: 'user', + parts: [{ type: 'text', text: 'Continue' }], + metadata: {}, + }, + ]); + const where = jest.fn().mockReturnValue({ orderBy }); + + mockGetOwnedThread.mockResolvedValue({ id: 17, uuid: 'thread-uuid' }); + mockMessageQuery.mockReturnValueOnce({ where }); + + const result = await AgentMessageStore.listMessages('thread-uuid', 'sample-user'); + + expect(result).toEqual([ + expect.objectContaining({ + id: '33333333-3333-4333-8333-333333333333', + role: 'user', + parts: [{ type: 'text', text: 'Continue' }], + }), + ]); + }); + }); + + describe('syncCanonicalMessages', () => { + it('persists canonical parts without uiMessage as the source of truth', async () => { + const insertedRow = { + id: 11, + uuid: '11111111-1111-4111-8111-111111111111', + threadId: 17, + role: 'user', + parts: [{ type: 'text', text: 'Hello' }], + uiMessage: null, + metadata: { clientMessageId: 'client-message-1' }, + }; + const insert = jest.fn().mockResolvedValue(insertedRow); + const orderBy = jest.fn().mockResolvedValue([insertedRow]); + const existingWhere = jest.fn().mockResolvedValue([]); + const reloadedWhere = jest.fn().mockReturnValue({ orderBy }); + + mockGetOwnedThread.mockResolvedValue({ id: 17, uuid: 'thread-uuid' }); + mockMessageQuery + .mockReturnValueOnce({ where: existingWhere }) + .mockReturnValueOnce({ insert }) + .mockReturnValueOnce({ where: reloadedWhere }); + + const result = await AgentMessageStore.syncCanonicalMessages('thread-uuid', 'sample-user', [ + { + id: 'client-message-1', + role: 'user', + parts: [{ type: 'text', text: 'Hello' }], + }, + ]); + + expect(insert).toHaveBeenCalledWith( + expect.objectContaining({ + uuid: '11111111-1111-4111-8111-111111111111', + threadId: 17, + role: 'user', + parts: [{ type: 'text', text: 'Hello' }], + uiMessage: null, + metadata: { clientMessageId: 'client-message-1' }, + }) + ); + expect(result[0]).toEqual( + expect.objectContaining({ + id: '11111111-1111-4111-8111-111111111111', + role: 'user', + parts: [{ type: 'text', text: 'Hello' }], + metadata: { clientMessageId: 'client-message-1' }, + }) + ); + }); + + it('reuses existing rows by canonical client message id', async () => { + const existingRow = { + id: 11, + uuid: '22222222-2222-4222-8222-222222222222', + runId: null, + role: 'user', + parts: [{ type: 'text', text: 'Before' }], + uiMessage: null, + metadata: { clientMessageId: 'client-message-1' }, + }; + const patchAndFetchById = jest.fn().mockResolvedValue(existingRow); + const orderBy = jest.fn().mockResolvedValue([ + { + ...existingRow, + parts: [{ type: 'text', text: 'After' }], + }, + ]); + const existingWhere = jest.fn().mockResolvedValue([existingRow]); + const reloadedWhere = jest.fn().mockReturnValue({ orderBy }); + + mockGetOwnedThread.mockResolvedValue({ id: 17, uuid: 'thread-uuid' }); + mockMessageQuery + .mockReturnValueOnce({ where: existingWhere }) + .mockReturnValueOnce({ patchAndFetchById }) + .mockReturnValueOnce({ where: reloadedWhere }); + + const result = await AgentMessageStore.syncCanonicalMessages('thread-uuid', 'sample-user', [ + { + id: 'client-message-1', + role: 'user', + parts: [{ type: 'text', text: 'After' }], + }, + ]); + + expect(patchAndFetchById).toHaveBeenCalledWith( + 11, + expect.objectContaining({ + parts: [{ type: 'text', text: 'After' }], + uiMessage: null, + metadata: { clientMessageId: 'client-message-1' }, + }) + ); + expect(result).toHaveLength(1); + }); + + it('strips non-canonical parts before persisting messages', async () => { + const insertedRow = { + id: 11, + uuid: '11111111-1111-4111-8111-111111111111', + threadId: 17, + role: 'assistant', + parts: [{ type: 'text', text: 'Done' }], + uiMessage: null, + metadata: {}, + }; + const insert = jest.fn().mockResolvedValue(insertedRow); + const orderBy = jest.fn().mockResolvedValue([insertedRow]); + const existingWhere = jest.fn().mockResolvedValue([]); + const reloadedWhere = jest.fn().mockReturnValue({ orderBy }); + + mockGetOwnedThread.mockResolvedValue({ id: 17, uuid: 'thread-uuid' }); + mockMessageQuery + .mockReturnValueOnce({ where: existingWhere }) + .mockReturnValueOnce({ insert }) + .mockReturnValueOnce({ where: reloadedWhere }); + + await AgentMessageStore.syncCanonicalMessages('thread-uuid', 'sample-user', [ + { + role: 'assistant', + parts: [ + { type: 'text', text: 'Done' }, + { + type: 'dynamic-tool', + toolCallId: 'tool-call-1', + toolName: 'workspace_edit_file', + state: 'output-available', + } as any, + ], + }, + ]); + + expect(insert).toHaveBeenCalledWith( + expect.objectContaining({ + parts: [{ type: 'text', text: 'Done' }], + uiMessage: null, + }) + ); + }); + }); + + describe('upsertCanonicalMessagesForThread', () => { + it('does not compare non-uuid client message IDs against the uuid column', async () => { + const existingLookupQuery = { + where: jest.fn(), + whereIn: jest.fn(), + orWhereIn: jest.fn(), + whereRaw: jest.fn(), + orWhereRaw: jest.fn(), + }; + existingLookupQuery.where.mockImplementation((arg: unknown) => { + if (typeof arg === 'function') { + arg(existingLookupQuery); + return Promise.resolve([]); + } + return existingLookupQuery; + }); + existingLookupQuery.whereIn.mockReturnValue(existingLookupQuery); + existingLookupQuery.orWhereIn.mockReturnValue(existingLookupQuery); + existingLookupQuery.whereRaw.mockReturnValue(existingLookupQuery); + existingLookupQuery.orWhereRaw.mockReturnValue(existingLookupQuery); + const insert = jest.fn().mockResolvedValue({ + id: 11, + uuid: '11111111-1111-4111-8111-111111111111', + metadata: { clientMessageId: 'short-client-message-id' }, + }); + + mockMessageQuery.mockReturnValueOnce(existingLookupQuery).mockReturnValueOnce({ insert }); + + await AgentMessageStore.upsertCanonicalMessagesForThread({ id: 17 }, [ + { + id: 'short-client-message-id', + role: 'user', + parts: [{ type: 'text', text: 'Hello' }], + }, + ]); + + expect(existingLookupQuery.whereIn).toHaveBeenCalledWith('clientMessageId', ['short-client-message-id']); + expect(existingLookupQuery.whereIn).not.toHaveBeenCalledWith('uuid', expect.anything()); + expect(existingLookupQuery.orWhereRaw).toHaveBeenCalledWith('"metadata"->>? = ANY(?::text[])', [ + 'clientMessageId', + ['short-client-message-id'], + ]); + expect(insert).toHaveBeenCalledWith( + expect.objectContaining({ + threadId: 17, + role: 'user', + clientMessageId: 'short-client-message-id', + metadata: { clientMessageId: 'short-client-message-id' }, + }) + ); + }); + }); + + describe('syncCanonicalMessagesFromUiMessages', () => { + it('stores only canonical conversational parts from harness UI messages', async () => { + const insertedRow = { + id: 11, + uuid: '33333333-3333-4333-8333-333333333333', + threadId: 17, + role: 'assistant', + parts: [{ type: 'text', text: 'Done' }], + uiMessage: null, + metadata: { runId: 'run-1' }, + }; + const insert = jest.fn().mockResolvedValue(insertedRow); + const orderBy = jest.fn().mockResolvedValue([insertedRow]); + const existingWhere = jest.fn().mockResolvedValue([]); + const reloadedWhere = jest.fn().mockReturnValue({ orderBy }); + + mockGetOwnedThread.mockResolvedValue({ id: 17, uuid: 'thread-uuid' }); + mockMessageQuery + .mockReturnValueOnce({ where: existingWhere }) + .mockReturnValueOnce({ insert }) + .mockReturnValueOnce({ where: reloadedWhere }); + + await AgentMessageStore.syncCanonicalMessagesFromUiMessages('thread-uuid', 'sample-user', [ + { + id: '33333333-3333-4333-8333-333333333333', + role: 'assistant', + metadata: { runId: 'run-1' }, + parts: [ + { type: 'text', text: 'Done' }, + { + type: 'dynamic-tool', + toolCallId: 'tool-1', + toolName: 'workspace_edit_file', + state: 'output-available', + output: { ok: true }, + }, + ], + } as any, + ]); + + expect(insert).toHaveBeenCalledWith( + expect.objectContaining({ + role: 'assistant', + parts: [{ type: 'text', text: 'Done' }], + uiMessage: null, + metadata: { runId: 'run-1' }, + }) + ); + }); + + it('does not persist assistant messages that only contain tool UI parts', async () => { + const existingWhere = jest.fn().mockResolvedValue([]); const orderBy = jest.fn().mockResolvedValue([]); - const select = jest.fn().mockReturnValue({ orderBy }); - const secondWhere = jest.fn().mockReturnValue({ select }); - const firstWhere = jest.fn().mockReturnValue({ where: secondWhere }); - const joinRelated = jest.fn().mockReturnValue({ where: firstWhere }); - const alias = jest.fn().mockReturnValue({ joinRelated }); - mockMessageQuery.mockReturnValue({ alias }); - - await AgentMessageStore.listRunMessages('2b3a4084-caf5-4eca-a4bb-6fdd7e93ae04', 'sample-user'); - - expect(alias).toHaveBeenCalledWith('message'); - expect(joinRelated).toHaveBeenCalledWith('run.thread.session'); - expect(firstWhere).toHaveBeenCalledWith('run.uuid', '2b3a4084-caf5-4eca-a4bb-6fdd7e93ae04'); - expect(secondWhere).toHaveBeenCalledWith('run:thread:session.userId', 'sample-user'); + const reloadedWhere = jest.fn().mockReturnValue({ orderBy }); + + mockGetOwnedThread.mockResolvedValue({ id: 17, uuid: 'thread-uuid' }); + mockMessageQuery.mockReturnValueOnce({ where: existingWhere }).mockReturnValueOnce({ where: reloadedWhere }); + + const result = await AgentMessageStore.syncCanonicalMessagesFromUiMessages('thread-uuid', 'sample-user', [ + { + id: '33333333-3333-4333-8333-333333333333', + role: 'assistant', + metadata: { runId: 'run-1' }, + parts: [ + { + type: 'dynamic-tool', + toolCallId: 'tool-1', + toolName: 'workspace_write_file', + state: 'output-available', + output: { ok: true }, + }, + ], + } as any, + ]); + + expect(mockMessageQuery).toHaveBeenCalledTimes(2); + expect(result).toEqual([]); }); }); }); diff --git a/src/server/services/agent/__tests__/PolicyService.test.ts b/src/server/services/agent/__tests__/PolicyService.test.ts index 80d1def6..c3deded9 100644 --- a/src/server/services/agent/__tests__/PolicyService.test.ts +++ b/src/server/services/agent/__tests__/PolicyService.test.ts @@ -15,6 +15,7 @@ */ import AgentPolicyService from '../PolicyService'; +import { DEFAULT_AGENT_APPROVAL_POLICY } from '../types'; describe('AgentPolicyService', () => { it('keeps read-only sandbox tools in the read capability', () => { @@ -40,4 +41,10 @@ describe('AgentPolicyService', () => { it('maps mutating external MCP tools to external_mcp_write without workspace heuristics', () => { expect(AgentPolicyService.capabilityForExternalMcpTool('editJiraIssue')).toBe('external_mcp_write'); }); + + it('requires approval for deployment mutations by default', () => { + expect(AgentPolicyService.modeForCapability(DEFAULT_AGENT_APPROVAL_POLICY, 'deploy_k8s_mutation')).toBe( + 'require_approval' + ); + }); }); diff --git a/src/server/services/agent/__tests__/ProviderRegistry.test.ts b/src/server/services/agent/__tests__/ProviderRegistry.test.ts index 184470f2..ce5ad68f 100644 --- a/src/server/services/agent/__tests__/ProviderRegistry.test.ts +++ b/src/server/services/agent/__tests__/ProviderRegistry.test.ts @@ -89,7 +89,6 @@ describe('resolveRequestedModelSelection', () => { describe('AgentProviderRegistry credential resolution', () => { const originalAnthropicKey = process.env.ANTHROPIC_API_KEY; - const originalEnableAuth = process.env.ENABLE_AUTH; beforeEach(() => { jest.clearAllMocks(); @@ -119,12 +118,6 @@ describe('AgentProviderRegistry credential resolution', () => { } else { process.env.ANTHROPIC_API_KEY = originalAnthropicKey; } - - if (originalEnableAuth === undefined) { - delete process.env.ENABLE_AUTH; - } else { - process.env.ENABLE_AUTH = originalEnableAuth; - } }); it('uses stored user keys and ignores process env fallback', async () => { @@ -224,66 +217,7 @@ describe('AgentProviderRegistry credential resolution', () => { ).rejects.toBeInstanceOf(MissingAgentProviderApiKeyError); }); - it('uses the request api key for local auth-disabled sessions when the provider matches', async () => { - process.env.ENABLE_AUTH = 'false'; - - await expect( - AgentProviderRegistry.getRequiredStoredApiKey({ - provider: 'anthropic', - userIdentity: { - userId: 'sample-user', - githubUsername: 'sample-user', - }, - requestApiKey: 'sample-browser-anthropic-key', - requestApiKeyProvider: 'anthropic', - }) - ).resolves.toBe('sample-browser-anthropic-key'); - - await expect( - AgentProviderRegistry.resolveCredentialEnvMap({ - repoFullName: 'example-org/example-repo', - userIdentity: { - userId: 'sample-user', - githubUsername: 'sample-user', - }, - requestApiKey: 'sample-browser-anthropic-key', - requestApiKeyProvider: 'anthropic', - }) - ).resolves.toEqual({ - ANTHROPIC_API_KEY: 'sample-browser-anthropic-key', - }); - - await expect( - AgentProviderRegistry.getRequiredStoredApiKey({ - provider: 'google', - userIdentity: { - userId: 'sample-user', - githubUsername: 'sample-user', - }, - requestApiKey: 'sample-browser-gemini-key', - requestApiKeyProvider: 'gemini', - }) - ).resolves.toBe('sample-browser-gemini-key'); - }); - - it('does not treat a mismatched local request key as valid for another provider', async () => { - process.env.ENABLE_AUTH = 'false'; - - await expect( - AgentProviderRegistry.getRequiredStoredApiKey({ - provider: 'gemini', - userIdentity: { - userId: 'sample-user', - githubUsername: 'sample-user', - }, - requestApiKey: 'sample-browser-anthropic-key', - requestApiKeyProvider: 'anthropic', - }) - ).rejects.toBeInstanceOf(MissingAgentProviderApiKeyError); - }); - - it('exposes provider-scoped local request keys in the available model list', async () => { - process.env.ENABLE_AUTH = 'false'; + it('lists only models backed by stored user keys', async () => { mockGetEffectiveConfig.mockResolvedValueOnce({ providers: [ { @@ -316,6 +250,9 @@ describe('AgentProviderRegistry credential resolution', () => { }, ], }); + (UserApiKeyService.getDecryptedKey as jest.Mock).mockImplementation(async (_userId: string, provider: string) => + provider === 'gemini' ? 'user-gemini-key' : null + ); await expect( AgentProviderRegistry.listAvailableModelsForUser({ @@ -324,8 +261,6 @@ describe('AgentProviderRegistry credential resolution', () => { userId: 'sample-user', githubUsername: 'sample-user', }, - requestApiKey: 'sample-browser-gemini-key', - requestApiKeyProvider: 'gemini', }) ).resolves.toEqual([ { diff --git a/src/server/services/agent/__tests__/RunAdmissionService.test.ts b/src/server/services/agent/__tests__/RunAdmissionService.test.ts new file mode 100644 index 00000000..7abc4970 --- /dev/null +++ b/src/server/services/agent/__tests__/RunAdmissionService.test.ts @@ -0,0 +1,221 @@ +/** + * Copyright 2026 GoodRx, Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +jest.mock('server/models/AgentRun', () => ({ + __esModule: true, + default: { + query: jest.fn(), + transaction: jest.fn(), + }, +})); + +jest.mock('server/models/AgentSession', () => ({ + __esModule: true, + default: { + query: jest.fn(), + }, +})); + +jest.mock('server/models/AgentThread', () => ({ + __esModule: true, + default: { + query: jest.fn(), + }, +})); + +jest.mock('server/services/agent/MessageStore', () => ({ + __esModule: true, + default: { + findCanonicalMessageByClientMessageId: jest.fn(), + insertUserMessageForRun: jest.fn(), + }, +})); + +jest.mock('server/services/agent/RunEventService', () => ({ + __esModule: true, + default: { + appendStatusEvent: jest.fn(), + }, +})); + +jest.mock('server/lib/dependencies', () => ({})); + +import AgentRun from 'server/models/AgentRun'; +import AgentSession from 'server/models/AgentSession'; +import AgentThread from 'server/models/AgentThread'; +import AgentMessageStore from '../MessageStore'; +import AgentRunAdmissionService from '../RunAdmissionService'; +import AgentRunEventService from '../RunEventService'; + +const mockRunQuery = AgentRun.query as jest.Mock; +const mockRunTransaction = AgentRun.transaction as jest.Mock; +const mockSessionQuery = AgentSession.query as jest.Mock; +const mockThreadQuery = AgentThread.query as jest.Mock; +const mockFindCanonicalMessageByClientMessageId = AgentMessageStore.findCanonicalMessageByClientMessageId as jest.Mock; +const mockInsertUserMessageForRun = AgentMessageStore.insertUserMessageForRun as jest.Mock; +const mockAppendStatusEvent = AgentRunEventService.appendStatusEvent as jest.Mock; + +function buildActiveRunQuery(activeRun: unknown = null) { + const query = { + where: jest.fn(), + whereNotIn: jest.fn(), + orderBy: jest.fn(), + first: jest.fn().mockResolvedValue(activeRun), + }; + query.where.mockReturnValue(query); + query.whereNotIn.mockReturnValue(query); + query.orderBy.mockReturnValue(query); + return query; +} + +describe('AgentRunAdmissionService', () => { + beforeEach(() => { + jest.clearAllMocks(); + mockRunTransaction.mockImplementation(async (callback) => callback({ trx: true })); + mockSessionQuery.mockReturnValue({ + findById: jest.fn().mockReturnValue({ + forUpdate: jest.fn().mockResolvedValue({}), + }), + }); + mockThreadQuery.mockReturnValue({ + patchAndFetchById: jest.fn().mockResolvedValue({}), + }); + mockFindCanonicalMessageByClientMessageId.mockResolvedValue(undefined); + mockInsertUserMessageForRun.mockResolvedValue({ id: 31, uuid: 'message-1' }); + mockAppendStatusEvent.mockResolvedValue(undefined); + }); + + it('persists submitted message and queued run in the same transaction', async () => { + const queuedRun = { + id: 23, + uuid: 'run-1', + status: 'queued', + }; + const activeRunQuery = buildActiveRunQuery(); + const insertRunQuery = { + insertAndFetch: jest.fn().mockResolvedValue(queuedRun), + }; + mockRunQuery.mockReturnValueOnce(activeRunQuery).mockReturnValueOnce(insertRunQuery); + + const admission = await AgentRunAdmissionService.createQueuedRunWithMessage({ + thread: { id: 7, uuid: 'thread-1', metadata: {} } as Parameters< + typeof AgentRunAdmissionService.createQueuedRunWithMessage + >[0]['thread'], + session: { id: 17, uuid: 'session-1' } as Parameters< + typeof AgentRunAdmissionService.createQueuedRunWithMessage + >[0]['session'], + policy: { defaultMode: 'require_approval', rules: {} } as any, + message: { clientMessageId: 'client-message-1', parts: [{ type: 'text', text: 'Hi' }] }, + resolvedHarness: 'lifecycle_ai_sdk', + resolvedProvider: 'openai', + resolvedModel: 'gpt-5.4', + runtimeOptions: { maxIterations: 12 }, + }); + + expect(admission).toEqual({ + run: queuedRun, + message: { id: 31, uuid: 'message-1' }, + created: true, + }); + expect(mockInsertUserMessageForRun).toHaveBeenCalledWith( + expect.objectContaining({ id: 7, uuid: 'thread-1' }), + queuedRun, + { clientMessageId: 'client-message-1', parts: [{ type: 'text', text: 'Hi' }] }, + { trx: true } + ); + expect(insertRunQuery.insertAndFetch).toHaveBeenCalledWith( + expect.objectContaining({ + threadId: 7, + sessionId: 17, + status: 'queued', + resolvedHarness: 'lifecycle_ai_sdk', + resolvedProvider: 'openai', + resolvedModel: 'gpt-5.4', + policySnapshot: expect.objectContaining({ + runtimeOptions: { maxIterations: 12 }, + }), + }) + ); + expect(mockAppendStatusEvent).toHaveBeenCalledWith('run-1', 'run.queued', { + threadId: 'thread-1', + sessionId: 'session-1', + }); + }); + + it('does not persist messages when another run is active', async () => { + const activeRunQuery = buildActiveRunQuery({ id: 99, uuid: 'run-active' }); + mockRunQuery.mockReturnValueOnce(activeRunQuery); + + await expect( + AgentRunAdmissionService.createQueuedRunWithMessage({ + thread: { id: 7, uuid: 'thread-1', metadata: {} } as Parameters< + typeof AgentRunAdmissionService.createQueuedRunWithMessage + >[0]['thread'], + session: { id: 17, uuid: 'session-1' } as Parameters< + typeof AgentRunAdmissionService.createQueuedRunWithMessage + >[0]['session'], + policy: { defaultMode: 'require_approval', rules: {} } as any, + message: { clientMessageId: 'client-message-1', parts: [{ type: 'text', text: 'Hi' }] }, + resolvedHarness: 'lifecycle_ai_sdk', + resolvedProvider: 'openai', + resolvedModel: 'gpt-5.4', + }) + ).rejects.toThrow('Wait for the current agent run to finish before starting another run.'); + + expect(mockInsertUserMessageForRun).not.toHaveBeenCalled(); + expect(mockAppendStatusEvent).not.toHaveBeenCalled(); + }); + + it('returns the existing run for duplicate client message ids', async () => { + const existingMessage = { + id: 31, + uuid: 'message-1', + runId: 23, + }; + const existingRun = { + id: 23, + uuid: 'run-1', + status: 'queued', + }; + const findRunQuery = { + findById: jest.fn().mockResolvedValue(existingRun), + }; + mockFindCanonicalMessageByClientMessageId.mockResolvedValueOnce(existingMessage); + mockRunQuery.mockReturnValueOnce(findRunQuery); + + const admission = await AgentRunAdmissionService.createQueuedRunWithMessage({ + thread: { id: 7, uuid: 'thread-1', metadata: {} } as Parameters< + typeof AgentRunAdmissionService.createQueuedRunWithMessage + >[0]['thread'], + session: { id: 17, uuid: 'session-1' } as Parameters< + typeof AgentRunAdmissionService.createQueuedRunWithMessage + >[0]['session'], + policy: { defaultMode: 'require_approval', rules: {} } as any, + message: { clientMessageId: 'client-message-1', parts: [{ type: 'text', text: 'Hi' }] }, + resolvedHarness: 'lifecycle_ai_sdk', + resolvedProvider: 'openai', + resolvedModel: 'gpt-5.4', + }); + + expect(admission).toEqual({ + run: existingRun, + message: existingMessage, + created: false, + }); + expect(mockInsertUserMessageForRun).not.toHaveBeenCalled(); + expect(mockAppendStatusEvent).not.toHaveBeenCalled(); + }); +}); diff --git a/src/server/services/agent/__tests__/RunEventService.test.ts b/src/server/services/agent/__tests__/RunEventService.test.ts new file mode 100644 index 00000000..5b7361ef --- /dev/null +++ b/src/server/services/agent/__tests__/RunEventService.test.ts @@ -0,0 +1,874 @@ +/** + * Copyright 2026 GoodRx, Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +jest.mock('server/models/AgentRun', () => ({ + __esModule: true, + default: { + query: jest.fn(), + transaction: jest.fn(), + }, +})); + +jest.mock('server/models/AgentRunEvent', () => ({ + __esModule: true, + default: { + knex: jest.fn(), + query: jest.fn(), + }, +})); + +jest.mock('server/lib/agentSession/runtimeConfig', () => { + const actual = jest.requireActual('server/lib/agentSession/runtimeConfig'); + return { + __esModule: true, + ...actual, + resolveAgentSessionDurabilityConfig: jest.fn().mockResolvedValue({ + runExecutionLeaseMs: 30 * 60 * 1000, + queuedRunDispatchStaleMs: 30 * 1000, + dispatchRecoveryLimit: 50, + maxDurablePayloadBytes: 64 * 1024, + payloadPreviewBytes: 16 * 1024, + fileChangePreviewChars: 4000, + }), + }; +}); + +import AgentRun from 'server/models/AgentRun'; +import AgentRunEvent from 'server/models/AgentRunEvent'; +import AgentRunEventService from '../RunEventService'; +import { AgentRunOwnershipLostError } from '../AgentRunOwnershipLostError'; + +const mockRunQuery = AgentRun.query as jest.Mock; +const mockRunTransaction = AgentRun.transaction as jest.Mock; +const mockRunEventKnex = AgentRunEvent.knex as jest.Mock; +const mockRunEventQuery = AgentRunEvent.query as jest.Mock; + +describe('AgentRunEventService', () => { + beforeEach(() => { + jest.clearAllMocks(); + mockRunTransaction.mockImplementation(async (callback) => callback({ trx: true })); + mockRunEventKnex.mockReturnValue({ + raw: jest.fn().mockResolvedValue(undefined), + }); + }); + + afterEach(() => { + jest.restoreAllMocks(); + }); + + it('loads run events after a sequence cursor with one extra row for hasMore', async () => { + const limit = jest.fn().mockResolvedValue([ + { + uuid: 'event-5', + sequence: 5, + }, + { + uuid: 'event-6', + sequence: 6, + }, + { + uuid: 'event-7', + sequence: 7, + }, + ]); + const orderById = jest.fn().mockReturnValue({ limit }); + const orderBySequence = jest.fn().mockReturnValue({ orderBy: orderById }); + const whereSequence = jest.fn().mockReturnValue({ orderBy: orderBySequence }); + const whereRun = jest.fn().mockReturnValue({ where: whereSequence }); + + mockRunEventQuery.mockReturnValue({ where: whereRun }); + + const result = await AgentRunEventService.listRunEventsPageForRun( + { + id: 17, + uuid: 'run-1', + threadId: 11, + sessionId: 13, + threadUuid: 'thread-1', + sessionUuid: 'session-1', + status: 'running', + } as any, + { + afterSequence: 4, + limit: 2, + } + ); + + expect(whereRun).toHaveBeenCalledWith({ runId: 17 }); + expect(whereSequence).toHaveBeenCalledWith('sequence', '>', 4); + expect(orderBySequence).toHaveBeenCalledWith('sequence', 'asc'); + expect(orderById).toHaveBeenCalledWith('id', 'asc'); + expect(limit).toHaveBeenCalledWith(3); + expect(result).toEqual({ + events: [ + expect.objectContaining({ + uuid: 'event-5', + runUuid: 'run-1', + threadUuid: 'thread-1', + sessionUuid: 'session-1', + sequence: 5, + }), + expect.objectContaining({ + uuid: 'event-6', + runUuid: 'run-1', + threadUuid: 'thread-1', + sessionUuid: 'session-1', + sequence: 6, + }), + ], + nextSequence: 6, + hasMore: true, + run: { + id: 'run-1', + status: 'running', + }, + limit: 2, + maxLimit: 500, + }); + }); + + it('persists canonical run-event payloads without private UI replay chunks', async () => { + const insert = jest.fn().mockResolvedValue(undefined); + const latestFirst = jest.fn().mockResolvedValue(null); + const runFindOne = jest.fn().mockResolvedValue({ id: 17, uuid: 'run-1' }); + const runForUpdate = jest.fn().mockResolvedValue({ id: 17, uuid: 'run-1' }); + const runFindById = jest.fn().mockReturnValue({ forUpdate: runForUpdate }); + + mockRunQuery.mockReturnValueOnce({ findOne: runFindOne }).mockReturnValueOnce({ findById: runFindById }); + mockRunEventQuery + .mockReturnValueOnce({ + where: jest.fn().mockReturnValue({ + orderBy: jest.fn().mockReturnValue({ + first: latestFirst, + }), + }), + }) + .mockReturnValueOnce({ insert }); + + await AgentRunEventService.appendEventsForChunks('run-1', [ + { type: 'start', messageId: 'assistant-1' } as any, + { type: 'text-start', id: 'text-1' } as any, + { type: 'text-delta', id: 'text-1', delta: 'Hello' } as any, + { type: 'finish', finishReason: 'stop' } as any, + ]); + + expect(insert).toHaveBeenCalledWith([ + expect.objectContaining({ + runId: 17, + sequence: 1, + eventType: 'message.created', + payload: { + messageId: 'assistant-1', + metadata: {}, + }, + }), + expect.objectContaining({ + sequence: 2, + eventType: 'message.part.started', + payload: { + partType: 'text', + partId: 'text-1', + }, + }), + expect.objectContaining({ + sequence: 3, + eventType: 'message.delta', + payload: { + partType: 'text', + partId: 'text-1', + delta: 'Hello', + }, + }), + expect.objectContaining({ + sequence: 4, + eventType: 'run.finished', + payload: { + finishReason: 'stop', + metadata: {}, + }, + }), + ]); + expect(runFindById).toHaveBeenCalledWith(17); + expect(runForUpdate).toHaveBeenCalledTimes(1); + expect(JSON.stringify(insert.mock.calls[0][0])).not.toContain('__uiReplayChunk'); + expect(mockRunEventKnex().raw).toHaveBeenCalledWith('select pg_notify(?, ?)', [ + 'agent_run_events', + JSON.stringify({ + runId: 'run-1', + latestSequence: 4, + }), + ]); + }); + + it('truncates oversized durable event payloads', async () => { + const insert = jest.fn().mockResolvedValue(undefined); + const latestFirst = jest.fn().mockResolvedValue(null); + const runFindOne = jest.fn().mockResolvedValue({ id: 17, uuid: 'run-1' }); + const runForUpdate = jest.fn().mockResolvedValue({ id: 17, uuid: 'run-1' }); + const runFindById = jest.fn().mockReturnValue({ forUpdate: runForUpdate }); + + mockRunQuery.mockReturnValueOnce({ findOne: runFindOne }).mockReturnValueOnce({ findById: runFindById }); + mockRunEventQuery + .mockReturnValueOnce({ + where: jest.fn().mockReturnValue({ + orderBy: jest.fn().mockReturnValue({ + first: latestFirst, + }), + }), + }) + .mockReturnValueOnce({ insert }); + + await AgentRunEventService.appendEventsForChunks('run-1', [ + { + type: 'tool-output-available', + toolCallId: 'tool-call-1', + output: { + content: 'x'.repeat(70 * 1024), + }, + } as any, + ]); + + expect(insert).toHaveBeenCalledWith([ + expect.objectContaining({ + eventType: 'tool.call.completed', + payload: expect.objectContaining({ + toolCallId: 'tool-call-1', + status: 'completed', + output: expect.objectContaining({ + truncated: true, + originalJsonBytes: expect.any(Number), + preview: expect.any(String), + }), + }), + }), + ]); + }); + + it('persists approval request events with the pending action link when present', async () => { + const insert = jest.fn().mockResolvedValue(undefined); + const latestFirst = jest.fn().mockResolvedValue(null); + const runFindOne = jest.fn().mockResolvedValue({ id: 17, uuid: 'run-1' }); + const runForUpdate = jest.fn().mockResolvedValue({ id: 17, uuid: 'run-1' }); + const runFindById = jest.fn().mockReturnValue({ forUpdate: runForUpdate }); + + mockRunQuery.mockReturnValueOnce({ findOne: runFindOne }).mockReturnValueOnce({ findById: runFindById }); + mockRunEventQuery + .mockReturnValueOnce({ + where: jest.fn().mockReturnValue({ + orderBy: jest.fn().mockReturnValue({ + first: latestFirst, + }), + }), + }) + .mockReturnValueOnce({ insert }); + + await AgentRunEventService.appendEventsForChunks('run-1', [ + { + type: 'tool-approval-request', + actionId: 'action-1', + approvalId: 'approval-1', + toolCallId: 'tool-call-1', + } as any, + ]); + + expect(insert).toHaveBeenCalledWith([ + expect.objectContaining({ + eventType: 'approval.requested', + payload: { + actionId: 'action-1', + approvalId: 'approval-1', + toolCallId: 'tool-call-1', + }, + }), + ]); + }); + + it('persists stream events when the locked run owner matches', async () => { + const insert = jest.fn().mockResolvedValue(undefined); + const latestFirst = jest.fn().mockResolvedValue(null); + const runFindOne = jest.fn().mockResolvedValue({ id: 17, uuid: 'run-1' }); + const runForUpdate = jest.fn().mockResolvedValue({ + id: 17, + uuid: 'run-1', + status: 'running', + executionOwner: 'worker-1', + }); + const runFindById = jest.fn().mockReturnValue({ forUpdate: runForUpdate }); + + mockRunQuery.mockReturnValueOnce({ findOne: runFindOne }).mockReturnValueOnce({ findById: runFindById }); + mockRunEventQuery + .mockReturnValueOnce({ + where: jest.fn().mockReturnValue({ + orderBy: jest.fn().mockReturnValue({ + first: latestFirst, + }), + }), + }) + .mockReturnValueOnce({ insert }); + + await AgentRunEventService.appendEventsForChunksForExecutionOwner('run-1', 'worker-1', [ + { type: 'text-delta', id: 'text-1', delta: 'Hello' } as any, + ]); + + expect(runFindById).toHaveBeenCalledWith(17); + expect(runForUpdate).toHaveBeenCalledTimes(1); + expect(insert).toHaveBeenCalledWith([ + expect.objectContaining({ + runId: 17, + sequence: 1, + eventType: 'message.delta', + payload: { + partType: 'text', + partId: 'text-1', + delta: 'Hello', + }, + }), + ]); + expect(mockRunEventKnex().raw).toHaveBeenCalledWith('select pg_notify(?, ?)', [ + 'agent_run_events', + JSON.stringify({ + runId: 'run-1', + latestSequence: 1, + }), + ]); + }); + + it('throws ownership loss without writing stream events when the locked run owner is stale', async () => { + const runFindOne = jest.fn().mockResolvedValue({ id: 17, uuid: 'run-1' }); + const runForUpdate = jest.fn().mockResolvedValue({ + id: 17, + uuid: 'run-1', + status: 'running', + executionOwner: 'worker-2', + }); + const runFindById = jest.fn().mockReturnValue({ forUpdate: runForUpdate }); + + mockRunQuery.mockReturnValueOnce({ findOne: runFindOne }).mockReturnValueOnce({ findById: runFindById }); + + let thrownError: unknown; + try { + await AgentRunEventService.appendEventsForChunksForExecutionOwner('run-1', 'worker-1', [ + { type: 'text-delta', id: 'text-1', delta: 'Hello' } as any, + ]); + } catch (error) { + thrownError = error; + } + + expect(thrownError).toBeInstanceOf(AgentRunOwnershipLostError); + expect(thrownError).toMatchObject({ + runUuid: 'run-1', + expectedExecutionOwner: 'worker-1', + currentStatus: 'running', + currentExecutionOwner: 'worker-2', + }); + + expect(mockRunEventQuery).not.toHaveBeenCalled(); + expect(mockRunEventKnex().raw).not.toHaveBeenCalled(); + }); + + it('projects UI replay chunks from canonical run events', () => { + const chunks = AgentRunEventService.projectUiChunksFromEvents([ + { + eventType: 'message.created', + payload: { + messageId: 'assistant-1', + metadata: { provider: 'openai' }, + }, + }, + { + eventType: 'message.part.started', + payload: { + partType: 'text', + partId: 'text-1', + }, + }, + { + eventType: 'message.delta', + payload: { + partType: 'text', + partId: 'text-1', + delta: 'Hello', + }, + }, + { + eventType: 'message.part.completed', + payload: { + partType: 'text', + partId: 'text-1', + }, + }, + { + eventType: 'tool.call.started', + payload: { + toolCallId: 'tool-call-1', + toolName: 'workspace_read_file', + inputStatus: 'available', + input: { path: '/workspace/README.md' }, + }, + }, + { + eventType: 'tool.call.completed', + payload: { + toolCallId: 'tool-call-1', + status: 'completed', + output: { ok: true }, + }, + }, + { + eventType: 'approval.requested', + payload: { + actionId: 'action-1', + approvalId: 'approval-1', + toolCallId: 'tool-call-2', + }, + }, + { + eventType: 'tool.file_change', + payload: { + id: 'file-change-1', + data: { + id: 'change-1', + toolCallId: 'tool-call-2', + sourceTool: 'workspace_edit_file', + path: 'README.md', + displayPath: 'README.md', + kind: 'edited', + stage: 'awaiting-approval', + additions: 1, + deletions: 0, + truncated: false, + }, + }, + }, + { + eventType: 'run.finished', + payload: { + finishReason: 'stop', + metadata: {}, + }, + }, + ] as any); + + expect(chunks).toEqual([ + { type: 'start', messageId: 'assistant-1', messageMetadata: { provider: 'openai' } }, + { type: 'text-start', id: 'text-1' }, + { type: 'text-delta', id: 'text-1', delta: 'Hello' }, + { type: 'text-end', id: 'text-1' }, + { + type: 'tool-input-available', + toolCallId: 'tool-call-1', + toolName: 'workspace_read_file', + input: { path: '/workspace/README.md' }, + }, + { + type: 'tool-output-available', + toolCallId: 'tool-call-1', + output: { ok: true }, + }, + { + type: 'tool-approval-request', + actionId: 'action-1', + approvalId: 'approval-1', + toolCallId: 'tool-call-2', + }, + { + type: 'data-file-change', + id: 'file-change-1', + data: { + id: 'change-1', + toolCallId: 'tool-call-2', + sourceTool: 'workspace_edit_file', + path: 'README.md', + displayPath: 'README.md', + kind: 'edited', + stage: 'awaiting-approval', + additions: 1, + deletions: 0, + truncated: false, + }, + }, + { type: 'finish', finishReason: 'stop', messageMetadata: {} }, + ]); + }); + + it('projects run.failed events to UI error chunks for stream compatibility', () => { + const chunks = AgentRunEventService.projectUiChunksFromEvents([ + { + eventType: 'run.failed', + payload: { + status: 'failed', + error: { + message: 'Sample run failure.', + }, + }, + }, + ] as any); + + expect(chunks).toEqual([ + { + type: 'error', + errorText: 'Sample run failure.', + }, + ]); + }); + + it('serializes canonical run event payloads as public data', () => { + const serialized = AgentRunEventService.serializeRunEvent({ + uuid: 'event-1', + runUuid: 'run-1', + threadUuid: 'thread-1', + sessionUuid: 'session-1', + runId: 17, + sequence: 4, + eventType: 'message.delta', + payload: { + partType: 'text', + partId: 'text-1', + delta: 'Hello', + }, + createdAt: '2026-04-24T00:00:00.000Z', + updatedAt: '2026-04-24T00:00:00.000Z', + } as any); + + expect(serialized).toEqual({ + id: 'event-1', + runId: 'run-1', + threadId: 'thread-1', + sessionId: 'session-1', + sequence: 4, + eventType: 'message.delta', + version: 1, + payload: { + partType: 'text', + partId: 'text-1', + delta: 'Hello', + }, + createdAt: '2026-04-24T00:00:00.000Z', + updatedAt: '2026-04-24T00:00:00.000Z', + }); + }); + + it('streams canonical SSE frames after the requested sequence cursor', async () => { + const terminalEvent = { + uuid: 'event-7', + runUuid: 'run-1', + threadUuid: 'thread-1', + sessionUuid: 'session-1', + runId: 17, + sequence: 7, + eventType: 'run.completed', + payload: { + status: 'completed', + }, + createdAt: '2026-04-24T00:00:01.000Z', + updatedAt: '2026-04-24T00:00:01.000Z', + } as any; + const listRunEventsPage = jest + .spyOn(AgentRunEventService, 'listRunEventsPage') + .mockResolvedValueOnce({ + events: [ + { + uuid: 'event-6', + runUuid: 'run-1', + threadUuid: 'thread-1', + sessionUuid: 'session-1', + runId: 17, + sequence: 6, + eventType: 'message.delta', + payload: { + partType: 'text', + partId: 'text-1', + delta: 'Hello', + }, + createdAt: '2026-04-24T00:00:00.000Z', + updatedAt: '2026-04-24T00:00:00.000Z', + } as any, + ], + nextSequence: 6, + hasMore: false, + run: { + id: 'run-1', + status: 'running', + }, + limit: 100, + maxLimit: 500, + }) + .mockResolvedValueOnce({ + events: [terminalEvent], + nextSequence: 7, + hasMore: false, + run: { + id: 'run-1', + status: 'completed', + }, + limit: 100, + maxLimit: 500, + }); + const waitForRunEventNotification = jest.spyOn(AgentRunEventService, 'waitForRunEventNotification'); + mockRunQuery.mockReturnValue({ + findOne: jest.fn().mockResolvedValue({ + uuid: 'run-1', + status: 'completed', + }), + }); + + const text = await new Response(AgentRunEventService.createCanonicalRunEventStream('run-1', 5)).text(); + + expect(listRunEventsPage).toHaveBeenNthCalledWith(1, 'run-1', { + afterSequence: 5, + limit: 100, + }); + expect(listRunEventsPage).toHaveBeenNthCalledWith(2, 'run-1', { + afterSequence: 6, + limit: 100, + }); + expect(waitForRunEventNotification).not.toHaveBeenCalled(); + expect(text).toContain( + 'data: {"id":"event-6","runId":"run-1","threadId":"thread-1","sessionId":"session-1","sequence":6,"eventType":"message.delta","version":1,"payload":{"partType":"text","partId":"text-1","delta":"Hello"},"createdAt":"2026-04-24T00:00:00.000Z","updatedAt":"2026-04-24T00:00:00.000Z"}' + ); + expect(text).toContain('id: 7\nevent: run.completed'); + }); + + it('backs off when the run-event notification listener is unavailable', async () => { + jest.useFakeTimers(); + try { + const acquireConnection = jest.fn().mockRejectedValue(new Error('listen unavailable')); + mockRunEventKnex.mockReturnValue({ + client: { + acquireConnection, + releaseConnection: jest.fn(), + }, + }); + + const promise = AgentRunEventService.waitForRunEventNotification('run-1', 7, 25); + const settled = jest.fn(); + void promise.then(settled); + + await Promise.resolve(); + await Promise.resolve(); + expect(settled).not.toHaveBeenCalled(); + + await jest.advanceTimersByTimeAsync(24); + expect(settled).not.toHaveBeenCalled(); + + await jest.advanceTimersByTimeAsync(1); + await expect(promise).resolves.toBe(false); + expect(acquireConnection).toHaveBeenCalledTimes(1); + } finally { + jest.useRealTimers(); + } + }); + + it('drains events written after terminal status before closing the canonical stream', async () => { + const terminalEvent = { + uuid: 'event-2', + runUuid: 'run-1', + threadUuid: 'thread-1', + sessionUuid: 'session-1', + runId: 17, + sequence: 2, + eventType: 'run.completed', + payload: { + status: 'completed', + }, + createdAt: null, + updatedAt: null, + } as any; + const listRunEventsPage = jest + .spyOn(AgentRunEventService, 'listRunEventsPage') + .mockResolvedValueOnce({ + events: [], + nextSequence: 1, + hasMore: false, + run: { + id: 'run-1', + status: 'running', + }, + limit: 100, + maxLimit: 500, + }) + .mockResolvedValueOnce({ + events: [terminalEvent], + nextSequence: 2, + hasMore: false, + run: { + id: 'run-1', + status: 'completed', + }, + limit: 100, + maxLimit: 500, + }); + mockRunQuery.mockReturnValue({ + findOne: jest.fn().mockResolvedValue({ + uuid: 'run-1', + status: 'completed', + }), + }); + + const text = await new Response(AgentRunEventService.createCanonicalRunEventStream('run-1', 1)).text(); + + expect(listRunEventsPage).toHaveBeenNthCalledWith(1, 'run-1', { + afterSequence: 1, + limit: 100, + }); + expect(listRunEventsPage).toHaveBeenNthCalledWith(2, 'run-1', { + afterSequence: 1, + limit: 100, + }); + expect(text).toContain('id: 2\nevent: run.completed'); + }); + + it('waits once for the terminal event when terminal status is visible first', async () => { + const terminalEvent = { + uuid: 'event-2', + runUuid: 'run-1', + threadUuid: 'thread-1', + sessionUuid: 'session-1', + runId: 17, + sequence: 2, + eventType: 'run.completed', + payload: { + status: 'completed', + }, + createdAt: null, + updatedAt: null, + } as any; + const listRunEventsPage = jest + .spyOn(AgentRunEventService, 'listRunEventsPage') + .mockResolvedValueOnce({ + events: [], + nextSequence: 1, + hasMore: false, + run: { + id: 'run-1', + status: 'running', + }, + limit: 100, + maxLimit: 500, + }) + .mockResolvedValueOnce({ + events: [], + nextSequence: 1, + hasMore: false, + run: { + id: 'run-1', + status: 'completed', + }, + limit: 100, + maxLimit: 500, + }) + .mockResolvedValueOnce({ + events: [terminalEvent], + nextSequence: 2, + hasMore: false, + run: { + id: 'run-1', + status: 'completed', + }, + limit: 100, + maxLimit: 500, + }); + const waitForRunEventNotification = jest + .spyOn(AgentRunEventService, 'waitForRunEventNotification') + .mockResolvedValue(true); + mockRunQuery.mockReturnValue({ + findOne: jest.fn().mockResolvedValue({ + uuid: 'run-1', + status: 'completed', + }), + }); + + const text = await new Response( + AgentRunEventService.createCanonicalRunEventStream('run-1', 1, { pollIntervalMs: 10 }) + ).text(); + + expect(waitForRunEventNotification).toHaveBeenCalledWith('run-1', 1, 10); + expect(listRunEventsPage).toHaveBeenNthCalledWith(3, 'run-1', { + afterSequence: 1, + limit: 100, + }); + expect(text).toContain('id: 2\nevent: run.completed'); + }); + + it('keeps following terminal runs until a terminal event is available', async () => { + const terminalEvent = { + uuid: 'event-2', + runUuid: 'run-1', + threadUuid: 'thread-1', + sessionUuid: 'session-1', + runId: 17, + sequence: 2, + eventType: 'run.completed', + payload: { + status: 'completed', + }, + createdAt: null, + updatedAt: null, + } as any; + const listRunEventsPage = jest + .spyOn(AgentRunEventService, 'listRunEventsPage') + .mockResolvedValueOnce({ + events: [], + nextSequence: 1, + hasMore: false, + run: { + id: 'run-1', + status: 'running', + }, + limit: 100, + maxLimit: 500, + }) + .mockResolvedValueOnce({ + events: [], + nextSequence: 1, + hasMore: false, + run: { + id: 'run-1', + status: 'completed', + }, + limit: 100, + maxLimit: 500, + }) + .mockResolvedValueOnce({ + events: [terminalEvent], + nextSequence: 2, + hasMore: false, + run: { + id: 'run-1', + status: 'completed', + }, + limit: 100, + maxLimit: 500, + }); + const waitForRunEventNotification = jest + .spyOn(AgentRunEventService, 'waitForRunEventNotification') + .mockResolvedValue(false); + mockRunQuery.mockReturnValue({ + findOne: jest.fn().mockResolvedValue({ + uuid: 'run-1', + status: 'completed', + }), + }); + + const text = await new Response( + AgentRunEventService.createCanonicalRunEventStream('run-1', 1, { pollIntervalMs: 10 }) + ).text(); + + expect(waitForRunEventNotification).toHaveBeenCalledWith('run-1', 1, 10); + expect(listRunEventsPage).toHaveBeenCalledTimes(3); + expect(text).toContain('id: 2\nevent: run.completed'); + }); +}); diff --git a/src/server/services/agent/__tests__/RunExecutor.test.ts b/src/server/services/agent/__tests__/RunExecutor.test.ts index 7a4c9ce0..a6340d75 100644 --- a/src/server/services/agent/__tests__/RunExecutor.test.ts +++ b/src/server/services/agent/__tests__/RunExecutor.test.ts @@ -49,24 +49,41 @@ jest.mock('server/services/agent/CapabilityService', () => ({ }, })); -const mockCreateRun = jest.fn().mockResolvedValue({ id: 11, uuid: 'run-1', status: 'running' }); +const mockCreateQueuedRun = jest.fn().mockResolvedValue({ id: 11, uuid: 'run-1', status: 'queued' }); +const mockClaimQueuedRunForExecution = jest + .fn() + .mockResolvedValue({ id: 11, uuid: 'run-1', status: 'starting', executionOwner: 'worker-1' }); const mockRegisterAbortController = jest.fn(); -const mockPatchRun = jest.fn().mockResolvedValue(undefined); +const mockClearAbortController = jest.fn(); +const mockPatchProgressForExecutionOwner = jest.fn().mockResolvedValue(undefined); +const mockHeartbeatRunExecution = jest.fn().mockResolvedValue(undefined); const mockGetRunByUuid = jest.fn(); -const mockPatchStatus = jest.fn(); const mockMarkFailed = jest.fn(); -const mockMarkCompleted = jest.fn(); +const mockMarkFailedForExecutionOwner = jest.fn(); +const mockStartRunForExecutionOwner = jest.fn(); +let mockLastFinalizeResult: unknown; +const mockFinalizeRunForExecutionOwner = jest.fn(async (_runId, _owner, finalize) => { + mockLastFinalizeResult = await finalize({ + run: { id: 11, uuid: 'run-1', status: 'running', executionOwner: 'worker-1' }, + trx: { trx: true }, + }); + return { id: 11, uuid: 'run-1', status: (mockLastFinalizeResult as { status: string }).status }; +}); jest.mock('server/services/agent/RunService', () => ({ __esModule: true, default: { - createRun: (...args: unknown[]) => mockCreateRun(...args), + createQueuedRun: (...args: unknown[]) => mockCreateQueuedRun(...args), + claimQueuedRunForExecution: (...args: unknown[]) => mockClaimQueuedRunForExecution(...args), registerAbortController: (...args: unknown[]) => mockRegisterAbortController(...args), - patchRun: (...args: unknown[]) => mockPatchRun(...args), + clearAbortController: (...args: unknown[]) => mockClearAbortController(...args), + patchProgressForExecutionOwner: (...args: unknown[]) => mockPatchProgressForExecutionOwner(...args), + heartbeatRunExecution: (...args: unknown[]) => mockHeartbeatRunExecution(...args), getRunByUuid: (...args: unknown[]) => mockGetRunByUuid(...args), - patchStatus: (...args: unknown[]) => mockPatchStatus(...args), markFailed: (...args: unknown[]) => mockMarkFailed(...args), - markCompleted: (...args: unknown[]) => mockMarkCompleted(...args), + markFailedForExecutionOwner: (...args: unknown[]) => mockMarkFailedForExecutionOwner(...args), + startRunForExecutionOwner: (...args: unknown[]) => mockStartRunForExecutionOwner(...args), + finalizeRunForExecutionOwner: (...args: unknown[]) => mockFinalizeRunForExecutionOwner(...args), }, })); @@ -103,16 +120,43 @@ jest.mock('server/services/agent/ApprovalService', () => ({ __esModule: true, default: { syncApprovalRequestsFromMessages: jest.fn(), + syncApprovalRequestStateFromMessages: jest.fn(), + }, +})); + +const mockEnqueueRun = jest.fn(); + +jest.mock('server/services/agent/RunQueueService', () => ({ + __esModule: true, + default: { + enqueueRun: (...args: unknown[]) => mockEnqueueRun(...args), }, })); jest.mock('server/services/agent/MessageStore', () => ({ __esModule: true, default: { - syncMessages: jest.fn(), + syncCanonicalMessagesFromUiMessages: jest.fn(), + upsertCanonicalUiMessagesForThread: jest.fn(), }, })); +jest.mock('server/lib/agentSession/runtimeConfig', () => { + const actual = jest.requireActual('server/lib/agentSession/runtimeConfig'); + return { + __esModule: true, + ...actual, + resolveAgentSessionDurabilityConfig: jest.fn().mockResolvedValue({ + runExecutionLeaseMs: 30 * 60 * 1000, + queuedRunDispatchStaleMs: 30 * 1000, + dispatchRecoveryLimit: 50, + maxDurablePayloadBytes: 64 * 1024, + payloadPreviewBytes: 16 * 1024, + fileChangePreviewChars: 4000, + }), + }; +}); + const mockToolExecutionInsert = jest.fn(); const mockToolExecutionFirst = jest.fn(); const mockToolExecutionPatchAndFetchById = jest.fn(); @@ -159,7 +203,9 @@ import AgentSessionService from 'server/services/agentSession'; import { SessionWorkspaceGatewayUnavailableError } from 'server/services/agent/errors'; const mockSyncApprovalRequests = ApprovalService.syncApprovalRequestsFromMessages as jest.Mock; -const mockSyncMessages = AgentMessageStore.syncMessages as jest.Mock; +const mockSyncApprovalRequestState = ApprovalService.syncApprovalRequestStateFromMessages as jest.Mock; +const mockSyncCanonicalMessagesFromUiMessages = AgentMessageStore.syncCanonicalMessagesFromUiMessages as jest.Mock; +const mockUpsertCanonicalUiMessagesForThread = AgentMessageStore.upsertCanonicalUiMessagesForThread as jest.Mock; const mockMarkSessionRuntimeFailure = AgentSessionService.markSessionRuntimeFailure as jest.Mock; describe('AgentRunExecutor', () => { @@ -173,12 +219,32 @@ describe('AgentRunExecutor', () => { binding: null, }); mockBuildToolSet.mockResolvedValue({}); - mockCreateRun.mockResolvedValue({ id: 11, uuid: 'run-1', status: 'running' }); - mockPatchRun.mockResolvedValue(undefined); + mockCreateQueuedRun.mockResolvedValue({ id: 11, uuid: 'run-1', status: 'queued' }); + mockClaimQueuedRunForExecution.mockResolvedValue({ + id: 11, + uuid: 'run-1', + status: 'starting', + executionOwner: 'worker-1', + }); + mockPatchProgressForExecutionOwner.mockResolvedValue(undefined); + mockHeartbeatRunExecution.mockResolvedValue(undefined); mockGetRunByUuid.mockResolvedValue({ id: 11, uuid: 'run-1', status: 'running' }); - mockPatchStatus.mockResolvedValue(undefined); mockMarkFailed.mockResolvedValue(undefined); - mockMarkCompleted.mockResolvedValue(undefined); + mockMarkFailedForExecutionOwner.mockResolvedValue(undefined); + mockStartRunForExecutionOwner.mockImplementation(async (runId, owner) => ({ + id: 11, + uuid: runId, + status: 'running', + executionOwner: owner, + })); + mockLastFinalizeResult = null; + mockFinalizeRunForExecutionOwner.mockImplementation(async (_runId, _owner, finalize) => { + mockLastFinalizeResult = await finalize({ + run: { id: 11, uuid: 'run-1', status: 'running', executionOwner: 'worker-1' }, + trx: { trx: true }, + }); + return { id: 11, uuid: 'run-1', status: (mockLastFinalizeResult as { status: string }).status }; + }); mockGetSessionAppendSystemPrompt.mockResolvedValue('Append prompt'); mockTouchActivity.mockResolvedValue(undefined); mockMarkSessionRuntimeFailure.mockResolvedValue(undefined); @@ -192,8 +258,14 @@ describe('AgentRunExecutor', () => { }); mockPendingActionFirst.mockResolvedValue(null); mockToolExecutionFirst.mockResolvedValue(undefined); - mockSyncMessages.mockImplementation(async (_threadId, _userId, messages) => messages); + mockSyncCanonicalMessagesFromUiMessages.mockImplementation(async (_threadId, _userId, messages) => messages); + mockUpsertCanonicalUiMessagesForThread.mockResolvedValue(undefined); mockSyncApprovalRequests.mockResolvedValue([]); + mockSyncApprovalRequestState.mockResolvedValue({ + pendingActions: [], + resolvedActionCount: 0, + }); + mockEnqueueRun.mockResolvedValue(undefined); }); it('builds agent instructions from the control-plane and session prompts', async () => { @@ -323,10 +395,33 @@ describe('AgentRunExecutor', () => { }) ).rejects.toThrow('tool setup failed'); - expect(mockCreateRun).not.toHaveBeenCalled(); + expect(mockCreateQueuedRun).not.toHaveBeenCalled(); expect(mockMarkFailed).not.toHaveBeenCalled(); }); + it('marks an existing queued run failed when setup fails before execution starts', async () => { + mockBuildToolSet.mockRejectedValueOnce(new Error('tool setup failed')); + + await expect( + AgentRunExecutor.execute({ + session: { uuid: 'sess-1' } as any, + thread: { id: 7, uuid: 'thread-1' } as any, + userIdentity: { userId: 'sample-user' } as any, + messages: [], + existingRun: { id: 11, uuid: 'queued-run-1', status: 'queued', executionOwner: 'worker-1' } as any, + }) + ).rejects.toThrow('tool setup failed'); + + expect(mockCreateQueuedRun).not.toHaveBeenCalled(); + expect(mockMarkFailedForExecutionOwner).toHaveBeenCalledWith( + 'queued-run-1', + 'worker-1', + expect.objectContaining({ message: 'tool setup failed' }), + expect.any(Object), + { dispatchAttemptId: undefined } + ); + }); + it('records a runtime session failure when the workspace gateway is unavailable', async () => { mockBuildToolSet.mockRejectedValueOnce( new SessionWorkspaceGatewayUnavailableError({ @@ -348,7 +443,7 @@ describe('AgentRunExecutor', () => { 'sess-1', expect.any(SessionWorkspaceGatewayUnavailableError) ); - expect(mockCreateRun).not.toHaveBeenCalled(); + expect(mockCreateQueuedRun).not.toHaveBeenCalled(); }); it('marks the run failed if agent construction throws after the run is created', async () => { @@ -365,11 +460,14 @@ describe('AgentRunExecutor', () => { }) ).rejects.toThrow('agent init failed'); - expect(mockCreateRun).toHaveBeenCalled(); - expect(mockMarkFailed).toHaveBeenCalledWith( + expect(mockCreateQueuedRun).toHaveBeenCalled(); + expect(mockClaimQueuedRunForExecution).toHaveBeenCalledWith('run-1', expect.stringMatching(/^direct:/)); + expect(mockMarkFailedForExecutionOwner).toHaveBeenCalledWith( 'run-1', + expect.stringMatching(/^direct:/), expect.objectContaining({ message: 'agent init failed' }), - expect.any(Object) + expect.any(Object), + { dispatchAttemptId: undefined } ); }); @@ -394,25 +492,126 @@ describe('AgentRunExecutor', () => { isAborted: false, }); - expect(mockMarkFailed).toHaveBeenCalledWith( - 'run-1', + expect(mockLastFinalizeResult).toEqual( expect.objectContaining({ - code: 'max_iterations_exceeded', - details: expect.objectContaining({ - finishReason: 'tool-calls', - maxIterations: 8, + status: 'failed', + error: expect.objectContaining({ + code: 'max_iterations_exceeded', + details: expect.objectContaining({ + finishReason: 'tool-calls', + maxIterations: 8, + }), }), - }), - expect.any(Object), - { - finishReason: 'tool-calls', - } + }) ); - expect(mockMarkCompleted).not.toHaveBeenCalled(); + expect(mockMarkFailedForExecutionOwner).not.toHaveBeenCalled(); + }); + + it('marks the run waiting when finalization syncs pending approvals', async () => { + mockSyncApprovalRequestState.mockResolvedValueOnce({ + pendingActions: [{ id: 99 }], + resolvedActionCount: 0, + }); + + const execution = await AgentRunExecutor.execute({ + session: { uuid: 'sess-1', id: 17 } as any, + thread: { id: 7, uuid: 'thread-1' } as any, + userIdentity: { userId: 'sample-user' } as any, + messages: [], + }); + + await execution.onStreamFinish({ + messages: [ + { + id: 'assistant-1', + role: 'assistant', + parts: [], + metadata: { runId: 'run-1' }, + } as any, + ], + finishReason: 'tool-calls', + isAborted: false, + }); + + expect(mockLastFinalizeResult).toEqual( + expect.objectContaining({ + status: 'waiting_for_approval', + }) + ); + expect(mockEnqueueRun).not.toHaveBeenCalled(); + }); + + it('keeps the owner heartbeat active until stream finalization or dispose', async () => { + jest.useFakeTimers(); + + try { + const execution = await AgentRunExecutor.execute({ + session: { uuid: 'sess-1', id: 17 } as any, + thread: { id: 7, uuid: 'thread-1' } as any, + userIdentity: { userId: 'sample-user' } as any, + messages: [], + }); + + expect(mockHeartbeatRunExecution).not.toHaveBeenCalled(); + + jest.advanceTimersByTime(60_000); + await Promise.resolve(); + + expect(mockHeartbeatRunExecution).toHaveBeenCalledWith('run-1', expect.stringMatching(/^direct:/)); + + mockHeartbeatRunExecution.mockClear(); + execution.dispose(); + jest.advanceTimersByTime(60_000); + await Promise.resolve(); + + expect(mockHeartbeatRunExecution).not.toHaveBeenCalled(); + } finally { + jest.useRealTimers(); + } + }); + + it('requeues the run when finalization finds already resolved approvals', async () => { + mockSyncApprovalRequestState.mockResolvedValueOnce({ + pendingActions: [], + resolvedActionCount: 1, + }); + + const execution = await AgentRunExecutor.execute({ + session: { uuid: 'sess-1', id: 17 } as any, + thread: { id: 7, uuid: 'thread-1' } as any, + userIdentity: { userId: 'sample-user' } as any, + messages: [], + requestGitHubToken: 'sample-gh-token', + }); + + await execution.onStreamFinish({ + messages: [ + { + id: 'assistant-1', + role: 'assistant', + parts: [], + metadata: { runId: 'run-1' }, + } as any, + ], + finishReason: 'tool-calls', + isAborted: false, + }); + + expect(mockLastFinalizeResult).toEqual( + expect.objectContaining({ + status: 'queued', + patch: expect.objectContaining({ + queuedAt: expect.any(String), + }), + }) + ); + expect(mockEnqueueRun).toHaveBeenCalledWith('run-1', 'approval_resolved', { + githubToken: 'sample-gh-token', + }); }); it('marks the run failed if stream finalization persistence throws', async () => { - mockSyncMessages.mockRejectedValueOnce(new Error('message sync failed')); + mockUpsertCanonicalUiMessagesForThread.mockRejectedValueOnce(new Error('message sync failed')); const execution = await AgentRunExecutor.execute({ session: { uuid: 'sess-1', id: 17 } as any, @@ -429,14 +628,12 @@ describe('AgentRunExecutor', () => { }) ).rejects.toThrow('message sync failed'); - expect(mockMarkFailed).toHaveBeenCalledWith( + expect(mockMarkFailedForExecutionOwner).toHaveBeenCalledWith( 'run-1', + expect.stringMatching(/^direct:/), expect.objectContaining({ message: 'message sync failed' }), expect.any(Object), - { - finishReason: 'stop', - } + { dispatchAttemptId: undefined } ); - expect(mockMarkCompleted).not.toHaveBeenCalled(); }); }); diff --git a/src/server/services/agent/__tests__/RunQueueService.test.ts b/src/server/services/agent/__tests__/RunQueueService.test.ts new file mode 100644 index 00000000..661ca960 --- /dev/null +++ b/src/server/services/agent/__tests__/RunQueueService.test.ts @@ -0,0 +1,146 @@ +/** + * Copyright 2026 GoodRx, Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +jest.mock('server/lib/queueManager', () => { + const mockState = { + queue: { + add: jest.fn(), + }, + registerCalls: [] as unknown[][], + }; + (global as any).__runQueueServiceMockState = mockState; + + const queue = { + add: mockState.queue.add, + }; + const manager = { + registerQueue: jest.fn((...args: unknown[]) => { + mockState.registerCalls.push(args); + return queue; + }), + }; + return { + __esModule: true, + default: { + getInstance: jest.fn(() => manager), + }, + }; +}); + +jest.mock('server/lib/redisClient', () => ({ + __esModule: true, + default: { + getInstance: jest.fn(() => ({ + getConnection: jest.fn(() => ({ + duplicate: jest.fn(), + })), + })), + }, +})); + +jest.mock('server/lib/encryption', () => ({ + encrypt: jest.fn((value: string) => `encrypted:${value}`), +})); + +jest.mock('server/lib/logger', () => ({ + extractContextForQueue: jest.fn(() => ({ + correlationId: 'correlation-1', + sender: 'sample-user', + })), +})); + +import { encrypt } from 'server/lib/encryption'; +import { extractContextForQueue } from 'server/lib/logger'; +import AgentRunQueueService from '../RunQueueService'; + +const mockState = (global as any).__runQueueServiceMockState as { + queue: { + add: jest.Mock; + }; + registerCalls: unknown[][]; +}; +const mockQueueAdd = mockState.queue.add; +const mockEncrypt = encrypt as jest.Mock; +const mockExtractContextForQueue = extractContextForQueue as jest.Mock; + +describe('AgentRunQueueService', () => { + beforeEach(() => { + mockQueueAdd.mockClear(); + mockEncrypt.mockClear(); + mockExtractContextForQueue.mockClear(); + }); + + it('registers the execute queue with bounded failed-job retention', () => { + expect(mockState.registerCalls).toEqual( + expect.arrayContaining([ + [ + 'agent_run_execute', + expect.objectContaining({ + defaultJobOptions: expect.objectContaining({ + attempts: 1, + removeOnComplete: true, + removeOnFail: 100, + }), + }), + ], + ]) + ); + }); + + it('enqueues a dispatch attempt id in the job payload and return value', async () => { + mockQueueAdd.mockResolvedValue(undefined); + + const result = await AgentRunQueueService.enqueueRun('run-1', 'submit', { + githubToken: ' token-1 ', + }); + + expect(result.dispatchAttemptId).toEqual(expect.any(String)); + expect(result.dispatchAttemptId).not.toHaveLength(0); + expect(mockQueueAdd).toHaveBeenCalledWith( + 'execute-run', + expect.objectContaining({ + runId: 'run-1', + reason: 'submit', + dispatchAttemptId: result.dispatchAttemptId, + encryptedGithubToken: 'encrypted:token-1', + correlationId: 'correlation-1', + sender: 'sample-user', + }), + expect.objectContaining({ + jobId: `agent-run:run-1:${result.dispatchAttemptId}`, + }) + ); + }); + + it('uses a distinct BullMQ job id for each dispatch attempt and keeps reason out of the uniqueness boundary', async () => { + mockQueueAdd.mockResolvedValue(undefined); + + const first = await AgentRunQueueService.enqueueRun('run-1', 'resume'); + const second = await AgentRunQueueService.enqueueRun('run-1', 'resume'); + + const firstOptions = mockQueueAdd.mock.calls[0][2]; + const secondOptions = mockQueueAdd.mock.calls[1][2]; + + expect(first.dispatchAttemptId).not.toEqual(second.dispatchAttemptId); + expect(firstOptions.jobId).toBe(`agent-run:run-1:${first.dispatchAttemptId}`); + expect(secondOptions.jobId).toBe(`agent-run:run-1:${second.dispatchAttemptId}`); + expect(firstOptions.jobId).not.toEqual(secondOptions.jobId); + expect(firstOptions.jobId).not.toContain(':resume'); + expect(secondOptions.jobId).not.toContain(':resume'); + expect(mockQueueAdd.mock.calls[0][1]).toEqual(expect.objectContaining({ reason: 'resume' })); + expect(mockQueueAdd.mock.calls[1][1]).toEqual(expect.objectContaining({ reason: 'resume' })); + }); +}); diff --git a/src/server/services/agent/__tests__/RunService.test.ts b/src/server/services/agent/__tests__/RunService.test.ts index 8e48ec62..99f28f46 100644 --- a/src/server/services/agent/__tests__/RunService.test.ts +++ b/src/server/services/agent/__tests__/RunService.test.ts @@ -15,6 +15,14 @@ */ jest.mock('server/models/AgentRun', () => ({ + __esModule: true, + default: { + query: jest.fn(), + transaction: jest.fn(), + }, +})); + +jest.mock('server/models/AgentSession', () => ({ __esModule: true, default: { query: jest.fn(), @@ -23,15 +31,65 @@ jest.mock('server/models/AgentRun', () => ({ jest.mock('server/lib/dependencies', () => ({})); +jest.mock('../RunEventService', () => ({ + __esModule: true, + default: { + appendStatusEvent: jest.fn(), + appendStatusEventForRunInTransaction: jest.fn(), + appendChunkEventsForRunInTransaction: jest.fn(), + notifyRunEventsInserted: jest.fn(), + }, +})); + +jest.mock('server/lib/agentSession/runtimeConfig', () => { + const actual = jest.requireActual('server/lib/agentSession/runtimeConfig'); + return { + __esModule: true, + ...actual, + resolveAgentSessionDurabilityConfig: jest.fn().mockResolvedValue({ + runExecutionLeaseMs: 30 * 60 * 1000, + queuedRunDispatchStaleMs: 30 * 1000, + dispatchRecoveryLimit: 50, + maxDurablePayloadBytes: 64 * 1024, + payloadPreviewBytes: 16 * 1024, + fileChangePreviewChars: 4000, + }), + }; +}); + import AgentRunService from '../RunService'; import AgentRun from 'server/models/AgentRun'; +import AgentSession from 'server/models/AgentSession'; +import AgentRunEventService from '../RunEventService'; +import { AgentRunOwnershipLostError } from '../AgentRunOwnershipLostError'; +import { resolveAgentSessionDurabilityConfig } from 'server/lib/agentSession/runtimeConfig'; const mockRunQuery = AgentRun.query as jest.Mock; +const mockRunTransaction = AgentRun.transaction as jest.Mock; +const mockSessionQuery = AgentSession.query as jest.Mock; +const mockAppendStatusEvent = AgentRunEventService.appendStatusEvent as jest.Mock; +const mockAppendStatusEventForRunInTransaction = AgentRunEventService.appendStatusEventForRunInTransaction as jest.Mock; +const mockAppendChunkEventsForRunInTransaction = AgentRunEventService.appendChunkEventsForRunInTransaction as jest.Mock; +const mockNotifyRunEventsInserted = AgentRunEventService.notifyRunEventsInserted as jest.Mock; +const mockResolveDurabilityConfig = resolveAgentSessionDurabilityConfig as jest.Mock; const VALID_RUN_UUID = '123e4567-e89b-12d3-a456-426614174000'; describe('AgentRunService', () => { beforeEach(() => { jest.clearAllMocks(); + mockRunTransaction.mockImplementation(async (callback) => callback({ trx: true })); + mockResolveDurabilityConfig.mockResolvedValue({ + runExecutionLeaseMs: 30 * 60 * 1000, + queuedRunDispatchStaleMs: 30 * 1000, + dispatchRecoveryLimit: 50, + maxDurablePayloadBytes: 64 * 1024, + payloadPreviewBytes: 16 * 1024, + fileChangePreviewChars: 4000, + }); + }); + + afterEach(() => { + jest.restoreAllMocks(); }); describe('getOwnedRun', () => { @@ -88,4 +146,454 @@ describe('AgentRunService', () => { expect(findOne).toHaveBeenCalledWith({ uuid: VALID_RUN_UUID }); }); }); + + describe('listRunsNeedingDispatch', () => { + it('finds stale queued runs and expired execution leases', async () => { + const staleQueuedBuilder: any = { + where: jest.fn().mockReturnThis(), + }; + const expiredLeaseBuilder: any = { + whereIn: jest.fn().mockReturnThis(), + whereNotNull: jest.fn().mockReturnThis(), + where: jest.fn().mockReturnThis(), + }; + const runs = [{ uuid: VALID_RUN_UUID }]; + const query: any = { + where: jest.fn((callback) => { + callback(staleQueuedBuilder); + return query; + }), + orWhere: jest.fn((callback) => { + callback(expiredLeaseBuilder); + return query; + }), + orderBy: jest.fn().mockReturnThis(), + limit: jest.fn().mockResolvedValue(runs), + }; + mockRunQuery.mockReturnValue(query); + + const now = new Date('2026-04-24T12:00:00.000Z'); + await expect( + AgentRunService.listRunsNeedingDispatch({ + now, + queuedStaleMs: 10_000, + limit: 25, + }) + ).resolves.toBe(runs); + + expect(staleQueuedBuilder.where).toHaveBeenNthCalledWith(1, 'status', 'queued'); + expect(staleQueuedBuilder.where).toHaveBeenNthCalledWith(2, 'queuedAt', '<', '2026-04-24T11:59:50.000Z'); + expect(expiredLeaseBuilder.whereIn).toHaveBeenCalledWith('status', ['starting', 'running']); + expect(expiredLeaseBuilder.whereNotNull).toHaveBeenCalledWith('leaseExpiresAt'); + expect(expiredLeaseBuilder.where).toHaveBeenCalledWith('leaseExpiresAt', '<=', '2026-04-24T12:00:00.000Z'); + expect(query.orderBy).toHaveBeenCalledWith('updatedAt', 'asc'); + expect(query.limit).toHaveBeenCalledWith(25); + }); + }); + + describe('cancelRun', () => { + it('records cancellation through the shared status patch path', async () => { + const runningRun = { + id: 1, + uuid: VALID_RUN_UUID, + status: 'running', + }; + const cancelledRun = { + ...runningRun, + status: 'cancelled', + }; + const getOwnedRun = jest + .spyOn(AgentRunService, 'getOwnedRun') + .mockResolvedValueOnce(runningRun as Awaited>) + .mockResolvedValueOnce(cancelledRun as Awaited>); + const patchStatus = jest + .spyOn(AgentRunService, 'patchStatus') + .mockResolvedValue(cancelledRun as Awaited>); + + await expect(AgentRunService.cancelRun(VALID_RUN_UUID, 'sample-user')).resolves.toBe(cancelledRun); + + expect(patchStatus).toHaveBeenCalledWith( + VALID_RUN_UUID, + 'cancelled', + expect.objectContaining({ + cancelledAt: expect.any(String), + completedAt: expect.any(String), + }) + ); + expect(getOwnedRun).toHaveBeenCalledTimes(2); + }); + }); + + describe('markFailed', () => { + it('uses the same serialized error on the run and run.failed event payload', async () => { + const findOne = jest.fn().mockResolvedValue({ + id: 17, + uuid: VALID_RUN_UUID, + }); + const patchAndFetchById = jest.fn().mockImplementation((_id, patch) => + Promise.resolve({ + id: 17, + uuid: VALID_RUN_UUID, + status: 'failed', + ...patch, + }) + ); + mockRunQuery.mockReturnValueOnce({ findOne }).mockReturnValueOnce({ patchAndFetchById }); + const error = Object.assign(new Error('Sample run failure.'), { + name: 'SampleRunError', + code: 'sample_failure', + details: { + reason: 'sample', + }, + }); + + const failedRun = await AgentRunService.markFailed(VALID_RUN_UUID, error, { + totalTokens: 12, + }); + const patch = patchAndFetchById.mock.calls[0][1]; + const eventPayload = mockAppendStatusEvent.mock.calls[0][2]; + + expect(failedRun.error).toEqual(patch.error); + expect(eventPayload.error).toEqual(failedRun.error); + expect(mockAppendStatusEvent).toHaveBeenCalledWith( + VALID_RUN_UUID, + 'run.failed', + expect.objectContaining({ + status: 'failed', + error: failedRun.error, + usageSummary: { + totalTokens: 12, + }, + }) + ); + }); + }); + + describe('patchStatus', () => { + it('emits canonical run status event names for approval waits and resumed runs', async () => { + const findOne = jest.fn().mockResolvedValue({ + id: 17, + uuid: VALID_RUN_UUID, + }); + const patchAndFetchById = jest.fn().mockImplementation((_id, patch) => + Promise.resolve({ + id: 17, + uuid: VALID_RUN_UUID, + usageSummary: {}, + error: null, + ...patch, + }) + ); + mockRunQuery + .mockReturnValueOnce({ findOne }) + .mockReturnValueOnce({ patchAndFetchById }) + .mockReturnValueOnce({ findOne }) + .mockReturnValueOnce({ patchAndFetchById }); + + await AgentRunService.patchStatus(VALID_RUN_UUID, 'waiting_for_approval'); + await AgentRunService.patchStatus(VALID_RUN_UUID, 'queued'); + + expect(mockAppendStatusEvent).toHaveBeenNthCalledWith( + 1, + VALID_RUN_UUID, + 'run.waiting_for_approval', + expect.objectContaining({ + status: 'waiting_for_approval', + }) + ); + expect(mockAppendStatusEvent).toHaveBeenNthCalledWith( + 2, + VALID_RUN_UUID, + 'run.queued', + expect.objectContaining({ + status: 'queued', + }) + ); + }); + }); + + describe('owner-aware execution helpers', () => { + it('updates a matching owner terminal status and emits one status event after the transition', async () => { + const ownedRun = { + id: 17, + uuid: VALID_RUN_UUID, + status: 'running', + executionOwner: 'worker-1', + }; + const completedRun = { + ...ownedRun, + status: 'completed', + executionOwner: null, + leaseExpiresAt: null, + heartbeatAt: null, + usageSummary: { + totalTokens: 12, + }, + }; + const findOne = jest.fn().mockReturnValue({ + forUpdate: jest.fn().mockResolvedValue(ownedRun), + }); + const patchAndFetchById = jest.fn().mockResolvedValue(completedRun); + mockAppendStatusEventForRunInTransaction.mockResolvedValue(12); + + mockRunQuery.mockReturnValueOnce({ findOne }).mockReturnValueOnce({ patchAndFetchById }); + + await expect( + AgentRunService.markCompletedForExecutionOwner(VALID_RUN_UUID, 'worker-1', { + totalTokens: 12, + }) + ).resolves.toBe(completedRun); + + expect(patchAndFetchById).toHaveBeenCalledWith( + 17, + expect.objectContaining({ + status: 'completed', + executionOwner: null, + leaseExpiresAt: null, + heartbeatAt: null, + usageSummary: { + totalTokens: 12, + }, + }) + ); + expect(mockAppendStatusEventForRunInTransaction).toHaveBeenCalledTimes(1); + expect(mockAppendStatusEventForRunInTransaction).toHaveBeenCalledWith( + completedRun, + 'run.completed', + expect.objectContaining({ + status: 'completed', + usageSummary: { + totalTokens: 12, + }, + }), + { trx: true } + ); + expect(mockNotifyRunEventsInserted).toHaveBeenCalledWith(VALID_RUN_UUID, 12); + }); + + it('throws ownership loss without patching or appending a status event when the owner is stale', async () => { + const run = { + id: 17, + uuid: VALID_RUN_UUID, + status: 'running', + executionOwner: 'worker-2', + }; + const findOne = jest.fn().mockReturnValue({ + forUpdate: jest.fn().mockResolvedValue(run), + }); + + mockRunQuery.mockReturnValueOnce({ findOne }); + + let error: unknown; + try { + await AgentRunService.markCompletedForExecutionOwner(VALID_RUN_UUID, 'worker-1', { + totalTokens: 12, + }); + } catch (caught) { + error = caught; + } + + expect(error).toBeInstanceOf(AgentRunOwnershipLostError); + expect(error).toMatchObject({ + runUuid: VALID_RUN_UUID, + expectedExecutionOwner: 'worker-1', + currentStatus: 'running', + currentExecutionOwner: 'worker-2', + }); + + expect(mockRunQuery).toHaveBeenCalledTimes(1); + expect(mockAppendStatusEvent).not.toHaveBeenCalled(); + expect(mockAppendStatusEventForRunInTransaction).not.toHaveBeenCalled(); + }); + + it('does not append stream chunks when the owner is stale', async () => { + const run = { + id: 17, + uuid: VALID_RUN_UUID, + status: 'running', + executionOwner: 'worker-2', + }; + const findOne = jest.fn().mockReturnValue({ + forUpdate: jest.fn().mockResolvedValue(run), + }); + const beforeAppendChunks = jest.fn(); + + mockRunQuery.mockReturnValueOnce({ findOne }); + + await expect( + AgentRunService.appendStreamChunksForExecutionOwner( + VALID_RUN_UUID, + 'worker-1', + [{ type: 'text-delta', id: 'text-1', delta: 'stale' } as any], + { beforeAppendChunks } + ) + ).rejects.toBeInstanceOf(AgentRunOwnershipLostError); + + expect(beforeAppendChunks).not.toHaveBeenCalled(); + expect(mockAppendChunkEventsForRunInTransaction).not.toHaveBeenCalled(); + }); + + it('does not run final message sync when the owner is stale', async () => { + const run = { + id: 17, + uuid: VALID_RUN_UUID, + status: 'running', + executionOwner: 'worker-2', + }; + const findOne = jest.fn().mockReturnValue({ + forUpdate: jest.fn().mockResolvedValue(run), + }); + const finalize = jest.fn(); + + mockRunQuery.mockReturnValueOnce({ findOne }); + + await expect( + AgentRunService.finalizeRunForExecutionOwner(VALID_RUN_UUID, 'worker-1', finalize) + ).rejects.toBeInstanceOf(AgentRunOwnershipLostError); + + expect(finalize).not.toHaveBeenCalled(); + expect(mockAppendStatusEventForRunInTransaction).not.toHaveBeenCalled(); + }); + + it('releases ownership when finalization queues a resolved approval continuation', async () => { + const ownedRun = { + id: 17, + uuid: VALID_RUN_UUID, + status: 'running', + executionOwner: 'worker-1', + }; + const queuedRun = { + ...ownedRun, + status: 'queued', + executionOwner: null, + leaseExpiresAt: null, + heartbeatAt: null, + }; + const findOne = jest.fn().mockReturnValue({ + forUpdate: jest.fn().mockResolvedValue(ownedRun), + }); + const patchAndFetchById = jest.fn().mockResolvedValue(queuedRun); + const finalize = jest.fn().mockResolvedValue({ + status: 'queued', + patch: { + queuedAt: '2026-04-24T12:00:00.000Z', + }, + }); + mockAppendStatusEventForRunInTransaction.mockResolvedValue(31); + + mockRunQuery.mockReturnValueOnce({ findOne }).mockReturnValueOnce({ patchAndFetchById }); + + await expect(AgentRunService.finalizeRunForExecutionOwner(VALID_RUN_UUID, 'worker-1', finalize)).resolves.toBe( + queuedRun + ); + + expect(patchAndFetchById).toHaveBeenCalledWith( + 17, + expect.objectContaining({ + status: 'queued', + queuedAt: '2026-04-24T12:00:00.000Z', + executionOwner: null, + leaseExpiresAt: null, + heartbeatAt: null, + }) + ); + expect(mockAppendStatusEventForRunInTransaction).toHaveBeenCalledWith( + queuedRun, + 'run.queued', + expect.objectContaining({ + status: 'queued', + executionOwner: 'worker-1', + }), + { trx: true } + ); + expect(mockNotifyRunEventsInserted).toHaveBeenCalledWith(VALID_RUN_UUID, 31); + }); + }); + + describe('heartbeatRunExecution', () => { + it('throws ownership loss when the conditional heartbeat update matches no rows', async () => { + const patch = jest.fn().mockResolvedValue(0); + const heartbeatQuery = { + where: jest.fn().mockReturnThis(), + whereNotIn: jest.fn().mockReturnThis(), + patch, + }; + const findOne = jest.fn().mockResolvedValue({ + uuid: VALID_RUN_UUID, + status: 'running', + executionOwner: 'worker-2', + }); + mockRunQuery.mockReturnValueOnce(heartbeatQuery).mockReturnValueOnce({ findOne }); + + await expect(AgentRunService.heartbeatRunExecution(VALID_RUN_UUID, 'worker-1')).rejects.toMatchObject({ + runUuid: VALID_RUN_UUID, + expectedExecutionOwner: 'worker-1', + currentStatus: 'running', + currentExecutionOwner: 'worker-2', + }); + }); + }); + + describe('claimQueuedRunForExecution', () => { + it('claims a queued run under a session row lock', async () => { + const run = { + id: 17, + uuid: VALID_RUN_UUID, + sessionId: 23, + status: 'queued', + }; + const findOne = jest.fn().mockReturnValue({ + forUpdate: jest.fn().mockResolvedValue(run), + }); + const patchAndFetchById = jest.fn().mockResolvedValue({ + ...run, + status: 'starting', + executionOwner: 'worker-1', + }); + mockRunQuery.mockReturnValueOnce({ findOne }).mockReturnValueOnce({ patchAndFetchById }); + mockSessionQuery.mockReturnValue({ + findById: jest.fn().mockReturnValue({ + forUpdate: jest.fn().mockResolvedValue({ id: 23 }), + }), + }); + + await expect( + AgentRunService.claimQueuedRunForExecution(VALID_RUN_UUID, 'worker-1', 30 * 60 * 1000) + ).resolves.toEqual( + expect.objectContaining({ + status: 'starting', + executionOwner: 'worker-1', + }) + ); + + expect(findOne).toHaveBeenCalledWith({ uuid: VALID_RUN_UUID }); + expect(patchAndFetchById).toHaveBeenCalledWith( + 17, + expect.objectContaining({ + status: 'starting', + executionOwner: 'worker-1', + leaseExpiresAt: expect.any(String), + heartbeatAt: expect.any(String), + }) + ); + }); + + it('skips a run that is already owned and not stale', async () => { + const run = { + id: 17, + uuid: VALID_RUN_UUID, + sessionId: 23, + status: 'running', + leaseExpiresAt: new Date(Date.now() + 60_000).toISOString(), + }; + const findOne = jest.fn().mockReturnValue({ + forUpdate: jest.fn().mockResolvedValue(run), + }); + mockRunQuery.mockReturnValueOnce({ findOne }); + + await expect( + AgentRunService.claimQueuedRunForExecution(VALID_RUN_UUID, 'worker-1', 30 * 60 * 1000) + ).resolves.toBeNull(); + }); + }); }); diff --git a/src/server/services/agent/__tests__/SessionReadService.test.ts b/src/server/services/agent/__tests__/SessionReadService.test.ts new file mode 100644 index 00000000..a8c4e440 --- /dev/null +++ b/src/server/services/agent/__tests__/SessionReadService.test.ts @@ -0,0 +1,234 @@ +/** + * Copyright 2026 GoodRx, Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +jest.mock('server/models/AgentSession', () => ({ + __esModule: true, + default: { + query: jest.fn(), + }, +})); + +jest.mock('server/models/AgentSource', () => ({ + __esModule: true, + default: { + query: jest.fn(), + }, +})); + +jest.mock('server/models/AgentSandbox', () => ({ + __esModule: true, + default: { + query: jest.fn(), + }, +})); + +jest.mock('server/models/AgentSandboxExposure', () => ({ + __esModule: true, + default: { + query: jest.fn(), + }, +})); + +jest.mock('server/models/AgentThread', () => ({ + __esModule: true, + default: { + query: jest.fn(), + }, +})); + +jest.mock('server/services/agent/SandboxService', () => ({ + __esModule: true, + default: { + serializeSandboxExposure: jest.fn((exposure) => ({ + id: exposure.uuid, + kind: exposure.kind, + status: exposure.status, + targetPort: exposure.targetPort, + url: exposure.url, + metadata: exposure.metadata || {}, + lastVerifiedAt: exposure.lastVerifiedAt, + endedAt: exposure.endedAt, + createdAt: exposure.createdAt || null, + updatedAt: exposure.updatedAt || null, + })), + }, +})); + +jest.mock('server/services/agent/ThreadService', () => ({ + __esModule: true, + default: { + serializeThread: jest.fn(), + }, +})); + +jest.mock('server/lib/dependencies', () => ({})); + +import AgentSession from 'server/models/AgentSession'; +import AgentSource from 'server/models/AgentSource'; +import AgentSandbox from 'server/models/AgentSandbox'; +import AgentSandboxExposure from 'server/models/AgentSandboxExposure'; +import AgentThread from 'server/models/AgentThread'; +import AgentSessionReadService from '../SessionReadService'; + +const mockSessionQuery = AgentSession.query as jest.Mock; +const mockSourceQuery = AgentSource.query as jest.Mock; +const mockSandboxQuery = AgentSandbox.query as jest.Mock; +const mockSandboxExposureQuery = AgentSandboxExposure.query as jest.Mock; +const mockThreadQuery = AgentThread.query as jest.Mock; + +function buildSession(overrides: Record = {}) { + return { + id: 17, + uuid: 'session-1', + status: 'active', + userId: 'sample-user', + ownerGithubUsername: 'sample-user', + defaultThreadId: 9, + defaultModel: 'gpt-5.4', + defaultHarness: 'lifecycle_ai_sdk', + buildUuid: 'build-1', + buildKind: 'environment', + sessionKind: 'environment', + workspaceStatus: 'ready', + lastActivity: '2026-04-24T12:00:00.000Z', + endedAt: null, + createdAt: '2026-04-24T12:00:00.000Z', + updatedAt: '2026-04-24T12:05:00.000Z', + workspaceRepos: [{ repo: 'example-org/example-repo', branch: 'main', mountPath: '/workspace/example-repo' }], + selectedServices: [{ name: 'sample-service' }], + ...overrides, + }; +} + +function buildPagedSessionQuery(results: unknown[], total: number) { + const query = { + where: jest.fn(), + whereIn: jest.fn(), + orderBy: jest.fn(), + page: jest.fn().mockResolvedValue({ results, total }), + }; + query.where.mockReturnValue(query); + query.whereIn.mockReturnValue(query); + query.orderBy.mockReturnValue(query); + return query; +} + +function buildOrderedQuery(rows: T[], orderCalls = 1) { + const query = { + whereIn: jest.fn(), + where: jest.fn(), + whereNull: jest.fn(), + orderBy: jest.fn(), + }; + query.whereIn.mockReturnValue(query); + query.where.mockReturnValue(query); + query.whereNull.mockReturnValue(query); + for (let index = 0; index < orderCalls - 1; index += 1) { + query.orderBy.mockReturnValueOnce(query); + } + query.orderBy.mockResolvedValueOnce(rows); + return query; +} + +describe('AgentSessionReadService', () => { + beforeEach(() => { + jest.clearAllMocks(); + }); + + it('lists owned sessions with capped pagination and batched related reads', async () => { + const session = buildSession(); + const source = { + id: 3, + uuid: 'source-1', + sessionId: 17, + adapter: 'lifecycle_environment', + status: 'ready', + input: {}, + sandboxRequirements: { filesystem: 'persistent' }, + error: null, + preparedAt: '2026-04-24T12:00:00.000Z', + cleanedUpAt: null, + createdAt: '2026-04-24T12:00:00.000Z', + updatedAt: '2026-04-24T12:00:00.000Z', + }; + const sandbox = { + id: 4, + uuid: 'sandbox-1', + sessionId: 17, + generation: 1, + provider: 'lifecycle_kubernetes', + status: 'ready', + capabilitySnapshot: {}, + suspendedAt: null, + endedAt: null, + error: null, + createdAt: '2026-04-24T12:00:00.000Z', + updatedAt: '2026-04-24T12:00:00.000Z', + }; + const exposure = { + id: 5, + uuid: 'exposure-1', + sandboxId: 4, + kind: 'editor', + status: 'ready', + targetPort: null, + url: '/api/agent-session/workspace-editor/session-1/', + metadata: {}, + lastVerifiedAt: '2026-04-24T12:00:00.000Z', + endedAt: null, + createdAt: '2026-04-24T12:00:00.000Z', + updatedAt: '2026-04-24T12:00:00.000Z', + }; + const defaultThread = { id: 9, uuid: 'thread-1', sessionId: 17 }; + + const sessionQuery = buildPagedSessionQuery([session], 101); + const sourceQuery = { whereIn: jest.fn().mockResolvedValue([source]) }; + const sandboxQuery = buildOrderedQuery([sandbox], 2); + const defaultThreadQuery = { whereIn: jest.fn().mockResolvedValue([defaultThread]) }; + const fallbackThreadQuery = buildOrderedQuery([], 1); + const exposureQuery = buildOrderedQuery([exposure], 1); + mockSessionQuery.mockReturnValueOnce(sessionQuery); + mockSourceQuery.mockReturnValueOnce(sourceQuery); + mockSandboxQuery.mockReturnValueOnce(sandboxQuery); + mockThreadQuery.mockReturnValueOnce(defaultThreadQuery).mockReturnValueOnce(fallbackThreadQuery); + mockSandboxExposureQuery.mockReturnValueOnce(exposureQuery); + + const result = await AgentSessionReadService.listOwnedSessionRecords('sample-user', { + page: 2, + limit: 1000, + }); + + expect(sessionQuery.page).toHaveBeenCalledWith(1, 100); + expect(result.metadata.pagination).toEqual({ + current: 2, + total: 2, + items: 101, + limit: 100, + }); + expect(result.records).toHaveLength(1); + expect(result.records[0].session.defaultThreadId).toBe('thread-1'); + expect(result.records[0].source.id).toBe('source-1'); + expect(result.records[0].sandbox.exposures).toEqual([ + expect.objectContaining({ + id: 'exposure-1', + kind: 'editor', + }), + ]); + expect(sourceQuery.whereIn).toHaveBeenCalledWith('sessionId', [17]); + expect(sandboxQuery.whereIn).toHaveBeenCalledWith('sessionId', [17]); + expect(exposureQuery.whereIn).toHaveBeenCalledWith('sandboxId', [4]); + }); +}); diff --git a/src/server/services/agent/__tests__/SourceService.test.ts b/src/server/services/agent/__tests__/SourceService.test.ts new file mode 100644 index 00000000..283d0fe5 --- /dev/null +++ b/src/server/services/agent/__tests__/SourceService.test.ts @@ -0,0 +1,140 @@ +/** + * Copyright 2026 GoodRx, Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +jest.mock('server/models/AgentSource', () => ({ + __esModule: true, + default: { + query: jest.fn(), + }, +})); + +jest.mock('server/lib/dependencies', () => ({})); + +import AgentSource from 'server/models/AgentSource'; +import AgentSourceService from '../SourceService'; + +const mockSourceQuery = AgentSource.query as jest.Mock; + +describe('AgentSourceService', () => { + beforeEach(() => { + jest.clearAllMocks(); + }); + + it('creates a source row for new session ownership', async () => { + const insertAndFetch = jest.fn().mockResolvedValue({ id: 1 }); + mockSourceQuery.mockReturnValueOnce({ insertAndFetch }); + + await AgentSourceService.createSessionSource({ + id: 3, + buildUuid: 'build-1', + buildKind: 'environment', + sessionKind: 'environment', + status: 'starting', + workspaceStatus: 'provisioning', + workspaceRepos: [{ repo: 'example-org/example-repo', mountPath: '/workspace/example-repo', primary: true }], + selectedServices: [], + updatedAt: '2026-04-24T12:00:00.000Z', + endedAt: null, + } as Parameters[0]); + + expect(insertAndFetch).toHaveBeenCalledWith( + expect.objectContaining({ + sessionId: 3, + adapter: 'lifecycle_environment', + status: 'ready', + input: { + buildUuid: 'build-1', + buildKind: 'environment', + sessionKind: 'environment', + }, + sandboxRequirements: expect.objectContaining({ + filesystem: 'persistent', + suspendMode: 'none', + }), + }) + ); + }); + + it('records source cleanup when sessions end', async () => { + const existingSource = { + id: 7, + status: 'ready', + cleanedUpAt: null, + error: null, + }; + const patchAndFetchById = jest.fn().mockResolvedValue({ + ...existingSource, + status: 'cleaned_up', + cleanedUpAt: '2026-04-24T12:00:00.000Z', + }); + + mockSourceQuery + .mockReturnValueOnce({ + findOne: jest.fn().mockResolvedValue(existingSource), + }) + .mockReturnValueOnce({ + patchAndFetchById, + }); + + await AgentSourceService.recordSessionState({ + id: 3, + status: 'ended', + workspaceStatus: 'ended', + endedAt: '2026-04-24T12:00:00.000Z', + updatedAt: '2026-04-24T12:00:00.000Z', + } as Parameters[0]); + + expect(patchAndFetchById).toHaveBeenCalledWith(7, { + status: 'cleaned_up', + cleanedUpAt: '2026-04-24T12:00:00.000Z', + }); + }); + + it('records source failure only for terminal session errors', async () => { + const existingSource = { + id: 8, + status: 'ready', + cleanedUpAt: null, + error: null, + }; + const patchAndFetchById = jest.fn().mockResolvedValue({ + ...existingSource, + status: 'failed', + error: { message: 'Source failed' }, + }); + + mockSourceQuery + .mockReturnValueOnce({ + findOne: jest.fn().mockResolvedValue(existingSource), + }) + .mockReturnValueOnce({ + patchAndFetchById, + }); + + await AgentSourceService.recordSessionState({ + id: 4, + status: 'error', + workspaceStatus: 'failed', + endedAt: null, + updatedAt: '2026-04-24T12:00:00.000Z', + } as Parameters[0]); + + expect(patchAndFetchById).toHaveBeenCalledWith(8, { + status: 'failed', + error: { message: 'Source failed' }, + }); + }); +}); diff --git a/src/server/services/agent/__tests__/sandboxExecSafety.test.ts b/src/server/services/agent/__tests__/sandboxExecSafety.test.ts index f5fe2b6d..14200660 100644 --- a/src/server/services/agent/__tests__/sandboxExecSafety.test.ts +++ b/src/server/services/agent/__tests__/sandboxExecSafety.test.ts @@ -14,7 +14,7 @@ * limitations under the License. */ -import { isReadOnlyWorkspaceCommand } from '../sandboxExecSafety'; +import { getUnsafeWorkspaceMutationReason, isReadOnlyWorkspaceCommand } from '../sandboxExecSafety'; describe('isReadOnlyWorkspaceCommand', () => { it('allows simple read-only git inspection commands', () => { @@ -39,4 +39,31 @@ describe('isReadOnlyWorkspaceCommand', () => { expect(isReadOnlyWorkspaceCommand('git status; git diff')).toBe(false); expect(isReadOnlyWorkspaceCommand('git status $(whoami)')).toBe(false); }); + + it('rejects output redirection except dev-null inspection noise', () => { + expect(isReadOnlyWorkspaceCommand('cat package.json > package-copy.json')).toBe(false); + expect(isReadOnlyWorkspaceCommand('find /workspace -type f 2>/dev/null | head -20')).toBe(true); + }); +}); + +describe('getUnsafeWorkspaceMutationReason', () => { + it('rejects broad node kill commands that can terminate the workspace gateway', () => { + expect(getUnsafeWorkspaceMutationReason('kill -9 $(pidof node)')).toContain('workspace gateway'); + expect(getUnsafeWorkspaceMutationReason('pkill -f node')).toContain('workspace gateway'); + expect(getUnsafeWorkspaceMutationReason("ps aux | grep node | awk '{print $2}' | xargs kill -9")).toContain( + 'workspace gateway' + ); + }); + + it('allows targeted process management commands', () => { + expect(getUnsafeWorkspaceMutationReason('kill 4242')).toBeNull(); + expect(getUnsafeWorkspaceMutationReason('lsof -ti tcp:3000 | xargs kill -9')).toBeNull(); + }); +}); + +describe('workspace mutation command safety', () => { + it('allows GitHub CLI commands and git pushes through the approved mutation tool', () => { + expect(getUnsafeWorkspaceMutationReason('gh repo create sample --private')).toBeNull(); + expect(getUnsafeWorkspaceMutationReason('git push -u origin main')).toBeNull(); + }); }); diff --git a/src/server/services/agent/__tests__/sandboxToolCatalog.test.ts b/src/server/services/agent/__tests__/sandboxToolCatalog.test.ts index 428804b4..830732d9 100644 --- a/src/server/services/agent/__tests__/sandboxToolCatalog.test.ts +++ b/src/server/services/agent/__tests__/sandboxToolCatalog.test.ts @@ -71,7 +71,7 @@ describe('sandboxToolCatalog', () => { ).toEqual([ '- inspect files, services, and git state: workspace.read_file, workspace.glob, workspace.grep, workspace.exec, session.get_workspace_state, session.list_ports, session.list_processes, session.get_service_status, git.status, git.diff', '- change workspace files directly: workspace.write_file, workspace.edit_file', - '- run mutating or networked shell commands: workspace.exec_mutation', + '- run mutating or networked shell commands that are not direct file edits: workspace.exec_mutation', '- manage git changes: git.add, git.commit, git.branch', '- discover and learn equipped skills: skills.list, skills.learn', '- do not claim a tool is unavailable unless it is not equipped here or a real tool call fails', diff --git a/src/server/services/agent/__tests__/streamState.test.ts b/src/server/services/agent/__tests__/streamChunks.test.ts similarity index 76% rename from src/server/services/agent/__tests__/streamState.test.ts rename to src/server/services/agent/__tests__/streamChunks.test.ts index 9cbe1dcf..27eeca40 100644 --- a/src/server/services/agent/__tests__/streamState.test.ts +++ b/src/server/services/agent/__tests__/streamChunks.test.ts @@ -14,9 +14,9 @@ * limitations under the License. */ -import { sanitizeAgentRunStreamChunks, sanitizeAgentRunStreamState } from '../streamState'; +import { sanitizeAgentRunStreamChunks } from '../streamChunks'; -describe('agent stream replay sanitization', () => { +describe('agent stream chunk sanitization', () => { it('removes duplicate fileChanges from tool-output chunks when canonical file-change chunks exist', () => { const chunks = sanitizeAgentRunStreamChunks([ { @@ -64,24 +64,4 @@ describe('agent stream replay sanitization', () => { expect(text).not.toContain('fileChanges'); expect(text).toContain('"path": "file.ts"'); }); - - it('drops redundant top-level finishReason when the finish chunk already records it', () => { - const streamState = sanitizeAgentRunStreamState({ - finishReason: 'stop', - chunks: [ - { - type: 'finish', - finishReason: 'stop', - }, - ], - }); - - expect(streamState.finishReason).toBeUndefined(); - expect(streamState.chunks).toEqual([ - { - type: 'finish', - finishReason: 'stop', - }, - ]); - }); }); diff --git a/src/server/services/agent/canonicalMessages.ts b/src/server/services/agent/canonicalMessages.ts new file mode 100644 index 00000000..2011e9da --- /dev/null +++ b/src/server/services/agent/canonicalMessages.ts @@ -0,0 +1,231 @@ +/** + * Copyright 2026 GoodRx, Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +import { v4 as uuid } from 'uuid'; +import type { AgentUIMessage } from './types'; + +export type CanonicalAgentMessagePart = + | { type: 'text'; text: string } + | { type: 'reasoning'; text: string } + | { type: 'file_ref'; path?: string | null; url?: string | null; mediaType?: string | null; title?: string | null } + | { type: 'source_ref'; url?: string | null; title?: string | null; sourceType?: string | null }; + +export type CanonicalAgentInputMessage = { + id?: string; + clientMessageId?: string | null; + role: 'user' | 'assistant' | 'system'; + parts: CanonicalAgentMessagePart[]; +}; + +export type CanonicalAgentRunMessageInput = { + clientMessageId?: string | null; + parts: CanonicalAgentMessagePart[]; +}; + +export type CanonicalAgentMessage = { + id: string; + clientMessageId: string | null; + threadId: string; + runId: string | null; + role: 'user' | 'assistant'; + parts: CanonicalAgentMessagePart[]; + createdAt: string | null; +}; + +export type AgentRunRuntimeOptions = { + maxIterations?: number; +}; + +function normalizeText(value: unknown): string | null { + return typeof value === 'string' && value.trim() ? value : null; +} + +export function normalizeCanonicalAgentMessagePart(value: unknown): CanonicalAgentMessagePart | null { + if (!value || typeof value !== 'object') { + return null; + } + + const part = value as Record; + switch (part.type) { + case 'text': { + const text = normalizeText(part.text); + return text ? { type: 'text', text } : null; + } + case 'reasoning': { + const text = normalizeText(part.text); + return text ? { type: 'reasoning', text } : null; + } + case 'file_ref': { + const path = normalizeText(part.path); + const url = normalizeText(part.url); + if (!path && !url) { + return null; + } + + return { + type: 'file_ref', + path, + url, + mediaType: normalizeText(part.mediaType), + title: normalizeText(part.title), + }; + } + case 'source_ref': { + const url = normalizeText(part.url); + const title = normalizeText(part.title); + if (!url && !title) { + return null; + } + + return { + type: 'source_ref', + url, + title, + sourceType: normalizeText(part.sourceType), + }; + } + default: + return null; + } +} + +export function isCanonicalAgentMessagePart(value: unknown): value is CanonicalAgentMessagePart { + return normalizeCanonicalAgentMessagePart(value) !== null; +} + +export function normalizeCanonicalAgentMessageParts(value: unknown): CanonicalAgentMessagePart[] { + if (!Array.isArray(value)) { + return []; + } + + const parts: CanonicalAgentMessagePart[] = []; + for (const part of value) { + pushPart(parts, normalizeCanonicalAgentMessagePart(part)); + } + + return parts; +} + +function pushPart(parts: CanonicalAgentMessagePart[], part: CanonicalAgentMessagePart | null): void { + if (part) { + parts.push(part); + } +} + +export function getCanonicalPartsFromUiMessage(message: AgentUIMessage): CanonicalAgentMessagePart[] { + const parts: CanonicalAgentMessagePart[] = []; + + for (const rawPart of message.parts || []) { + if (!rawPart || typeof rawPart !== 'object') { + continue; + } + + const part = rawPart as Record; + const partType = typeof part.type === 'string' ? part.type : ''; + + if (partType === 'text') { + pushPart( + parts, + (() => { + const text = normalizeText(part.text); + return text ? { type: 'text', text } : null; + })() + ); + continue; + } + + if (partType === 'reasoning') { + pushPart( + parts, + (() => { + const text = normalizeText(part.text); + return text ? { type: 'reasoning', text } : null; + })() + ); + continue; + } + + if (partType === 'file') { + pushPart( + parts, + normalizeCanonicalAgentMessagePart({ + type: 'file_ref', + path: normalizeText(part.filename) || normalizeText(part.path), + url: normalizeText(part.url), + mediaType: normalizeText(part.mediaType), + title: normalizeText(part.filename) || normalizeText(part.title), + }) + ); + continue; + } + + if (partType === 'source-url' || partType === 'source-document') { + pushPart( + parts, + normalizeCanonicalAgentMessagePart({ + type: 'source_ref', + url: normalizeText(part.url), + title: normalizeText(part.title), + sourceType: partType === 'source-document' ? 'document' : 'url', + }) + ); + } + } + + return parts; +} + +export function toUiMessageFromCanonicalInput( + message: CanonicalAgentInputMessage, + metadata?: Record +): AgentUIMessage { + const parts = [] as AgentUIMessage['parts']; + + for (const part of message.parts) { + if (part.type === 'text') { + parts.push({ type: 'text', text: part.text } as AgentUIMessage['parts'][number]); + continue; + } + + if (part.type === 'reasoning') { + parts.push({ type: 'reasoning', text: part.text } as AgentUIMessage['parts'][number]); + continue; + } + + if (part.type === 'file_ref') { + parts.push({ + type: 'file', + ...(part.path ? { path: part.path, filename: part.title || part.path } : {}), + ...(part.url ? { url: part.url } : {}), + ...(part.mediaType ? { mediaType: part.mediaType } : {}), + } as AgentUIMessage['parts'][number]); + continue; + } + + parts.push({ + type: part.sourceType === 'document' ? 'source-document' : 'source-url', + ...(part.url ? { url: part.url } : {}), + ...(part.title ? { title: part.title } : {}), + } as AgentUIMessage['parts'][number]); + } + + return { + id: typeof message.id === 'string' && message.id.trim() ? message.id : uuid(), + role: message.role, + parts, + metadata: metadata || {}, + } as AgentUIMessage; +} diff --git a/src/server/services/agent/fileChanges.ts b/src/server/services/agent/fileChanges.ts index f84ec410..237fd930 100644 --- a/src/server/services/agent/fileChanges.ts +++ b/src/server/services/agent/fileChanges.ts @@ -17,8 +17,9 @@ import { createHash } from 'node:crypto'; import { type DynamicToolUIPart, type ToolUIPart, type UITools } from 'ai'; import type { AgentFileChangeArtifact, AgentFileChangeData, AgentFileChangeStage, AgentUIMessage } from './types'; +import { DEFAULT_AGENT_SESSION_FILE_CHANGE_PREVIEW_CHARS } from 'server/lib/agentSession/runtimeConfig'; -const MAX_FILE_CHANGE_PREVIEW_CHARS = 4000; +const MAX_FILE_CHANGE_PREVIEW_CHARS = DEFAULT_AGENT_SESSION_FILE_CHANGE_PREVIEW_CHARS; type ToolLikePart = ToolUIPart | DynamicToolUIPart; @@ -74,14 +75,12 @@ function trimWorkspacePrefix(path: string): string { return path.replace(/^\/workspace\//, '').replace(/^\.\//, ''); } -function trimPreview(value: string | null | undefined): string | null { +function trimPreview(value: string | null | undefined, maxChars = MAX_FILE_CHANGE_PREVIEW_CHARS): string | null { if (!value) { return null; } - return value.length > MAX_FILE_CHANGE_PREVIEW_CHARS - ? `${value.slice(0, MAX_FILE_CHANGE_PREVIEW_CHARS)}\n\n[truncated]` - : value; + return value.length > maxChars ? `${value.slice(0, maxChars)}\n\n[truncated]` : value; } function countChangedLines(value: string): number { @@ -133,11 +132,13 @@ function mapArtifactToData({ toolCallId, sourceTool, stage, + previewChars, }: { artifact: AgentFileChangeArtifact; toolCallId: string; sourceTool: string; stage: AgentFileChangeStage; + previewChars?: number; }): AgentFileChangeData { return { ...artifact, @@ -147,8 +148,8 @@ function mapArtifactToData({ displayPath: trimWorkspacePrefix(artifact.path), stage, unifiedDiff: artifact.unifiedDiff ?? null, - beforeTextPreview: trimPreview(artifact.beforeTextPreview), - afterTextPreview: trimPreview(artifact.afterTextPreview), + beforeTextPreview: trimPreview(artifact.beforeTextPreview, previewChars), + afterTextPreview: trimPreview(artifact.afterTextPreview, previewChars), summary: artifact.summary ?? summarizeChange(artifact.kind, artifact.path), }; } @@ -232,10 +233,12 @@ export function buildProposedFileChanges({ toolCallId, sourceTool, input, + previewChars, }: { toolCallId: string; sourceTool: string; input: Record; + previewChars?: number; }): AgentFileChangeData[] { const toolKey = normalizeToolKey(sourceTool); @@ -268,6 +271,7 @@ export function buildProposedFileChanges({ toolCallId, sourceTool, stage: 'awaiting-approval', + previewChars, }), ]; } @@ -301,6 +305,7 @@ export function buildProposedFileChanges({ toolCallId, sourceTool, stage: 'awaiting-approval', + previewChars, }), ]; } @@ -314,12 +319,14 @@ export function buildResultFileChanges({ input, result, failed, + previewChars, }: { toolCallId: string; sourceTool: string; input: Record; result: unknown; failed: boolean; + previewChars?: number; }): AgentFileChangeData[] { const payload = unwrapToolPayload(result); const artifacts = @@ -336,6 +343,7 @@ export function buildResultFileChanges({ toolCallId, sourceTool, stage: failed ? 'failed' : 'applied', + previewChars, }) ); } @@ -348,6 +356,7 @@ export function buildResultFileChanges({ toolCallId, sourceTool, input, + previewChars, }).map((change) => ({ ...change, stage: 'failed', diff --git a/src/server/services/agent/payloadLimits.ts b/src/server/services/agent/payloadLimits.ts new file mode 100644 index 00000000..328fe379 --- /dev/null +++ b/src/server/services/agent/payloadLimits.ts @@ -0,0 +1,69 @@ +/** + * Copyright 2026 GoodRx, Inc. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +import { + DEFAULT_AGENT_SESSION_MAX_DURABLE_PAYLOAD_BYTES, + DEFAULT_AGENT_SESSION_PAYLOAD_PREVIEW_BYTES, +} from 'server/lib/agentSession/runtimeConfig'; + +export const MAX_AGENT_DURABLE_PAYLOAD_BYTES = DEFAULT_AGENT_SESSION_MAX_DURABLE_PAYLOAD_BYTES; +const PREVIEW_BYTES = DEFAULT_AGENT_SESSION_PAYLOAD_PREVIEW_BYTES; + +export interface DurablePayloadLimits { + maxDurablePayloadBytes?: number; + payloadPreviewBytes?: number; +} + +function jsonByteLength(value: unknown): number { + return Buffer.byteLength(JSON.stringify(value) ?? 'null', 'utf8'); +} + +export function limitDurablePayloadValue(value: unknown, limits: DurablePayloadLimits = {}): unknown { + const maxBytes = limits.maxDurablePayloadBytes ?? MAX_AGENT_DURABLE_PAYLOAD_BYTES; + const previewBytes = limits.payloadPreviewBytes ?? PREVIEW_BYTES; + const serialized = JSON.stringify(value) ?? 'null'; + const bytes = Buffer.byteLength(serialized, 'utf8'); + if (bytes <= maxBytes) { + return value; + } + + return { + truncated: true, + originalJsonBytes: bytes, + preview: serialized.slice(0, previewBytes), + }; +} + +export function limitDurablePayloadRecord( + payload: Record, + limits: DurablePayloadLimits = {} +): Record { + const maxBytes = limits.maxDurablePayloadBytes ?? MAX_AGENT_DURABLE_PAYLOAD_BYTES; + if (jsonByteLength(payload) <= maxBytes) { + return payload; + } + + const limited: Record = {}; + for (const [key, value] of Object.entries(payload)) { + limited[key] = limitDurablePayloadValue(value, limits); + } + + if (jsonByteLength(limited) <= maxBytes) { + return limited; + } + + return limitDurablePayloadValue(payload, limits) as Record; +} diff --git a/src/server/services/agent/providerConfig.ts b/src/server/services/agent/providerConfig.ts index 5fb40792..d3dcc38e 100644 --- a/src/server/services/agent/providerConfig.ts +++ b/src/server/services/agent/providerConfig.ts @@ -15,8 +15,6 @@ */ export const SUPPORTED_AGENT_PROVIDER_NAMES = ['anthropic', 'openai', 'gemini', 'google'] as const; -export const AGENT_API_KEY_HEADER = 'x-lifecycle-agent-api-key'; -export const AGENT_API_KEY_PROVIDER_HEADER = 'x-lifecycle-agent-api-key-provider'; export type SupportedAgentProviderName = (typeof SUPPORTED_AGENT_PROVIDER_NAMES)[number]; export const STORED_AGENT_PROVIDER_NAMES = ['anthropic', 'openai', 'gemini'] as const; diff --git a/src/server/services/agent/sandboxExecSafety.ts b/src/server/services/agent/sandboxExecSafety.ts index 1e2b2e74..b08c57ea 100644 --- a/src/server/services/agent/sandboxExecSafety.ts +++ b/src/server/services/agent/sandboxExecSafety.ts @@ -49,6 +49,44 @@ const READ_ONLY_SEGMENT_PATTERNS: RegExp[] = [ const BLOCKED_SHELL_OPERATORS = /&&|\|\||;|`|\$\(/; const DEV_NULL_REDIRECTION = /\s+\d?>\s*\/dev\/null/g; +const OUTPUT_REDIRECTION = /(^|[^0-9])>>?\s*(?!&)|[0-9]>>?\s*(?!\/dev\/null)/; +const UNSAFE_WORKSPACE_MUTATION_PATTERNS: Array<{ + pattern: RegExp; + reason: string; +}> = [ + { + pattern: /\bkill\b[^\n]*\$\(\s*pidof\s+node\s*\)/i, + reason: + 'This command targets every node process and can terminate the workspace gateway. Inspect the process list and stop only the specific app process instead.', + }, + { + pattern: /\bkill\b[^\n]*\$\(\s*pgrep(?:\s+-f)?[^\n]*\b(?:node|workspace-gateway)\b[^\n]*\)/i, + reason: + 'This command targets generic node or workspace-gateway processes and can terminate the workspace gateway. Stop only the specific app process instead.', + }, + { + pattern: /\b(?:pidof|pgrep(?:\s+-f)?)\b[^\n]*\b(?:node|workspace-gateway)\b[^\n]*\bxargs\s+kill\b/i, + reason: + 'This command kills PIDs discovered from generic node or workspace-gateway process searches and can terminate the workspace gateway.', + }, + { + pattern: /\bps\b[^\n]*\bgrep\s+(?:-w\s+)?(?:node|workspace-gateway)\b[^\n]*\bxargs\s+kill\b/i, + reason: + 'This command kills PIDs discovered from generic node or workspace-gateway process searches and can terminate the workspace gateway.', + }, + { + pattern: /\bpkill\b[^\n]*\b(?:node|workspace-gateway)\b/i, + reason: 'This command kills generic node or workspace-gateway processes and can terminate the workspace gateway.', + }, + { + pattern: /\bkillall\b[^\n]*\b(?:node|workspace-gateway)\b/i, + reason: 'This command kills generic node or workspace-gateway processes and can terminate the workspace gateway.', + }, + { + pattern: /\bkill\b[^\n]*(?:\$\$|\$PPID|\$BASHPID|\b1\b)/, + reason: 'This command targets the current shell, its parent, or PID 1 and can terminate the workspace gateway.', + }, +]; function normalizeCommandSegment(segment: string): string { return segment.replace(DEV_NULL_REDIRECTION, '').trim(); @@ -73,6 +111,10 @@ export function isReadOnlyWorkspaceCommand(command: string): boolean { return false; } + if (OUTPUT_REDIRECTION.test(normalized.replace(DEV_NULL_REDIRECTION, ''))) { + return false; + } + const segments = normalized .split('|') .map((segment) => segment.trim()) @@ -84,3 +126,25 @@ export function isReadOnlyWorkspaceCommand(command: string): boolean { return segments.every(isReadOnlyCommandSegment); } + +export function getUnsafeWorkspaceMutationReason(command: string): string | null { + const normalized = command.trim(); + if (!normalized) { + return null; + } + + for (const entry of UNSAFE_WORKSPACE_MUTATION_PATTERNS) { + if (entry.pattern.test(normalized)) { + return entry.reason; + } + } + + return null; +} + +export function assertSafeWorkspaceMutationCommand(command: string): void { + const reason = getUnsafeWorkspaceMutationReason(command); + if (reason) { + throw new Error(reason); + } +} diff --git a/src/server/services/agent/sandboxToolCatalog.ts b/src/server/services/agent/sandboxToolCatalog.ts index e622a2c7..f23460a0 100644 --- a/src/server/services/agent/sandboxToolCatalog.ts +++ b/src/server/services/agent/sandboxToolCatalog.ts @@ -231,7 +231,7 @@ const PROMPT_CATEGORY_COPY: Record | null): Record { - if (!streamState || typeof streamState !== 'object') { - return {}; - } - - const nextState = cloneValue(streamState); - const rawChunks = Array.isArray(nextState.chunks) ? nextState.chunks : []; - const sanitizedChunks = sanitizeAgentRunStreamChunks(rawChunks as AgentUiMessageChunk[]); - - if (sanitizedChunks.length > 0) { - nextState.chunks = sanitizedChunks; - } - - const finishChunk = sanitizedChunks.find((chunk) => isRecord(chunk) && chunk.type === 'finish'); - if ( - typeof nextState.finishReason === 'string' && - isRecord(finishChunk) && - finishChunk.finishReason === nextState.finishReason - ) { - delete nextState.finishReason; - } - - return nextState; -} diff --git a/src/server/services/agent/toolKeys.ts b/src/server/services/agent/toolKeys.ts index aab553bd..8e2ca407 100644 --- a/src/server/services/agent/toolKeys.ts +++ b/src/server/services/agent/toolKeys.ts @@ -18,6 +18,9 @@ export const SESSION_WORKSPACE_SERVER_SLUG = 'sandbox'; export const SESSION_WORKSPACE_SERVER_NAME = 'Session Workspace'; export const SESSION_WORKSPACE_READONLY_TOOL_NAME = 'workspace.exec'; export const SESSION_WORKSPACE_MUTATION_TOOL_NAME = 'workspace.exec_mutation'; +export const LIFECYCLE_BUILTIN_SERVER_SLUG = 'lifecycle'; +export const LIFECYCLE_BUILTIN_SERVER_NAME = 'Lifecycle'; +export const CHAT_PUBLISH_HTTP_TOOL_NAME = 'publish_http'; export function buildAgentToolKey(serverSlug: string, toolName: string): string { return `mcp__${serverSlug}__${toolName}`.replace(/[^a-zA-Z0-9_]/g, '_'); @@ -27,6 +30,7 @@ export function buildWorkspaceReadonlyExecDescription(serverName: string): strin return ( `Run a read-only workspace inspection command through ${serverName}. ` + 'Use this for safe file, git, and directory inspection only. ' + + 'Do not chain commands with &&, ||, or ;. Run separate inspection commands instead. ' + 'Examples: git remote -v, git status --short --branch, ls -la, find . -name "*.ts", rg pattern src.' ); } @@ -34,7 +38,8 @@ export function buildWorkspaceReadonlyExecDescription(serverName: string): strin export function buildWorkspaceMutationExecDescription(serverName: string): string { return ( `Run a mutating or networked workspace command through ${serverName}. ` + - 'Use this for installs, starting processes, git pushes or commits, gh commands, and other state-changing operations. ' + + 'Use this for installs, starting processes, GitHub CLI operations, git pushes, local git commits, and other state-changing operations that are not direct file-content edits. ' + + 'When creating or changing file contents, use workspace.write_file or workspace.edit_file so the file changes can be reviewed. ' + 'This path is intended for commands that require approval.' ); } diff --git a/src/server/services/agent/types.ts b/src/server/services/agent/types.ts index 6e9f37a3..eb806731 100644 --- a/src/server/services/agent/types.ts +++ b/src/server/services/agent/types.ts @@ -31,6 +31,7 @@ export type AgentCapabilityKey = (typeof AGENT_CAPABILITY_KEYS)[number]; export type AgentApprovalMode = 'allow' | 'require_approval' | 'deny'; export type AgentRunStatus = | 'queued' + | 'starting' | 'running' | 'waiting_for_approval' | 'waiting_for_input' @@ -103,6 +104,7 @@ export interface AgentApprovalPolicy { } export interface AgentUIMessageMetadata { + clientMessageId?: string; sessionId?: string; threadId?: string; runId?: string; diff --git a/src/server/services/agentPrewarm.ts b/src/server/services/agentPrewarm.ts index 41eb0268..15b78f73 100644 --- a/src/server/services/agentPrewarm.ts +++ b/src/server/services/agentPrewarm.ts @@ -412,7 +412,13 @@ export default class AgentPrewarmService extends BaseService { } as unknown as Partial); try { - await createAgentPvc(plan.namespace, pvcName, '10Gi', plan.buildUuid).catch((error: unknown) => { + await createAgentPvc( + plan.namespace, + pvcName, + runtimeConfig.workspaceStorage.defaultSize, + plan.buildUuid, + runtimeConfig.workspaceStorage.accessMode + ).catch((error: unknown) => { const httpError = error as k8s.HttpError; if (httpError?.statusCode === 409 || httpError?.response?.statusCode === 409) { return null; diff --git a/src/server/services/agentSandboxSession.ts b/src/server/services/agentSandboxSession.ts index fefef87a..7b4ac3a4 100644 --- a/src/server/services/agentSandboxSession.ts +++ b/src/server/services/agentSandboxSession.ts @@ -24,6 +24,7 @@ import { mergeAgentSessionResources, type ResolvedAgentSessionReadinessConfig, type ResolvedAgentSessionResources, + type ResolvedAgentSessionWorkspaceStorageIntent, } from 'server/lib/agentSession/runtimeConfig'; import { BuildEnvironmentVariables } from 'server/lib/buildEnvVariables'; import { getLogger } from 'server/lib/logger'; @@ -133,8 +134,6 @@ export interface LaunchSandboxSessionOptions { userId: string; userIdentity?: RequestUserIdentity; githubToken?: string | null; - requestApiKey?: string | null; - requestApiKeyProvider?: string | null; baseBuildUuid: string; services?: RequestedSandboxServices; model?: string; @@ -145,6 +144,8 @@ export interface LaunchSandboxSessionOptions { keepAttachedServicesOnSessionNode?: boolean; readiness: ResolvedAgentSessionReadinessConfig; resources: ResolvedAgentSessionResources; + workspaceStorage?: ResolvedAgentSessionWorkspaceStorageIntent; + redisTtlSeconds?: number; onProgress?: (stage: SandboxLaunchStage, message: string) => Promise | void; } @@ -262,8 +263,6 @@ export default class AgentSandboxSessionService extends BaseService { buildUuid: sandboxBuild.uuid, buildKind: BuildKind.SANDBOX, githubToken: opts.githubToken, - requestApiKey: opts.requestApiKey, - requestApiKeyProvider: opts.requestApiKeyProvider, services: selectedSandboxServices.map(({ selectedService, sandboxDeploy }) => ({ name: selectedService.name, deployId: sandboxDeploy.id, @@ -287,6 +286,8 @@ export default class AgentSandboxSessionService extends BaseService { selectedServices.map((service) => service.devConfig.agentSession?.readiness) ), resources: mergeAgentSessionResources(opts.resources, lifecycleConfig.environment?.agentSession?.resources), + workspaceStorage: opts.workspaceStorage, + redisTtlSeconds: opts.redisTtlSeconds, userIdentity: opts.userIdentity, }); const sessionOpenMs = elapsedMs(openSessionStartedAt); diff --git a/src/server/services/agentSession.ts b/src/server/services/agentSession.ts index 051024f8..29d53f5b 100644 --- a/src/server/services/agentSession.ts +++ b/src/server/services/agentSession.ts @@ -19,6 +19,7 @@ import * as k8s from '@kubernetes/client-node'; import { Writable } from 'stream'; import { v4 as uuid } from 'uuid'; import type Database from 'server/database'; +import AgentRun from 'server/models/AgentRun'; import AgentSession from 'server/models/AgentSession'; import Build from 'server/models/Build'; import Configuration from 'server/models/Configuration'; @@ -39,18 +40,26 @@ import { } from 'server/lib/agentSession/editorServiceFactory'; import { ensureAgentSessionServiceAccount } from 'server/lib/agentSession/serviceAccountFactory'; import { isGvisorAvailable } from 'server/lib/agentSession/gvisorCheck'; +import { createOrUpdateChatPreview } from 'server/lib/agentSession/chatPreviewFactory'; import { DevModeManager } from 'server/lib/agentSession/devModeManager'; import type { DevModeResourceSnapshot } from 'server/lib/agentSession/devModeManager'; +import { createOrUpdateNamespace, deleteNamespace } from 'server/lib/kubernetes'; import { buildAgentNetworkPolicy } from 'server/lib/kubernetes/networkPolicyFactory'; import { DevConfig } from 'server/models/yaml/YamlService'; import RedisClient from 'server/lib/redisClient'; import { extractContextForQueue, getLogger } from 'server/lib/logger'; -import { BuildKind, FeatureFlags } from 'shared/constants'; +import { AgentChatStatus, AgentSessionKind, AgentWorkspaceStatus, BuildKind, FeatureFlags } from 'shared/constants'; import type { RequestUserIdentity } from 'server/lib/get-user'; import { + DEFAULT_AGENT_SESSION_REDIS_TTL_SECONDS, + DEFAULT_AGENT_SESSION_WORKSPACE_STORAGE_ACCESS_MODE, + DEFAULT_AGENT_SESSION_WORKSPACE_STORAGE_SIZE, + resolveAgentSessionRuntimeConfig, + resolveAgentSessionWorkspaceStorageIntent, resolveKeepAttachedServicesOnSessionNode, type ResolvedAgentSessionReadinessConfig, type ResolvedAgentSessionResources, + type ResolvedAgentSessionWorkspaceStorageIntent, } from 'server/lib/agentSession/runtimeConfig'; import { cleanupForwardedAgentEnvSecrets, resolveForwardedAgentEnv } from 'server/lib/agentSession/forwardedEnv'; import { EMPTY_AGENT_SESSION_SKILL_PLAN, resolveAgentSessionSkillPlan } from 'server/lib/agentSession/skillPlan'; @@ -89,8 +98,11 @@ import { } from 'server/services/ai/mcp/sessionPod'; import AgentPrewarmService from './agentPrewarm'; import AgentSessionConfigService from './agentSessionConfig'; +import AgentChatSessionService from './agent/ChatSessionService'; import AgentPolicyService from './agent/PolicyService'; import AgentProviderRegistry from './agent/ProviderRegistry'; +import AgentSandboxService from './agent/SandboxService'; +import AgentSourceService from './agent/SourceService'; import { buildSessionWorkspacePromptLines } from './agent/sandboxToolCatalog'; import { loadAgentSessionServiceCandidates, @@ -102,10 +114,17 @@ import type { AgentSessionSkillRef } from 'server/models/yaml/YamlService'; const logger = () => getLogger(); const SESSION_REDIS_PREFIX = 'lifecycle:agent:session:'; -const SESSION_REDIS_TTL = 7200; const ACTIVE_ENVIRONMENT_SESSION_UNIQUE_INDEX = 'agent_sessions_active_environment_build_unique'; const DEV_MODE_REDEPLOY_GRAPH = '[deployable.[repository], repository, service, build.[pullRequest.[repository]]]'; const SESSION_DEPLOY_GRAPH = '[deployable, repository, service]'; +const AGENT_RUN_TERMINAL_STATUSES = ['completed', 'failed', 'cancelled']; + +export class ActiveAgentRunSuspensionError extends Error { + constructor() { + super('Cannot suspend a chat runtime while an agent run is active'); + this.name = 'ActiveAgentRunSuspensionError'; + } +} const agentNetworkPolicySetupByNamespace = new Map>(); type AgentSessionSummaryRecordBase = AgentSession & { @@ -130,11 +149,65 @@ type ActiveEnvironmentSessionSummary = { ownedByCurrentUser: boolean; }; +function resolveSessionKindFromBuildKind(buildKind: BuildKind): AgentSessionKind { + return buildKind === BuildKind.SANDBOX ? AgentSessionKind.SANDBOX : AgentSessionKind.ENVIRONMENT; +} + +function canSessionAcceptMessages( + session: Pick +): boolean { + if (session.chatStatus !== AgentChatStatus.READY) { + return false; + } + + if (session.sessionKind === AgentSessionKind.CHAT) { + return true; + } + + return session.workspaceStatus === AgentWorkspaceStatus.READY; +} + +function getSessionMessageBlockReason( + session: Pick +): string { + if (canSessionAcceptMessages(session)) { + return ''; + } + + if ( + session.sessionKind !== AgentSessionKind.CHAT && + (session.workspaceStatus === AgentWorkspaceStatus.PROVISIONING || session.status === 'starting') + ) { + return 'Wait for the session to finish starting before sending a message.'; + } + + return 'This session is no longer available for new messages.'; +} + +function warmDefaultThread(sessionUuid: string, userId: string): void { + // Default-thread creation stays best-effort so chat readiness does not + // depend on secondary DB work; ThreadService.listThreadsForSession() will + // create or retry the default thread on first access if this warm-up fails. + void (async () => { + const AgentThreadService = (await import('server/services/agent/ThreadService')).default; + await AgentThreadService.getDefaultThreadForSession(sessionUuid, userId); + })().catch((error: unknown) => { + logger().warn( + { error, sessionId: sessionUuid }, + `Session: default thread creation skipped sessionId=${sessionUuid}` + ); + }); +} + export function buildAgentSessionPodName(sessionUuid: string, buildUuid?: string | null): string { const identifier = buildUuid ?? sessionUuid.slice(0, 8); return normalizeKubernetesLabelValue(`agent-${identifier}`.toLowerCase()).replace(/[_.]/g, '-'); } +export function buildChatSessionNamespace(sessionUuid: string): string { + return normalizeKubernetesLabelValue(`chat-${sessionUuid.slice(0, 8)}`.toLowerCase()).replace(/[_.]/g, '-'); +} + export class ActiveEnvironmentSessionError extends Error { activeSession: ActiveEnvironmentSessionSummary; @@ -656,12 +729,24 @@ function isUniqueConstraintError(error: unknown, constraintName: string): boolea return knexError?.code === '23505' && knexError?.constraint === constraintName; } +function getRequestedWorkspaceStorageSize(input: unknown): string | null { + if (!input || typeof input !== 'object' || Array.isArray(input)) { + return null; + } + + const workspace = (input as { workspace?: unknown }).workspace; + if (!workspace || typeof workspace !== 'object' || Array.isArray(workspace)) { + return null; + } + + const storageSize = (workspace as { storageSize?: unknown }).storageSize; + return typeof storageSize === 'string' && storageSize.trim() ? storageSize.trim() : null; +} + export interface CreateSessionOptions { userId: string; userIdentity?: RequestUserIdentity; githubToken?: string | null; - requestApiKey?: string | null; - requestApiKeyProvider?: string | null; buildUuid?: string; buildKind?: BuildKind; services?: Array<{ @@ -690,9 +775,34 @@ export interface CreateSessionOptions { keepAttachedServicesOnSessionNode?: boolean; readiness?: ResolvedAgentSessionReadinessConfig; resources?: ResolvedAgentSessionResources; + workspaceStorage?: ResolvedAgentSessionWorkspaceStorageIntent; + redisTtlSeconds?: number; +} + +export interface CreateChatSessionOptions { + userId: string; + userIdentity?: RequestUserIdentity; + model?: string; +} + +export interface CreateChatRuntimeOptions { + sessionId: string; + userId: string; + userIdentity?: RequestUserIdentity; + githubToken?: string | null; } export default class AgentSessionService { + static canAcceptMessages(session: Pick): boolean { + return canSessionAcceptMessages(session); + } + + static getMessageBlockReason( + session: Pick + ): string { + return getSessionMessageBlockReason(session); + } + static async getSessionStartupFailure(sessionId: string): Promise { const redis = RedisClient.getInstance().getRedis(); const failure = await getAgentSessionStartupFailure(redis, sessionId); @@ -723,6 +833,8 @@ export default class AgentSessionService { .findById(session.id) .patch({ status: 'error', + chatStatus: AgentChatStatus.ERROR, + workspaceStatus: AgentWorkspaceStatus.FAILED, endedAt: new Date().toISOString(), } as unknown as Partial) .catch(() => {}); @@ -867,10 +979,299 @@ export default class AgentSessionService { }; } + static async createChatSession(opts: CreateChatSessionOptions): Promise { + return AgentChatSessionService.createChatSession(opts); + } + + static async provisionChatRuntime(opts: CreateChatRuntimeOptions): Promise { + const session = await AgentSession.query().findOne({ uuid: opts.sessionId, userId: opts.userId }); + if (!session) { + throw new Error('Session not found'); + } + + if (session.sessionKind !== AgentSessionKind.CHAT) { + throw new Error('Runtime provisioning is only supported for chat sessions'); + } + + if (session.status !== 'active') { + throw new Error('Only active chat sessions can provision a workspace runtime'); + } + + if (session.workspaceStatus === AgentWorkspaceStatus.PROVISIONING) { + throw new Error('Workspace runtime is already provisioning'); + } + + if ( + session.workspaceStatus === AgentWorkspaceStatus.READY && + session.namespace && + session.podName && + session.pvcName + ) { + return session; + } + + const runtimeConfig = await resolveAgentSessionRuntimeConfig(); + const source = await AgentSourceService.getSessionSource(session.id).catch(() => null); + const workspaceStorage = resolveAgentSessionWorkspaceStorageIntent({ + requestedSize: getRequestedWorkspaceStorageSize(source?.input), + storage: runtimeConfig.workspaceStorage, + }); + const namespace = buildChatSessionNamespace(session.uuid); + const podName = buildAgentSessionPodName(session.uuid); + const pvcName = `agent-pvc-${session.uuid.slice(0, 8)}`; + const apiKeySecretName = `agent-secret-${session.uuid.slice(0, 8)}`; + const redis = RedisClient.getInstance().getRedis(); + + if (session.namespace && session.namespace !== namespace) { + await deleteNamespace(session.namespace).catch(() => {}); + } + + await createOrUpdateNamespace({ + name: namespace, + buildUUID: session.uuid, + staticEnv: false, + ttl: true, + author: opts.userIdentity?.githubUsername || session.ownerGithubUsername || null, + }); + + const provisioningPatch = { + namespace, + podName, + pvcName, + status: 'active', + chatStatus: AgentChatStatus.READY, + workspaceStatus: AgentWorkspaceStatus.PROVISIONING, + workspaceRepos: [], + selectedServices: [], + devModeSnapshots: {}, + forwardedAgentSecretProviders: [], + skillPlan: session.skillPlan || EMPTY_AGENT_SESSION_SKILL_PLAN, + } as unknown as Partial; + await AgentSession.query().findById(session.id).patch(provisioningPatch); + const provisioningSession = { + ...session, + ...provisioningPatch, + } as AgentSession; + await AgentSandboxService.recordSessionSandboxState(provisioningSession, { workspaceStorage }); + + try { + const sessionPodMcpConfigJson = serializeSessionWorkspaceGatewayServers([]); + const [, , agentServiceAccountName, useGvisor] = await Promise.all([ + createAgentPvc(namespace, pvcName, workspaceStorage.storageSize, undefined, workspaceStorage.accessMode), + createAgentApiKeySecret( + namespace, + apiKeySecretName, + {}, + opts.githubToken, + undefined, + {}, + { + [SESSION_POD_MCP_CONFIG_SECRET_KEY]: sessionPodMcpConfigJson, + } + ), + ensureAgentSessionServiceAccount(namespace), + isGvisorAvailable(), + createSessionWorkspaceService(namespace, podName), + ensureAgentNetworkPolicy(namespace), + ]); + + await createSessionWorkspacePod({ + podName, + namespace, + pvcName, + workspaceImage: runtimeConfig.workspaceImage, + workspaceEditorImage: runtimeConfig.workspaceEditorImage, + workspaceGatewayImage: runtimeConfig.workspaceGatewayImage, + apiKeySecretName, + hasGitHubToken: Boolean(opts.githubToken), + workspacePath: SESSION_WORKSPACE_ROOT, + workspaceRepos: [], + skillPlan: session.skillPlan || EMPTY_AGENT_SESSION_SKILL_PLAN, + forwardedAgentEnv: {}, + forwardedAgentSecretRefs: [], + useGvisor, + userIdentity: opts.userIdentity, + nodeSelector: runtimeConfig.nodeSelector, + readiness: runtimeConfig.readiness, + serviceAccountName: agentServiceAccountName, + resources: runtimeConfig.resources, + }); + + await Promise.all([ + redis.setex( + `${SESSION_REDIS_PREFIX}${session.uuid}`, + runtimeConfig.cleanup.redisTtlSeconds, + JSON.stringify({ podName, namespace, status: 'active' }) + ), + AgentSession.query() + .findById(session.id) + .patch({ + status: 'active', + chatStatus: AgentChatStatus.READY, + workspaceStatus: AgentWorkspaceStatus.READY, + namespace, + podName, + pvcName, + } as unknown as Partial), + ]); + + logger().info(`Session: runtime ready sessionId=${session.uuid} namespace=${namespace} podName=${podName}`); + + const readySession = await AgentSession.query().findOne({ uuid: session.uuid }); + if (!readySession) { + throw new Error('Session not found after runtime provisioning'); + } + + await AgentSandboxService.recordSessionSandboxState(readySession, { workspaceStorage }); + return readySession; + } catch (error) { + logger().warn( + { error, sessionId: session.uuid, namespace }, + `Session: runtime provision failed sessionId=${session.uuid}` + ); + + await Promise.all([ + deleteAgentRuntimeResources(namespace, podName, apiKeySecretName).catch(() => {}), + deleteAgentPvc(namespace, pvcName).catch(() => {}), + deleteNamespace(namespace).catch(() => {}), + redis.del(`${SESSION_REDIS_PREFIX}${session.uuid}`).catch(() => {}), + ]); + + await AgentSandboxService.recordSessionSandboxState( + { + ...session, + namespace, + podName, + pvcName, + status: 'active', + workspaceStatus: AgentWorkspaceStatus.FAILED, + } as AgentSession, + { workspaceStorage } + ).catch(() => {}); + + await AgentSession.query() + .findById(session.id) + .patch({ + status: 'active', + chatStatus: AgentChatStatus.READY, + workspaceStatus: AgentWorkspaceStatus.FAILED, + namespace: null, + podName: null, + pvcName: null, + workspaceRepos: [], + selectedServices: [], + devModeSnapshots: {}, + } as unknown as Partial); + + throw error; + } + } + + static async publishChatHttpPort({ + sessionId, + userId, + port, + }: { + sessionId: string; + userId: string; + port: number; + }): Promise<{ + url: string; + host: string | null; + path: string; + port: number; + serviceName: string; + ingressName: string; + }> { + const session = await AgentSession.query().findOne({ uuid: sessionId, userId }); + if (!session) { + throw new Error('Session not found'); + } + + if (session.sessionKind !== AgentSessionKind.CHAT) { + throw new Error('HTTP publishing is only supported for chat sessions'); + } + + if (session.workspaceStatus !== AgentWorkspaceStatus.READY || !session.namespace || !session.podName) { + throw new Error('Workspace runtime is not ready yet'); + } + + const publication = await createOrUpdateChatPreview({ + sessionUuid: session.uuid, + namespace: session.namespace, + podName: session.podName, + port, + }); + + logger().info( + `Session: preview ready sessionId=${session.uuid} namespace=${session.namespace} port=${port} url=${publication.url}` + ); + + return publication; + } + + static async suspendChatRuntime({ sessionId, userId }: { sessionId: string; userId: string }): Promise { + const session = await AgentSession.query().findOne({ uuid: sessionId, userId }); + if (!session) { + throw new Error('Session not found'); + } + + if (session.sessionKind !== AgentSessionKind.CHAT) { + throw new Error('Runtime suspension is only supported for chat sessions'); + } + + if (session.status !== 'active') { + throw new Error('Only active chat sessions can be suspended'); + } + + if (session.workspaceStatus === AgentWorkspaceStatus.HIBERNATED) { + return session; + } + + const activeRun = await AgentRun.query() + .where({ sessionId: session.id }) + .whereNotIn('status', AGENT_RUN_TERMINAL_STATUSES) + .first(); + if (activeRun) { + throw new ActiveAgentRunSuspensionError(); + } + + if (session.workspaceStatus !== AgentWorkspaceStatus.READY || !session.namespace || !session.pvcName) { + throw new Error('Workspace runtime is not ready'); + } + + const redis = RedisClient.getInstance().getRedis(); + const apiKeySecretName = `agent-secret-${session.uuid.slice(0, 8)}`; + if (session.podName) { + await deleteAgentRuntimeResources(session.namespace, session.podName, apiKeySecretName); + } + + await Promise.all([ + redis.del(`${SESSION_REDIS_PREFIX}${session.uuid}`).catch(() => {}), + clearAgentSessionStartupFailure(redis, session.uuid).catch(() => {}), + ]); + + const suspendedSession = await AgentSession.query().patchAndFetchById(session.id, { + status: 'active', + chatStatus: AgentChatStatus.READY, + workspaceStatus: AgentWorkspaceStatus.HIBERNATED, + podName: null, + } as unknown as Partial); + await AgentSandboxService.recordSessionSandboxState(suspendedSession); + + logger().info(`Session: runtime suspended sessionId=${session.uuid} namespace=${session.namespace}`); + return suspendedSession; + } + + static async resumeChatRuntime(opts: CreateChatRuntimeOptions): Promise { + return this.provisionChatRuntime(opts); + } + static async createSession(opts: CreateSessionOptions) { const sessionStartedAt = Date.now(); const sessionUuid = uuid(); const buildKind = opts.buildKind || BuildKind.ENVIRONMENT; + const sessionKind = resolveSessionKindFromBuildKind(buildKind); const podName = buildAgentSessionPodName(sessionUuid, opts.buildUuid); const apiKeySecretName = `agent-secret-${sessionUuid.slice(0, 8)}`; const requestedModelId = opts.model?.trim() || undefined; @@ -883,6 +1284,12 @@ export default class AgentSessionService { let sessionPersisted = false; let session: AgentSession | null = null; let resolvedModelId = requestedModelId || 'unresolved-model'; + const workspaceStorage = opts.workspaceStorage ?? { + requestedSize: null, + storageSize: DEFAULT_AGENT_SESSION_WORKSPACE_STORAGE_SIZE, + accessMode: DEFAULT_AGENT_SESSION_WORKSPACE_STORAGE_ACCESS_MODE, + }; + const redisTtlSeconds = opts.redisTtlSeconds ?? DEFAULT_AGENT_SESSION_REDIS_TTL_SECONDS; const redis = RedisClient.getInstance().getRedis(); const templatedServices = await resolveTemplatedDevConfigEnvs(opts.buildUuid, opts.namespace, opts.services); const { @@ -914,18 +1321,14 @@ export default class AgentSessionService { }); resolvedModelId = selection.modelId; const resolvedServiceNames = (resolvedServices || []).map((service) => service.name); - const [, providerApiKeys, sessionPodServers, compatiblePrewarm, forwardedAgentEnv] = await Promise.all([ + const [, providerApiKeys, sessionPodServers, resolvedCompatiblePrewarm, forwardedAgentEnv] = await Promise.all([ AgentProviderRegistry.getRequiredStoredApiKey({ provider: selection.provider, userIdentity: providerUserIdentity, - requestApiKey: opts.requestApiKey, - requestApiKeyProvider: opts.requestApiKeyProvider, }), AgentProviderRegistry.resolveCredentialEnvMap({ repoFullName: primaryWorkspaceRepo?.repo, userIdentity: providerUserIdentity, - requestApiKey: opts.requestApiKey, - requestApiKeyProvider: opts.requestApiKeyProvider, }), primaryWorkspaceRepo?.repo ? new McpConfigService().resolveSessionPodServersForRepo( @@ -943,6 +1346,7 @@ export default class AgentSessionService { ), resolveForwardedAgentEnv(resolvedServices, opts.namespace, sessionUuid, opts.buildUuid), ]); + const compatiblePrewarm = workspaceStorage.requestedSize ? null : resolvedCompatiblePrewarm; const sessionPodMcpConfigJson = serializeSessionWorkspaceGatewayServers(sessionPodServers); const pvcName = compatiblePrewarm?.pvcName || `agent-pvc-${sessionUuid.slice(0, 8)}`; const forwardedPlainAgentEnv = Object.fromEntries( @@ -963,30 +1367,49 @@ export default class AgentSessionService { try { const keepAttachedServicesOnSessionNode = opts.keepAttachedServicesOnSessionNode !== false; - session = await AgentSession.query().insertAndFetch({ - uuid: sessionUuid, - buildUuid: opts.buildUuid || null, - buildKind, - userId: opts.userId, - ownerGithubUsername: opts.userIdentity?.githubUsername || null, - podName, - namespace: opts.namespace, - pvcName, - model: resolvedModelId, - status: 'starting', - keepAttachedServicesOnSessionNode, - devModeSnapshots, - forwardedAgentSecretProviders: forwardedAgentEnv.secretProviders, - workspaceRepos, - selectedServices, - skillPlan, - } as unknown as Partial); + session = await AgentSession.transaction(async (trx) => { + const createdSession = await AgentSession.query(trx).insertAndFetch({ + uuid: sessionUuid, + buildUuid: opts.buildUuid || null, + buildKind, + sessionKind, + userId: opts.userId, + ownerGithubUsername: opts.userIdentity?.githubUsername || null, + podName, + namespace: opts.namespace, + pvcName, + model: resolvedModelId, + defaultModel: resolvedModelId, + defaultHarness: 'lifecycle_ai_sdk', + status: 'starting', + chatStatus: AgentChatStatus.READY, + workspaceStatus: AgentWorkspaceStatus.PROVISIONING, + keepAttachedServicesOnSessionNode, + devModeSnapshots, + forwardedAgentSecretProviders: forwardedAgentEnv.secretProviders, + workspaceRepos, + selectedServices, + skillPlan, + } as unknown as Partial); + + await AgentSourceService.createSessionSource(createdSession, { trx, workspaceStorage }); + await AgentSandboxService.recordSessionSandboxState(createdSession, { trx, workspaceStorage }); + return createdSession; + }); sessionPersisted = true; const combinedInstallCommand = buildCombinedInstallCommand(resolvedServices); const infraSetupStartedAt = Date.now(); const [, , agentServiceAccountName, useGvisor] = await Promise.all([ - compatiblePrewarm ? Promise.resolve(null) : createAgentPvc(opts.namespace, pvcName, '10Gi', opts.buildUuid), + compatiblePrewarm + ? Promise.resolve(null) + : createAgentPvc( + opts.namespace, + pvcName, + workspaceStorage.storageSize, + opts.buildUuid, + workspaceStorage.accessMode + ), createAgentApiKeySecret( opts.namespace, apiKeySecretName, @@ -1118,7 +1541,7 @@ export default class AgentSessionService { enabledServices.map((service) => Deploy.query().findById(service.deployId).patch({ devMode: true, - devModeSessionId: session.id, + devModeSessionId: session!.id, }) ) ); @@ -1128,13 +1551,15 @@ export default class AgentSessionService { await Promise.all([ redis.setex( `${SESSION_REDIS_PREFIX}${sessionUuid}`, - SESSION_REDIS_TTL, + redisTtlSeconds, JSON.stringify({ podName, namespace: opts.namespace, status: 'active' }) ), AgentSession.query() .findById(session.id) .patch({ status: 'active', + chatStatus: AgentChatStatus.READY, + workspaceStatus: AgentWorkspaceStatus.READY, } as unknown as Partial), ]); const finalizeMs = elapsedMs(finalizeStartedAt); @@ -1142,7 +1567,10 @@ export default class AgentSessionService { session = { ...session, status: 'active', + chatStatus: AgentChatStatus.READY, + workspaceStatus: AgentWorkspaceStatus.READY, } as AgentSession; + await AgentSandboxService.recordSessionSandboxState(session, { workspaceStorage }); await clearAgentSessionStartupFailure(redis, sessionUuid).catch(() => {}); @@ -1156,19 +1584,7 @@ export default class AgentSessionService { }` ); - const readySession = session; - // Default-thread creation stays best-effort here so chat readiness does not - // depend on secondary DB work; ThreadService.listThreadsForSession() will - // create or retry the default thread on first access if this async warm-up fails. - void (async () => { - const AgentThreadService = (await import('server/services/agent/ThreadService')).default; - await AgentThreadService.getDefaultThreadForSession(readySession.uuid, opts.userId); - })().catch((error: unknown) => { - logger().warn( - { error, sessionId: readySession.uuid }, - `Session: default thread creation skipped sessionId=${readySession.uuid}` - ); - }); + warmDefaultThread(session.uuid, opts.userId); return session!; } catch (err) { @@ -1208,16 +1624,32 @@ export default class AgentSessionService { const endedAt = new Date().toISOString(); if (sessionPersisted) { - await AgentSession.query() + const failedPatch = { + status: 'error', + chatStatus: AgentChatStatus.ERROR, + workspaceStatus: AgentWorkspaceStatus.FAILED, + endedAt, + } as unknown as Partial; + const patched = await AgentSession.query() .findById(session!.id) - .patch({ - status: 'error', - endedAt, - } as unknown as Partial) - .catch(() => {}); + .patch(failedPatch) + .then( + () => true, + () => false + ); + if (patched) { + const failedSession = { + ...session!, + ...failedPatch, + } as AgentSession; + await Promise.all([ + AgentSourceService.recordSessionState(failedSession).catch(() => {}), + AgentSandboxService.recordSessionSandboxState(failedSession, { workspaceStorage }).catch(() => {}), + ]); + } } else { - await AgentSession.query() - .insert({ + const failedSession = await AgentSession.query() + .insertAndFetch({ uuid: sessionUuid, userId: opts.userId, ownerGithubUsername: opts.userIdentity?.githubUsername || null, @@ -1225,16 +1657,27 @@ export default class AgentSessionService { namespace: opts.namespace, pvcName, model: resolvedModelId, + defaultModel: resolvedModelId, + defaultHarness: 'lifecycle_ai_sdk', status: 'error', buildUuid: opts.buildUuid || null, buildKind, + sessionKind, + chatStatus: AgentChatStatus.ERROR, + workspaceStatus: AgentWorkspaceStatus.FAILED, endedAt, devModeSnapshots: {}, forwardedAgentSecretProviders: forwardedAgentEnv.secretProviders, workspaceRepos, selectedServices, } as unknown as Partial) - .catch(() => {}); + .catch(() => null); + if (failedSession) { + await Promise.all([ + AgentSourceService.createSessionSource(failedSession, { workspaceStorage }).catch(() => {}), + AgentSandboxService.recordSessionSandboxState(failedSession, { workspaceStorage }).catch(() => {}), + ]); + } } const revertPromise = @@ -1285,9 +1728,58 @@ export default class AgentSessionService { const apiKeySecretName = `agent-secret-${session.uuid.slice(0, 8)}`; const redis = RedisClient.getInstance().getRedis(); + const markSessionEnded = async (extraPatch: Partial = {}) => { + const endedPatch = { + status: 'ended', + chatStatus: AgentChatStatus.ENDED, + workspaceStatus: AgentWorkspaceStatus.ENDED, + endedAt: new Date().toISOString(), + ...extraPatch, + } as unknown as Partial; + await AgentSession.query().findById(session.id).patch(endedPatch); + const endedSession = { + ...session, + ...endedPatch, + } as AgentSession; + + await Promise.all([ + AgentSourceService.recordSessionState(endedSession).catch(() => {}), + AgentSandboxService.recordSessionSandboxState(endedSession).catch(() => {}), + ]); + }; logger().info(`Session: ending sessionId=${sessionId} status=${session.status} namespace=${session.namespace}`); + if (session.sessionKind === AgentSessionKind.CHAT && session.namespace) { + await Promise.all([ + deleteNamespace(session.namespace).catch(() => {}), + clearAgentSessionStartupFailure(redis, session.uuid).catch(() => {}), + ]); + + await markSessionEnded({ + devModeSnapshots: {}, + }); + + await redis.del(`${SESSION_REDIS_PREFIX}${session.uuid}`); + + logger().info(`Session: ended sessionId=${sessionId} namespace=${session.namespace}`); + return; + } + + if (!session.namespace || !session.podName || !session.pvcName) { + await markSessionEnded({ + devModeSnapshots: {}, + }); + + await Promise.all([ + redis.del(`${SESSION_REDIS_PREFIX}${session.uuid}`), + clearAgentSessionStartupFailure(redis, session.uuid).catch(() => {}), + ]); + + logger().info(`Session: ended sessionId=${sessionId} namespace=none`); + return; + } + const build = session.buildUuid ? await Build.query() .findOne({ uuid: session.buildUuid }) @@ -1295,10 +1787,7 @@ export default class AgentSessionService { : null; if (build?.kind === BuildKind.SANDBOX) { - await AgentSession.query().findById(session.id).patch({ - status: 'ended', - endedAt: new Date().toISOString(), - }); + await markSessionEnded(); await Promise.all([ redis.del(`${SESSION_REDIS_PREFIX}${session.uuid}`), @@ -1347,9 +1836,7 @@ export default class AgentSessionService { } triggerDevModeDeployRestore(session.namespace, session.devModeSnapshots, devModeDeploys); - await AgentSession.query().findById(session.id).patch({ - status: 'ended', - endedAt: new Date().toISOString(), + await markSessionEnded({ devModeSnapshots: {}, }); @@ -1380,6 +1867,14 @@ export default class AgentSessionService { throw new Error('Session build context is missing'); } + if (!session.namespace || !session.podName || !session.pvcName) { + throw new Error('Session runtime is not ready for service attachment'); + } + + const namespace = session.namespace; + const podName = session.podName; + const pvcName = session.pvcName; + const workspaceRepos = session.workspaceRepos || []; if (workspaceRepos.length !== 1) { throw new Error('Connecting services after startup is only supported for single-repo sessions'); @@ -1445,7 +1940,7 @@ export default class AgentSessionService { ); const templatedServices = await resolveTemplatedDevConfigEnvs( session.buildUuid || undefined, - session.namespace, + namespace, candidateServices ); const { services: resolvedServices, selectedServices } = applyWorkspaceReposToServices( @@ -1459,30 +1954,28 @@ export default class AgentSessionService { const installCommand = buildCombinedInstallCommand(resolvedServices); logger().info( - `Session: services attaching sessionId=${sessionId} namespace=${session.namespace} services=${ + `Session: services attaching sessionId=${sessionId} namespace=${namespace} services=${ (resolvedServices || []).map((service) => service.name).join(',') || 'none' }` ); if (installCommand) { - await runCommandInSessionWorkspace(session.namespace, session.podName, installCommand); + await runCommandInSessionWorkspace(namespace, podName, installCommand); } if ((skillPlan.skills || []).length > 0) { await runCommandInSessionWorkspace( - session.namespace, - session.podName, + namespace, + podName, generateSkillBootstrapCommand(skillPlan, { useGitHubToken: true }) ); } const keepAttachedServicesOnSessionNode = await resolveSessionAttachmentPlacementPolicy(session); - const agentNodeName = keepAttachedServicesOnSessionNode - ? await resolveAgentPodNodeName(session.namespace, session.podName) - : null; + const agentNodeName = keepAttachedServicesOnSessionNode ? await resolveAgentPodNodeName(namespace, podName) : null; if (keepAttachedServicesOnSessionNode && !agentNodeName) { - throw new Error(`Session workspace pod ${session.podName} did not report a scheduled node`); + throw new Error(`Session workspace pod ${podName} did not report a scheduled node`); } const enabledDevModeDeployIds: number[] = []; @@ -1491,8 +1984,8 @@ export default class AgentSessionService { try { const enabledServices = await enableServicesInDevModeParallel({ - namespace: session.namespace, - pvcName: session.pvcName, + namespace, + pvcName, services: resolvedServices || [], requiredNodeName: keepAttachedServicesOnSessionNode ? agentNodeName || undefined : undefined, }).catch((error) => { @@ -1531,10 +2024,10 @@ export default class AgentSessionService { logger().info( { sessionId, - namespace: session.namespace, + namespace, services: serviceNames, }, - `Session: services attached sessionId=${sessionId} namespace=${session.namespace} services=${ + `Session: services attached sessionId=${sessionId} namespace=${namespace} services=${ serviceNames.join(',') || 'none' }` ); @@ -1553,7 +2046,7 @@ export default class AgentSessionService { } if (deploysToRevert.length > 0) { - await restoreDevModeDeploys(session.namespace, addedSnapshots, deploysToRevert).catch(() => {}); + await restoreDevModeDeploys(namespace, addedSnapshots, deploysToRevert).catch(() => {}); } } @@ -1577,7 +2070,9 @@ export default class AgentSessionService { configuredPrompt?: string ): Promise { const [session, effectiveConfig, approvalPolicy] = await Promise.all([ - AgentSession.query().findOne({ uuid: sessionId }).select('id', 'namespace', 'buildUuid', 'skillPlan'), + AgentSession.query() + .findOne({ uuid: sessionId }) + .select('id', 'namespace', 'buildUuid', 'skillPlan', 'sessionKind'), AgentSessionConfigService.getInstance().getEffectiveConfig(repoFullName), AgentPolicyService.getEffectivePolicy(repoFullName), ]); @@ -1588,13 +2083,17 @@ export default class AgentSessionService { return resolvedConfiguredPrompt; } + if (!session.namespace) { + return resolvedConfiguredPrompt; + } + try { const context = await resolveAgentSessionPromptContext({ sessionDbId: session.id, namespace: session.namespace, buildUuid: session.buildUuid, }); - const toolLines = repoFullName + const toolLines = session.namespace ? buildSessionWorkspacePromptLines({ approvalPolicy, toolRules: effectiveConfig.toolRules, diff --git a/src/server/services/agentSessionConfig.ts b/src/server/services/agentSessionConfig.ts index 52599520..ba8f7e40 100644 --- a/src/server/services/agentSessionConfig.ts +++ b/src/server/services/agentSessionConfig.ts @@ -43,7 +43,14 @@ import { import { McpConfigService } from 'server/services/ai/mcp/config'; import { normalizeAuthConfig, requiresUserConnection } from 'server/services/ai/mcp/connectionConfig'; import AgentPolicyService from './agent/PolicyService'; -import { buildAgentToolKey, SESSION_WORKSPACE_SERVER_NAME, SESSION_WORKSPACE_SERVER_SLUG } from './agent/toolKeys'; +import { + buildAgentToolKey, + CHAT_PUBLISH_HTTP_TOOL_NAME, + LIFECYCLE_BUILTIN_SERVER_NAME, + LIFECYCLE_BUILTIN_SERVER_SLUG, + SESSION_WORKSPACE_SERVER_NAME, + SESSION_WORKSPACE_SERVER_SLUG, +} from './agent/toolKeys'; import type { McpDiscoveredTool } from 'server/services/ai/mcp/types'; import { getSessionWorkspaceToolSortKey, @@ -119,6 +126,22 @@ function normalizeStringRecord(value: unknown): Record | undefin return Object.keys(normalized).length > 0 ? normalized : undefined; } +function normalizeStringArray(value: unknown): string[] | undefined { + if (!Array.isArray(value)) { + return undefined; + } + + const normalized = Array.from( + new Set(value.filter((entry): entry is string => typeof entry === 'string').map((entry) => entry.trim())) + ).filter(Boolean); + + return normalized.length > 0 ? normalized : undefined; +} + +function normalizeWorkspaceStorageAccessMode(value: unknown): 'ReadWriteOnce' | 'ReadWriteMany' | undefined { + return value === 'ReadWriteOnce' || value === 'ReadWriteMany' ? value : undefined; +} + function normalizeResourceRequirements(value: unknown) { if (!value || typeof value !== 'object' || Array.isArray(value)) { return undefined; @@ -137,7 +160,7 @@ function normalizeResourceRequirements(value: unknown) { }; } -function validateRequiredRuntimeImages(config: Partial): void { +function validateRequiredRuntimeImages(config: Partial): void { const missingFields: string[] = []; if (!normalizeOptionalString(config.workspaceImage)) { @@ -228,6 +251,51 @@ function normalizeRuntimeSettings(value: unknown): AgentSessionRuntimeSettingsVa const workspaceGatewayResources = normalizeResourceRequirements( (value as { resources?: { workspaceGateway?: unknown } }).resources?.workspaceGateway ); + const workspaceStorageDefaultSize = normalizeOptionalString( + (value as { workspaceStorage?: { defaultSize?: unknown } }).workspaceStorage?.defaultSize + ); + const workspaceStorageAllowedSizes = normalizeStringArray( + (value as { workspaceStorage?: { allowedSizes?: unknown } }).workspaceStorage?.allowedSizes + ); + const workspaceStorageAllowClientOverride = normalizeBoolean( + (value as { workspaceStorage?: { allowClientOverride?: unknown } }).workspaceStorage?.allowClientOverride + ); + const workspaceStorageAccessMode = normalizeWorkspaceStorageAccessMode( + (value as { workspaceStorage?: { accessMode?: unknown } }).workspaceStorage?.accessMode + ); + const cleanupActiveIdleSuspendMs = normalizePositiveInteger( + (value as { cleanup?: { activeIdleSuspendMs?: unknown } }).cleanup?.activeIdleSuspendMs + ); + const cleanupStartingTimeoutMs = normalizePositiveInteger( + (value as { cleanup?: { startingTimeoutMs?: unknown } }).cleanup?.startingTimeoutMs + ); + const cleanupHibernatedRetentionMs = normalizePositiveInteger( + (value as { cleanup?: { hibernatedRetentionMs?: unknown } }).cleanup?.hibernatedRetentionMs + ); + const cleanupIntervalMs = normalizePositiveInteger( + (value as { cleanup?: { intervalMs?: unknown } }).cleanup?.intervalMs + ); + const cleanupRedisTtlSeconds = normalizePositiveInteger( + (value as { cleanup?: { redisTtlSeconds?: unknown } }).cleanup?.redisTtlSeconds + ); + const durabilityRunExecutionLeaseMs = normalizePositiveInteger( + (value as { durability?: { runExecutionLeaseMs?: unknown } }).durability?.runExecutionLeaseMs + ); + const durabilityQueuedRunDispatchStaleMs = normalizePositiveInteger( + (value as { durability?: { queuedRunDispatchStaleMs?: unknown } }).durability?.queuedRunDispatchStaleMs + ); + const durabilityDispatchRecoveryLimit = normalizePositiveInteger( + (value as { durability?: { dispatchRecoveryLimit?: unknown } }).durability?.dispatchRecoveryLimit + ); + const durabilityMaxDurablePayloadBytes = normalizePositiveInteger( + (value as { durability?: { maxDurablePayloadBytes?: unknown } }).durability?.maxDurablePayloadBytes + ); + const durabilityPayloadPreviewBytes = normalizePositiveInteger( + (value as { durability?: { payloadPreviewBytes?: unknown } }).durability?.payloadPreviewBytes + ); + const durabilityFileChangePreviewChars = normalizePositiveInteger( + (value as { durability?: { fileChangePreviewChars?: unknown } }).durability?.fileChangePreviewChars + ); return { ...(workspaceImage ? { workspaceImage } : {}), @@ -258,6 +326,67 @@ function normalizeRuntimeSettings(value: unknown): AgentSessionRuntimeSettingsVa }, } : {}), + ...(workspaceStorageDefaultSize || + workspaceStorageAllowedSizes || + workspaceStorageAllowClientOverride !== undefined || + workspaceStorageAccessMode + ? { + workspaceStorage: { + ...(workspaceStorageDefaultSize ? { defaultSize: workspaceStorageDefaultSize } : {}), + ...(workspaceStorageAllowedSizes ? { allowedSizes: workspaceStorageAllowedSizes } : {}), + ...(workspaceStorageAllowClientOverride !== undefined + ? { allowClientOverride: workspaceStorageAllowClientOverride } + : {}), + ...(workspaceStorageAccessMode ? { accessMode: workspaceStorageAccessMode } : {}), + }, + } + : {}), + ...(cleanupActiveIdleSuspendMs !== undefined || + cleanupStartingTimeoutMs !== undefined || + cleanupHibernatedRetentionMs !== undefined || + cleanupIntervalMs !== undefined || + cleanupRedisTtlSeconds !== undefined + ? { + cleanup: { + ...(cleanupActiveIdleSuspendMs !== undefined ? { activeIdleSuspendMs: cleanupActiveIdleSuspendMs } : {}), + ...(cleanupStartingTimeoutMs !== undefined ? { startingTimeoutMs: cleanupStartingTimeoutMs } : {}), + ...(cleanupHibernatedRetentionMs !== undefined + ? { hibernatedRetentionMs: cleanupHibernatedRetentionMs } + : {}), + ...(cleanupIntervalMs !== undefined ? { intervalMs: cleanupIntervalMs } : {}), + ...(cleanupRedisTtlSeconds !== undefined ? { redisTtlSeconds: cleanupRedisTtlSeconds } : {}), + }, + } + : {}), + ...(durabilityRunExecutionLeaseMs !== undefined || + durabilityQueuedRunDispatchStaleMs !== undefined || + durabilityDispatchRecoveryLimit !== undefined || + durabilityMaxDurablePayloadBytes !== undefined || + durabilityPayloadPreviewBytes !== undefined || + durabilityFileChangePreviewChars !== undefined + ? { + durability: { + ...(durabilityRunExecutionLeaseMs !== undefined + ? { runExecutionLeaseMs: durabilityRunExecutionLeaseMs } + : {}), + ...(durabilityQueuedRunDispatchStaleMs !== undefined + ? { queuedRunDispatchStaleMs: durabilityQueuedRunDispatchStaleMs } + : {}), + ...(durabilityDispatchRecoveryLimit !== undefined + ? { dispatchRecoveryLimit: durabilityDispatchRecoveryLimit } + : {}), + ...(durabilityMaxDurablePayloadBytes !== undefined + ? { maxDurablePayloadBytes: durabilityMaxDurablePayloadBytes } + : {}), + ...(durabilityPayloadPreviewBytes !== undefined + ? { payloadPreviewBytes: durabilityPayloadPreviewBytes } + : {}), + ...(durabilityFileChangePreviewChars !== undefined + ? { fileChangePreviewChars: durabilityFileChangePreviewChars } + : {}), + }, + } + : {}), }; } @@ -359,6 +488,9 @@ export default class AgentSessionConfigService extends BaseService { delete nextDefaults.scheduling; delete nextDefaults.readiness; delete nextDefaults.resources; + delete nextDefaults.workspaceStorage; + delete nextDefaults.cleanup; + delete nextDefaults.durability; if (normalized.workspaceImage) { nextDefaults.workspaceImage = normalized.workspaceImage; @@ -378,6 +510,15 @@ export default class AgentSessionConfigService extends BaseService { if (normalized.resources) { nextDefaults.resources = normalized.resources; } + if (normalized.workspaceStorage) { + nextDefaults.workspaceStorage = normalized.workspaceStorage; + } + if (normalized.cleanup) { + nextDefaults.cleanup = normalized.cleanup; + } + if (normalized.durability) { + nextDefaults.durability = normalized.durability; + } validateRequiredRuntimeImages(nextDefaults); @@ -487,6 +628,7 @@ export default class AgentSessionConfigService extends BaseService { serverName, sourceType, sourceScope, + capabilityKey: capabilityKeyOverride, annotations, }: { toolName: string; @@ -495,13 +637,15 @@ export default class AgentSessionConfigService extends BaseService { serverName: string; sourceType: 'builtin' | 'mcp'; sourceScope: string; + capabilityKey?: AgentSessionToolInventoryEntry['capabilityKey']; annotations?: McpDiscoveredTool['annotations']; }) => { const toolKey = buildAgentToolKey(serverSlug, toolName); const capabilityKey = - sourceType === 'builtin' + capabilityKeyOverride || + (sourceType === 'builtin' ? AgentPolicyService.capabilityForSessionWorkspaceTool(toolName, annotations) - : AgentPolicyService.capabilityForExternalMcpTool(toolName, annotations); + : AgentPolicyService.capabilityForExternalMcpTool(toolName, annotations)); const approvalMode = AgentPolicyService.modeForCapability(approvalPolicy, capabilityKey); const scopeRuleMode = toRuleSelection(activeScopeConfig.toolRules || [], toolKey); const effectiveRuleMode = toRuleSelection(effectiveConfig.toolRules, toolKey); @@ -541,6 +685,16 @@ export default class AgentSessionConfigService extends BaseService { }); } + appendEntry({ + toolName: CHAT_PUBLISH_HTTP_TOOL_NAME, + description: 'Expose a running HTTP app from the chat workspace and return its reachable URL.', + serverSlug: LIFECYCLE_BUILTIN_SERVER_SLUG, + serverName: LIFECYCLE_BUILTIN_SERVER_NAME, + sourceType: 'builtin', + sourceScope: 'session', + capabilityKey: 'deploy_k8s_mutation', + }); + for (const config of mcpDefinitions) { const tools = await this.listDiscoveredToolsForDefinition(config); for (const tool of tools) { diff --git a/src/server/services/ai/mcp/config.ts b/src/server/services/ai/mcp/config.ts index 91637474..2bd7d335 100644 --- a/src/server/services/ai/mcp/config.ts +++ b/src/server/services/ai/mcp/config.ts @@ -274,15 +274,15 @@ export class McpConfigService { }); } - async resolveServersForRepo( - repoFullName: string, + async resolveServers( + repoFullName?: string, disabledSlugs?: string[], userIdentity?: RequestUserIdentity | null ): Promise { const configs = await this.listEffectiveConfigs(repoFullName); const disabled = new Set(disabledSlugs ?? []); const filteredConfigs = configs.filter((config) => !disabled.has(config.slug)); - const sharedScopes = ['global', repoFullName]; + const sharedScopes = ['global', ...(repoFullName ? [repoFullName] : [])]; const definitionFingerprints = buildDefinitionFingerprintMap(filteredConfigs); const userConnections = userIdentity @@ -395,6 +395,14 @@ export class McpConfigService { }); } + async resolveServersForRepo( + repoFullName: string, + disabledSlugs?: string[], + userIdentity?: RequestUserIdentity | null + ): Promise { + return this.resolveServers(repoFullName, disabledSlugs, userIdentity); + } + async resolveSessionPodServersForRepo( repoFullName: string, disabledSlugs?: string[], diff --git a/src/server/services/globalConfig.ts b/src/server/services/globalConfig.ts index e05e67b3..e955e59d 100644 --- a/src/server/services/globalConfig.ts +++ b/src/server/services/globalConfig.ts @@ -276,6 +276,12 @@ export default class GlobalConfigService extends BaseService { async setConfig(key: string, value: any): Promise { try { await this.db.knex('global_config').insert({ key, config: value }).onConflict('key').merge(); + this.clearMemoryCache(); + try { + await this.redis.del(REDIS_CACHE_KEY); + } catch (cacheError) { + getLogger().warn({ error: cacheError }, `Config: cache clear failed key=${key}`); + } getLogger().info(`Config: set key=${key}`); } catch (err: any) { getLogger().error({ error: err }, `Config: set failed key=${key}`); diff --git a/src/server/services/types/agentSessionConfig.ts b/src/server/services/types/agentSessionConfig.ts index e4c1e5cb..e67de88a 100644 --- a/src/server/services/types/agentSessionConfig.ts +++ b/src/server/services/types/agentSessionConfig.ts @@ -15,6 +15,7 @@ */ import type { AgentApprovalMode, AgentCapabilityKey } from 'server/services/agent/types'; +import type { AgentSessionWorkspaceStorageAccessMode } from './globalConfig'; export type AgentSessionToolRuleMode = AgentApprovalMode; export type AgentSessionToolRuleSelection = AgentSessionToolRuleMode | 'inherit'; @@ -52,6 +53,30 @@ export interface AgentSessionResourceRequirementsValue { limits?: Record; } +export interface AgentSessionWorkspaceStorageSettingsValue { + defaultSize?: string; + allowedSizes?: string[]; + allowClientOverride?: boolean; + accessMode?: AgentSessionWorkspaceStorageAccessMode; +} + +export interface AgentSessionCleanupSettingsValue { + activeIdleSuspendMs?: number; + startingTimeoutMs?: number; + hibernatedRetentionMs?: number; + intervalMs?: number; + redisTtlSeconds?: number; +} + +export interface AgentSessionDurabilitySettingsValue { + runExecutionLeaseMs?: number; + queuedRunDispatchStaleMs?: number; + dispatchRecoveryLimit?: number; + maxDurablePayloadBytes?: number; + payloadPreviewBytes?: number; + fileChangePreviewChars?: number; +} + export interface AgentSessionRuntimeSettingsValue { workspaceImage?: string; workspaceEditorImage?: string; @@ -66,6 +91,9 @@ export interface AgentSessionRuntimeSettingsValue { editor?: AgentSessionResourceRequirementsValue; workspaceGateway?: AgentSessionResourceRequirementsValue; }; + workspaceStorage?: AgentSessionWorkspaceStorageSettingsValue; + cleanup?: AgentSessionCleanupSettingsValue; + durability?: AgentSessionDurabilitySettingsValue; } export interface AgentSessionToolInventoryEntry { diff --git a/src/server/services/types/globalConfig.ts b/src/server/services/types/globalConfig.ts index b90cc9f2..4dc4f017 100644 --- a/src/server/services/types/globalConfig.ts +++ b/src/server/services/types/globalConfig.ts @@ -82,6 +82,32 @@ export type AgentSessionResourcesConfig = { workspaceGateway?: ResourceRequirements | null; }; +export type AgentSessionWorkspaceStorageAccessMode = 'ReadWriteOnce' | 'ReadWriteMany'; + +export type AgentSessionWorkspaceStorageConfig = { + defaultSize?: string | null; + allowedSizes?: string[] | null; + allowClientOverride?: boolean | null; + accessMode?: AgentSessionWorkspaceStorageAccessMode | null; +}; + +export type AgentSessionCleanupConfig = { + activeIdleSuspendMs?: number | string | null; + startingTimeoutMs?: number | string | null; + hibernatedRetentionMs?: number | string | null; + intervalMs?: number | string | null; + redisTtlSeconds?: number | string | null; +}; + +export type AgentSessionDurabilityConfig = { + runExecutionLeaseMs?: number | string | null; + queuedRunDispatchStaleMs?: number | string | null; + dispatchRecoveryLimit?: number | string | null; + maxDurablePayloadBytes?: number | string | null; + payloadPreviewBytes?: number | string | null; + fileChangePreviewChars?: number | string | null; +}; + export type AgentSessionDefaults = { workspaceImage?: string | null; workspaceEditorImage?: string | null; @@ -89,6 +115,9 @@ export type AgentSessionDefaults = { scheduling?: AgentSessionSchedulingConfig; readiness?: AgentSessionReadinessConfig; resources?: AgentSessionResourcesConfig; + workspaceStorage?: AgentSessionWorkspaceStorageConfig; + cleanup?: AgentSessionCleanupConfig; + durability?: AgentSessionDurabilityConfig; controlPlane?: AgentSessionControlPlaneConfig; }; diff --git a/src/shared/config.ts b/src/shared/config.ts index 8681fa74..01902ec4 100644 --- a/src/shared/config.ts +++ b/src/shared/config.ts @@ -119,6 +119,8 @@ export const QUEUE_NAMES = { AGENT_SESSION_CLEANUP: 'agent_session_cleanup', AGENT_SESSION_PREWARM: 'agent_session_prewarm', AGENT_SANDBOX_SESSION_LAUNCH: 'agent_sandbox_session_launch', + AGENT_RUN_EXECUTE: 'agent_run_execute', + AGENT_RUN_RECOVERY: 'agent_run_recovery', } as const; export const GITHUB_APP_INSTALLATION_ID = getServerRuntimeConfig('GITHUB_APP_INSTALLATION_ID', 'YOUR_VALUE_HERE'); diff --git a/src/shared/constants.ts b/src/shared/constants.ts index e1216361..31d0e676 100644 --- a/src/shared/constants.ts +++ b/src/shared/constants.ts @@ -78,6 +78,27 @@ export enum BuildKind { SANDBOX = 'sandbox', } +export enum AgentSessionKind { + ENVIRONMENT = 'environment', + SANDBOX = 'sandbox', + CHAT = 'chat', +} + +export enum AgentChatStatus { + READY = 'ready', + ENDED = 'ended', + ERROR = 'error', +} + +export enum AgentWorkspaceStatus { + NONE = 'none', + PROVISIONING = 'provisioning', + READY = 'ready', + HIBERNATED = 'hibernated', + FAILED = 'failed', + ENDED = 'ended', +} + export enum PullRequestStatus { OPEN = 'open', CLOSED = 'closed', diff --git a/src/shared/openApiSpec.test.ts b/src/shared/openApiSpec.test.ts new file mode 100644 index 00000000..416b4eba --- /dev/null +++ b/src/shared/openApiSpec.test.ts @@ -0,0 +1,68 @@ +import swaggerJSDoc from 'swagger-jsdoc'; +import { openApiSpecificationForV2Api } from './openApiSpec'; + +const swaggerSpec = swaggerJSDoc(openApiSpecificationForV2Api) as any; +const schemas = swaggerSpec.components.schemas; + +function getJsonErrorSchema(path: string, method: string, status: string) { + return swaggerSpec.paths[path][method].responses[status].content['application/json'].schema; +} + +describe('OpenAPI v2 agent session contract', () => { + it('documents canonical run events with public context and a version', () => { + const eventSchema = schemas.AgentRunMessagePartEvent; + + expect(eventSchema.required).toEqual( + expect.arrayContaining(['id', 'runId', 'threadId', 'sessionId', 'sequence', 'eventType', 'version', 'payload']) + ); + expect(eventSchema.properties.threadId).toEqual({ type: 'string', description: 'Public thread UUID.' }); + expect(eventSchema.properties.sessionId).toEqual({ type: 'string', description: 'Public session UUID.' }); + expect(eventSchema.properties.version).toEqual({ type: 'integer', enum: [1] }); + }); + + it('keeps canonical reference message parts aligned with runtime validation', () => { + const messagePartSchema = schemas.CanonicalAgentMessagePart; + const fileRefSchema = messagePartSchema.oneOf.find((entry: any) => entry.properties.type.enum[0] === 'file_ref'); + const sourceRefSchema = messagePartSchema.oneOf.find( + (entry: any) => entry.properties.type.enum[0] === 'source_ref' + ); + + expect(fileRefSchema.anyOf).toEqual( + expect.arrayContaining([ + expect.objectContaining({ required: ['path'], properties: { path: { type: 'string', minLength: 1 } } }), + expect.objectContaining({ required: ['url'], properties: { url: { type: 'string', minLength: 1 } } }), + ]) + ); + expect(sourceRefSchema.anyOf).toEqual( + expect.arrayContaining([ + expect.objectContaining({ required: ['url'], properties: { url: { type: 'string', minLength: 1 } } }), + expect.objectContaining({ required: ['title'], properties: { title: { type: 'string', minLength: 1 } } }), + ]) + ); + }); + + it('removes migration bridges without preserving UI-shaped success contracts', () => { + expect(swaggerSpec.paths['/api/v2/ai/agent/threads/{threadId}/conversation']).toBeUndefined(); + expect(swaggerSpec.paths['/api/v2/ai/agent/runs/{runId}/stream']).toBeUndefined(); + expect(swaggerSpec.paths['/api/v2/ai/agent/pending-actions/{actionId}/approve']).toBeUndefined(); + expect(swaggerSpec.paths['/api/v2/ai/agent/pending-actions/{actionId}/deny']).toBeUndefined(); + expect(schemas.AgentUIMessage).toBeUndefined(); + expect(schemas.AgentUIMessagePart).toBeUndefined(); + expect(schemas.AgentUIMessageMetadata).toBeUndefined(); + }); + + it('documents JSON error responses for changed canonical endpoints', () => { + expect(getJsonErrorSchema('/api/v2/ai/agent/threads/{threadId}/messages', 'get', '400')).toEqual({ + $ref: '#/components/schemas/ApiErrorResponse', + }); + expect(getJsonErrorSchema('/api/v2/ai/agent/runs/{runId}/events/stream', 'get', '400')).toEqual({ + $ref: '#/components/schemas/ApiErrorResponse', + }); + expect(getJsonErrorSchema('/api/v2/ai/config/agent-session', 'get', '401')).toEqual({ + $ref: '#/components/schemas/ApiErrorResponse', + }); + expect(getJsonErrorSchema('/api/v2/ai/config/agent-session/runtime', 'put', '401')).toEqual({ + $ref: '#/components/schemas/ApiErrorResponse', + }); + }); +}); diff --git a/src/shared/openApiSpec.ts b/src/shared/openApiSpec.ts index 9b097e53..0b37e0c4 100644 --- a/src/shared/openApiSpec.ts +++ b/src/shared/openApiSpec.ts @@ -1,5 +1,87 @@ import { OAS3Options } from 'swagger-jsdoc'; -import { BuildKind, BuildStatus, DeployStatus, DeployTypes } from './constants'; +import { + AgentChatStatus, + AgentSessionKind, + AgentWorkspaceStatus, + BuildKind, + BuildStatus, + DeployStatus, + DeployTypes, +} from './constants'; + +const agentRunEventBaseProperties = { + id: { type: 'string' }, + runId: { type: 'string', description: 'Public run UUID.' }, + threadId: { type: 'string', description: 'Public thread UUID.' }, + sessionId: { type: 'string', description: 'Public session UUID.' }, + sequence: { type: 'integer' }, + version: { type: 'integer', enum: [1] }, + createdAt: { type: 'string', format: 'date-time', nullable: true }, + updatedAt: { type: 'string', format: 'date-time', nullable: true }, +}; + +const agentRunEventRequired = [ + 'id', + 'runId', + 'threadId', + 'sessionId', + 'sequence', + 'eventType', + 'version', + 'payload', + 'createdAt', + 'updatedAt', +]; + +function agentRunEventSchema(eventTypes: string[], payload: Record) { + return { + type: 'object', + properties: { + ...agentRunEventBaseProperties, + eventType: { type: 'string', enum: eventTypes }, + payload: { + oneOf: [payload, { $ref: '#/components/schemas/AgentRunTruncatedValue' }], + }, + }, + required: agentRunEventRequired, + additionalProperties: false, + }; +} + +const agentRunEventDiscriminatorMapping = { + 'message.created': '#/components/schemas/AgentRunMessageCreatedEvent', + 'message.metadata': '#/components/schemas/AgentRunMessageMetadataEvent', + 'message.part.started': '#/components/schemas/AgentRunMessagePartEvent', + 'message.delta': '#/components/schemas/AgentRunMessagePartEvent', + 'message.part.completed': '#/components/schemas/AgentRunMessagePartEvent', + 'message.source': '#/components/schemas/AgentRunMessageSourceEvent', + 'message.file': '#/components/schemas/AgentRunMessageFileEvent', + 'tool.call.input.started': '#/components/schemas/AgentRunToolInputStartedEvent', + 'tool.call.input.delta': '#/components/schemas/AgentRunToolInputDeltaEvent', + 'tool.call.started': '#/components/schemas/AgentRunToolStartedEvent', + 'tool.call.completed': '#/components/schemas/AgentRunToolCompletedEvent', + 'tool.file_change': '#/components/schemas/AgentRunToolFileChangeEvent', + 'approval.requested': '#/components/schemas/AgentRunApprovalRequestedEvent', + 'approval.resolved': '#/components/schemas/AgentRunApprovalResolvedEvent', + 'approval.responded': '#/components/schemas/AgentRunApprovalRespondedEvent', + 'run.queued': '#/components/schemas/AgentRunStatusEvent', + 'run.started': '#/components/schemas/AgentRunStatusEvent', + 'run.waiting_for_approval': '#/components/schemas/AgentRunStatusEvent', + 'run.completed': '#/components/schemas/AgentRunStatusEvent', + 'run.failed': '#/components/schemas/AgentRunStatusEvent', + 'run.cancelled': '#/components/schemas/AgentRunStatusEvent', + 'run.updated': '#/components/schemas/AgentRunStatusEvent', + 'run.step.started': '#/components/schemas/AgentRunStepEvent', + 'run.step.completed': '#/components/schemas/AgentRunStepEvent', + 'run.finished': '#/components/schemas/AgentRunFinishedEvent', + 'run.error': '#/components/schemas/AgentRunErrorEvent', + 'run.aborted': '#/components/schemas/AgentRunAbortedEvent', +}; + +const agentRunEventPayloadMetadata = { + type: 'object', + additionalProperties: true, +}; export const openApiSpecificationForV2Api: OAS3Options = { definition: { @@ -76,6 +158,8 @@ export const openApiSpecificationForV2Api: OAS3Options = { type: 'object', properties: { pagination: { $ref: '#/components/schemas/PaginationMetadata' }, + limit: { type: 'integer' }, + maxLimit: { type: 'integer' }, }, }, @@ -342,45 +426,227 @@ export const openApiSpecificationForV2Api: OAS3Options = { }, }, additionalProperties: false, + example: { + maxIterations: 12, + workspaceToolDiscoveryTimeoutMs: 30000, + workspaceToolExecutionTimeoutMs: 120000, + toolRules: [ + { + toolKey: 'mcp__sandbox__workspace_edit_file', + mode: 'require_approval', + }, + ], + }, }, - AgentUIMessageMetadata: { + CanonicalAgentMessagePart: { + oneOf: [ + { + type: 'object', + properties: { + type: { type: 'string', enum: ['text'] }, + text: { type: 'string', minLength: 1 }, + }, + required: ['type', 'text'], + additionalProperties: false, + }, + { + type: 'object', + properties: { + type: { type: 'string', enum: ['reasoning'] }, + text: { type: 'string', minLength: 1 }, + }, + required: ['type', 'text'], + additionalProperties: false, + }, + { + type: 'object', + properties: { + type: { type: 'string', enum: ['file_ref'] }, + path: { type: 'string', minLength: 1, nullable: true }, + url: { type: 'string', minLength: 1, nullable: true }, + mediaType: { type: 'string', minLength: 1, nullable: true }, + title: { type: 'string', minLength: 1, nullable: true }, + }, + required: ['type'], + anyOf: [ + { + type: 'object', + properties: { path: { type: 'string', minLength: 1 } }, + required: ['path'], + }, + { + type: 'object', + properties: { url: { type: 'string', minLength: 1 } }, + required: ['url'], + }, + ], + additionalProperties: false, + }, + { + type: 'object', + properties: { + type: { type: 'string', enum: ['source_ref'] }, + url: { type: 'string', minLength: 1, nullable: true }, + title: { type: 'string', minLength: 1, nullable: true }, + sourceType: { type: 'string', minLength: 1, nullable: true }, + }, + required: ['type'], + anyOf: [ + { + type: 'object', + properties: { url: { type: 'string', minLength: 1 } }, + required: ['url'], + }, + { + type: 'object', + properties: { title: { type: 'string', minLength: 1 } }, + required: ['title'], + }, + ], + additionalProperties: false, + }, + ], + }, + + AgentMessage: { type: 'object', properties: { - sessionId: { type: 'string', nullable: true }, - threadId: { type: 'string', nullable: true }, + id: { type: 'string' }, + clientMessageId: { type: 'string', nullable: true }, + threadId: { type: 'string' }, runId: { type: 'string', nullable: true }, - provider: { type: 'string', nullable: true }, - model: { type: 'string', nullable: true }, + role: { type: 'string', enum: ['user', 'assistant'] }, + parts: { + type: 'array', + items: { $ref: '#/components/schemas/CanonicalAgentMessagePart' }, + minItems: 1, + }, createdAt: { type: 'string', format: 'date-time', nullable: true }, - completedAt: { type: 'string', format: 'date-time', nullable: true }, - usage: { type: 'object', nullable: true, additionalProperties: true }, }, - additionalProperties: true, + required: ['id', 'clientMessageId', 'threadId', 'runId', 'role', 'parts', 'createdAt'], + additionalProperties: false, + example: { + id: 'message-1', + clientMessageId: 'client-message-1', + threadId: 'thread-1', + runId: 'run-1', + role: 'user', + parts: [{ type: 'text', text: 'Check the sample service.' }], + createdAt: '2026-04-25T00:00:00.000Z', + }, }, - AgentUIMessagePart: { + AgentThreadMessagesResponse: { type: 'object', properties: { - type: { type: 'string' }, + thread: { $ref: '#/components/schemas/AgentThread' }, + messages: { + type: 'array', + items: { $ref: '#/components/schemas/AgentMessage' }, + }, + pagination: { + type: 'object', + properties: { + hasMore: { type: 'boolean' }, + nextBeforeMessageId: { type: 'string', nullable: true }, + }, + required: ['hasMore', 'nextBeforeMessageId'], + additionalProperties: false, + }, }, - required: ['type'], - additionalProperties: true, + required: ['thread', 'messages', 'pagination'], + additionalProperties: false, }, - AgentUIMessage: { + AgentRunRuntimeOptions: { type: 'object', properties: { - id: { type: 'string' }, - role: { type: 'string', enum: ['system', 'user', 'assistant', 'tool'] }, - metadata: { $ref: '#/components/schemas/AgentUIMessageMetadata' }, + maxIterations: { + type: 'integer', + minimum: 1, + maximum: 100, + }, + }, + additionalProperties: false, + example: { + workspaceImage: 'registry.example.test/lifecycle/workspace:sample', + workspaceEditorImage: 'registry.example.test/lifecycle/editor:sample', + workspaceGatewayImage: 'registry.example.test/lifecycle/gateway:sample', + scheduling: { + keepAttachedServicesOnSessionNode: true, + }, + readiness: { + timeoutMs: 120000, + pollMs: 2000, + }, + workspaceStorage: { + defaultSize: '10Gi', + allowedSizes: ['10Gi', '20Gi'], + allowClientOverride: true, + accessMode: 'ReadWriteOnce', + }, + cleanup: { + activeIdleSuspendMs: 1800000, + startingTimeoutMs: 900000, + hibernatedRetentionMs: 86400000, + intervalMs: 300000, + redisTtlSeconds: 7200, + }, + durability: { + runExecutionLeaseMs: 1800000, + queuedRunDispatchStaleMs: 30000, + dispatchRecoveryLimit: 50, + maxDurablePayloadBytes: 65536, + payloadPreviewBytes: 16384, + fileChangePreviewChars: 4000, + }, + }, + }, + + CreateAgentThreadRunMessage: { + type: 'object', + properties: { + clientMessageId: { type: 'string' }, parts: { type: 'array', - items: { $ref: '#/components/schemas/AgentUIMessagePart' }, + items: { $ref: '#/components/schemas/CanonicalAgentMessagePart' }, + minItems: 1, + }, + }, + required: ['parts'], + additionalProperties: false, + }, + + CreateAgentThreadRunRequest: { + type: 'object', + properties: { + message: { $ref: '#/components/schemas/CreateAgentThreadRunMessage' }, + model: { + type: 'object', + properties: { + provider: { type: 'string' }, + id: { type: 'string' }, + }, + additionalProperties: false, + }, + runtimeOptions: { $ref: '#/components/schemas/AgentRunRuntimeOptions' }, + }, + required: ['message'], + additionalProperties: false, + example: { + message: { + clientMessageId: 'client-message-1', + parts: [{ type: 'text', text: 'Check the sample service.' }], + }, + model: { + provider: 'openai', + id: 'gpt-5.2', + }, + runtimeOptions: { + maxIterations: 12, }, }, - required: ['id', 'role', 'parts'], - additionalProperties: true, }, AgentThread: { @@ -399,28 +665,152 @@ export const openApiSpecificationForV2Api: OAS3Options = { required: ['id', 'isDefault', 'metadata'], }, + AgentSessionDefaults: { + type: 'object', + properties: { + model: { type: 'string' }, + harness: { type: 'string', nullable: true }, + }, + required: ['model', 'harness'], + }, + + AgentSource: { + type: 'object', + properties: { + id: { type: 'string' }, + adapter: { type: 'string' }, + status: { type: 'string', enum: ['requested', 'preparing', 'ready', 'failed', 'cleaned_up'] }, + input: { type: 'object', additionalProperties: true }, + sandboxRequirements: { type: 'object', additionalProperties: true }, + error: { type: 'object', additionalProperties: true, nullable: true }, + preparedAt: { type: 'string', format: 'date-time', nullable: true }, + cleanedUpAt: { type: 'string', format: 'date-time', nullable: true }, + createdAt: { type: 'string', format: 'date-time', nullable: true }, + updatedAt: { type: 'string', format: 'date-time', nullable: true }, + }, + required: ['id', 'adapter', 'status', 'input', 'sandboxRequirements', 'error'], + }, + + AgentSandboxExposure: { + type: 'object', + properties: { + id: { type: 'string' }, + kind: { type: 'string' }, + status: { type: 'string', enum: ['provisioning', 'ready', 'failed', 'ended'] }, + targetPort: { type: 'integer', nullable: true }, + url: { type: 'string', nullable: true }, + metadata: { type: 'object', additionalProperties: true }, + lastVerifiedAt: { type: 'string', format: 'date-time', nullable: true }, + endedAt: { type: 'string', format: 'date-time', nullable: true }, + createdAt: { type: 'string', format: 'date-time', nullable: true }, + updatedAt: { type: 'string', format: 'date-time', nullable: true }, + }, + required: ['id', 'kind', 'status', 'metadata'], + }, + + AgentSandbox: { + type: 'object', + properties: { + id: { type: 'string', nullable: true }, + generation: { type: 'integer', nullable: true }, + provider: { type: 'string', nullable: true }, + status: { + type: 'string', + enum: ['none', 'provisioning', 'ready', 'suspending', 'suspended', 'resuming', 'failed', 'ended'], + }, + capabilitySnapshot: { type: 'object', additionalProperties: true }, + exposures: { + type: 'array', + items: { $ref: '#/components/schemas/AgentSandboxExposure' }, + }, + suspendedAt: { type: 'string', format: 'date-time', nullable: true }, + endedAt: { type: 'string', format: 'date-time', nullable: true }, + error: { type: 'object', additionalProperties: true, nullable: true }, + createdAt: { type: 'string', format: 'date-time', nullable: true }, + updatedAt: { type: 'string', format: 'date-time', nullable: true }, + }, + required: ['id', 'generation', 'provider', 'status', 'capabilitySnapshot', 'exposures', 'error'], + }, + AgentSessionSummary: { + type: 'object', + properties: { + session: { + type: 'object', + properties: { + id: { type: 'string' }, + status: { type: 'string', enum: ['ready', 'ended', 'error'] }, + userId: { type: 'string' }, + ownerGithubUsername: { type: 'string', nullable: true }, + defaults: { $ref: '#/components/schemas/AgentSessionDefaults' }, + defaultThreadId: { type: 'string', nullable: true }, + lastActivity: { type: 'string', format: 'date-time', nullable: true }, + endedAt: { type: 'string', format: 'date-time', nullable: true }, + createdAt: { type: 'string', format: 'date-time', nullable: true }, + updatedAt: { type: 'string', format: 'date-time', nullable: true }, + }, + required: ['id', 'status', 'userId', 'ownerGithubUsername', 'defaults', 'defaultThreadId'], + }, + source: { $ref: '#/components/schemas/AgentSource' }, + sandbox: { $ref: '#/components/schemas/AgentSandbox' }, + }, + required: ['session', 'source', 'sandbox'], + }, + + AgentAdminSessionSummary: { type: 'object', properties: { id: { type: 'string' }, + sessionKind: { + type: 'string', + enum: Object.values(AgentSessionKind), + }, buildUuid: { type: 'string', nullable: true }, baseBuildUuid: { type: 'string', nullable: true }, - buildKind: { $ref: '#/components/schemas/BuildKind' }, + buildKind: { + type: 'string', + enum: Object.values(BuildKind), + nullable: true, + }, userId: { type: 'string' }, ownerGithubUsername: { type: 'string', nullable: true }, - podName: { type: 'string' }, - namespace: { type: 'string' }, - pvcName: { type: 'string' }, + podName: { type: 'string', nullable: true }, + namespace: { type: 'string', nullable: true }, + pvcName: { type: 'string', nullable: true }, model: { type: 'string' }, - status: { type: 'string', enum: ['starting', 'active', 'ended', 'error'] }, + status: { + type: 'string', + enum: ['starting', 'active', 'ended', 'error'], + }, + chatStatus: { + type: 'string', + enum: Object.values(AgentChatStatus), + }, + workspaceStatus: { + type: 'string', + enum: Object.values(AgentWorkspaceStatus), + }, repo: { type: 'string', nullable: true }, branch: { type: 'string', nullable: true }, primaryRepo: { type: 'string', nullable: true }, primaryBranch: { type: 'string', nullable: true }, - services: { type: 'array', items: { type: 'string' } }, - workspaceRepos: { type: 'array', items: { type: 'object', additionalProperties: true } }, - selectedServices: { type: 'array', items: { type: 'object', additionalProperties: true } }, - startupFailure: { type: 'object', additionalProperties: true, nullable: true }, + services: { + type: 'array', + items: { type: 'string' }, + }, + workspaceRepos: { + type: 'array', + items: { type: 'object', additionalProperties: true }, + }, + selectedServices: { + type: 'array', + items: { type: 'object', additionalProperties: true }, + }, + startupFailure: { + type: 'object', + additionalProperties: true, + nullable: true, + }, lastActivity: { type: 'string', format: 'date-time', nullable: true }, endedAt: { type: 'string', format: 'date-time', nullable: true }, threadCount: { type: 'integer' }, @@ -428,10 +818,11 @@ export const openApiSpecificationForV2Api: OAS3Options = { lastRunAt: { type: 'string', format: 'date-time', nullable: true }, createdAt: { type: 'string', format: 'date-time', nullable: true }, updatedAt: { type: 'string', format: 'date-time', nullable: true }, - editorUrl: { type: 'string' }, + editorUrl: { type: 'string', nullable: true }, }, required: [ 'id', + 'sessionKind', 'buildUuid', 'baseBuildUuid', 'buildKind', @@ -442,6 +833,8 @@ export const openApiSpecificationForV2Api: OAS3Options = { 'pvcName', 'model', 'status', + 'chatStatus', + 'workspaceStatus', 'repo', 'branch', 'primaryRepo', @@ -461,38 +854,446 @@ export const openApiSpecificationForV2Api: OAS3Options = { ], }, + AgentRunStatus: { + type: 'string', + enum: [ + 'queued', + 'starting', + 'running', + 'waiting_for_approval', + 'waiting_for_input', + 'completed', + 'failed', + 'cancelled', + ], + }, + AgentRun: { type: 'object', properties: { id: { type: 'string' }, threadId: { type: 'string', nullable: true }, sessionId: { type: 'string', nullable: true }, - status: { - type: 'string', - enum: [ - 'queued', - 'running', - 'waiting_for_approval', - 'waiting_for_input', - 'completed', - 'failed', - 'cancelled', - ], - }, + status: { $ref: '#/components/schemas/AgentRunStatus' }, + requestedHarness: { type: 'string', nullable: true }, + resolvedHarness: { type: 'string', nullable: true }, + requestedProvider: { type: 'string', nullable: true }, + requestedModel: { type: 'string', nullable: true }, + resolvedProvider: { type: 'string', nullable: true }, + resolvedModel: { type: 'string', nullable: true }, provider: { type: 'string' }, model: { type: 'string' }, + sandboxRequirement: { type: 'object', additionalProperties: true }, + sandboxGeneration: { type: 'integer', nullable: true }, queuedAt: { type: 'string', format: 'date-time', nullable: true }, startedAt: { type: 'string', format: 'date-time', nullable: true }, completedAt: { type: 'string', format: 'date-time', nullable: true }, cancelledAt: { type: 'string', format: 'date-time', nullable: true }, usageSummary: { type: 'object', additionalProperties: true }, policySnapshot: { type: 'object', additionalProperties: true }, - streamState: { type: 'object', additionalProperties: true }, - error: { type: 'object', additionalProperties: true, nullable: true }, + error: { + allOf: [{ $ref: '#/components/schemas/AgentRunError' }], + nullable: true, + }, createdAt: { type: 'string', format: 'date-time', nullable: true }, updatedAt: { type: 'string', format: 'date-time', nullable: true }, }, - required: ['id', 'status', 'provider', 'model', 'usageSummary', 'policySnapshot', 'streamState'], + required: [ + 'id', + 'status', + 'requestedHarness', + 'resolvedHarness', + 'requestedProvider', + 'requestedModel', + 'resolvedProvider', + 'resolvedModel', + 'provider', + 'model', + 'sandboxRequirement', + 'sandboxGeneration', + 'usageSummary', + 'policySnapshot', + ], + }, + + CreateAgentThreadRunResponse: { + type: 'object', + properties: { + run: { $ref: '#/components/schemas/AgentRun' }, + message: { $ref: '#/components/schemas/AgentMessage' }, + links: { + type: 'object', + properties: { + events: { type: 'string' }, + eventStream: { type: 'string' }, + pendingActions: { type: 'string' }, + }, + required: ['events', 'eventStream', 'pendingActions'], + additionalProperties: false, + }, + }, + required: ['run', 'message', 'links'], + additionalProperties: false, + }, + + AgentRunError: { + type: 'object', + properties: { + message: { type: 'string' }, + name: { type: 'string', nullable: true }, + code: { type: 'string', nullable: true }, + stack: { type: 'string', nullable: true }, + details: { type: 'object', additionalProperties: true, nullable: true }, + }, + required: ['message'], + additionalProperties: true, + }, + + AgentRunTruncatedValue: { + type: 'object', + properties: { + truncated: { type: 'boolean', enum: [true] }, + originalJsonBytes: { type: 'integer' }, + preview: { type: 'string' }, + }, + required: ['truncated', 'originalJsonBytes', 'preview'], + additionalProperties: false, + }, + + AgentRunMessageCreatedEvent: agentRunEventSchema(['message.created'], { + type: 'object', + properties: { + messageId: { type: 'string' }, + metadata: agentRunEventPayloadMetadata, + }, + required: ['messageId', 'metadata'], + additionalProperties: false, + }), + + AgentRunMessageMetadataEvent: agentRunEventSchema(['message.metadata'], { + type: 'object', + properties: { + metadata: agentRunEventPayloadMetadata, + }, + required: ['metadata'], + additionalProperties: false, + }), + + AgentRunMessagePartEvent: agentRunEventSchema( + ['message.part.started', 'message.delta', 'message.part.completed'], + { + type: 'object', + properties: { + partType: { type: 'string', enum: ['text', 'reasoning'] }, + partId: { type: 'string' }, + delta: { type: 'string' }, + providerMetadata: agentRunEventPayloadMetadata, + }, + required: ['partType', 'partId'], + additionalProperties: false, + } + ), + + AgentRunMessageSourceEvent: agentRunEventSchema(['message.source'], { + type: 'object', + properties: { + sourceType: { type: 'string', enum: ['url', 'document'] }, + sourceId: { type: 'string' }, + url: { type: 'string' }, + mediaType: { type: 'string' }, + title: { type: 'string' }, + filename: { type: 'string' }, + providerMetadata: agentRunEventPayloadMetadata, + }, + required: ['sourceType', 'sourceId'], + additionalProperties: false, + }), + + AgentRunMessageFileEvent: agentRunEventSchema(['message.file'], { + type: 'object', + properties: { + url: { type: 'string' }, + mediaType: { type: 'string' }, + providerMetadata: agentRunEventPayloadMetadata, + }, + required: ['url', 'mediaType'], + additionalProperties: false, + }), + + AgentRunToolInputStartedEvent: agentRunEventSchema(['tool.call.input.started'], { + type: 'object', + properties: { + toolCallId: { type: 'string' }, + toolName: { type: 'string' }, + providerExecuted: { type: 'boolean' }, + providerMetadata: agentRunEventPayloadMetadata, + dynamic: { type: 'boolean' }, + title: { type: 'string' }, + }, + required: ['toolCallId', 'toolName'], + additionalProperties: false, + }), + + AgentRunToolInputDeltaEvent: agentRunEventSchema(['tool.call.input.delta'], { + type: 'object', + properties: { + toolCallId: { type: 'string' }, + inputTextDelta: { type: 'string' }, + }, + required: ['toolCallId', 'inputTextDelta'], + additionalProperties: false, + }), + + AgentRunToolInputEvent: { + oneOf: [ + { $ref: '#/components/schemas/AgentRunToolInputStartedEvent' }, + { $ref: '#/components/schemas/AgentRunToolInputDeltaEvent' }, + ], + discriminator: { + propertyName: 'eventType', + mapping: { + 'tool.call.input.started': '#/components/schemas/AgentRunToolInputStartedEvent', + 'tool.call.input.delta': '#/components/schemas/AgentRunToolInputDeltaEvent', + }, + }, + }, + + AgentRunToolStartedEvent: agentRunEventSchema(['tool.call.started'], { + type: 'object', + properties: { + toolCallId: { type: 'string' }, + toolName: { type: 'string' }, + inputStatus: { type: 'string', enum: ['available', 'error'] }, + input: { nullable: true }, + errorText: { type: 'string', nullable: true }, + providerExecuted: { type: 'boolean' }, + providerMetadata: agentRunEventPayloadMetadata, + dynamic: { type: 'boolean' }, + title: { type: 'string' }, + }, + required: ['toolCallId', 'toolName', 'inputStatus', 'input', 'errorText'], + additionalProperties: false, + }), + + AgentRunToolCompletedEvent: agentRunEventSchema(['tool.call.completed'], { + type: 'object', + properties: { + toolCallId: { type: 'string' }, + output: { nullable: true }, + errorText: { type: 'string', nullable: true }, + status: { type: 'string', enum: ['completed', 'denied', 'failed'] }, + providerExecuted: { type: 'boolean' }, + providerMetadata: agentRunEventPayloadMetadata, + dynamic: { type: 'boolean' }, + preliminary: { type: 'boolean' }, + }, + required: ['toolCallId', 'output', 'errorText', 'status'], + additionalProperties: false, + }), + + AgentRunToolCallEvent: { + oneOf: [ + { $ref: '#/components/schemas/AgentRunToolStartedEvent' }, + { $ref: '#/components/schemas/AgentRunToolCompletedEvent' }, + ], + discriminator: { + propertyName: 'eventType', + mapping: { + 'tool.call.started': '#/components/schemas/AgentRunToolStartedEvent', + 'tool.call.completed': '#/components/schemas/AgentRunToolCompletedEvent', + }, + }, + }, + + AgentRunToolFileChangeEvent: agentRunEventSchema(['tool.file_change'], { + type: 'object', + properties: { + id: { type: 'string' }, + data: { type: 'object', additionalProperties: true }, + transient: { type: 'boolean' }, + }, + required: ['data'], + additionalProperties: false, + }), + + AgentRunApprovalRequestedEvent: agentRunEventSchema(['approval.requested'], { + type: 'object', + properties: { + actionId: { type: 'string' }, + approvalId: { type: 'string' }, + toolCallId: { type: 'string' }, + }, + required: ['approvalId', 'toolCallId'], + additionalProperties: false, + }), + + AgentRunApprovalResolvedEvent: agentRunEventSchema(['approval.resolved'], { + type: 'object', + properties: { + actionId: { type: 'string' }, + approvalId: { type: 'string' }, + toolCallId: { type: 'string', nullable: true }, + approved: { type: 'boolean' }, + reason: { type: 'string', nullable: true }, + }, + required: ['actionId', 'approvalId', 'toolCallId', 'approved', 'reason'], + additionalProperties: false, + }), + + AgentRunApprovalRespondedEvent: agentRunEventSchema(['approval.responded'], { + type: 'object', + properties: { + actionId: { type: 'string' }, + approvalId: { type: 'string' }, + toolCallId: { type: 'string', nullable: true }, + approved: { type: 'boolean' }, + reason: { type: 'string', nullable: true }, + }, + required: ['actionId', 'approvalId', 'toolCallId', 'approved', 'reason'], + additionalProperties: false, + }), + + AgentRunApprovalEvent: { + oneOf: [ + { $ref: '#/components/schemas/AgentRunApprovalRequestedEvent' }, + { $ref: '#/components/schemas/AgentRunApprovalResolvedEvent' }, + { $ref: '#/components/schemas/AgentRunApprovalRespondedEvent' }, + ], + discriminator: { + propertyName: 'eventType', + mapping: { + 'approval.requested': '#/components/schemas/AgentRunApprovalRequestedEvent', + 'approval.resolved': '#/components/schemas/AgentRunApprovalResolvedEvent', + 'approval.responded': '#/components/schemas/AgentRunApprovalRespondedEvent', + }, + }, + }, + + AgentRunStatusEvent: agentRunEventSchema( + [ + 'run.queued', + 'run.started', + 'run.waiting_for_approval', + 'run.completed', + 'run.failed', + 'run.cancelled', + 'run.updated', + ], + { + type: 'object', + properties: { + threadId: { type: 'string' }, + sessionId: { type: 'string' }, + status: { $ref: '#/components/schemas/AgentRunStatus' }, + error: { + allOf: [{ $ref: '#/components/schemas/AgentRunError' }], + nullable: true, + }, + usageSummary: { type: 'object', additionalProperties: true }, + }, + additionalProperties: false, + } + ), + + AgentRunStepEvent: agentRunEventSchema(['run.step.started', 'run.step.completed'], { + type: 'object', + properties: {}, + additionalProperties: false, + }), + + AgentRunFinishedEvent: agentRunEventSchema(['run.finished'], { + type: 'object', + properties: { + finishReason: { type: 'string' }, + metadata: agentRunEventPayloadMetadata, + }, + required: ['finishReason', 'metadata'], + additionalProperties: false, + }), + + AgentRunErrorEvent: agentRunEventSchema(['run.error'], { + type: 'object', + properties: { + errorText: { type: 'string' }, + }, + required: ['errorText'], + additionalProperties: false, + }), + + AgentRunAbortedEvent: agentRunEventSchema(['run.aborted'], { + type: 'object', + properties: { + reason: { type: 'string' }, + }, + required: ['reason'], + additionalProperties: false, + }), + + AgentRunEvent: { + oneOf: [ + { $ref: '#/components/schemas/AgentRunMessageCreatedEvent' }, + { $ref: '#/components/schemas/AgentRunMessageMetadataEvent' }, + { $ref: '#/components/schemas/AgentRunMessagePartEvent' }, + { $ref: '#/components/schemas/AgentRunMessageSourceEvent' }, + { $ref: '#/components/schemas/AgentRunMessageFileEvent' }, + { $ref: '#/components/schemas/AgentRunToolInputStartedEvent' }, + { $ref: '#/components/schemas/AgentRunToolInputDeltaEvent' }, + { $ref: '#/components/schemas/AgentRunToolStartedEvent' }, + { $ref: '#/components/schemas/AgentRunToolCompletedEvent' }, + { $ref: '#/components/schemas/AgentRunToolFileChangeEvent' }, + { $ref: '#/components/schemas/AgentRunApprovalRequestedEvent' }, + { $ref: '#/components/schemas/AgentRunApprovalResolvedEvent' }, + { $ref: '#/components/schemas/AgentRunApprovalRespondedEvent' }, + { $ref: '#/components/schemas/AgentRunStatusEvent' }, + { $ref: '#/components/schemas/AgentRunStepEvent' }, + { $ref: '#/components/schemas/AgentRunFinishedEvent' }, + { $ref: '#/components/schemas/AgentRunErrorEvent' }, + { $ref: '#/components/schemas/AgentRunAbortedEvent' }, + ], + discriminator: { + propertyName: 'eventType', + mapping: agentRunEventDiscriminatorMapping, + }, + example: { + id: 'event-1', + runId: 'run-1', + threadId: 'thread-1', + sessionId: 'session-1', + sequence: 3, + eventType: 'message.delta', + version: 1, + payload: { + partType: 'text', + partId: 'text-1', + delta: 'Hello', + }, + createdAt: '2026-04-25T00:00:02.000Z', + updatedAt: '2026-04-25T00:00:02.000Z', + }, + }, + + AgentPendingActionArgumentSummary: { + type: 'object', + properties: { + name: { type: 'string' }, + value: { type: 'string' }, + }, + required: ['name', 'value'], + additionalProperties: false, + }, + + AgentPendingActionFileChangePreview: { + type: 'object', + properties: { + path: { type: 'string' }, + action: { type: 'string' }, + summary: { type: 'string' }, + additions: { type: 'integer', nullable: true }, + deletions: { type: 'integer', nullable: true }, + truncated: { type: 'boolean' }, + }, + required: ['path', 'action', 'summary', 'additions', 'deletions', 'truncated'], + additionalProperties: false, }, AgentPendingAction: { @@ -503,16 +1304,67 @@ export const openApiSpecificationForV2Api: OAS3Options = { runId: { type: 'string', nullable: true }, kind: { type: 'string' }, status: { type: 'string', enum: ['pending', 'approved', 'denied'] }, - capabilityKey: { type: 'string' }, title: { type: 'string' }, description: { type: 'string' }, - payload: { type: 'object', additionalProperties: true }, - resolution: { type: 'object', additionalProperties: true, nullable: true }, - resolvedAt: { type: 'string', format: 'date-time', nullable: true }, - createdAt: { type: 'string', format: 'date-time', nullable: true }, - updatedAt: { type: 'string', format: 'date-time', nullable: true }, + requestedAt: { type: 'string', format: 'date-time', nullable: true }, + expiresAt: { type: 'string', format: 'date-time', nullable: true }, + toolName: { type: 'string', nullable: true }, + argumentsSummary: { + type: 'array', + items: { $ref: '#/components/schemas/AgentPendingActionArgumentSummary' }, + }, + commandPreview: { type: 'string', nullable: true }, + fileChangePreview: { + type: 'array', + items: { $ref: '#/components/schemas/AgentPendingActionFileChangePreview' }, + }, + riskLabels: { + type: 'array', + items: { type: 'string' }, + }, + }, + required: [ + 'id', + 'kind', + 'status', + 'threadId', + 'runId', + 'title', + 'description', + 'requestedAt', + 'expiresAt', + 'toolName', + 'argumentsSummary', + 'commandPreview', + 'fileChangePreview', + 'riskLabels', + ], + additionalProperties: false, + example: { + id: 'action-1', + threadId: 'thread-1', + runId: 'run-1', + kind: 'tool_approval', + status: 'pending', + title: 'Approve workspace edit', + description: 'A workspace edit requires approval.', + requestedAt: '2026-04-25T00:00:03.000Z', + expiresAt: null, + toolName: 'mcp__sandbox__workspace_edit_file', + argumentsSummary: [{ name: 'path', value: 'sample-file.txt' }], + commandPreview: null, + fileChangePreview: [ + { + path: 'sample-file.txt', + action: 'edited', + summary: 'edited sample-file.txt', + additions: 1, + deletions: 0, + truncated: false, + }, + ], + riskLabels: ['Workspace write'], }, - required: ['id', 'kind', 'status', 'capabilityKey', 'title', 'description', 'payload'], }, AgentToolExecution: { @@ -525,6 +1377,7 @@ export const openApiSpecificationForV2Api: OAS3Options = { source: { type: 'string' }, serverSlug: { type: 'string', nullable: true }, toolName: { type: 'string' }, + toolCallId: { type: 'string', nullable: true }, args: { type: 'object', additionalProperties: true }, result: { type: 'object', additionalProperties: true, nullable: true }, status: { type: 'string', enum: ['queued', 'running', 'completed', 'failed', 'cancelled'] }, @@ -544,6 +1397,7 @@ export const openApiSpecificationForV2Api: OAS3Options = { 'source', 'serverSlug', 'toolName', + 'toolCallId', 'args', 'result', 'status', @@ -555,6 +1409,27 @@ export const openApiSpecificationForV2Api: OAS3Options = { 'createdAt', 'updatedAt', ], + additionalProperties: false, + example: { + id: 'tool-execution-1', + threadId: 'thread-1', + runId: 'run-1', + pendingActionId: 'action-1', + source: 'mcp', + serverSlug: 'sandbox', + toolName: 'workspace.edit_file', + toolCallId: 'tool-call-1', + args: { path: 'sample-file.txt' }, + result: null, + status: 'completed', + safetyLevel: null, + approved: true, + startedAt: '2026-04-25T00:00:04.000Z', + completedAt: '2026-04-25T00:00:05.000Z', + durationMs: 1000, + createdAt: '2026-04-25T00:00:04.000Z', + updatedAt: '2026-04-25T00:00:05.000Z', + }, }, AgentAdminThreadSummary: { @@ -576,7 +1451,7 @@ export const openApiSpecificationForV2Api: OAS3Options = { AgentAdminSessionDetail: { type: 'object', properties: { - session: { $ref: '#/components/schemas/AgentSessionSummary' }, + session: { $ref: '#/components/schemas/AgentAdminSessionSummary' }, threads: { type: 'array', items: { $ref: '#/components/schemas/AgentAdminThreadSummary' }, @@ -588,16 +1463,20 @@ export const openApiSpecificationForV2Api: OAS3Options = { AgentAdminThreadConversation: { type: 'object', properties: { - session: { $ref: '#/components/schemas/AgentSessionSummary' }, + session: { $ref: '#/components/schemas/AgentAdminSessionSummary' }, thread: { $ref: '#/components/schemas/AgentAdminThreadSummary' }, messages: { type: 'array', - items: { $ref: '#/components/schemas/AgentUIMessage' }, + items: { $ref: '#/components/schemas/AgentMessage' }, }, runs: { type: 'array', items: { $ref: '#/components/schemas/AgentRun' }, }, + events: { + type: 'array', + items: { $ref: '#/components/schemas/AgentRunEvent' }, + }, pendingActions: { type: 'array', items: { $ref: '#/components/schemas/AgentPendingAction' }, @@ -607,7 +1486,7 @@ export const openApiSpecificationForV2Api: OAS3Options = { items: { $ref: '#/components/schemas/AgentToolExecution' }, }, }, - required: ['session', 'thread', 'messages', 'runs', 'pendingActions', 'toolExecutions'], + required: ['session', 'thread', 'messages', 'runs', 'events', 'pendingActions', 'toolExecutions'], }, AgentAdminMcpServerCoverage: { @@ -1207,6 +2086,21 @@ export const openApiSpecificationForV2Api: OAS3Options = { enum: Object.values(BuildKind), }, + AgentSessionKind: { + type: 'string', + enum: Object.values(AgentSessionKind), + }, + + AgentChatStatus: { + type: 'string', + enum: Object.values(AgentChatStatus), + }, + + AgentWorkspaceStatus: { + type: 'string', + enum: Object.values(AgentWorkspaceStatus), + }, + /** * @description The main Build object. */ @@ -1708,6 +2602,46 @@ export const openApiSpecificationForV2Api: OAS3Options = { additionalProperties: false, }, + AgentSessionWorkspaceStorageSettings: { + type: 'object', + properties: { + defaultSize: { type: 'string', minLength: 1, maxLength: 64 }, + allowedSizes: { + type: 'array', + items: { type: 'string', minLength: 1, maxLength: 64 }, + uniqueItems: true, + }, + allowClientOverride: { type: 'boolean' }, + accessMode: { type: 'string', enum: ['ReadWriteOnce', 'ReadWriteMany'] }, + }, + additionalProperties: false, + }, + + AgentSessionCleanupSettings: { + type: 'object', + properties: { + activeIdleSuspendMs: { type: 'integer', minimum: 1 }, + startingTimeoutMs: { type: 'integer', minimum: 1 }, + hibernatedRetentionMs: { type: 'integer', minimum: 1 }, + intervalMs: { type: 'integer', minimum: 1 }, + redisTtlSeconds: { type: 'integer', minimum: 1 }, + }, + additionalProperties: false, + }, + + AgentSessionDurabilitySettings: { + type: 'object', + properties: { + runExecutionLeaseMs: { type: 'integer', minimum: 1 }, + queuedRunDispatchStaleMs: { type: 'integer', minimum: 1 }, + dispatchRecoveryLimit: { type: 'integer', minimum: 1 }, + maxDurablePayloadBytes: { type: 'integer', minimum: 1 }, + payloadPreviewBytes: { type: 'integer', minimum: 1 }, + fileChangePreviewChars: { type: 'integer', minimum: 1 }, + }, + additionalProperties: false, + }, + AgentSessionRuntimeSettings: { type: 'object', properties: { @@ -1747,6 +2681,9 @@ export const openApiSpecificationForV2Api: OAS3Options = { }, additionalProperties: false, }, + workspaceStorage: { $ref: '#/components/schemas/AgentSessionWorkspaceStorageSettings' }, + cleanup: { $ref: '#/components/schemas/AgentSessionCleanupSettings' }, + durability: { $ref: '#/components/schemas/AgentSessionDurabilitySettings' }, }, additionalProperties: false, }, @@ -2983,7 +3920,7 @@ export const openApiSpecificationForV2Api: OAS3Options = { metadata: { $ref: '#/components/schemas/ResponseMetadata' }, data: { type: 'array', - items: { $ref: '#/components/schemas/AgentSessionSummary' }, + items: { $ref: '#/components/schemas/AgentAdminSessionSummary' }, }, }, required: ['data'], diff --git a/sysops/workspace-gateway/index.mjs b/sysops/workspace-gateway/index.mjs index e172362c..5ed054d4 100644 --- a/sysops/workspace-gateway/index.mjs +++ b/sysops/workspace-gateway/index.mjs @@ -28,6 +28,7 @@ const MAX_GREP_RESULTS = parsePositiveInt(process.env.LIFECYCLE_SANDBOX_MAX_GREP const MAX_COMMAND_OUTPUT_CHARS = parsePositiveInt(process.env.LIFECYCLE_SANDBOX_MAX_COMMAND_OUTPUT_CHARS, 24_000); const MAX_FILE_CHANGE_PREVIEW_CHARS = parsePositiveInt(process.env.LIFECYCLE_SANDBOX_MAX_FILE_CHANGE_PREVIEW_CHARS, 4000); const MAX_FILE_CHANGE_DIFF_CHARS = parsePositiveInt(process.env.LIFECYCLE_SANDBOX_MAX_FILE_CHANGE_DIFF_CHARS, 16_000); +const MAX_EXEC_FILE_CHANGES = parsePositiveInt(process.env.LIFECYCLE_SANDBOX_MAX_EXEC_FILE_CHANGES, 50); const STATE_FILE = process.env.LIFECYCLE_SANDBOX_STATE_FILE || ''; const PORTS_FILE = process.env.LIFECYCLE_SANDBOX_PORTS_FILE || ''; const PROCESSES_FILE = process.env.LIFECYCLE_SANDBOX_PROCESSES_FILE || ''; @@ -316,6 +317,178 @@ async function buildFileChangeArtifact({ }; } +async function findNearestGitRoot(startPath) { + let current = resolve(startPath); + + while (isWithinWorkspace(current)) { + if (await fileExists(resolve(current, '.git'))) { + return current; + } + + if (current === WORKSPACE_ROOT) { + break; + } + + current = dirname(current); + } + + return null; +} + +function snapshotHasPathUnderRoot(snapshot, workspaceRootPath) { + return [...snapshot.keys()].some( + (path) => + path === workspaceRootPath || path.startsWith(`${workspaceRootPath}/`) + ); +} + +function normalizeGitStatusPath(value) { + return value.split(sep).join('/').replace(/^\.\/+/, ''); +} + +function gitStatusPathCoversFile(statusPath, repoRelativePath) { + const normalizedStatusPath = normalizeGitStatusPath(statusPath); + const normalizedRepoRelativePath = normalizeGitStatusPath(repoRelativePath); + + return ( + normalizedStatusPath === normalizedRepoRelativePath || + (normalizedStatusPath.endsWith('/') && normalizedRepoRelativePath.startsWith(normalizedStatusPath)) + ); +} + +function parseGitStatusPaths(stdout) { + const entries = stdout.split('\0').filter(Boolean); + const paths = []; + + for (let index = 0; index < entries.length; index += 1) { + const entry = entries[index]; + const status = entry.slice(0, 2); + const path = entry.slice(3); + + if (path) { + paths.push(path); + } + + if (status.includes('R') || status.includes('C')) { + index += 1; + } + } + + return paths; +} + +async function readGitStatusPaths(repoRoot) { + try { + const result = await execFile('/usr/bin/git', [ + '-C', + repoRoot, + 'status', + '--porcelain=v1', + '-z', + '--untracked-files=normal', + ]); + + return parseGitStatusPaths(result.stdout); + } catch { + return null; + } +} + +async function readGitHeadText(repoRoot, repoRelativePath) { + try { + const result = await execFile( + '/usr/bin/git', + ['-C', repoRoot, 'show', `HEAD:${repoRelativePath}`], + { + maxBuffer: 10 * 1024 * 1024, + } + ); + + return result.stdout; + } catch { + return null; + } +} + +async function getNewGitRepoInfoForPath(path, beforeSnapshot, cache) { + const absolutePath = resolveWorkspacePath(path); + const repoRoot = await findNearestGitRoot(dirname(absolutePath)); + + if (!repoRoot) { + return null; + } + + const workspaceRootPath = toWorkspaceRelativePath(repoRoot); + const existedBefore = workspaceRootPath + ? snapshotHasPathUnderRoot(beforeSnapshot, workspaceRootPath) + : beforeSnapshot.size > 0; + if (existedBefore) { + return null; + } + + if (!cache.has(repoRoot)) { + cache.set( + repoRoot, + readGitStatusPaths(repoRoot).then((statusPaths) => ({ + repoRoot, + statusPaths, + })) + ); + } + + return cache.get(repoRoot); +} + +async function normalizeSnapshotChangeCandidate({ + path, + beforeSnapshot, + afterSnapshot, + newGitRepoInfoCache, +}) { + const hadBefore = beforeSnapshot.has(path); + const hasAfter = afterSnapshot.has(path); + + if (!hadBefore && hasAfter) { + const newRepoInfo = await getNewGitRepoInfoForPath( + path, + beforeSnapshot, + newGitRepoInfoCache + ); + if (newRepoInfo?.statusPaths) { + const repoRelativePath = normalizeGitStatusPath( + relative(newRepoInfo.repoRoot, resolveWorkspacePath(path)) + ); + const repoReportsPath = newRepoInfo.statusPaths.some((statusPath) => + gitStatusPathCoversFile(statusPath, repoRelativePath) + ); + + if (!repoReportsPath) { + return null; + } + + const baselineText = await readGitHeadText( + newRepoInfo.repoRoot, + repoRelativePath + ); + if (baselineText !== null) { + return { + path, + kind: 'edited', + before: baselineText, + after: afterSnapshot.get(path) || '', + }; + } + } + } + + return { + path, + kind: !hadBefore ? 'created' : !hasAfter ? 'deleted' : 'edited', + before: beforeSnapshot.get(path) || '', + after: afterSnapshot.get(path) || '', + }; +} + function isWithinWorkspace(candidate) { const normalized = resolve(candidate); return normalized === WORKSPACE_ROOT || normalized.startsWith(`${WORKSPACE_ROOT}${sep}`); @@ -518,26 +691,129 @@ async function editWorkspaceFile({ path, oldText, newText, replaceAll = false }) }; } -async function runWorkspaceCommand({ command, cwd = '.', timeoutMs = 30000 }) { - const resolvedCwd = resolveWorkspacePath(cwd); - const { stdout, stderr } = await execFile('/bin/bash', ['-lc', command], { - cwd: resolvedCwd, - timeout: timeoutMs, - maxBuffer: 10 * 1024 * 1024, - env: { - ...process.env, - HOME: process.env.HOME || WORKSPACE_ROOT, - }, - }); +async function snapshotWorkspaceTextFiles() { + const snapshot = new Map(); + let files = []; + try { + files = await collectFilesUnderPath('.'); + } catch { + return snapshot; + } + + for (const absolutePath of files) { + try { + if (isReservedWorkspacePath(absolutePath)) { + continue; + } + const text = await readFile(absolutePath, 'utf8'); + snapshot.set(toWorkspaceRelativePath(absolutePath), text); + } catch { + // Ignore files that disappear or cannot be decoded while snapshotting. + } + } + + return snapshot; +} + +async function buildFileChangesFromSnapshots(beforeSnapshot, afterSnapshot) { + const changedPaths = [...new Set([...beforeSnapshot.keys(), ...afterSnapshot.keys()])] + .filter((path) => beforeSnapshot.get(path) !== afterSnapshot.get(path)) + .sort(); + const newGitRepoInfoCache = new Map(); + const candidates = []; + + for (const path of changedPaths) { + const candidate = await normalizeSnapshotChangeCandidate({ + path, + beforeSnapshot, + afterSnapshot, + newGitRepoInfoCache, + }); + + if (candidate) { + candidates.push(candidate); + } + } + + const limitedCandidates = candidates.slice(0, MAX_EXEC_FILE_CHANGES); + const fileChanges = []; + + for (const candidate of limitedCandidates) { + fileChanges.push( + await buildFileChangeArtifact({ + path: candidate.path, + kind: candidate.kind, + before: candidate.before, + after: candidate.after, + }) + ); + } return { - cwd: toWorkspaceRelativePath(resolvedCwd), - stdout: truncateText(stdout), - stderr: truncateText(stderr), - success: true, + fileChanges, + fileChangesTruncated: candidates.length > limitedCandidates.length, }; } +function getCommandErrorFileChanges(error) { + return isRecord(error) && Array.isArray(error.fileChanges) ? error.fileChanges : []; +} + +function getCommandErrorFileChangesTruncated(error) { + return isRecord(error) && error.fileChangesTruncated === true; +} + +function commandErrorText(message, error) { + const fileChanges = getCommandErrorFileChanges(error); + const fileChangesTruncated = getCommandErrorFileChangesTruncated(error); + + return textResult({ + ok: false, + error: message, + details: error instanceof Error ? error.message : String(error), + ...(fileChanges.length > 0 ? { fileChanges } : {}), + ...(fileChangesTruncated ? { fileChangesTruncated } : {}), + }); +} + +async function runWorkspaceCommand({ command, cwd = '.', timeoutMs = 30000, captureFileChanges = false }) { + const resolvedCwd = resolveWorkspacePath(cwd); + const beforeSnapshot = captureFileChanges ? await snapshotWorkspaceTextFiles() : null; + + try { + const { stdout, stderr } = await execFile('/bin/bash', ['-lc', command], { + cwd: resolvedCwd, + timeout: timeoutMs, + maxBuffer: 10 * 1024 * 1024, + env: { + ...process.env, + HOME: process.env.HOME || WORKSPACE_ROOT, + }, + }); + + const changes = beforeSnapshot + ? await buildFileChangesFromSnapshots(beforeSnapshot, await snapshotWorkspaceTextFiles()) + : { fileChanges: [], fileChangesTruncated: false }; + + return { + cwd: toWorkspaceRelativePath(resolvedCwd), + stdout: truncateText(stdout), + stderr: truncateText(stderr), + success: true, + ...(changes.fileChanges.length > 0 ? { fileChanges: changes.fileChanges } : {}), + ...(changes.fileChangesTruncated ? { fileChangesTruncated: true } : {}), + }; + } catch (error) { + if (beforeSnapshot && isRecord(error)) { + const changes = await buildFileChangesFromSnapshots(beforeSnapshot, await snapshotWorkspaceTextFiles()); + error.fileChanges = changes.fileChanges; + error.fileChangesTruncated = changes.fileChangesTruncated; + } + + throw error; + } +} + async function walkFiles(rootDir, relativePrefix = '', results = [], limit = MAX_LIST_RESULTS) { if (results.length >= limit) { return results; @@ -1106,10 +1382,11 @@ function buildServer() { command: z.string().min(1).describe('Command to run with bash -lc'), cwd: z.string().optional().describe('Working directory relative to the workspace'), timeoutMs: z.number().int().positive().max(120000).optional().describe('Command timeout in milliseconds'), + captureFileChanges: z.boolean().optional().describe('Internal Lifecycle flag for file-change capture'), }, annotations: { destructiveHint: true, openWorldHint: true }, }, - async ({ command, cwd, timeoutMs }) => { + async ({ command, cwd, timeoutMs, captureFileChanges }) => { try { return textResult({ ok: true, @@ -1118,10 +1395,11 @@ function buildServer() { command, cwd: cwd || '.', timeoutMs: timeoutMs || 30000, + captureFileChanges: captureFileChanges === true, }), }); } catch (error) { - return errorText('Unable to run workspace command', error instanceof Error ? error.message : String(error)); + return commandErrorText('Unable to run workspace command', error); } } );