diff --git a/Tiltfile b/Tiltfile index 930035c2..50f24377 100644 --- a/Tiltfile +++ b/Tiltfile @@ -239,7 +239,7 @@ if use_local_keycloak_charts: local_resource( 'lifecycle-keycloak-chart-deps', - cmd='helm dependency build {}'.format(lifecycle_keycloak_chart), + cmd='helm dependency update {}'.format(lifecycle_keycloak_chart), deps=[ '{}/Chart.yaml'.format(lifecycle_keycloak_chart), '{}/values.yaml'.format(lifecycle_keycloak_chart), diff --git a/src/app/api/v2/ai/admin/agent/capabilities/route.test.ts b/src/app/api/v2/ai/admin/agent/capabilities/route.test.ts new file mode 100644 index 00000000..74fc0b5e --- /dev/null +++ b/src/app/api/v2/ai/admin/agent/capabilities/route.test.ts @@ -0,0 +1,313 @@ +/** + * 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 { AgentRuntimeConfigValidationError } from 'server/lib/validation/agentRuntimeConfigValidator'; + +const mockGetUser = jest.fn(); +const mockListCapabilityInventory = jest.fn(); +const mockGetGlobalConfig = jest.fn(); +const mockGetRepoConfig = jest.fn(); +const mockGetEffectiveConfig = jest.fn(); +const mockUpdateGlobalCapabilityPolicy = jest.fn(); +const mockUpdateRepoCapabilityPolicy = jest.fn(); + +jest.mock('server/lib/get-user', () => ({ + getUser: (...args: unknown[]) => mockGetUser(...args), +})); + +jest.mock('server/services/agentSessionConfig', () => ({ + __esModule: true, + default: { + getInstance: jest.fn(() => ({ + listCapabilityInventory: (...args: unknown[]) => mockListCapabilityInventory(...args), + })), + }, +})); + +jest.mock('server/services/agentRuntime/config/agentRuntimeConfig', () => ({ + __esModule: true, + default: { + getInstance: jest.fn(() => ({ + getGlobalConfig: (...args: unknown[]) => mockGetGlobalConfig(...args), + getRepoConfig: (...args: unknown[]) => mockGetRepoConfig(...args), + getEffectiveConfig: (...args: unknown[]) => mockGetEffectiveConfig(...args), + updateGlobalCapabilityPolicy: (...args: unknown[]) => mockUpdateGlobalCapabilityPolicy(...args), + updateRepoCapabilityPolicy: (...args: unknown[]) => mockUpdateRepoCapabilityPolicy(...args), + })), + }, +})); + +import { GET, PUT } from './route'; + +function makeRequest(url: string, body?: unknown): NextRequest { + return { + headers: new Headers([['x-request-id', 'req-test']]), + nextUrl: new URL(url), + json: jest.fn().mockResolvedValue(body), + } as unknown as NextRequest; +} + +const capabilityRow = { + capabilityId: 'workspace_shell', + label: 'Workspace shell', + description: 'Run shell commands in the workspace.', + category: 'workspace', + defaultAvailability: 'all_users', + effectiveAvailability: 'admin_only', + approvalMode: 'require_approval', + runtimeCapabilityKey: 'shell_exec', + userSelectable: true, + toolCount: 1, + resourceCount: 0, + resourceGrants: [], + tools: [], +}; + +describe('/api/v2/ai/admin/agent/capabilities', () => { + const originalEnableAuth = process.env.ENABLE_AUTH; + + beforeEach(() => { + jest.clearAllMocks(); + process.env.ENABLE_AUTH = 'true'; + mockGetUser.mockReturnValue({ + sub: 'sample-admin', + realm_access: { + roles: ['admin'], + }, + }); + mockListCapabilityInventory.mockResolvedValue([capabilityRow]); + mockGetGlobalConfig.mockResolvedValue({ + capabilityPolicy: { + availability: { + workspace_shell: 'admin_only', + }, + }, + }); + mockGetRepoConfig.mockResolvedValue(null); + mockGetEffectiveConfig.mockResolvedValue({ + capabilityPolicy: { + availability: { + workspace_shell: 'admin_only', + }, + }, + }); + mockUpdateGlobalCapabilityPolicy.mockResolvedValue({}); + mockUpdateRepoCapabilityPolicy.mockResolvedValue({}); + }); + + afterEach(() => { + if (originalEnableAuth === undefined) { + delete process.env.ENABLE_AUTH; + } else { + process.env.ENABLE_AUTH = originalEnableAuth; + } + }); + + it('rejects non-admin users before loading inventory', async () => { + mockGetUser.mockReturnValue({ + sub: 'sample-user', + realm_access: { + roles: ['user'], + }, + }); + + const response = await GET(makeRequest('http://localhost/api/v2/ai/admin/agent/capabilities')); + const body = await response.json(); + + expect(response.status).toBe(403); + expect(body.error.message).toBe('Forbidden: insufficient permissions'); + expect(mockListCapabilityInventory).not.toHaveBeenCalled(); + }); + + it('returns global capability inventory and effective policy', async () => { + const response = await GET(makeRequest('http://localhost/api/v2/ai/admin/agent/capabilities')); + const body = await response.json(); + + expect(response.status).toBe(200); + expect(mockListCapabilityInventory).toHaveBeenCalledWith('global'); + expect(body.data).toEqual( + expect.objectContaining({ + scope: 'global', + scopeType: 'global', + capabilityPolicy: { + availability: { + workspace_shell: 'admin_only', + }, + }, + effectiveCapabilityPolicy: { + availability: { + workspace_shell: 'admin_only', + }, + }, + capabilities: [capabilityRow], + }) + ); + }); + + it('returns repo capability inventory with inherited policy metadata', async () => { + mockGetRepoConfig.mockResolvedValue({ + capabilityPolicy: { + availability: { + workspace_shell: 'all_users', + }, + }, + }); + mockGetEffectiveConfig.mockResolvedValue({ + capabilityPolicy: { + availability: { + workspace_shell: 'all_users', + }, + }, + }); + + const response = await GET( + makeRequest('http://localhost/api/v2/ai/admin/agent/capabilities?scope=Example-Org/Example-Repo') + ); + const body = await response.json(); + + expect(response.status).toBe(200); + expect(mockGetRepoConfig).toHaveBeenCalledWith('example-org/example-repo'); + expect(mockListCapabilityInventory).toHaveBeenCalledWith('example-org/example-repo'); + expect(body.data).toEqual( + expect.objectContaining({ + scope: 'example-org/example-repo', + scopeType: 'repo', + repoFullName: 'example-org/example-repo', + capabilityPolicy: { + availability: { + workspace_shell: 'all_users', + }, + }, + inheritedCapabilityPolicy: { + availability: { + workspace_shell: 'admin_only', + }, + }, + effectiveCapabilityPolicy: { + availability: { + workspace_shell: 'all_users', + }, + }, + }) + ); + }); + + it('updates global capability policy and returns refreshed inventory', async () => { + const body = { + capabilityPolicy: { + availability: { + workspace_shell: 'disabled', + }, + }, + }; + + const response = await PUT(makeRequest('http://localhost/api/v2/ai/admin/agent/capabilities', body)); + + expect(response.status).toBe(200); + expect(mockUpdateGlobalCapabilityPolicy).toHaveBeenCalledWith(body.capabilityPolicy); + expect(mockUpdateRepoCapabilityPolicy).not.toHaveBeenCalled(); + expect(mockListCapabilityInventory).toHaveBeenCalledWith('global'); + }); + + it('updates repo capability policy and returns refreshed inventory', async () => { + const body = { + capabilityPolicy: { + availability: { + workspace_shell: 'all_users', + }, + }, + }; + + const response = await PUT( + makeRequest('http://localhost/api/v2/ai/admin/agent/capabilities?scope=Example-Org/Example-Repo', body) + ); + + expect(response.status).toBe(200); + expect(mockUpdateRepoCapabilityPolicy).toHaveBeenCalledWith('example-org/example-repo', body.capabilityPolicy); + expect(mockUpdateGlobalCapabilityPolicy).not.toHaveBeenCalled(); + expect(mockListCapabilityInventory).toHaveBeenCalledWith('example-org/example-repo'); + }); + + it('rejects malformed repo scope', async () => { + const response = await GET(makeRequest('http://localhost/api/v2/ai/admin/agent/capabilities?scope=repo')); + const body = await response.json(); + + expect(response.status).toBe(400); + expect(body.error.message).toContain('Repo capability scope'); + expect(mockListCapabilityInventory).not.toHaveBeenCalled(); + }); + + it('rejects invalid capability ids from service validation', async () => { + mockUpdateGlobalCapabilityPolicy.mockRejectedValueOnce( + new AgentRuntimeConfigValidationError('Unknown capability id "sample_unknown".') + ); + + const response = await PUT( + makeRequest('http://localhost/api/v2/ai/admin/agent/capabilities', { + capabilityPolicy: { + availability: { + sample_unknown: 'disabled', + }, + }, + }) + ); + const body = await response.json(); + + expect(response.status).toBe(400); + expect(body.error.message).toBe('Unknown capability id "sample_unknown".'); + }); + + it('rejects invalid capability availability values from service validation', async () => { + mockUpdateGlobalCapabilityPolicy.mockRejectedValueOnce( + new AgentRuntimeConfigValidationError('Capability "workspace_shell" has invalid availability "sometimes".') + ); + + const response = await PUT( + makeRequest('http://localhost/api/v2/ai/admin/agent/capabilities', { + capabilityPolicy: { + availability: { + workspace_shell: 'sometimes', + }, + }, + }) + ); + const body = await response.json(); + + expect(response.status).toBe(400); + expect(body.error.message).toBe('Capability "workspace_shell" has invalid availability "sometimes".'); + }); + + it.each([ + ['null', null], + ['array', []], + ['scalar', 'disabled'], + ])('rejects malformed capability policy availability: %s', async (_label, availability) => { + const response = await PUT( + makeRequest('http://localhost/api/v2/ai/admin/agent/capabilities', { + capabilityPolicy: { + availability, + }, + }) + ); + const body = await response.json(); + + expect(response.status).toBe(400); + expect(body.error.message).toBe('capabilityPolicy.availability must be an object.'); + expect(mockUpdateGlobalCapabilityPolicy).not.toHaveBeenCalled(); + expect(mockUpdateRepoCapabilityPolicy).not.toHaveBeenCalled(); + }); +}); diff --git a/src/app/api/v2/ai/admin/agent/capabilities/route.ts b/src/app/api/v2/ai/admin/agent/capabilities/route.ts new file mode 100644 index 00000000..828f758f --- /dev/null +++ b/src/app/api/v2/ai/admin/agent/capabilities/route.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. + */ + +import { NextRequest } from 'next/server'; +import { createApiHandler } from 'server/lib/createApiHandler'; +import { errorResponse, successResponse } from 'server/lib/response'; +import { normalizeRepoFullName } from 'server/lib/normalizeRepoFullName'; +import { AgentRuntimeConfigValidationError } from 'server/lib/validation/agentRuntimeConfigValidator'; +import AgentRuntimeConfigService from 'server/services/agentRuntime/config/agentRuntimeConfig'; +import AgentSessionConfigService from 'server/services/agentSessionConfig'; +import type { CapabilityPolicyConfig } from 'server/services/types/agentRuntimeConfig'; + +export const dynamic = 'force-dynamic'; + +type ParsedScope = + | { + scopeType: 'global'; + scope: 'global'; + repoFullName?: undefined; + } + | { + scopeType: 'repo'; + scope: string; + repoFullName: string; + }; + +function isRecord(value: unknown): value is Record { + return Boolean(value) && typeof value === 'object' && !Array.isArray(value); +} + +function parseScope(req: NextRequest): ParsedScope | Error { + const rawScope = req.nextUrl.searchParams.get('scope')?.trim() || 'global'; + if (rawScope === 'global') { + return { scopeType: 'global', scope: 'global' }; + } + + const repoFullName = normalizeRepoFullName(rawScope); + if (!/^[^/\s]+\/[^/\s]+$/.test(repoFullName)) { + return new Error( + 'Repo capability scope must be "global" or a repository full name like "example-org/example-repo".' + ); + } + + return { + scopeType: 'repo', + scope: repoFullName, + repoFullName, + }; +} + +function readCapabilityPolicy(body: unknown): CapabilityPolicyConfig | Error { + if (!isRecord(body)) { + return new Error('Request body must be an object.'); + } + + const capabilityPolicy = body.capabilityPolicy; + if (!isRecord(capabilityPolicy)) { + return new Error('Request body must include capabilityPolicy.'); + } + + if (capabilityPolicy.availability !== undefined && !isRecord(capabilityPolicy.availability)) { + return new Error('capabilityPolicy.availability must be an object.'); + } + + return capabilityPolicy as CapabilityPolicyConfig; +} + +async function buildResponse(scope: ParsedScope) { + const agentRuntimeConfigService = AgentRuntimeConfigService.getInstance(); + const [globalConfig, repoConfig, effectiveConfig, capabilities] = await Promise.all([ + agentRuntimeConfigService.getGlobalConfig(), + scope.repoFullName ? agentRuntimeConfigService.getRepoConfig(scope.repoFullName) : Promise.resolve(null), + agentRuntimeConfigService.getEffectiveConfig(scope.repoFullName), + AgentSessionConfigService.getInstance().listCapabilityInventory(scope.scope), + ]); + const capabilityPolicy = + scope.scopeType === 'global' ? globalConfig.capabilityPolicy || {} : repoConfig?.capabilityPolicy || {}; + + return { + scope: scope.scope, + scopeType: scope.scopeType, + ...(scope.repoFullName ? { repoFullName: scope.repoFullName } : {}), + capabilityPolicy, + ...(scope.scopeType === 'repo' ? { inheritedCapabilityPolicy: globalConfig.capabilityPolicy || {} } : {}), + effectiveCapabilityPolicy: effectiveConfig.capabilityPolicy || {}, + capabilities, + }; +} + +/** + * @openapi + * /api/v2/ai/admin/agent/capabilities: + * get: + * summary: Get agent capability policy inventory + * description: Returns catalog-backed capability rows and effective capability policy for an admin scope. + * tags: + * - Agent Admin + * operationId: getAdminAgentCapabilities + * parameters: + * - in: query + * name: scope + * schema: + * type: string + * default: global + * description: Use `global` or a repository full name such as `example-org/example-repo`. + * responses: + * '200': + * description: Capability governance inventory. + * content: + * application/json: + * schema: + * $ref: '#/components/schemas/GetAdminAgentCapabilitiesSuccessResponse' + * '400': + * description: Invalid scope. + * content: + * application/json: + * schema: + * $ref: '#/components/schemas/ApiErrorResponse' + * '401': + * description: Unauthorized. + * content: + * application/json: + * schema: + * $ref: '#/components/schemas/ApiErrorResponse' + * '403': + * description: Forbidden. + * content: + * application/json: + * schema: + * $ref: '#/components/schemas/ApiErrorResponse' + * put: + * summary: Update agent capability policy + * description: Updates the capability policy for the selected global or repository scope. + * tags: + * - Agent Admin + * operationId: updateAdminAgentCapabilities + * parameters: + * - in: query + * name: scope + * schema: + * type: string + * default: global + * description: Use `global` or a repository full name such as `example-org/example-repo`. + * requestBody: + * required: true + * content: + * application/json: + * schema: + * $ref: '#/components/schemas/UpdateAdminAgentCapabilitiesRequest' + * responses: + * '200': + * description: Updated capability governance inventory. + * content: + * application/json: + * schema: + * $ref: '#/components/schemas/GetAdminAgentCapabilitiesSuccessResponse' + * '400': + * description: Validation error. + * content: + * application/json: + * schema: + * $ref: '#/components/schemas/ApiErrorResponse' + * '401': + * description: Unauthorized. + * content: + * application/json: + * schema: + * $ref: '#/components/schemas/ApiErrorResponse' + * '403': + * description: Forbidden. + * content: + * application/json: + * schema: + * $ref: '#/components/schemas/ApiErrorResponse' + */ +const getHandler = async (req: NextRequest) => { + const scope = parseScope(req); + if (scope instanceof Error) { + return errorResponse(scope, { status: 400 }, req); + } + + return successResponse(await buildResponse(scope), { status: 200 }, req); +}; + +const putHandler = async (req: NextRequest) => { + const scope = parseScope(req); + if (scope instanceof Error) { + return errorResponse(scope, { status: 400 }, req); + } + + let body: unknown; + try { + body = await req.json(); + } catch { + return errorResponse(new Error('Invalid JSON in request body'), { status: 400 }, req); + } + + const capabilityPolicy = readCapabilityPolicy(body); + if (capabilityPolicy instanceof Error) { + return errorResponse(capabilityPolicy, { status: 400 }, req); + } + + const agentRuntimeConfigService = AgentRuntimeConfigService.getInstance(); + try { + if (scope.scopeType === 'global') { + await agentRuntimeConfigService.updateGlobalCapabilityPolicy(capabilityPolicy); + } else { + await agentRuntimeConfigService.updateRepoCapabilityPolicy(scope.repoFullName, capabilityPolicy); + } + } catch (error) { + if (error instanceof AgentRuntimeConfigValidationError) { + return errorResponse(error, { status: 400 }, req); + } + throw error; + } + + return successResponse(await buildResponse(scope), { status: 200 }, req); +}; + +export const GET = createApiHandler(getHandler, { roles: ['admin'] }); +export const PUT = createApiHandler(putHandler, { roles: ['admin'] }); diff --git a/src/app/api/v2/ai/admin/agent/creation-policy/route.test.ts b/src/app/api/v2/ai/admin/agent/creation-policy/route.test.ts new file mode 100644 index 00000000..2aeb10e0 --- /dev/null +++ b/src/app/api/v2/ai/admin/agent/creation-policy/route.test.ts @@ -0,0 +1,191 @@ +/** + * 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 { AgentRuntimeConfigValidationError } from 'server/lib/validation/agentRuntimeConfigValidator'; + +const mockGetUser = jest.fn(); +const mockGetGlobalConfig = jest.fn(); +const mockUpdateGlobalCustomAgentCreationPolicy = jest.fn(); + +jest.mock('server/lib/get-user', () => ({ + getUser: (...args: unknown[]) => mockGetUser(...args), +})); + +jest.mock('server/services/agentRuntime/config/agentRuntimeConfig', () => ({ + __esModule: true, + default: { + getInstance: jest.fn(() => ({ + getGlobalConfig: (...args: unknown[]) => mockGetGlobalConfig(...args), + updateGlobalCustomAgentCreationPolicy: (...args: unknown[]) => mockUpdateGlobalCustomAgentCreationPolicy(...args), + })), + }, +})); + +import { GET, PUT } from './route'; + +function makeRequest(url: string, body?: unknown): NextRequest { + return { + headers: new Headers([['x-request-id', 'req-test']]), + nextUrl: new URL(url), + json: jest.fn().mockResolvedValue(body), + } as unknown as NextRequest; +} + +describe('/api/v2/ai/admin/agent/creation-policy', () => { + const originalEnableAuth = process.env.ENABLE_AUTH; + + beforeEach(() => { + jest.clearAllMocks(); + process.env.ENABLE_AUTH = 'true'; + mockGetUser.mockReturnValue({ + sub: 'sample-admin', + realm_access: { + roles: ['admin'], + }, + }); + mockGetGlobalConfig.mockResolvedValue({}); + mockUpdateGlobalCustomAgentCreationPolicy.mockResolvedValue({}); + }); + + afterEach(() => { + if (originalEnableAuth === undefined) { + delete process.env.ENABLE_AUTH; + } else { + process.env.ENABLE_AUTH = originalEnableAuth; + } + }); + + it('rejects non-admin users before loading policy', async () => { + mockGetUser.mockReturnValue({ + sub: 'sample-user', + realm_access: { + roles: ['user'], + }, + }); + + const response = await GET(makeRequest('http://localhost/api/v2/ai/admin/agent/creation-policy')); + const body = await response.json(); + + expect(response.status).toBe(403); + expect(body.error.message).toBe('Forbidden: insufficient permissions'); + expect(mockGetGlobalConfig).not.toHaveBeenCalled(); + }); + + it('returns an empty policy when no custom-agent creation policy is configured', async () => { + const response = await GET(makeRequest('http://localhost/api/v2/ai/admin/agent/creation-policy')); + const body = await response.json(); + + expect(response.status).toBe(200); + expect(body.data).toEqual({ + customAgentCreationPolicy: {}, + }); + }); + + it('returns the configured custom-agent creation policy', async () => { + mockGetGlobalConfig.mockResolvedValueOnce({ + customAgentCreationPolicy: { + mode: 'allowlist', + allowedGithubUsernames: ['sample-user'], + capabilityAvailability: { + workspace_shell: 'reserved', + }, + }, + }); + + const response = await GET(makeRequest('http://localhost/api/v2/ai/admin/agent/creation-policy')); + const body = await response.json(); + + expect(response.status).toBe(200); + expect(body.data.customAgentCreationPolicy).toEqual({ + mode: 'allowlist', + allowedGithubUsernames: ['sample-user'], + capabilityAvailability: { + workspace_shell: 'reserved', + }, + }); + }); + + it('updates policy and returns the normalized saved policy', async () => { + const policy = { + mode: 'allowlist', + allowedUserIds: ['sample-user-id'], + capabilityAvailability: { + workspace_shell: 'reserved', + }, + }; + mockGetGlobalConfig.mockResolvedValueOnce({ + customAgentCreationPolicy: policy, + }); + + const response = await PUT( + makeRequest('http://localhost/api/v2/ai/admin/agent/creation-policy', { + customAgentCreationPolicy: policy, + }) + ); + const body = await response.json(); + + expect(response.status).toBe(200); + expect(mockUpdateGlobalCustomAgentCreationPolicy).toHaveBeenCalledWith(policy); + expect(body.data.customAgentCreationPolicy).toEqual(policy); + }); + + it.each([ + ['missing policy', {}], + ['null policy', { customAgentCreationPolicy: null }], + ['array policy', { customAgentCreationPolicy: [] }], + ])('rejects malformed body: %s', async (_label, body) => { + const response = await PUT(makeRequest('http://localhost/api/v2/ai/admin/agent/creation-policy', body)); + const responseBody = await response.json(); + + expect(response.status).toBe(400); + expect(responseBody.error.message).toBe('Request body must include customAgentCreationPolicy.'); + expect(mockUpdateGlobalCustomAgentCreationPolicy).not.toHaveBeenCalled(); + }); + + it('rejects malformed capability availability before service mutation', async () => { + const response = await PUT( + makeRequest('http://localhost/api/v2/ai/admin/agent/creation-policy', { + customAgentCreationPolicy: { + capabilityAvailability: [], + }, + }) + ); + const body = await response.json(); + + expect(response.status).toBe(400); + expect(body.error.message).toBe('customAgentCreationPolicy.capabilityAvailability must be an object.'); + expect(mockUpdateGlobalCustomAgentCreationPolicy).not.toHaveBeenCalled(); + }); + + it('maps service validation errors to 400', async () => { + mockUpdateGlobalCustomAgentCreationPolicy.mockRejectedValueOnce( + new AgentRuntimeConfigValidationError('Invalid custom agent creation mode "sometimes".') + ); + + const response = await PUT( + makeRequest('http://localhost/api/v2/ai/admin/agent/creation-policy', { + customAgentCreationPolicy: { + mode: 'sometimes', + }, + }) + ); + const body = await response.json(); + + expect(response.status).toBe(400); + expect(body.error.message).toBe('Invalid custom agent creation mode "sometimes".'); + }); +}); diff --git a/src/app/api/v2/ai/admin/agent/creation-policy/route.ts b/src/app/api/v2/ai/admin/agent/creation-policy/route.ts new file mode 100644 index 00000000..d159bce8 --- /dev/null +++ b/src/app/api/v2/ai/admin/agent/creation-policy/route.ts @@ -0,0 +1,154 @@ +/** + * 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 { createApiHandler } from 'server/lib/createApiHandler'; +import { errorResponse, successResponse } from 'server/lib/response'; +import { AgentRuntimeConfigValidationError } from 'server/lib/validation/agentRuntimeConfigValidator'; +import AgentRuntimeConfigService from 'server/services/agentRuntime/config/agentRuntimeConfig'; +import type { CustomAgentCreationPolicyConfig } from 'server/services/types/agentRuntimeConfig'; + +export const dynamic = 'force-dynamic'; + +function isRecord(value: unknown): value is Record { + return Boolean(value) && typeof value === 'object' && !Array.isArray(value); +} + +function readCustomAgentCreationPolicy(body: unknown): CustomAgentCreationPolicyConfig | Error { + if (!isRecord(body)) { + return new Error('Request body must be an object.'); + } + + const customAgentCreationPolicy = body.customAgentCreationPolicy; + if (!isRecord(customAgentCreationPolicy)) { + return new Error('Request body must include customAgentCreationPolicy.'); + } + + if ( + customAgentCreationPolicy.capabilityAvailability !== undefined && + !isRecord(customAgentCreationPolicy.capabilityAvailability) + ) { + return new Error('customAgentCreationPolicy.capabilityAvailability must be an object.'); + } + + return customAgentCreationPolicy as CustomAgentCreationPolicyConfig; +} + +async function buildResponse() { + const config = await AgentRuntimeConfigService.getInstance().getGlobalConfig(); + + return { + customAgentCreationPolicy: config.customAgentCreationPolicy || {}, + }; +} + +/** + * @openapi + * /api/v2/ai/admin/agent/creation-policy: + * get: + * summary: Get custom-agent creation policy + * description: Returns the global policy that controls who can create custom agents and which capabilities are available during custom-agent creation. + * tags: + * - Agent Admin + * operationId: getAdminCustomAgentCreationPolicy + * responses: + * '200': + * description: Custom-agent creation policy. + * content: + * application/json: + * schema: + * $ref: '#/components/schemas/GetAdminCustomAgentCreationPolicySuccessResponse' + * '401': + * description: Unauthorized. + * content: + * application/json: + * schema: + * $ref: '#/components/schemas/ApiErrorResponse' + * '403': + * description: Forbidden. + * content: + * application/json: + * schema: + * $ref: '#/components/schemas/ApiErrorResponse' + * put: + * summary: Update custom-agent creation policy + * description: Updates the global policy that controls who can create custom agents and which capabilities are available during custom-agent creation. + * tags: + * - Agent Admin + * operationId: updateAdminCustomAgentCreationPolicy + * requestBody: + * required: true + * content: + * application/json: + * schema: + * $ref: '#/components/schemas/UpdateAdminCustomAgentCreationPolicyRequest' + * responses: + * '200': + * description: Updated custom-agent creation policy. + * content: + * application/json: + * schema: + * $ref: '#/components/schemas/GetAdminCustomAgentCreationPolicySuccessResponse' + * '400': + * description: Validation error. + * content: + * application/json: + * schema: + * $ref: '#/components/schemas/ApiErrorResponse' + * '401': + * description: Unauthorized. + * content: + * application/json: + * schema: + * $ref: '#/components/schemas/ApiErrorResponse' + * '403': + * description: Forbidden. + * content: + * application/json: + * schema: + * $ref: '#/components/schemas/ApiErrorResponse' + */ +const getHandler = async (req: NextRequest) => { + return successResponse(await buildResponse(), { status: 200 }, req); +}; + +const putHandler = async (req: NextRequest) => { + let body: unknown; + try { + body = await req.json(); + } catch { + return errorResponse(new Error('Invalid JSON in request body'), { status: 400 }, req); + } + + const customAgentCreationPolicy = readCustomAgentCreationPolicy(body); + if (customAgentCreationPolicy instanceof Error) { + return errorResponse(customAgentCreationPolicy, { status: 400 }, req); + } + + try { + await AgentRuntimeConfigService.getInstance().updateGlobalCustomAgentCreationPolicy(customAgentCreationPolicy); + } catch (error) { + if (error instanceof AgentRuntimeConfigValidationError) { + return errorResponse(error, { status: 400 }, req); + } + throw error; + } + + return successResponse(await buildResponse(), { status: 200 }, req); +}; + +export const GET = createApiHandler(getHandler, { roles: ['admin'] }); +export const PUT = createApiHandler(putHandler, { roles: ['admin'] }); diff --git a/src/app/api/v2/ai/admin/agent/tools/route.ts b/src/app/api/v2/ai/admin/agent/tools/route.ts index c2e3067e..663160da 100644 --- a/src/app/api/v2/ai/admin/agent/tools/route.ts +++ b/src/app/api/v2/ai/admin/agent/tools/route.ts @@ -49,6 +49,12 @@ export const dynamic = 'force-dynamic'; * application/json: * schema: * $ref: '#/components/schemas/ApiErrorResponse' + * '403': + * description: Forbidden + * content: + * application/json: + * schema: + * $ref: '#/components/schemas/ApiErrorResponse' */ const getHandler = async (req: NextRequest) => { const userIdentity = getRequestUserIdentity(req); @@ -61,4 +67,4 @@ const getHandler = async (req: NextRequest) => { return successResponse(data, { status: 200 }, req); }; -export const GET = createApiHandler(getHandler); +export const GET = createApiHandler(getHandler, { roles: ['admin'] }); diff --git a/src/app/api/v2/ai/admin/feedback/[id]/conversation/route.test.ts b/src/app/api/v2/ai/admin/feedback/[id]/conversation/route.test.ts deleted file mode 100644 index e6a1bffe..00000000 --- a/src/app/api/v2/ai/admin/feedback/[id]/conversation/route.test.ts +++ /dev/null @@ -1,159 +0,0 @@ -/** - * Copyright 2025 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/models/MessageFeedback', () => { - const model: { query: jest.Mock } = { query: jest.fn() }; - return { __esModule: true, default: model }; -}); - -jest.mock('server/models/ConversationFeedback', () => { - const model: { query: jest.Mock } = { query: jest.fn() }; - return { __esModule: true, default: model }; -}); - -jest.mock('server/models/Conversation', () => { - const model: { query: jest.Mock } = { query: jest.fn() }; - return { __esModule: true, default: model }; -}); - -import { GET } from './route'; -import MessageFeedback from 'server/models/MessageFeedback'; -import ConversationFeedback from 'server/models/ConversationFeedback'; -import Conversation from 'server/models/Conversation'; - -const MockMessageFeedback = MessageFeedback as unknown as { query: jest.Mock }; -const MockConversationFeedback = ConversationFeedback as unknown as { query: jest.Mock }; -const MockConversation = Conversation as unknown as { query: jest.Mock }; - -function makeRequest(): NextRequest { - return { - headers: new Headers([['x-request-id', 'req-test']]), - } as unknown as NextRequest; -} - -describe('GET /api/v2/ai/admin/feedback/[id]/conversation', () => { - beforeEach(() => { - jest.clearAllMocks(); - }); - - it('returns 400 for invalid feedback id format', async () => { - const response = await GET(makeRequest(), { params: { id: 'invalid' } }); - - expect(response.status).toBe(400); - await expect(response.json()).resolves.toMatchObject({ - error: { message: 'Invalid feedback id format' }, - }); - }); - - it('returns message conversation replay with feedback comment', async () => { - const findByIdMessageFeedback = jest.fn().mockResolvedValue({ - id: 9, - buildUuid: 'uuid-1', - messageId: 42, - repo: 'org/repo', - rating: 'up', - text: 'Very helpful', - userIdentifier: 'sample-user', - createdAt: '2026-02-27T10:00:00.000Z', - }); - MockMessageFeedback.query.mockReturnValue({ findById: findByIdMessageFeedback }); - - const findByIdConversation = jest.fn().mockReturnThis(); - const withGraphFetched = jest.fn().mockReturnThis(); - const modifiers = jest.fn().mockResolvedValue({ - messageCount: 2, - model: 'claude-sonnet', - messages: [ - { - id: 1, - role: 'user', - content: 'Why is my build failing?', - timestamp: '1700000000000', - metadata: {}, - }, - { - id: 2, - role: 'assistant', - content: 'The dockerfilePath is wrong.', - timestamp: 1700000001000, - metadata: { toolCalls: [] }, - }, - ], - }); - MockConversation.query.mockReturnValue({ - findById: findByIdConversation, - withGraphFetched, - modifiers, - }); - - const response = await GET(makeRequest(), { params: { id: 'message-9' } }); - const body = await response.json(); - - expect(response.status).toBe(200); - expect(body.data.feedbackText).toBe('Very helpful'); - expect(body.data.feedbackUserIdentifier).toBe('sample-user'); - expect(body.data.feedbackType).toBe('message'); - expect(body.data.ratedMessageId).toBe(42); - expect(body.data.conversation.messages).toHaveLength(2); - expect(body.data.conversation.messages[0].timestamp).toBe(1700000000000); - expect(MockConversationFeedback.query).not.toHaveBeenCalled(); - }); - - it('returns conversation feedback replay payload', async () => { - const findByIdConversationFeedback = jest.fn().mockResolvedValue({ - id: 5, - buildUuid: 'uuid-2', - repo: 'org/repo', - rating: 'down', - text: 'Session feedback comment', - userIdentifier: 'sample-user', - createdAt: '2026-02-27T10:05:00.000Z', - }); - MockConversationFeedback.query.mockReturnValue({ findById: findByIdConversationFeedback }); - - const findByIdConversation = jest.fn().mockReturnThis(); - const withGraphFetched = jest.fn().mockReturnThis(); - const modifiers = jest.fn().mockResolvedValue({ - messageCount: 1, - model: null, - messages: [ - { - id: 10, - role: 'assistant', - content: 'Sample response', - timestamp: 1700000002000, - metadata: {}, - }, - ], - }); - MockConversation.query.mockReturnValue({ - findById: findByIdConversation, - withGraphFetched, - modifiers, - }); - - const response = await GET(makeRequest(), { params: { id: 'conversation-5' } }); - const body = await response.json(); - - expect(response.status).toBe(200); - expect(body.data.feedbackType).toBe('conversation'); - expect(body.data.feedbackText).toBe('Session feedback comment'); - expect(body.data.feedbackUserIdentifier).toBe('sample-user'); - expect(body.data.ratedMessageId).toBeNull(); - }); -}); diff --git a/src/app/api/v2/ai/admin/feedback/[id]/conversation/route.ts b/src/app/api/v2/ai/admin/feedback/[id]/conversation/route.ts deleted file mode 100644 index c41c38cb..00000000 --- a/src/app/api/v2/ai/admin/feedback/[id]/conversation/route.ts +++ /dev/null @@ -1,174 +0,0 @@ -/** - * Copyright 2025 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 { createApiHandler } from 'server/lib/createApiHandler'; -import { errorResponse, successResponse } from 'server/lib/response'; -import MessageFeedback from 'server/models/MessageFeedback'; -import ConversationFeedback from 'server/models/ConversationFeedback'; -import Conversation from 'server/models/Conversation'; - -/** - * @openapi - * /api/v2/ai/admin/feedback/{id}/conversation: - * get: - * summary: Get AI feedback conversation replay - * description: > - * Returns the full persisted conversation for a message or conversation - * feedback record so admins can replay the exchange that led to the rating. - * tags: - * - AI Feedback Admin - * operationId: getAdminFeedbackConversation - * parameters: - * - in: path - * name: id - * required: true - * schema: - * type: string - * description: > - * Composite feedback identifier in the form `message-123` or - * `conversation-456`. - * responses: - * '200': - * description: Feedback conversation replay. - * content: - * application/json: - * schema: - * $ref: '#/components/schemas/GetAdminFeedbackConversationSuccessResponse' - * '400': - * description: Invalid feedback identifier. - * content: - * application/json: - * schema: - * $ref: '#/components/schemas/ApiErrorResponse' - * '404': - * description: Feedback record or conversation not found. - * content: - * application/json: - * schema: - * $ref: '#/components/schemas/ApiErrorResponse' - * '500': - * description: Server error - * content: - * application/json: - * schema: - * $ref: '#/components/schemas/ApiErrorResponse' - */ -const getHandler = async (req: NextRequest, { params }: { params: { id: string } }) => { - const { id } = params; - - if (!id) { - return errorResponse(new Error('Missing required parameter: id'), { status: 400 }, req); - } - - const dashIndex = id.indexOf('-'); - if (dashIndex === -1) { - return errorResponse(new Error('Invalid feedback id format'), { status: 400 }, req); - } - - const feedbackType = id.slice(0, dashIndex) as 'message' | 'conversation'; - const recordId = parseInt(id.slice(dashIndex + 1), 10); - - if (feedbackType !== 'message' && feedbackType !== 'conversation') { - return errorResponse(new Error('Invalid feedback type: must be "message" or "conversation"'), { status: 400 }, req); - } - - if (isNaN(recordId)) { - return errorResponse(new Error('Invalid feedback record id'), { status: 400 }, req); - } - - let buildUuid: string; - let ratedMessageId: number | null = null; - let repo: string; - let feedbackRating: 'up' | 'down'; - let feedbackText: string | null = null; - let feedbackUserIdentifier: string | null = null; - let feedbackCreatedAt: string | Date; - - if (feedbackType === 'message') { - const record = await MessageFeedback.query().findById(recordId); - if (!record) { - return errorResponse(new Error('Feedback record not found'), { status: 404 }, req); - } - buildUuid = record.buildUuid; - ratedMessageId = record.messageId; - repo = record.repo; - feedbackRating = record.rating; - feedbackText = record.text ?? null; - feedbackUserIdentifier = record.userIdentifier ?? null; - feedbackCreatedAt = record.createdAt; - } else { - const record = await ConversationFeedback.query().findById(recordId); - if (!record) { - return errorResponse(new Error('Feedback record not found'), { status: 404 }, req); - } - buildUuid = record.buildUuid; - repo = record.repo; - feedbackRating = record.rating; - feedbackText = record.text ?? null; - feedbackUserIdentifier = record.userIdentifier ?? null; - feedbackCreatedAt = record.createdAt; - } - - const conversation = await Conversation.query() - .findById(buildUuid) - .withGraphFetched('messages(orderByTimestamp)') - .modifiers({ - orderByTimestamp(builder: any) { - builder.orderBy('timestamp', 'asc'); - }, - }); - - if (!conversation) { - return errorResponse(new Error('Conversation not found'), { status: 404 }, req); - } - - const messages = ((conversation as any).messages || []).map((msg: any) => { - // PostgreSQL BIGINT values can come back as strings; normalize for UI consumers. - const numericTimestamp = typeof msg.timestamp === 'number' ? msg.timestamp : Number(msg.timestamp); - - return { - id: msg.id, - role: msg.role, - content: msg.content, - timestamp: Number.isNaN(numericTimestamp) ? Date.parse(String(msg.timestamp)) : numericTimestamp, - metadata: msg.metadata || {}, - }; - }); - - return successResponse( - { - feedbackType, - feedbackId: recordId, - buildUuid, - repo, - ratedMessageId, - feedbackRating, - feedbackText, - feedbackUserIdentifier, - feedbackCreatedAt, - conversation: { - messageCount: conversation.messageCount, - model: conversation.model || null, - messages, - }, - }, - { status: 200 }, - req - ); -}; - -export const GET = createApiHandler(getHandler); diff --git a/src/app/api/v2/ai/admin/feedback/route.test.ts b/src/app/api/v2/ai/admin/feedback/route.test.ts deleted file mode 100644 index e5450644..00000000 --- a/src/app/api/v2/ai/admin/feedback/route.test.ts +++ /dev/null @@ -1,218 +0,0 @@ -/** - * Copyright 2025 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/dependencies', () => ({ - defaultDb: { - knex: Object.assign(jest.fn(), { - raw: jest.fn(), - from: jest.fn(), - queryBuilder: jest.fn(), - }), - }, -})); - -import { GET } from './route'; -import { defaultDb } from 'server/lib/dependencies'; - -const mockKnex = defaultDb.knex as jest.Mock & { - raw: jest.Mock; - from: jest.Mock; - queryBuilder: jest.Mock; -}; - -function makeRequest(url: string): NextRequest { - return { - headers: new Headers([['x-request-id', 'req-test']]), - nextUrl: new URL(url), - } as unknown as NextRequest; -} - -function createMessageFeedbackQueryBuilder() { - const query: any = {}; - query.leftJoin = jest.fn().mockReturnValue(query); - query.select = jest.fn().mockReturnValue(query); - query.where = jest.fn().mockReturnValue(query); - query.clone = jest.fn().mockReturnValue({ - as: jest.fn().mockReturnValue('feedback_count_query'), - }); - query.as = jest.fn().mockReturnValue('feedback_rows_query'); - return query; -} - -function createConversationFeedbackQueryBuilder() { - const query: any = {}; - query.select = jest.fn().mockReturnValue(query); - query.where = jest.fn().mockReturnValue(query); - query.clone = jest.fn().mockReturnValue({ - as: jest.fn().mockReturnValue('feedback_count_query'), - }); - query.as = jest.fn().mockReturnValue('feedback_rows_query'); - return query; -} - -describe('GET /api/v2/ai/admin/feedback', () => { - beforeEach(() => { - jest.clearAllMocks(); - mockKnex.raw.mockImplementation((sql: string, bindings?: unknown[]) => ({ sql, bindings })); - mockKnex.queryBuilder.mockReturnValue({ - unionAll: jest.fn(), - }); - }); - - it('returns message feedback rows with safe Unicode truncation and default sorting', async () => { - const messageQueryBuilder = createMessageFeedbackQueryBuilder(); - - mockKnex.mockImplementation((table: string) => { - if (table === 'message_feedback as mf') { - return messageQueryBuilder; - } - throw new Error(`Unexpected table query: ${table}`); - }); - - const countBuilder = { - count: jest.fn().mockReturnThis(), - first: jest.fn().mockResolvedValue({ count: '1' }), - }; - const rowsBuilder = { - select: jest.fn().mockReturnThis(), - orderBy: jest.fn().mockReturnThis(), - offset: jest.fn().mockReturnThis(), - limit: jest.fn().mockResolvedValue([ - { - id: 'message-1', - feedbackType: 'message', - buildUuid: 'uuid-1', - rating: 'up', - text: null, - userIdentifier: 'sample-user', - repo: 'org/repo', - prNumber: null, - messageId: 11, - messagePreview: null, - messageContent: '😀'.repeat(200), - messageMetadata: { - debugMetrics: { - inputTokens: 1_000_000, - outputTokens: 1_000_000, - inputCostPerMillion: 1.5, - outputCostPerMillion: 2.5, - }, - }, - createdAt: '2026-02-27T10:00:00.000Z', - }, - ]), - }; - - mockKnex.from.mockReturnValueOnce(countBuilder).mockReturnValueOnce(rowsBuilder); - - const response = await GET(makeRequest('http://localhost/api/v2/ai/admin/feedback?type=message&page=1&limit=25')); - const body = await response.json(); - const preview = body.data[0].messagePreview as string; - const previewWithoutEllipsis = preview.endsWith('…') ? preview.slice(0, -1) : preview; - - expect(response.status).toBe(200); - expect(body.data).toHaveLength(1); - expect(preview.endsWith('…')).toBe(true); - expect(/[\uD800-\uDBFF]$/.test(previewWithoutEllipsis)).toBe(false); - expect(body.data[0].userIdentifier).toBe('sample-user'); - expect(body.data[0].costUsd).toBeCloseTo(4, 6); - expect(rowsBuilder.orderBy).toHaveBeenCalledWith('createdAt', 'desc'); - }); - - it('returns conversation feedback rows with aggregated session cost', async () => { - const conversationQueryBuilder = createConversationFeedbackQueryBuilder(); - - const conversationMessagesBuilder = { - select: jest.fn().mockReturnThis(), - whereIn: jest.fn().mockResolvedValue([ - { - buildUuid: 'uuid-2', - metadata: { - debugMetrics: { - inputTokens: 2_000_000, - outputTokens: 0, - inputCostPerMillion: 1, - outputCostPerMillion: 1, - }, - }, - }, - { - buildUuid: 'uuid-2', - metadata: { - debugMetrics: { - inputTokens: 1_000_000, - outputTokens: 1_000_000, - inputCostPerMillion: 1, - outputCostPerMillion: 1, - }, - }, - }, - ]), - }; - - mockKnex.mockImplementation((table: string) => { - if (table === 'conversation_feedback as cf') { - return conversationQueryBuilder; - } - if (table === 'conversation_messages as cm') { - return conversationMessagesBuilder; - } - throw new Error(`Unexpected table query: ${table}`); - }); - - const countBuilder = { - count: jest.fn().mockReturnThis(), - first: jest.fn().mockResolvedValue({ count: '1' }), - }; - const rowsBuilder = { - select: jest.fn().mockReturnThis(), - orderBy: jest.fn().mockReturnThis(), - offset: jest.fn().mockReturnThis(), - limit: jest.fn().mockResolvedValue([ - { - id: 'conversation-1', - feedbackType: 'conversation', - buildUuid: 'uuid-2', - rating: 'up', - text: 'Great session', - userIdentifier: 'sample-user', - repo: 'org/repo', - prNumber: null, - messageId: null, - messagePreview: null, - messageContent: null, - messageMetadata: null, - createdAt: '2026-02-27T11:00:00.000Z', - }, - ]), - }; - - mockKnex.from.mockReturnValueOnce(countBuilder).mockReturnValueOnce(rowsBuilder); - - const response = await GET( - makeRequest('http://localhost/api/v2/ai/admin/feedback?type=conversation&page=1&limit=25') - ); - const body = await response.json(); - - expect(response.status).toBe(200); - expect(body.data).toHaveLength(1); - expect(body.data[0].userIdentifier).toBe('sample-user'); - expect(body.data[0].costUsd).toBeCloseTo(4, 6); - expect(conversationMessagesBuilder.whereIn).toHaveBeenCalledWith('cm.buildUuid', ['uuid-2']); - }); -}); diff --git a/src/app/api/v2/ai/admin/feedback/route.ts b/src/app/api/v2/ai/admin/feedback/route.ts deleted file mode 100644 index 346d15d3..00000000 --- a/src/app/api/v2/ai/admin/feedback/route.ts +++ /dev/null @@ -1,492 +0,0 @@ -/** - * Copyright 2025 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 { Knex } from 'knex'; -import { NextRequest } from 'next/server'; -import { createApiHandler } from 'server/lib/createApiHandler'; -import { defaultDb } from 'server/lib/dependencies'; -import { successResponse } from 'server/lib/response'; - -interface FeedbackEntry { - id: string; - feedbackType: 'message' | 'conversation'; - buildUuid: string; - rating: 'up' | 'down'; - text: string | null; - userIdentifier: string | null; - repo: string; - prNumber: number | null; - messageId: number | null; - messagePreview: string | null; - costUsd: number | null; - createdAt: string; -} - -interface FeedbackQueryRow { - id: string; - feedbackType: 'message' | 'conversation'; - buildUuid: string; - rating: 'up' | 'down'; - text: string | null; - userIdentifier: string | null; - repo: string; - prNumber: number | null; - messageId: number | null; - messagePreview: string | null; - createdAt: string; - messageContent?: string | null; - messageMetadata?: unknown; -} - -interface FeedbackFilters { - repo?: string; - rating?: string; - from?: string; - to?: string; - sortBy?: 'createdAt'; - sortDirection?: 'asc' | 'desc'; -} - -const PREVIEW_MAX_LENGTH = 140; - -interface DebugMetricsPayload { - inputTokens?: unknown; - outputTokens?: unknown; - inputCostPerMillion?: unknown; - outputCostPerMillion?: unknown; -} - -interface ConversationMessageCostRow { - buildUuid: string; - metadata: unknown; -} - -/** - * @openapi - * /api/v2/ai/admin/feedback: - * get: - * summary: List AI feedback records - * description: > - * Returns paginated message-level and conversation-level AI feedback for - * admin review. Results can be filtered by repository, rating, type, and - * date range, and are ordered by creation time. - * tags: - * - AI Feedback Admin - * operationId: getAdminFeedbackList - * parameters: - * - 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 - * description: Number of feedback records per page. - * - in: query - * name: repo - * schema: - * type: string - * description: Case-insensitive repository or build UUID search term. - * - in: query - * name: rating - * schema: - * type: string - * enum: [up, down] - * description: Filter by user rating. - * - in: query - * name: type - * schema: - * type: string - * enum: [message, conversation] - * description: Filter by feedback record type. - * - in: query - * name: from - * schema: - * type: string - * format: date-time - * description: Include feedback created on or after this timestamp. - * - in: query - * name: to - * schema: - * type: string - * format: date-time - * description: Include feedback created on or before this timestamp. - * - in: query - * name: sortBy - * schema: - * type: string - * enum: [createdAt] - * default: createdAt - * description: Sort field. Only `createdAt` is currently supported. - * - in: query - * name: sortDirection - * schema: - * type: string - * enum: [asc, desc] - * default: desc - * description: Sort direction. - * responses: - * '200': - * description: Paginated feedback records. - * content: - * application/json: - * schema: - * $ref: '#/components/schemas/GetAdminFeedbackListSuccessResponse' - * '500': - * description: Server error - * content: - * application/json: - * schema: - * $ref: '#/components/schemas/ApiErrorResponse' - */ - -function toSingleLine(value: string): string { - return value.replace(/\s+/g, ' ').trim(); -} - -function truncatePreview(value: string): string { - const graphemes = Array.from(value); - if (graphemes.length <= PREVIEW_MAX_LENGTH) { - return value; - } - return graphemes.slice(0, PREVIEW_MAX_LENGTH - 1).join('') + '…'; -} - -function parseFiniteNumber(value: unknown): number | null { - if (typeof value === 'number') { - return Number.isFinite(value) ? value : null; - } - - if (typeof value === 'string' && value.trim()) { - const parsed = Number(value); - return Number.isFinite(parsed) ? parsed : null; - } - - return null; -} - -function extractDebugMetrics(metadata: unknown): DebugMetricsPayload | null { - if (!metadata || typeof metadata !== 'object' || Array.isArray(metadata)) { - return null; - } - - const candidate = (metadata as Record).debugMetrics; - if (!candidate || typeof candidate !== 'object' || Array.isArray(candidate)) { - return null; - } - - return candidate as DebugMetricsPayload; -} - -function computeCostFromDebugMetrics(debugMetrics: DebugMetricsPayload | null): number | null { - if (!debugMetrics) { - return null; - } - - const inputTokens = parseFiniteNumber(debugMetrics.inputTokens); - const outputTokens = parseFiniteNumber(debugMetrics.outputTokens); - const inputCostPerMillion = parseFiniteNumber(debugMetrics.inputCostPerMillion); - const outputCostPerMillion = parseFiniteNumber(debugMetrics.outputCostPerMillion); - - let totalCost = 0; - let hasCost = false; - - if (inputTokens != null && inputCostPerMillion != null) { - totalCost += (inputTokens / 1_000_000) * inputCostPerMillion; - hasCost = true; - } - - if (outputTokens != null && outputCostPerMillion != null) { - totalCost += (outputTokens / 1_000_000) * outputCostPerMillion; - hasCost = true; - } - - return hasCost ? totalCost : null; -} - -function computeMessageCost(metadata: unknown): number | null { - return computeCostFromDebugMetrics(extractDebugMetrics(metadata)); -} - -async function computeSessionCostByBuildUuid(db: Knex, buildUuids: string[]): Promise> { - const uniqueBuildUuids = [...new Set(buildUuids.filter(Boolean))]; - const costByBuildUuid = new Map(); - if (uniqueBuildUuids.length === 0) { - return costByBuildUuid; - } - - const rows = (await db('conversation_messages as cm') - .select({ - buildUuid: 'cm.buildUuid', - metadata: 'cm.metadata', - }) - .whereIn('cm.buildUuid', uniqueBuildUuids)) as ConversationMessageCostRow[]; - - const rollingCost = new Map(); - const hasCost = new Set(); - - for (const row of rows) { - const messageCost = computeMessageCost(row.metadata); - if (messageCost == null) { - continue; - } - rollingCost.set(row.buildUuid, (rollingCost.get(row.buildUuid) || 0) + messageCost); - hasCost.add(row.buildUuid); - } - - for (const buildUuid of uniqueBuildUuids) { - costByBuildUuid.set(buildUuid, hasCost.has(buildUuid) ? rollingCost.get(buildUuid) || 0 : null); - } - - return costByBuildUuid; -} - -function extractSummaryFromStructuredContent(content: string): string | null { - const trimmed = content.trim(); - if (!trimmed.startsWith('{')) { - return null; - } - - try { - const parsed = JSON.parse(trimmed) as { - type?: string; - summary?: string; - services?: Array<{ - serviceName?: string; - issue?: string; - rootCause?: string; - suggestedFix?: string; - }>; - }; - - if (typeof parsed.summary === 'string' && parsed.summary.trim()) { - return parsed.summary; - } - - const firstService = parsed.services?.[0]; - if (firstService) { - const detail = firstService.issue || firstService.rootCause || firstService.suggestedFix; - if (detail && firstService.serviceName) { - return `${firstService.serviceName}: ${detail}`; - } - if (detail) { - return detail; - } - } - } catch { - const regexMatch = trimmed.match(/"summary"\s*:\s*"([^"]+)"/); - if (regexMatch?.[1]) { - return regexMatch[1]; - } - } - - return null; -} - -function buildPreview(row: FeedbackQueryRow): string | null { - if (row.feedbackType === 'conversation') { - return row.text ? truncatePreview(toSingleLine(row.text)) : null; - } - - const source = row.messageContent || row.messagePreview || row.text; - if (!source) { - return null; - } - - const structuredSummary = extractSummaryFromStructuredContent(source); - if (structuredSummary) { - return truncatePreview(toSingleLine(structuredSummary)); - } - - return truncatePreview(toSingleLine(source)); -} - -function applyFilters(query: Knex.QueryBuilder, tableAlias: string, filters: FeedbackFilters) { - if (filters.repo) { - const search = `%${filters.repo.trim()}%`; - query.where((builder) => { - builder.whereILike(`${tableAlias}.repo`, search).orWhereILike(`${tableAlias}.buildUuid`, search); - }); - } - if (filters.rating) { - query.where(`${tableAlias}.rating`, filters.rating); - } - if (filters.from) { - query.where(`${tableAlias}.createdAt`, '>=', filters.from); - } - if (filters.to) { - query.where(`${tableAlias}.createdAt`, '<=', filters.to); - } -} - -function buildMessageFeedbackQuery(db: Knex, filters: FeedbackFilters): Knex.QueryBuilder { - const query = db('message_feedback as mf') - .leftJoin('conversation_messages as cm', 'mf.messageId', 'cm.id') - .select({ - buildUuid: 'mf.buildUuid', - rating: 'mf.rating', - text: 'mf.text', - userIdentifier: 'mf.userIdentifier', - repo: 'mf.repo', - prNumber: 'mf.prNumber', - messageId: 'mf.messageId', - createdAt: 'mf.createdAt', - }) - .select(db.raw("concat('message-', ??) as ??", ['mf.id', 'id'])) - .select(db.raw("'message' as ??", ['feedbackType'])) - .select(db.raw('NULL::text as ??', ['messagePreview'])) - .select({ messageContent: 'cm.content' }) - .select({ messageMetadata: 'cm.metadata' }); - - applyFilters(query, 'mf', filters); - return query; -} - -function buildConversationFeedbackQuery(db: Knex, filters: FeedbackFilters): Knex.QueryBuilder { - const query = db('conversation_feedback as cf') - .select({ - buildUuid: 'cf.buildUuid', - rating: 'cf.rating', - text: 'cf.text', - userIdentifier: 'cf.userIdentifier', - repo: 'cf.repo', - prNumber: 'cf.prNumber', - }) - .select(db.raw('NULL::integer as ??', ['messageId'])) - .select({ createdAt: 'cf.createdAt' }) - .select(db.raw("concat('conversation-', ??) as ??", ['cf.id', 'id'])) - .select(db.raw("'conversation' as ??", ['feedbackType'])) - .select(db.raw('NULL::text as ??', ['messagePreview'])) - .select(db.raw('NULL::text as ??', ['messageContent'])) - .select(db.raw('NULL::jsonb as ??', ['messageMetadata'])); - - applyFilters(query, 'cf', filters); - return query; -} - -function buildUnifiedFeedbackQuery(db: Knex, filters: FeedbackFilters, type?: string): Knex.QueryBuilder | null { - const queries: Knex.QueryBuilder[] = []; - - if (!type || type === 'message') { - queries.push(buildMessageFeedbackQuery(db, filters)); - } - if (!type || type === 'conversation') { - queries.push(buildConversationFeedbackQuery(db, filters)); - } - - if (queries.length === 0) { - return null; - } - - if (queries.length === 1) { - return queries[0]; - } - - return db.queryBuilder().unionAll(queries, true); -} - -const getHandler = async (req: NextRequest) => { - const searchParams = req.nextUrl.searchParams; - const page = Math.max(1, parseInt(searchParams.get('page') || '1', 10)); - const limit = Math.max(1, parseInt(searchParams.get('limit') || '25', 10)); - const repo = searchParams.get('repo') || undefined; - const rating = searchParams.get('rating') || undefined; - const type = searchParams.get('type') || undefined; - const from = searchParams.get('from') || undefined; - const to = searchParams.get('to') || undefined; - const sortBy = searchParams.get('sortBy'); - const sortDirection = searchParams.get('sortDirection'); - const normalizedSortBy: FeedbackFilters['sortBy'] = sortBy === 'createdAt' ? 'createdAt' : 'createdAt'; - const normalizedSortDirection: FeedbackFilters['sortDirection'] = sortDirection === 'asc' ? 'asc' : 'desc'; - const filters: FeedbackFilters = { - repo, - rating, - from, - to, - sortBy: normalizedSortBy, - sortDirection: normalizedSortDirection, - }; - - const db = defaultDb.knex; - const unifiedQuery = buildUnifiedFeedbackQuery(db, filters, type); - - let data: FeedbackEntry[] = []; - let totalCount = 0; - - if (unifiedQuery) { - const countResult = await db - .from(unifiedQuery.clone().as('feedback')) - .count<{ count: string }>('* as count') - .first(); - totalCount = Number(countResult?.count || 0); - - const offset = (page - 1) * limit; - const rows = (await db - .from(unifiedQuery.as('feedback')) - .select('*') - .orderBy(filters.sortBy || 'createdAt', filters.sortDirection || 'desc') - .offset(offset) - .limit(limit)) as FeedbackQueryRow[]; - const conversationBuildUuids = rows - .filter((row) => row.feedbackType === 'conversation') - .map((row) => row.buildUuid); - const sessionCostByBuildUuid = await computeSessionCostByBuildUuid(db, conversationBuildUuids); - - data = rows.map((row) => ({ - id: row.id, - feedbackType: row.feedbackType, - buildUuid: row.buildUuid, - rating: row.rating, - text: row.text, - userIdentifier: row.userIdentifier, - repo: row.repo, - prNumber: row.prNumber, - messageId: row.messageId, - messagePreview: buildPreview(row), - costUsd: - row.feedbackType === 'message' - ? computeMessageCost(row.messageMetadata) - : sessionCostByBuildUuid.get(row.buildUuid) ?? null, - createdAt: row.createdAt, - })); - } - - const totalPages = Math.ceil(totalCount / limit); - - return successResponse( - data, - { - status: 200, - metadata: { - pagination: { - page, - totalPages, - totalCount, - limit, - }, - }, - }, - req - ); -}; - -export const GET = createApiHandler(getHandler); 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 index 5f24c4d8..c255519f 100644 --- 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 @@ -83,18 +83,10 @@ jest.mock('server/services/agent/SourceService', () => ({ }, })); -jest.mock('server/services/agent/CapabilityService', () => ({ +jest.mock('server/services/agent/RunPlanResolver', () => ({ __esModule: true, default: { - resolveSessionContext: jest.fn(), - }, -})); - -jest.mock('server/services/agent/ProviderRegistry', () => ({ - __esModule: true, - MissingAgentProviderApiKeyError: class MissingAgentProviderApiKeyError extends Error {}, - default: { - resolveSelection: jest.fn(), + resolveForRunAdmission: jest.fn(), }, })); @@ -176,8 +168,7 @@ 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 AgentRunPlanResolver from 'server/services/agent/RunPlanResolver'; import AgentRunAdmissionService from 'server/services/agent/RunAdmissionService'; import AgentRunQueueService from 'server/services/agent/RunQueueService'; import AgentRunService from 'server/services/agent/RunService'; @@ -200,8 +191,7 @@ const mockGetOwnedThreadWithSession = AgentThreadService.getOwnedThreadWithSessi 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 mockResolveForRunAdmission = AgentRunPlanResolver.resolveForRunAdmission as jest.Mock; const mockCreateQueuedRunWithMessage = AgentRunAdmissionService.createQueuedRunWithMessage as jest.Mock; const mockEnqueueRun = AgentRunQueueService.enqueueRun as jest.Mock; const mockGetOwnedRun = AgentRunService.getOwnedRun as jest.Mock; @@ -375,12 +365,25 @@ function simulateApprovalRequest() { commandPreview: null, fileChangePreview: [ { + id: 'tool-call-1:sample-file.txt', + toolCallId: 'tool-call-1', + sourceTool: 'workspace_edit_file', path: 'sample-file.txt', - action: 'edited', + displayPath: 'sample-file.txt', + kind: 'edited', + stage: 'awaiting-approval', summary: 'Updated sample-file.txt', additions: 1, deletions: 1, truncated: false, + unifiedDiff: null, + beforeTextPreview: null, + afterTextPreview: null, + encoding: null, + oldSizeBytes: null, + newSizeBytes: null, + oldSha256: null, + newSha256: null, }, ], riskLabels: ['Workspace write'], @@ -464,16 +467,71 @@ describe('canonical agent session API acceptance flow', () => { }); mockCanAcceptMessages.mockReturnValue(true); mockGetSessionSource.mockResolvedValue({ + uuid: 'source-1', + adapter: 'blank_workspace', status: 'ready', sandboxRequirements: { filesystem: 'persistent' }, }); - mockResolveSessionContext.mockResolvedValue({ + mockResolveForRunAdmission.mockResolvedValue({ approvalPolicy: { defaultMode: 'require_approval', rules: {} }, - repoFullName: 'example-org/example-repo', - }); - mockResolveSelection.mockResolvedValue({ - provider: 'openai', - modelId: 'gpt-5.4', + requestedHarness: null, + requestedProvider: null, + requestedModel: null, + resolvedHarness: 'lifecycle_ai_sdk', + resolvedProvider: 'openai', + resolvedModel: 'gpt-5.4', + sandboxRequirement: { filesystem: 'persistent' }, + runtimeOptions: {}, + runPlanSnapshot: { + version: 1, + capturedAt: '2026-05-03T00:00:00.000Z', + agent: { + id: 'system.freeform', + label: 'Free-form', + ownerKind: 'system', + version: 1, + sourceKind: 'freeform_chat', + resourcePolicy: { + sourceKinds: ['build_context_chat', 'workspace_session', 'freeform_chat'], + workspaceRequired: false, + sandboxRequired: false, + }, + modelPreference: null, + }, + source: { + id: 'source-1', + adapter: 'blank_workspace', + status: 'ready', + sessionKind: 'chat', + freshness: { + capturedAt: '2026-05-03T00:00:00.000Z', + freshnessSource: 'source', + }, + }, + model: { + requestedProvider: null, + requestedModel: null, + resolvedProvider: 'openai', + resolvedModel: 'gpt-5.4', + }, + runtime: { + requestedHarness: null, + resolvedHarness: 'lifecycle_ai_sdk', + sandboxRequirement: { filesystem: 'persistent' }, + runtimeOptions: {}, + approvalPolicy: { defaultMode: 'require_approval', rules: {} }, + }, + prompt: { + instructionRefs: [], + renderedSummary: 'Sample prompt summary', + renderedHash: 'sha256:sample-rendered-prompt', + }, + capabilities: { + provisionalCapabilityIds: [], + resolvedCapabilityAccess: [], + }, + warnings: [], + }, }); mockCreateQueuedRunWithMessage.mockImplementation(async ({ message }) => { const storedMessage = { @@ -591,6 +649,17 @@ describe('canonical agent session API acceptance flow', () => { const runId = runBody.data.run.id; expect(runResponse.status).toBe(201); + expect(mockResolveForRunAdmission).toHaveBeenCalledWith( + expect.objectContaining({ + thread: expect.objectContaining({ uuid: 'thread-1' }), + session: expect.objectContaining({ uuid: 'session-1' }), + source: expect.objectContaining({ uuid: 'source-1', adapter: 'blank_workspace', status: 'ready' }), + userIdentity: sampleUser, + requestedProvider: null, + requestedModel: null, + runtimeOptions: {}, + }) + ); expect(runBody.data).toEqual( expect.objectContaining({ run: { diff --git a/src/app/api/v2/ai/agent/api-keys/route.test.ts b/src/app/api/v2/ai/agent/api-keys/route.test.ts index 95cd6be3..8682aa18 100644 --- a/src/app/api/v2/ai/agent/api-keys/route.test.ts +++ b/src/app/api/v2/ai/agent/api-keys/route.test.ts @@ -30,7 +30,7 @@ jest.mock('server/services/userApiKey', () => ({ }, })); -jest.mock('server/services/aiAgentConfig', () => ({ +jest.mock('server/services/agentRuntime/config/agentRuntimeConfig', () => ({ __esModule: true, default: { getInstance: jest.fn(() => ({ diff --git a/src/app/api/v2/ai/agent/api-keys/route.ts b/src/app/api/v2/ai/agent/api-keys/route.ts index 54c52ad7..7c3f9bff 100644 --- a/src/app/api/v2/ai/agent/api-keys/route.ts +++ b/src/app/api/v2/ai/agent/api-keys/route.ts @@ -18,7 +18,7 @@ import { NextRequest } from 'next/server'; import { createApiHandler } from 'server/lib/createApiHandler'; import { successResponse, errorResponse } from 'server/lib/response'; import { getRequestUserIdentity } from 'server/lib/get-user'; -import AIAgentConfigService from 'server/services/aiAgentConfig'; +import AgentRuntimeConfigService from 'server/services/agentRuntime/config/agentRuntimeConfig'; import UserApiKeyService from 'server/services/userApiKey'; import { STORED_AGENT_PROVIDER_NAMES, @@ -45,7 +45,7 @@ function getSearchParam(req: NextRequest, key: string): string | null { async function getConfiguredProviders(): Promise { try { - const config = await AIAgentConfigService.getInstance().getEffectiveConfig(); + const config = await AgentRuntimeConfigService.getInstance().getEffectiveConfig(); const configuredProviders = (config.providers || []) .map((provider: { name?: unknown; enabled?: unknown }) => provider.enabled !== false && typeof provider.name === 'string' ? normalizeProvider(provider.name) : null diff --git a/src/app/api/v2/ai/agent/build-context-chats/route.test.ts b/src/app/api/v2/ai/agent/build-context-chats/route.test.ts new file mode 100644 index 00000000..10a5834e --- /dev/null +++ b/src/app/api/v2/ai/agent/build-context-chats/route.test.ts @@ -0,0 +1,270 @@ +/** + * 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/BuildContextChatService', () => { + class BuildContextChatBuildNotFoundError extends Error { + constructor(readonly buildUuid: string) { + super(`Build not found: ${buildUuid}`); + this.name = 'BuildContextChatBuildNotFoundError'; + } + } + + return { + __esModule: true, + default: { + launchBuildContextChat: jest.fn(), + }, + BuildContextChatBuildNotFoundError, + }; +}); + +jest.mock('server/services/agent/ProviderRegistry', () => ({ + AgentModelSelectionError: class AgentModelSelectionError extends Error { + constructor(message: string) { + super(message); + this.name = 'AgentModelSelectionError'; + } + }, + MissingAgentProviderApiKeyError: class MissingAgentProviderApiKeyError extends Error { + constructor(readonly provider: string) { + super(`No stored API key is configured for provider "${provider}".`); + this.name = 'MissingAgentProviderApiKeyError'; + } + }, +})); + +jest.mock('server/services/agent/SessionReadService', () => ({ + __esModule: true, + default: { + serializeSessionRecord: jest.fn(), + serializeThread: jest.fn(), + }, +})); + +import { POST } from './route'; +import { getRequestUserIdentity } from 'server/lib/get-user'; +import BuildContextChatService, { + BuildContextChatBuildNotFoundError, +} from 'server/services/agent/BuildContextChatService'; +import { AgentModelSelectionError, MissingAgentProviderApiKeyError } from 'server/services/agent/ProviderRegistry'; +import AgentSessionReadService from 'server/services/agent/SessionReadService'; + +const mockGetRequestUserIdentity = getRequestUserIdentity as jest.Mock; +const mockLaunchBuildContextChat = BuildContextChatService.launchBuildContextChat as jest.Mock; +const mockSerializeSessionRecord = AgentSessionReadService.serializeSessionRecord as jest.Mock; +const mockSerializeThread = AgentSessionReadService.serializeThread 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/build-context-chats'), + } as unknown as NextRequest; +} + +function mockSuccessfulLaunch(overrides: { created?: boolean; reused?: boolean } = {}) { + const session = { id: 17, uuid: 'session-1' }; + const thread = { id: 29, uuid: 'thread-1' }; + mockLaunchBuildContextChat.mockResolvedValue({ + session, + thread, + created: overrides.created ?? true, + reused: overrides.reused ?? false, + buildContext: { + buildUuid: 'build-1', + buildKind: 'environment', + namespace: 'env-sample', + baseBuildUuid: 'base-build-1', + pullRequest: { + fullName: 'example-org/example-repo', + branchName: 'feature/sample', + pullRequestNumber: 123, + }, + contextFreshAt: '2026-04-30T00:00:00.000Z', + }, + }); + mockSerializeSessionRecord.mockResolvedValue({ + session: { + id: 'session-1', + status: 'ready', + userId: 'sample-user', + ownerGithubUsername: 'sample-user', + defaults: { model: 'gpt-5.4', harness: 'lifecycle_ai_sdk' }, + defaultThreadId: 'thread-1', + }, + source: { id: 'source-1', adapter: 'blank_workspace', status: 'ready' }, + sandbox: { id: null, status: 'none' }, + }); + mockSerializeThread.mockResolvedValue({ + id: 'thread-1', + sessionId: 'session-1', + title: 'Default thread', + isDefault: true, + metadata: {}, + session: { id: 'session-1' }, + }); +} + +describe('POST /api/v2/ai/agent/build-context-chats', () => { + beforeEach(() => { + jest.clearAllMocks(); + mockGetRequestUserIdentity.mockReturnValue({ + userId: 'sample-user', + githubUsername: 'sample-user', + }); + mockSuccessfulLaunch(); + }); + + it('returns 401 for unauthenticated route access', async () => { + mockGetRequestUserIdentity.mockReturnValueOnce(null); + + const response = await POST(makeRequest({ buildUuid: 'build-1' })); + + expect(response.status).toBe(401); + expect(mockLaunchBuildContextChat).not.toHaveBeenCalled(); + }); + + it('returns 400 for non-object launch bodies', async () => { + const response = await POST(makeRequest(null)); + + expect(response.status).toBe(400); + expect(mockLaunchBuildContextChat).not.toHaveBeenCalled(); + }); + + it.each([ + ['missing buildUuid', {}], + ['blank buildUuid', { buildUuid: ' ' }], + ['non-string defaults.model', { buildUuid: 'build-1', defaults: { model: 123 } }], + ['unsupported defaults key', { buildUuid: 'build-1', defaults: { model: 'gpt-5.4', provider: 'openai' } }], + ['unsupported source field', { buildUuid: 'build-1', source: { adapter: 'blank_workspace' } }], + ['unsupported workspace field', { buildUuid: 'build-1', workspace: {} }], + ['unsupported thread field', { buildUuid: 'build-1', thread: {} }], + ['unsupported sandbox field', { buildUuid: 'build-1', sandbox: {} }], + ['unsupported services field', { buildUuid: 'build-1', services: [] }], + ['unsupported message field', { buildUuid: 'build-1', message: 'hello' }], + ])('returns 400 for %s', async (_name, body) => { + const response = await POST(makeRequest(body)); + + expect(response.status).toBe(400); + expect(mockLaunchBuildContextChat).not.toHaveBeenCalled(); + }); + + it('maps invalid or unknown buildUuid values to 404', async () => { + mockLaunchBuildContextChat.mockRejectedValueOnce(new BuildContextChatBuildNotFoundError('missing-build')); + + const response = await POST(makeRequest({ buildUuid: 'missing-build' })); + const body = await response.json(); + + expect(response.status).toBe(404); + expect(body.error.message).toBe('Build not found: missing-build'); + }); + + it('maps missing provider API keys to 400', async () => { + mockLaunchBuildContextChat.mockRejectedValueOnce(new MissingAgentProviderApiKeyError('openai')); + + const response = await POST(makeRequest({ buildUuid: 'build-1' })); + + expect(response.status).toBe(400); + }); + + it('maps invalid requested model errors to 400', async () => { + mockLaunchBuildContextChat.mockRejectedValueOnce(new AgentModelSelectionError('Model sample-model is not enabled')); + + const response = await POST(makeRequest({ buildUuid: 'build-1', defaults: { model: 'sample-model' } })); + const body = await response.json(); + + expect(response.status).toBe(400); + expect(body.error.message).toBe('Model sample-model is not enabled'); + }); + + it('delegates launch requests and returns 201 for created chats with execution links', async () => { + const response = await POST(makeRequest({ buildUuid: ' build-1 ', defaults: { model: ' gpt-5.4 ' } })); + const body = await response.json(); + + expect(response.status).toBe(201); + expect(mockLaunchBuildContextChat).toHaveBeenCalledWith({ + buildUuid: 'build-1', + userId: 'sample-user', + userIdentity: { + userId: 'sample-user', + githubUsername: 'sample-user', + }, + model: 'gpt-5.4', + }); + expect(mockSerializeSessionRecord).toHaveBeenCalledWith({ id: 17, uuid: 'session-1' }); + expect(mockSerializeThread).toHaveBeenCalledWith({ id: 29, uuid: 'thread-1' }, { id: 17, uuid: 'session-1' }); + expect(body.data).toEqual( + expect.objectContaining({ + created: true, + reused: false, + buildContext: { + buildUuid: 'build-1', + buildKind: 'environment', + namespace: 'env-sample', + baseBuildUuid: 'base-build-1', + repo: 'example-org/example-repo', + branch: 'feature/sample', + pullRequestNumber: 123, + contextFreshAt: '2026-04-30T00:00:00.000Z', + }, + session: { + session: { + id: 'session-1', + status: 'ready', + userId: 'sample-user', + ownerGithubUsername: 'sample-user', + defaults: { model: 'gpt-5.4', harness: 'lifecycle_ai_sdk' }, + defaultThreadId: 'thread-1', + }, + source: { id: 'source-1', adapter: 'blank_workspace', status: 'ready' }, + sandbox: { id: null, status: 'none' }, + }, + thread: { + id: 'thread-1', + sessionId: 'session-1', + title: 'Default thread', + isDefault: true, + metadata: {}, + session: { id: 'session-1' }, + }, + links: { + messages: '/api/v2/ai/agent/threads/thread-1/messages', + runs: '/api/v2/ai/agent/threads/thread-1/runs', + events: '/api/v2/ai/agent/runs/{runId}/events', + eventStream: '/api/v2/ai/agent/runs/{runId}/events/stream', + pendingActions: '/api/v2/ai/agent/threads/thread-1/pending-actions', + }, + }) + ); + }); + + it('returns 200 for reused chats', async () => { + mockSuccessfulLaunch({ created: false, reused: true }); + + const response = await POST(makeRequest({ buildUuid: 'build-1' })); + const body = await response.json(); + + expect(response.status).toBe(200); + expect(body.data.created).toBe(false); + expect(body.data.reused).toBe(true); + }); +}); diff --git a/src/app/api/v2/ai/agent/build-context-chats/route.ts b/src/app/api/v2/ai/agent/build-context-chats/route.ts new file mode 100644 index 00000000..c9afbd8f --- /dev/null +++ b/src/app/api/v2/ai/agent/build-context-chats/route.ts @@ -0,0 +1,211 @@ +/** + * 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 BuildContextChatService, { + BuildContextChatBuildNotFoundError, +} from 'server/services/agent/BuildContextChatService'; +import { AgentModelSelectionError, MissingAgentProviderApiKeyError } from 'server/services/agent/ProviderRegistry'; +import AgentSessionReadService from 'server/services/agent/SessionReadService'; + +interface CreateBuildContextChatBody { + buildUuid: string; + defaults?: { + model?: string; + }; +} + +const ALLOWED_TOP_LEVEL_KEYS = ['buildUuid', 'defaults']; +const ALLOWED_DEFAULT_KEYS = ['model']; + +function unknownKeys(value: Record, allowedKeys: string[]) { + return Object.keys(value).filter((key) => !allowedKeys.includes(key)); +} + +function parseCreateBuildContextChatBody(body: unknown): CreateBuildContextChatBody { + if (!body || typeof body !== 'object' || Array.isArray(body)) { + throw new Error('Request body must be an object'); + } + + const requestBody = body as Record; + const unsupportedKeys = unknownKeys(requestBody, ALLOWED_TOP_LEVEL_KEYS); + if (unsupportedKeys.length > 0) { + throw new Error(`Unsupported build-context chat fields: ${unsupportedKeys.join(', ')}`); + } + + if (typeof requestBody.buildUuid !== 'string' || !requestBody.buildUuid.trim()) { + throw new Error('buildUuid must be a non-empty string'); + } + + if (requestBody.defaults === undefined) { + return { + buildUuid: requestBody.buildUuid.trim(), + }; + } + + if (!requestBody.defaults || typeof requestBody.defaults !== 'object' || Array.isArray(requestBody.defaults)) { + throw new Error('defaults must be an object'); + } + + const defaults = requestBody.defaults as Record; + const unsupportedDefaultKeys = unknownKeys(defaults, ALLOWED_DEFAULT_KEYS); + if (unsupportedDefaultKeys.length > 0) { + throw new Error(`Unsupported build-context chat defaults fields: ${unsupportedDefaultKeys.join(', ')}`); + } + + if (defaults.model !== undefined && typeof defaults.model !== 'string') { + throw new Error('defaults.model must be a string'); + } + + const requestedModel = defaults.model?.trim() || undefined; + return { + buildUuid: requestBody.buildUuid.trim(), + ...(requestedModel ? { defaults: { model: requestedModel } } : {}), + }; +} + +/** + * @openapi + * /api/v2/ai/agent/build-context-chats: + * post: + * summary: Create or reuse a build-context agent chat + * description: Creates or reuses a build-context chat session and default thread. This endpoint does not submit a message or run. + * tags: + * - Agent Sessions + * operationId: createBuildContextAgentChat + * requestBody: + * required: true + * content: + * application/json: + * schema: + * $ref: '#/components/schemas/CreateBuildContextAgentChatRequest' + * responses: + * '200': + * description: Existing build-context chat reused + * content: + * application/json: + * schema: + * allOf: + * - $ref: '#/components/schemas/SuccessApiResponse' + * - type: object + * required: [data] + * properties: + * data: + * $ref: '#/components/schemas/BuildContextAgentChatResponse' + * '201': + * description: Build-context chat created + * content: + * application/json: + * schema: + * allOf: + * - $ref: '#/components/schemas/SuccessApiResponse' + * - type: object + * required: [data] + * properties: + * data: + * $ref: '#/components/schemas/BuildContextAgentChatResponse' + * '400': + * description: Invalid request or model configuration + * content: + * application/json: + * schema: + * $ref: '#/components/schemas/ApiErrorResponse' + * '401': + * description: Unauthorized + * content: + * application/json: + * schema: + * $ref: '#/components/schemas/ApiErrorResponse' + * '404': + * description: Build not found + * content: + * application/json: + * schema: + * $ref: '#/components/schemas/ApiErrorResponse' + */ +const postHandler = async (req: NextRequest) => { + const userIdentity = getRequestUserIdentity(req); + if (!userIdentity) { + return errorResponse(new Error('Unauthorized'), { status: 401 }, req); + } + + let requestBody: CreateBuildContextChatBody; + try { + requestBody = parseCreateBuildContextChatBody(await req.json().catch(() => null)); + } catch (error) { + return errorResponse(error, { status: 400 }, req); + } + + try { + const result = await BuildContextChatService.launchBuildContextChat({ + buildUuid: requestBody.buildUuid, + userId: userIdentity.userId, + userIdentity, + model: requestBody.defaults?.model, + }); + const [session, thread] = await Promise.all([ + AgentSessionReadService.serializeSessionRecord(result.session), + AgentSessionReadService.serializeThread(result.thread, result.session), + ]); + const threadId = thread.id; + + return successResponse( + { + session, + thread, + created: result.created, + reused: result.reused, + buildContext: { + buildUuid: result.buildContext.buildUuid, + buildKind: result.buildContext.buildKind, + namespace: result.buildContext.namespace, + baseBuildUuid: result.buildContext.baseBuildUuid, + repo: result.buildContext.pullRequest?.fullName ?? null, + branch: result.buildContext.pullRequest?.branchName ?? null, + pullRequestNumber: result.buildContext.pullRequest?.pullRequestNumber ?? null, + contextFreshAt: result.buildContext.contextFreshAt, + }, + links: { + messages: `/api/v2/ai/agent/threads/${threadId}/messages`, + runs: `/api/v2/ai/agent/threads/${threadId}/runs`, + events: '/api/v2/ai/agent/runs/{runId}/events', + eventStream: '/api/v2/ai/agent/runs/{runId}/events/stream', + pendingActions: `/api/v2/ai/agent/threads/${threadId}/pending-actions`, + }, + }, + { status: result.created ? 201 : 200 }, + req + ); + } catch (error) { + if (error instanceof BuildContextChatBuildNotFoundError) { + return errorResponse(error, { status: 404 }, req); + } + if (error instanceof MissingAgentProviderApiKeyError) { + return errorResponse(error, { status: 400 }, req); + } + if (error instanceof AgentModelSelectionError) { + return errorResponse(error, { status: 400 }, req); + } + + throw error; + } +}; + +export const POST = createApiHandler(postHandler); diff --git a/src/app/api/v2/ai/agent/definition-capabilities/route.test.ts b/src/app/api/v2/ai/agent/definition-capabilities/route.test.ts new file mode 100644 index 00000000..6ac21a9c --- /dev/null +++ b/src/app/api/v2/ai/agent/definition-capabilities/route.test.ts @@ -0,0 +1,189 @@ +/** + * 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 mockGetRequestUserIdentity = jest.fn(); +const mockGetUserDefinitionCreationStatus = jest.fn(); +const mockListUserSelectableCapabilities = jest.fn(); + +jest.mock('server/lib/get-user', () => ({ + getRequestUserIdentity: (...args: unknown[]) => mockGetRequestUserIdentity(...args), +})); + +jest.mock('server/services/agent/CustomAgentDefinitionService', () => ({ + __esModule: true, + customAgentDefinitionService: { + getUserDefinitionCreationStatus: (...args: unknown[]) => mockGetUserDefinitionCreationStatus(...args), + listUserSelectableCapabilities: (...args: unknown[]) => mockListUserSelectableCapabilities(...args), + }, +})); + +import { GET } from './route'; + +const capabilityRows = [ + { + capabilityId: 'read_context', + label: 'Read/context', + description: 'Read session context.', + category: 'read', + toolCount: 0, + resourceCount: 1, + requiresWorkspace: false, + tools: [], + resources: [{ name: 'Session context' }], + }, + { + capabilityId: 'workspace_shell', + label: 'Command tools', + description: 'Run shell commands inside a development workspace.', + category: 'workspace', + toolCount: 1, + resourceCount: 1, + requiresWorkspace: true, + tools: [{ name: 'Workspace exec', description: null }], + resources: [{ name: 'Workspace shell' }], + }, +]; + +function makeRequest(url = 'http://localhost/api/v2/ai/agent/definition-capabilities'): NextRequest { + return { + headers: new Headers([['x-request-id', 'req-test']]), + nextUrl: new URL(url), + } as unknown as NextRequest; +} + +describe('/api/v2/ai/agent/definition-capabilities', () => { + beforeEach(() => { + jest.clearAllMocks(); + mockGetRequestUserIdentity.mockReturnValue({ + userId: 'sample-user', + githubUsername: 'sample-user', + }); + mockGetUserDefinitionCreationStatus.mockResolvedValue({ + canCreate: true, + creationUnavailableReason: null, + }); + mockListUserSelectableCapabilities.mockResolvedValue(capabilityRows); + }); + + it('returns 401 for unauthenticated requests', async () => { + mockGetRequestUserIdentity.mockReturnValueOnce(null); + + const response = await GET(makeRequest()); + + expect(response.status).toBe(401); + expect(mockListUserSelectableCapabilities).not.toHaveBeenCalled(); + }); + + it('defaults to current_workspace_when_available and delegates userSelectable inventory loading', async () => { + const response = await GET(makeRequest()); + const body = await response.json(); + + expect(response.status).toBe(200); + expect(mockGetUserDefinitionCreationStatus).toHaveBeenCalledWith({ + userIdentity: { + userId: 'sample-user', + githubUsername: 'sample-user', + }, + }); + expect(mockListUserSelectableCapabilities).toHaveBeenCalledWith({ + userIdentity: { + userId: 'sample-user', + githubUsername: 'sample-user', + }, + resourceBehavior: 'current_workspace_when_available', + }); + expect(body.data).toEqual({ + resourceBehavior: 'current_workspace_when_available', + canCreate: true, + creationUnavailableReason: null, + capabilities: capabilityRows, + }); + }); + + it('returns blocked creation status without loading capabilities', async () => { + mockGetUserDefinitionCreationStatus.mockResolvedValueOnce({ + canCreate: false, + creationUnavailableReason: 'creation_restricted', + }); + + const response = await GET(makeRequest()); + const body = await response.json(); + + expect(response.status).toBe(200); + expect(mockListUserSelectableCapabilities).not.toHaveBeenCalled(); + expect(body.data).toEqual({ + resourceBehavior: 'current_workspace_when_available', + canCreate: false, + creationUnavailableReason: 'creation_restricted', + capabilities: [], + }); + }); + + it('passes chat_only resourceBehavior through to omit source-incompatible capabilities in the service', async () => { + mockListUserSelectableCapabilities.mockResolvedValueOnce([capabilityRows[0]]); + + const response = await GET( + makeRequest('http://localhost/api/v2/ai/agent/definition-capabilities?resourceBehavior=chat_only') + ); + const body = await response.json(); + + expect(response.status).toBe(200); + expect(mockListUserSelectableCapabilities).toHaveBeenCalledWith( + expect.objectContaining({ resourceBehavior: 'chat_only' }) + ); + expect(body.data.capabilities.map((capability: { capabilityId: string }) => capability.capabilityId)).toEqual([ + 'read_context', + ]); + }); + + it('rejects unsupported resourceBehavior values', async () => { + const response = await GET( + makeRequest('http://localhost/api/v2/ai/agent/definition-capabilities?resourceBehavior=workspace_only') + ); + + expect(response.status).toBe(400); + expect(mockGetUserDefinitionCreationStatus).not.toHaveBeenCalled(); + expect(mockListUserSelectableCapabilities).not.toHaveBeenCalled(); + }); + + it('does not include admin_only, system_only, disabled, hidden count, toolKey, or serverSlug response fields', async () => { + mockListUserSelectableCapabilities.mockResolvedValueOnce([ + { + ...capabilityRows[0], + hiddenRestrictedCount: 2, + toolKey: 'mcp__sample__read_secret', + serverSlug: 'sample-internal-mcp', + defaultAvailability: 'admin_only', + }, + ]); + + const response = await GET(makeRequest()); + const body = await response.json(); + const serialized = JSON.stringify(body.data); + + expect(response.status).toBe(200); + expect(body.data.capabilities).toEqual([capabilityRows[0]]); + expect(serialized).toContain('requiresWorkspace'); + expect(serialized).not.toContain('admin_only'); + expect(serialized).not.toContain('system_only'); + expect(serialized).not.toContain('disabled'); + expect(serialized).not.toContain('hiddenRestricted'); + expect(serialized).not.toContain('toolKey'); + expect(serialized).not.toContain('serverSlug'); + }); +}); diff --git a/src/app/api/v2/ai/agent/definition-capabilities/route.ts b/src/app/api/v2/ai/agent/definition-capabilities/route.ts new file mode 100644 index 00000000..4db7b587 --- /dev/null +++ b/src/app/api/v2/ai/agent/definition-capabilities/route.ts @@ -0,0 +1,128 @@ +/** + * 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 { getRequestUserIdentity } from 'server/lib/get-user'; +import { errorResponse, successResponse } from 'server/lib/response'; +import { + customAgentDefinitionService, + type UserAgentDefinitionCapability, +} from 'server/services/agent/CustomAgentDefinitionService'; +import type { UserAgentDefinitionResourceBehavior } from 'server/services/agent/agentDefinitionTypes'; + +export const dynamic = 'force-dynamic'; + +const DEFAULT_RESOURCE_BEHAVIOR: UserAgentDefinitionResourceBehavior = 'current_workspace_when_available'; +const RESOURCE_BEHAVIORS = new Set([ + 'chat_only', + 'current_workspace_when_available', +]); + +function parseResourceBehavior(req: NextRequest): UserAgentDefinitionResourceBehavior | Error { + const value = req.nextUrl.searchParams.get('resourceBehavior')?.trim() || DEFAULT_RESOURCE_BEHAVIOR; + if (!RESOURCE_BEHAVIORS.has(value as UserAgentDefinitionResourceBehavior)) { + return new Error('resourceBehavior must be chat_only or current_workspace_when_available.'); + } + + return value as UserAgentDefinitionResourceBehavior; +} + +function serializeCapability(capability: UserAgentDefinitionCapability): UserAgentDefinitionCapability { + return { + capabilityId: capability.capabilityId, + label: capability.label, + description: capability.description, + category: capability.category, + toolCount: capability.toolCount, + resourceCount: capability.resourceCount, + requiresWorkspace: capability.requiresWorkspace, + tools: capability.tools.map((tool) => ({ + name: tool.name, + description: tool.description, + })), + resources: capability.resources.map((resource) => ({ + name: resource.name, + description: resource.description, + })), + }; +} + +/** + * @openapi + * /api/v2/ai/agent/definition-capabilities: + * get: + * summary: List capabilities available for custom agent definitions + * description: Returns creator eligibility and only capabilities the current user can select for private custom agents. + * tags: + * - Agent Platform + * operationId: listUserAgentDefinitionCapabilities + * parameters: + * - in: query + * name: resourceBehavior + * schema: + * $ref: '#/components/schemas/UserAgentDefinitionResourceBehavior' + * description: Resource behavior to use when filtering source-compatible capabilities. + * responses: + * '200': + * description: User-visible capability inventory. + * content: + * application/json: + * schema: + * $ref: '#/components/schemas/UserAgentDefinitionCapabilitiesSuccessResponse' + * '400': + * description: Invalid resource behavior. + * content: + * application/json: + * schema: + * $ref: '#/components/schemas/ApiErrorResponse' + * '401': + * description: Unauthorized. + * content: + * application/json: + * schema: + * $ref: '#/components/schemas/ApiErrorResponse' + */ +const getHandler = async (req: NextRequest) => { + const userIdentity = getRequestUserIdentity(req); + if (!userIdentity) { + return errorResponse(new Error('Unauthorized'), { status: 401 }, req); + } + + const resourceBehavior = parseResourceBehavior(req); + if (resourceBehavior instanceof Error) { + return errorResponse(resourceBehavior, { status: 400 }, req); + } + + const creationStatus = await customAgentDefinitionService.getUserDefinitionCreationStatus({ userIdentity }); + if (!creationStatus.canCreate) { + return successResponse({ resourceBehavior, ...creationStatus, capabilities: [] }, { status: 200 }, req); + } + + const capabilities = await customAgentDefinitionService.listUserSelectableCapabilities({ + userIdentity, + resourceBehavior, + }); + + return successResponse( + { resourceBehavior, ...creationStatus, capabilities: capabilities.map(serializeCapability) }, + { status: 200 }, + req + ); +}; + +export const GET = createApiHandler(getHandler); diff --git a/src/app/api/v2/ai/agent/definitions/[definitionId]/route.test.ts b/src/app/api/v2/ai/agent/definitions/[definitionId]/route.test.ts new file mode 100644 index 00000000..02168b1f --- /dev/null +++ b/src/app/api/v2/ai/agent/definitions/[definitionId]/route.test.ts @@ -0,0 +1,273 @@ +/** + * 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 mockGetRequestUserIdentity = jest.fn(); +const mockGetUserDefinition = jest.fn(); +const mockUpdateUserDefinition = jest.fn(); +const mockArchiveUserDefinition = jest.fn(); + +jest.mock('server/lib/get-user', () => ({ + getRequestUserIdentity: (...args: unknown[]) => mockGetRequestUserIdentity(...args), +})); + +jest.mock('server/services/agent/CustomAgentDefinitionService', () => { + class CustomAgentDefinitionServiceError extends Error { + constructor(public readonly code: string, message: string) { + super(message); + this.name = 'CustomAgentDefinitionServiceError'; + } + } + + return { + __esModule: true, + CustomAgentDefinitionServiceError, + customAgentDefinitionService: { + getUserDefinition: (...args: unknown[]) => mockGetUserDefinition(...args), + updateUserDefinition: (...args: unknown[]) => mockUpdateUserDefinition(...args), + archiveUserDefinition: (...args: unknown[]) => mockArchiveUserDefinition(...args), + }, + serializeUserAgentDefinition: (definition: any) => ({ + id: definition.id, + version: definition.version, + name: definition.name, + description: definition.description, + instructions: definition.instructionAddendum || '', + capabilityIds: definition.optionalCapabilityRefs?.length + ? definition.optionalCapabilityRefs + : definition.capabilityRefs, + modelPreference: definition.modelPreference || null, + resourceBehavior: definition.resourcePolicy?.sourceKinds?.includes('workspace_session') + ? 'current_workspace_when_available' + : 'chat_only', + status: definition.status === 'archived' ? 'archived' : 'active', + }), + }; +}); + +import { DELETE, GET, PATCH } from './route'; +import { CustomAgentDefinitionServiceError } from 'server/services/agent/CustomAgentDefinitionService'; + +const sampleDefinition = { + id: 'custom.sample-agent', + version: 1, + owner: { kind: 'user', userId: 'sample-user', organizationId: null }, + name: 'Sample agent', + description: 'Helps with sample workflows.', + instructionRefs: [], + instructionAddendum: 'Answer with concise steps.', + capabilityRefs: ['read_context'], + requiredCapabilityRefs: [], + optionalCapabilityRefs: ['read_context'], + resourcePolicy: { + sourceKinds: ['freeform_chat'], + workspaceRequired: false, + sandboxRequired: false, + }, + modelPreference: null, + status: 'active', + codeOwned: false, + readOnly: false, +}; + +const publicDefinition = { + id: 'custom.sample-agent', + version: 1, + name: 'Sample agent', + description: 'Helps with sample workflows.', + instructions: 'Answer with concise steps.', + capabilityIds: ['read_context'], + modelPreference: null, + resourceBehavior: 'chat_only', + status: 'active', +}; + +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/definitions/custom.sample-agent'), + } as unknown as NextRequest; +} + +const params = { params: { definitionId: 'custom.sample-agent' } }; + +describe('/api/v2/ai/agent/definitions/[definitionId]', () => { + beforeEach(() => { + jest.clearAllMocks(); + mockGetRequestUserIdentity.mockReturnValue({ + userId: 'sample-user', + githubUsername: 'sample-user', + }); + mockGetUserDefinition.mockResolvedValue(sampleDefinition); + mockUpdateUserDefinition.mockResolvedValue({ ...sampleDefinition, version: 2 }); + mockArchiveUserDefinition.mockResolvedValue({ ...sampleDefinition, status: 'archived' }); + }); + + it('GET returns 401 for unauthenticated requests', async () => { + mockGetRequestUserIdentity.mockReturnValueOnce(null); + + const response = await GET(makeRequest(), params); + + expect(response.status).toBe(401); + expect(mockGetUserDefinition).not.toHaveBeenCalled(); + }); + + it('GET returns 404 for non-owned or archived rows', async () => { + mockGetUserDefinition.mockRejectedValueOnce(new CustomAgentDefinitionServiceError('not_found', 'Agent not found.')); + + const response = await GET(makeRequest(), params); + const body = await response.json(); + + expect(response.status).toBe(404); + expect(mockGetUserDefinition).toHaveBeenCalledWith('custom.sample-agent', 'sample-user'); + expect(body.error.message).toBe('Agent not found.'); + }); + + it('PATCH rejects ownerKind, ownerUserId, codeOwned, readOnly, and other unsupported fields', async () => { + const response = await PATCH( + makeRequest({ + name: 'Sample agent', + instructions: 'Answer briefly.', + resourceBehavior: 'chat_only', + capabilityIds: ['read_context'], + ownerKind: 'system', + ownerUserId: 'other-user', + codeOwned: true, + readOnly: true, + unsupportedField: true, + }), + params + ); + const body = await response.json(); + + expect(response.status).toBe(400); + expect(body.error.message).toContain('Unsupported agent definition fields'); + expect(mockUpdateUserDefinition).not.toHaveBeenCalled(); + }); + + it('PATCH delegates normalized update input to CustomAgentDefinitionService', async () => { + const response = await PATCH( + makeRequest({ + name: 'Updated helper', + description: 'Summarizes release notes.', + instructions: 'Prefer bullets.', + capabilityIds: ['read_context'], + modelPreference: null, + resourceBehavior: 'current_workspace_when_available', + }), + params + ); + + expect(response.status).toBe(200); + const body = await response.json(); + expect(mockUpdateUserDefinition).toHaveBeenCalledWith( + 'custom.sample-agent', + expect.objectContaining({ + userId: 'sample-user', + githubUsername: 'sample-user', + }), + { + name: 'Updated helper', + description: 'Summarizes release notes.', + instructionAddendum: 'Prefer bullets.', + capabilityRefs: ['read_context'], + modelPreference: null, + resourceBehavior: 'current_workspace_when_available', + } + ); + expect(body.data.definition).toEqual({ + ...publicDefinition, + version: 2, + }); + }); + + it.each([ + ['unknown_capability', 'sample_unknown_capability'], + ['admin_only', 'external_mcp_write'], + ['system_only', 'approval_controls'], + ['disabled', 'read_context'], + ['source_incompatible', 'workspace_shell'], + ])( + 'PATCH maps %s capability validation failures to deterministic sanitized 400 responses', + async (code, capabilityId) => { + mockUpdateUserDefinition.mockRejectedValueOnce( + new CustomAgentDefinitionServiceError( + code as any, + 'Some selected capabilities are no longer available. Review the list and save again.' + ) + ); + + const response = await PATCH( + makeRequest({ + name: 'Sample agent', + instructions: 'Answer briefly.', + capabilityIds: [capabilityId], + resourceBehavior: 'chat_only', + }), + params + ); + const body = await response.json(); + + expect(response.status).toBe(400); + expect(body.error.message).toBe( + 'Some selected capabilities are no longer available. Review the list and save again.' + ); + expect(body.error.message).not.toContain(capabilityId); + expect(body.error.message).not.toContain('mcp__'); + expect(body.error.message).not.toContain('toolKey'); + expect(body.error.message).not.toContain('serverSlug'); + } + ); + + it('PATCH maps creation policy denials to 403 with sanitized response messages', async () => { + mockUpdateUserDefinition.mockRejectedValueOnce( + new CustomAgentDefinitionServiceError( + 'creation_unavailable', + 'Custom agent creation is not available. Ask an admin for access.' + ) + ); + + const response = await PATCH( + makeRequest({ + name: 'Sample agent', + instructions: 'Answer briefly.', + capabilityIds: ['read_context'], + resourceBehavior: 'chat_only', + }), + params + ); + const body = await response.json(); + + expect(response.status).toBe(403); + expect(body.error.message).toBe('Custom agent creation is not available. Ask an admin for access.'); + expect(body.error.message).not.toContain('roles'); + expect(body.error.message).not.toContain('customAgentCreationPolicy'); + }); + + it('DELETE archives the current user definition and returns archived state', async () => { + const response = await DELETE(makeRequest(), params); + const body = await response.json(); + + expect(response.status).toBe(200); + expect(mockArchiveUserDefinition).toHaveBeenCalledWith('custom.sample-agent', 'sample-user'); + expect(body.data).toEqual({ + archived: true, + definition: { ...publicDefinition, status: 'archived' }, + }); + }); +}); diff --git a/src/app/api/v2/ai/agent/definitions/[definitionId]/route.ts b/src/app/api/v2/ai/agent/definitions/[definitionId]/route.ts new file mode 100644 index 00000000..b25d638d --- /dev/null +++ b/src/app/api/v2/ai/agent/definitions/[definitionId]/route.ts @@ -0,0 +1,384 @@ +/** + * 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 { getRequestUserIdentity } from 'server/lib/get-user'; +import { errorResponse, successResponse } from 'server/lib/response'; +import { + CustomAgentDefinitionServiceError, + customAgentDefinitionService, + serializeUserAgentDefinition, +} from 'server/services/agent/CustomAgentDefinitionService'; +import type { + AgentDefinitionModelPreference, + UserAgentDefinitionResourceBehavior, + UserAgentDefinitionUpsertInput, +} from 'server/services/agent/agentDefinitionTypes'; +import type { AgentCapabilityCatalogId } from 'server/services/agent/capabilityCatalog'; + +export const dynamic = 'force-dynamic'; + +const ALLOWED_REQUEST_FIELDS = new Set([ + 'name', + 'description', + 'instructions', + 'capabilityIds', + 'modelPreference', + 'resourceBehavior', +]); +const RESOURCE_BEHAVIORS = new Set([ + 'chat_only', + 'current_workspace_when_available', +]); + +function mapDefinitionError(error: unknown, req: NextRequest) { + if (error instanceof CustomAgentDefinitionServiceError) { + if (error.code === 'not_found') { + return errorResponse(error, { status: 404 }, req); + } + + if (error.code === 'model_unavailable') { + return errorResponse(error, { status: 409 }, req); + } + + if (error.code === 'creation_unavailable') { + return errorResponse(error, { status: 403 }, req); + } + + return errorResponse(error, { status: 400 }, req); + } + + throw error; +} + +async function readRequestBody(req: NextRequest): Promise | Error> { + let body: unknown; + try { + body = await req.json(); + } catch { + return new Error('Invalid JSON in request body'); + } + + if (!body || typeof body !== 'object' || Array.isArray(body)) { + return new Error('Request body must be an object.'); + } + + return body as Record; +} + +function readNullableString(value: unknown, fieldName: string): string | null | Error { + if (value === undefined || value === null) { + return null; + } + + if (typeof value !== 'string') { + return new Error(`${fieldName} must be a string.`); + } + + return value; +} + +function readRequiredString(value: unknown, fieldName: string): string | Error { + if (typeof value !== 'string') { + return new Error(`${fieldName} must be a string.`); + } + + return value; +} + +function readCapabilityIds(value: unknown): AgentCapabilityCatalogId[] | Error { + if (value === undefined) { + return []; + } + + if (!Array.isArray(value) || value.some((item) => typeof item !== 'string')) { + return new Error('capabilityIds must be an array of strings.'); + } + + return value as AgentCapabilityCatalogId[]; +} + +function readModelPreference(value: unknown): AgentDefinitionModelPreference | null | Error { + if (value === undefined || value === null) { + return null; + } + + if (typeof value !== 'object' || Array.isArray(value)) { + return new Error('modelPreference must be an object or null.'); + } + + const modelPreference = value as Record; + const unknownKeys = Object.keys(modelPreference).filter((key) => key !== 'provider' && key !== 'model'); + if (unknownKeys.length > 0) { + return new Error(`Unsupported modelPreference fields: ${unknownKeys.join(', ')}`); + } + + const provider = readNullableString(modelPreference.provider, 'modelPreference.provider'); + if (provider instanceof Error) { + return provider; + } + + const model = readNullableString(modelPreference.model, 'modelPreference.model'); + if (model instanceof Error) { + return model; + } + + return { provider, model }; +} + +function readResourceBehavior(value: unknown): UserAgentDefinitionResourceBehavior | Error { + if (typeof value !== 'string' || !RESOURCE_BEHAVIORS.has(value as UserAgentDefinitionResourceBehavior)) { + return new Error('resourceBehavior must be chat_only or current_workspace_when_available.'); + } + + return value as UserAgentDefinitionResourceBehavior; +} + +function parseUpsertBody(body: Record): UserAgentDefinitionUpsertInput | Error { + const unsupportedFields = Object.keys(body).filter((key) => !ALLOWED_REQUEST_FIELDS.has(key)); + if (unsupportedFields.length > 0) { + return new Error(`Unsupported agent definition fields: ${unsupportedFields.join(', ')}`); + } + + const name = readRequiredString(body.name, 'name'); + if (name instanceof Error) { + return name; + } + + const instructions = readRequiredString(body.instructions, 'instructions'); + if (instructions instanceof Error) { + return instructions; + } + + const description = readNullableString(body.description, 'description'); + if (description instanceof Error) { + return description; + } + + const capabilityRefs = readCapabilityIds(body.capabilityIds); + if (capabilityRefs instanceof Error) { + return capabilityRefs; + } + + const modelPreference = readModelPreference(body.modelPreference); + if (modelPreference instanceof Error) { + return modelPreference; + } + + const resourceBehavior = readResourceBehavior(body.resourceBehavior); + if (resourceBehavior instanceof Error) { + return resourceBehavior; + } + + return { + name, + description, + instructionAddendum: instructions, + capabilityRefs, + modelPreference, + resourceBehavior, + }; +} + +/** + * @openapi + * /api/v2/ai/agent/definitions/{definitionId}: + * get: + * summary: Get a current-user custom agent definition + * description: Returns one private custom agent definition owned by the authenticated user. + * tags: + * - Agent Platform + * operationId: getUserAgentDefinition + * parameters: + * - in: path + * name: definitionId + * required: true + * schema: + * type: string + * responses: + * '200': + * description: Custom agent definition. + * content: + * application/json: + * schema: + * $ref: '#/components/schemas/UserAgentDefinitionSuccessResponse' + * '401': + * description: Unauthorized. + * content: + * application/json: + * schema: + * $ref: '#/components/schemas/ApiErrorResponse' + * '404': + * description: Definition not found. + * content: + * application/json: + * schema: + * $ref: '#/components/schemas/ApiErrorResponse' + * patch: + * summary: Update a current-user custom agent definition + * description: Updates one private custom agent definition using the same validated definition payload as creation. + * tags: + * - Agent Platform + * operationId: updateUserAgentDefinition + * parameters: + * - in: path + * name: definitionId + * required: true + * schema: + * type: string + * requestBody: + * required: true + * content: + * application/json: + * schema: + * $ref: '#/components/schemas/UserAgentDefinitionUpsertRequest' + * responses: + * '200': + * description: Custom agent definition updated. + * content: + * application/json: + * schema: + * $ref: '#/components/schemas/UserAgentDefinitionSuccessResponse' + * '400': + * description: Invalid request or capability selection. + * content: + * application/json: + * schema: + * $ref: '#/components/schemas/ApiErrorResponse' + * '401': + * description: Unauthorized. + * content: + * application/json: + * schema: + * $ref: '#/components/schemas/ApiErrorResponse' + * '403': + * description: Custom-agent management is unavailable for the caller or selected capabilities. + * content: + * application/json: + * schema: + * $ref: '#/components/schemas/ApiErrorResponse' + * '404': + * description: Definition not found. + * content: + * application/json: + * schema: + * $ref: '#/components/schemas/ApiErrorResponse' + * '409': + * description: Selected model is unavailable. + * content: + * application/json: + * schema: + * $ref: '#/components/schemas/ApiErrorResponse' + * delete: + * summary: Archive a current-user custom agent definition + * description: Archives one private custom agent definition without affecting existing run history. + * tags: + * - Agent Platform + * operationId: deleteUserAgentDefinition + * parameters: + * - in: path + * name: definitionId + * required: true + * schema: + * type: string + * responses: + * '200': + * description: Custom agent definition archived. + * content: + * application/json: + * schema: + * $ref: '#/components/schemas/DeleteUserAgentDefinitionSuccessResponse' + * '401': + * description: Unauthorized. + * content: + * application/json: + * schema: + * $ref: '#/components/schemas/ApiErrorResponse' + * '404': + * description: Definition not found. + * content: + * application/json: + * schema: + * $ref: '#/components/schemas/ApiErrorResponse' + */ +const getHandler = async (req: NextRequest, { params }: { params: { definitionId: string } }) => { + const userIdentity = getRequestUserIdentity(req); + if (!userIdentity) { + return errorResponse(new Error('Unauthorized'), { status: 401 }, req); + } + + try { + const definition = await customAgentDefinitionService.getUserDefinition(params.definitionId, userIdentity.userId); + return successResponse({ definition: serializeUserAgentDefinition(definition) }, { status: 200 }, req); + } catch (error) { + return mapDefinitionError(error, req); + } +}; + +const patchHandler = async (req: NextRequest, { params }: { params: { definitionId: string } }) => { + const userIdentity = getRequestUserIdentity(req); + if (!userIdentity) { + return errorResponse(new Error('Unauthorized'), { status: 401 }, req); + } + + const body = await readRequestBody(req); + if (body instanceof Error) { + return errorResponse(body, { status: 400 }, req); + } + + const input = parseUpsertBody(body); + if (input instanceof Error) { + return errorResponse(input, { status: 400 }, req); + } + + try { + const definition = await customAgentDefinitionService.updateUserDefinition( + params.definitionId, + userIdentity, + input + ); + return successResponse({ definition: serializeUserAgentDefinition(definition) }, { status: 200 }, req); + } catch (error) { + return mapDefinitionError(error, req); + } +}; + +const deleteHandler = async (req: NextRequest, { params }: { params: { definitionId: string } }) => { + const userIdentity = getRequestUserIdentity(req); + if (!userIdentity) { + return errorResponse(new Error('Unauthorized'), { status: 401 }, req); + } + + try { + const definition = await customAgentDefinitionService.archiveUserDefinition( + params.definitionId, + userIdentity.userId + ); + return successResponse( + { archived: true, definition: serializeUserAgentDefinition(definition) }, + { status: 200 }, + req + ); + } catch (error) { + return mapDefinitionError(error, req); + } +}; + +export const GET = createApiHandler(getHandler); +export const PATCH = createApiHandler(patchHandler); +export const DELETE = createApiHandler(deleteHandler); diff --git a/src/app/api/v2/ai/agent/definitions/route.test.ts b/src/app/api/v2/ai/agent/definitions/route.test.ts new file mode 100644 index 00000000..2a65f392 --- /dev/null +++ b/src/app/api/v2/ai/agent/definitions/route.test.ts @@ -0,0 +1,269 @@ +/** + * 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 mockGetRequestUserIdentity = jest.fn(); +const mockListUserDefinitions = jest.fn(); +const mockCreateUserDefinition = jest.fn(); + +jest.mock('server/lib/get-user', () => ({ + getRequestUserIdentity: (...args: unknown[]) => mockGetRequestUserIdentity(...args), +})); + +jest.mock('server/services/agent/CustomAgentDefinitionService', () => { + class CustomAgentDefinitionServiceError extends Error { + constructor(public readonly code: string, message: string) { + super(message); + this.name = 'CustomAgentDefinitionServiceError'; + } + } + + return { + __esModule: true, + CustomAgentDefinitionServiceError, + customAgentDefinitionService: { + listUserDefinitions: (...args: unknown[]) => mockListUserDefinitions(...args), + createUserDefinition: (...args: unknown[]) => mockCreateUserDefinition(...args), + }, + serializeUserAgentDefinition: (definition: any) => ({ + id: definition.id, + version: definition.version, + name: definition.name, + description: definition.description, + instructions: definition.instructionAddendum || '', + capabilityIds: definition.optionalCapabilityRefs?.length + ? definition.optionalCapabilityRefs + : definition.capabilityRefs, + modelPreference: definition.modelPreference || null, + resourceBehavior: definition.resourcePolicy?.sourceKinds?.includes('workspace_session') + ? 'current_workspace_when_available' + : 'chat_only', + status: definition.status === 'archived' ? 'archived' : 'active', + }), + }; +}); + +import { GET, POST } from './route'; +import { CustomAgentDefinitionServiceError } from 'server/services/agent/CustomAgentDefinitionService'; + +const sampleDefinition = { + id: 'custom.sample-agent', + version: 1, + owner: { kind: 'user', userId: 'sample-user', organizationId: null }, + name: 'Sample agent', + description: 'Helps with sample workflows.', + instructionRefs: [], + instructionAddendum: 'Answer with concise steps.', + capabilityRefs: ['read_context'], + requiredCapabilityRefs: [], + optionalCapabilityRefs: ['read_context'], + resourcePolicy: { + sourceKinds: ['freeform_chat'], + workspaceRequired: false, + sandboxRequired: false, + }, + modelPreference: null, + status: 'active', + codeOwned: false, + readOnly: false, +}; + +const publicDefinition = { + id: 'custom.sample-agent', + version: 1, + name: 'Sample agent', + description: 'Helps with sample workflows.', + instructions: 'Answer with concise steps.', + capabilityIds: ['read_context'], + modelPreference: null, + resourceBehavior: 'chat_only', + status: 'active', +}; + +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/definitions'), + } as unknown as NextRequest; +} + +describe('/api/v2/ai/agent/definitions', () => { + beforeEach(() => { + jest.clearAllMocks(); + mockGetRequestUserIdentity.mockReturnValue({ + userId: 'sample-user', + githubUsername: 'sample-user', + }); + mockListUserDefinitions.mockResolvedValue([sampleDefinition]); + mockCreateUserDefinition.mockResolvedValue(sampleDefinition); + }); + + it('GET returns 401 for unauthenticated requests', async () => { + mockGetRequestUserIdentity.mockReturnValueOnce(null); + + const response = await GET(makeRequest()); + + expect(response.status).toBe(401); + expect(mockListUserDefinitions).not.toHaveBeenCalled(); + }); + + it('GET returns only current-user custom agents through the service', async () => { + const response = await GET(makeRequest()); + const body = await response.json(); + + expect(response.status).toBe(200); + expect(mockListUserDefinitions).toHaveBeenCalledWith({ userId: 'sample-user' }); + expect(body.data).toEqual({ definitions: [publicDefinition] }); + expect(JSON.stringify(body.data)).not.toContain('sample-user'); + expect(JSON.stringify(body.data)).not.toContain('other-user'); + }); + + it('POST rejects ownerKind, ownerUserId, codeOwned, readOnly, and other unsupported fields', async () => { + const response = await POST( + makeRequest({ + name: 'Sample agent', + instructions: 'Answer briefly.', + resourceBehavior: 'chat_only', + capabilityIds: ['read_context'], + ownerKind: 'system', + ownerUserId: 'other-user', + codeOwned: true, + readOnly: true, + unsupportedField: true, + }) + ); + const body = await response.json(); + + expect(response.status).toBe(400); + expect(body.error.message).toContain('Unsupported agent definition fields'); + expect(mockCreateUserDefinition).not.toHaveBeenCalled(); + }); + + it('POST delegates normalized create input to CustomAgentDefinitionService', async () => { + const response = await POST( + makeRequest({ + name: ' Release helper ', + description: ' Summarizes release notes. ', + instructions: ' Keep the response brief. ', + capabilityIds: ['read_context'], + modelPreference: { provider: 'openai', model: 'sample-model' }, + resourceBehavior: 'chat_only', + }) + ); + + expect(response.status).toBe(201); + const body = await response.json(); + expect(mockCreateUserDefinition).toHaveBeenCalledWith( + expect.objectContaining({ + userId: 'sample-user', + githubUsername: 'sample-user', + }), + { + name: ' Release helper ', + description: ' Summarizes release notes. ', + instructionAddendum: ' Keep the response brief. ', + capabilityRefs: ['read_context'], + modelPreference: { provider: 'openai', model: 'sample-model' }, + resourceBehavior: 'chat_only', + } + ); + expect(body.data.definition).toEqual(publicDefinition); + }); + + it.each([ + ['unknown_capability', 'sample_unknown_capability'], + ['admin_only', 'external_mcp_write'], + ['system_only', 'approval_controls'], + ['disabled', 'read_context'], + ['source_incompatible', 'workspace_shell'], + ])( + 'POST maps %s capability validation errors to 400 with sanitized response messages', + async (code, capabilityId) => { + mockCreateUserDefinition.mockRejectedValueOnce( + new CustomAgentDefinitionServiceError( + code as any, + 'Some selected capabilities are no longer available. Review the list and save again.' + ) + ); + + const response = await POST( + makeRequest({ + name: 'Sample agent', + instructions: 'Answer briefly.', + capabilityIds: [capabilityId], + resourceBehavior: 'chat_only', + }) + ); + const body = await response.json(); + + expect(response.status).toBe(400); + expect(body.error.message).toBe( + 'Some selected capabilities are no longer available. Review the list and save again.' + ); + expect(body.error.message).not.toContain(capabilityId); + expect(body.error.message).not.toContain('mcp__'); + expect(body.error.message).not.toContain('toolKey'); + expect(body.error.message).not.toContain('serverSlug'); + } + ); + + it('POST maps unavailable model errors to 409 with sanitized response messages', async () => { + mockCreateUserDefinition.mockRejectedValueOnce( + new CustomAgentDefinitionServiceError('model_unavailable', 'Selected model is unavailable. Choose another model.') + ); + + const response = await POST( + makeRequest({ + name: 'Sample agent', + instructions: 'Answer briefly.', + modelPreference: { provider: 'internal-provider', model: 'internal-model' }, + resourceBehavior: 'chat_only', + }) + ); + const body = await response.json(); + + expect(response.status).toBe(409); + expect(body.error.message).toBe('Selected model is unavailable. Choose another model.'); + expect(body.error.message).not.toContain('internal-provider'); + expect(body.error.message).not.toContain('internal-model'); + }); + + it('POST maps creation policy denials to 403 with sanitized response messages', async () => { + mockCreateUserDefinition.mockRejectedValueOnce( + new CustomAgentDefinitionServiceError( + 'creation_unavailable', + 'Custom agent creation is not available. Ask an admin for access.' + ) + ); + + const response = await POST( + makeRequest({ + name: 'Sample agent', + instructions: 'Answer briefly.', + capabilityIds: ['read_context'], + resourceBehavior: 'chat_only', + }) + ); + const body = await response.json(); + + expect(response.status).toBe(403); + expect(body.error.message).toBe('Custom agent creation is not available. Ask an admin for access.'); + expect(body.error.message).not.toContain('roles'); + expect(body.error.message).not.toContain('customAgentCreationPolicy'); + }); +}); diff --git a/src/app/api/v2/ai/agent/definitions/route.ts b/src/app/api/v2/ai/agent/definitions/route.ts new file mode 100644 index 00000000..56cae4ee --- /dev/null +++ b/src/app/api/v2/ai/agent/definitions/route.ts @@ -0,0 +1,299 @@ +/** + * 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 { getRequestUserIdentity } from 'server/lib/get-user'; +import { errorResponse, successResponse } from 'server/lib/response'; +import { + CustomAgentDefinitionServiceError, + customAgentDefinitionService, + serializeUserAgentDefinition, +} from 'server/services/agent/CustomAgentDefinitionService'; +import type { + AgentDefinitionModelPreference, + UserAgentDefinitionResourceBehavior, + UserAgentDefinitionUpsertInput, +} from 'server/services/agent/agentDefinitionTypes'; +import type { AgentCapabilityCatalogId } from 'server/services/agent/capabilityCatalog'; + +export const dynamic = 'force-dynamic'; + +const ALLOWED_REQUEST_FIELDS = new Set([ + 'name', + 'description', + 'instructions', + 'capabilityIds', + 'modelPreference', + 'resourceBehavior', +]); +const RESOURCE_BEHAVIORS = new Set([ + 'chat_only', + 'current_workspace_when_available', +]); + +function mapDefinitionError(error: unknown, req: NextRequest) { + if (error instanceof CustomAgentDefinitionServiceError) { + if (error.code === 'not_found') { + return errorResponse(error, { status: 404 }, req); + } + + if (error.code === 'model_unavailable') { + return errorResponse(error, { status: 409 }, req); + } + + if (error.code === 'creation_unavailable') { + return errorResponse(error, { status: 403 }, req); + } + + return errorResponse(error, { status: 400 }, req); + } + + throw error; +} + +async function readRequestBody(req: NextRequest): Promise | Error> { + let body: unknown; + try { + body = await req.json(); + } catch { + return new Error('Invalid JSON in request body'); + } + + if (!body || typeof body !== 'object' || Array.isArray(body)) { + return new Error('Request body must be an object.'); + } + + return body as Record; +} + +function readNullableString(value: unknown, fieldName: string): string | null | Error { + if (value === undefined || value === null) { + return null; + } + + if (typeof value !== 'string') { + return new Error(`${fieldName} must be a string.`); + } + + return value; +} + +function readRequiredString(value: unknown, fieldName: string): string | Error { + if (typeof value !== 'string') { + return new Error(`${fieldName} must be a string.`); + } + + return value; +} + +function readCapabilityIds(value: unknown): AgentCapabilityCatalogId[] | Error { + if (value === undefined) { + return []; + } + + if (!Array.isArray(value) || value.some((item) => typeof item !== 'string')) { + return new Error('capabilityIds must be an array of strings.'); + } + + return value as AgentCapabilityCatalogId[]; +} + +function readModelPreference(value: unknown): AgentDefinitionModelPreference | null | Error { + if (value === undefined || value === null) { + return null; + } + + if (typeof value !== 'object' || Array.isArray(value)) { + return new Error('modelPreference must be an object or null.'); + } + + const modelPreference = value as Record; + const unknownKeys = Object.keys(modelPreference).filter((key) => key !== 'provider' && key !== 'model'); + if (unknownKeys.length > 0) { + return new Error(`Unsupported modelPreference fields: ${unknownKeys.join(', ')}`); + } + + const provider = readNullableString(modelPreference.provider, 'modelPreference.provider'); + if (provider instanceof Error) { + return provider; + } + + const model = readNullableString(modelPreference.model, 'modelPreference.model'); + if (model instanceof Error) { + return model; + } + + return { provider, model }; +} + +function readResourceBehavior(value: unknown): UserAgentDefinitionResourceBehavior | Error { + if (typeof value !== 'string' || !RESOURCE_BEHAVIORS.has(value as UserAgentDefinitionResourceBehavior)) { + return new Error('resourceBehavior must be chat_only or current_workspace_when_available.'); + } + + return value as UserAgentDefinitionResourceBehavior; +} + +function parseUpsertBody(body: Record): UserAgentDefinitionUpsertInput | Error { + const unsupportedFields = Object.keys(body).filter((key) => !ALLOWED_REQUEST_FIELDS.has(key)); + if (unsupportedFields.length > 0) { + return new Error(`Unsupported agent definition fields: ${unsupportedFields.join(', ')}`); + } + + const name = readRequiredString(body.name, 'name'); + if (name instanceof Error) { + return name; + } + + const instructions = readRequiredString(body.instructions, 'instructions'); + if (instructions instanceof Error) { + return instructions; + } + + const description = readNullableString(body.description, 'description'); + if (description instanceof Error) { + return description; + } + + const capabilityRefs = readCapabilityIds(body.capabilityIds); + if (capabilityRefs instanceof Error) { + return capabilityRefs; + } + + const modelPreference = readModelPreference(body.modelPreference); + if (modelPreference instanceof Error) { + return modelPreference; + } + + const resourceBehavior = readResourceBehavior(body.resourceBehavior); + if (resourceBehavior instanceof Error) { + return resourceBehavior; + } + + return { + name, + description, + instructionAddendum: instructions, + capabilityRefs, + modelPreference, + resourceBehavior, + }; +} + +/** + * @openapi + * /api/v2/ai/agent/definitions: + * get: + * summary: List current-user custom agent definitions + * description: Lists the authenticated user's private custom agent definitions for the unified agent platform. + * tags: + * - Agent Platform + * operationId: listUserAgentDefinitions + * responses: + * '200': + * description: Current user's active custom agent definitions. + * content: + * application/json: + * schema: + * $ref: '#/components/schemas/ListUserAgentDefinitionsSuccessResponse' + * '401': + * description: Unauthorized. + * content: + * application/json: + * schema: + * $ref: '#/components/schemas/ApiErrorResponse' + * post: + * summary: Create a current-user custom agent definition + * description: Creates a private custom agent definition from user-selectable capabilities and resource behavior. + * tags: + * - Agent Platform + * operationId: createUserAgentDefinition + * requestBody: + * required: true + * content: + * application/json: + * schema: + * $ref: '#/components/schemas/UserAgentDefinitionUpsertRequest' + * responses: + * '201': + * description: Custom agent definition created. + * content: + * application/json: + * schema: + * $ref: '#/components/schemas/UserAgentDefinitionSuccessResponse' + * '400': + * description: Invalid request or capability selection. + * content: + * application/json: + * schema: + * $ref: '#/components/schemas/ApiErrorResponse' + * '401': + * description: Unauthorized. + * content: + * application/json: + * schema: + * $ref: '#/components/schemas/ApiErrorResponse' + * '403': + * description: Custom-agent creation is unavailable for the caller or selected capabilities. + * content: + * application/json: + * schema: + * $ref: '#/components/schemas/ApiErrorResponse' + * '409': + * description: Selected model is unavailable. + * content: + * application/json: + * schema: + * $ref: '#/components/schemas/ApiErrorResponse' + */ +const getHandler = async (req: NextRequest) => { + const userIdentity = getRequestUserIdentity(req); + if (!userIdentity) { + return errorResponse(new Error('Unauthorized'), { status: 401 }, req); + } + + const definitions = await customAgentDefinitionService.listUserDefinitions({ userId: userIdentity.userId }); + return successResponse({ definitions: definitions.map(serializeUserAgentDefinition) }, { status: 200 }, req); +}; + +const postHandler = async (req: NextRequest) => { + const userIdentity = getRequestUserIdentity(req); + if (!userIdentity) { + return errorResponse(new Error('Unauthorized'), { status: 401 }, req); + } + + const body = await readRequestBody(req); + if (body instanceof Error) { + return errorResponse(body, { status: 400 }, req); + } + + const input = parseUpsertBody(body); + if (input instanceof Error) { + return errorResponse(input, { status: 400 }, req); + } + + try { + const definition = await customAgentDefinitionService.createUserDefinition(userIdentity, input); + return successResponse({ definition: serializeUserAgentDefinition(definition) }, { status: 201 }, req); + } catch (error) { + return mapDefinitionError(error, req); + } +}; + +export const GET = createApiHandler(getHandler); +export const POST = createApiHandler(postHandler); diff --git a/src/app/api/v2/ai/agent/github-token/route.test.ts b/src/app/api/v2/ai/agent/github-token/route.test.ts index 2dc88f2f..1bc07820 100644 --- a/src/app/api/v2/ai/agent/github-token/route.test.ts +++ b/src/app/api/v2/ai/agent/github-token/route.test.ts @@ -43,8 +43,11 @@ function makeRequest(): NextRequest { } describe('GET /api/v2/ai/agent/github-token', () => { + const originalEnableAuth = process.env.ENABLE_AUTH; + beforeEach(() => { jest.clearAllMocks(); + process.env.ENABLE_AUTH = 'true'; mockGetUser.mockReturnValue({ sub: 'user-123', realm_access: { @@ -57,6 +60,14 @@ describe('GET /api/v2/ai/agent/github-token', () => { }); }); + afterEach(() => { + if (originalEnableAuth === undefined) { + delete process.env.ENABLE_AUTH; + } else { + process.env.ENABLE_AUTH = originalEnableAuth; + } + }); + it('returns 401 when no user is available', async () => { mockGetUser.mockReturnValue(null); mockGetRequestUserIdentity.mockReturnValue(null); diff --git a/src/app/api/v2/ai/agent/mcp-connections/[slug]/oauth/callback/route.test.ts b/src/app/api/v2/ai/agent/mcp-connections/[slug]/oauth/callback/route.test.ts index 68eafd96..df373bd0 100644 --- a/src/app/api/v2/ai/agent/mcp-connections/[slug]/oauth/callback/route.test.ts +++ b/src/app/api/v2/ai/agent/mcp-connections/[slug]/oauth/callback/route.test.ts @@ -27,12 +27,17 @@ jest.mock('@ai-sdk/mcp', () => ({ auth: (...args: unknown[]) => mockAuth(...args), })); -jest.mock('server/services/ai/mcp/config', () => ({ - McpConfigService: jest.fn().mockImplementation(() => ({ - getBySlugAndScope: (...args: unknown[]) => mockGetBySlugAndScope(...args), - discoverTools: (...args: unknown[]) => mockDiscoverTools(...args), - })), -})); +jest.mock('server/services/agentRuntime/mcp/config', () => { + const actual = jest.requireActual('server/services/agentRuntime/mcp/config'); + return { + __esModule: true, + ...actual, + McpConfigService: jest.fn().mockImplementation(() => ({ + getBySlugAndScope: (...args: unknown[]) => mockGetBySlugAndScope(...args), + discoverTools: (...args: unknown[]) => mockDiscoverTools(...args), + })), + }; +}); jest.mock('server/services/userMcpConnection', () => ({ __esModule: true, @@ -42,7 +47,7 @@ jest.mock('server/services/userMcpConnection', () => ({ }, })); -jest.mock('server/services/ai/mcp/oauthFlow', () => ({ +jest.mock('server/services/agentRuntime/mcp/oauthFlow', () => ({ __esModule: true, extractMcpOAuthFlowId: (state: string | null | undefined) => { if (!state) { @@ -62,7 +67,7 @@ jest.mock('server/lib/logger', () => ({ })); import { GET } from './route'; -import { buildMcpDefinitionFingerprint } from 'server/services/ai/mcp/connectionConfig'; +import { buildMcpDefinitionFingerprint } from 'server/services/agentRuntime/mcp/connectionConfig'; function makeRequest( url = 'http://localhost/api/v2/ai/agent/mcp-connections/sample-oauth/oauth/callback?code=sample-code&state=flow-123.sample-state' @@ -176,6 +181,35 @@ describe('GET /api/v2/ai/agent/mcp-connections/[slug]/oauth/callback', () => { expect(html).toContain('https://app.example.com'); }); + it('redacts OAuth and token secrets from failed discovery output', async () => { + mockDiscoverTools.mockRejectedValueOnce( + new Error( + 'Discovery failed for code sample-code token sample-access-token state flow-123.sample-state verifier sample-code-verifier' + ) + ); + + const response = await GET(makeRequest(), { + params: Promise.resolve({ slug: 'sample-oauth' }), + }); + const html = await response.text(); + const persisted = mockUpsertConnection.mock.calls[mockUpsertConnection.mock.calls.length - 1]?.[0]; + + expect(response.status).toBe(422); + expect(persisted).toEqual( + expect.objectContaining({ + slug: 'sample-oauth', + scope: 'global', + discoveredTools: [], + validationError: 'Discovery failed for code ****** token ****** state ****** verifier ******', + }) + ); + expect(html).toContain('Discovery failed for code ****** token ****** state ****** verifier ******'); + expect(html).not.toContain('sample-code'); + expect(html).not.toContain('sample-access-token'); + expect(html).not.toContain('flow-123.sample-state'); + expect(html).not.toContain('sample-code-verifier'); + }); + it('rejects expired or reused flows before completing OAuth', async () => { mockConsumeFlow.mockResolvedValueOnce(null); diff --git a/src/app/api/v2/ai/agent/mcp-connections/[slug]/oauth/callback/route.ts b/src/app/api/v2/ai/agent/mcp-connections/[slug]/oauth/callback/route.ts index 1bc2fbf5..4d73a66d 100644 --- a/src/app/api/v2/ai/agent/mcp-connections/[slug]/oauth/callback/route.ts +++ b/src/app/api/v2/ai/agent/mcp-connections/[slug]/oauth/callback/route.ts @@ -23,11 +23,14 @@ import { buildMcpDefinitionFingerprint, mergeCompiledConnectionConfig, normalizeAuthConfig, -} from 'server/services/ai/mcp/connectionConfig'; -import { McpConfigService } from 'server/services/ai/mcp/config'; -import McpOAuthFlowService, { extractMcpOAuthFlowId, type McpOAuthFlowRecord } from 'server/services/ai/mcp/oauthFlow'; -import { PersistentOAuthClientProvider } from 'server/services/ai/mcp/oauthProvider'; -import type { McpDiscoveredTool, McpStoredUserConnectionState } from 'server/services/ai/mcp/types'; +} from 'server/services/agentRuntime/mcp/connectionConfig'; +import { McpConfigService, sanitizeMcpErrorMessage } from 'server/services/agentRuntime/mcp/config'; +import McpOAuthFlowService, { + extractMcpOAuthFlowId, + type McpOAuthFlowRecord, +} from 'server/services/agentRuntime/mcp/oauthFlow'; +import { PersistentOAuthClientProvider } from 'server/services/agentRuntime/mcp/oauthProvider'; +import type { McpDiscoveredTool, McpStoredUserConnectionState } from 'server/services/agentRuntime/mcp/types'; import UserMcpConnectionService from 'server/services/userMcpConnection'; type OAuthConnectionState = Extract; @@ -418,11 +421,10 @@ const getHandler = async (req: NextRequest, { params }: { params: Promise<{ slug validatedAt: existing?.validatedAt, interactive: false, }); - const transport = applyCompiledConnectionConfigToTransport( - config.transport, - mergeCompiledConnectionConfig(config.sharedConfig || {}, undefined), - { authProvider: provider } - ); + const compiledConfig = mergeCompiledConnectionConfig(config.sharedConfig || {}, undefined); + const transport = applyCompiledConnectionConfigToTransport(config.transport, compiledConfig, { + authProvider: provider, + }); if (transport.type === 'stdio') { const invalidTransportMessage = `OAuth MCP connection '${flow.slug}' must use HTTP or SSE transport.`; await clearPendingFlowState(flow, invalidTransportMessage); @@ -482,7 +484,19 @@ const getHandler = async (req: NextRequest, { params }: { params: Promise<{ slug } ); } catch (error) { - const message = error instanceof Error ? error.message : String(error); + const message = sanitizeMcpErrorMessage(error, [ + { + values: { + code, + callbackState, + codeVerifier: provider.currentState.codeVerifier, + oauthState: provider.currentState.oauthState, + }, + compiledConfig, + transport, + extraSecrets: [provider.currentState.tokens, provider.currentState.clientInformation], + }, + ]); await persistOAuthConnectionState({ flow, state: provider.currentState, diff --git a/src/app/api/v2/ai/agent/mcp-connections/[slug]/oauth/start/route.test.ts b/src/app/api/v2/ai/agent/mcp-connections/[slug]/oauth/start/route.test.ts index f4cc4731..3c16d69c 100644 --- a/src/app/api/v2/ai/agent/mcp-connections/[slug]/oauth/start/route.test.ts +++ b/src/app/api/v2/ai/agent/mcp-connections/[slug]/oauth/start/route.test.ts @@ -27,11 +27,16 @@ jest.mock('@ai-sdk/mcp', () => ({ auth: (...args: unknown[]) => mockAuth(...args), })); -jest.mock('server/services/ai/mcp/config', () => ({ - McpConfigService: jest.fn().mockImplementation(() => ({ - getBySlugAndScope: (...args: unknown[]) => mockGetBySlugAndScope(...args), - })), -})); +jest.mock('server/services/agentRuntime/mcp/config', () => { + const actual = jest.requireActual('server/services/agentRuntime/mcp/config'); + return { + __esModule: true, + ...actual, + McpConfigService: jest.fn().mockImplementation(() => ({ + getBySlugAndScope: (...args: unknown[]) => mockGetBySlugAndScope(...args), + })), + }; +}); jest.mock('server/services/userMcpConnection', () => ({ __esModule: true, @@ -41,7 +46,7 @@ jest.mock('server/services/userMcpConnection', () => ({ }, })); -jest.mock('server/services/ai/mcp/oauthFlow', () => ({ +jest.mock('server/services/agentRuntime/mcp/oauthFlow', () => ({ __esModule: true, default: { create: (...args: unknown[]) => mockCreateFlow(...args), @@ -166,6 +171,40 @@ describe('POST /api/v2/ai/agent/mcp-connections/[slug]/oauth/start', () => { expect(mockInvalidateFlow).toHaveBeenCalledWith('flow-123'); }); + it('redacts MCP secrets when OAuth authorization setup fails', async () => { + mockGetBySlugAndScope.mockResolvedValueOnce({ + id: 7, + slug: 'sample-oauth', + scope: 'global', + enabled: true, + timeout: 30000, + preset: 'oauth-http', + transport: { type: 'http', url: 'https://mcp.example.com/v1/mcp?api_key=query/secret+value', headers: {} }, + sharedConfig: { + headers: { Authorization: 'Bearer shared-header-secret' }, + }, + authConfig: { + mode: 'oauth', + provider: 'generic-oauth2.1', + scope: 'sample.read', + }, + } as const); + mockAuth.mockRejectedValueOnce( + new Error( + 'OAuth start failed Authorization=Bearer shared-header-secret query=query/secret+value encoded=query%2Fsecret%2Bvalue' + ) + ); + + const response = await POST(makeRequest(), { + params: Promise.resolve({ slug: 'sample-oauth' }), + }); + const body = await response.json(); + + expect(response.status).toBe(422); + expect(mockInvalidateFlow).toHaveBeenCalledWith('flow-123'); + expect(body.error.message).toBe('OAuth start failed Authorization=****** query=****** encoded=******'); + }); + it('drops stale saved client metadata when the persisted redirect URI no longer matches', async () => { mockGetDecryptedConnection.mockResolvedValueOnce({ state: { diff --git a/src/app/api/v2/ai/agent/mcp-connections/[slug]/oauth/start/route.ts b/src/app/api/v2/ai/agent/mcp-connections/[slug]/oauth/start/route.ts index 665e7d83..c9099e7b 100644 --- a/src/app/api/v2/ai/agent/mcp-connections/[slug]/oauth/start/route.ts +++ b/src/app/api/v2/ai/agent/mcp-connections/[slug]/oauth/start/route.ts @@ -25,11 +25,11 @@ import { buildMcpDefinitionFingerprint, mergeCompiledConnectionConfig, normalizeAuthConfig, -} from 'server/services/ai/mcp/connectionConfig'; -import { McpConfigService } from 'server/services/ai/mcp/config'; -import McpOAuthFlowService from 'server/services/ai/mcp/oauthFlow'; -import { PersistentOAuthClientProvider } from 'server/services/ai/mcp/oauthProvider'; -import type { McpStoredUserConnectionState } from 'server/services/ai/mcp/types'; +} from 'server/services/agentRuntime/mcp/connectionConfig'; +import { McpConfigService, sanitizeMcpErrorMessage } from 'server/services/agentRuntime/mcp/config'; +import McpOAuthFlowService from 'server/services/agentRuntime/mcp/oauthFlow'; +import { PersistentOAuthClientProvider } from 'server/services/agentRuntime/mcp/oauthProvider'; +import type { McpStoredUserConnectionState } from 'server/services/agentRuntime/mcp/types'; import UserMcpConnectionService from 'server/services/userMcpConnection'; type OAuthConnectionState = Extract; @@ -203,11 +203,10 @@ const postHandler = async (req: NextRequest, { params }: { params: Promise<{ slu validatedAt: existing?.validatedAt, interactive: true, }); - const transport = applyCompiledConnectionConfigToTransport( - config.transport, - mergeCompiledConnectionConfig(config.sharedConfig || {}, undefined), - { authProvider: provider } - ); + const compiledConfig = mergeCompiledConnectionConfig(config.sharedConfig || {}, undefined); + const transport = applyCompiledConnectionConfigToTransport(config.transport, compiledConfig, { + authProvider: provider, + }); if (transport.type === 'stdio') { return NextResponse.json( { @@ -239,7 +238,18 @@ const postHandler = async (req: NextRequest, { params }: { params: Promise<{ slu ); } catch (error) { await McpOAuthFlowService.invalidate(flow.flowId); - throw error; + const message = sanitizeMcpErrorMessage(error, [ + { + values: { + oauthState: provider.currentState.oauthState, + codeVerifier: provider.currentState.codeVerifier, + }, + compiledConfig, + transport, + extraSecrets: [provider.currentState.tokens, provider.currentState.clientInformation], + }, + ]); + return errorResponse(new Error(message), { status: 422 }, req); } }; diff --git a/src/app/api/v2/ai/agent/mcp-connections/[slug]/route.test.ts b/src/app/api/v2/ai/agent/mcp-connections/[slug]/route.test.ts index fcc7113b..2af2362f 100644 --- a/src/app/api/v2/ai/agent/mcp-connections/[slug]/route.test.ts +++ b/src/app/api/v2/ai/agent/mcp-connections/[slug]/route.test.ts @@ -23,12 +23,17 @@ const mockDeleteConnection = jest.fn(); const mockGetMaskedState = jest.fn(); const mockGetRequestUserIdentity = jest.fn(); -jest.mock('server/services/ai/mcp/config', () => ({ - McpConfigService: jest.fn().mockImplementation(() => ({ - getBySlugAndScope: (...args: unknown[]) => mockGetBySlugAndScope(...args), - discoverTools: (...args: unknown[]) => mockDiscoverTools(...args), - })), -})); +jest.mock('server/services/agentRuntime/mcp/config', () => { + const actual = jest.requireActual('server/services/agentRuntime/mcp/config'); + return { + __esModule: true, + ...actual, + McpConfigService: jest.fn().mockImplementation(() => ({ + getBySlugAndScope: (...args: unknown[]) => mockGetBySlugAndScope(...args), + discoverTools: (...args: unknown[]) => mockDiscoverTools(...args), + })), + }; +}); jest.mock('server/services/userMcpConnection', () => ({ __esModule: true, @@ -48,8 +53,8 @@ jest.mock('server/lib/logger', () => ({ })); import { DELETE, PUT } from './route'; -import { buildMcpDefinitionFingerprint } from 'server/services/ai/mcp/connectionConfig'; -import type { McpAuthConfig, McpTransportConfig } from 'server/services/ai/mcp/types'; +import { buildMcpDefinitionFingerprint } from 'server/services/agentRuntime/mcp/connectionConfig'; +import type { McpAuthConfig, McpTransportConfig } from 'server/services/agentRuntime/mcp/types'; function makeRequest( body: unknown, @@ -155,22 +160,24 @@ describe('MCP user connection route', () => { }); it('stores validation errors on the user connection without mutating shared discovery', async () => { - mockDiscoverTools.mockRejectedValue(new Error('HTTP 401 Unauthorized')); + mockDiscoverTools.mockRejectedValue( + new Error('HTTP 401 Unauthorized for Authorization Bearer invalid-token and token invalid-token') + ); const response = await PUT(makeRequest({ values: { apiToken: 'invalid-token' } }), { params: Promise.resolve({ slug: 'sample-connector' }), }); + const body = await response.json(); expect(response.status).toBe(422); expect(mockUpsertConnection).toHaveBeenCalledWith( expect.objectContaining({ discoveredTools: [], - validationError: 'HTTP 401 Unauthorized', + validationError: 'HTTP 401 Unauthorized for Authorization ****** and token ******', }) ); - await expect(response.json()).resolves.toMatchObject({ - error: { message: 'HTTP 401 Unauthorized' }, - }); + expect(body.error.message).toBe('HTTP 401 Unauthorized for Authorization ****** and token ******'); + expect(body.error.message).not.toContain('invalid-token'); }); it('deletes the current user connection and returns the empty state', async () => { diff --git a/src/app/api/v2/ai/agent/mcp-connections/[slug]/route.ts b/src/app/api/v2/ai/agent/mcp-connections/[slug]/route.ts index 88fb320d..a00147df 100644 --- a/src/app/api/v2/ai/agent/mcp-connections/[slug]/route.ts +++ b/src/app/api/v2/ai/agent/mcp-connections/[slug]/route.ts @@ -26,8 +26,8 @@ import { mergeCompiledConnectionConfig, normalizeAuthConfig, normalizeUserConnectionValues, -} from 'server/services/ai/mcp/connectionConfig'; -import { McpConfigService } from 'server/services/ai/mcp/config'; +} from 'server/services/agentRuntime/mcp/connectionConfig'; +import { McpConfigService, sanitizeMcpErrorMessage } from 'server/services/agentRuntime/mcp/config'; import UserMcpConnectionService from 'server/services/userMcpConnection'; /** @@ -211,7 +211,16 @@ const putHandler = async (req: NextRequest, { params }: { params: Promise<{ slug validatedAt, }); } catch (error) { - const message = error instanceof Error ? error.message : String(error); + const message = sanitizeMcpErrorMessage(error, [ + { + values, + compiledConfig: compiledUserConfig, + }, + { + compiledConfig, + transport, + }, + ]); await UserMcpConnectionService.upsertConnection({ userId: userIdentity.userId, ownerGithubUsername: userIdentity.githubUsername, diff --git a/src/app/api/v2/ai/agent/mcp-connections/route.ts b/src/app/api/v2/ai/agent/mcp-connections/route.ts index d2007c47..06511d11 100644 --- a/src/app/api/v2/ai/agent/mcp-connections/route.ts +++ b/src/app/api/v2/ai/agent/mcp-connections/route.ts @@ -18,7 +18,7 @@ import { NextRequest } from 'next/server'; import { createApiHandler } from 'server/lib/createApiHandler'; import { errorResponse, successResponse } from 'server/lib/response'; import { getRequestUserIdentity } from 'server/lib/get-user'; -import { McpConfigService } from 'server/services/ai/mcp/config'; +import { McpConfigService } from 'server/services/agentRuntime/mcp/config'; export const dynamic = 'force-dynamic'; diff --git a/src/app/api/v2/ai/agent/runs/[runId]/route.ts b/src/app/api/v2/ai/agent/runs/[runId]/route.ts index 7b1d965e..23a685cb 100644 --- a/src/app/api/v2/ai/agent/runs/[runId]/route.ts +++ b/src/app/api/v2/ai/agent/runs/[runId]/route.ts @@ -26,8 +26,9 @@ import AgentRunService from 'server/services/agent/RunService'; * /api/v2/ai/agent/runs/{runId}: * get: * summary: Get an agent run + * description: Returns run status and the public runPlan summary for an owned run. * tags: - * - Agent Sessions + * - Agent Platform * operationId: getAgentRun * parameters: * - in: path diff --git a/src/app/api/v2/ai/agent/runtime-config/repos/[...fullName]/route.test.ts b/src/app/api/v2/ai/agent/runtime-config/repos/[...fullName]/route.test.ts new file mode 100644 index 00000000..8fed17ae --- /dev/null +++ b/src/app/api/v2/ai/agent/runtime-config/repos/[...fullName]/route.test.ts @@ -0,0 +1,105 @@ +/** + * 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 mockGetUser = jest.fn(); +const mockSetRepoConfig = jest.fn(); +const mockUpdateRepoAdditiveRules = jest.fn(); +const mockDeleteRepoConfig = jest.fn(); + +jest.mock('server/lib/get-user', () => ({ + getUser: (...args: unknown[]) => mockGetUser(...args), +})); + +jest.mock('server/services/agentRuntime/config/agentRuntimeConfig', () => ({ + __esModule: true, + default: { + getInstance: jest.fn(() => ({ + setRepoConfig: (...args: unknown[]) => mockSetRepoConfig(...args), + updateRepoAdditiveRules: (...args: unknown[]) => mockUpdateRepoAdditiveRules(...args), + deleteRepoConfig: (...args: unknown[]) => mockDeleteRepoConfig(...args), + })), + }, +})); + +import { DELETE, PATCH, PUT } from './route'; + +function makeRequest(body?: unknown): NextRequest { + return { + headers: new Headers([['x-request-id', 'req-test']]), + nextUrl: new URL('http://localhost/api/v2/ai/agent/runtime-config/repos/example-org/example-repo'), + json: jest.fn().mockResolvedValue(body ?? {}), + } as unknown as NextRequest; +} + +const params = { params: { fullName: ['example-org', 'example-repo'] } }; + +describe('/api/v2/ai/agent/runtime-config/repos/[...fullName]', () => { + const originalEnableAuth = process.env.ENABLE_AUTH; + + beforeEach(() => { + jest.clearAllMocks(); + process.env.ENABLE_AUTH = 'true'; + mockGetUser.mockReturnValue({ + sub: 'sample-user', + realm_access: { + roles: ['user'], + }, + }); + }); + + afterEach(() => { + if (originalEnableAuth === undefined) { + delete process.env.ENABLE_AUTH; + } else { + process.env.ENABLE_AUTH = originalEnableAuth; + } + }); + + it('rejects non-admin repo config replacement before parsing the body', async () => { + const request = makeRequest({}); + + const response = await PUT(request, params); + const body = await response.json(); + + expect(response.status).toBe(403); + expect(body.error.message).toBe('Forbidden: insufficient permissions'); + expect(request.json).not.toHaveBeenCalled(); + expect(mockSetRepoConfig).not.toHaveBeenCalled(); + }); + + it('rejects non-admin repo additive-rule patches before parsing the body', async () => { + const request = makeRequest({ additiveRules: [] }); + + const response = await PATCH(request, params); + const body = await response.json(); + + expect(response.status).toBe(403); + expect(body.error.message).toBe('Forbidden: insufficient permissions'); + expect(request.json).not.toHaveBeenCalled(); + expect(mockUpdateRepoAdditiveRules).not.toHaveBeenCalled(); + }); + + it('rejects non-admin repo config deletion before mutating config', async () => { + const response = await DELETE(makeRequest(), params); + const body = await response.json(); + + expect(response.status).toBe(403); + expect(body.error.message).toBe('Forbidden: insufficient permissions'); + expect(mockDeleteRepoConfig).not.toHaveBeenCalled(); + }); +}); diff --git a/src/app/api/v2/ai/agent-config/repos/[...fullName]/route.ts b/src/app/api/v2/ai/agent/runtime-config/repos/[...fullName]/route.ts similarity index 72% rename from src/app/api/v2/ai/agent-config/repos/[...fullName]/route.ts rename to src/app/api/v2/ai/agent/runtime-config/repos/[...fullName]/route.ts index f7e4fa30..6dd0bd93 100644 --- a/src/app/api/v2/ai/agent-config/repos/[...fullName]/route.ts +++ b/src/app/api/v2/ai/agent/runtime-config/repos/[...fullName]/route.ts @@ -18,14 +18,14 @@ import { NextRequest } from 'next/server'; import { createApiHandler } from 'server/lib/createApiHandler'; import { errorResponse, successResponse } from 'server/lib/response'; import { getLogger } from 'server/lib/logger'; -import AIAgentConfigService from 'server/services/aiAgentConfig'; +import AgentRuntimeConfigService from 'server/services/agentRuntime/config/agentRuntimeConfig'; import JsonSchema from 'jsonschema'; import { - aiAgentAdditiveRulesUpdateSchema, - aiAgentRepoOverrideSchema, -} from 'server/lib/validation/aiAgentConfigSchemas'; + agentRuntimeAdditiveRulesUpdateSchema, + agentRuntimeRepoOverrideSchema, +} from 'server/lib/validation/agentRuntimeConfigSchemas'; import { normalizeRepoFullName } from 'server/lib/normalizeRepoFullName'; -import { AIAgentConfigValidationError } from 'server/lib/validation/aiAgentConfigValidator'; +import { AgentRuntimeConfigValidationError } from 'server/lib/validation/agentRuntimeConfigValidator'; export const dynamic = 'force-dynamic'; @@ -57,16 +57,16 @@ function parseFullNameParams(segments: string[]): { fullName: string; isEffectiv /** * @openapi - * /api/v2/ai/agent-config/repos/{owner}/{repo}: + * /api/v2/ai/agent/runtime-config/repos/{owner}/{repo}: * get: - * summary: Get repository AI agent config override + * summary: Get repository agent runtime config override * description: > - * Returns the raw per-repository AI agent configuration override (not merged + * Returns the raw per-repository Agent runtime configuration override (not merged * with the global config). This shows only the fields that are explicitly * overridden for this repository. Returns 404 if no override exists. * tags: - * - AI Agent Config - * operationId: getRepoAIAgentConfig + * - Agent Runtime Config + * operationId: getRepoAgentRuntimeConfig * parameters: * - in: path * name: owner @@ -84,11 +84,11 @@ function parseFullNameParams(segments: string[]): { fullName: string; isEffectiv * example: example-repo * responses: * '200': - * description: Repository AI agent config override + * description: Repository agent runtime config override * content: * application/json: * schema: - * $ref: '#/components/schemas/GetRepoAIAgentConfigSuccessResponse' + * $ref: '#/components/schemas/GetRepoAgentRuntimeConfigSuccessResponse' * '404': * description: No config override found for the repository * content: @@ -102,9 +102,9 @@ function parseFullNameParams(segments: string[]): { fullName: string; isEffectiv * schema: * $ref: '#/components/schemas/ApiErrorResponse' * - * /api/v2/ai/agent-config/repos/{owner}/{repo}/effective: + * /api/v2/ai/agent/runtime-config/repos/{owner}/{repo}/effective: * get: - * summary: Get effective (merged) AI agent config for repository + * summary: Get effective (merged) agent runtime config for repository * description: > * Returns the fully merged configuration for a repository. The merge applies * the repository override on top of the global config using these rules: @@ -115,8 +115,8 @@ function parseFullNameParams(segments: string[]): { fullName: string; isEffectiv * always taken from the global config (not overridable per-repo). * If no repo override exists, this returns the global config unchanged. * tags: - * - AI Agent Config - * operationId: getEffectiveAIAgentConfig + * - Agent Runtime Config + * operationId: getEffectiveAgentRuntimeConfig * parameters: * - in: path * name: owner @@ -134,11 +134,11 @@ function parseFullNameParams(segments: string[]): { fullName: string; isEffectiv * example: example-repo * responses: * '200': - * description: Effective merged AI agent configuration + * description: Effective merged Agent runtime configuration * content: * application/json: * schema: - * $ref: '#/components/schemas/GetEffectiveAIAgentConfigSuccessResponse' + * $ref: '#/components/schemas/GetEffectiveAgentRuntimeConfigSuccessResponse' * '500': * description: Server error * content: @@ -154,11 +154,11 @@ const getHandler = async (req: NextRequest, { params }: { params: { fullName: st return errorResponse(new Error('Invalid repository fullName. Expected format: owner/repo'), { status: 400 }, req); } - const service = AIAgentConfigService.getInstance(); + const service = AgentRuntimeConfigService.getInstance(); if (parsed.isEffective) { const result = await service.getEffectiveConfig(parsed.fullName); - getLogger().info('AIAgentConfig: effective config read repo=' + parsed.fullName + ' via=api'); + getLogger().info('AgentRuntimeConfig: effective config read repo=' + parsed.fullName + ' via=api'); return successResponse({ repoFullName: parsed.fullName, effectiveConfig: result }, { status: 200 }, req); } @@ -171,25 +171,25 @@ const getHandler = async (req: NextRequest, { params }: { params: { fullName: st ); } - getLogger().info('AIAgentConfig: repo config read repo=' + parsed.fullName + ' via=api'); + getLogger().info('AgentRuntimeConfig: repo config read repo=' + parsed.fullName + ' via=api'); return successResponse({ repoFullName: parsed.fullName, config: result }, { status: 200 }, req); }; /** * @openapi - * /api/v2/ai/agent-config/repos/{owner}/{repo}: + * /api/v2/ai/agent/runtime-config/repos/{owner}/{repo}: * put: - * summary: Create or update repository AI agent config override + * summary: Create or update repository agent runtime config override * description: > - * Upserts a per-repository AI agent configuration override. Only include + * Upserts a per-repository Agent runtime configuration override. Only include * the fields you want to override — omitted fields will continue to inherit * from the global config. Validation enforces the same limits as the global * config (prompt length, pattern count, core tool restrictions). The request - * body is validated against the AIAgentRepoOverride schema. + * body is validated against the AgentRuntimeRepoOverride schema. * Changes take effect immediately and invalidate cached effective configs. * tags: - * - AI Agent Config - * operationId: putRepoAIAgentConfig + * - Agent Runtime Config + * operationId: putRepoAgentRuntimeConfig * parameters: * - in: path * name: owner @@ -210,14 +210,14 @@ const getHandler = async (req: NextRequest, { params }: { params: { fullName: st * content: * application/json: * schema: - * $ref: '#/components/schemas/AIAgentRepoOverride' + * $ref: '#/components/schemas/AgentRuntimeRepoOverride' * responses: * '200': - * description: Updated repository AI agent config override + * description: Updated repository agent runtime config override * content: * application/json: * schema: - * $ref: '#/components/schemas/GetRepoAIAgentConfigSuccessResponse' + * $ref: '#/components/schemas/GetRepoAgentRuntimeConfigSuccessResponse' * '400': * description: > * Validation error. Possible reasons: JSON schema violation, @@ -227,6 +227,18 @@ const getHandler = async (req: NextRequest, { params }: { params: { fullName: st * application/json: * schema: * $ref: '#/components/schemas/ApiErrorResponse' + * '401': + * description: Unauthorized + * content: + * application/json: + * schema: + * $ref: '#/components/schemas/ApiErrorResponse' + * '403': + * description: Forbidden + * content: + * application/json: + * schema: + * $ref: '#/components/schemas/ApiErrorResponse' * '500': * description: Server error * content: @@ -239,8 +251,8 @@ const getHandler = async (req: NextRequest, { params }: { params: { fullName: st * Updates only the additiveRules field for a repository override while preserving * any other repository-specific settings that already exist. * tags: - * - AI Agent Config - * operationId: patchRepoAIAgentAdditiveRules + * - Agent Runtime Config + * operationId: patchRepoAgentRuntimeAdditiveRules * parameters: * - in: path * name: owner @@ -257,20 +269,32 @@ const getHandler = async (req: NextRequest, { params }: { params: { fullName: st * content: * application/json: * schema: - * $ref: '#/components/schemas/AIAgentAdditiveRulesUpdateRequest' + * $ref: '#/components/schemas/AgentRuntimeAdditiveRulesUpdateRequest' * responses: * '200': - * description: Updated repository AI agent config override + * description: Updated repository agent runtime config override * content: * application/json: * schema: - * $ref: '#/components/schemas/GetRepoAIAgentConfigSuccessResponse' + * $ref: '#/components/schemas/GetRepoAgentRuntimeConfigSuccessResponse' * '400': * description: Validation error * content: * application/json: * schema: * $ref: '#/components/schemas/ApiErrorResponse' + * '401': + * description: Unauthorized + * content: + * application/json: + * schema: + * $ref: '#/components/schemas/ApiErrorResponse' + * '403': + * description: Forbidden + * content: + * application/json: + * schema: + * $ref: '#/components/schemas/ApiErrorResponse' * '500': * description: Server error * content: @@ -302,26 +326,26 @@ const putHandler = async (req: NextRequest, { params }: { params: { fullName: st } const validator = new JsonSchema.Validator(); - const result = validator.validate(body, aiAgentRepoOverrideSchema); + const result = validator.validate(body, agentRuntimeRepoOverrideSchema); if (!result.valid) { const messages = result.errors.map((e) => e.stack).join('; '); return errorResponse(new Error('Validation failed: ' + messages), { status: 400 }, req); } - const service = AIAgentConfigService.getInstance(); + const service = AgentRuntimeConfigService.getInstance(); try { await service.setRepoConfig(parsed.fullName, body as any); } catch (error) { - if (error instanceof AIAgentConfigValidationError) { + if (error instanceof AgentRuntimeConfigValidationError) { return errorResponse(error, { status: 400 }, req); } throw error; } const updated = await service.getRepoConfig(parsed.fullName); - getLogger().info('AIAgentConfig: repo config updated repo=' + parsed.fullName + ' via=api'); + getLogger().info('AgentRuntimeConfig: repo config updated repo=' + parsed.fullName + ' via=api'); return successResponse({ repoFullName: parsed.fullName, config: updated }, { status: 200 }, req); }; @@ -349,24 +373,24 @@ const patchHandler = async (req: NextRequest, { params }: { params: { fullName: } const validator = new JsonSchema.Validator(); - const result = validator.validate(body, aiAgentAdditiveRulesUpdateSchema); + const result = validator.validate(body, agentRuntimeAdditiveRulesUpdateSchema); if (!result.valid) { const messages = result.errors.map((e) => e.stack).join('; '); return errorResponse(new Error('Validation failed: ' + messages), { status: 400 }, req); } - const service = AIAgentConfigService.getInstance(); + const service = AgentRuntimeConfigService.getInstance(); try { const updated = await service.updateRepoAdditiveRules( parsed.fullName, (body as { additiveRules: string[] }).additiveRules ); - getLogger().info('AIAgentConfig: repo additive rules updated repo=' + parsed.fullName + ' via=api'); + getLogger().info('AgentRuntimeConfig: repo additive rules updated repo=' + parsed.fullName + ' via=api'); return successResponse({ repoFullName: parsed.fullName, config: updated }, { status: 200 }, req); } catch (error) { - if (error instanceof AIAgentConfigValidationError) { + if (error instanceof AgentRuntimeConfigValidationError) { return errorResponse(error, { status: 400 }, req); } throw error; @@ -375,16 +399,16 @@ const patchHandler = async (req: NextRequest, { params }: { params: { fullName: /** * @openapi - * /api/v2/ai/agent-config/repos/{owner}/{repo}: + * /api/v2/ai/agent/runtime-config/repos/{owner}/{repo}: * delete: - * summary: Delete repository AI agent config override + * summary: Delete repository agent runtime config override * description: > - * Soft-deletes the per-repository AI agent configuration override. + * Soft-deletes the per-repository Agent runtime configuration override. * After deletion, the repository reverts to using the global config only. * The override can be recreated later with a PUT request. * tags: - * - AI Agent Config - * operationId: deleteRepoAIAgentConfig + * - Agent Runtime Config + * operationId: deleteRepoAgentRuntimeConfig * parameters: * - in: path * name: owner @@ -413,6 +437,18 @@ const patchHandler = async (req: NextRequest, { params }: { params: { fullName: * application/json: * schema: * $ref: '#/components/schemas/ApiErrorResponse' + * '401': + * description: Unauthorized + * content: + * application/json: + * schema: + * $ref: '#/components/schemas/ApiErrorResponse' + * '403': + * description: Forbidden + * content: + * application/json: + * schema: + * $ref: '#/components/schemas/ApiErrorResponse' * '500': * description: Server error * content: @@ -432,13 +468,13 @@ const deleteHandler = async (req: NextRequest, { params }: { params: { fullName: return errorResponse(new Error('Cannot DELETE the effective config endpoint.'), { status: 400 }, req); } - const service = AIAgentConfigService.getInstance(); + const service = AgentRuntimeConfigService.getInstance(); await service.deleteRepoConfig(parsed.fullName); - getLogger().info('AIAgentConfig: repo config deleted repo=' + parsed.fullName + ' via=api'); + getLogger().info('AgentRuntimeConfig: repo config deleted repo=' + parsed.fullName + ' via=api'); return successResponse({ repoFullName: parsed.fullName, deleted: true }, { status: 200 }, req); }; export const GET = createApiHandler(getHandler); -export const PUT = createApiHandler(putHandler); -export const PATCH = createApiHandler(patchHandler); -export const DELETE = createApiHandler(deleteHandler); +export const PUT = createApiHandler(putHandler, { roles: ['admin'] }); +export const PATCH = createApiHandler(patchHandler, { roles: ['admin'] }); +export const DELETE = createApiHandler(deleteHandler, { roles: ['admin'] }); diff --git a/src/app/api/v2/ai/agent-config/repos/route.ts b/src/app/api/v2/ai/agent/runtime-config/repos/route.ts similarity index 73% rename from src/app/api/v2/ai/agent-config/repos/route.ts rename to src/app/api/v2/ai/agent/runtime-config/repos/route.ts index 7e17d93a..f92e5851 100644 --- a/src/app/api/v2/ai/agent-config/repos/route.ts +++ b/src/app/api/v2/ai/agent/runtime-config/repos/route.ts @@ -18,28 +18,28 @@ import { NextRequest } from 'next/server'; import { createApiHandler } from 'server/lib/createApiHandler'; import { successResponse } from 'server/lib/response'; import { getLogger } from 'server/lib/logger'; -import AIAgentConfigService from 'server/services/aiAgentConfig'; +import AgentRuntimeConfigService from 'server/services/agentRuntime/config/agentRuntimeConfig'; /** * @openapi - * /api/v2/ai/agent-config/repos: + * /api/v2/ai/agent/runtime-config/repos: * get: - * summary: List all repository AI agent config overrides + * summary: List all repository agent runtime config overrides * description: > - * Returns an array of all repository-level AI agent configuration overrides. + * Returns an array of all repository-level Agent runtime configuration overrides. * These are partial configs that override specific fields of the global config. * To see the fully merged config for a specific repo, use the /effective endpoint instead. * Soft-deleted overrides are excluded from the response. * tags: - * - AI Agent Config - * operationId: listRepoAIAgentConfigs + * - Agent Runtime Config + * operationId: listRepoAgentRuntimeConfigs * responses: * '200': * description: Array of repository config overrides * content: * application/json: * schema: - * $ref: '#/components/schemas/ListRepoAIAgentConfigsSuccessResponse' + * $ref: '#/components/schemas/ListRepoAgentRuntimeConfigsSuccessResponse' * '500': * description: Server error * content: @@ -48,9 +48,9 @@ import AIAgentConfigService from 'server/services/aiAgentConfig'; * $ref: '#/components/schemas/ApiErrorResponse' */ const getHandler = async (req: NextRequest) => { - const service = AIAgentConfigService.getInstance(); + const service = AgentRuntimeConfigService.getInstance(); const configs = await service.listRepoConfigs(); - getLogger().info('AIAgentConfig: repo configs listed via=api count=' + configs.length); + getLogger().info('AgentRuntimeConfig: repo configs listed via=api count=' + configs.length); return successResponse(configs, { status: 200 }, req); }; diff --git a/src/app/api/v2/ai/agent/runtime-config/route.test.ts b/src/app/api/v2/ai/agent/runtime-config/route.test.ts new file mode 100644 index 00000000..a7c3a74d --- /dev/null +++ b/src/app/api/v2/ai/agent/runtime-config/route.test.ts @@ -0,0 +1,95 @@ +/** + * 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 mockGetUser = jest.fn(); +const mockSetGlobalConfig = jest.fn(); +const mockUpdateGlobalAdditiveRules = jest.fn(); +const mockUpdateGlobalApprovalPolicy = jest.fn(); + +jest.mock('server/lib/get-user', () => ({ + getUser: (...args: unknown[]) => mockGetUser(...args), +})); + +jest.mock('server/services/agentRuntime/config/agentRuntimeConfig', () => ({ + __esModule: true, + default: { + getInstance: jest.fn(() => ({ + setGlobalConfig: (...args: unknown[]) => mockSetGlobalConfig(...args), + updateGlobalAdditiveRules: (...args: unknown[]) => mockUpdateGlobalAdditiveRules(...args), + updateGlobalApprovalPolicy: (...args: unknown[]) => mockUpdateGlobalApprovalPolicy(...args), + })), + }, +})); + +import { PATCH, PUT } from './route'; + +function makeRequest(body: unknown): NextRequest { + return { + headers: new Headers([['x-request-id', 'req-test']]), + nextUrl: new URL('http://localhost/api/v2/ai/agent/runtime-config'), + json: jest.fn().mockResolvedValue(body), + } as unknown as NextRequest; +} + +describe('/api/v2/ai/agent/runtime-config', () => { + const originalEnableAuth = process.env.ENABLE_AUTH; + + beforeEach(() => { + jest.clearAllMocks(); + process.env.ENABLE_AUTH = 'true'; + mockGetUser.mockReturnValue({ + sub: 'sample-user', + realm_access: { + roles: ['user'], + }, + }); + }); + + afterEach(() => { + if (originalEnableAuth === undefined) { + delete process.env.ENABLE_AUTH; + } else { + process.env.ENABLE_AUTH = originalEnableAuth; + } + }); + + it('rejects non-admin global config replacement before parsing the body', async () => { + const request = makeRequest({}); + + const response = await PUT(request); + const body = await response.json(); + + expect(response.status).toBe(403); + expect(body.error.message).toBe('Forbidden: insufficient permissions'); + expect(request.json).not.toHaveBeenCalled(); + expect(mockSetGlobalConfig).not.toHaveBeenCalled(); + }); + + it('rejects non-admin global config patches before parsing the body', async () => { + const request = makeRequest({ additiveRules: [] }); + + const response = await PATCH(request); + const body = await response.json(); + + expect(response.status).toBe(403); + expect(body.error.message).toBe('Forbidden: insufficient permissions'); + expect(request.json).not.toHaveBeenCalled(); + expect(mockUpdateGlobalAdditiveRules).not.toHaveBeenCalled(); + expect(mockUpdateGlobalApprovalPolicy).not.toHaveBeenCalled(); + }); +}); diff --git a/src/app/api/v2/ai/agent-config/route.ts b/src/app/api/v2/ai/agent/runtime-config/route.ts similarity index 64% rename from src/app/api/v2/ai/agent-config/route.ts rename to src/app/api/v2/ai/agent/runtime-config/route.ts index 0a173bb9..80abc05a 100644 --- a/src/app/api/v2/ai/agent-config/route.ts +++ b/src/app/api/v2/ai/agent/runtime-config/route.ts @@ -18,37 +18,37 @@ import { NextRequest } from 'next/server'; import { createApiHandler } from 'server/lib/createApiHandler'; import { errorResponse, successResponse } from 'server/lib/response'; import { getLogger } from 'server/lib/logger'; -import AIAgentConfigService from 'server/services/aiAgentConfig'; +import AgentRuntimeConfigService from 'server/services/agentRuntime/config/agentRuntimeConfig'; import JsonSchema from 'jsonschema'; import { - aiAgentAdditiveRulesUpdateSchema, - aiAgentApprovalPolicyUpdateSchema, - aiAgentConfigPatchSchema, - aiAgentConfigSchema, -} from 'server/lib/validation/aiAgentConfigSchemas'; -import { AIAgentConfigValidationError } from 'server/lib/validation/aiAgentConfigValidator'; -import type { ApprovalPolicyConfig } from 'server/services/types/aiAgentConfig'; + agentRuntimeAdditiveRulesUpdateSchema, + agentRuntimeApprovalPolicyUpdateSchema, + agentRuntimeConfigPatchSchema, + agentRuntimeConfigSchema, +} from 'server/lib/validation/agentRuntimeConfigSchemas'; +import { AgentRuntimeConfigValidationError } from 'server/lib/validation/agentRuntimeConfigValidator'; +import type { ApprovalPolicyConfig } from 'server/services/types/agentRuntimeConfig'; /** * @openapi - * /api/v2/ai/agent-config: + * /api/v2/ai/agent/runtime-config: * get: - * summary: Get global AI agent configuration + * summary: Get global Agent runtime configuration * description: > - * Returns the current global AI agent configuration. This is the base configuration + * Returns the current global Agent runtime configuration. This is the base configuration * that applies to all repositories unless overridden by a per-repository config. * The response includes provider settings, model definitions, session limits, * tool exclusions, and performance tuning parameters. * tags: - * - AI Agent Config - * operationId: getGlobalAIAgentConfig + * - Agent Runtime Config + * operationId: getGlobalAgentRuntimeConfig * responses: * '200': - * description: Global AI agent configuration + * description: Global Agent runtime configuration * content: * application/json: * schema: - * $ref: '#/components/schemas/GetGlobalAIAgentConfigSuccessResponse' + * $ref: '#/components/schemas/GetGlobalAgentRuntimeConfigSuccessResponse' * '500': * description: Server error * content: @@ -57,39 +57,39 @@ import type { ApprovalPolicyConfig } from 'server/services/types/aiAgentConfig'; * $ref: '#/components/schemas/ApiErrorResponse' */ const getHandler = async (req: NextRequest) => { - const service = AIAgentConfigService.getInstance(); + const service = AgentRuntimeConfigService.getInstance(); const config = await service.getGlobalConfig(); - getLogger().info('AIAgentConfig: global config read via=api'); + getLogger().info('AgentRuntimeConfig: global config read via=api'); return successResponse(config, { status: 200 }, req); }; /** * @openapi - * /api/v2/ai/agent-config: + * /api/v2/ai/agent/runtime-config: * put: - * summary: Update global AI agent configuration + * summary: Update global Agent runtime configuration * description: > - * Validates and replaces the global AI agent configuration. The full AIAgentConfig + * Validates and replaces the global Agent runtime configuration. The full AgentRuntimeConfig * object must be provided (this is a full replacement, not a patch). Validation * enforces limits on systemPromptOverride length (max 50,000 chars), excludedFilePatterns * count, and prevents exclusion of core tools. The updated configuration is returned. * Changes take effect immediately and invalidate cached effective configs. * tags: - * - AI Agent Config - * operationId: updateGlobalAIAgentConfig + * - Agent Runtime Config + * operationId: updateGlobalAgentRuntimeConfig * requestBody: * required: true * content: * application/json: * schema: - * $ref: '#/components/schemas/AIAgentConfig' + * $ref: '#/components/schemas/AgentRuntimeConfig' * responses: * '200': - * description: Updated global AI agent configuration + * description: Updated global Agent runtime configuration * content: * application/json: * schema: - * $ref: '#/components/schemas/GetGlobalAIAgentConfigSuccessResponse' + * $ref: '#/components/schemas/GetGlobalAgentRuntimeConfigSuccessResponse' * '400': * description: > * Validation error. Possible reasons: JSON schema violation, @@ -99,6 +99,18 @@ const getHandler = async (req: NextRequest) => { * application/json: * schema: * $ref: '#/components/schemas/ApiErrorResponse' + * '401': + * description: Unauthorized + * content: + * application/json: + * schema: + * $ref: '#/components/schemas/ApiErrorResponse' + * '403': + * description: Forbidden + * content: + * application/json: + * schema: + * $ref: '#/components/schemas/ApiErrorResponse' * '500': * description: Server error * content: @@ -106,35 +118,47 @@ const getHandler = async (req: NextRequest) => { * schema: * $ref: '#/components/schemas/ApiErrorResponse' * patch: - * summary: Patch global AI agent configuration + * summary: Patch global Agent runtime configuration * description: > - * Updates one patchable global AI agent configuration section without replacing + * Updates one patchable global Agent runtime configuration section without replacing * the rest of the configuration. Supported patch targets are additive rules and * approval policy. Approval policy updates replace that section with the provided * value, so omitting defaultMode or rules clears them. This avoids revalidating * unrelated provider/model settings. * tags: - * - AI Agent Config - * operationId: patchGlobalAIAgentConfig + * - Agent Runtime Config + * operationId: patchGlobalAgentRuntimeConfig * requestBody: * required: true * content: * application/json: * schema: - * $ref: '#/components/schemas/AIAgentConfigPatchRequest' + * $ref: '#/components/schemas/AgentRuntimeConfigPatchRequest' * responses: * '200': - * description: Updated global AI agent configuration + * description: Updated global Agent runtime configuration * content: * application/json: * schema: - * $ref: '#/components/schemas/GetGlobalAIAgentConfigSuccessResponse' + * $ref: '#/components/schemas/GetGlobalAgentRuntimeConfigSuccessResponse' * '400': * description: Validation error * content: * application/json: * schema: * $ref: '#/components/schemas/ApiErrorResponse' + * '401': + * description: Unauthorized + * content: + * application/json: + * schema: + * $ref: '#/components/schemas/ApiErrorResponse' + * '403': + * description: Forbidden + * content: + * application/json: + * schema: + * $ref: '#/components/schemas/ApiErrorResponse' * '500': * description: Server error * content: @@ -151,19 +175,19 @@ const putHandler = async (req: NextRequest) => { } const validator = new JsonSchema.Validator(); - const result = validator.validate(body, aiAgentConfigSchema); + const result = validator.validate(body, agentRuntimeConfigSchema); if (!result.valid) { const messages = result.errors.map((e) => e.stack).join('; '); return errorResponse(new Error('Validation failed: ' + messages), { status: 400 }, req); } - const service = AIAgentConfigService.getInstance(); + const service = AgentRuntimeConfigService.getInstance(); try { await service.setGlobalConfig(body as any); } catch (error) { - if (error instanceof AIAgentConfigValidationError) { + if (error instanceof AgentRuntimeConfigValidationError) { return errorResponse(error, { status: 400 }, req); } throw error; @@ -182,18 +206,18 @@ const patchHandler = async (req: NextRequest) => { } const validator = new JsonSchema.Validator(); - const result = validator.validate(body, aiAgentConfigPatchSchema); + const result = validator.validate(body, agentRuntimeConfigPatchSchema); if (!result.valid) { const messages = result.errors.map((e) => e.stack).join('; '); return errorResponse(new Error('Validation failed: ' + messages), { status: 400 }, req); } - const service = AIAgentConfigService.getInstance(); + const service = AgentRuntimeConfigService.getInstance(); try { if ('additiveRules' in (body as Record)) { - const additiveRulesResult = validator.validate(body, aiAgentAdditiveRulesUpdateSchema); + const additiveRulesResult = validator.validate(body, agentRuntimeAdditiveRulesUpdateSchema); if (!additiveRulesResult.valid) { const messages = additiveRulesResult.errors.map((e) => e.stack).join('; '); return errorResponse(new Error('Validation failed: ' + messages), { status: 400 }, req); @@ -205,7 +229,7 @@ const patchHandler = async (req: NextRequest) => { return successResponse(updatedConfig, { status: 200 }, req); } - const approvalPolicyResult = validator.validate(body, aiAgentApprovalPolicyUpdateSchema); + const approvalPolicyResult = validator.validate(body, agentRuntimeApprovalPolicyUpdateSchema); if (!approvalPolicyResult.valid) { const messages = approvalPolicyResult.errors.map((e) => e.stack).join('; '); return errorResponse(new Error('Validation failed: ' + messages), { status: 400 }, req); @@ -216,7 +240,7 @@ const patchHandler = async (req: NextRequest) => { ); return successResponse(updatedConfig, { status: 200 }, req); } catch (error) { - if (error instanceof AIAgentConfigValidationError) { + if (error instanceof AgentRuntimeConfigValidationError) { return errorResponse(error, { status: 400 }, req); } throw error; @@ -224,5 +248,5 @@ const patchHandler = async (req: NextRequest) => { }; export const GET = createApiHandler(getHandler); -export const PUT = createApiHandler(putHandler); -export const PATCH = createApiHandler(patchHandler); +export const PUT = createApiHandler(putHandler, { roles: ['admin'] }); +export const PATCH = createApiHandler(patchHandler, { roles: ['admin'] }); diff --git a/src/app/api/v2/ai/agent/runtime-controls/preview/route.test.ts b/src/app/api/v2/ai/agent/runtime-controls/preview/route.test.ts new file mode 100644 index 00000000..2c8a80d7 --- /dev/null +++ b/src/app/api/v2/ai/agent/runtime-controls/preview/route.test.ts @@ -0,0 +1,128 @@ +/** + * 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 mockGetRequestUserIdentity = jest.fn(); +const mockGetEntryPreview = jest.fn(); + +jest.mock('server/lib/get-user', () => ({ + getRequestUserIdentity: (...args: unknown[]) => mockGetRequestUserIdentity(...args), +})); + +jest.mock('server/services/agent/ThreadRuntimeControlsService', () => { + class AgentThreadRuntimeControlsError extends Error { + constructor(public readonly code: string, message: string) { + super(message); + this.name = 'AgentThreadRuntimeControlsError'; + } + } + + return { + __esModule: true, + default: { + getEntryPreview: (...args: unknown[]) => mockGetEntryPreview(...args), + }, + AgentThreadRuntimeControlsError, + }; +}); + +import { POST } from './route'; +import { AgentThreadRuntimeControlsError } from 'server/services/agent/ThreadRuntimeControlsService'; + +const previewState = { + tools: { + required: [], + optional: [ + { + id: 'rtc_optional', + label: 'Workspace files', + description: 'Work with files.', + required: false, + selected: true, + available: true, + }, + ], + selectedChoiceIds: ['rtc_optional'], + }, + mcp: { + connections: [], + selectedChoiceIds: [], + }, + canEdit: true, + disabledReason: null, +}; + +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/runtime-controls/preview'), + } as unknown as NextRequest; +} + +describe('/api/v2/ai/agent/runtime-controls/preview', () => { + beforeEach(() => { + jest.clearAllMocks(); + mockGetRequestUserIdentity.mockReturnValue({ + userId: 'sample-user', + githubUsername: 'sample-user', + }); + mockGetEntryPreview.mockResolvedValue(previewState); + }); + + it('returns sanitized /new runtime choices without a thread', async () => { + const body = { + agentId: 'custom.sample-agent', + source: { adapter: 'lifecycle_fork', input: { repo: 'example-org/example-repo' } }, + defaults: { provider: 'openai', model: 'sample-model' }, + runtimeControlChoices: { toolChoiceIds: ['rtc_optional'], mcpChoiceIds: [] }, + }; + + const response = await POST(makeRequest(body)); + const payload = await response.json(); + + expect(response.status).toBe(200); + expect(mockGetEntryPreview).toHaveBeenCalledWith({ + userIdentity: { userId: 'sample-user', githubUsername: 'sample-user' }, + ...body, + }); + expect(payload.data).toEqual(previewState); + }); + + it('returns 400 for malformed and unknown choices', async () => { + const malformed = await POST(makeRequest({ runtimeControlChoices: 'workspace_files' })); + expect(malformed.status).toBe(400); + + mockGetEntryPreview.mockRejectedValueOnce( + new AgentThreadRuntimeControlsError('unknown_choice', 'Unknown runtime control choice.') + ); + const unknown = await POST(makeRequest({ runtimeControlChoices: { toolChoiceIds: ['workspace_files'] } })); + expect(unknown.status).toBe(400); + }); + + it('returns 403 for unavailable choices and 401 without identity', async () => { + mockGetEntryPreview.mockRejectedValueOnce( + new AgentThreadRuntimeControlsError('policy_denied', 'Runtime control choice is unavailable.') + ); + const denied = await POST(makeRequest({ runtimeControlChoices: { toolChoiceIds: ['rtc_optional'] } })); + expect(denied.status).toBe(403); + + mockGetRequestUserIdentity.mockReturnValueOnce(null); + const unauthorized = await POST(makeRequest()); + expect(unauthorized.status).toBe(401); + }); +}); diff --git a/src/app/api/v2/ai/agent/runtime-controls/preview/route.ts b/src/app/api/v2/ai/agent/runtime-controls/preview/route.ts new file mode 100644 index 00000000..d423eb1e --- /dev/null +++ b/src/app/api/v2/ai/agent/runtime-controls/preview/route.ts @@ -0,0 +1,224 @@ +/** + * 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 { getRequestUserIdentity } from 'server/lib/get-user'; +import { errorResponse, successResponse } from 'server/lib/response'; +import AgentThreadRuntimeControlsService, { + AgentThreadRuntimeControlsError, + type AgentRuntimeControlsEntryDefaultsInput, + type AgentRuntimeControlsEntrySourceInput, + type AgentThreadRuntimeControlChoiceInput, +} from 'server/services/agent/ThreadRuntimeControlsService'; + +type RuntimeControlsPreviewBody = { + agentId?: string | null; + source?: AgentRuntimeControlsEntrySourceInput; + defaults?: AgentRuntimeControlsEntryDefaultsInput; + runtimeControlChoices?: AgentThreadRuntimeControlChoiceInput; +}; + +function isPlainObject(value: unknown): value is Record { + return value != null && typeof value === 'object' && !Array.isArray(value); +} + +function parseStringArray(value: unknown, fieldName: string): string[] | undefined | Error { + if (value === undefined) { + return undefined; + } + + if (!Array.isArray(value)) { + return new Error(`${fieldName} must be an array of choice ids.`); + } + + for (const item of value) { + if (typeof item !== 'string' || !item.trim()) { + return new Error(`${fieldName} must contain only choice ids.`); + } + } + + return value.map((item) => item.trim()); +} + +function parseRuntimeControlChoices(value: unknown): AgentThreadRuntimeControlChoiceInput | undefined | Error { + if (value === undefined) { + return undefined; + } + + if (!isPlainObject(value)) { + return new Error('runtimeControlChoices must be an object.'); + } + + const unknownKeys = Object.keys(value).filter( + (key) => key !== 'agentId' && key !== 'toolChoiceIds' && key !== 'mcpChoiceIds' + ); + if (unknownKeys.length > 0) { + return new Error(`Unsupported runtime-control fields: ${unknownKeys.join(', ')}.`); + } + + const agentId = + value.agentId === undefined || value.agentId === null + ? undefined + : typeof value.agentId === 'string' && value.agentId.trim() + ? value.agentId.trim() + : new Error('runtimeControlChoices.agentId must be a non-empty string.'); + if (agentId instanceof Error) { + return agentId; + } + + const toolChoiceIds = parseStringArray(value.toolChoiceIds, 'runtimeControlChoices.toolChoiceIds'); + if (toolChoiceIds instanceof Error) { + return toolChoiceIds; + } + + const mcpChoiceIds = parseStringArray(value.mcpChoiceIds, 'runtimeControlChoices.mcpChoiceIds'); + if (mcpChoiceIds instanceof Error) { + return mcpChoiceIds; + } + + return { agentId, toolChoiceIds, mcpChoiceIds }; +} + +function parsePreviewBody(body: unknown): RuntimeControlsPreviewBody | Error { + if (!isPlainObject(body)) { + return new Error('Request body must be an object.'); + } + + const unknownKeys = Object.keys(body).filter( + (key) => key !== 'agentId' && key !== 'source' && key !== 'defaults' && key !== 'runtimeControlChoices' + ); + if (unknownKeys.length > 0) { + return new Error(`Unsupported runtime-control preview fields: ${unknownKeys.join(', ')}.`); + } + + const agentId = + body.agentId === undefined || body.agentId === null + ? undefined + : typeof body.agentId === 'string' && body.agentId.trim() + ? body.agentId.trim() + : new Error('agentId must be a non-empty string.'); + if (agentId instanceof Error) { + return agentId; + } + + if (body.source !== undefined && !isPlainObject(body.source)) { + return new Error('source must be an object.'); + } + + if (body.defaults !== undefined && !isPlainObject(body.defaults)) { + return new Error('defaults must be an object.'); + } + + const runtimeControlChoices = parseRuntimeControlChoices(body.runtimeControlChoices); + if (runtimeControlChoices instanceof Error) { + return runtimeControlChoices; + } + + return { + agentId, + source: body.source as AgentRuntimeControlsEntrySourceInput | undefined, + defaults: body.defaults as AgentRuntimeControlsEntryDefaultsInput | undefined, + runtimeControlChoices, + }; +} + +function mapRuntimeControlsError(error: unknown, req: NextRequest) { + if (error instanceof AgentThreadRuntimeControlsError) { + const statusByCode: Record = { + invalid_input: 400, + unknown_choice: 400, + policy_denied: 403, + not_found: 404, + active_run: 409, + }; + return errorResponse(error, { status: statusByCode[error.code] }, req); + } + + throw error; +} + +/** + * @openapi + * /api/v2/ai/agent/runtime-controls/preview: + * post: + * summary: Preview runtime control choices for a new run or session + * description: Previews available runtime control choices for the selected agent, source, defaults, and draft runtime-control choices. + * tags: + * - Agent Platform + * operationId: agentRuntimeControlsPreview + * requestBody: + * required: true + * content: + * application/json: + * schema: + * $ref: '#/components/schemas/AgentRuntimeControlsPreviewRequest' + * responses: + * '200': + * description: Runtime controls preview + * content: + * application/json: + * schema: + * allOf: + * - $ref: '#/components/schemas/SuccessApiResponse' + * - type: object + * required: [data] + * properties: + * data: + * $ref: '#/components/schemas/AgentThreadRuntimeControlsState' + * '400': + * description: Invalid runtime-control choices + * content: + * application/json: + * schema: + * $ref: '#/components/schemas/ApiErrorResponse' + * '401': + * description: Unauthorized + * content: + * application/json: + * schema: + * $ref: '#/components/schemas/ApiErrorResponse' + * '403': + * description: Runtime-control choice is unavailable + * content: + * application/json: + * schema: + * $ref: '#/components/schemas/ApiErrorResponse' + */ +const postHandler = async (req: NextRequest) => { + const userIdentity = getRequestUserIdentity(req); + if (!userIdentity) { + return errorResponse(new Error('Unauthorized'), { status: 401 }, req); + } + + const parsedBody = parsePreviewBody(await req.json().catch(() => ({}))); + if (parsedBody instanceof Error) { + return errorResponse(parsedBody, { status: 400 }, req); + } + + try { + const state = await AgentThreadRuntimeControlsService.getEntryPreview({ + userIdentity, + ...parsedBody, + }); + return successResponse(state, { status: 200 }, req); + } catch (error) { + return mapRuntimeControlsError(error, req); + } +}; + +export const POST = createApiHandler(postHandler); diff --git a/src/app/api/v2/ai/agent/sessions/[sessionId]/threads/route.test.ts b/src/app/api/v2/ai/agent/sessions/[sessionId]/threads/route.test.ts new file mode 100644 index 00000000..ee0a9ce9 --- /dev/null +++ b/src/app/api/v2/ai/agent/sessions/[sessionId]/threads/route.test.ts @@ -0,0 +1,182 @@ +/** + * 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/dependencies', () => ({ + defaultDb: {}, + defaultRedis: {}, +})); + +jest.mock('server/lib/get-user', () => ({ + getRequestUserIdentity: jest.fn(), +})); + +jest.mock('server/services/agent/ThreadService', () => ({ + __esModule: true, + default: { + createThread: jest.fn(), + listThreadsForSession: jest.fn(), + serializeThread: jest.fn((thread, sessionId) => ({ + id: thread.uuid, + sessionId, + title: thread.title, + isDefault: thread.isDefault, + metadata: thread.metadata ?? {}, + })), + }, +})); + +import { GET, POST } from './route'; +import { getRequestUserIdentity } from 'server/lib/get-user'; +import AgentThreadService from 'server/services/agent/ThreadService'; + +const mockGetRequestUserIdentity = getRequestUserIdentity as jest.Mock; +const mockCreateThread = AgentThreadService.createThread as jest.Mock; +const mockListThreadsForSession = AgentThreadService.listThreadsForSession as jest.Mock; +const mockSerializeThread = AgentThreadService.serializeThread 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/sessions/session-1/threads'), + } as unknown as NextRequest; +} + +describe('/api/v2/ai/agent/sessions/[sessionId]/threads', () => { + beforeEach(() => { + jest.clearAllMocks(); + mockGetRequestUserIdentity.mockReturnValue({ + userId: 'sample-user', + githubUsername: 'sample-user', + }); + }); + + it('lists owned session threads', async () => { + mockListThreadsForSession.mockResolvedValue([ + { + uuid: 'thread-1', + title: 'Default thread', + isDefault: true, + metadata: {}, + }, + ]); + + const response = await GET(makeRequest(), { + params: { sessionId: 'session-1' }, + }); + const body = await response.json(); + + expect(response.status).toBe(200); + expect(mockListThreadsForSession).toHaveBeenCalledWith('session-1', 'sample-user'); + expect(body.data.threads).toEqual([ + { + id: 'thread-1', + sessionId: 'session-1', + title: 'Default thread', + isDefault: true, + metadata: {}, + }, + ]); + }); + + it('maps missing sessions during list to 404', async () => { + mockListThreadsForSession.mockRejectedValueOnce(new Error('Agent session not found')); + + const response = await GET(makeRequest(), { + params: { sessionId: 'missing-session' }, + }); + const body = await response.json(); + + expect(response.status).toBe(404); + expect(body.error.message).toBe('Agent session not found'); + }); + + it('creates a thread in an active owned session', async () => { + mockCreateThread.mockResolvedValue({ + uuid: 'thread-2', + title: 'New chat', + isDefault: false, + metadata: {}, + }); + + const response = await POST(makeRequest({ title: 'New chat' }), { + params: { sessionId: 'session-1' }, + }); + const body = await response.json(); + + expect(response.status).toBe(201); + expect(mockCreateThread).toHaveBeenCalledWith('session-1', 'sample-user', 'New chat'); + expect(body.data).toEqual({ + id: 'thread-2', + sessionId: 'session-1', + title: 'New chat', + isDefault: false, + metadata: {}, + }); + }); + + it('rejects new threads for inactive sessions', async () => { + mockCreateThread.mockRejectedValueOnce(new Error('Cannot create a thread for an inactive session')); + + const response = await POST(makeRequest({ title: 'New chat' }), { + params: { sessionId: 'session-1' }, + }); + const body = await response.json(); + + expect(response.status).toBe(409); + expect(body.error.message).toBe('Cannot create a thread for an inactive session'); + }); + + it('maps runtime-unavailable sessions during create to 409', async () => { + mockCreateThread.mockRejectedValueOnce(new Error('This session is no longer available for new messages.')); + + const response = await POST(makeRequest({ title: 'New chat' }), { + params: { sessionId: 'session-1' }, + }); + const body = await response.json(); + + expect(response.status).toBe(409); + expect(body.error.message).toBe('This session is no longer available for new messages.'); + }); + + it('maps missing sessions during create to 404', async () => { + mockCreateThread.mockRejectedValueOnce(new Error('Agent session not found')); + + const response = await POST(makeRequest({ title: 'New chat' }), { + params: { sessionId: 'missing-session' }, + }); + const body = await response.json(); + + expect(response.status).toBe(404); + expect(body.error.message).toBe('Agent session not found'); + }); + + it('rejects unauthenticated requests', async () => { + mockGetRequestUserIdentity.mockReturnValueOnce(null); + + const response = await POST(makeRequest({ title: 'New chat' }), { + params: { sessionId: 'session-1' }, + }); + const body = await response.json(); + + expect(response.status).toBe(401); + expect(body.error.message).toBe('Unauthorized'); + expect(mockCreateThread).not.toHaveBeenCalled(); + expect(mockSerializeThread).not.toHaveBeenCalled(); + }); +}); diff --git a/src/app/api/v2/ai/agent/sessions/[sessionId]/threads/route.ts b/src/app/api/v2/ai/agent/sessions/[sessionId]/threads/route.ts index 4091cd92..7cd20e74 100644 --- a/src/app/api/v2/ai/agent/sessions/[sessionId]/threads/route.ts +++ b/src/app/api/v2/ai/agent/sessions/[sessionId]/threads/route.ts @@ -54,6 +54,10 @@ import AgentThreadService from 'server/services/agent/ThreadService'; * type: array * items: * $ref: '#/components/schemas/AgentThread' + * '401': + * description: Unauthorized + * '404': + * description: Agent session not found * post: * summary: Create a new thread in an agent session * tags: @@ -87,6 +91,12 @@ import AgentThreadService from 'server/services/agent/ThreadService'; * properties: * data: * $ref: '#/components/schemas/AgentThread' + * '401': + * description: Unauthorized + * '404': + * description: Agent session not found + * '409': + * description: Session cannot create new threads in its current state */ const getHandler = async (req: NextRequest, { params }: { params: { sessionId: string } }) => { const userIdentity = getRequestUserIdentity(req); @@ -94,14 +104,22 @@ const getHandler = async (req: NextRequest, { params }: { params: { sessionId: s return errorResponse(new Error('Unauthorized'), { status: 401 }, req); } - const threads = await AgentThreadService.listThreadsForSession(params.sessionId, userIdentity.userId); - return successResponse( - { - threads: threads.map((thread) => AgentThreadService.serializeThread(thread, params.sessionId)), - }, - { status: 200 }, - req - ); + try { + const threads = await AgentThreadService.listThreadsForSession(params.sessionId, userIdentity.userId); + return successResponse( + { + threads: threads.map((thread) => AgentThreadService.serializeThread(thread, params.sessionId)), + }, + { status: 200 }, + req + ); + } catch (error) { + if (error instanceof Error && error.message === 'Agent session not found') { + return errorResponse(error, { status: 404 }, req); + } + + throw error; + } }; const postHandler = async (req: NextRequest, { params }: { params: { sessionId: string } }) => { @@ -111,9 +129,29 @@ const postHandler = async (req: NextRequest, { params }: { params: { sessionId: } const body = await req.json().catch(() => ({})); - const thread = await AgentThreadService.createThread(params.sessionId, userIdentity.userId, body?.title); + try { + const thread = await AgentThreadService.createThread(params.sessionId, userIdentity.userId, body?.title); + + return successResponse(AgentThreadService.serializeThread(thread, params.sessionId), { status: 201 }, req); + } catch (error) { + if (error instanceof Error && error.message === 'Agent session not found') { + return errorResponse(error, { status: 404 }, req); + } - return successResponse(AgentThreadService.serializeThread(thread, params.sessionId), { status: 201 }, req); + if (error instanceof Error && error.message === 'Cannot create a thread for an inactive session') { + return errorResponse(error, { status: 409 }, req); + } + + if ( + error instanceof Error && + (error.message === 'Wait for the session to finish starting before sending a message.' || + error.message === 'This session is no longer available for new messages.') + ) { + return errorResponse(error, { status: 409 }, req); + } + + throw error; + } }; export const GET = createApiHandler(getHandler); diff --git a/src/app/api/v2/ai/agent/sessions/route.test.ts b/src/app/api/v2/ai/agent/sessions/route.test.ts new file mode 100644 index 00000000..bc1ff65c --- /dev/null +++ b/src/app/api/v2/ai/agent/sessions/route.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. + */ + +import { NextRequest } from 'next/server'; + +const mockGetRequestUserIdentity = jest.fn(); +const mockCreateChatSession = jest.fn(); +const mockSerializeSessionRecord = jest.fn(); +const mockResolveAgentSessionRuntimeConfig = jest.fn(); +const mockResolveAgentSessionWorkspaceStorageIntent = jest.fn(); + +jest.mock('server/lib/get-user', () => ({ + getRequestUserIdentity: (...args: unknown[]) => mockGetRequestUserIdentity(...args), +})); + +jest.mock('server/services/agent/ChatSessionService', () => ({ + __esModule: true, + default: { + createChatSession: (...args: unknown[]) => mockCreateChatSession(...args), + }, +})); + +jest.mock('server/services/agent/SessionReadService', () => ({ + __esModule: true, + default: { + listOwnedSessionRecords: jest.fn(), + serializeSessionRecord: (...args: unknown[]) => mockSerializeSessionRecord(...args), + }, + DEFAULT_AGENT_SESSION_LIST_LIMIT: 25, + MAX_AGENT_SESSION_LIST_LIMIT: 100, +})); + +jest.mock('server/lib/agentSession/runtimeConfig', () => { + class AgentSessionRuntimeConfigError extends Error {} + class AgentSessionWorkspaceStorageConfigError extends Error {} + + return { + resolveAgentSessionRuntimeConfig: (...args: unknown[]) => mockResolveAgentSessionRuntimeConfig(...args), + resolveAgentSessionWorkspaceStorageIntent: (...args: unknown[]) => + mockResolveAgentSessionWorkspaceStorageIntent(...args), + AgentSessionRuntimeConfigError, + AgentSessionWorkspaceStorageConfigError, + }; +}); + +jest.mock('server/services/agent/ProviderRegistry', () => { + class MissingAgentProviderApiKeyError extends Error {} + return { + __esModule: true, + default: {}, + MissingAgentProviderApiKeyError, + }; +}); + +jest.mock('server/services/agent/ThreadRuntimeControlsService', () => { + class AgentThreadRuntimeControlsError extends Error { + constructor(public readonly code: string, message: string) { + super(message); + this.name = 'AgentThreadRuntimeControlsError'; + } + } + + return { + __esModule: true, + AgentThreadRuntimeControlsError, + }; +}); + +import { POST } from './route'; +import { AgentThreadRuntimeControlsError } from 'server/services/agent/ThreadRuntimeControlsService'; + +const userIdentity = { + userId: 'sample-user', + githubUsername: 'sample-user', + preferredUsername: 'sample-user', + email: 'sample-user@example.com', + firstName: 'Sample', + lastName: 'User', + displayName: 'Sample User', + gitUserName: 'Sample User', + gitUserEmail: 'sample-user@example.com', +}; + +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/sessions'), + } as unknown as NextRequest; +} + +describe('/api/v2/ai/agent/sessions runtimeControlChoices', () => { + beforeEach(() => { + jest.clearAllMocks(); + mockGetRequestUserIdentity.mockReturnValue(userIdentity); + mockResolveAgentSessionRuntimeConfig.mockResolvedValue({ + workspaceStorage: {}, + }); + mockResolveAgentSessionWorkspaceStorageIntent.mockReturnValue(undefined); + mockCreateChatSession.mockResolvedValue({ uuid: 'session-1' }); + mockSerializeSessionRecord.mockResolvedValue({ + session: { + id: 'session-1', + status: 'ready', + userId: 'sample-user', + ownerGithubUsername: 'sample-user', + defaults: { provider: 'openai', model: 'sample-model', harness: null }, + defaultThreadId: 'thread-1', + }, + source: {}, + sandbox: {}, + }); + }); + + it('passes runtimeControlChoices to chat session creation', async () => { + const body = { + defaults: { provider: 'openai', model: 'sample-model' }, + source: { adapter: 'blank_workspace', input: {} }, + runtimeControlChoices: { + agentId: 'custom.sample-agent', + toolChoiceIds: ['rtc_tool_choice'], + mcpChoiceIds: [], + }, + }; + + const response = await POST(makeRequest(body)); + + expect(response.status).toBe(201); + expect(mockCreateChatSession).toHaveBeenCalledWith( + expect.objectContaining({ + userId: 'sample-user', + userIdentity, + provider: 'openai', + model: 'sample-model', + runtimeControlChoices: body.runtimeControlChoices, + }) + ); + }); + + it('preserves current behavior when runtimeControlChoices is absent', async () => { + const response = await POST( + makeRequest({ + defaults: { provider: 'openai', model: 'sample-model' }, + source: { adapter: 'blank_workspace', input: {} }, + }) + ); + + expect(response.status).toBe(201); + expect(mockCreateChatSession).toHaveBeenCalledWith( + expect.not.objectContaining({ + runtimeControlChoices: expect.anything(), + }) + ); + }); + + it('maps invalid bootstrap runtime choices to 403', async () => { + mockCreateChatSession.mockRejectedValueOnce( + new AgentThreadRuntimeControlsError('policy_denied', 'Runtime control choice is unavailable.') + ); + + const response = await POST( + makeRequest({ + defaults: { provider: 'openai', model: 'sample-model' }, + source: { adapter: 'blank_workspace', input: {} }, + runtimeControlChoices: { + toolChoiceIds: ['rtc_denied'], + mcpChoiceIds: [], + }, + }) + ); + + expect(response.status).toBe(403); + }); +}); diff --git a/src/app/api/v2/ai/agent/sessions/route.ts b/src/app/api/v2/ai/agent/sessions/route.ts index 36f2bb29..18ecbdae 100644 --- a/src/app/api/v2/ai/agent/sessions/route.ts +++ b/src/app/api/v2/ai/agent/sessions/route.ts @@ -28,6 +28,10 @@ import { DEFAULT_AGENT_SESSION_LIST_LIMIT, MAX_AGENT_SESSION_LIST_LIMIT, } from 'server/services/agent/SessionReadService'; +import { + AgentThreadRuntimeControlsError, + type AgentThreadRuntimeControlChoiceInput, +} from 'server/services/agent/ThreadRuntimeControlsService'; import { AgentSessionKind, BuildKind } from 'shared/constants'; interface RequestedAgentSessionServiceRef { @@ -50,6 +54,7 @@ interface ResolvedSessionService { interface CreateSessionBody { defaults?: { + provider?: string; model?: string; harness?: string; }; @@ -60,6 +65,11 @@ interface CreateSessionBody { workspace?: { storageSize?: string; }; + runtimeControlChoices?: AgentThreadRuntimeControlChoiceInput; +} + +function isPlainObject(value: unknown): value is Record { + return value != null && typeof value === 'object' && !Array.isArray(value); } function repoNameFromRepoUrl(repoUrl?: string | null) { @@ -92,6 +102,78 @@ function parseRequestedWorkspaceStorageSize(body: CreateSessionBody): string | u return workspace.storageSize.trim(); } +function parseRuntimeControlChoiceIds(value: unknown, fieldName: string): string[] | undefined | Error { + if (value === undefined) { + return undefined; + } + + if (!Array.isArray(value)) { + return new Error(`${fieldName} must be an array of choice ids.`); + } + + for (const item of value) { + if (typeof item !== 'string' || !item.trim()) { + return new Error(`${fieldName} must contain only choice ids.`); + } + } + + return value.map((item) => item.trim()); +} + +function parseRuntimeControlChoices(value: unknown): AgentThreadRuntimeControlChoiceInput | undefined | Error { + if (value === undefined) { + return undefined; + } + + if (!isPlainObject(value)) { + return new Error('runtimeControlChoices must be an object'); + } + + const unknownKeys = Object.keys(value).filter( + (key) => key !== 'agentId' && key !== 'toolChoiceIds' && key !== 'mcpChoiceIds' + ); + if (unknownKeys.length > 0) { + return new Error(`Unsupported runtime-control fields: ${unknownKeys.join(', ')}`); + } + + const agentId = + value.agentId === undefined || value.agentId === null + ? undefined + : typeof value.agentId === 'string' && value.agentId.trim() + ? value.agentId.trim() + : new Error('runtimeControlChoices.agentId must be a non-empty string'); + if (agentId instanceof Error) { + return agentId; + } + + const toolChoiceIds = parseRuntimeControlChoiceIds(value.toolChoiceIds, 'runtimeControlChoices.toolChoiceIds'); + if (toolChoiceIds instanceof Error) { + return toolChoiceIds; + } + + const mcpChoiceIds = parseRuntimeControlChoiceIds(value.mcpChoiceIds, 'runtimeControlChoices.mcpChoiceIds'); + if (mcpChoiceIds instanceof Error) { + return mcpChoiceIds; + } + + return { agentId, toolChoiceIds, mcpChoiceIds }; +} + +function mapRuntimeControlsError(error: unknown, req: NextRequest) { + if (error instanceof AgentThreadRuntimeControlsError) { + const statusByCode: Record = { + invalid_input: 400, + unknown_choice: 400, + policy_denied: 403, + not_found: 404, + active_run: 409, + }; + return errorResponse(error, { status: statusByCode[error.code] }, req); + } + + return null; +} + async function resolveLifecycleConfigForSession({ buildContext, repoUrl, @@ -201,7 +283,7 @@ async function resolveRequestedServices( * get: * summary: List agent sessions for the authenticated user * tags: - * - Agent Sessions + * - Agent Platform * operationId: getAgentSessions * parameters: * - in: query @@ -251,8 +333,9 @@ async function resolveRequestedServices( * $ref: '#/components/schemas/ApiErrorResponse' * post: * summary: Create a new agent session + * description: Creates an agent session and accepts optional runtimeControlChoices for the first run's runtime-control choices. * tags: - * - Agent Sessions + * - Agent Platform * operationId: createAgentSession * requestBody: * required: true @@ -267,6 +350,8 @@ async function resolveRequestedServices( * properties: * model: * type: string + * provider: + * type: string * harness: * type: string * source: @@ -284,6 +369,8 @@ async function resolveRequestedServices( * storageSize: * type: string * description: Optional workspace PVC size. Accepted only when admin runtime settings allow client overrides. + * runtimeControlChoices: + * $ref: '#/components/schemas/AgentRuntimeControlChoicesInput' * sandbox: * type: object * properties: @@ -372,8 +459,14 @@ const postHandler = async (req: NextRequest) => { const body = (await req.json()) as CreateSessionBody; let requestedWorkspaceStorageSize: string | undefined; + let runtimeControlChoices: AgentThreadRuntimeControlChoiceInput | undefined; try { requestedWorkspaceStorageSize = parseRequestedWorkspaceStorageSize(body); + const parsedRuntimeControlChoices = parseRuntimeControlChoices(body.runtimeControlChoices); + if (parsedRuntimeControlChoices instanceof Error) { + throw parsedRuntimeControlChoices; + } + runtimeControlChoices = parsedRuntimeControlChoices; } catch (err) { return errorResponse(err, { status: 400 }, req); } @@ -389,6 +482,8 @@ const postHandler = async (req: NextRequest) => { ? (sourceInput as { services: unknown[] }).services : undefined; const requestedModel = body.defaults?.model; + const requestedProvider = + typeof body.defaults?.provider === 'string' ? body.defaults.provider.trim() || undefined : undefined; const sessionKind = body.source?.adapter === 'blank_workspace' ? AgentSessionKind.CHAT @@ -418,8 +513,10 @@ const postHandler = async (req: NextRequest) => { const session = await AgentChatSessionService.createChatSession({ userId: userIdentity.userId, userIdentity, + provider: requestedProvider, model: requestedModel, workspaceStorage, + runtimeControlChoices, }); return successResponse(await AgentSessionReadService.serializeSessionRecord(session), { status: 201 }, req); @@ -433,6 +530,10 @@ const postHandler = async (req: NextRequest) => { if (err instanceof AgentSessionRuntimeConfigError || err instanceof AgentSessionWorkspaceStorageConfigError) { return errorResponse(err, { status: 400 }, req); } + const runtimeControlsResponse = mapRuntimeControlsError(err, req); + if (runtimeControlsResponse) { + return runtimeControlsResponse; + } return errorResponse(err, { status: 500 }, req); } @@ -520,6 +621,7 @@ const postHandler = async (req: NextRequest) => { buildUuid, buildKind, services: resolvedServices, + provider: requestedProvider, model: requestedModel, environmentSkillRefs: lifecycleConfig?.environment?.agentSession?.skills, repoUrl, diff --git a/src/app/api/v2/ai/agent/threads/[threadId]/agent/route.test.ts b/src/app/api/v2/ai/agent/threads/[threadId]/agent/route.test.ts new file mode 100644 index 00000000..71f88d58 --- /dev/null +++ b/src/app/api/v2/ai/agent/threads/[threadId]/agent/route.test.ts @@ -0,0 +1,183 @@ +/** + * 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 mockGetRequestUserIdentity = jest.fn(); +const mockGetThreadAgentState = jest.fn(); +const mockSwitchThreadAgent = jest.fn(); + +jest.mock('server/lib/get-user', () => ({ + getRequestUserIdentity: (...args: unknown[]) => mockGetRequestUserIdentity(...args), +})); + +jest.mock('server/services/agent/AgentSelectionService', () => { + class AgentThreadAgentSwitchError extends Error { + constructor( + public readonly reason: string, + message: string, + public readonly details: Record = {} + ) { + super(message); + this.name = 'AgentThreadAgentSwitchError'; + } + } + + return { + __esModule: true, + default: { + getThreadAgentState: (...args: unknown[]) => mockGetThreadAgentState(...args), + switchThreadAgent: (...args: unknown[]) => mockSwitchThreadAgent(...args), + }, + AgentThreadAgentSwitchError, + }; +}); + +import { GET, PATCH } from './route'; +import { AgentThreadAgentSwitchError } from 'server/services/agent/AgentSelectionService'; + +const agentState = { + selectedId: null, + defaultId: 'system.freeform', + currentId: 'system.freeform', + groups: [ + { + id: 'built_in', + label: 'Built in', + agents: [ + { id: 'system.debug', ownerKind: 'system', label: 'Debug', group: 'built_in', available: true }, + { id: 'system.develop', ownerKind: 'system', label: 'Develop', group: 'built_in', available: false }, + { id: 'system.freeform', ownerKind: 'system', label: 'Free-form', group: 'built_in', available: true }, + ], + }, + { + id: 'my_agents', + label: 'My agents', + agents: [ + { + id: 'custom.sample-agent', + ownerKind: 'user', + label: 'Sample custom agent', + group: 'my_agents', + available: true, + }, + ], + }, + ], +}; + +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/agent'), + } as unknown as NextRequest; +} + +describe('/api/v2/ai/agent/threads/[threadId]/agent', () => { + beforeEach(() => { + jest.clearAllMocks(); + mockGetRequestUserIdentity.mockReturnValue({ + userId: 'sample-user', + githubUsername: 'sample-user', + }); + mockGetThreadAgentState.mockResolvedValue(agentState); + mockSwitchThreadAgent.mockResolvedValue({ + previousAgent: agentState.groups[0].agents[2], + nextAgent: agentState.groups[1].agents[0], + switched: true, + state: agentState, + }); + }); + + it('GET returns built_in and my_agents selection state', async () => { + const response = await GET(makeRequest(), { params: { threadId: 'thread-1' } }); + const body = await response.json(); + + expect(response.status).toBe(200); + expect(mockGetThreadAgentState).toHaveBeenCalledWith({ + threadId: 'thread-1', + userIdentity: { userId: 'sample-user', githubUsername: 'sample-user' }, + }); + expect(body.data.groups.map((group: { id: string }) => group.id)).toEqual(['built_in', 'my_agents']); + }); + + it('PATCH accepts only agentId and delegates the switch', async () => { + const response = await PATCH(makeRequest({ agentId: 'custom.sample-agent' }), { + params: { threadId: 'thread-1' }, + }); + + expect(response.status).toBe(200); + expect(mockSwitchThreadAgent).toHaveBeenCalledWith({ + threadId: 'thread-1', + userIdentity: { userId: 'sample-user', githubUsername: 'sample-user' }, + agentId: 'custom.sample-agent', + }); + }); + + it('returns 401 without identity', async () => { + mockGetRequestUserIdentity.mockReturnValueOnce(null); + + const response = await GET(makeRequest(), { params: { threadId: 'thread-1' } }); + + expect(response.status).toBe(401); + }); + + it('returns 404 for non-owned thread/session errors', async () => { + mockGetThreadAgentState.mockRejectedValueOnce(new Error('Agent thread not found')); + + const response = await GET(makeRequest(), { params: { threadId: 'thread-1' } }); + + expect(response.status).toBe(404); + }); + + it('returns 400 for non-string or unsupported agent switch bodies', async () => { + const response = await PATCH(makeRequest({ agentId: 'custom.sample-agent', another: true }), { + params: { threadId: 'thread-1' }, + }); + const body = await response.json(); + + expect(response.status).toBe(400); + expect(body.error.message).toContain('Unsupported switch request fields'); + + const missingIdResponse = await PATCH(makeRequest({ agentId: 7 }), { params: { threadId: 'thread-1' } }); + expect(missingIdResponse.status).toBe(400); + }); + + it('returns 400 for another user or unknown custom agent ids', async () => { + mockSwitchThreadAgent.mockRejectedValueOnce(new AgentThreadAgentSwitchError('unknown_agent', 'Unknown agent.')); + + const response = await PATCH(makeRequest({ agentId: 'custom.another-user-agent' }), { + params: { threadId: 'thread-1' }, + }); + + expect(response.status).toBe(400); + }); + + it('returns 409 for active run switch failures', async () => { + mockSwitchThreadAgent.mockRejectedValueOnce( + new AgentThreadAgentSwitchError('active_run', 'Wait for the current run to finish before switching agents.') + ); + + const response = await PATCH(makeRequest({ agentId: 'custom.sample-agent' }), { + params: { threadId: 'thread-1' }, + }); + const body = await response.json(); + + expect(response.status).toBe(409); + expect(body.error.message).toBe('Wait for the current run to finish before switching agents.'); + }); +}); diff --git a/src/app/api/v2/ai/agent/threads/[threadId]/agent/route.ts b/src/app/api/v2/ai/agent/threads/[threadId]/agent/route.ts new file mode 100644 index 00000000..daed2bee --- /dev/null +++ b/src/app/api/v2/ai/agent/threads/[threadId]/agent/route.ts @@ -0,0 +1,151 @@ +/** + * 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 AgentSelectionService, { AgentThreadAgentSwitchError } from 'server/services/agent/AgentSelectionService'; + +function mapAgentSelectionError(error: unknown, req: NextRequest) { + if (error instanceof AgentThreadAgentSwitchError) { + return errorResponse(error, { status: error.reason === 'unknown_agent' ? 400 : 409 }, req); + } + + if ( + error instanceof Error && + (error.message === 'Agent thread not found' || error.message === 'Agent session not found') + ) { + return errorResponse(error, { status: 404 }, req); + } + + throw error; +} + +/** + * @openapi + * /api/v2/ai/agent/threads/{threadId}/agent: + * get: + * summary: Get the current agent selection state for a thread + * description: Returns the selected agent that will be used when future runs are created for this thread. + * tags: + * - Agent Platform + * operationId: getAgentThreadSelection + * parameters: + * - in: path + * name: threadId + * required: true + * schema: + * type: string + * responses: + * '200': + * description: Agent selection state + * content: + * application/json: + * schema: + * allOf: + * - $ref: '#/components/schemas/SuccessApiResponse' + * - type: object + * required: [data] + * properties: + * data: + * $ref: '#/components/schemas/AgentSelectionState' + * patch: + * summary: Switch the agent used for future runs in a thread + * description: Updates the thread's selected agent for future runs without changing already-created runs. + * tags: + * - Agent Platform + * operationId: switchAgentThreadSelection + * parameters: + * - in: path + * name: threadId + * required: true + * schema: + * type: string + * requestBody: + * required: true + * content: + * application/json: + * schema: + * $ref: '#/components/schemas/SwitchAgentSelectionRequest' + * responses: + * '200': + * description: Agent selection switched or already selected + * content: + * application/json: + * schema: + * allOf: + * - $ref: '#/components/schemas/SuccessApiResponse' + * - type: object + * required: [data] + * properties: + * data: + * $ref: '#/components/schemas/SwitchAgentSelectionResponse' + */ +const getHandler = async (req: NextRequest, { params }: { params: { threadId: string } }) => { + const userIdentity = getRequestUserIdentity(req); + if (!userIdentity) { + return errorResponse(new Error('Unauthorized'), { status: 401 }, req); + } + + try { + const state = await AgentSelectionService.getThreadAgentState({ threadId: params.threadId, userIdentity }); + return successResponse(state, { status: 200 }, req); + } catch (error) { + return mapAgentSelectionError(error, req); + } +}; + +const patchHandler = 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 unknownKeys = Object.keys(requestBody).filter((key) => key !== 'agentId'); + if (unknownKeys.length > 0) { + return errorResponse( + new Error(`Unsupported switch request fields: ${unknownKeys.join(', ')}`), + { status: 400 }, + req + ); + } + + if (typeof requestBody.agentId !== 'string' || !requestBody.agentId.trim()) { + return errorResponse(new Error('agentId must be a non-empty string'), { status: 400 }, req); + } + + try { + const result = await AgentSelectionService.switchThreadAgent({ + threadId: params.threadId, + userIdentity, + agentId: requestBody.agentId.trim(), + }); + return successResponse(result, { status: 200 }, req); + } catch (error) { + return mapAgentSelectionError(error, req); + } +}; + +export const GET = createApiHandler(getHandler); +export const PATCH = createApiHandler(patchHandler); 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 index 4bc379ee..2c0411d9 100644 --- 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 @@ -84,7 +84,29 @@ describe('GET /api/v2/ai/agent/threads/[threadId]/pending-actions', () => { 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' }], + fileChangePreview: [ + { + id: 'tool-call-1:sample-file.txt', + toolCallId: 'tool-call-1', + sourceTool: 'workspace_edit_file', + path: 'sample-file.txt', + displayPath: 'sample-file.txt', + kind: 'edited', + stage: 'awaiting-approval', + summary: 'Updated sample-file.txt', + additions: 1, + deletions: 1, + truncated: false, + unifiedDiff: null, + beforeTextPreview: null, + afterTextPreview: null, + encoding: null, + oldSizeBytes: null, + newSizeBytes: null, + oldSha256: null, + newSha256: null, + }, + ], riskLabels: ['Workspace write'], }); @@ -107,7 +129,29 @@ describe('GET /api/v2/ai/agent/threads/[threadId]/pending-actions', () => { 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' }], + fileChangePreview: [ + { + id: 'tool-call-1:sample-file.txt', + toolCallId: 'tool-call-1', + sourceTool: 'workspace_edit_file', + path: 'sample-file.txt', + displayPath: 'sample-file.txt', + kind: 'edited', + stage: 'awaiting-approval', + summary: 'Updated sample-file.txt', + additions: 1, + deletions: 1, + truncated: false, + unifiedDiff: null, + beforeTextPreview: null, + afterTextPreview: null, + encoding: null, + oldSizeBytes: null, + newSizeBytes: null, + oldSha256: null, + newSha256: null, + }, + ], riskLabels: ['Workspace write'], }, ]); 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 index d806a3ad..45e58f34 100644 --- 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 @@ -24,13 +24,6 @@ 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: { @@ -53,10 +46,10 @@ jest.mock('server/services/agent/MessageStore', () => ({ }, })); -jest.mock('server/services/agent/ProviderRegistry', () => ({ +jest.mock('server/services/agent/RunPlanResolver', () => ({ __esModule: true, default: { - resolveSelection: jest.fn(), + resolveForRunAdmission: jest.fn(), }, })); @@ -104,9 +97,8 @@ jest.mock('server/services/agentSession', () => ({ 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 AgentRunPlanResolver from 'server/services/agent/RunPlanResolver'; import AgentRunQueueService from 'server/services/agent/RunQueueService'; import AgentRunService from 'server/services/agent/RunService'; import AgentSourceService from 'server/services/agent/SourceService'; @@ -115,9 +107,8 @@ 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 mockResolveForRunAdmission = AgentRunPlanResolver.resolveForRunAdmission as jest.Mock; const mockEnqueueRun = AgentRunQueueService.enqueueRun as jest.Mock; const mockMarkQueuedRunDispatchFailed = AgentRunService.markQueuedRunDispatchFailed as jest.Mock; const mockGetSessionSource = AgentSourceService.getSessionSource as jest.Mock; @@ -125,6 +116,67 @@ const mockGetOwnedThreadWithSession = AgentThreadService.getOwnedThreadWithSessi const mockCanAcceptMessages = AgentSessionService.canAcceptMessages as jest.Mock; const mockTouchActivity = AgentSessionService.touchActivity as jest.Mock; +const customAgentRunPlanSnapshot = { + version: 1, + capturedAt: '2026-05-01T00:00:00.000Z', + agent: { + id: 'custom.sample-agent', + label: 'Sample custom agent', + ownerKind: 'user', + version: 3, + sourceKind: 'freeform_chat', + modelPreference: { + provider: 'anthropic', + model: 'claude-sonnet-4.6', + }, + }, + source: { + id: 'source-1', + adapter: 'blank_workspace', + status: 'ready', + sessionKind: 'chat', + freshness: { + capturedAt: '2026-05-01T00:00:00.000Z', + freshnessSource: 'source', + }, + }, + model: { + requestedProvider: 'anthropic', + requestedModel: 'claude-sonnet-4.6', + resolvedProvider: 'anthropic', + resolvedModel: 'claude-sonnet-4.6', + }, + runtime: { + requestedHarness: null, + resolvedHarness: 'lifecycle_ai_sdk', + sandboxRequirement: { filesystem: 'persistent' }, + runtimeOptions: { maxIterations: 9 }, + approvalPolicy: { + defaultMode: 'require_approval', + rules: { read: 'allow' }, + }, + }, + prompt: { + instructionRefs: [], + instructionAddendum: 'Use the sample custom instructions.', + renderedSummary: 'Sample custom agent description', + renderedHash: 'sha256:sample-custom-agent-prompt', + }, + capabilities: { + provisionalCapabilityIds: ['read_context'], + resolvedCapabilityAccess: [ + { + capabilityId: 'read_context', + availability: 'all_users', + allowed: true, + runtimeCapabilityKey: 'read', + approvalMode: 'allow', + }, + ], + }, + warnings: [], +} as const; + function makeRequest(body: Record): NextRequest { return { json: jest.fn().mockResolvedValue(body), @@ -152,16 +204,71 @@ describe('POST /api/v2/ai/agent/threads/[threadId]/runs', () => { }); mockCanAcceptMessages.mockReturnValue(true); mockGetSessionSource.mockResolvedValue({ + uuid: 'source-1', + adapter: 'blank_workspace', status: 'ready', sandboxRequirements: { filesystem: 'persistent' }, }); - mockResolveSessionContext.mockResolvedValue({ - repoFullName: 'example-org/example-repo', - approvalPolicy: 'on-request', - }); - mockResolveSelection.mockResolvedValue({ - provider: 'openai', - modelId: 'gpt-5.4', + mockResolveForRunAdmission.mockResolvedValue({ + approvalPolicy: { defaultMode: 'require_approval', rules: {} }, + requestedHarness: null, + requestedProvider: null, + requestedModel: null, + resolvedHarness: 'lifecycle_ai_sdk', + resolvedProvider: 'openai', + resolvedModel: 'gpt-5.4', + sandboxRequirement: { filesystem: 'persistent' }, + runtimeOptions: { maxIterations: 12 }, + runPlanSnapshot: { + version: 1, + capturedAt: '2026-05-01T00:00:00.000Z', + agent: { + id: 'system.freeform', + label: 'Free-form', + ownerKind: 'system', + version: 1, + sourceKind: 'freeform_chat', + resourcePolicy: { + sourceKinds: ['build_context_chat', 'workspace_session', 'freeform_chat'], + workspaceRequired: false, + sandboxRequired: false, + }, + modelPreference: null, + }, + source: { + id: 'source-1', + adapter: 'blank_workspace', + status: 'ready', + sessionKind: 'chat', + freshness: { + capturedAt: '2026-05-01T00:00:00.000Z', + freshnessSource: 'source', + }, + }, + model: { + requestedProvider: null, + requestedModel: null, + resolvedProvider: 'openai', + resolvedModel: 'gpt-5.4', + }, + runtime: { + requestedHarness: null, + resolvedHarness: 'lifecycle_ai_sdk', + sandboxRequirement: { filesystem: 'persistent' }, + runtimeOptions: { maxIterations: 12 }, + approvalPolicy: { defaultMode: 'require_approval', rules: {} }, + }, + prompt: { + instructionRefs: [], + renderedSummary: 'Sample prompt summary', + renderedHash: 'sha256:sample-rendered-prompt', + }, + capabilities: { + provisionalCapabilityIds: [], + resolvedCapabilityAccess: [], + }, + warnings: [], + }, }); mockCreateQueuedRunWithMessage.mockResolvedValue({ run: { @@ -181,6 +288,7 @@ describe('POST /api/v2/ai/agent/threads/[threadId]/runs', () => { }); it('rejects run admission when no explicit or session model exists', async () => { + mockResolveForRunAdmission.mockRejectedValueOnce(new Error('Agent run model is required')); mockGetOwnedThreadWithSession.mockResolvedValueOnce({ thread: { id: 7, uuid: 'thread-1' }, session: { @@ -207,6 +315,28 @@ describe('POST /api/v2/ai/agent/threads/[threadId]/runs', () => { expect(mockCreateQueuedRunWithMessage).not.toHaveBeenCalled(); }); + it('rejects run admission policy failures without queueing a run', async () => { + mockResolveForRunAdmission.mockRejectedValueOnce( + new Error('Agent capability "read_context" is unavailable: creator_capability_reserved.') + ); + + 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 capability "read_context" is unavailable: creator_capability_reserved.'); + expect(mockCreateQueuedRunWithMessage).not.toHaveBeenCalled(); + expect(mockEnqueueRun).not.toHaveBeenCalled(); + }); + it('resolves explicit-or-default values before queueing', async () => { const response = await POST( makeRequest({ @@ -220,11 +350,17 @@ describe('POST /api/v2/ai/agent/threads/[threadId]/runs', () => { ); expect(response.status).toBe(201); - expect(mockResolveSelection).toHaveBeenCalledWith({ - repoFullName: 'example-org/example-repo', - requestedProvider: undefined, - requestedModelId: 'gpt-5.4', - }); + expect(mockResolveForRunAdmission).toHaveBeenCalledWith( + expect.objectContaining({ + thread: expect.objectContaining({ id: 7, uuid: 'thread-1' }), + session: expect.objectContaining({ id: 17, uuid: 'session-1' }), + source: expect.objectContaining({ uuid: 'source-1', status: 'ready' }), + userIdentity: { userId: 'sample-user', githubUsername: 'sample-user' }, + requestedProvider: null, + requestedModel: null, + runtimeOptions: { maxIterations: 12 }, + }) + ); expect(mockCreateQueuedRunWithMessage).toHaveBeenCalledWith( expect.objectContaining({ message: { @@ -238,8 +374,15 @@ describe('POST /api/v2/ai/agent/threads/[threadId]/runs', () => { resolvedProvider: 'openai', resolvedModel: 'gpt-5.4', runtimeOptions: { maxIterations: 12 }, + runPlanSnapshot: expect.objectContaining({ + version: 1, + agent: expect.objectContaining({ id: 'system.freeform' }), + }), }) ); + expect(mockResolveForRunAdmission.mock.invocationCallOrder[0]).toBeLessThan( + mockCreateQueuedRunWithMessage.mock.invocationCallOrder[0] + ); expect(mockEnqueueRun).toHaveBeenCalledWith('run-1', 'submit', { githubToken: 'sample-gh-token' }); const body = await response.json(); expect(body.data).toEqual( @@ -255,6 +398,57 @@ describe('POST /api/v2/ai/agent/threads/[threadId]/runs', () => { ); }); + it('passes a custom-agent runPlanSnapshot through queued run admission and response links', async () => { + mockResolveForRunAdmission.mockResolvedValueOnce({ + approvalPolicy: customAgentRunPlanSnapshot.runtime.approvalPolicy, + requestedHarness: null, + requestedProvider: 'anthropic', + requestedModel: 'claude-sonnet-4.6', + resolvedHarness: 'lifecycle_ai_sdk', + resolvedProvider: 'anthropic', + resolvedModel: 'claude-sonnet-4.6', + sandboxRequirement: { filesystem: 'persistent' }, + runtimeOptions: { maxIterations: 9 }, + runPlanSnapshot: customAgentRunPlanSnapshot, + }); + + 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(201); + expect(mockCreateQueuedRunWithMessage).toHaveBeenCalledWith( + expect.objectContaining({ + policy: customAgentRunPlanSnapshot.runtime.approvalPolicy, + requestedProvider: 'anthropic', + requestedModel: 'claude-sonnet-4.6', + resolvedProvider: 'anthropic', + resolvedModel: 'claude-sonnet-4.6', + runtimeOptions: { maxIterations: 9 }, + runPlanSnapshot: customAgentRunPlanSnapshot, + }) + ); + expect(mockEnqueueRun).toHaveBeenCalledWith('run-1', 'submit', { githubToken: 'sample-gh-token' }); + 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({ @@ -366,7 +560,7 @@ describe('POST /api/v2/ai/agent/threads/[threadId]/runs', () => { expect(response.status).toBe(400); expect(body.error.message).toBe('model must contain only provider and id fields'); - expect(mockResolveSelection).not.toHaveBeenCalled(); + expect(mockResolveForRunAdmission).not.toHaveBeenCalled(); expect(mockCreateQueuedRunWithMessage).not.toHaveBeenCalled(); }); @@ -387,6 +581,58 @@ describe('POST /api/v2/ai/agent/threads/[threadId]/runs', () => { expect(mockCreateQueuedRunWithMessage).not.toHaveBeenCalled(); }); + it('rejects public agent selection', async () => { + const response = await POST( + makeRequest({ + message: { + parts: [{ type: 'text', text: 'Hi' }], + }, + agent: { id: 'system.freeform' }, + }), + { params: { threadId: 'thread-1' } } + ); + const body = await response.json(); + + expect(response.status).toBe(400); + expect(body.error.message).toBe('Unsupported run request fields: agent'); + expect(mockCreateQueuedRunWithMessage).not.toHaveBeenCalled(); + }); + + it('rejects public agentId selection', async () => { + const response = await POST( + makeRequest({ + message: { + parts: [{ type: 'text', text: 'Hi' }], + }, + agentId: 'system.freeform', + }), + { params: { threadId: 'thread-1' } } + ); + const body = await response.json(); + + expect(response.status).toBe(400); + expect(body.error.message).toBe('Unsupported run request fields: agentId'); + expect(mockCreateQueuedRunWithMessage).not.toHaveBeenCalled(); + }); + + it('rejects public run plan snapshots', async () => { + const response = await POST( + makeRequest({ + message: { + parts: [{ type: 'text', text: 'Hi' }], + }, + runPlanSnapshot: customAgentRunPlanSnapshot, + }), + { params: { threadId: 'thread-1' } } + ); + const body = await response.json(); + + expect(response.status).toBe(400); + expect(body.error.message).toBe('Unsupported run request fields: runPlanSnapshot'); + expect(mockResolveForRunAdmission).not.toHaveBeenCalled(); + expect(mockCreateQueuedRunWithMessage).not.toHaveBeenCalled(); + }); + it('returns an idempotent response and emits a fresh dispatch signal for a queued run', async () => { mockCreateQueuedRunWithMessage.mockResolvedValueOnce({ run: { 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 index 17469b7b..65658fa5 100644 --- a/src/app/api/v2/ai/agent/threads/[threadId]/runs/route.ts +++ b/src/app/api/v2/ai/agent/threads/[threadId]/runs/route.ts @@ -20,11 +20,10 @@ 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 AgentRunPlanResolver from 'server/services/agent/RunPlanResolver'; import AgentThreadService from 'server/services/agent/ThreadService'; import { normalizeCanonicalAgentMessagePart, @@ -163,8 +162,9 @@ function normalizeRuntimeOptions(value: unknown): AgentRunRuntimeOptions | null * /api/v2/ai/agent/threads/{threadId}/runs: * post: * summary: Create and enqueue a managed run for an agent thread + * description: Creates a run and resolves its run plan server-side from the thread's selected agent, runtime-control choices, requested model, source, and policy. Request body supports only message, model, and runtimeOptions. * tags: - * - Agent Sessions + * - Agent Platform * operationId: createAgentThreadRun * parameters: * - in: path @@ -291,32 +291,16 @@ const postHandler = async (req: NextRequest, { params }: { params: { threadId: s 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; + let runPlan; try { - selection = await AgentProviderRegistry.resolveSelection({ - repoFullName, - requestedProvider: requestedProvider || undefined, - requestedModelId: resolvedModelRequest, + runPlan = await AgentRunPlanResolver.resolveForRunAdmission({ + thread, + session, + source, + userIdentity, + requestedProvider: modelRequest.requestedProvider, + requestedModel: modelRequest.requestedModel, + runtimeOptions, }); } catch (error) { return errorResponse(error instanceof Error ? error : new Error('Invalid agent run model'), { status: 400 }, req); @@ -327,16 +311,17 @@ const postHandler = async (req: NextRequest, { params }: { params: { threadId: s admission = await AgentRunAdmissionService.createQueuedRunWithMessage({ thread, session, - policy: approvalPolicy, + policy: runPlan.approvalPolicy, message, - requestedHarness, - requestedProvider, - requestedModel, - resolvedHarness, - resolvedProvider: selection.provider, - resolvedModel: selection.modelId, - sandboxRequirement: source.sandboxRequirements || {}, - runtimeOptions, + requestedHarness: runPlan.requestedHarness, + requestedProvider: runPlan.requestedProvider, + requestedModel: runPlan.requestedModel, + resolvedHarness: runPlan.resolvedHarness, + resolvedProvider: runPlan.resolvedProvider, + resolvedModel: runPlan.resolvedModel, + sandboxRequirement: runPlan.sandboxRequirement, + runtimeOptions: runPlan.runtimeOptions, + runPlanSnapshot: runPlan.runPlanSnapshot, }); } catch (error) { if (AgentRunService.isActiveRunConflictError(error)) { diff --git a/src/app/api/v2/ai/agent/threads/[threadId]/runtime-controls/route.test.ts b/src/app/api/v2/ai/agent/threads/[threadId]/runtime-controls/route.test.ts new file mode 100644 index 00000000..fee4f109 --- /dev/null +++ b/src/app/api/v2/ai/agent/threads/[threadId]/runtime-controls/route.test.ts @@ -0,0 +1,197 @@ +/** + * 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 mockGetRequestUserIdentity = jest.fn(); +const mockGetState = jest.fn(); +const mockPatchChoices = jest.fn(); + +jest.mock('server/lib/get-user', () => ({ + getRequestUserIdentity: (...args: unknown[]) => mockGetRequestUserIdentity(...args), +})); + +jest.mock('server/services/agent/ThreadRuntimeControlsService', () => { + class AgentThreadRuntimeControlsError extends Error { + constructor(public readonly code: string, message: string) { + super(message); + this.name = 'AgentThreadRuntimeControlsError'; + } + } + + return { + __esModule: true, + default: { + getState: (...args: unknown[]) => mockGetState(...args), + patchChoices: (...args: unknown[]) => mockPatchChoices(...args), + }, + AgentThreadRuntimeControlsError, + }; +}); + +import { GET, PATCH } from './route'; +import { AgentThreadRuntimeControlsError } from 'server/services/agent/ThreadRuntimeControlsService'; + +const runtimeControlsState = { + tools: { + required: [ + { + id: 'rtc_required', + label: 'Read/context', + description: 'Read safe context.', + required: true, + selected: true, + available: true, + }, + ], + optional: [ + { + id: 'rtc_optional', + label: 'Workspace files', + description: 'Work with files.', + required: false, + selected: true, + available: true, + }, + ], + selectedChoiceIds: ['rtc_required', 'rtc_optional'], + }, + mcp: { + connections: [ + { + id: 'rtc_mcp', + label: 'Sample MCP', + description: 'Provides sample context.', + required: false, + selected: true, + available: true, + }, + ], + selectedChoiceIds: ['rtc_mcp'], + }, + canEdit: true, + disabledReason: null, +}; + +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/runtime-controls'), + } as unknown as NextRequest; +} + +function makeInvalidJsonRequest(): NextRequest { + return { + json: jest.fn().mockRejectedValue(new SyntaxError('Unexpected token')), + headers: new Headers([['x-request-id', 'req-test']]), + nextUrl: new URL('http://localhost/api/v2/ai/agent/threads/thread-1/runtime-controls'), + } as unknown as NextRequest; +} + +describe('/api/v2/ai/agent/threads/[threadId]/runtime-controls', () => { + beforeEach(() => { + jest.clearAllMocks(); + mockGetRequestUserIdentity.mockReturnValue({ + userId: 'sample-user', + githubUsername: 'sample-user', + }); + mockGetState.mockResolvedValue(runtimeControlsState); + mockPatchChoices.mockResolvedValue(runtimeControlsState); + }); + + it('GET returns sanitized runtime-control state', async () => { + const response = await GET(makeRequest(), { params: { threadId: 'thread-1' } }); + const body = await response.json(); + + expect(response.status).toBe(200); + expect(mockGetState).toHaveBeenCalledWith({ + threadId: 'thread-1', + userIdentity: { userId: 'sample-user', githubUsername: 'sample-user' }, + }); + expect(body.data).toEqual(runtimeControlsState); + expect(JSON.stringify(body.data)).not.toContain('workspace_files'); + expect(JSON.stringify(body.data)).not.toContain('sample-mcp'); + }); + + it('PATCH accepts choice arrays and returns updated state', async () => { + const response = await PATCH( + makeRequest({ + toolChoiceIds: ['rtc_optional'], + mcpChoiceIds: ['rtc_mcp'], + }), + { params: { threadId: 'thread-1' } } + ); + + expect(response.status).toBe(200); + expect(mockPatchChoices).toHaveBeenCalledWith({ + threadId: 'thread-1', + userIdentity: { userId: 'sample-user', githubUsername: 'sample-user' }, + toolChoiceIds: ['rtc_optional'], + mcpChoiceIds: ['rtc_mcp'], + }); + }); + + it('returns 400 for malformed bodies and unknown choices', async () => { + const malformed = await PATCH(makeRequest({ toolChoiceIds: 'workspace_files' }), { + params: { threadId: 'thread-1' }, + }); + expect(malformed.status).toBe(400); + + const invalidJson = await PATCH(makeInvalidJsonRequest(), { + params: { threadId: 'thread-1' }, + }); + expect(invalidJson.status).toBe(400); + + mockPatchChoices.mockRejectedValueOnce( + new AgentThreadRuntimeControlsError('unknown_choice', 'Unknown runtime control choice.') + ); + const unknown = await PATCH(makeRequest({ toolChoiceIds: ['workspace_files'], mcpChoiceIds: [] }), { + params: { threadId: 'thread-1' }, + }); + expect(unknown.status).toBe(400); + }); + + it('maps policy, ownership, and active-run service errors', async () => { + mockPatchChoices.mockRejectedValueOnce( + new AgentThreadRuntimeControlsError('policy_denied', 'Runtime control choice is unavailable.') + ); + const denied = await PATCH(makeRequest({ toolChoiceIds: ['rtc_optional'], mcpChoiceIds: [] }), { + params: { threadId: 'thread-1' }, + }); + expect(denied.status).toBe(403); + + mockGetState.mockRejectedValueOnce(new AgentThreadRuntimeControlsError('not_found', 'Agent thread not found')); + const missing = await GET(makeRequest(), { params: { threadId: 'thread-1' } }); + expect(missing.status).toBe(404); + + mockPatchChoices.mockRejectedValueOnce( + new AgentThreadRuntimeControlsError('active_run', 'Change after this response finishes.') + ); + const active = await PATCH(makeRequest({ toolChoiceIds: [], mcpChoiceIds: [] }), { + params: { threadId: 'thread-1' }, + }); + expect(active.status).toBe(409); + }); + + it('returns 401 without identity', async () => { + mockGetRequestUserIdentity.mockReturnValueOnce(null); + + const response = await GET(makeRequest(), { params: { threadId: 'thread-1' } }); + + expect(response.status).toBe(401); + }); +}); diff --git a/src/app/api/v2/ai/agent/threads/[threadId]/runtime-controls/route.ts b/src/app/api/v2/ai/agent/threads/[threadId]/runtime-controls/route.ts new file mode 100644 index 00000000..59828c75 --- /dev/null +++ b/src/app/api/v2/ai/agent/threads/[threadId]/runtime-controls/route.ts @@ -0,0 +1,232 @@ +/** + * 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 AgentThreadRuntimeControlsService, { + AgentThreadRuntimeControlsError, +} from 'server/services/agent/ThreadRuntimeControlsService'; + +type RuntimeControlsPatchBody = { + toolChoiceIds?: string[]; + mcpChoiceIds?: string[]; +}; + +function isPlainObject(value: unknown): value is Record { + return value != null && typeof value === 'object' && !Array.isArray(value); +} + +function readChoiceIds(value: unknown, fieldName: string): string[] | undefined | Error { + if (value === undefined) { + return undefined; + } + + if (!Array.isArray(value)) { + return new Error(`${fieldName} must be an array of choice ids.`); + } + + for (const item of value) { + if (typeof item !== 'string' || !item.trim()) { + return new Error(`${fieldName} must contain only choice ids.`); + } + } + + return value.map((item) => item.trim()); +} + +function parsePatchBody(body: unknown): RuntimeControlsPatchBody | Error { + if (!isPlainObject(body)) { + return new Error('Request body must be an object.'); + } + + const unknownKeys = Object.keys(body).filter((key) => key !== 'toolChoiceIds' && key !== 'mcpChoiceIds'); + if (unknownKeys.length > 0) { + return new Error(`Unsupported runtime-control fields: ${unknownKeys.join(', ')}.`); + } + + const toolChoiceIds = readChoiceIds(body.toolChoiceIds, 'toolChoiceIds'); + if (toolChoiceIds instanceof Error) { + return toolChoiceIds; + } + + const mcpChoiceIds = readChoiceIds(body.mcpChoiceIds, 'mcpChoiceIds'); + if (mcpChoiceIds instanceof Error) { + return mcpChoiceIds; + } + + return { toolChoiceIds, mcpChoiceIds }; +} + +function mapRuntimeControlsError(error: unknown, req: NextRequest) { + if (error instanceof AgentThreadRuntimeControlsError) { + const statusByCode: Record = { + invalid_input: 400, + unknown_choice: 400, + policy_denied: 403, + not_found: 404, + active_run: 409, + }; + return errorResponse(error, { status: statusByCode[error.code] }, req); + } + + throw error; +} + +/** + * @openapi + * /api/v2/ai/agent/threads/{threadId}/runtime-controls: + * get: + * summary: Get runtime control choices for a thread + * description: Returns available and selected runtime control choices that apply to future runs in the thread. + * tags: + * - Agent Platform + * operationId: getAgentThreadRuntimeControls + * parameters: + * - in: path + * name: threadId + * required: true + * schema: + * type: string + * responses: + * '200': + * description: Thread runtime controls + * content: + * application/json: + * schema: + * allOf: + * - $ref: '#/components/schemas/SuccessApiResponse' + * - type: object + * required: [data] + * properties: + * data: + * $ref: '#/components/schemas/AgentThreadRuntimeControlsState' + * '401': + * description: Unauthorized + * content: + * application/json: + * schema: + * $ref: '#/components/schemas/ApiErrorResponse' + * '404': + * description: Thread not found + * content: + * application/json: + * schema: + * $ref: '#/components/schemas/ApiErrorResponse' + * patch: + * summary: Update runtime control choices for future runs in a thread + * description: Updates selected runtime control choices for future runs without changing active or completed runs. + * tags: + * - Agent Platform + * operationId: patchAgentThreadRuntimeControls + * parameters: + * - in: path + * name: threadId + * required: true + * schema: + * type: string + * requestBody: + * required: true + * content: + * application/json: + * schema: + * $ref: '#/components/schemas/AgentThreadRuntimeControlsPatchRequest' + * responses: + * '200': + * description: Updated thread runtime controls + * content: + * application/json: + * schema: + * allOf: + * - $ref: '#/components/schemas/SuccessApiResponse' + * - type: object + * required: [data] + * properties: + * data: + * $ref: '#/components/schemas/AgentThreadRuntimeControlsState' + * '400': + * description: Invalid runtime-control choices + * content: + * application/json: + * schema: + * $ref: '#/components/schemas/ApiErrorResponse' + * '403': + * description: Runtime-control choice is unavailable + * content: + * application/json: + * schema: + * $ref: '#/components/schemas/ApiErrorResponse' + * '404': + * description: Thread not found + * content: + * application/json: + * schema: + * $ref: '#/components/schemas/ApiErrorResponse' + * '409': + * description: Runtime controls cannot change while a run is active + * content: + * application/json: + * schema: + * $ref: '#/components/schemas/ApiErrorResponse' + */ +const getHandler = async (req: NextRequest, { params }: { params: { threadId: string } }) => { + const userIdentity = getRequestUserIdentity(req); + if (!userIdentity) { + return errorResponse(new Error('Unauthorized'), { status: 401 }, req); + } + + try { + const state = await AgentThreadRuntimeControlsService.getState({ threadId: params.threadId, userIdentity }); + return successResponse(state, { status: 200 }, req); + } catch (error) { + return mapRuntimeControlsError(error, req); + } +}; + +const patchHandler = async (req: NextRequest, { params }: { params: { threadId: string } }) => { + const userIdentity = getRequestUserIdentity(req); + if (!userIdentity) { + return errorResponse(new Error('Unauthorized'), { status: 401 }, req); + } + + let body: unknown; + try { + body = await req.json(); + } catch { + return errorResponse(new Error('Request body must be valid JSON.'), { status: 400 }, req); + } + + const parsedBody = parsePatchBody(body); + if (parsedBody instanceof Error) { + return errorResponse(parsedBody, { status: 400 }, req); + } + + try { + const state = await AgentThreadRuntimeControlsService.patchChoices({ + threadId: params.threadId, + userIdentity, + ...parsedBody, + }); + return successResponse(state, { status: 200 }, req); + } catch (error) { + return mapRuntimeControlsError(error, req); + } +}; + +export const GET = createApiHandler(getHandler); +export const PATCH = createApiHandler(patchHandler); diff --git a/src/app/api/v2/ai/agent/threads/[threadId]/usage/route.test.ts b/src/app/api/v2/ai/agent/threads/[threadId]/usage/route.test.ts new file mode 100644 index 00000000..f148b657 --- /dev/null +++ b/src/app/api/v2/ai/agent/threads/[threadId]/usage/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/AgentUsageService', () => ({ + __esModule: true, + default: { + getOwnedThreadUsage: jest.fn(), + }, +})); + +import { GET } from './route'; +import { getRequestUserIdentity } from 'server/lib/get-user'; +import AgentUsageService from 'server/services/agent/AgentUsageService'; + +const mockGetRequestUserIdentity = getRequestUserIdentity as jest.Mock; +const mockGetOwnedThreadUsage = AgentUsageService.getOwnedThreadUsage 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/usage'), + } as unknown as NextRequest; +} + +describe('GET /api/v2/ai/agent/threads/[threadId]/usage', () => { + beforeEach(() => { + jest.clearAllMocks(); + mockGetRequestUserIdentity.mockReturnValue({ + userId: 'sample-user', + githubUsername: 'sample-user', + }); + mockGetOwnedThreadUsage.mockResolvedValue({ + threadId: 'thread-1', + sessionId: 'session-1', + usageSummary: { + totalTokens: 42, + inputTokens: 30, + outputTokens: 12, + }, + usageByModel: [ + { + provider: 'openai', + model: 'gpt-5.4', + totalTokens: 42, + inputTokens: 30, + outputTokens: 12, + runCount: 1, + reportedRunCount: 1, + missingUsageRunCount: 0, + }, + ], + usageCompleteness: { + runCount: 1, + reportedRunCount: 1, + missingUsageRunCount: 0, + complete: true, + }, + }); + }); + + it('returns thread usage for the authenticated owner', async () => { + const response = await GET(makeRequest(), { params: { threadId: 'thread-1' } }); + const body = await response.json(); + + expect(response.status).toBe(200); + expect(mockGetOwnedThreadUsage).toHaveBeenCalledWith('thread-1', 'sample-user'); + expect(body.data).toEqual( + expect.objectContaining({ + threadId: 'thread-1', + sessionId: 'session-1', + usageSummary: { + totalTokens: 42, + inputTokens: 30, + outputTokens: 12, + }, + }) + ); + }); + + it('returns 401 without a request identity', async () => { + mockGetRequestUserIdentity.mockReturnValueOnce(null); + + const response = await GET(makeRequest(), { params: { threadId: 'thread-1' } }); + + expect(response.status).toBe(401); + expect(mockGetOwnedThreadUsage).not.toHaveBeenCalled(); + }); + + it('maps missing thread or session ownership to 404', async () => { + mockGetOwnedThreadUsage.mockRejectedValueOnce(new Error('Agent thread not found')); + + const response = await GET(makeRequest(), { params: { threadId: 'missing-thread' } }); + const body = await response.json(); + + 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]/usage/route.ts b/src/app/api/v2/ai/agent/threads/[threadId]/usage/route.ts new file mode 100644 index 00000000..0bffead5 --- /dev/null +++ b/src/app/api/v2/ai/agent/threads/[threadId]/usage/route.ts @@ -0,0 +1,80 @@ +/** + * 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 AgentUsageService from 'server/services/agent/AgentUsageService'; + +/** + * @openapi + * /api/v2/ai/agent/threads/{threadId}/usage: + * get: + * summary: Get exact token usage for an agent thread + * description: Returns provider-reported token usage aggregated from agent runs in the thread, grouped by the resolved provider and model used for each run. + * tags: + * - Agent Platform + * operationId: getAgentThreadUsage + * parameters: + * - in: path + * name: threadId + * required: true + * schema: + * type: string + * responses: + * '200': + * description: Exact provider-reported usage for the thread + * content: + * application/json: + * schema: + * $ref: '#/components/schemas/GetAgentThreadUsageSuccessResponse' + * '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' + */ +const getHandler = async (req: NextRequest, { params }: { params: { threadId: string } }) => { + const userIdentity = getRequestUserIdentity(req); + if (!userIdentity) { + return errorResponse(new Error('Unauthorized'), { status: 401 }, req); + } + + try { + const usage = await AgentUsageService.getOwnedThreadUsage(params.threadId, userIdentity.userId); + return successResponse(usage, { status: 200 }, req); + } 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; + } +}; + +export const GET = createApiHandler(getHandler); diff --git a/src/app/api/v2/ai/chat/[buildUuid]/feedback/route.test.ts b/src/app/api/v2/ai/chat/[buildUuid]/feedback/route.test.ts deleted file mode 100644 index 351b62c1..00000000 --- a/src/app/api/v2/ai/chat/[buildUuid]/feedback/route.test.ts +++ /dev/null @@ -1,141 +0,0 @@ -/** - * Copyright 2025 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/dependencies', () => ({ - defaultDb: {}, - defaultRedis: {}, -})); - -jest.mock('server/services/ai/conversation/storage', () => ({ - __esModule: true, - default: jest.fn(), -})); - -jest.mock('server/services/ai/conversation/persistence', () => ({ - __esModule: true, - default: jest.fn().mockImplementation(() => ({})), -})); - -jest.mock('server/services/ai/feedback/FeedbackService', () => ({ - __esModule: true, - default: jest.fn(), -})); - -jest.mock('server/services/ai/feedback/resolveFeedbackContext', () => ({ - resolveFeedbackContext: jest.fn(), -})); - -import { POST } from './route'; -import { MAX_FEEDBACK_TEXT_LENGTH } from 'server/services/ai/feedback/constants'; -import FeedbackService from 'server/services/ai/feedback/FeedbackService'; -import { resolveFeedbackContext } from 'server/services/ai/feedback/resolveFeedbackContext'; - -const mockSubmitConversationFeedback = jest.fn(); -const mockResolveFeedbackContext = resolveFeedbackContext as jest.Mock; -const MockFeedbackService = FeedbackService as unknown as jest.Mock; - -function makeRequest(body: unknown, userClaims?: Record): NextRequest { - const headers = new Headers([['x-request-id', 'req-test']]); - if (userClaims) { - headers.set('x-user', Buffer.from(JSON.stringify(userClaims), 'utf8').toString('base64url')); - } - return { - headers, - json: jest.fn().mockResolvedValue(body), - } as unknown as NextRequest; -} - -describe('POST /api/v2/ai/chat/[buildUuid]/feedback', () => { - beforeEach(() => { - jest.clearAllMocks(); - MockFeedbackService.mockImplementation(() => ({ - submitConversationFeedback: mockSubmitConversationFeedback, - })); - mockResolveFeedbackContext.mockResolvedValue({ repo: 'org/repo', prNumber: 123 }); - mockSubmitConversationFeedback.mockResolvedValue({ id: 1 }); - }); - - it('returns 400 when text is not a string', async () => { - const response = await POST(makeRequest({ rating: 'up', text: 123 }), { - params: { buildUuid: 'uuid-1' }, - }); - - expect(response.status).toBe(400); - await expect(response.json()).resolves.toMatchObject({ - error: { message: 'Invalid text: must be a string' }, - }); - expect(mockSubmitConversationFeedback).not.toHaveBeenCalled(); - }); - - it('returns 400 when text exceeds max length', async () => { - const response = await POST(makeRequest({ rating: 'up', text: 'a'.repeat(MAX_FEEDBACK_TEXT_LENGTH + 1) }), { - params: { buildUuid: 'uuid-1' }, - }); - - expect(response.status).toBe(400); - await expect(response.json()).resolves.toMatchObject({ - error: { - message: `Invalid text: exceeds max length of ${MAX_FEEDBACK_TEXT_LENGTH} characters`, - }, - }); - expect(mockSubmitConversationFeedback).not.toHaveBeenCalled(); - }); - - it('returns 404 when repository context cannot be resolved', async () => { - mockResolveFeedbackContext.mockResolvedValue({ repo: '' }); - - const response = await POST(makeRequest({ rating: 'up' }), { - params: { buildUuid: 'uuid-2' }, - }); - - expect(response.status).toBe(404); - expect(mockSubmitConversationFeedback).not.toHaveBeenCalled(); - }); - - it('submits feedback successfully', async () => { - const response = await POST(makeRequest({ rating: 'down', text: 'Needs work' }), { - params: { buildUuid: 'uuid-3' }, - }); - - expect(response.status).toBe(201); - expect(mockSubmitConversationFeedback).toHaveBeenCalledWith({ - buildUuid: 'uuid-3', - rating: 'down', - text: 'Needs work', - userIdentifier: undefined, - repo: 'org/repo', - prNumber: 123, - }); - }); - - it('derives userIdentifier from auth claims when present', async () => { - const response = await POST(makeRequest({ rating: 'up' }, { github_username: 'sample-user' }), { - params: { buildUuid: 'uuid-4' }, - }); - - expect(response.status).toBe(201); - expect(mockSubmitConversationFeedback).toHaveBeenCalledWith({ - buildUuid: 'uuid-4', - rating: 'up', - text: undefined, - userIdentifier: 'sample-user', - repo: 'org/repo', - prNumber: 123, - }); - }); -}); diff --git a/src/app/api/v2/ai/chat/[buildUuid]/feedback/route.ts b/src/app/api/v2/ai/chat/[buildUuid]/feedback/route.ts deleted file mode 100644 index b8ddecc0..00000000 --- a/src/app/api/v2/ai/chat/[buildUuid]/feedback/route.ts +++ /dev/null @@ -1,129 +0,0 @@ -/** - * Copyright 2025 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 { createApiHandler } from 'server/lib/createApiHandler'; -import { errorResponse, successResponse } from 'server/lib/response'; -import { defaultDb, defaultRedis } from 'server/lib/dependencies'; -import { getUser } from 'server/lib/get-user'; -import AIAgentConversationService from 'server/services/ai/conversation/storage'; -import ConversationPersistenceService from 'server/services/ai/conversation/persistence'; -import FeedbackService from 'server/services/ai/feedback/FeedbackService'; -import { MAX_FEEDBACK_TEXT_LENGTH } from 'server/services/ai/feedback/constants'; -import { resolveFeedbackContext } from 'server/services/ai/feedback/resolveFeedbackContext'; -import { resolveUserIdentifierFromPayload } from 'server/services/ai/feedback/userIdentifier'; - -/** - * @openapi - * /api/v2/ai/chat/{buildUuid}/feedback: - * post: - * summary: Submit conversation feedback - * description: > - * Submit a thumbs up/down rating for an entire conversation session. - * Triggers conversation persistence from Redis to Postgres if not already persisted. - * tags: - * - AI Chat - * operationId: submitConversationFeedback - * parameters: - * - in: path - * name: buildUuid - * required: true - * schema: - * type: string - * description: The UUID of the build whose conversation to rate. - * requestBody: - * required: true - * content: - * application/json: - * schema: - * type: object - * required: - * - rating - * properties: - * rating: - * type: string - * enum: [up, down] - * text: - * type: string - * responses: - * '201': - * description: Feedback created successfully - * '400': - * description: Invalid or missing rating - * '404': - * description: Conversation not found - * '500': - * description: Server error - */ -const postHandler = async (req: NextRequest, { params }: { params: { buildUuid: string } }) => { - const { buildUuid } = params; - - if (!buildUuid) { - return errorResponse(new Error('Missing required parameter: buildUuid'), { status: 400 }, req); - } - - const body = await req.json(); - const { rating, text } = body; - - if (!rating || (rating !== 'up' && rating !== 'down')) { - return errorResponse(new Error('Invalid rating: must be "up" or "down"'), { status: 400 }, req); - } - - if (text !== undefined && typeof text !== 'string') { - return errorResponse(new Error('Invalid text: must be a string'), { status: 400 }, req); - } - - if (typeof text === 'string' && Array.from(text).length > MAX_FEEDBACK_TEXT_LENGTH) { - return errorResponse( - new Error(`Invalid text: exceeds max length of ${MAX_FEEDBACK_TEXT_LENGTH} characters`), - { status: 400 }, - req - ); - } - - const conversationService = new AIAgentConversationService(defaultDb, defaultRedis); - const persistenceService = new ConversationPersistenceService(conversationService); - const feedbackService = new FeedbackService(persistenceService); - const userIdentifier = resolveUserIdentifierFromPayload(getUser(req)); - - const { repo, prNumber } = await resolveFeedbackContext(buildUuid); - if (!repo) { - return errorResponse( - new Error(`Unable to resolve repository context for buildUuid=${buildUuid}`), - { status: 404 }, - req - ); - } - - try { - const record = await feedbackService.submitConversationFeedback({ - buildUuid, - rating, - text, - userIdentifier, - repo, - prNumber, - }); - return successResponse(record, { status: 201 }, req); - } catch (error) { - if (error instanceof Error && error.message.includes('not found')) { - return errorResponse(error, { status: 404 }, req); - } - throw error; - } -}; - -export const POST = createApiHandler(postHandler); diff --git a/src/app/api/v2/ai/chat/[buildUuid]/messages/[messageId]/feedback/route.test.ts b/src/app/api/v2/ai/chat/[buildUuid]/messages/[messageId]/feedback/route.test.ts deleted file mode 100644 index 56ec6441..00000000 --- a/src/app/api/v2/ai/chat/[buildUuid]/messages/[messageId]/feedback/route.test.ts +++ /dev/null @@ -1,137 +0,0 @@ -/** - * Copyright 2025 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/dependencies', () => ({ - defaultDb: {}, - defaultRedis: {}, -})); - -jest.mock('server/services/ai/conversation/storage', () => ({ - __esModule: true, - default: jest.fn(), -})); - -jest.mock('server/services/ai/conversation/persistence', () => ({ - __esModule: true, - default: jest.fn().mockImplementation(() => ({})), -})); - -jest.mock('server/services/ai/feedback/FeedbackService', () => ({ - __esModule: true, - default: jest.fn(), -})); - -jest.mock('server/services/ai/feedback/resolveFeedbackContext', () => ({ - resolveFeedbackContext: jest.fn(), -})); - -import { POST } from './route'; -import { MAX_FEEDBACK_TEXT_LENGTH } from 'server/services/ai/feedback/constants'; -import FeedbackService from 'server/services/ai/feedback/FeedbackService'; -import { resolveFeedbackContext } from 'server/services/ai/feedback/resolveFeedbackContext'; - -const mockSubmitMessageFeedback = jest.fn(); -const mockResolveFeedbackContext = resolveFeedbackContext as jest.Mock; -const MockFeedbackService = FeedbackService as unknown as jest.Mock; - -function makeRequest(body: unknown, userClaims?: Record): NextRequest { - const headers = new Headers([['x-request-id', 'req-test']]); - if (userClaims) { - headers.set('x-user', Buffer.from(JSON.stringify(userClaims), 'utf8').toString('base64url')); - } - return { - headers, - json: jest.fn().mockResolvedValue(body), - } as unknown as NextRequest; -} - -describe('POST /api/v2/ai/chat/[buildUuid]/messages/[messageId]/feedback', () => { - beforeEach(() => { - jest.clearAllMocks(); - MockFeedbackService.mockImplementation(() => ({ - submitMessageFeedback: mockSubmitMessageFeedback, - })); - mockResolveFeedbackContext.mockResolvedValue({ repo: 'org/repo', prNumber: 123 }); - mockSubmitMessageFeedback.mockResolvedValue({ id: 1 }); - }); - - it('returns 400 for invalid messageId', async () => { - const response = await POST(makeRequest({ rating: 'up' }), { - params: { buildUuid: 'uuid-1', messageId: '-1' }, - }); - - expect(response.status).toBe(400); - await expect(response.json()).resolves.toMatchObject({ - error: { message: 'Invalid messageId: must be a non-negative integer' }, - }); - expect(mockSubmitMessageFeedback).not.toHaveBeenCalled(); - }); - - it('returns 400 when text exceeds max length', async () => { - const response = await POST(makeRequest({ rating: 'up', text: 'a'.repeat(MAX_FEEDBACK_TEXT_LENGTH + 1) }), { - params: { buildUuid: 'uuid-1', messageId: '1' }, - }); - - expect(response.status).toBe(400); - await expect(response.json()).resolves.toMatchObject({ - error: { - message: `Invalid text: exceeds max length of ${MAX_FEEDBACK_TEXT_LENGTH} characters`, - }, - }); - expect(mockSubmitMessageFeedback).not.toHaveBeenCalled(); - }); - - it('accepts messageTimestamp when messageId path param is 0', async () => { - const response = await POST(makeRequest({ rating: 'down', messageTimestamp: 1700000000000 }), { - params: { buildUuid: 'uuid-2', messageId: '0' }, - }); - - expect(response.status).toBe(201); - expect(mockSubmitMessageFeedback).toHaveBeenCalledWith({ - buildUuid: 'uuid-2', - messageId: undefined, - messageTimestamp: 1700000000000, - rating: 'down', - text: undefined, - userIdentifier: undefined, - repo: 'org/repo', - prNumber: 123, - }); - }); - - it('derives userIdentifier from auth claims when present', async () => { - const response = await POST( - makeRequest({ rating: 'up', messageTimestamp: 1700000000000 }, { preferred_username: 'sample-user' }), - { - params: { buildUuid: 'uuid-3', messageId: '0' }, - } - ); - - expect(response.status).toBe(201); - expect(mockSubmitMessageFeedback).toHaveBeenCalledWith({ - buildUuid: 'uuid-3', - messageId: undefined, - messageTimestamp: 1700000000000, - rating: 'up', - text: undefined, - userIdentifier: 'sample-user', - repo: 'org/repo', - prNumber: 123, - }); - }); -}); diff --git a/src/app/api/v2/ai/chat/[buildUuid]/messages/[messageId]/feedback/route.ts b/src/app/api/v2/ai/chat/[buildUuid]/messages/[messageId]/feedback/route.ts deleted file mode 100644 index e7b23dcd..00000000 --- a/src/app/api/v2/ai/chat/[buildUuid]/messages/[messageId]/feedback/route.ts +++ /dev/null @@ -1,172 +0,0 @@ -/** - * Copyright 2025 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 { createApiHandler } from 'server/lib/createApiHandler'; -import { errorResponse, successResponse } from 'server/lib/response'; -import { defaultDb, defaultRedis } from 'server/lib/dependencies'; -import { getUser } from 'server/lib/get-user'; -import AIAgentConversationService from 'server/services/ai/conversation/storage'; -import ConversationPersistenceService from 'server/services/ai/conversation/persistence'; -import FeedbackService from 'server/services/ai/feedback/FeedbackService'; -import { MAX_FEEDBACK_TEXT_LENGTH } from 'server/services/ai/feedback/constants'; -import { resolveFeedbackContext } from 'server/services/ai/feedback/resolveFeedbackContext'; -import { resolveUserIdentifierFromPayload } from 'server/services/ai/feedback/userIdentifier'; - -/** - * @openapi - * /api/v2/ai/chat/{buildUuid}/messages/{messageId}/feedback: - * post: - * summary: Submit message feedback - * description: > - * Submit a thumbs up/down rating for a specific agent message. - * Triggers conversation persistence from Redis to Postgres if not already persisted. - * tags: - * - AI Chat - * operationId: submitMessageFeedback - * parameters: - * - in: path - * name: buildUuid - * required: true - * schema: - * type: string - * description: The UUID of the build. - * - in: path - * name: messageId - * required: true - * schema: - * type: integer - * description: The ID of the message to rate. - * requestBody: - * required: true - * content: - * application/json: - * schema: - * type: object - * required: - * - rating - * properties: - * rating: - * type: string - * enum: [up, down] - * text: - * type: string - * messageTimestamp: - * type: integer - * description: > - * Optional assistant message timestamp. Use this when a persisted - * message ID is not yet available. - * responses: - * '201': - * description: Feedback created successfully - * '400': - * description: Invalid or missing parameters - * '404': - * description: Conversation or message not found - * '500': - * description: Server error - */ -const postHandler = async (req: NextRequest, { params }: { params: { buildUuid: string; messageId: string } }) => { - const { buildUuid, messageId } = params; - - if (!buildUuid) { - return errorResponse(new Error('Missing required parameter: buildUuid'), { status: 400 }, req); - } - - const parsedMessageId = parseInt(messageId, 10); - if (isNaN(parsedMessageId) || parsedMessageId < 0) { - return errorResponse(new Error('Invalid messageId: must be a non-negative integer'), { status: 400 }, req); - } - - const body = await req.json(); - const { rating, text, messageTimestamp } = body; - - if (!rating || (rating !== 'up' && rating !== 'down')) { - return errorResponse(new Error('Invalid rating: must be "up" or "down"'), { status: 400 }, req); - } - - if (text !== undefined && typeof text !== 'string') { - return errorResponse(new Error('Invalid text: must be a string'), { status: 400 }, req); - } - - if (typeof text === 'string' && Array.from(text).length > MAX_FEEDBACK_TEXT_LENGTH) { - return errorResponse( - new Error(`Invalid text: exceeds max length of ${MAX_FEEDBACK_TEXT_LENGTH} characters`), - { status: 400 }, - req - ); - } - - let parsedMessageTimestamp: number | undefined; - if (messageTimestamp !== undefined) { - if (typeof messageTimestamp === 'number' && Number.isFinite(messageTimestamp)) { - parsedMessageTimestamp = messageTimestamp; - } else if (typeof messageTimestamp === 'string') { - const value = Number(messageTimestamp); - if (Number.isFinite(value)) { - parsedMessageTimestamp = value; - } - } - - if (parsedMessageTimestamp === undefined) { - return errorResponse(new Error('Invalid messageTimestamp: must be a number'), { status: 400 }, req); - } - } - - const resolvedMessageId = parsedMessageId > 0 ? parsedMessageId : undefined; - if (!resolvedMessageId && parsedMessageTimestamp === undefined) { - return errorResponse( - new Error('Either a valid messageId path param or messageTimestamp is required'), - { status: 400 }, - req - ); - } - - const conversationService = new AIAgentConversationService(defaultDb, defaultRedis); - const persistenceService = new ConversationPersistenceService(conversationService); - const feedbackService = new FeedbackService(persistenceService); - const userIdentifier = resolveUserIdentifierFromPayload(getUser(req)); - - const { repo, prNumber } = await resolveFeedbackContext(buildUuid); - if (!repo) { - return errorResponse( - new Error(`Unable to resolve repository context for buildUuid=${buildUuid}`), - { status: 404 }, - req - ); - } - - try { - const record = await feedbackService.submitMessageFeedback({ - buildUuid, - messageId: resolvedMessageId, - messageTimestamp: parsedMessageTimestamp, - rating, - text, - userIdentifier, - repo, - prNumber, - }); - return successResponse(record, { status: 201 }, req); - } catch (error) { - if (error instanceof Error && error.message.includes('not found')) { - return errorResponse(error, { status: 404 }, req); - } - throw error; - } -}; - -export const POST = createApiHandler(postHandler); diff --git a/src/app/api/v2/ai/chat/[buildUuid]/messages/route.test.ts b/src/app/api/v2/ai/chat/[buildUuid]/messages/route.test.ts deleted file mode 100644 index f492bc68..00000000 --- a/src/app/api/v2/ai/chat/[buildUuid]/messages/route.test.ts +++ /dev/null @@ -1,96 +0,0 @@ -/** - * Copyright 2025 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 mockGetConversation = jest.fn(); - -jest.mock('server/lib/dependencies', () => ({ - defaultDb: {}, - defaultRedis: {}, -})); - -jest.mock('server/services/ai/conversation/storage', () => ({ - __esModule: true, - default: jest.fn().mockImplementation(() => ({ - getConversation: mockGetConversation, - })), -})); - -jest.mock('server/models/ConversationMessage', () => { - const model: { query: jest.Mock } = { query: jest.fn() }; - return { __esModule: true, default: model }; -}); - -import { GET } from './route'; -import ConversationMessage from 'server/models/ConversationMessage'; - -const MockConversationMessage = ConversationMessage as unknown as { - query: jest.Mock; -}; - -function makeRequest(): NextRequest { - return { - headers: new Headers([['x-request-id', 'req-test']]), - } as unknown as NextRequest; -} - -describe('GET /api/v2/ai/chat/[buildUuid]/messages', () => { - beforeEach(() => { - jest.clearAllMocks(); - }); - - it('maps messageId for bigint timestamps and preserves duplicate ordering', async () => { - mockGetConversation.mockResolvedValue({ - messages: [ - { - role: 'assistant', - content: 'first', - timestamp: 1700000000001, - }, - { - role: 'assistant', - content: 'second', - timestamp: 1700000000001, - }, - { - role: 'assistant', - content: 'third', - timestamp: 1700000000002, - }, - ], - lastActivity: 1700000005000, - }); - - const orderBy = jest.fn().mockResolvedValue([ - { id: 501, role: 'assistant', timestamp: '1700000000001' }, - { id: 502, role: 'assistant', timestamp: '1700000000001' }, - { id: 503, role: 'assistant', timestamp: '1700000000002' }, - ]); - const select = jest.fn().mockReturnValue({ orderBy }); - const where = jest.fn().mockReturnValue({ select }); - MockConversationMessage.query.mockReturnValue({ where }); - - const response = await GET(makeRequest(), { params: { buildUuid: 'uuid-1' } }); - const body = await response.json(); - - expect(response.status).toBe(200); - expect(body.data.messages).toHaveLength(3); - expect(body.data.messages[0].messageId).toBe(501); - expect(body.data.messages[1].messageId).toBe(502); - expect(body.data.messages[2].messageId).toBe(503); - }); -}); diff --git a/src/app/api/v2/ai/chat/[buildUuid]/messages/route.ts b/src/app/api/v2/ai/chat/[buildUuid]/messages/route.ts deleted file mode 100644 index 6db71df2..00000000 --- a/src/app/api/v2/ai/chat/[buildUuid]/messages/route.ts +++ /dev/null @@ -1,101 +0,0 @@ -/** - * Copyright 2025 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 { createApiHandler } from 'server/lib/createApiHandler'; -import { errorResponse, successResponse } from 'server/lib/response'; -import { defaultDb, defaultRedis } from 'server/lib/dependencies'; -import AIAgentConversationService from 'server/services/ai/conversation/storage'; -import ConversationMessage from 'server/models/ConversationMessage'; - -/** - * @openapi - * /api/v2/ai/chat/{buildUuid}/messages: - * get: - * summary: Get conversation messages - * description: > - * Returns the full conversation history for a given build UUID, including - * user and assistant messages. Assistant messages may include tool call - * activity history, evidence items, and debug data collected during streaming. - * Returns an empty array if no conversation exists for the build. - * tags: - * - AI Chat - * operationId: getAIChatMessages - * parameters: - * - in: path - * name: buildUuid - * required: true - * schema: - * type: string - * description: The UUID of the build to retrieve messages for. - * example: white-poetry-596195 - * responses: - * '200': - * description: Conversation messages with optional activity history and debug data. - * content: - * application/json: - * schema: - * $ref: '#/components/schemas/GetAIMessagesSuccessResponse' - * '400': - * description: Missing or invalid buildUuid - * content: - * application/json: - * schema: - * $ref: '#/components/schemas/ApiErrorResponse' - * '500': - * description: Server error - * content: - * application/json: - * schema: - * $ref: '#/components/schemas/ApiErrorResponse' - */ -const getHandler = async (req: NextRequest, { params }: { params: { buildUuid: string } }) => { - const { buildUuid } = params; - - if (!buildUuid) { - return errorResponse(new Error('Missing required parameter: buildUuid'), { status: 400 }, req); - } - - const conversationService = new AIAgentConversationService(defaultDb, defaultRedis); - const conversation = await conversationService.getConversation(buildUuid); - - const pgMessages = await ConversationMessage.query() - .where({ buildUuid }) - .select('id', 'timestamp', 'role') - .orderBy('id', 'asc'); - - const messageIdsByKey = new Map(); - for (const pgMessage of pgMessages) { - const key = `${pgMessage.role}:${String(pgMessage.timestamp)}`; - const existing = messageIdsByKey.get(key); - if (existing) { - existing.push(pgMessage.id); - } else { - messageIdsByKey.set(key, [pgMessage.id]); - } - } - - const messages = (conversation?.messages || []).map((msg) => { - const key = `${msg.role}:${String(msg.timestamp)}`; - const ids = messageIdsByKey.get(key); - const matchedMessageId = ids?.shift(); - return matchedMessageId != null ? { ...msg, messageId: matchedMessageId } : msg; - }); - - return successResponse({ messages, lastActivity: conversation?.lastActivity || null }, { status: 200 }, req); -}; - -export const GET = createApiHandler(getHandler); diff --git a/src/app/api/v2/ai/chat/[buildUuid]/route.ts b/src/app/api/v2/ai/chat/[buildUuid]/route.ts deleted file mode 100644 index 2f92b773..00000000 --- a/src/app/api/v2/ai/chat/[buildUuid]/route.ts +++ /dev/null @@ -1,678 +0,0 @@ -/** - * Copyright 2025 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 { createStreamHandler } from 'server/lib/createStreamHandler'; -import { defaultDb, defaultRedis } from 'server/lib/dependencies'; -import AIAgentContextService from 'server/services/ai/context/gatherer'; -import AIAgentConversationService from 'server/services/ai/conversation/storage'; -import AIAgentService from 'server/services/aiAgent'; -import AIAgentConfigService from 'server/services/aiAgentConfig'; -import { getLogger, withLogContext } from 'server/lib/logger'; -import { extractJsonFromResponse } from 'server/services/ai/utils/jsonExtraction'; -import { sanitizeForJson } from 'server/services/ai/utils/sanitize'; -import { normalizeInvestigationPayload } from 'server/services/ai/utils/normalizePayload'; -import { authorizeToolForFixTarget, FixTargetScope } from 'server/services/ai/utils/fixTargetAuthorization'; -import { - createClassifiedError, - ErrorCategory, - getUserErrorMessage, - getSuggestedAction, - isAuthError, -} from 'server/services/ai/errors'; -import { isBrokenCircuitError } from 'cockatiel'; -import type { ConfirmationDetails } from 'server/services/ai/types/tool'; -import type { AIChatSSEEvent, SSEErrorEvent } from 'shared/types/aiChat'; - -export const dynamic = 'force-dynamic'; - -/** - * @openapi - * /api/v2/ai/chat/{buildUuid}: - * post: - * summary: Stream AI chat response - * description: > - * Sends a message to the AI agent and streams the response as Server-Sent Events (SSE). - * The buildUuid identifies the ephemeral environment context. - * - * - * **Streaming protocol:** - * The response uses the `text/event-stream` content type. Each event is a JSON object - * sent as `data: {JSON}\n\n`. Clients should use the EventSource API or a streaming - * fetch with an AbortController. The connection is kept alive via the `Connection: keep-alive` - * header. Clients can cancel by aborting the request. - * - * - * **SSE event types (by `type` field):** - * - * - `chunk` — Streamed text fragment of the AI response. Concatenate all chunks to build - * the full response. See SSEChunkEvent schema. - * - * - `tool_call` — The AI is invoking a tool. Contains a `toolCallId` that correlates - * with a later `processing` event. See SSEToolCallEvent schema. - * - * - `processing` — A tool call has completed. Messages starting with a checkmark (✓) - * indicate success. See SSEProcessingEvent schema. - * - * - `thinking` — The AI is reasoning before producing output. See SSEThinkingEvent schema. - * - * - `error` — A non-fatal processing error during investigation. See SSEActivityErrorEvent schema. - * - * - `evidence_file` — A source file found as evidence. See SSEEvidenceFileEvent schema. - * - * - `evidence_commit` — A git commit found as evidence. See SSEEvidenceCommitEvent schema. - * - * - `evidence_resource` — A Kubernetes resource found as evidence. See SSEEvidenceResourceEvent schema. - * - * - `debug_context` — System prompt and model selection info. See SSEDebugContextEvent schema. - * - * - `debug_tool_call` — Raw tool invocation data. See SSEDebugToolCallEvent schema. - * - * - `debug_tool_result` — Raw tool result data. See SSEDebugToolResultEvent schema. - * - * - `debug_metrics` — Aggregate token/cost metrics. See SSEDebugMetricsEvent schema. - * - * - `complete_json` — Structured JSON response (e.g. investigation_complete). - * Sent before the `complete` event. See SSECompleteJsonEvent schema. - * - * - `complete` — Final event signaling end of stream. See SSECompleteEvent schema. - * - * - * **Error handling:** - * Because the HTTP 200 status is committed before streaming begins, errors during - * processing are delivered as SSE events with `error: true` instead of HTTP error codes. - * See the SSEErrorEvent schema for the error payload format, which includes a `category` - * field for classification and a `suggestedAction` field for client retry logic. - * - * - * **Typical event sequence:** - * `chunk*` → `tool_call` → `processing` → `chunk*` → ... → `complete_json`? → `complete` - * - * tags: - * - AI Chat - * operationId: streamAIChat - * parameters: - * - in: path - * name: buildUuid - * required: true - * schema: - * type: string - * description: The UUID of the build to chat with. - * requestBody: - * required: true - * content: - * application/json: - * schema: - * type: object - * required: - * - message - * properties: - * message: - * type: string - * description: The user message to send to the AI agent. - * example: Why is the web service CrashLoopBackOff? - * clearHistory: - * type: boolean - * description: Whether to clear conversation history before processing. - * default: false - * provider: - * type: string - * description: > - * The LLM provider to use. When omitted the default provider from the - * effective configuration is used. - * example: anthropic - * modelId: - * type: string - * description: > - * The specific model ID to use. Must belong to the given provider. - * When omitted the default model is used. - * example: claude-sonnet-4-20250514 - * isSystemAction: - * type: boolean - * description: > - * Whether this message is a system-initiated action (e.g. automatic - * investigation triggered by a deployment event). - * default: false - * mode: - * type: string - * enum: [investigate, fix] - * description: > - * Operation mode. "investigate" (default) is read-only analysis. - * "fix" enables the agent to take corrective actions via tools. - * default: investigate - * fixTarget: - * type: object - * description: > - * Optional service-scoped fix target. When provided in fix mode, - * mutating tool calls are constrained to this selected issue. - * The additional action hint fields are used by fix-target - * authorization to distinguish file edits from PR label - * updates and Kubernetes patch actions. - * properties: - * serviceName: - * type: string - * suggestedFix: - * type: string - * filePath: - * type: string - * files: - * type: array - * items: - * type: object - * properties: - * path: - * type: string - * autoFixAction: - * type: string - * actionType: - * type: string - * fixType: - * type: string - * tool: - * type: string - * toolName: - * type: string - * responses: - * '200': - * description: > - * SSE event stream. Each line is formatted as `data: {JSON}\n\n`. - * Events use a `type` discriminator field — see the SSE event schemas - * (SSEChunkEvent, SSEToolCallEvent, SSEProcessingEvent, SSEThinkingEvent, - * SSEActivityErrorEvent, SSEEvidenceFileEvent, SSEEvidenceCommitEvent, - * SSEEvidenceResourceEvent, SSEDebugContextEvent, SSEDebugToolCallEvent, - * SSEDebugToolResultEvent, SSEDebugMetricsEvent, SSECompleteJsonEvent, - * SSECompleteEvent). Errors arrive as SSEErrorEvent with `error: true`. - * headers: - * Content-Type: - * schema: - * type: string - * example: text/event-stream; charset=utf-8 - * Cache-Control: - * schema: - * type: string - * example: no-cache, no-transform - * Connection: - * schema: - * type: string - * example: keep-alive - * X-Accel-Buffering: - * schema: - * type: string - * example: 'no' - * content: - * text/event-stream: - * schema: - * type: string - * description: > - * Newline-delimited SSE events. Each event is `data: \n\n`. - * examples: - * chunk: - * summary: Text chunk event - * value: 'data: {"type":"chunk","content":"The web service is failing because..."}\n\n' - * tool_call: - * summary: Tool invocation event - * value: 'data: {"type":"tool_call","message":"Reading pod logs...","toolCallId":"tc_1"}\n\n' - * processing: - * summary: Tool completion event - * value: 'data: {"type":"processing","message":"✓ Pod logs retrieved","toolCallId":"tc_1","details":{"toolDurationMs":1200}}\n\n' - * complete: - * summary: Stream completion event - * value: 'data: {"type":"complete","totalInvestigationTimeMs":8500}\n\n' - * error: - * summary: Stream error event - * value: 'data: {"error":true,"userMessage":"Rate limit exceeded","category":"rate-limited","suggestedAction":"retry","retryAfter":30,"modelName":"claude-sonnet-4-20250514"}\n\n' - * '400': - * description: Missing required fields or invalid JSON - * content: - * application/json: - * schema: - * $ref: '#/components/schemas/ApiErrorResponse' - */ -const postHandler = async (req: NextRequest, { params }: { params: { buildUuid: string } }): Promise => { - const { buildUuid } = params; - - let body: any; - try { - body = await req.json(); - } catch { - return new Response(JSON.stringify({ error: 'Invalid JSON in request body' }), { - status: 400, - headers: { 'Content-Type': 'application/json' }, - }); - } - - if (!buildUuid || !body.message) { - return new Response(JSON.stringify({ error: 'Missing required fields: buildUuid and message' }), { - status: 400, - headers: { 'Content-Type': 'application/json' }, - }); - } - - const { readable, writable } = new TransformStream(); - const writer = writable.getWriter(); - const encoder = new TextEncoder(); - let writerClosed = false; - - const sendEvent = (data: AIChatSSEEvent | SSEErrorEvent) => { - if (writerClosed) return; - try { - writer.write(encoder.encode(`data: ${JSON.stringify(data)}\n\n`)).catch(() => {}); - } catch { - writerClosed = true; - } - }; - - const MAX_STORED_JSON_LENGTH = 8000; - const truncateForStorage = (value: unknown): unknown => { - try { - const json = JSON.stringify(value); - if (json.length <= MAX_STORED_JSON_LENGTH) return value; - return { _truncated: true, preview: json.slice(0, MAX_STORED_JSON_LENGTH), originalLength: json.length }; - } catch { - return { _truncated: true, preview: '[non-serializable]', originalLength: 0 }; - } - }; - - req.signal.addEventListener('abort', () => { - writerClosed = true; - try { - writer.close(); - } catch { - // already closed - } - }); - - (async () => { - try { - const aiAgentConfigService = AIAgentConfigService.getInstance(); - const aiAgentConfig = await aiAgentConfigService.getEffectiveConfig(); - - if (!aiAgentConfig?.enabled) { - sendEvent({ - error: true, - userMessage: 'AI Agent is not enabled', - category: 'deterministic', - suggestedAction: 'check-config', - retryAfter: null, - modelName: 'AI model', - code: 'AI_AGENT_DISABLED', - }); - return; - } - - const { message, clearHistory, provider, modelId, isSystemAction, mode: requestedMode, fixTarget } = body; - - getLogger().info( - `AI: v2 chat request received provider=${provider} modelId=${modelId} hasProvider=${!!provider} hasModelId=${!!modelId}` - ); - - await withLogContext({ buildUuid }, async () => { - const aiAgentContextService = new AIAgentContextService(defaultDb, defaultRedis); - const conversationService = new AIAgentConversationService(defaultDb, defaultRedis); - const llmService = new AIAgentService(defaultDb, defaultRedis); - - if (clearHistory) { - await conversationService.clearConversation(buildUuid); - } - - const conversation = await conversationService.getConversation(buildUuid); - const conversationHistory = conversation?.messages || []; - - let context; - try { - context = await aiAgentContextService.gatherFullContext(buildUuid); - } catch (error: any) { - getLogger().error({ error }, 'AI: context gather failed'); - sendEvent({ - error: true, - userMessage: `Build not found: ${error.message}`, - category: 'deterministic', - suggestedAction: null, - retryAfter: null, - modelName: 'AI model', - code: 'CONTEXT_ERROR', - }); - return; - } - - const repoFullName = context.lifecycleContext?.pullRequest?.fullName; - - const mode = requestedMode === 'fix' ? 'fix' : 'investigate'; - - try { - if (provider && modelId) { - await llmService.initializeWithMode(mode, provider, modelId, repoFullName); - } else { - await llmService.initialize(repoFullName); - } - } catch (error: any) { - getLogger().error({ error }, 'AI: init failed'); - sendEvent({ - error: true, - userMessage: error.message, - category: 'deterministic', - suggestedAction: 'check-config', - retryAfter: null, - modelName: modelId || 'AI model', - code: 'LLM_INIT_ERROR', - }); - return; - } - - let aiResponse = ''; - let isJsonResponse = false; - let totalInvestigationTimeMs = 0; - const collectedActivities: Array<{ - type: string; - message: string; - status?: 'pending' | 'completed' | 'failed'; - details?: { toolDurationMs?: number; totalDurationMs?: number }; - toolCallId?: string; - resultPreview?: string; - }> = []; - const collectedEvidence: Array> = []; - let collectedDebugContext: any = null; - const collectedDebugToolData = new Map< - string, - { - toolCallId: string; - toolName: string; - toolArgs: Record; - toolResult?: unknown; - toolDurationMs?: number; - } - >(); - let collectedDebugMetrics: any = null; - try { - getLogger().info(`AI: using mode=${mode} (requestedMode=${requestedMode})`); - - const onToolConfirmation = - mode === 'fix' - ? async (details: ConfirmationDetails) => { - return true; - } - : undefined; - const onToolAuthorization = - mode === 'fix' && fixTarget - ? async ( - tool: { name: string; description: string; category: string; safetyLevel: string }, - args: Record - ) => authorizeToolForFixTarget(fixTarget as FixTargetScope, { ...tool, args }) - : undefined; - - const result = await llmService.processQueryStream( - message, - context, - conversationHistory, - (chunk) => { - sendEvent({ type: 'chunk', content: chunk }); - }, - (activity) => { - sendEvent(activity as AIChatSSEEvent); - if (activity.type === 'tool_call') { - collectedActivities.push({ - type: activity.type, - message: activity.message, - status: 'pending', - toolCallId: activity.toolCallId, - }); - } else if (activity.type === 'processing') { - const isFailed = !activity.message.startsWith('\u2713'); - const matchIdx = collectedActivities.findIndex( - (a) => a.status === 'pending' && activity.toolCallId && a.toolCallId === activity.toolCallId - ); - if (matchIdx !== -1) { - collectedActivities[matchIdx] = { - ...collectedActivities[matchIdx], - message: activity.message, - status: isFailed ? 'failed' : 'completed', - details: activity.details, - toolCallId: activity.toolCallId, - resultPreview: activity.resultPreview, - }; - } else { - collectedActivities.push({ - type: activity.type, - message: activity.message, - status: isFailed ? 'failed' : 'completed', - details: activity.details, - toolCallId: activity.toolCallId, - resultPreview: activity.resultPreview, - }); - } - } else { - collectedActivities.push({ type: activity.type, message: activity.message }); - } - }, - (evidenceEvent) => { - sendEvent(evidenceEvent); - collectedEvidence.push(evidenceEvent as unknown as Record); - }, - onToolConfirmation, - onToolAuthorization, - mode, - (event: AIChatSSEEvent) => { - try { - sendEvent(event); - } catch { - /* SSE send must not block collection */ - } - try { - const e = event as any; - if (e.type === 'debug_context') { - collectedDebugContext = { - systemPrompt: e.systemPrompt, - maskingStats: e.maskingStats, - provider: e.provider, - modelId: e.modelId, - }; - } else if (e.type === 'debug_tool_call') { - collectedDebugToolData.set(e.toolCallId, { - toolCallId: e.toolCallId, - toolName: e.toolName, - toolArgs: e.toolArgs, - }); - } else if (e.type === 'debug_tool_result') { - const existing = collectedDebugToolData.get(e.toolCallId); - if (existing) { - existing.toolResult = e.toolResult; - existing.toolDurationMs = e.toolDurationMs; - } else { - collectedDebugToolData.set(e.toolCallId, { - toolCallId: e.toolCallId, - toolName: e.toolName, - toolArgs: {}, - toolResult: e.toolResult, - toolDurationMs: e.toolDurationMs, - }); - } - } else if (e.type === 'debug_metrics') { - collectedDebugMetrics = { - iterations: e.iterations, - totalToolCalls: e.totalToolCalls, - totalDurationMs: e.totalDurationMs, - inputTokens: e.inputTokens, - outputTokens: e.outputTokens, - inputCostPerMillion: e.inputCostPerMillion, - outputCostPerMillion: e.outputCostPerMillion, - }; - } - } catch { - /* debug collection must never disrupt the stream */ - } - } - ); - - aiResponse = result.response; - isJsonResponse = result.isJson; - totalInvestigationTimeMs = result.totalInvestigationTimeMs; - let preambleText: string | undefined = result.preamble; - let completeJsonEmitted = false; - - if (aiResponse.includes('"investigation_complete"')) { - const extracted = extractJsonFromResponse(aiResponse, buildUuid); - if (extracted.isJson) { - aiResponse = extracted.response; - isJsonResponse = true; - if (extracted.preamble && !preambleText) { - preambleText = extracted.preamble; - } - } - } - - if (isJsonResponse) { - try { - let parsed = JSON.parse(aiResponse); - getLogger().info( - `AI: JSON response type=${parsed.type} hasPullRequest=${!!context.lifecycleContext - ?.pullRequest} fullName=${context.lifecycleContext?.pullRequest?.fullName} branch=${ - context.lifecycleContext?.pullRequest?.branch - }` - ); - if (parsed.type === 'investigation_complete') { - parsed = normalizeInvestigationPayload(parsed, { availableTools: result.availableTools }); - - if (context.lifecycleContext?.pullRequest) { - const fullName = context.lifecycleContext.pullRequest.fullName; - const branch = context.lifecycleContext.pullRequest.branch; - if (fullName && branch) { - const [owner, name] = fullName.split('/'); - parsed.repository = { owner, name, branch }; - getLogger().info(`AI: added repository to response owner=${owner} name=${name} branch=${branch}`); - } - } - - const sanitized = sanitizeForJson(parsed); - aiResponse = JSON.stringify(sanitized, null, 2); - - JSON.parse(aiResponse); - } - } catch (e) { - getLogger().error( - { error: e instanceof Error ? e.message : String(e), responseLength: aiResponse.length }, - 'AI: JSON validation failed' - ); - aiResponse = - '\u26a0\ufe0f Investigation completed but response formatting failed. Please try asking a more specific question.'; - isJsonResponse = false; - } - - if (isJsonResponse && !completeJsonEmitted) { - completeJsonEmitted = true; - sendEvent({ - type: 'complete_json', - content: aiResponse, - totalInvestigationTimeMs, - ...(preambleText ? { preamble: preambleText } : {}), - }); - } - } - } catch (error: any) { - getLogger().error({ error }, 'AI: query failed'); - const currentModel = modelId || 'AI model'; - - if (isBrokenCircuitError(error)) { - const ctx = { modelName: currentModel, providerName: provider || 'unknown' }; - sendEvent({ - error: true, - userMessage: getUserErrorMessage(ErrorCategory.TRANSIENT, ctx), - category: 'transient', - suggestedAction: 'switch-model', - retryAfter: null, - modelName: currentModel, - code: 'CIRCUIT_BREAKER_OPEN', - }); - } else { - const classified = createClassifiedError(provider || 'unknown', error); - const ctx = { - modelName: currentModel, - providerName: classified.providerName, - retryAfter: classified.retryAfter, - isAuthError: isAuthError(error), - }; - sendEvent({ - error: true, - userMessage: getUserErrorMessage(classified.category, ctx), - category: classified.category as SSEErrorEvent['category'], - suggestedAction: getSuggestedAction(classified.category, ctx.isAuthError), - retryAfter: classified.retryAfter ?? null, - modelName: currentModel, - code: 'LLM_API_ERROR', - }); - } - return; - } - - await conversationService.addMessage(buildUuid, { - role: 'user', - content: message, - timestamp: Date.now(), - isSystemAction, - }); - - const assistantTimestamp = Date.now(); - await conversationService.addMessage(buildUuid, { - role: 'assistant', - content: aiResponse, - timestamp: assistantTimestamp, - activityHistory: collectedActivities.length > 0 ? collectedActivities : undefined, - evidenceItems: collectedEvidence.length > 0 ? collectedEvidence : undefined, - totalInvestigationTimeMs, - debugContext: collectedDebugContext || undefined, - debugToolData: - collectedDebugToolData.size > 0 - ? Array.from(collectedDebugToolData.values()).map((td) => ({ - ...td, - toolArgs: truncateForStorage(td.toolArgs), - toolResult: td.toolResult !== undefined ? truncateForStorage(td.toolResult) : undefined, - })) - : undefined, - debugMetrics: collectedDebugMetrics || undefined, - }); - - sendEvent({ type: 'complete', totalInvestigationTimeMs, assistantTimestamp }); - }); - } catch (error: any) { - getLogger().error({ error }, 'AI: chat request failed'); - sendEvent({ - error: true, - userMessage: error?.message || 'Internal error', - category: 'ambiguous', - suggestedAction: 'retry', - retryAfter: null, - modelName: 'AI model', - }); - } finally { - writerClosed = true; - try { - await writer.close(); - } catch { - // Writer may already be closed by abort handler - } - } - })(); - - return new Response(readable, { - headers: { - 'Content-Type': 'text/event-stream; charset=utf-8', - 'Cache-Control': 'no-cache, no-transform', - Connection: 'keep-alive', - 'X-Accel-Buffering': 'no', - }, - }); -}; - -export const POST = createStreamHandler(postHandler); diff --git a/src/app/api/v2/ai/chat/[buildUuid]/session/route.ts b/src/app/api/v2/ai/chat/[buildUuid]/session/route.ts deleted file mode 100644 index eff517e7..00000000 --- a/src/app/api/v2/ai/chat/[buildUuid]/session/route.ts +++ /dev/null @@ -1,77 +0,0 @@ -/** - * Copyright 2025 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 { createApiHandler } from 'server/lib/createApiHandler'; -import { errorResponse, successResponse } from 'server/lib/response'; -import { defaultDb, defaultRedis } from 'server/lib/dependencies'; -import AIAgentConversationService from 'server/services/ai/conversation/storage'; - -/** - * @openapi - * /api/v2/ai/chat/{buildUuid}/session: - * delete: - * summary: Clear conversation session - * description: > - * Deletes all conversation messages for a given build UUID and resets the session. - * Sessions are stored in Redis and expire automatically after the configured - * sessionTTL (default 3600 seconds). Use this endpoint to force-clear a session - * before its TTL, for example when starting a fresh investigation. - * tags: - * - AI Chat - * operationId: deleteAIChatSession - * parameters: - * - in: path - * name: buildUuid - * required: true - * schema: - * type: string - * description: The UUID of the build whose session to clear. - * example: white-poetry-596195 - * responses: - * '200': - * description: Session cleared successfully - * content: - * application/json: - * schema: - * $ref: '#/components/schemas/DeleteAISessionSuccessResponse' - * '400': - * description: Missing or invalid buildUuid - * content: - * application/json: - * schema: - * $ref: '#/components/schemas/ApiErrorResponse' - * '500': - * description: Server error - * content: - * application/json: - * schema: - * $ref: '#/components/schemas/ApiErrorResponse' - */ -const deleteHandler = async (req: NextRequest, { params }: { params: { buildUuid: string } }) => { - const { buildUuid } = params; - - if (!buildUuid) { - return errorResponse(new Error('Missing required parameter: buildUuid'), { status: 400 }, req); - } - - const conversationService = new AIAgentConversationService(defaultDb, defaultRedis); - const clearedCount = await conversationService.clearConversation(buildUuid); - - return successResponse({ success: true, messagesCleared: clearedCount }, { status: 200 }, req); -}; - -export const DELETE = createApiHandler(deleteHandler); diff --git a/src/app/api/v2/ai/config/mcp-presets/route.test.ts b/src/app/api/v2/ai/config/mcp-presets/route.test.ts index 267448fc..855aabb6 100644 --- a/src/app/api/v2/ai/config/mcp-presets/route.test.ts +++ b/src/app/api/v2/ai/config/mcp-presets/route.test.ts @@ -18,7 +18,7 @@ import { NextRequest } from 'next/server'; const mockListMcpPresets = jest.fn(); -jest.mock('server/services/ai/mcp/presets', () => ({ +jest.mock('server/services/agentRuntime/mcp/presets', () => ({ listMcpPresets: (...args: unknown[]) => mockListMcpPresets(...args), })); diff --git a/src/app/api/v2/ai/config/mcp-presets/route.ts b/src/app/api/v2/ai/config/mcp-presets/route.ts index e8552595..7584cf96 100644 --- a/src/app/api/v2/ai/config/mcp-presets/route.ts +++ b/src/app/api/v2/ai/config/mcp-presets/route.ts @@ -17,7 +17,7 @@ import { NextRequest } from 'next/server'; import { createApiHandler } from 'server/lib/createApiHandler'; import { successResponse } from 'server/lib/response'; -import { listMcpPresets } from 'server/services/ai/mcp/presets'; +import { listMcpPresets } from 'server/services/agentRuntime/mcp/presets'; /** * @openapi diff --git a/src/app/api/v2/ai/config/mcp-servers/[slug]/route.ts b/src/app/api/v2/ai/config/mcp-servers/[slug]/route.ts index b015a261..71c7a1d6 100644 --- a/src/app/api/v2/ai/config/mcp-servers/[slug]/route.ts +++ b/src/app/api/v2/ai/config/mcp-servers/[slug]/route.ts @@ -17,7 +17,7 @@ import { NextRequest, NextResponse } from 'next/server'; import { createApiHandler } from 'server/lib/createApiHandler'; import { successResponse } from 'server/lib/response'; -import { McpConfigService, redactSharedConfigSecrets } from 'server/services/ai/mcp/config'; +import { McpConfigService, redactMcpConfigSecrets } from 'server/services/agentRuntime/mcp/config'; import 'server/lib/dependencies'; /** @@ -81,7 +81,7 @@ const getHandler = async (req: NextRequest, { params }: { params: Promise<{ slug ); } - const result = redactSharedConfigSecrets(config.toJSON ? config.toJSON() : config); + const result = redactMcpConfigSecrets(config.toJSON ? config.toJSON() : config); return successResponse(result, { status: 200 }, req); }; @@ -151,7 +151,7 @@ const putHandler = async (req: NextRequest, { params }: { params: Promise<{ slug try { const config = await service.update(slug, scope, body); - const result = redactSharedConfigSecrets(config.toJSON ? config.toJSON() : config); + const result = redactMcpConfigSecrets(config.toJSON ? config.toJSON() : config); return successResponse(result, { status: 200 }, req); } catch (error) { const message = error instanceof Error ? error.message : String(error); diff --git a/src/app/api/v2/ai/config/mcp-servers/route.ts b/src/app/api/v2/ai/config/mcp-servers/route.ts index afcad0b2..b61fd5a7 100644 --- a/src/app/api/v2/ai/config/mcp-servers/route.ts +++ b/src/app/api/v2/ai/config/mcp-servers/route.ts @@ -17,7 +17,7 @@ import { NextRequest, NextResponse } from 'next/server'; import { createApiHandler } from 'server/lib/createApiHandler'; import { successResponse } from 'server/lib/response'; -import { McpConfigService, redactSharedConfigSecrets } from 'server/services/ai/mcp/config'; +import { McpConfigService, redactMcpConfigSecrets } from 'server/services/agentRuntime/mcp/config'; import 'server/lib/dependencies'; /** @@ -64,7 +64,7 @@ const getHandler = async (req: NextRequest) => { const scope = req.nextUrl.searchParams.get('scope') || 'global'; const service = new McpConfigService(); const configs = await service.listByScope(scope); - const redacted = configs.map((config) => redactSharedConfigSecrets(config.toJSON ? config.toJSON() : config)); + const redacted = configs.map((config) => redactMcpConfigSecrets(config.toJSON ? config.toJSON() : config)); return successResponse(redacted, { status: 200 }, req); }; @@ -149,7 +149,7 @@ const postHandler = async (req: NextRequest) => { try { const config = await service.create(input); - const result = redactSharedConfigSecrets(config.toJSON ? config.toJSON() : config); + const result = redactMcpConfigSecrets(config.toJSON ? config.toJSON() : config); return successResponse(result, { status: 201 }, req); } catch (error) { const message = error instanceof Error ? error.message : String(error); diff --git a/src/app/api/v2/ai/config/route.ts b/src/app/api/v2/ai/config/route.ts index df460529..ef4b0657 100644 --- a/src/app/api/v2/ai/config/route.ts +++ b/src/app/api/v2/ai/config/route.ts @@ -17,16 +17,16 @@ import { NextRequest } from 'next/server'; import { createApiHandler } from 'server/lib/createApiHandler'; import { successResponse } from 'server/lib/response'; -import AIAgentConfigService from 'server/services/aiAgentConfig'; +import AgentRuntimeConfigService from 'server/services/agentRuntime/config/agentRuntimeConfig'; import { getProviderEnvVarCandidates, normalizeStoredAgentProviderName } from 'server/services/agent/providerConfig'; /** * @openapi * /api/v2/ai/config: * get: - * summary: Check AI agent configuration status + * summary: Check Agent runtime configuration status * description: > - * Lightweight status check that returns whether the AI agent is enabled, + * Lightweight status check that returns whether the agent runtime is enabled, * which LLM provider is active, and whether the provider API key is set. * Use this endpoint to gate UI features before making heavier API calls. * When `enabled` is false, the chat endpoint will reject messages. @@ -48,14 +48,14 @@ import { getProviderEnvVarCandidates, normalizeStoredAgentProviderName } from 's * $ref: '#/components/schemas/ApiErrorResponse' */ const getHandler = async (req: NextRequest) => { - const aiAgentConfigService = AIAgentConfigService.getInstance(); - const aiAgentConfig = await aiAgentConfigService.getEffectiveConfig(); + const agentRuntimeConfigService = AgentRuntimeConfigService.getInstance(); + const agentRuntimeConfig = await agentRuntimeConfigService.getEffectiveConfig(); - if (!aiAgentConfig?.enabled) { + if (!agentRuntimeConfig?.enabled) { return successResponse({ enabled: false }, { status: 200 }, req); } - const enabledProvider = aiAgentConfig.providers?.find((p: any) => p.enabled !== false); + const enabledProvider = agentRuntimeConfig.providers?.find((p: any) => p.enabled !== false); const provider = normalizeStoredAgentProviderName(enabledProvider?.name) || 'anthropic'; const apiKeySet = getProviderEnvVarCandidates(provider, enabledProvider?.apiKeyEnvVar).some((envVar) => Boolean(process.env[envVar]) diff --git a/src/app/api/v2/ai/models/route.ts b/src/app/api/v2/ai/models/route.ts deleted file mode 100644 index 5812d8aa..00000000 --- a/src/app/api/v2/ai/models/route.ts +++ /dev/null @@ -1,77 +0,0 @@ -/** - * Copyright 2025 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 { createApiHandler } from 'server/lib/createApiHandler'; -import { successResponse } from 'server/lib/response'; -import { getLogger } from 'server/lib/logger'; -import AIAgentConfigService from 'server/services/aiAgentConfig'; -import { transformProviderModels } from 'server/services/ai/utils/modelTransformation'; - -export const dynamic = 'force-dynamic'; - -/** - * @openapi - * /api/v2/ai/models: - * get: - * summary: Get available AI models - * description: > - * Returns a list of enabled AI models from the effective configuration. - * Each model entry includes the provider, model ID, display name, max tokens, - * and whether it is the default. Returns an empty array if the AI agent is disabled. - * Use this endpoint to populate a model-picker in the UI before calling the chat endpoint. - * tags: - * - AI Chat - * operationId: getAIModels - * responses: - * '200': - * description: List of available AI models. Empty array when AI is disabled. - * content: - * application/json: - * schema: - * $ref: '#/components/schemas/GetAIModelsSuccessResponse' - * '500': - * description: Server error - * content: - * application/json: - * schema: - * $ref: '#/components/schemas/ApiErrorResponse' - */ -const getHandler = async (req: NextRequest) => { - const aiAgentConfigService = AIAgentConfigService.getInstance(); - const aiAgentConfig = await aiAgentConfigService.getEffectiveConfig(); - - if (!aiAgentConfig?.enabled) { - return successResponse({ models: [] }, { status: 200 }, req); - } - - if (!aiAgentConfig.providers || !Array.isArray(aiAgentConfig.providers)) { - getLogger().warn('AI: config missing providers array'); - return successResponse({ models: [] }, { status: 200 }, req); - } - - const models = transformProviderModels(aiAgentConfig.providers); - - getLogger().info( - `AI: models endpoint returning ${models.length} models: ${models - .map((m: any) => `${m.displayName}[provider=${m.provider}]`) - .join(', ')}` - ); - - return successResponse({ models }, { status: 200 }, req); -}; - -export const GET = createApiHandler(getHandler); diff --git a/src/server/db/migrations/022_agent_runtime_contract_unification.ts b/src/server/db/migrations/022_agent_runtime_contract_unification.ts new file mode 100644 index 00000000..068fe65a --- /dev/null +++ b/src/server/db/migrations/022_agent_runtime_contract_unification.ts @@ -0,0 +1,216 @@ +/** + * 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_ARRAY = '[]'; +const EMPTY_OBJECT = '{}'; +const BUILD_CONTEXT_CHAT_INDEX_NAME = 'agent_sessions_active_build_context_chat_unique'; +const AGENT_DEFINITIONS_OWNER_STATUS_UPDATED_AT_INDEX = 'agent_definitions_owner_status_updated_at_idx'; +const AGENT_DEFINITIONS_OWNER_DEFINITION_ID_INDEX = 'agent_definitions_owner_definition_id_idx'; +const LEGACY_GLOBAL_CONFIG_KEY = 'aiAgent'; +const RUNTIME_GLOBAL_CONFIG_KEY = 'agentRuntime'; +const LEGACY_REPO_CONFIG_TABLE = 'ai_agent_repo_config'; +const RUNTIME_REPO_CONFIG_TABLE = 'agent_runtime_repo_config'; + +async function migrateGlobalConfigKey(knex: Knex, fromKey: string, toKey: string, description: string): Promise { + const fromRow = await knex('global_config').where({ key: fromKey }).first(); + const toRow = await knex('global_config').where({ key: toKey }).first(); + + if (fromRow && !toRow) { + await knex('global_config').where({ key: fromKey }).update({ + key: toKey, + description, + updatedAt: knex.fn.now(), + }); + return; + } + + if (fromRow && toRow) { + await knex('global_config').where({ key: fromKey }).delete(); + } +} + +async function copyRepoConfigRows(knex: Knex, fromTable: string, toTable: string): Promise { + await knex.raw( + ` + INSERT INTO ?? ("repositoryFullName", config, "createdAt", "updatedAt", "deletedAt") + SELECT "repositoryFullName", config, "createdAt", "updatedAt", "deletedAt" + FROM ?? + ON CONFLICT ("repositoryFullName") DO UPDATE SET + config = EXCLUDED.config, + "updatedAt" = EXCLUDED."updatedAt", + "deletedAt" = EXCLUDED."deletedAt" + `, + [toTable, fromTable] + ); +} + +async function renameRepoConfigTable(knex: Knex, fromTable: string, toTable: string): Promise { + const hasFromTable = await knex.schema.hasTable(fromTable); + const hasToTable = await knex.schema.hasTable(toTable); + + if (hasFromTable && !hasToTable) { + await knex.schema.renameTable(fromTable, toTable); + return; + } + + if (hasFromTable && hasToTable) { + await copyRepoConfigRows(knex, fromTable, toTable); + await knex.schema.dropTable(fromTable); + } +} + +async function renamePostgresIdentifier( + knex: Knex, + type: 'constraint' | 'index' | 'sequence', + tableName: string, + from: string, + to: string +): Promise { + if (type === 'constraint') { + const fromExists = await knex('pg_constraint as c') + .join('pg_class as rel', 'rel.oid', 'c.conrelid') + .where('c.conname', from) + .where('rel.relname', tableName) + .first('c.oid'); + const toExists = await knex('pg_constraint as c') + .join('pg_class as rel', 'rel.oid', 'c.conrelid') + .where('c.conname', to) + .where('rel.relname', tableName) + .first('c.oid'); + + if (fromExists && !toExists) { + await knex.raw('ALTER TABLE ?? RENAME CONSTRAINT ?? TO ??', [tableName, from, to]); + } + return; + } + + const relkind = type === 'index' ? 'i' : 'S'; + const fromExists = await knex('pg_class').where({ relname: from, relkind }).first('oid'); + const toExists = await knex('pg_class').where({ relname: to, relkind }).first('oid'); + + if (!fromExists || toExists) { + return; + } + + await knex.raw(`ALTER ${type === 'index' ? 'INDEX' : 'SEQUENCE'} ?? RENAME TO ??`, [from, to]); +} + +async function renameRepoConfigPostgresObjects(knex: Knex, direction: 'up' | 'down'): Promise { + const fromPrefix = direction === 'up' ? 'ai_agent_repo_config' : 'agent_runtime_repo_config'; + const toPrefix = direction === 'up' ? 'agent_runtime_repo_config' : 'ai_agent_repo_config'; + const tableName = direction === 'up' ? RUNTIME_REPO_CONFIG_TABLE : LEGACY_REPO_CONFIG_TABLE; + + await renamePostgresIdentifier(knex, 'constraint', tableName, `${fromPrefix}_pkey`, `${toPrefix}_pkey`); + await renamePostgresIdentifier( + knex, + 'constraint', + tableName, + `${fromPrefix}_repositoryfullname_unique`, + `${toPrefix}_repositoryfullname_unique` + ); + await renamePostgresIdentifier( + knex, + 'index', + tableName, + `${fromPrefix}_repositoryfullname_index`, + `${toPrefix}_repositoryfullname_index` + ); + await renamePostgresIdentifier(knex, 'sequence', tableName, `${fromPrefix}_id_seq`, `${toPrefix}_id_seq`); +} + +async function createAgentDefinitionsTable(knex: Knex): Promise { + await knex.schema.createTable('agent_definitions', (table) => { + table.increments('id').primary(); + table.string('definitionId', 255).notNullable().unique(); + table.integer('version').notNullable().defaultTo(1); + table.string('ownerKind', 32).notNullable(); + table.string('ownerUserId', 255).nullable(); + table.string('ownerOrganizationId', 255).nullable(); + table.string('name', 255).notNullable(); + table.text('description').nullable(); + table.jsonb('instructionRefs').notNullable().defaultTo(EMPTY_ARRAY); + table.text('instructionAddendum').nullable(); + table.jsonb('capabilityRefs').notNullable().defaultTo(EMPTY_ARRAY); + table.jsonb('requiredCapabilityRefs').notNullable().defaultTo(EMPTY_ARRAY); + table.jsonb('optionalCapabilityRefs').notNullable().defaultTo(EMPTY_ARRAY); + table.jsonb('resourcePolicy').notNullable().defaultTo(EMPTY_OBJECT); + table.jsonb('modelPreference').nullable(); + table.string('status', 32).notNullable().defaultTo('active'); + table.boolean('codeOwned').notNullable().defaultTo(false); + table.boolean('readOnly').notNullable().defaultTo(false); + table.timestamp('createdAt').notNullable().defaultTo(knex.fn.now()); + table.timestamp('updatedAt').notNullable().defaultTo(knex.fn.now()); + + table.index(['definitionId']); + table.index(['ownerKind']); + table.index(['status']); + table.index(['ownerKind', 'ownerUserId', 'status', 'updatedAt'], AGENT_DEFINITIONS_OWNER_STATUS_UPDATED_AT_INDEX); + table.index(['ownerKind', 'ownerUserId', 'definitionId'], AGENT_DEFINITIONS_OWNER_DEFINITION_ID_INDEX); + }); +} + +export async function up(knex: Knex): Promise { + await knex.raw(` + create unique index if not exists "${BUILD_CONTEXT_CHAT_INDEX_NAME}" + on agent_sessions ("userId", "buildUuid") + where "buildUuid" is not null + and "sessionKind" = 'chat' + and "status" = 'active' + and "chatStatus" = 'ready' + `); + + await knex.schema.alterTable('agent_runs', (table) => { + table.jsonb('runPlanSnapshot').nullable(); + }); + + await createAgentDefinitionsTable(knex); + + await migrateGlobalConfigKey( + knex, + LEGACY_GLOBAL_CONFIG_KEY, + RUNTIME_GLOBAL_CONFIG_KEY, + 'Global configuration for agent runtime providers, limits, approvals, capabilities, and custom-agent creation.' + ); + + await renameRepoConfigTable(knex, LEGACY_REPO_CONFIG_TABLE, RUNTIME_REPO_CONFIG_TABLE); + await renameRepoConfigPostgresObjects(knex, 'up'); +} + +export async function down(knex: Knex): Promise { + await migrateGlobalConfigKey( + knex, + RUNTIME_GLOBAL_CONFIG_KEY, + LEGACY_GLOBAL_CONFIG_KEY, + 'Global configuration for AI agent providers, limits, approvals, capabilities, and custom-agent creation.' + ); + + await renameRepoConfigTable(knex, RUNTIME_REPO_CONFIG_TABLE, LEGACY_REPO_CONFIG_TABLE); + await renameRepoConfigPostgresObjects(knex, 'down'); + + await knex.schema.dropTableIfExists('agent_definitions'); + + await knex.schema.alterTable('agent_runs', (table) => { + table.dropColumn('runPlanSnapshot'); + }); + + await knex.raw(`drop index if exists "${BUILD_CONTEXT_CHAT_INDEX_NAME}"`); +} diff --git a/src/server/lib/__tests__/get-user.test.ts b/src/server/lib/__tests__/get-user.test.ts index bdab371d..b8978c90 100644 --- a/src/server/lib/__tests__/get-user.test.ts +++ b/src/server/lib/__tests__/get-user.test.ts @@ -82,6 +82,7 @@ describe('get-user helpers', () => { given_name: 'Sample', family_name: 'User', email: 'sample-user@example.com', + realm_access: { roles: ['user', 'admin', 'offline_access'] }, }) ); @@ -93,6 +94,7 @@ describe('get-user helpers', () => { displayName: 'Sample User', gitUserName: 'Sample User', gitUserEmail: 'sample-user@example.com', + roles: ['user', 'admin'], }) ); }); @@ -106,6 +108,7 @@ describe('get-user helpers', () => { displayName: 'local-dev-user', gitUserName: 'local-dev-user', gitUserEmail: 'local-dev-user@users.noreply.github.com', + roles: ['admin'], }) ); }); diff --git a/src/server/lib/agentSession/__tests__/podFactory.test.ts b/src/server/lib/agentSession/__tests__/podFactory.test.ts index 22928b06..4fac4ed6 100644 --- a/src/server/lib/agentSession/__tests__/podFactory.test.ts +++ b/src/server/lib/agentSession/__tests__/podFactory.test.ts @@ -15,7 +15,7 @@ */ import * as k8s from '@kubernetes/client-node'; -import { SESSION_POD_MCP_CONFIG_ENV } from 'server/services/ai/mcp/sessionPod'; +import { SESSION_POD_MCP_CONFIG_ENV } from 'server/services/agentRuntime/mcp/sessionPod'; import { SESSION_WORKSPACE_HOME_VOLUME_NAME, SESSION_WORKSPACE_SHARED_HOME_DIR } from '../configSeeder'; import { SESSION_WORKSPACE_EDITOR_PROJECT_FILE } from '../workspace'; diff --git a/src/server/lib/agentSession/__tests__/systemPrompt.test.ts b/src/server/lib/agentSession/__tests__/systemPrompt.test.ts index 2b840154..7682a449 100644 --- a/src/server/lib/agentSession/__tests__/systemPrompt.test.ts +++ b/src/server/lib/agentSession/__tests__/systemPrompt.test.ts @@ -28,6 +28,7 @@ import Deploy from 'server/models/Deploy'; import { fetchLifecycleConfig, getDeployingServicesByName } from 'server/models/yaml'; import { buildAgentSessionDynamicSystemPrompt, + buildLifecycleDebuggingProfilePrompt, combineAgentSessionAppendSystemPrompt, resolveAgentSessionPromptContext, } from '../systemPrompt'; @@ -35,11 +36,16 @@ import { describe('agent session system prompt', () => { beforeEach(() => { jest.clearAllMocks(); + jest.useFakeTimers().setSystemTime(new Date('2026-04-30T12:00:00.000Z')); (AgentSession.query as jest.Mock) = jest.fn().mockReturnValue({ findById: jest.fn().mockResolvedValue(null), }); }); + afterEach(() => { + jest.useRealTimers(); + }); + it('builds a compact dynamic session context prompt', () => { expect( buildAgentSessionDynamicSystemPrompt({ @@ -79,13 +85,99 @@ describe('agent session system prompt', () => { ).toBe('Use concise responses.\n\nSession context:\n- namespace: env-sample'); }); + it('builds a Lifecycle debugging profile without the legacy JSON-only output contract', () => { + const profile = buildLifecycleDebuggingProfilePrompt(); + + expect(profile).toContain('Compare desired config state with actual runtime state'); + expect(profile).toContain('Investigate build failures before deploy failures'); + expect(profile).toContain('Cite specific evidence before diagnosing a root cause'); + expect(profile).toContain('Say when there is not enough evidence'); + expect(profile).toContain( + 'Only perform mutating fixes through approval-gated actions when those tools are available' + ); + expect(profile).not.toContain('output_schema'); + expect(profile).not.toContain('fixesApplied'); + expect(profile).not.toContain('investigation_complete'); + }); + + it('builds diagnostic prompt sections for build-context chats without sensitive legacy fields', () => { + const prompt = buildAgentSessionDynamicSystemPrompt({ + buildUuid: 'sample-build-1', + gatheredAt: '2026-04-30T12:00:00.000Z', + build: { + uuid: 'sample-build-1', + status: 'deploy_failed', + statusMessage: 'web deploy failed', + namespace: 'env-sample-123456', + sha: 'abc123', + }, + pullRequest: { + fullName: 'example-org/example-repo', + branchName: 'feature/sample', + pullRequestNumber: 42, + url: 'https://github.com/example-org/example-repo/pull/42', + status: 'open', + labels: ['lifecycle-deploy'], + latestCommit: 'abc123', + repositoryUrl: 'https://github.com/example-org/example-repo', + }, + services: [], + diagnosticServices: [ + { + name: 'next-web', + status: 'deploy_failed', + statusMessage: 'CrashLoopBackOff', + publicUrl: 'https://next-web-sample.lifecycle.dev.example.com', + repo: 'example-org/example-repo', + branch: 'feature/sample', + dockerImage: 'registry.example.test/next-web:abc123', + buildPipelineId: 'build-pipeline-1', + deployPipelineId: 'deploy-pipeline-1', + }, + ], + }); + + expect(prompt).toContain('Lifecycle debugging profile:'); + expect(prompt).toContain('Build context:'); + expect(prompt).toContain( + '- buildUuid=sample-build-1: status=deploy_failed, statusMessage=web deploy failed, namespace=env-sample-123456, sha=abc123' + ); + expect(prompt).toContain('Pull request:'); + expect(prompt).toContain( + '- repo=example-org/example-repo, branch=feature/sample, number=42, url=https://github.com/example-org/example-repo/pull/42, status=open, labels=lifecycle-deploy, latestCommit=abc123, repositoryUrl=https://github.com/example-org/example-repo' + ); + expect(prompt).toContain('Diagnostic services:'); + expect(prompt).toContain( + '- next-web: status=deploy_failed, statusMessage=CrashLoopBackOff, repo=example-org/example-repo, branch=feature/sample, publicUrl=https://next-web-sample.lifecycle.dev.example.com, dockerImage=registry.example.test/next-web:abc123, buildPipelineId=build-pipeline-1, deployPipelineId=deploy-pipeline-1' + ); + expect(prompt).toContain('Context freshness:'); + expect(prompt).toContain('- gatheredAt: 2026-04-30T12:00:00.000Z'); + expect(prompt).not.toContain('secret'); + expect(prompt).not.toContain('MCP token'); + expect(prompt).not.toContain('conversation_messages'); + expect(prompt).not.toContain('server/services/ai/context'); + expect(prompt).not.toContain('server/services/ai/prompts'); + }); + it('resolves selected service public URLs and workdirs from deploy and lifecycle config metadata', async () => { const buildGraphQuery = { withGraphFetched: jest.fn().mockResolvedValue({ + uuid: 'sample-123456', + status: 'ready', + statusMessage: 'ready', + namespace: 'env-sample-123456', pullRequest: { fullName: 'example-org/example-repo', branchName: 'feature/sample', + pullRequestNumber: 42, + status: 'open', + labels: ['lifecycle-deploy'], + latestCommit: 'abc123', + repository: { + htmlUrl: 'https://github.com/example-org/example-repo', + }, }, + deploys: [], }), }; (Build.query as jest.Mock) = jest.fn().mockReturnValue({ @@ -125,15 +217,39 @@ describe('agent session system prompt', () => { ).resolves.toEqual({ namespace: 'env-sample-123456', buildUuid: 'sample-123456', + gatheredAt: '2026-04-30T12:00:00.000Z', + build: { + uuid: 'sample-123456', + status: 'ready', + statusMessage: 'ready', + namespace: 'env-sample-123456', + sha: undefined, + }, + pullRequest: { + fullName: 'example-org/example-repo', + branchName: 'feature/sample', + pullRequestNumber: 42, + url: 'https://github.com/example-org/example-repo/pull/42', + status: 'open', + labels: ['lifecycle-deploy'], + latestCommit: 'abc123', + repositoryUrl: 'https://github.com/example-org/example-repo', + }, services: [ { name: 'next-web', + status: undefined, + statusMessage: undefined, publicUrl: 'https://next-web-sample.lifecycle.dev.example.com', repo: 'example-org/example-repo', branch: 'feature/sample', + dockerImage: undefined, + buildPipelineId: undefined, + deployPipelineId: undefined, workDir: '/workspace/apps/next-web', }, ], + diagnosticServices: [], skillsAvailable: false, }); @@ -145,4 +261,122 @@ describe('agent session system prompt', () => { 'next-web' ); }); + + it('resolves build-context chat diagnostics without a workspace namespace', async () => { + const buildGraphQuery = { + withGraphFetched: jest.fn().mockResolvedValue({ + uuid: 'sample-build-1', + status: 'build_failed', + statusMessage: 'image build failed', + namespace: 'env-sample-123456', + sha: 'abc123', + pullRequest: { + fullName: 'example-org/example-repo', + branchName: 'feature/sample', + pullRequestNumber: 42, + status: 'open', + labels: ['lifecycle-deploy'], + latestCommit: 'abc123', + repository: { + htmlUrl: 'https://github.com/example-org/example-repo', + }, + }, + deploys: [ + { + id: 10, + uuid: 'next-web-deploy-1', + status: 'deploy_failed', + statusMessage: 'CrashLoopBackOff', + branchName: 'feature/sample', + publicUrl: 'next-web-sample.lifecycle.dev.example.com', + dockerImage: 'registry.example.test/next-web:abc123', + buildPipelineId: 'build-pipeline-1', + deployPipelineId: 'deploy-pipeline-1', + deployable: { name: 'next-web' }, + repository: { fullName: 'example-org/example-repo' }, + service: null, + }, + ], + }), + }; + (Build.query as jest.Mock) = jest.fn().mockReturnValue({ + findOne: jest.fn().mockReturnValue(buildGraphQuery), + }); + + const deployGraphQuery = { + withGraphFetched: jest.fn().mockResolvedValue([]), + }; + (Deploy.query as jest.Mock) = jest.fn().mockReturnValue({ + where: jest.fn().mockReturnValue(deployGraphQuery), + }); + + (fetchLifecycleConfig as jest.Mock).mockResolvedValue({ + services: [{ name: 'next-web', dev: { workDir: '/workspace/apps/next-web' } }], + }); + (getDeployingServicesByName as jest.Mock).mockReturnValue({ + name: 'next-web', + dev: { workDir: '/workspace/apps/next-web' }, + }); + + await expect( + resolveAgentSessionPromptContext({ + sessionDbId: 123, + namespace: null, + buildUuid: 'sample-build-1', + }) + ).resolves.toEqual({ + namespace: null, + buildUuid: 'sample-build-1', + gatheredAt: '2026-04-30T12:00:00.000Z', + build: { + uuid: 'sample-build-1', + status: 'build_failed', + statusMessage: 'image build failed', + namespace: 'env-sample-123456', + sha: 'abc123', + }, + pullRequest: { + fullName: 'example-org/example-repo', + branchName: 'feature/sample', + pullRequestNumber: 42, + url: 'https://github.com/example-org/example-repo/pull/42', + status: 'open', + labels: ['lifecycle-deploy'], + latestCommit: 'abc123', + repositoryUrl: 'https://github.com/example-org/example-repo', + }, + services: [ + { + name: 'next-web', + status: 'deploy_failed', + statusMessage: 'CrashLoopBackOff', + publicUrl: 'https://next-web-sample.lifecycle.dev.example.com', + repo: 'example-org/example-repo', + branch: 'feature/sample', + dockerImage: 'registry.example.test/next-web:abc123', + buildPipelineId: 'build-pipeline-1', + deployPipelineId: 'deploy-pipeline-1', + workDir: '/workspace/apps/next-web', + }, + ], + diagnosticServices: [ + { + name: 'next-web', + status: 'deploy_failed', + statusMessage: 'CrashLoopBackOff', + publicUrl: 'https://next-web-sample.lifecycle.dev.example.com', + repo: 'example-org/example-repo', + branch: 'feature/sample', + dockerImage: 'registry.example.test/next-web:abc123', + buildPipelineId: 'build-pipeline-1', + deployPipelineId: 'deploy-pipeline-1', + }, + ], + skillsAvailable: false, + }); + + expect(buildGraphQuery.withGraphFetched).toHaveBeenCalledWith( + '[pullRequest.[repository], deploys.[deployable, repository, service]]' + ); + }); }); diff --git a/src/server/lib/agentSession/podFactory.ts b/src/server/lib/agentSession/podFactory.ts index 869010ab..dd367ee5 100644 --- a/src/server/lib/agentSession/podFactory.ts +++ b/src/server/lib/agentSession/podFactory.ts @@ -28,7 +28,10 @@ import { buildLifecycleLabels } from 'server/lib/kubernetes/labels'; import type { RequestUserIdentity } from 'server/lib/get-user'; import { buildPodEnvWithSecrets } from 'server/lib/secretEnvBuilder'; import type { SecretRefWithEnvKey } from 'server/lib/secretRefs'; -import { SESSION_POD_MCP_CONFIG_ENV, SESSION_POD_MCP_CONFIG_SECRET_KEY } from 'server/services/ai/mcp/sessionPod'; +import { + SESSION_POD_MCP_CONFIG_ENV, + SESSION_POD_MCP_CONFIG_SECRET_KEY, +} from 'server/services/agentRuntime/mcp/sessionPod'; import { SESSION_WORKSPACE_EDITOR_PROJECT_FILE, SESSION_WORKSPACE_SUBPATH, diff --git a/src/server/lib/agentSession/systemPrompt.ts b/src/server/lib/agentSession/systemPrompt.ts index c1327557..6206c8ea 100644 --- a/src/server/lib/agentSession/systemPrompt.ts +++ b/src/server/lib/agentSession/systemPrompt.ts @@ -22,24 +22,52 @@ import type { LifecycleConfig } from 'server/models/yaml'; export interface AgentSessionPromptServiceContext { name: string; + status?: string; + statusMessage?: string; publicUrl?: string; repo?: string; branch?: string; + dockerImage?: string; + buildPipelineId?: string; + deployPipelineId?: string; workspacePath?: string; workDir?: string; } +export interface AgentSessionPromptBuildContext { + uuid: string; + status?: string; + statusMessage?: string; + namespace?: string; + sha?: string; +} + +export interface AgentSessionPromptPullRequestContext { + fullName?: string; + branchName?: string; + pullRequestNumber?: number; + url?: string; + status?: string; + labels?: string[]; + latestCommit?: string; + repositoryUrl?: string; +} + export interface AgentSessionPromptContext { - namespace: string; + namespace?: string | null; buildUuid?: string | null; + gatheredAt?: string; + build?: AgentSessionPromptBuildContext; + pullRequest?: AgentSessionPromptPullRequestContext; services: AgentSessionPromptServiceContext[]; + diagnosticServices?: AgentSessionPromptServiceContext[]; skillsAvailable?: boolean; toolLines?: string[]; } type SessionPromptLookupContext = { sessionDbId: number; - namespace: string; + namespace?: string | null; buildUuid?: string | null; }; @@ -56,22 +84,87 @@ function formatPublicUrl(value: unknown): string | undefined { return /^https?:\/\//.test(normalized) ? normalized : `https://${normalized}`; } +function normalizeStringArray(value: unknown): string[] | undefined { + if (!Array.isArray(value)) { + return undefined; + } + + const normalized = value.map((item) => normalizeOptionalString(item)).filter((item): item is string => Boolean(item)); + + return normalized.length ? normalized : undefined; +} + +function formatDetails(details: Array): string { + return details.filter((value): value is string => Boolean(value)).join(', '); +} + +export function buildLifecycleDebuggingProfilePrompt(): string { + return [ + 'Lifecycle debugging profile:', + '- Compare desired config state with actual runtime state before diagnosing.', + '- Investigate build failures before deploy failures.', + '- Cite specific evidence before diagnosing a root cause.', + '- Say when there is not enough evidence instead of fabricating a cause.', + '- Keep findings concise and lead with the highest-impact finding.', + '- Use available tools for fresh facts when the user says state changed or context is incomplete.', + '- Only perform mutating fixes through approval-gated actions when those tools are available.', + ].join('\n'); +} + export function buildAgentSessionDynamicSystemPrompt(context: AgentSessionPromptContext): string { - const lines = ['Session context:', `- namespace: ${context.namespace}`]; + const lines = ['Session context:']; + + if (context.namespace) { + lines.push(`- namespace: ${context.namespace}`); + } if (context.buildUuid) { lines.push(`- buildUuid: ${context.buildUuid}`); } + if (context.build) { + lines.push('', buildLifecycleDebuggingProfilePrompt(), '', 'Build context:'); + const details = formatDetails([ + context.build.status ? `status=${context.build.status}` : undefined, + context.build.statusMessage ? `statusMessage=${context.build.statusMessage}` : undefined, + context.build.namespace ? `namespace=${context.build.namespace}` : undefined, + context.build.sha ? `sha=${context.build.sha}` : undefined, + ]); + lines.push(`- buildUuid=${context.build.uuid}${details ? `: ${details}` : ''}`); + } + + if (context.pullRequest) { + const pr = context.pullRequest; + const details = formatDetails([ + pr.fullName ? `repo=${pr.fullName}` : undefined, + pr.branchName ? `branch=${pr.branchName}` : undefined, + pr.pullRequestNumber != null ? `number=${pr.pullRequestNumber}` : undefined, + pr.url ? `url=${pr.url}` : undefined, + pr.status ? `status=${pr.status}` : undefined, + pr.labels?.length ? `labels=${pr.labels.join('|')}` : undefined, + pr.latestCommit ? `latestCommit=${pr.latestCommit}` : undefined, + pr.repositoryUrl ? `repositoryUrl=${pr.repositoryUrl}` : undefined, + ]); + + if (details) { + lines.push('Pull request:', `- ${details}`); + } + } + if (context.services.length > 0) { lines.push('- selected services:'); const services = [...context.services].sort((left, right) => left.name.localeCompare(right.name)); for (const service of services) { const details = [ + service.status ? `status=${service.status}` : null, + service.statusMessage ? `statusMessage=${service.statusMessage}` : null, service.repo ? `repo=${service.repo}` : null, service.branch ? `branch=${service.branch}` : null, service.publicUrl ? `publicUrl=${service.publicUrl}` : null, + service.dockerImage ? `dockerImage=${service.dockerImage}` : null, + service.buildPipelineId ? `buildPipelineId=${service.buildPipelineId}` : null, + service.deployPipelineId ? `deployPipelineId=${service.deployPipelineId}` : null, service.workspacePath ? `workspacePath=${service.workspacePath}` : null, service.workDir ? `workDir=${service.workDir}` : null, ].filter((value): value is string => Boolean(value)); @@ -80,6 +173,36 @@ export function buildAgentSessionDynamicSystemPrompt(context: AgentSessionPrompt } } + if (context.diagnosticServices?.length) { + lines.push('Diagnostic services:'); + + const diagnosticServices = [...context.diagnosticServices].sort((left, right) => + left.name.localeCompare(right.name) + ); + for (const service of diagnosticServices) { + const details = formatDetails([ + service.status ? `status=${service.status}` : undefined, + service.statusMessage ? `statusMessage=${service.statusMessage}` : undefined, + service.repo ? `repo=${service.repo}` : undefined, + service.branch ? `branch=${service.branch}` : undefined, + service.publicUrl ? `publicUrl=${service.publicUrl}` : undefined, + service.dockerImage ? `dockerImage=${service.dockerImage}` : undefined, + service.buildPipelineId ? `buildPipelineId=${service.buildPipelineId}` : undefined, + service.deployPipelineId ? `deployPipelineId=${service.deployPipelineId}` : undefined, + ]); + + lines.push(`- ${service.name}${details ? `: ${details}` : ''}`); + } + } + + if (context.gatheredAt) { + lines.push( + 'Context freshness:', + `- gatheredAt: ${context.gatheredAt}`, + '- Treat these as Lifecycle database facts from gatheredAt; use available tools for fresh facts when state may have changed.' + ); + } + if (context.skillsAvailable) { lines.push('- equipped skills: use skills.list to discover them and skills.learn to load a skill before using it'); } @@ -103,6 +226,14 @@ export function combineAgentSessionAppendSystemPrompt( return parts.length > 0 ? parts.join('\n\n') : undefined; } +type BuildDiagnosticContext = { + source: { repo?: string; branch?: string }; + build?: AgentSessionPromptBuildContext; + pullRequest?: AgentSessionPromptPullRequestContext; + deploys: Deploy[]; + diagnosticServices: AgentSessionPromptServiceContext[]; +}; + async function fetchCachedLifecycleConfig( repositoryName: string, branchName: string, @@ -119,16 +250,83 @@ async function fetchCachedLifecycleConfig( return promise; } -async function resolveBuildSource(buildUuid?: string | null): Promise<{ repo?: string; branch?: string }> { +function buildPullRequestUrl(fullName?: string, pullRequestNumber?: number): string | undefined { + if (!fullName || pullRequestNumber == null) { + return undefined; + } + + return `https://github.com/${fullName}/pull/${pullRequestNumber}`; +} + +function formatDeployDiagnosticService( + deploy: Deploy, + buildSource: { repo?: string; branch?: string } +): AgentSessionPromptServiceContext | null { + const name = + normalizeOptionalString(deploy.deployable?.name) || + normalizeOptionalString(deploy.service?.name) || + normalizeOptionalString(deploy.uuid); + + if (!name) { + return null; + } + + return { + name, + status: normalizeOptionalString(deploy.status), + statusMessage: normalizeOptionalString(deploy.statusMessage), + publicUrl: formatPublicUrl(deploy.publicUrl), + repo: normalizeOptionalString(deploy.repository?.fullName) || buildSource.repo, + branch: normalizeOptionalString(deploy.branchName) || buildSource.branch, + dockerImage: normalizeOptionalString(deploy.dockerImage), + buildPipelineId: normalizeOptionalString(deploy.buildPipelineId), + deployPipelineId: normalizeOptionalString(deploy.deployPipelineId), + }; +} + +async function resolveBuildDiagnosticContext(buildUuid?: string | null): Promise { const normalizedBuildUuid = normalizeOptionalString(buildUuid); if (!normalizedBuildUuid) { - return {}; + return { source: {}, deploys: [], diagnosticServices: [] }; } - const build = await Build.query().findOne({ uuid: normalizedBuildUuid }).withGraphFetched('[pullRequest]'); + const build = await Build.query() + .findOne({ uuid: normalizedBuildUuid }) + .withGraphFetched('[pullRequest.[repository], deploys.[deployable, repository, service]]'); + const pullRequest = build?.pullRequest; + const pullRequestNumber = pullRequest?.pullRequestNumber; + const source = { + repo: normalizeOptionalString(pullRequest?.fullName), + branch: normalizeOptionalString(pullRequest?.branchName), + }; + return { - repo: normalizeOptionalString(build?.pullRequest?.fullName), - branch: normalizeOptionalString(build?.pullRequest?.branchName), + source, + build: build + ? { + uuid: build.uuid, + status: normalizeOptionalString(build.status), + statusMessage: normalizeOptionalString(build.statusMessage), + namespace: normalizeOptionalString(build.namespace), + sha: normalizeOptionalString(build.sha), + } + : undefined, + pullRequest: pullRequest + ? { + fullName: source.repo, + branchName: source.branch, + pullRequestNumber, + url: buildPullRequestUrl(source.repo, pullRequestNumber), + status: normalizeOptionalString(pullRequest.status), + labels: normalizeStringArray(pullRequest.labels), + latestCommit: normalizeOptionalString(pullRequest.latestCommit), + repositoryUrl: normalizeOptionalString(pullRequest.repository?.htmlUrl), + } + : undefined, + deploys: build?.deploys || [], + diagnosticServices: (build?.deploys || []) + .map((deploy) => formatDeployDiagnosticService(deploy, source)) + .filter((service): service is AgentSessionPromptServiceContext => Boolean(service)), }; } @@ -140,10 +338,11 @@ export async function resolveAgentSessionPromptContext( Deploy.query() .where({ devModeSessionId: lookup.sessionDbId }) .withGraphFetched('[deployable, repository, service]'), - resolveBuildSource(lookup.buildUuid), + resolveBuildDiagnosticContext(lookup.buildUuid), ]); const lifecycleConfigCache = new Map>(); - const deployById = new Map(deploys.filter((deploy) => deploy.id != null).map((deploy) => [deploy.id, deploy])); + const allDeploys = deploys.length > 0 ? deploys : buildSource.deploys; + const deployById = new Map(allDeploys.filter((deploy) => deploy.id != null).map((deploy) => [deploy.id, deploy])); let services: AgentSessionPromptServiceContext[]; @@ -163,7 +362,7 @@ export async function resolveAgentSessionPromptContext( } else { services = ( await Promise.all( - deploys.map(async (deploy): Promise => { + allDeploys.map(async (deploy): Promise => { const serviceName = normalizeOptionalString(deploy.deployable?.name) || normalizeOptionalString(deploy.service?.name) || @@ -173,8 +372,8 @@ export async function resolveAgentSessionPromptContext( return null; } - const repositoryName = normalizeOptionalString(deploy.repository?.fullName) || buildSource.repo; - const branchName = normalizeOptionalString(deploy.branchName) || buildSource.branch; + const repositoryName = normalizeOptionalString(deploy.repository?.fullName) || buildSource.source.repo; + const branchName = normalizeOptionalString(deploy.branchName) || buildSource.source.branch; let workDir: string | undefined; if (repositoryName && branchName) { @@ -185,9 +384,14 @@ export async function resolveAgentSessionPromptContext( return { name: serviceName, + status: normalizeOptionalString(deploy.status), + statusMessage: normalizeOptionalString(deploy.statusMessage), publicUrl: formatPublicUrl(deploy.publicUrl), repo: repositoryName, branch: branchName, + dockerImage: normalizeOptionalString(deploy.dockerImage), + buildPipelineId: normalizeOptionalString(deploy.buildPipelineId), + deployPipelineId: normalizeOptionalString(deploy.deployPipelineId), workDir, }; }) @@ -198,7 +402,11 @@ export async function resolveAgentSessionPromptContext( return { namespace: lookup.namespace, buildUuid: lookup.buildUuid, + gatheredAt: new Date().toISOString(), + build: buildSource.build, + pullRequest: buildSource.pullRequest, services, + diagnosticServices: buildSource.diagnosticServices, skillsAvailable: Boolean(session?.skillPlan?.skills?.length), }; } diff --git a/src/server/lib/createApiHandler.test.ts b/src/server/lib/createApiHandler.test.ts index 8f77aa15..5bd9c518 100644 --- a/src/server/lib/createApiHandler.test.ts +++ b/src/server/lib/createApiHandler.test.ts @@ -19,6 +19,16 @@ import { NextRequest, NextResponse } from 'next/server'; import { createApiHandler } from './createApiHandler'; describe('createApiHandler', () => { + const originalEnableAuth = process.env.ENABLE_AUTH; + + afterEach(() => { + if (originalEnableAuth === undefined) { + delete process.env.ENABLE_AUTH; + } else { + process.env.ENABLE_AUTH = originalEnableAuth; + } + }); + it('returns a standard error response for application errors', async () => { const handler = createApiHandler(async () => { throw new Error('sample failure'); @@ -52,4 +62,31 @@ describe('createApiHandler', () => { await expect(response.json()).resolves.toEqual({ ok: true }); expect(response.status).toBe(201); }); + + it('bypasses role checks when auth is disabled', async () => { + process.env.ENABLE_AUTH = 'false'; + + const handler = createApiHandler(async () => NextResponse.json({ ok: true }), { roles: ['admin'] }); + + const response = await handler(new NextRequest('http://localhost/api/v2/sample')); + + await expect(response.json()).resolves.toEqual({ ok: true }); + expect(response.status).toBe(200); + }); + + it('rejects role checks without a verified user when auth is enabled', async () => { + process.env.ENABLE_AUTH = 'true'; + + const handler = createApiHandler(async () => NextResponse.json({ ok: true }), { roles: ['admin'] }); + + const response = await handler(new NextRequest('http://localhost/api/v2/sample')); + + await expect(response.json()).resolves.toMatchObject({ + data: null, + error: { + message: 'Unauthorized', + }, + }); + expect(response.status).toBe(401); + }); }); diff --git a/src/server/lib/get-user.ts b/src/server/lib/get-user.ts index 78adf1e1..d0fa74e2 100644 --- a/src/server/lib/get-user.ts +++ b/src/server/lib/get-user.ts @@ -40,6 +40,18 @@ function normalizeClaim(value: unknown): string | null { return normalized || null; } +const REQUEST_USER_ROLES = ['user', 'admin'] as const; +type RequestUserRole = (typeof REQUEST_USER_ROLES)[number]; + +function normalizeRoles(value: unknown): RequestUserRole[] { + if (!Array.isArray(value)) { + return []; + } + + const allowed = new Set(REQUEST_USER_ROLES); + return value.filter((role): role is RequestUserRole => typeof role === 'string' && allowed.has(role)); +} + function getLocalDevUserId(): string { const configured = process.env.LOCAL_DEV_USER_ID?.trim(); return configured || 'local-dev-user'; @@ -63,6 +75,7 @@ export interface RequestUserIdentity { displayName: string; gitUserName: string; gitUserEmail: string; + roles: RequestUserRole[]; } function buildUserIdentity(payload: JWTPayload | null, userId: string): RequestUserIdentity { @@ -73,6 +86,10 @@ function buildUserIdentity(payload: JWTPayload | null, userId: string): RequestU const firstName = normalizeClaim(claims.given_name) || normalizeClaim(claims.firstName); const lastName = normalizeClaim(claims.family_name) || normalizeClaim(claims.lastName); const explicitName = normalizeClaim(claims.name); + const realmAccess = claims.realm_access as { roles?: unknown } | undefined; + const tokenRoles = normalizeRoles(realmAccess?.roles); + const roles: RequestUserRole[] = + tokenRoles.length > 0 ? tokenRoles : !payload && process.env.ENABLE_AUTH !== 'true' ? ['admin'] : []; const displayName = explicitName || [firstName, lastName].filter(Boolean).join(' ').trim() || @@ -92,6 +109,7 @@ function buildUserIdentity(payload: JWTPayload | null, userId: string): RequestU displayName, gitUserName, gitUserEmail, + roles, }; } diff --git a/src/server/lib/response.ts b/src/server/lib/response.ts index a8c7ed73..ce88b23b 100644 --- a/src/server/lib/response.ts +++ b/src/server/lib/response.ts @@ -26,7 +26,7 @@ interface Metadata { type SuccessStatusCode = 200 | 201; -type ErrorStatusCode = 400 | 401 | 404 | 409 | 500 | 502 | 503; +type ErrorStatusCode = 400 | 401 | 403 | 404 | 409 | 500 | 502 | 503; interface SuccessResponse { request_id: string; diff --git a/src/server/lib/roles.ts b/src/server/lib/roles.ts index b2054edd..c15d6a51 100644 --- a/src/server/lib/roles.ts +++ b/src/server/lib/roles.ts @@ -38,6 +38,10 @@ export function getUserRoles(payload: JWTPayload): LifecycleRole[] { export function requireRole(...roles: LifecycleRole[]) { return (handler: RouteHandler): RouteHandler => { return async (req: NextRequest, ...args: any[]) => { + if (process.env.ENABLE_AUTH !== 'true') { + return handler(req, ...args); + } + const user = getUser(req); if (!user) { diff --git a/src/server/lib/validation/__tests__/agentRuntimeConfigValidator.test.ts b/src/server/lib/validation/__tests__/agentRuntimeConfigValidator.test.ts new file mode 100644 index 00000000..03a9d37b --- /dev/null +++ b/src/server/lib/validation/__tests__/agentRuntimeConfigValidator.test.ts @@ -0,0 +1,181 @@ +/** + * 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 { + validateAgentRuntimeConfig, + validateAgentRuntimeRepoOverride, + AgentRuntimeConfigValidationError, +} from '../agentRuntimeConfigValidator'; +import type { AgentRuntimeConfig } from 'server/services/types/agentRuntimeConfig'; + +function makeConfig(): AgentRuntimeConfig { + return { + enabled: true, + providers: [ + { + name: 'anthropic', + enabled: true, + apiKeyEnvVar: 'ANTHROPIC_API_KEY', + models: [ + { + id: 'claude-sonnet-4-20250514', + displayName: 'Claude Sonnet 4', + enabled: true, + default: true, + maxTokens: 8192, + }, + ], + }, + ], + maxMessagesPerSession: 50, + sessionTTL: 3600, + }; +} + +describe('validateAgentRuntimeConfig', () => { + it('accepts a valid config', () => { + expect(() => validateAgentRuntimeConfig(makeConfig())).not.toThrow(); + }); + + it('rejects inline secrets in apiKeyEnvVar', () => { + const config = makeConfig(); + config.providers[0].apiKeyEnvVar = 'sk-ant-secret-value'; + + expect(() => validateAgentRuntimeConfig(config)).toThrow(AgentRuntimeConfigValidationError); + expect(() => validateAgentRuntimeConfig(config)).toThrow('apiKeyEnvVar must be an environment variable name'); + }); + + it('rejects duplicate providers', () => { + const config = makeConfig(); + config.providers.push({ + ...config.providers[0], + name: 'Anthropic', + }); + + expect(() => validateAgentRuntimeConfig(config)).toThrow('Duplicate provider "anthropic"'); + }); + + it('rejects enabled providers with no enabled models', () => { + const config = makeConfig(); + config.providers[0].models[0].enabled = false; + config.providers[0].models[0].default = false; + + expect(() => validateAgentRuntimeConfig(config)).toThrow('must have at least one enabled model'); + }); + + it('accepts valid capability availability policy', () => { + const config = makeConfig(); + config.capabilityPolicy = { + availability: { + workspace_shell: 'admin_only', + diagnostics_database: 'disabled', + }, + }; + + expect(() => validateAgentRuntimeConfig(config)).not.toThrow(); + }); + + it('rejects unknown capability ids', () => { + const config = makeConfig(); + config.capabilityPolicy = { + availability: { + sample_unknown: 'disabled', + } as any, + }; + + expect(() => validateAgentRuntimeConfig(config)).toThrow('Unknown capability id "sample_unknown".'); + }); + + it('rejects invalid capability availability values', () => { + const config = makeConfig(); + config.capabilityPolicy = { + availability: { + workspace_shell: 'sometimes', + } as any, + }; + + expect(() => validateAgentRuntimeConfig(config)).toThrow( + 'Capability "workspace_shell" has invalid availability "sometimes".' + ); + }); + + it('validates capability policy in repo overrides', () => { + expect(() => + validateAgentRuntimeRepoOverride({ + capabilityPolicy: { + availability: { + external_mcp_write: 'admin_only', + }, + }, + }) + ).not.toThrow(); + expect(() => + validateAgentRuntimeRepoOverride({ + capabilityPolicy: { + availability: { + sample_unknown: 'disabled', + } as any, + }, + }) + ).toThrow('Unknown capability id "sample_unknown".'); + }); + + it('accepts valid custom-agent creation policy', () => { + const config = makeConfig(); + config.customAgentCreationPolicy = { + mode: 'allowlist', + allowedUserIds: ['sample-user'], + allowedGithubUsernames: ['sample-gh-user'], + capabilityAvailability: { + external_mcp_write: 'reserved', + read_context: 'available', + }, + }; + + expect(() => validateAgentRuntimeConfig(config)).not.toThrow(); + }); + + it('rejects invalid custom-agent creation policy values', () => { + const config = makeConfig(); + config.customAgentCreationPolicy = { + mode: 'sometimes', + allowedUserIds: ['sample-user'], + } as any; + + expect(() => validateAgentRuntimeConfig(config)).toThrow('Invalid custom agent creation mode "sometimes".'); + + config.customAgentCreationPolicy = { + mode: 'enabled', + allowedGithubUsernames: ['sample-user'], + capabilityAvailability: { + sample_unknown: 'reserved', + }, + } as any; + + expect(() => validateAgentRuntimeConfig(config)).toThrow('Unknown creator capability id "sample_unknown".'); + + config.customAgentCreationPolicy = { + mode: 'enabled', + capabilityAvailability: { + read_context: 'maybe', + }, + } as any; + + expect(() => validateAgentRuntimeConfig(config)).toThrow( + 'Creator capability "read_context" has invalid availability "maybe".' + ); + }); +}); diff --git a/src/server/lib/validation/__tests__/aiAgentConfigValidator.test.ts b/src/server/lib/validation/__tests__/aiAgentConfigValidator.test.ts deleted file mode 100644 index 08544f85..00000000 --- a/src/server/lib/validation/__tests__/aiAgentConfigValidator.test.ts +++ /dev/null @@ -1,74 +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 { validateAIAgentConfig, AIAgentConfigValidationError } from '../aiAgentConfigValidator'; -import type { AIAgentConfig } from 'server/services/types/aiAgentConfig'; - -function makeConfig(): AIAgentConfig { - return { - enabled: true, - providers: [ - { - name: 'anthropic', - enabled: true, - apiKeyEnvVar: 'ANTHROPIC_API_KEY', - models: [ - { - id: 'claude-sonnet-4-20250514', - displayName: 'Claude Sonnet 4', - enabled: true, - default: true, - maxTokens: 8192, - }, - ], - }, - ], - maxMessagesPerSession: 50, - sessionTTL: 3600, - }; -} - -describe('validateAIAgentConfig', () => { - it('accepts a valid config', () => { - expect(() => validateAIAgentConfig(makeConfig())).not.toThrow(); - }); - - it('rejects inline secrets in apiKeyEnvVar', () => { - const config = makeConfig(); - config.providers[0].apiKeyEnvVar = 'sk-ant-secret-value'; - - expect(() => validateAIAgentConfig(config)).toThrow(AIAgentConfigValidationError); - expect(() => validateAIAgentConfig(config)).toThrow('apiKeyEnvVar must be an environment variable name'); - }); - - it('rejects duplicate providers', () => { - const config = makeConfig(); - config.providers.push({ - ...config.providers[0], - name: 'Anthropic', - }); - - expect(() => validateAIAgentConfig(config)).toThrow('Duplicate provider "anthropic"'); - }); - - it('rejects enabled providers with no enabled models', () => { - const config = makeConfig(); - config.providers[0].models[0].enabled = false; - config.providers[0].models[0].default = false; - - expect(() => validateAIAgentConfig(config)).toThrow('must have at least one enabled model'); - }); -}); diff --git a/src/server/lib/validation/aiAgentConfigSchemas.ts b/src/server/lib/validation/agentRuntimeConfigSchemas.ts similarity index 72% rename from src/server/lib/validation/aiAgentConfigSchemas.ts rename to src/server/lib/validation/agentRuntimeConfigSchemas.ts index ace0cc21..240e9b25 100644 --- a/src/server/lib/validation/aiAgentConfigSchemas.ts +++ b/src/server/lib/validation/agentRuntimeConfigSchemas.ts @@ -14,7 +14,54 @@ * limitations under the License. */ +import { AGENT_CAPABILITY_AVAILABILITIES, AGENT_CAPABILITY_CATALOG_IDS } from 'server/services/agent/capabilityCatalog'; + const approvalModeSchema = { type: 'string', enum: ['allow', 'require_approval', 'deny'] }; +const customAgentCreationModeSchema = { type: 'string', enum: ['enabled', 'disabled', 'admins_only', 'allowlist'] }; +const creatorCapabilityAvailabilitySchema = { type: 'string', enum: ['available', 'reserved'] }; + +export const capabilityPolicySchema = { + type: 'object', + properties: { + availability: { + type: 'object', + properties: Object.fromEntries( + AGENT_CAPABILITY_CATALOG_IDS.map((capabilityId) => [ + capabilityId, + { + type: 'string', + enum: [...AGENT_CAPABILITY_AVAILABILITIES], + }, + ]) + ), + additionalProperties: false, + }, + }, + additionalProperties: false, +}; + +export const customAgentCreationPolicySchema = { + type: 'object', + properties: { + mode: customAgentCreationModeSchema, + allowedUserIds: { + type: 'array', + items: { type: 'string' }, + }, + allowedGithubUsernames: { + type: 'array', + items: { type: 'string' }, + }, + capabilityAvailability: { + type: 'object', + properties: Object.fromEntries( + AGENT_CAPABILITY_CATALOG_IDS.map((capabilityId) => [capabilityId, creatorCapabilityAvailabilitySchema]) + ), + additionalProperties: false, + }, + }, + additionalProperties: false, +}; const approvalPolicyRulesSchema = { type: 'object', @@ -40,11 +87,13 @@ const approvalPolicySchema = { additionalProperties: false, }; -export const aiAgentConfigSchema = { +export const agentRuntimeConfigSchema = { type: 'object', properties: { enabled: { type: 'boolean' }, approvalPolicy: approvalPolicySchema, + capabilityPolicy: capabilityPolicySchema, + customAgentCreationPolicy: customAgentCreationPolicySchema, providers: { type: 'array', items: { @@ -94,13 +143,14 @@ export const aiAgentConfigSchema = { additionalProperties: false, }; -export const aiAgentRepoOverrideSchema = { +export const agentRuntimeRepoOverrideSchema = { type: 'object', properties: { enabled: { type: 'boolean' }, maxMessagesPerSession: { type: 'integer', minimum: 1 }, sessionTTL: { type: 'integer', minimum: 1 }, approvalPolicy: approvalPolicySchema, + capabilityPolicy: capabilityPolicySchema, additiveRules: { type: 'array', items: { type: 'string' } }, systemPromptOverride: { type: 'string', maxLength: 50000 }, excludedTools: { type: 'array', items: { type: 'string' } }, @@ -110,7 +160,7 @@ export const aiAgentRepoOverrideSchema = { additionalProperties: false, }; -export const aiAgentAdditiveRulesUpdateSchema = { +export const agentRuntimeAdditiveRulesUpdateSchema = { type: 'object', properties: { additiveRules: { type: 'array', items: { type: 'string' } }, @@ -119,7 +169,7 @@ export const aiAgentAdditiveRulesUpdateSchema = { additionalProperties: false, }; -export const aiAgentApprovalPolicyUpdateSchema = { +export const agentRuntimeApprovalPolicyUpdateSchema = { type: 'object', properties: { approvalPolicy: approvalPolicySchema, @@ -128,7 +178,7 @@ export const aiAgentApprovalPolicyUpdateSchema = { additionalProperties: false, }; -export const aiAgentConfigPatchSchema = { +export const agentRuntimeConfigPatchSchema = { type: 'object', properties: { additiveRules: { type: 'array', items: { type: 'string' } }, diff --git a/src/server/lib/validation/agentRuntimeConfigValidator.ts b/src/server/lib/validation/agentRuntimeConfigValidator.ts new file mode 100644 index 00000000..ba84f3cc --- /dev/null +++ b/src/server/lib/validation/agentRuntimeConfigValidator.ts @@ -0,0 +1,192 @@ +/** + * 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 { AgentRuntimeConfig, AgentRuntimeRepoOverride } from 'server/services/types/agentRuntimeConfig'; +import { isAgentCapabilityAvailability, isAgentCapabilityCatalogId } from 'server/services/agent/capabilityCatalog'; +import { + getProviderEnvVarCandidates, + isValidEnvVarName, + normalizeAgentProviderName, + type SupportedAgentProviderName, +} from 'server/services/agent/providerConfig'; +import { validateFileExclusionPatterns } from './filePatternValidator'; + +const CORE_TOOLS = ['query_database']; +const CUSTOM_AGENT_CREATION_MODES = ['enabled', 'disabled', 'admins_only', 'allowlist'] as const; +const CREATOR_CAPABILITY_AVAILABILITIES = ['available', 'reserved'] as const; + +type ModelConfig = NonNullable[number]['models'][number]; + +export class AgentRuntimeConfigValidationError extends Error {} + +function validateSharedConfigFields( + config: Pick< + AgentRuntimeConfig, + 'systemPromptOverride' | 'excludedTools' | 'excludedFilePatterns' | 'capabilityPolicy' | 'customAgentCreationPolicy' + > +): void { + if (config.systemPromptOverride !== undefined && config.systemPromptOverride.length > 50000) { + throw new AgentRuntimeConfigValidationError('systemPromptOverride exceeds maximum length of 50000 characters'); + } + + if (config.excludedTools && config.excludedTools.length > 0) { + for (const tool of config.excludedTools) { + if (CORE_TOOLS.includes(tool)) { + throw new AgentRuntimeConfigValidationError( + `Cannot exclude core tool: "${tool}". Core tools are required for agent operation.` + ); + } + } + } + + if (config.excludedFilePatterns && config.excludedFilePatterns.length > 0) { + validateFileExclusionPatterns(config.excludedFilePatterns); + } + + if (config.capabilityPolicy?.availability) { + for (const [capabilityId, availability] of Object.entries(config.capabilityPolicy.availability)) { + if (!isAgentCapabilityCatalogId(capabilityId)) { + throw new AgentRuntimeConfigValidationError(`Unknown capability id "${capabilityId}".`); + } + + if (!isAgentCapabilityAvailability(availability)) { + throw new AgentRuntimeConfigValidationError( + `Capability "${capabilityId}" has invalid availability "${availability}".` + ); + } + } + } + + const creationPolicy = config.customAgentCreationPolicy; + if (creationPolicy) { + if ( + creationPolicy.mode !== undefined && + !CUSTOM_AGENT_CREATION_MODES.includes(creationPolicy.mode as (typeof CUSTOM_AGENT_CREATION_MODES)[number]) + ) { + throw new AgentRuntimeConfigValidationError(`Invalid custom agent creation mode "${creationPolicy.mode}".`); + } + + for (const field of ['allowedUserIds', 'allowedGithubUsernames'] as const) { + const values = creationPolicy[field]; + if (values !== undefined && (!Array.isArray(values) || values.some((value) => typeof value !== 'string'))) { + throw new AgentRuntimeConfigValidationError(`customAgentCreationPolicy.${field} must be an array of strings.`); + } + } + + if ( + creationPolicy.capabilityAvailability !== undefined && + (!creationPolicy.capabilityAvailability || + typeof creationPolicy.capabilityAvailability !== 'object' || + Array.isArray(creationPolicy.capabilityAvailability)) + ) { + throw new AgentRuntimeConfigValidationError( + 'customAgentCreationPolicy.capabilityAvailability must be an object.' + ); + } + + if (creationPolicy.capabilityAvailability) { + for (const [capabilityId, availability] of Object.entries(creationPolicy.capabilityAvailability)) { + if (!isAgentCapabilityCatalogId(capabilityId)) { + throw new AgentRuntimeConfigValidationError(`Unknown creator capability id "${capabilityId}".`); + } + + if ( + !CREATOR_CAPABILITY_AVAILABILITIES.includes( + availability as (typeof CREATOR_CAPABILITY_AVAILABILITIES)[number] + ) + ) { + throw new AgentRuntimeConfigValidationError( + `Creator capability "${capabilityId}" has invalid availability "${availability}".` + ); + } + } + } + } +} + +function validateProviderModels( + providerName: SupportedAgentProviderName, + models: ModelConfig[], + providerEnabled: boolean +): void { + const seenModelIds = new Set(); + let defaultModelCount = 0; + let enabledModelCount = 0; + + for (const model of models) { + if (seenModelIds.has(model.id)) { + throw new AgentRuntimeConfigValidationError(`Provider "${providerName}" has duplicate model id "${model.id}".`); + } + + seenModelIds.add(model.id); + + if (model.enabled) { + enabledModelCount += 1; + } + + if (!model.default) { + continue; + } + + if (!model.enabled) { + throw new AgentRuntimeConfigValidationError( + `Provider "${providerName}" default model "${model.id}" must also be enabled.` + ); + } + + defaultModelCount += 1; + } + + if (providerEnabled && enabledModelCount === 0) { + throw new AgentRuntimeConfigValidationError(`Provider "${providerName}" must have at least one enabled model.`); + } + + if (defaultModelCount > 1) { + throw new AgentRuntimeConfigValidationError(`Provider "${providerName}" can have only one default model.`); + } +} + +export function validateAgentRuntimeConfig(config: AgentRuntimeConfig): void { + validateSharedConfigFields(config); + + const seenProviders = new Set(); + + for (const provider of config.providers || []) { + const providerName = normalizeAgentProviderName(provider.name); + if (!providerName) { + throw new AgentRuntimeConfigValidationError(`Unsupported provider "${provider.name}".`); + } + + if (seenProviders.has(providerName)) { + throw new AgentRuntimeConfigValidationError(`Duplicate provider "${providerName}" is not allowed.`); + } + + seenProviders.add(providerName); + + if (!isValidEnvVarName(provider.apiKeyEnvVar)) { + const exampleEnvVar = getProviderEnvVarCandidates(providerName)[0] || 'API_KEY_ENV_VAR'; + throw new AgentRuntimeConfigValidationError( + `Provider "${providerName}" apiKeyEnvVar must be an environment variable name like ${exampleEnvVar}.` + ); + } + + validateProviderModels(providerName, provider.models || [], provider.enabled !== false); + } +} + +export function validateAgentRuntimeRepoOverride(config: Partial): void { + validateSharedConfigFields(config); +} diff --git a/src/server/lib/validation/aiAgentConfigValidator.ts b/src/server/lib/validation/aiAgentConfigValidator.ts deleted file mode 100644 index 955392c8..00000000 --- a/src/server/lib/validation/aiAgentConfigValidator.ts +++ /dev/null @@ -1,126 +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 { AIAgentConfig, AIAgentRepoOverride } from 'server/services/types/aiAgentConfig'; -import { - getProviderEnvVarCandidates, - isValidEnvVarName, - normalizeAgentProviderName, - type SupportedAgentProviderName, -} from 'server/services/agent/providerConfig'; -import { validateFileExclusionPatterns } from './filePatternValidator'; - -const CORE_TOOLS = ['query_database']; - -type ModelConfig = NonNullable[number]['models'][number]; - -export class AIAgentConfigValidationError extends Error {} - -function validateSharedConfigFields( - config: Pick -): void { - if (config.systemPromptOverride !== undefined && config.systemPromptOverride.length > 50000) { - throw new AIAgentConfigValidationError('systemPromptOverride exceeds maximum length of 50000 characters'); - } - - if (config.excludedTools && config.excludedTools.length > 0) { - for (const tool of config.excludedTools) { - if (CORE_TOOLS.includes(tool)) { - throw new AIAgentConfigValidationError( - `Cannot exclude core tool: "${tool}". Core tools are required for agent operation.` - ); - } - } - } - - if (config.excludedFilePatterns && config.excludedFilePatterns.length > 0) { - validateFileExclusionPatterns(config.excludedFilePatterns); - } -} - -function validateProviderModels( - providerName: SupportedAgentProviderName, - models: ModelConfig[], - providerEnabled: boolean -): void { - const seenModelIds = new Set(); - let defaultModelCount = 0; - let enabledModelCount = 0; - - for (const model of models) { - if (seenModelIds.has(model.id)) { - throw new AIAgentConfigValidationError(`Provider "${providerName}" has duplicate model id "${model.id}".`); - } - - seenModelIds.add(model.id); - - if (model.enabled) { - enabledModelCount += 1; - } - - if (!model.default) { - continue; - } - - if (!model.enabled) { - throw new AIAgentConfigValidationError( - `Provider "${providerName}" default model "${model.id}" must also be enabled.` - ); - } - - defaultModelCount += 1; - } - - if (providerEnabled && enabledModelCount === 0) { - throw new AIAgentConfigValidationError(`Provider "${providerName}" must have at least one enabled model.`); - } - - if (defaultModelCount > 1) { - throw new AIAgentConfigValidationError(`Provider "${providerName}" can have only one default model.`); - } -} - -export function validateAIAgentConfig(config: AIAgentConfig): void { - validateSharedConfigFields(config); - - const seenProviders = new Set(); - - for (const provider of config.providers || []) { - const providerName = normalizeAgentProviderName(provider.name); - if (!providerName) { - throw new AIAgentConfigValidationError(`Unsupported provider "${provider.name}".`); - } - - if (seenProviders.has(providerName)) { - throw new AIAgentConfigValidationError(`Duplicate provider "${providerName}" is not allowed.`); - } - - seenProviders.add(providerName); - - if (!isValidEnvVarName(provider.apiKeyEnvVar)) { - const exampleEnvVar = getProviderEnvVarCandidates(providerName)[0] || 'API_KEY_ENV_VAR'; - throw new AIAgentConfigValidationError( - `Provider "${providerName}" apiKeyEnvVar must be an environment variable name like ${exampleEnvVar}.` - ); - } - - validateProviderModels(providerName, provider.models || [], provider.enabled !== false); - } -} - -export function validateAIAgentRepoOverride(config: Partial): void { - validateSharedConfigFields(config); -} diff --git a/src/server/models/AgentDefinition.ts b/src/server/models/AgentDefinition.ts new file mode 100644 index 00000000..bd59d5d2 --- /dev/null +++ b/src/server/models/AgentDefinition.ts @@ -0,0 +1,93 @@ +/** + * 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'; +import type { AgentCapabilityCatalogId } from 'server/services/agent/capabilityCatalog'; +import type { + AgentDefinitionModelPreference, + AgentDefinitionOwnerKind, + AgentDefinitionResourcePolicy, + AgentDefinitionStatus, +} from 'server/services/agent/agentDefinitionTypes'; + +export default class AgentDefinition extends Model { + definitionId!: string; + version!: number; + ownerKind!: AgentDefinitionOwnerKind; + ownerUserId!: string | null; + ownerOrganizationId!: string | null; + name!: string; + description!: string | null; + instructionRefs!: string[]; + instructionAddendum!: string | null; + capabilityRefs!: AgentCapabilityCatalogId[]; + requiredCapabilityRefs!: AgentCapabilityCatalogId[]; + optionalCapabilityRefs!: AgentCapabilityCatalogId[]; + resourcePolicy!: AgentDefinitionResourcePolicy; + modelPreference!: AgentDefinitionModelPreference | null; + status!: AgentDefinitionStatus; + codeOwned!: boolean; + readOnly!: boolean; + + static tableName = 'agent_definitions'; + static timestamps = true; + static idColumn = 'id'; + + static jsonSchema = { + type: 'object', + required: ['definitionId', 'version', 'ownerKind', 'name', 'instructionRefs', 'capabilityRefs', 'resourcePolicy'], + properties: { + id: { type: 'integer' }, + definitionId: { type: 'string', minLength: 1 }, + version: { type: 'integer', minimum: 1, default: 1 }, + ownerKind: { type: 'string', enum: ['system', 'admin', 'user'] }, + ownerUserId: { type: ['string', 'null'] }, + ownerOrganizationId: { type: ['string', 'null'] }, + name: { type: 'string', minLength: 1 }, + description: { type: ['string', 'null'] }, + instructionRefs: { type: 'array', items: { type: 'string' }, default: [] }, + instructionAddendum: { type: ['string', 'null'] }, + capabilityRefs: { type: 'array', items: { type: 'string' }, default: [] }, + requiredCapabilityRefs: { type: 'array', items: { type: 'string' }, default: [] }, + optionalCapabilityRefs: { type: 'array', items: { type: 'string' }, default: [] }, + resourcePolicy: { + type: 'object', + required: ['sourceKinds'], + properties: { + sourceKinds: { type: 'array', items: { type: 'string' }, default: [] }, + sandboxRequired: { type: 'boolean' }, + workspaceRequired: { type: 'boolean' }, + }, + default: { sourceKinds: [] }, + }, + modelPreference: { type: ['object', 'null'], default: null }, + status: { type: 'string', enum: ['active', 'disabled', 'archived'], default: 'active' }, + codeOwned: { type: 'boolean', default: false }, + readOnly: { type: 'boolean', default: false }, + }, + }; + + static get jsonAttributes() { + return [ + 'instructionRefs', + 'capabilityRefs', + 'requiredCapabilityRefs', + 'optionalCapabilityRefs', + 'resourcePolicy', + 'modelPreference', + ]; + } +} diff --git a/src/server/models/AgentRun.ts b/src/server/models/AgentRun.ts index 2d417be1..39d2573f 100644 --- a/src/server/models/AgentRun.ts +++ b/src/server/models/AgentRun.ts @@ -48,6 +48,7 @@ export default class AgentRun extends Model { cancelledAt!: string | null; usageSummary!: Record; policySnapshot!: Record; + runPlanSnapshot!: Record | null; error!: Record | null; static tableName = 'agent_runs'; @@ -97,12 +98,13 @@ export default class AgentRun extends Model { cancelledAt: { type: ['string', 'null'] }, usageSummary: { type: 'object', default: {} }, policySnapshot: { type: 'object', default: {} }, + runPlanSnapshot: { type: ['object', 'null'], default: null }, error: { type: ['object', 'null'], default: null }, }, }; static get jsonAttributes() { - return ['sandboxRequirement', 'usageSummary', 'policySnapshot', 'error']; + return ['sandboxRequirement', 'usageSummary', 'policySnapshot', 'runPlanSnapshot', 'error']; } static get relationMappings() { diff --git a/src/server/models/AIAgentRepoConfig.ts b/src/server/models/AgentRuntimeRepoConfig.ts similarity index 86% rename from src/server/models/AIAgentRepoConfig.ts rename to src/server/models/AgentRuntimeRepoConfig.ts index e269c477..e8fa287e 100644 --- a/src/server/models/AIAgentRepoConfig.ts +++ b/src/server/models/AgentRuntimeRepoConfig.ts @@ -16,10 +16,10 @@ import Model from './_Model'; -export default class AIAgentRepoConfig extends Model { +export default class AgentRuntimeRepoConfig extends Model { repositoryFullName: string; config: Record; - static tableName = 'ai_agent_repo_config'; + static tableName = 'agent_runtime_repo_config'; static timestamps = true; } diff --git a/src/server/models/Conversation.ts b/src/server/models/Conversation.ts deleted file mode 100644 index 1d977077..00000000 --- a/src/server/models/Conversation.ts +++ /dev/null @@ -1,47 +0,0 @@ -/** - * Copyright 2025 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 Conversation extends Model { - buildUuid!: string; - repo!: string; - model?: string | null; - messageCount!: number; - metadata!: Record; - - static tableName = 'conversations'; - static timestamps = true; - static idColumn = 'buildUuid'; - - static get jsonAttributes() { - return ['metadata']; - } - - static get relationMappings() { - const ConversationMessage = require('./ConversationMessage').default; - return { - messages: { - relation: Model.HasManyRelation, - modelClass: ConversationMessage, - join: { - from: 'conversations.buildUuid', - to: 'conversation_messages.buildUuid', - }, - }, - }; - } -} diff --git a/src/server/models/ConversationFeedback.ts b/src/server/models/ConversationFeedback.ts deleted file mode 100644 index 20ade7bf..00000000 --- a/src/server/models/ConversationFeedback.ts +++ /dev/null @@ -1,43 +0,0 @@ -/** - * Copyright 2025 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 ConversationFeedback extends Model { - buildUuid!: string; - rating!: 'up' | 'down'; - text?: string | null; - userIdentifier?: string | null; - repo!: string; - prNumber?: number | null; - - static tableName = 'conversation_feedback'; - static timestamps = true; - - static get relationMappings() { - const Conversation = require('./Conversation').default; - return { - conversation: { - relation: Model.BelongsToOneRelation, - modelClass: Conversation, - join: { - from: 'conversation_feedback.buildUuid', - to: 'conversations.buildUuid', - }, - }, - }; - } -} diff --git a/src/server/models/ConversationMessage.ts b/src/server/models/ConversationMessage.ts deleted file mode 100644 index 88698525..00000000 --- a/src/server/models/ConversationMessage.ts +++ /dev/null @@ -1,46 +0,0 @@ -/** - * Copyright 2025 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 ConversationMessage extends Model { - buildUuid!: string; - role!: 'user' | 'assistant' | 'system'; - content!: string; - timestamp!: number; - metadata!: Record; - - static tableName = 'conversation_messages'; - static timestamps = true; - - static get jsonAttributes() { - return ['metadata']; - } - - static get relationMappings() { - const Conversation = require('./Conversation').default; - return { - conversation: { - relation: Model.BelongsToOneRelation, - modelClass: Conversation, - join: { - from: 'conversation_messages.buildUuid', - to: 'conversations.buildUuid', - }, - }, - }; - } -} diff --git a/src/server/models/McpServerConfig.ts b/src/server/models/McpServerConfig.ts index 61a5f055..86723a30 100644 --- a/src/server/models/McpServerConfig.ts +++ b/src/server/models/McpServerConfig.ts @@ -20,7 +20,7 @@ import type { McpDiscoveredTool, McpSharedConnectionConfig, McpTransportConfig, -} from 'server/services/ai/mcp/types'; +} from 'server/services/agentRuntime/mcp/types'; export default class McpServerConfig extends Model { slug!: string; diff --git a/src/server/models/MessageFeedback.ts b/src/server/models/MessageFeedback.ts deleted file mode 100644 index 608530f3..00000000 --- a/src/server/models/MessageFeedback.ts +++ /dev/null @@ -1,53 +0,0 @@ -/** - * Copyright 2025 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 MessageFeedback extends Model { - buildUuid!: string; - messageId!: number; - rating!: 'up' | 'down'; - text?: string | null; - userIdentifier?: string | null; - repo!: string; - prNumber?: number | null; - - static tableName = 'message_feedback'; - static timestamps = true; - - static get relationMappings() { - const Conversation = require('./Conversation').default; - const ConversationMessage = require('./ConversationMessage').default; - return { - conversation: { - relation: Model.BelongsToOneRelation, - modelClass: Conversation, - join: { - from: 'message_feedback.buildUuid', - to: 'conversations.buildUuid', - }, - }, - message: { - relation: Model.BelongsToOneRelation, - modelClass: ConversationMessage, - join: { - from: 'message_feedback.messageId', - to: 'conversation_messages.id', - }, - }, - }; - } -} diff --git a/src/server/models/UserMcpConnection.ts b/src/server/models/UserMcpConnection.ts index 0716a435..5bef62cc 100644 --- a/src/server/models/UserMcpConnection.ts +++ b/src/server/models/UserMcpConnection.ts @@ -15,7 +15,7 @@ */ import Model from './_Model'; -import type { McpDiscoveredTool } from 'server/services/ai/mcp/types'; +import type { McpDiscoveredTool } from 'server/services/agentRuntime/mcp/types'; export default class UserMcpConnection extends Model { userId!: string; diff --git a/src/server/models/__tests__/AgentModelsValidation.test.ts b/src/server/models/__tests__/AgentModelsValidation.test.ts index b2322190..87d2e61c 100644 --- a/src/server/models/__tests__/AgentModelsValidation.test.ts +++ b/src/server/models/__tests__/AgentModelsValidation.test.ts @@ -15,6 +15,8 @@ */ import AgentMessage from 'server/models/AgentMessage'; +import AgentDefinition from 'server/models/AgentDefinition'; +import AgentRun from 'server/models/AgentRun'; import AgentThread from 'server/models/AgentThread'; describe('Agent model validation', () => { @@ -37,4 +39,85 @@ describe('Agent model validation', () => { }) ).not.toThrow(); }); + + test('allows agent runs with a null run plan snapshot', () => { + expect(() => + AgentRun.fromJson({ + threadId: 42, + sessionId: 7, + provider: 'sample-provider', + model: 'sample-model', + runPlanSnapshot: null, + }) + ).not.toThrow(); + }); + + test('allows agent runs with an object run plan snapshot', () => { + expect(() => + AgentRun.fromJson({ + threadId: 42, + sessionId: 7, + provider: 'sample-provider', + model: 'sample-model', + runPlanSnapshot: { + version: 1, + agent: { + id: 'system.freeform', + }, + }, + }) + ).not.toThrow(); + }); + + test('allows agent definitions with first-party preset JSON fields', () => { + expect(() => + AgentDefinition.fromJson({ + definitionId: 'system.freeform', + version: 1, + ownerKind: 'system', + name: 'Free-form', + instructionRefs: ['system:freeform'], + capabilityRefs: ['read_context'], + requiredCapabilityRefs: ['read_context'], + optionalCapabilityRefs: [], + resourcePolicy: { + sourceKinds: ['freeform_chat'], + workspaceRequired: false, + sandboxRequired: false, + }, + codeOwned: true, + readOnly: true, + status: 'active', + }) + ).not.toThrow(); + }); + + test('allows mutable user agent definitions to transition to archived', () => { + const userDefinition = { + definitionId: 'custom.sample-definition', + version: 1, + ownerKind: 'user', + ownerUserId: 'sample-user', + ownerOrganizationId: null, + name: 'Sample agent', + description: null, + instructionRefs: [], + instructionAddendum: 'Use a concise response style.', + capabilityRefs: ['read_context'], + requiredCapabilityRefs: [], + optionalCapabilityRefs: ['read_context'], + resourcePolicy: { + sourceKinds: ['freeform_chat'], + workspaceRequired: false, + sandboxRequired: false, + }, + modelPreference: null, + codeOwned: false, + readOnly: false, + status: 'active', + }; + + expect(() => AgentDefinition.fromJson(userDefinition)).not.toThrow(); + expect(() => AgentDefinition.fromJson({ ...userDefinition, status: 'archived' })).not.toThrow(); + }); }); diff --git a/src/server/models/index.ts b/src/server/models/index.ts index cea619af..2a831c7e 100644 --- a/src/server/models/index.ts +++ b/src/server/models/index.ts @@ -27,11 +27,7 @@ import Deployable from './Deployable'; import BotUser from './BotUser'; import GlobalConfig from './GlobalConfig'; import WebhookInvocations from './WebhookInvocations'; -import AIAgentRepoConfig from './AIAgentRepoConfig'; -import Conversation from './Conversation'; -import ConversationMessage from './ConversationMessage'; -import MessageFeedback from './MessageFeedback'; -import ConversationFeedback from './ConversationFeedback'; +import AgentRuntimeRepoConfig from './AgentRuntimeRepoConfig'; import AgentSession from './AgentSession'; import AgentSource from './AgentSource'; import AgentSandbox from './AgentSandbox'; @@ -60,11 +56,7 @@ export interface IModels { BotUser: typeof BotUser; GlobalConfig: typeof GlobalConfig; WebhookInvocations: typeof WebhookInvocations; - AIAgentRepoConfig: typeof AIAgentRepoConfig; - Conversation: typeof Conversation; - ConversationMessage: typeof ConversationMessage; - MessageFeedback: typeof MessageFeedback; - ConversationFeedback: typeof ConversationFeedback; + AgentRuntimeRepoConfig: typeof AgentRuntimeRepoConfig; AgentSession: typeof AgentSession; AgentSource: typeof AgentSource; AgentSandbox: typeof AgentSandbox; @@ -94,11 +86,7 @@ export { BotUser, GlobalConfig, WebhookInvocations, - AIAgentRepoConfig, - Conversation, - ConversationMessage, - MessageFeedback, - ConversationFeedback, + AgentRuntimeRepoConfig, AgentSession, AgentSource, AgentSandbox, diff --git a/src/server/services/__tests__/agentRuntimeConfig.test.ts b/src/server/services/__tests__/agentRuntimeConfig.test.ts new file mode 100644 index 00000000..4bf9bda3 --- /dev/null +++ b/src/server/services/__tests__/agentRuntimeConfig.test.ts @@ -0,0 +1,520 @@ +/** + * 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/globalConfig', () => ({ + __esModule: true, + default: { + getInstance: jest.fn(), + }, +})); + +jest.mock('server/lib/logger', () => ({ + getLogger: jest.fn(() => ({ + info: jest.fn(), + warn: jest.fn(), + error: jest.fn(), + debug: jest.fn(), + })), +})); + +import AgentRuntimeConfigService from 'server/services/agentRuntime/config/agentRuntimeConfig'; +import GlobalConfigService from 'server/services/globalConfig'; +import type { AgentRuntimeConfig } from 'server/services/types/agentRuntimeConfig'; + +function makeService() { + const repoUpsertQuery = { + insert: jest.fn().mockReturnThis(), + onConflict: jest.fn().mockReturnThis(), + merge: jest.fn().mockResolvedValue(undefined), + }; + + const knex = Object.assign(jest.fn().mockReturnValue(repoUpsertQuery), { + fn: { + now: jest.fn(() => 'now'), + }, + }); + + const db = { knex } as any; + const redis = { del: jest.fn().mockResolvedValue(undefined) } as any; + const redlock = {} as any; + const queueManager = {} as any; + + return { + service: new AgentRuntimeConfigService(db, redis, redlock, queueManager), + knex, + repoUpsertQuery, + redis, + }; +} + +describe('AgentRuntimeConfigService', () => { + beforeEach(() => { + jest.clearAllMocks(); + }); + + it('updates global additive rules without revalidating unrelated provider defaults', async () => { + const { service } = makeService(); + const currentConfig: AgentRuntimeConfig = { + enabled: true, + providers: [ + { + name: 'gemini', + enabled: true, + apiKeyEnvVar: 'GOOGLE_GENERATIVE_AI_API_KEY', + models: [ + { + id: 'gemini-1', + displayName: 'Gemini 1', + enabled: true, + default: true, + maxTokens: 8192, + }, + { + id: 'gemini-2', + displayName: 'Gemini 2', + enabled: true, + default: true, + maxTokens: 8192, + }, + ], + }, + ], + maxMessagesPerSession: 50, + sessionTTL: 3600, + }; + + const setConfig = jest.fn().mockResolvedValue(undefined); + (GlobalConfigService.getInstance as jest.Mock).mockReturnValue({ + getConfig: jest.fn().mockResolvedValue(currentConfig), + setConfig, + }); + + const result = await service.updateGlobalAdditiveRules(['test']); + + expect(setConfig).toHaveBeenCalledWith( + 'agentRuntime', + expect.objectContaining({ + providers: currentConfig.providers, + additiveRules: ['test'], + }) + ); + expect(result.additiveRules).toEqual(['test']); + expect(result.providers).toEqual(currentConfig.providers); + }); + + it('replaces global approval policy without revalidating unrelated provider defaults', async () => { + const { service } = makeService(); + const currentConfig: AgentRuntimeConfig = { + enabled: true, + providers: [ + { + name: 'gemini', + enabled: true, + apiKeyEnvVar: 'GOOGLE_GENERATIVE_AI_API_KEY', + models: [ + { + id: 'gemini-1', + displayName: 'Gemini 1', + enabled: true, + default: true, + maxTokens: 8192, + }, + { + id: 'gemini-2', + displayName: 'Gemini 2', + enabled: true, + default: true, + maxTokens: 8192, + }, + ], + }, + ], + maxMessagesPerSession: 50, + sessionTTL: 3600, + approvalPolicy: { + defaultMode: 'require_approval', + }, + }; + + const setConfig = jest.fn().mockResolvedValue(undefined); + (GlobalConfigService.getInstance as jest.Mock).mockReturnValue({ + getConfig: jest.fn().mockResolvedValue(currentConfig), + setConfig, + }); + + const result = await service.updateGlobalApprovalPolicy({ + defaultMode: 'require_approval', + rules: { + shell_exec: 'deny', + }, + }); + + expect(setConfig).toHaveBeenCalledWith( + 'agentRuntime', + expect.objectContaining({ + providers: currentConfig.providers, + approvalPolicy: { + defaultMode: 'require_approval', + rules: { + shell_exec: 'deny', + }, + }, + }) + ); + expect(result.approvalPolicy).toEqual({ + defaultMode: 'require_approval', + rules: { + shell_exec: 'deny', + }, + }); + expect(result.providers).toEqual(currentConfig.providers); + }); + + it('clears the global approval policy when given an empty replacement', async () => { + const { service } = makeService(); + const currentConfig: AgentRuntimeConfig = { + enabled: true, + providers: [ + { + name: 'openai', + enabled: true, + apiKeyEnvVar: 'OPENAI_API_KEY', + models: [ + { + id: 'gpt-5', + displayName: 'GPT-5', + enabled: true, + default: true, + maxTokens: 8192, + }, + ], + }, + ], + maxMessagesPerSession: 50, + sessionTTL: 3600, + approvalPolicy: { + defaultMode: 'deny', + rules: { + shell_exec: 'deny', + }, + }, + }; + + const setConfig = jest.fn().mockResolvedValue(undefined); + (GlobalConfigService.getInstance as jest.Mock).mockReturnValue({ + getConfig: jest.fn().mockResolvedValue(currentConfig), + setConfig, + }); + + const result = await service.updateGlobalApprovalPolicy({}); + + expect(setConfig).toHaveBeenCalledWith( + 'agentRuntime', + expect.objectContaining({ + providers: currentConfig.providers, + }) + ); + expect(setConfig.mock.calls[0]?.[1]).not.toHaveProperty('approvalPolicy'); + expect(result.approvalPolicy).toBeUndefined(); + expect(result.providers).toEqual(currentConfig.providers); + }); + + it('replaces global capability policy without revalidating unrelated provider defaults', async () => { + const { service } = makeService(); + const currentConfig: AgentRuntimeConfig = { + enabled: true, + providers: [ + { + name: 'gemini', + enabled: true, + apiKeyEnvVar: 'GOOGLE_GENERATIVE_AI_API_KEY', + models: [ + { + id: 'gemini-1', + displayName: 'Gemini 1', + enabled: true, + default: true, + maxTokens: 8192, + }, + { + id: 'gemini-2', + displayName: 'Gemini 2', + enabled: true, + default: true, + maxTokens: 8192, + }, + ], + }, + ], + maxMessagesPerSession: 50, + sessionTTL: 3600, + capabilityPolicy: { + availability: { + workspace_shell: 'admin_only', + }, + }, + }; + + const setConfig = jest.fn().mockResolvedValue(undefined); + (GlobalConfigService.getInstance as jest.Mock).mockReturnValue({ + getConfig: jest.fn().mockResolvedValue(currentConfig), + setConfig, + }); + + const result = await service.updateGlobalCapabilityPolicy({ + availability: { + diagnostics_database: 'disabled', + }, + }); + + expect(setConfig).toHaveBeenCalledWith( + 'agentRuntime', + expect.objectContaining({ + providers: currentConfig.providers, + capabilityPolicy: { + availability: { + diagnostics_database: 'disabled', + }, + }, + }) + ); + expect(result.capabilityPolicy).toEqual({ + availability: { + diagnostics_database: 'disabled', + }, + }); + expect(result.providers).toEqual(currentConfig.providers); + }); + + it('clears the global capability policy when given an empty replacement', async () => { + const { service } = makeService(); + const currentConfig: AgentRuntimeConfig = { + enabled: true, + providers: [ + { + name: 'openai', + enabled: true, + apiKeyEnvVar: 'OPENAI_API_KEY', + models: [ + { + id: 'gpt-5', + displayName: 'GPT-5', + enabled: true, + default: true, + maxTokens: 8192, + }, + ], + }, + ], + maxMessagesPerSession: 50, + sessionTTL: 3600, + capabilityPolicy: { + availability: { + workspace_shell: 'disabled', + }, + }, + }; + + const setConfig = jest.fn().mockResolvedValue(undefined); + (GlobalConfigService.getInstance as jest.Mock).mockReturnValue({ + getConfig: jest.fn().mockResolvedValue(currentConfig), + setConfig, + }); + + const result = await service.updateGlobalCapabilityPolicy({}); + + expect(setConfig).toHaveBeenCalledWith( + 'agentRuntime', + expect.objectContaining({ + providers: currentConfig.providers, + }) + ); + expect(setConfig.mock.calls[0]?.[1]).not.toHaveProperty('capabilityPolicy'); + expect(result.capabilityPolicy).toBeUndefined(); + expect(result.providers).toEqual(currentConfig.providers); + }); + + it('replaces custom-agent creation policy without revalidating unrelated provider defaults', async () => { + const { service } = makeService(); + const currentConfig: AgentRuntimeConfig = { + enabled: true, + providers: [ + { + name: 'gemini', + enabled: true, + apiKeyEnvVar: 'GOOGLE_GENERATIVE_AI_API_KEY', + models: [ + { + id: 'gemini-1', + displayName: 'Gemini 1', + enabled: true, + default: true, + maxTokens: 8192, + }, + { + id: 'gemini-2', + displayName: 'Gemini 2', + enabled: true, + default: true, + maxTokens: 8192, + }, + ], + }, + ], + maxMessagesPerSession: 50, + sessionTTL: 3600, + }; + + const setConfig = jest.fn().mockResolvedValue(undefined); + (GlobalConfigService.getInstance as jest.Mock).mockReturnValue({ + getConfig: jest.fn().mockResolvedValue(currentConfig), + setConfig, + }); + + const result = await service.updateGlobalCustomAgentCreationPolicy({ + mode: 'disabled', + }); + + expect(setConfig).toHaveBeenCalledWith( + 'agentRuntime', + expect.objectContaining({ + providers: currentConfig.providers, + customAgentCreationPolicy: { + mode: 'disabled', + }, + }) + ); + expect(result.customAgentCreationPolicy).toEqual({ + mode: 'disabled', + }); + expect(result.providers).toEqual(currentConfig.providers); + }); + + it('merges repo capability policy over global policy key by key', () => { + const { service } = makeService(); + + const result = (service as any).mergeConfigs( + { + enabled: true, + providers: [], + maxMessagesPerSession: 50, + sessionTTL: 3600, + capabilityPolicy: { + availability: { + workspace_shell: 'admin_only', + diagnostics_database: 'disabled', + }, + }, + }, + { + capabilityPolicy: { + availability: { + workspace_shell: 'all_users', + }, + }, + } + ); + + expect(result.capabilityPolicy).toEqual({ + availability: { + workspace_shell: 'all_users', + diagnostics_database: 'disabled', + }, + }); + }); + + it('updates repo additive rules while preserving other repo overrides', async () => { + const { service, knex, repoUpsertQuery, redis } = makeService(); + + jest.spyOn(service, 'getRepoConfig').mockResolvedValue({ + excludedTools: ['tool-a'], + approvalPolicy: { + defaultMode: 'require_approval', + }, + }); + + const result = await service.updateRepoAdditiveRules('Example-Org/Example-Repo', ['test']); + + expect(knex).toHaveBeenCalledWith('agent_runtime_repo_config'); + expect(repoUpsertQuery.insert).toHaveBeenCalledWith( + expect.objectContaining({ + repositoryFullName: 'example-org/example-repo', + config: JSON.stringify({ + excludedTools: ['tool-a'], + approvalPolicy: { + defaultMode: 'require_approval', + }, + additiveRules: ['test'], + }), + }) + ); + expect(redis.del).toHaveBeenCalledWith('agent_runtime_repo_config:example-org/example-repo'); + expect(result).toEqual({ + excludedTools: ['tool-a'], + approvalPolicy: { + defaultMode: 'require_approval', + }, + additiveRules: ['test'], + }); + }); + + it('updates repo capability policy while preserving other repo overrides', async () => { + const { service, knex, repoUpsertQuery, redis } = makeService(); + + jest.spyOn(service, 'getRepoConfig').mockResolvedValue({ + excludedTools: ['tool-a'], + approvalPolicy: { + defaultMode: 'require_approval', + }, + }); + + const result = await service.updateRepoCapabilityPolicy('Example-Org/Example-Repo', { + availability: { + workspace_shell: 'admin_only', + }, + }); + + expect(knex).toHaveBeenCalledWith('agent_runtime_repo_config'); + expect(repoUpsertQuery.insert).toHaveBeenCalledWith( + expect.objectContaining({ + repositoryFullName: 'example-org/example-repo', + config: JSON.stringify({ + excludedTools: ['tool-a'], + approvalPolicy: { + defaultMode: 'require_approval', + }, + capabilityPolicy: { + availability: { + workspace_shell: 'admin_only', + }, + }, + }), + }) + ); + expect(redis.del).toHaveBeenCalledWith('agent_runtime_repo_config:example-org/example-repo'); + expect(result).toEqual({ + excludedTools: ['tool-a'], + approvalPolicy: { + defaultMode: 'require_approval', + }, + capabilityPolicy: { + availability: { + workspace_shell: 'admin_only', + }, + }, + }); + }); +}); diff --git a/src/server/services/__tests__/agentSession.test.ts b/src/server/services/__tests__/agentSession.test.ts index 8d15c325..d680e67a 100644 --- a/src/server/services/__tests__/agentSession.test.ts +++ b/src/server/services/__tests__/agentSession.test.ts @@ -59,7 +59,7 @@ jest.mock('server/lib/kubernetes', () => ({ deleteNamespace: (...args: unknown[]) => mockDeleteNamespace(...args), })); jest.mock('server/lib/kubernetes/networkPolicyFactory'); -jest.mock('server/services/ai/mcp/config', () => ({ +jest.mock('server/services/agentRuntime/mcp/config', () => ({ __esModule: true, McpConfigService: jest.fn().mockImplementation(() => ({ resolveSessionPodServersForRepo: mockResolveSessionPodServersForRepo, @@ -112,7 +112,7 @@ jest.mock('server/services/agentPrewarm', () => ({ jest.mock('server/services/agent/ThreadService', () => ({ __esModule: true, default: { - getDefaultThreadForSession: mockGetDefaultThreadForSession, + getDefaultThreadForSession: (...args: unknown[]) => mockGetDefaultThreadForSession(...args), }, })); jest.mock('server/lib/nativeHelm/helm', () => ({ @@ -186,7 +186,7 @@ jest.mock('server/services/globalConfig', () => ({ default: { getInstance: jest.fn(() => ({ getConfig: jest.fn().mockImplementation(async (key: string) => { - if (key === 'aiAgent') { + if (key === 'agentRuntime') { return { enabled: true, providers: [ @@ -828,6 +828,237 @@ describe('AgentSessionService', () => { expect(session.namespace).toBe('chat-aaaaaaaa'); }); + it('preserves build-context repo metadata when chat runtime provisioning fails', async () => { + const workspaceRepos = [ + { + repo: 'example-org/example-repo', + repoUrl: 'https://github.com/example-org/example-repo.git', + branch: 'feature/sample', + revision: 'commit-sha-1', + mountPath: '/workspace', + primary: true, + }, + ]; + 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: [] }, + }; + mockSessionQuery.findOne.mockResolvedValue(chatSession); + (createSessionWorkspacePod as jest.Mock).mockRejectedValueOnce(new Error('pod creation failed')); + + await expect( + AgentSessionService.provisionChatRuntime({ + sessionId: 'aaaaaaaa-bbbb-cccc-dddd-eeeeeeeeeeee', + userId: 'user-123', + userIdentity: { + userId: 'user-123', + githubUsername: 'sample-user', + } as any, + githubToken: 'sample-gh-token', + }) + ).rejects.toThrow('pod creation failed'); + + expect(createSessionWorkspacePod).toHaveBeenCalledWith( + expect.objectContaining({ + workspaceRepos, + }) + ); + expect(mockSessionQuery.patch).toHaveBeenLastCalledWith( + expect.objectContaining({ + workspaceStatus: AgentWorkspaceStatus.FAILED, + namespace: null, + podName: null, + pvcName: null, + }) + ); + expect(mockSessionQuery.patch).toHaveBeenLastCalledWith( + expect.not.objectContaining({ + workspaceRepos: expect.any(Array), + selectedServices: expect.any(Array), + }) + ); + }); + + it('seeds repo-scoped stdio MCP servers when provisioning a build-context chat runtime', async () => { + const workspaceRepos = [ + { + repo: 'example-org/example-repo', + repoUrl: 'https://github.com/example-org/example-repo.git', + branch: 'feature/sample', + revision: 'commit-sha-1', + mountPath: '/workspace', + primary: true, + }, + ]; + 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); + mockResolveSessionPodServersForRepo.mockResolvedValueOnce([ + { + slug: 'sample-stdio', + name: 'Sample stdio', + transport: { + type: 'stdio', + command: 'sample-mcp', + args: ['--stdio'], + env: { + SAMPLE_TOKEN: 'sample-secret', + }, + }, + timeout: 30000, + defaultArgs: {}, + env: {}, + discoveredTools: [], + }, + ]); + + 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(mockResolveSessionPodServersForRepo).toHaveBeenCalledWith( + 'example-org/example-repo', + undefined, + expect.objectContaining({ + userId: 'user-123', + githubUsername: 'sample-user', + }) + ); + expect(createAgentApiKeySecret).toHaveBeenCalledWith( + 'chat-aaaaaaaa', + 'agent-secret-aaaaaaaa', + {}, + 'sample-gh-token', + undefined, + {}, + { + LIFECYCLE_SESSION_MCP_CONFIG_JSON: JSON.stringify([ + { + slug: 'sample-stdio', + name: 'Sample stdio', + transport: { + type: 'stdio', + command: 'sample-mcp', + args: ['--stdio'], + env: { + SAMPLE_TOKEN: 'sample-secret', + }, + }, + timeout: 30000, + }, + ]), + } + ); + }); + + it.each([AgentSessionKind.ENVIRONMENT, AgentSessionKind.SANDBOX])( + 'rejects %s sessions during chat runtime suspension', + async (sessionKind) => { + mockSessionQuery.findOne.mockResolvedValueOnce({ + id: 321, + uuid: 'sample-session-id', + userId: 'sample-user', + sessionKind, + status: 'active', + workspaceStatus: AgentWorkspaceStatus.READY, + namespace: 'sample-namespace', + podName: 'sample-pod', + pvcName: 'sample-pvc', + }); + + await expect( + AgentSessionService.suspendChatRuntime({ + sessionId: 'sample-session-id', + userId: 'sample-user', + }) + ).rejects.toThrow('Runtime suspension is only supported for chat sessions'); + + expect(deleteSessionWorkspacePod).not.toHaveBeenCalled(); + expect(deleteSessionWorkspaceService).not.toHaveBeenCalled(); + expect(deleteAgentApiKeySecret).not.toHaveBeenCalled(); + } + ); + + it.each([AgentSessionKind.ENVIRONMENT, AgentSessionKind.SANDBOX])( + 'rejects %s sessions during chat runtime resume provisioning', + async (sessionKind) => { + mockSessionQuery.findOne.mockResolvedValueOnce({ + id: 321, + uuid: 'sample-session-id', + userId: 'sample-user', + sessionKind, + status: 'active', + workspaceStatus: AgentWorkspaceStatus.HIBERNATED, + namespace: 'sample-namespace', + podName: null, + pvcName: 'sample-pvc', + }); + + await expect( + AgentSessionService.resumeChatRuntime({ + sessionId: 'sample-session-id', + userId: 'sample-user', + userIdentity: { + userId: 'sample-user', + githubUsername: 'sample-user', + } as any, + githubToken: 'sample-gh-token', + }) + ).rejects.toThrow('Runtime provisioning is only supported for chat sessions'); + + expect(mockCreateOrUpdateNamespace).not.toHaveBeenCalled(); + expect(createAgentPvc).not.toHaveBeenCalled(); + expect(createSessionWorkspacePod).not.toHaveBeenCalled(); + } + ); + it('publishes a chat session HTTP port through ingress', async () => { mockSessionQuery.findOne.mockResolvedValue({ id: 321, @@ -3220,6 +3451,47 @@ describe('AgentSessionService', () => { ); }); + it('appends dynamic build context for chat sessions without a namespace', async () => { + mockGetEffectiveAgentSessionConfig.mockResolvedValue({ + appendSystemPrompt: 'Use concise responses.', + }); + (systemPrompt.resolveAgentSessionPromptContext as jest.Mock).mockResolvedValue({ + namespace: null, + buildUuid: 'build-123', + services: [], + build: { uuid: 'build-123', status: 'build_failed' }, + }); + (systemPrompt.buildAgentSessionDynamicSystemPrompt as jest.Mock).mockReturnValue( + 'Session context:\n- buildUuid: build-123\nBuild context:\n- buildUuid=build-123: status=build_failed' + ); + (systemPrompt.combineAgentSessionAppendSystemPrompt as jest.Mock).mockReturnValue('combined build prompt'); + + (AgentSession.query as jest.Mock) = jest.fn().mockReturnValue({ + findOne: jest.fn().mockReturnValue({ + select: jest.fn().mockResolvedValue({ + id: 123, + namespace: null, + buildUuid: 'build-123', + skillPlan: { skills: [] }, + }), + }), + }); + + await expect(AgentSessionService.getSessionAppendSystemPrompt('sess-1')).resolves.toBe('combined build prompt'); + expect(systemPrompt.resolveAgentSessionPromptContext).toHaveBeenCalledWith({ + sessionDbId: 123, + namespace: null, + buildUuid: 'build-123', + }); + expect(systemPrompt.buildAgentSessionDynamicSystemPrompt).toHaveBeenCalledWith( + expect.objectContaining({ + namespace: null, + buildUuid: 'build-123', + toolLines: [], + }) + ); + }); + it('returns the configured control-plane prompt when the session cannot be found', async () => { mockGetEffectiveAgentSessionConfig.mockResolvedValue({ appendSystemPrompt: 'Use concise responses.', diff --git a/src/server/services/__tests__/agentSessionConfig.test.ts b/src/server/services/__tests__/agentSessionConfig.test.ts index 11ba4466..c5733adb 100644 --- a/src/server/services/__tests__/agentSessionConfig.test.ts +++ b/src/server/services/__tests__/agentSessionConfig.test.ts @@ -14,7 +14,7 @@ * limitations under the License. */ -jest.mock('server/services/ai/mcp/config', () => ({ +jest.mock('server/services/agentRuntime/mcp/config', () => ({ McpConfigService: jest.fn().mockImplementation(() => ({ listEffectiveDefinitions: jest.fn().mockResolvedValue([]), })), @@ -22,6 +22,9 @@ jest.mock('server/services/ai/mcp/config', () => ({ const mockGlobalConfigGetConfig = jest.fn(); const mockGlobalConfigSetConfig = jest.fn(); +const mockAgentRuntimeGetGlobalConfig = jest.fn(); +const mockAgentRuntimeGetRepoConfig = jest.fn(); +const mockAgentRuntimeGetEffectiveConfig = jest.fn(); jest.mock('server/services/globalConfig', () => ({ __esModule: true, @@ -33,6 +36,17 @@ jest.mock('server/services/globalConfig', () => ({ }, })); +jest.mock('server/services/agentRuntime/config/agentRuntimeConfig', () => ({ + __esModule: true, + default: { + getInstance: jest.fn(() => ({ + getGlobalConfig: (...args: unknown[]) => mockAgentRuntimeGetGlobalConfig(...args), + getRepoConfig: (...args: unknown[]) => mockAgentRuntimeGetRepoConfig(...args), + getEffectiveConfig: (...args: unknown[]) => mockAgentRuntimeGetEffectiveConfig(...args), + })), + }, +})); + import AgentSessionConfigService from 'server/services/agentSessionConfig'; import AgentPolicyService from 'server/services/agent/PolicyService'; import { DEFAULT_AGENT_APPROVAL_POLICY } from 'server/services/agent/types'; @@ -52,6 +66,9 @@ describe('AgentSessionConfigService', () => { jest.clearAllMocks(); mockGlobalConfigGetConfig.mockResolvedValue(undefined); mockGlobalConfigSetConfig.mockResolvedValue(undefined); + mockAgentRuntimeGetGlobalConfig.mockResolvedValue({}); + mockAgentRuntimeGetRepoConfig.mockResolvedValue({}); + mockAgentRuntimeGetEffectiveConfig.mockResolvedValue({}); }); it('lists admin-visible built-in tools in tool inventory', async () => { @@ -61,6 +78,9 @@ describe('AgentSessionConfigService', () => { jest.spyOn(service, 'getEffectiveConfig').mockResolvedValue({ systemPrompt: 'base', appendSystemPrompt: 'append', + maxIterations: 8, + workspaceToolDiscoveryTimeoutMs: 3000, + workspaceToolExecutionTimeoutMs: 15000, toolRules: [], }); jest.spyOn(AgentPolicyService, 'getEffectivePolicy').mockResolvedValue(DEFAULT_AGENT_APPROVAL_POLICY); @@ -100,6 +120,108 @@ describe('AgentSessionConfigService', () => { ); }); + it('lists global capability inventory with grouped tools and policy availability', async () => { + const service = makeService(); + + jest.spyOn(service, 'getGlobalConfig').mockResolvedValue({}); + jest.spyOn(service, 'getEffectiveConfig').mockResolvedValue({ + systemPrompt: 'base', + appendSystemPrompt: 'append', + maxIterations: 8, + workspaceToolDiscoveryTimeoutMs: 3000, + workspaceToolExecutionTimeoutMs: 15000, + toolRules: [], + }); + jest.spyOn(AgentPolicyService, 'getEffectivePolicy').mockResolvedValue(DEFAULT_AGENT_APPROVAL_POLICY); + mockAgentRuntimeGetGlobalConfig.mockResolvedValue({ + capabilityPolicy: { + availability: { + workspace_shell: 'admin_only', + }, + }, + }); + mockAgentRuntimeGetEffectiveConfig.mockResolvedValue({ + capabilityPolicy: { + availability: { + workspace_shell: 'admin_only', + }, + }, + }); + + const entries = await service.listCapabilityInventory('global'); + const shell = entries.find((entry) => entry.capabilityId === 'workspace_shell'); + + expect(shell).toEqual( + expect.objectContaining({ + capabilityId: 'workspace_shell', + configuredAvailability: 'admin_only', + effectiveAvailability: 'admin_only', + blockedReason: 'admin_only', + runtimeCapabilityKey: 'shell_exec', + resourceGrants: ['workspace_shell'], + }) + ); + expect(shell?.tools.map((tool) => tool.toolName)).toEqual(expect.arrayContaining(['workspace.exec_mutation'])); + }); + + it('lists repo capability inventory with inherited and repo-specific availability', async () => { + const service = makeService(); + + jest.spyOn(service, 'getGlobalConfig').mockResolvedValue({}); + jest.spyOn(service, 'getRepoConfig').mockResolvedValue({}); + jest.spyOn(service, 'getEffectiveConfig').mockResolvedValue({ + systemPrompt: 'base', + appendSystemPrompt: 'append', + maxIterations: 8, + workspaceToolDiscoveryTimeoutMs: 3000, + workspaceToolExecutionTimeoutMs: 15000, + toolRules: [], + }); + jest.spyOn(AgentPolicyService, 'getEffectivePolicy').mockResolvedValue(DEFAULT_AGENT_APPROVAL_POLICY); + mockAgentRuntimeGetGlobalConfig.mockResolvedValue({ + capabilityPolicy: { + availability: { + external_mcp_write: 'admin_only', + workspace_shell: 'disabled', + }, + }, + }); + mockAgentRuntimeGetRepoConfig.mockResolvedValue({ + capabilityPolicy: { + availability: { + workspace_shell: 'all_users', + }, + }, + }); + mockAgentRuntimeGetEffectiveConfig.mockResolvedValue({ + capabilityPolicy: { + availability: { + external_mcp_write: 'admin_only', + workspace_shell: 'all_users', + }, + }, + }); + + const entries = await service.listCapabilityInventory('Example-Org/Example-Repo'); + const shell = entries.find((entry) => entry.capabilityId === 'workspace_shell'); + const externalWrite = entries.find((entry) => entry.capabilityId === 'external_mcp_write'); + + expect(shell).toEqual( + expect.objectContaining({ + configuredAvailability: 'all_users', + inheritedAvailability: 'disabled', + effectiveAvailability: 'all_users', + }) + ); + expect(externalWrite).toEqual( + expect.objectContaining({ + inheritedAvailability: 'admin_only', + effectiveAvailability: 'admin_only', + resourceGrants: ['mcp_write'], + }) + ); + }); + it('merges repo control-plane numeric overrides over global defaults', async () => { const service = makeService(); @@ -173,6 +295,9 @@ describe('AgentSessionConfigService', () => { jest.spyOn(service, 'getEffectiveConfig').mockResolvedValue({ systemPrompt: 'base', appendSystemPrompt: 'append', + maxIterations: 8, + workspaceToolDiscoveryTimeoutMs: 3000, + workspaceToolExecutionTimeoutMs: 15000, toolRules: [ { toolKey: 'mcp__sandbox__workspace_read_file', diff --git a/src/server/services/__tests__/aiAgentConfig.test.ts b/src/server/services/__tests__/aiAgentConfig.test.ts deleted file mode 100644 index db0e4ac2..00000000 --- a/src/server/services/__tests__/aiAgentConfig.test.ts +++ /dev/null @@ -1,269 +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. - */ - -jest.mock('server/services/globalConfig', () => ({ - __esModule: true, - default: { - getInstance: jest.fn(), - }, -})); - -jest.mock('server/lib/logger', () => ({ - getLogger: jest.fn(() => ({ - info: jest.fn(), - warn: jest.fn(), - error: jest.fn(), - debug: jest.fn(), - })), -})); - -import AIAgentConfigService from 'server/services/aiAgentConfig'; -import GlobalConfigService from 'server/services/globalConfig'; -import type { AIAgentConfig } from 'server/services/types/aiAgentConfig'; - -function makeService() { - const repoUpsertQuery = { - insert: jest.fn().mockReturnThis(), - onConflict: jest.fn().mockReturnThis(), - merge: jest.fn().mockResolvedValue(undefined), - }; - - const knex = Object.assign(jest.fn().mockReturnValue(repoUpsertQuery), { - fn: { - now: jest.fn(() => 'now'), - }, - }); - - const db = { knex } as any; - const redis = { del: jest.fn().mockResolvedValue(undefined) } as any; - const redlock = {} as any; - const queueManager = {} as any; - - return { - service: new AIAgentConfigService(db, redis, redlock, queueManager), - knex, - repoUpsertQuery, - redis, - }; -} - -describe('AIAgentConfigService', () => { - beforeEach(() => { - jest.clearAllMocks(); - }); - - it('updates global additive rules without revalidating unrelated provider defaults', async () => { - const { service } = makeService(); - const currentConfig: AIAgentConfig = { - enabled: true, - providers: [ - { - name: 'gemini', - enabled: true, - apiKeyEnvVar: 'GOOGLE_GENERATIVE_AI_API_KEY', - models: [ - { - id: 'gemini-1', - displayName: 'Gemini 1', - enabled: true, - default: true, - maxTokens: 8192, - }, - { - id: 'gemini-2', - displayName: 'Gemini 2', - enabled: true, - default: true, - maxTokens: 8192, - }, - ], - }, - ], - maxMessagesPerSession: 50, - sessionTTL: 3600, - }; - - const setConfig = jest.fn().mockResolvedValue(undefined); - (GlobalConfigService.getInstance as jest.Mock).mockReturnValue({ - getConfig: jest.fn().mockResolvedValue(currentConfig), - setConfig, - }); - - const result = await service.updateGlobalAdditiveRules(['test']); - - expect(setConfig).toHaveBeenCalledWith( - 'aiAgent', - expect.objectContaining({ - providers: currentConfig.providers, - additiveRules: ['test'], - }) - ); - expect(result.additiveRules).toEqual(['test']); - expect(result.providers).toEqual(currentConfig.providers); - }); - - it('replaces global approval policy without revalidating unrelated provider defaults', async () => { - const { service } = makeService(); - const currentConfig: AIAgentConfig = { - enabled: true, - providers: [ - { - name: 'gemini', - enabled: true, - apiKeyEnvVar: 'GOOGLE_GENERATIVE_AI_API_KEY', - models: [ - { - id: 'gemini-1', - displayName: 'Gemini 1', - enabled: true, - default: true, - maxTokens: 8192, - }, - { - id: 'gemini-2', - displayName: 'Gemini 2', - enabled: true, - default: true, - maxTokens: 8192, - }, - ], - }, - ], - maxMessagesPerSession: 50, - sessionTTL: 3600, - approvalPolicy: { - defaultMode: 'require_approval', - }, - }; - - const setConfig = jest.fn().mockResolvedValue(undefined); - (GlobalConfigService.getInstance as jest.Mock).mockReturnValue({ - getConfig: jest.fn().mockResolvedValue(currentConfig), - setConfig, - }); - - const result = await service.updateGlobalApprovalPolicy({ - defaultMode: 'require_approval', - rules: { - shell_exec: 'deny', - }, - }); - - expect(setConfig).toHaveBeenCalledWith( - 'aiAgent', - expect.objectContaining({ - providers: currentConfig.providers, - approvalPolicy: { - defaultMode: 'require_approval', - rules: { - shell_exec: 'deny', - }, - }, - }) - ); - expect(result.approvalPolicy).toEqual({ - defaultMode: 'require_approval', - rules: { - shell_exec: 'deny', - }, - }); - expect(result.providers).toEqual(currentConfig.providers); - }); - - it('clears the global approval policy when given an empty replacement', async () => { - const { service } = makeService(); - const currentConfig: AIAgentConfig = { - enabled: true, - providers: [ - { - name: 'openai', - enabled: true, - apiKeyEnvVar: 'OPENAI_API_KEY', - models: [ - { - id: 'gpt-5', - displayName: 'GPT-5', - enabled: true, - default: true, - maxTokens: 8192, - }, - ], - }, - ], - maxMessagesPerSession: 50, - sessionTTL: 3600, - approvalPolicy: { - defaultMode: 'deny', - rules: { - shell_exec: 'deny', - }, - }, - }; - - const setConfig = jest.fn().mockResolvedValue(undefined); - (GlobalConfigService.getInstance as jest.Mock).mockReturnValue({ - getConfig: jest.fn().mockResolvedValue(currentConfig), - setConfig, - }); - - const result = await service.updateGlobalApprovalPolicy({}); - - expect(setConfig).toHaveBeenCalledWith( - 'aiAgent', - expect.objectContaining({ - providers: currentConfig.providers, - }) - ); - expect(setConfig.mock.calls[0]?.[1]).not.toHaveProperty('approvalPolicy'); - expect(result.approvalPolicy).toBeUndefined(); - expect(result.providers).toEqual(currentConfig.providers); - }); - - it('updates repo additive rules while preserving other repo overrides', async () => { - const { service, knex, repoUpsertQuery, redis } = makeService(); - - jest.spyOn(service, 'getRepoConfig').mockResolvedValue({ - excludedTools: ['tool-a'], - approvalPolicy: { - defaultMode: 'require_approval', - }, - }); - - const result = await service.updateRepoAdditiveRules('Example-Org/Example-Repo', ['test']); - - expect(knex).toHaveBeenCalledWith('ai_agent_repo_config'); - expect(repoUpsertQuery.insert).toHaveBeenCalledWith( - expect.objectContaining({ - repositoryFullName: 'example-org/example-repo', - config: JSON.stringify({ - excludedTools: ['tool-a'], - approvalPolicy: { - defaultMode: 'require_approval', - }, - additiveRules: ['test'], - }), - }) - ); - expect(redis.del).toHaveBeenCalledWith('ai_agent_repo_config:example-org/example-repo'); - expect(result).toEqual({ - excludedTools: ['tool-a'], - approvalPolicy: { - defaultMode: 'require_approval', - }, - additiveRules: ['test'], - }); - }); -}); diff --git a/src/server/services/agent/AdminService.ts b/src/server/services/agent/AdminService.ts index fbf0f595..381670f9 100644 --- a/src/server/services/agent/AdminService.ts +++ b/src/server/services/agent/AdminService.ts @@ -36,13 +36,13 @@ import type { McpDiscoveredTool, McpSharedConnectionConfig, McpTransportConfig, -} from 'server/services/ai/mcp/types'; +} from 'server/services/agentRuntime/mcp/types'; import { buildMcpDefinitionFingerprint, normalizeAuthConfig, requiresUserConnection, -} from 'server/services/ai/mcp/connectionConfig'; -import { redactSharedConfigSecrets } from 'server/services/ai/mcp/config'; +} from 'server/services/agentRuntime/mcp/connectionConfig'; +import { redactMcpConfigSecrets } from 'server/services/agentRuntime/mcp/config'; type SessionStatus = AgentSession['status']; @@ -478,14 +478,18 @@ export default class AgentAdminService { return configs.map((config) => { const rows = connectionRowsBySlug.get(config.slug) || []; const connectionRequired = requiresUserConnection(normalizeAuthConfig(config.authConfig)); + const redactedConfig = redactMcpConfigSecrets({ + transport: config.transport, + sharedConfig: config.sharedConfig || {}, + }); return { slug: config.slug, name: config.name, description: config.description ?? null, scope: config.scope, preset: config.preset ?? null, - transport: config.transport, - sharedConfig: redactSharedConfigSecrets({ sharedConfig: config.sharedConfig || {} }).sharedConfig || {}, + transport: redactedConfig.transport || config.transport, + sharedConfig: redactedConfig.sharedConfig || {}, authConfig: normalizeAuthConfig(config.authConfig), enabled: config.enabled, timeout: config.timeout, diff --git a/src/server/services/agent/AgentDefinitionRegistry.ts b/src/server/services/agent/AgentDefinitionRegistry.ts new file mode 100644 index 00000000..799e208e --- /dev/null +++ b/src/server/services/agent/AgentDefinitionRegistry.ts @@ -0,0 +1,171 @@ +/** + * 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 AgentDefinition from 'server/models/AgentDefinition'; +import type AgentSession from 'server/models/AgentSession'; +import type AgentSource from 'server/models/AgentSource'; +import { AgentSessionKind, AgentWorkspaceStatus } from 'shared/constants'; +import type { AgentDefinitionContract } from './agentDefinitionTypes'; +import { + isSystemAgentDefinitionId, + SYSTEM_AGENT_DEFINITIONS, + SYSTEM_AGENT_DEFINITION_IDS, + type SystemAgentDefinitionId, +} from './systemAgentDefinitions'; + +export type AgentDefinitionSummary = { + id: string; + version: number; + ownerKind: AgentDefinitionContract['owner']['kind']; + name: string; + description?: string | null; + status: AgentDefinitionContract['status']; + capabilityRefs: AgentDefinitionContract['capabilityRefs']; + requiredCapabilityRefs: AgentDefinitionContract['capabilityRefs']; + optionalCapabilityRefs: AgentDefinitionContract['capabilityRefs']; + resourcePolicy: AgentDefinitionContract['resourcePolicy']; + codeOwned: boolean; + readOnly: boolean; +}; + +function readString(value: unknown): string | null { + return typeof value === 'string' && value.trim() ? value.trim() : null; +} + +function toPersistenceRow(definition: AgentDefinitionContract): Partial { + return { + definitionId: definition.id, + version: definition.version, + ownerKind: definition.owner.kind, + ownerUserId: definition.owner.userId || null, + ownerOrganizationId: definition.owner.organizationId || null, + name: definition.name, + description: definition.description || null, + instructionRefs: definition.instructionRefs || [], + instructionAddendum: definition.instructionAddendum || null, + capabilityRefs: definition.capabilityRefs || [], + requiredCapabilityRefs: definition.requiredCapabilityRefs || definition.capabilityRefs || [], + optionalCapabilityRefs: definition.optionalCapabilityRefs || [], + resourcePolicy: definition.resourcePolicy, + modelPreference: definition.modelPreference || null, + status: definition.status, + codeOwned: Boolean(definition.codeOwned), + readOnly: Boolean(definition.readOnly), + }; +} + +export function agentDefinitionRowToContract(row: AgentDefinition): AgentDefinitionContract { + return { + id: row.definitionId, + version: row.version, + owner: { + kind: row.ownerKind, + userId: row.ownerUserId, + organizationId: row.ownerOrganizationId, + }, + name: row.name, + description: row.description, + instructionRefs: row.instructionRefs || [], + instructionAddendum: row.instructionAddendum, + capabilityRefs: row.capabilityRefs || [], + requiredCapabilityRefs: row.requiredCapabilityRefs || row.capabilityRefs || [], + optionalCapabilityRefs: row.optionalCapabilityRefs || [], + resourcePolicy: row.resourcePolicy, + modelPreference: row.modelPreference, + status: row.status, + codeOwned: row.codeOwned, + readOnly: row.readOnly, + }; +} + +export async function ensureSystemAgentDefinitionsSeeded(): Promise { + const rows = await Promise.all( + SYSTEM_AGENT_DEFINITION_IDS.map(async (agentId) => { + const row = toPersistenceRow(SYSTEM_AGENT_DEFINITIONS[agentId]); + return AgentDefinition.upsert(row, ['definitionId']) as Promise; + }) + ); + + return rows.map(agentDefinitionRowToContract); +} + +export async function listSystemAgentDefinitions(): Promise { + const rows = await AgentDefinition.query() + .whereIn('definitionId', [...SYSTEM_AGENT_DEFINITION_IDS]) + .where({ ownerKind: 'system' }) + .orderBy('definitionId', 'asc'); + + return rows.map(agentDefinitionRowToContract); +} + +export async function getSystemAgentDefinition(agentId: SystemAgentDefinitionId): Promise { + const row = await AgentDefinition.query().findOne({ + definitionId: agentId, + ownerKind: 'system', + }); + + if (!row) { + throw new Error(`System agent definition not found: ${agentId}`); + } + + return agentDefinitionRowToContract(row); +} + +export function serializeAgentDefinitionSummary(definition: AgentDefinitionContract): AgentDefinitionSummary { + return { + id: definition.id, + version: definition.version, + ownerKind: definition.owner.kind, + name: definition.name, + description: definition.description || null, + status: definition.status, + capabilityRefs: definition.capabilityRefs, + requiredCapabilityRefs: definition.requiredCapabilityRefs || definition.capabilityRefs, + optionalCapabilityRefs: definition.optionalCapabilityRefs || [], + resourcePolicy: definition.resourcePolicy, + codeOwned: Boolean(definition.codeOwned), + readOnly: Boolean(definition.readOnly), + }; +} + +export function assertAgentDefinitionMutable(definition: AgentDefinitionContract): void { + if (definition.codeOwned || definition.readOnly) { + throw new Error(`Agent definition "${definition.id}" is code-owned and read-only.`); + } +} + +export function inferDefaultSystemAgentDefinitionId( + session: AgentSession, + source: AgentSource +): SystemAgentDefinitionId { + if (session.sessionKind === AgentSessionKind.CHAT) { + if (readString(source.input?.buildUuid)) { + return 'system.debug'; + } + + if (session.workspaceStatus === AgentWorkspaceStatus.READY) { + return 'system.develop'; + } + + return 'system.freeform'; + } + + return 'system.develop'; +} + +export function normalizeSystemAgentDefinitionId(value: unknown): SystemAgentDefinitionId | null { + return isSystemAgentDefinitionId(value) ? value : null; +} diff --git a/src/server/services/agent/AgentSelectionService.ts b/src/server/services/agent/AgentSelectionService.ts new file mode 100644 index 00000000..bec43227 --- /dev/null +++ b/src/server/services/agent/AgentSelectionService.ts @@ -0,0 +1,371 @@ +/** + * 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 { Transaction } from 'objection'; +import AgentRun from 'server/models/AgentRun'; +import AgentThread from 'server/models/AgentThread'; +import type AgentSession from 'server/models/AgentSession'; +import type AgentSource from 'server/models/AgentSource'; +import type { RequestUserIdentity } from 'server/lib/get-user'; +import AgentCapabilityService from './CapabilityService'; +import * as AgentDefinitionRegistry from './AgentDefinitionRegistry'; +import { customAgentDefinitionService } from './CustomAgentDefinitionService'; +import AgentMessageStore from './MessageStore'; +import AgentPolicyService from './PolicyService'; +import { TERMINAL_RUN_STATUSES } from './RunService'; +import AgentSourceService from './SourceService'; +import AgentThreadService from './ThreadService'; +import type { AgentDefinitionContract } from './agentDefinitionTypes'; +import type { AgentCapabilitySourceKind } from './capabilityCatalog'; +import { + SYSTEM_AGENT_DEFINITION_IDS, + sourceKindForSystemAgentDefinitionId, + type SystemAgentDefinitionId, +} from './systemAgentDefinitions'; + +export type AgentSelectionGroupId = 'built_in' | 'my_agents'; + +export type AgentSelectionUnavailableReason = + | 'unknown_agent' + | 'active_run' + | 'disabled_agent' + | 'requires_workspace' + | 'source_incompatible' + | 'disabled_by_policy'; + +export class AgentThreadAgentSwitchError extends Error { + constructor( + public readonly reason: AgentSelectionUnavailableReason, + message: string, + public readonly details: Record = {} + ) { + super(message); + this.name = 'AgentThreadAgentSwitchError'; + } +} + +export type AgentSelectionSummary = { + id: string; + ownerKind: AgentDefinitionContract['owner']['kind']; + label: string; + description: string | null; + available: boolean; + unavailableReason: AgentSelectionUnavailableReason | null; + unavailableMessage: string | null; + group: AgentSelectionGroupId; +}; + +export type AgentSelectionGroup = { + id: AgentSelectionGroupId; + label: string; + agents: AgentSelectionSummary[]; +}; + +export type AgentSelectionState = { + selectedId: string | null; + defaultId: SystemAgentDefinitionId; + currentId: string; + groups: AgentSelectionGroup[]; +}; + +export type SwitchThreadAgentInput = { + threadId: string; + userIdentity: RequestUserIdentity; + agentId: string; +}; + +export type SwitchThreadAgentResult = { + previousAgent: AgentSelectionSummary; + nextAgent: AgentSelectionSummary; + switched: boolean; + state: AgentSelectionState; +}; + +type ValidationContext = { + sourceKind: AgentCapabilitySourceKind; + capabilityPolicy: Awaited>['capabilityPolicy']; + customAgentCreationPolicy: Awaited< + ReturnType + >['customAgentCreationPolicy']; + approvalPolicy: Awaited>['approvalPolicy']; + activeRun: boolean; +}; + +function orderSystemDefinitions(definitions: AgentDefinitionContract[]): AgentDefinitionContract[] { + const byId = new Map(definitions.map((definition) => [definition.id, definition])); + return SYSTEM_AGENT_DEFINITION_IDS.flatMap((agentId) => { + const definition = byId.get(agentId); + return definition ? [definition] : []; + }); +} + +function actorLabel(userIdentity: RequestUserIdentity): string { + return ( + userIdentity.displayName || + userIdentity.preferredUsername || + userIdentity.githubUsername || + userIdentity.email || + 'You' + ); +} + +async function hasActiveRun(sessionId: number, trx?: Transaction): Promise { + const activeRun = await AgentRun.query(trx).where({ sessionId }).whereNotIn('status', TERMINAL_RUN_STATUSES).first(); + + return Boolean(activeRun); +} + +function validateDefinition( + definition: AgentDefinitionContract, + context: ValidationContext +): Pick { + if (context.activeRun) { + return { + available: false, + unavailableReason: 'active_run', + unavailableMessage: 'Wait for the current run to finish before switching agents.', + }; + } + + if (definition.status !== 'active') { + return { + available: false, + unavailableReason: 'disabled_agent', + unavailableMessage: `${definition.name} is unavailable.`, + }; + } + + if (definition.resourcePolicy.workspaceRequired && context.sourceKind !== 'workspace_session') { + return { + available: false, + unavailableReason: 'requires_workspace', + unavailableMessage: 'Requires a prepared workspace.', + }; + } + + if (!definition.resourcePolicy.sourceKinds.includes(context.sourceKind)) { + return { + available: false, + unavailableReason: 'source_incompatible', + unavailableMessage: `${definition.name} is unavailable for this conversation.`, + }; + } + + const requiredCapabilityRefs = definition.requiredCapabilityRefs || definition.capabilityRefs; + const blockedCapability = AgentPolicyService.resolveCapabilitySetAccess(requiredCapabilityRefs, { + capabilityPolicy: context.capabilityPolicy, + customAgentCreationPolicy: context.customAgentCreationPolicy, + approvalPolicy: context.approvalPolicy, + definitionOwnerKind: definition.owner.kind, + sourceKind: context.sourceKind, + }).find((capability) => !capability.allowed); + + if (blockedCapability) { + return { + available: false, + unavailableReason: 'disabled_by_policy', + unavailableMessage: `${definition.name} is unavailable because a required capability is disabled.`, + }; + } + + return { + available: true, + unavailableReason: null, + unavailableMessage: null, + }; +} + +function summarizeDefinition( + definition: AgentDefinitionContract, + context: ValidationContext, + group: AgentSelectionGroupId +): AgentSelectionSummary { + return { + id: definition.id, + ownerKind: definition.owner.kind, + label: definition.name, + description: definition.description || null, + group, + ...validateDefinition(definition, context), + }; +} + +function flattenGroups(groups: AgentSelectionGroup[]): AgentSelectionSummary[] { + return groups.flatMap((group) => group.agents); +} + +export default class AgentSelectionService { + static async getThreadAgentState({ + threadId, + userIdentity, + }: { + threadId: string; + userIdentity: RequestUserIdentity; + }): Promise { + const { thread, session } = await AgentThreadService.getOwnedThreadWithSession(threadId, userIdentity.userId); + return this.getThreadAgentStateForThread({ thread, session, userIdentity }); + } + + static async getThreadAgentStateForThread({ + thread, + session, + userIdentity, + }: { + thread: AgentThread; + session: AgentSession; + userIdentity: RequestUserIdentity; + }): Promise { + const source = await AgentSourceService.getSessionSource(session.id); + if (!source || source.status !== 'ready') { + throw new AgentThreadAgentSwitchError('source_incompatible', 'Session source is not ready yet.'); + } + + return this.buildThreadAgentState({ thread, session, source, userIdentity }); + } + + static async switchThreadAgent({ + threadId, + userIdentity, + agentId, + }: SwitchThreadAgentInput): Promise { + const { thread, session } = await AgentThreadService.getOwnedThreadWithSession(threadId, userIdentity.userId); + const source = await AgentSourceService.getSessionSource(session.id); + if (!source || source.status !== 'ready') { + throw new AgentThreadAgentSwitchError('source_incompatible', 'Session source is not ready yet.'); + } + + const state = await this.buildThreadAgentState({ thread, session, source, userIdentity }); + const agents = flattenGroups(state.groups); + const previousAgent = agents.find((agent) => agent.id === state.currentId); + const nextAgent = agents.find((agent) => agent.id === agentId); + if (!previousAgent || !nextAgent) { + throw new AgentThreadAgentSwitchError('unknown_agent', 'Unknown agent.', { agentId }); + } + + if (!nextAgent.available) { + throw new AgentThreadAgentSwitchError( + nextAgent.unavailableReason || 'source_incompatible', + nextAgent.unavailableMessage || `${nextAgent.label} is unavailable.`, + { agentId } + ); + } + + if (state.currentId === agentId) { + return { + previousAgent, + nextAgent, + switched: false, + state, + }; + } + + return AgentThread.transaction(async (trx) => { + if (await hasActiveRun(session.id, trx)) { + throw new AgentThreadAgentSwitchError( + 'active_run', + 'Wait for the current run to finish before switching agents.', + { agentId } + ); + } + + const patchedThread = await AgentThread.query(trx).patchAndFetchById(thread.id, { + metadata: { + ...(thread.metadata || {}), + ...AgentThreadService.buildSelectedAgentDefinitionMetadataPatch(agentId), + }, + } as Partial); + + await AgentMessageStore.createAgentSwitchEvent({ + thread: patchedThread, + actor: { + userId: userIdentity.userId, + label: actorLabel(userIdentity), + }, + beforeAgent: { + id: previousAgent.id, + label: previousAgent.label, + }, + afterAgent: { + id: nextAgent.id, + label: nextAgent.label, + }, + trx, + }); + + return { + previousAgent, + nextAgent, + switched: true, + state: { + ...state, + selectedId: agentId, + currentId: agentId, + }, + }; + }); + } + + private static async buildThreadAgentState({ + thread, + session, + source, + userIdentity, + }: { + thread: AgentThread; + session: AgentSession; + source: AgentSource; + userIdentity: RequestUserIdentity; + }): Promise { + await AgentDefinitionRegistry.ensureSystemAgentDefinitionsSeeded(); + const systemDefinitions = orderSystemDefinitions(await AgentDefinitionRegistry.listSystemAgentDefinitions()); + const customDefinitions = ( + await customAgentDefinitionService.listUserDefinitions({ userId: userIdentity.userId }) + ).filter((definition) => definition.status === 'active'); + const defaultId = AgentDefinitionRegistry.inferDefaultSystemAgentDefinitionId(session, source); + const selectedId = AgentThreadService.getSelectedAgentDefinitionId(thread); + const { approvalPolicy, capabilityPolicy, customAgentCreationPolicy } = + await AgentCapabilityService.resolveSessionContext(session.uuid, userIdentity); + const activeRun = await hasActiveRun(session.id); + const context: ValidationContext = { + sourceKind: sourceKindForSystemAgentDefinitionId(defaultId), + capabilityPolicy, + customAgentCreationPolicy, + approvalPolicy, + activeRun, + }; + const groups: AgentSelectionGroup[] = [ + { + id: 'built_in', + label: 'Built in', + agents: systemDefinitions.map((definition) => summarizeDefinition(definition, context, 'built_in')), + }, + { + id: 'my_agents', + label: 'My agents', + agents: customDefinitions.map((definition) => summarizeDefinition(definition, context, 'my_agents')), + }, + ]; + const agents = flattenGroups(groups); + const selectedAgent = selectedId ? agents.find((agent) => agent.id === selectedId) : null; + + return { + selectedId: selectedAgent ? selectedId : null, + defaultId, + currentId: selectedAgent ? selectedId! : defaultId, + groups, + }; + } +} diff --git a/src/server/services/agent/AgentUsageService.ts b/src/server/services/agent/AgentUsageService.ts new file mode 100644 index 00000000..e2942134 --- /dev/null +++ b/src/server/services/agent/AgentUsageService.ts @@ -0,0 +1,291 @@ +/** + * 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 type { AgentRunStatus } from './types'; +import AgentThreadService from './ThreadService'; + +const OPTIONAL_USAGE_FIELDS = [ + 'inputTokens', + 'outputTokens', + 'reasoningTokens', + 'cachedInputTokens', + 'cacheCreationInputTokens', + 'cacheReadInputTokens', + 'nonCachedInputTokens', + 'textOutputTokens', +] as const; + +type OptionalUsageField = (typeof OPTIONAL_USAGE_FIELDS)[number]; +type UsageRecord = Partial>; + +const MISSING_USAGE_STATUSES: AgentRunStatus[] = [ + 'waiting_for_approval', + 'waiting_for_input', + 'completed', + 'failed', + 'cancelled', +]; + +export interface AgentUsageRunRecord { + sessionId?: number; + status: AgentRunStatus; + resolvedProvider?: string | null; + resolvedModel?: string | null; + provider?: string | null; + model?: string | null; + usageSummary?: Record | null; +} + +export interface AgentUsageSummary { + totalTokens: number; + inputTokens?: number; + outputTokens?: number; + reasoningTokens?: number; + cachedInputTokens?: number; + cacheCreationInputTokens?: number; + cacheReadInputTokens?: number; + nonCachedInputTokens?: number; + textOutputTokens?: number; +} + +export interface AgentUsageByModel extends AgentUsageSummary { + provider: string; + model: string; + runCount: number; + reportedRunCount: number; + missingUsageRunCount: number; +} + +export interface AgentUsageCompleteness { + runCount: number; + reportedRunCount: number; + missingUsageRunCount: number; + complete: boolean; +} + +export interface AgentUsageAggregate { + usageSummary: AgentUsageSummary; + usageByModel: AgentUsageByModel[]; + usageCompleteness: AgentUsageCompleteness; +} + +export interface AgentThreadUsageAggregate extends AgentUsageAggregate { + threadId: string; + sessionId: string; +} + +export interface AgentSessionUsageAggregate extends AgentUsageAggregate { + sessionId: string; +} + +interface UsageBucket { + provider: string; + model: string; + usageSummary: AgentUsageSummary; + runCount: number; + reportedRunCount: number; + missingUsageRunCount: number; +} + +function readUsageRecord(value: unknown): UsageRecord { + if (!value || typeof value !== 'object' || Array.isArray(value)) { + return {}; + } + + return value as UsageRecord; +} + +function readFiniteNumber(value: unknown): number | undefined { + return typeof value === 'number' && Number.isFinite(value) ? value : undefined; +} + +function readExactTotal(usageSummary: UsageRecord): number | undefined { + const totalTokens = readFiniteNumber(usageSummary.totalTokens); + if (totalTokens !== undefined) { + return totalTokens; + } + + const inputTokens = readFiniteNumber(usageSummary.inputTokens); + const outputTokens = readFiniteNumber(usageSummary.outputTokens); + if (inputTokens === undefined || outputTokens === undefined) { + return undefined; + } + + const computedTotal = inputTokens + outputTokens; + return Number.isFinite(computedTotal) ? computedTotal : undefined; +} + +function addOptionalUsageFields(target: AgentUsageSummary, usageSummary: UsageRecord): void { + for (const field of OPTIONAL_USAGE_FIELDS) { + const amount = readFiniteNumber(usageSummary[field]); + if (amount !== undefined) { + target[field] = (target[field] ?? 0) + amount; + } + } +} + +function readAttributionValue(primary: unknown, fallback: unknown, unknownValue: string): string { + if (typeof primary === 'string' && primary.trim()) { + return primary.trim(); + } + + if (typeof fallback === 'string' && fallback.trim()) { + return fallback.trim(); + } + + return unknownValue; +} + +function readAttribution(run: AgentUsageRunRecord): { provider: string; model: string } { + return { + provider: readAttributionValue(run.resolvedProvider, run.provider, 'unknown_provider'), + model: readAttributionValue(run.resolvedModel, run.model, 'unknown_model'), + }; +} + +function shouldCountMissingUsage(run: AgentUsageRunRecord): boolean { + return MISSING_USAGE_STATUSES.includes(run.status); +} + +function serializeUsageBucket(bucket: UsageBucket): AgentUsageByModel { + return { + provider: bucket.provider, + model: bucket.model, + ...bucket.usageSummary, + runCount: bucket.runCount, + reportedRunCount: bucket.reportedRunCount, + missingUsageRunCount: bucket.missingUsageRunCount, + }; +} + +export default class AgentUsageService { + static aggregateRuns(runs: AgentUsageRunRecord[]): AgentUsageAggregate { + const usageSummary: AgentUsageSummary = { + totalTokens: 0, + }; + const bucketsByAttribution = new Map(); + let reportedRunCount = 0; + let missingUsageRunCount = 0; + + for (const run of runs) { + const attribution = readAttribution(run); + const bucketKey = `${attribution.provider}\0${attribution.model}`; + let bucket = bucketsByAttribution.get(bucketKey); + if (!bucket) { + bucket = { + provider: attribution.provider, + model: attribution.model, + usageSummary: { + totalTokens: 0, + }, + runCount: 0, + reportedRunCount: 0, + missingUsageRunCount: 0, + }; + bucketsByAttribution.set(bucketKey, bucket); + } + + bucket.runCount += 1; + const runUsageSummary = readUsageRecord(run.usageSummary); + const exactTotal = readExactTotal(runUsageSummary); + if (exactTotal !== undefined) { + usageSummary.totalTokens += exactTotal; + bucket.usageSummary.totalTokens += exactTotal; + reportedRunCount += 1; + bucket.reportedRunCount += 1; + } else if (shouldCountMissingUsage(run)) { + missingUsageRunCount += 1; + bucket.missingUsageRunCount += 1; + } + + addOptionalUsageFields(usageSummary, runUsageSummary); + addOptionalUsageFields(bucket.usageSummary, runUsageSummary); + } + + return { + usageSummary, + usageByModel: [...bucketsByAttribution.values()].map(serializeUsageBucket), + usageCompleteness: { + runCount: runs.length, + reportedRunCount, + missingUsageRunCount, + complete: missingUsageRunCount === 0, + }, + }; + } + + static async aggregateThreadUsage(threadId: number): Promise { + const runs = await AgentRun.query().where({ threadId }).orderBy('createdAt', 'asc').orderBy('id', 'asc'); + return this.aggregateRuns(runs); + } + + static async aggregateSessionUsage(sessionId: number): Promise { + const runs = await AgentRun.query().where({ sessionId }).orderBy('createdAt', 'asc').orderBy('id', 'asc'); + return this.aggregateRuns(runs); + } + + static async aggregateSessionsUsage(sessionIds: number[]): Promise> { + const usageBySessionId = new Map(); + for (const sessionId of sessionIds) { + usageBySessionId.set(sessionId, this.aggregateRuns([])); + } + + if (sessionIds.length === 0) { + return usageBySessionId; + } + + const runs = await AgentRun.query() + .whereIn('sessionId', sessionIds) + .orderBy('sessionId', 'asc') + .orderBy('createdAt', 'asc') + .orderBy('id', 'asc'); + const runsBySessionId = new Map(); + + for (const run of runs) { + const existing = runsBySessionId.get(run.sessionId) || []; + existing.push(run); + runsBySessionId.set(run.sessionId, existing); + } + + for (const [sessionId, sessionRuns] of runsBySessionId.entries()) { + usageBySessionId.set(sessionId, this.aggregateRuns(sessionRuns)); + } + + return usageBySessionId; + } + + static async getOwnedThreadUsage(threadId: string, userId: string): Promise { + const { thread, session } = await AgentThreadService.getOwnedThreadWithSession(threadId, userId); + const aggregate = await this.aggregateThreadUsage(thread.id); + + return { + threadId: thread.uuid, + sessionId: session.uuid, + ...aggregate, + }; + } + + static async getOwnedSessionUsage(sessionId: string, userId: string): Promise { + const session = await AgentThreadService.getOwnedSession(sessionId, userId); + const aggregate = await this.aggregateSessionUsage(session.id); + + return { + sessionId: session.uuid, + ...aggregate, + }; + } +} diff --git a/src/server/services/agent/ApprovalService.ts b/src/server/services/agent/ApprovalService.ts index 169ce934..ca58c417 100644 --- a/src/server/services/agent/ApprovalService.ts +++ b/src/server/services/agent/ApprovalService.ts @@ -41,6 +41,11 @@ import { type ToolLikePart = ToolUIPart | DynamicToolUIPart; const SESSION_WORKSPACE_TOOL_KEY_PREFIX = `mcp__${SESSION_WORKSPACE_SERVER_SLUG}__`; +const FORCE_APPROVAL_TOOL_CAPABILITIES: Record = { + [buildAgentToolKey(LIFECYCLE_BUILTIN_SERVER_SLUG, 'update_file')]: 'git_write', + [buildAgentToolKey(LIFECYCLE_BUILTIN_SERVER_SLUG, 'update_pr_labels')]: 'git_write', + [buildAgentToolKey(LIFECYCLE_BUILTIN_SERVER_SLUG, 'patch_k8s_resource')]: 'deploy_k8s_mutation', +}; const ARGUMENT_PREVIEW_MAX_LENGTH = 160; const PENDING_ACTION_RESPONSE_FIELDS = new Set(['approved', 'reason']); @@ -76,6 +81,24 @@ function readString(value: unknown): string | null { return typeof value === 'string' && value.trim() ? value : null; } +function readNumber(value: unknown): number | null { + return typeof value === 'number' && Number.isFinite(value) ? value : null; +} + +function readFileChangeKind(value: unknown): AgentFileChangeData['kind'] | null { + return value === 'created' || value === 'edited' || value === 'deleted' ? value : null; +} + +function readFileChangeStage(value: unknown): AgentFileChangeData['stage'] | null { + return value === 'awaiting-approval' || + value === 'approved' || + value === 'applied' || + value === 'denied' || + value === 'failed' + ? value + : null; +} + function truncatePreview(value: string): string { return value.length > ARGUMENT_PREVIEW_MAX_LENGTH ? `${value.slice(0, ARGUMENT_PREVIEW_MAX_LENGTH - 3)}...` : value; } @@ -106,7 +129,7 @@ function summarizeArguments(input: unknown): Array<{ name: string; value: string } return Object.entries(input) - .filter(([name]) => !['content', 'oldText', 'newText', 'command', 'cmd'].includes(name)) + .filter(([name]) => !['content', 'new_content', 'oldText', 'newText', 'command', 'cmd'].includes(name)) .slice(0, 6) .map(([name, value]) => ({ name, @@ -134,14 +157,15 @@ function getCommandPreview(input: unknown): string | null { return null; } -function summarizeFileChanges(value: unknown): Array<{ - path: string; - action: string; - summary: string; - additions: number | null; - deletions: number | null; - truncated: boolean; -}> { +function summarizeFileChanges({ + value, + fallbackToolCallId, + fallbackSourceTool, +}: { + value: unknown; + fallbackToolCallId: string | null; + fallbackSourceTool: string | null; +}): AgentFileChangeData[] { if (!Array.isArray(value)) { return []; } @@ -149,16 +173,38 @@ function summarizeFileChanges(value: unknown): Array<{ 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'; + .flatMap((change) => { + const path = readString(change.path); + const kind = readFileChangeKind(change.kind); + if (!path || !kind) { + return []; + } + + const displayPath = readString(change.displayPath) || path; + const toolCallId = readString(change.toolCallId) || fallbackToolCallId || `${path}:file-change`; + const sourceTool = readString(change.sourceTool) || fallbackSourceTool || 'tool'; + const stage = readFileChangeStage(change.stage) || 'awaiting-approval'; + return { + id: readString(change.id) || `${toolCallId}:${path}`, + toolCallId, + sourceTool, path, - action, - summary: readString(change.summary) || `${action} ${path}`, - additions: typeof change.additions === 'number' ? change.additions : null, - deletions: typeof change.deletions === 'number' ? change.deletions : null, + displayPath, + kind, + stage, + additions: readNumber(change.additions) ?? 0, + deletions: readNumber(change.deletions) ?? 0, truncated: change.truncated === true, + unifiedDiff: readString(change.unifiedDiff), + beforeTextPreview: readString(change.beforeTextPreview), + afterTextPreview: readString(change.afterTextPreview), + summary: readString(change.summary) || `${kind} ${displayPath}`, + encoding: readString(change.encoding), + oldSizeBytes: readNumber(change.oldSizeBytes), + newSizeBytes: readNumber(change.newSizeBytes), + oldSha256: readString(change.oldSha256), + newSha256: readString(change.newSha256), }; }); } @@ -187,6 +233,11 @@ function getRiskLabels(capabilityKey: string | null | undefined): string[] { } function resolveApprovalCapabilityKey(toolName: string, fallback: AgentCapabilityKey): AgentCapabilityKey { + const forcedApprovalCapabilityKey = FORCE_APPROVAL_TOOL_CAPABILITIES[toolName]; + if (forcedApprovalCapabilityKey) { + return forcedApprovalCapabilityKey; + } + if (toolName === buildAgentToolKey(LIFECYCLE_BUILTIN_SERVER_SLUG, CHAT_PUBLISH_HTTP_TOOL_NAME)) { return 'deploy_k8s_mutation'; } @@ -217,7 +268,17 @@ function shouldPersistApprovalRequest({ const capabilityKey = resolveApprovalCapabilityKey(toolName, fallbackCapabilityKey); const toolRule = toolRules?.find((rule) => rule.toolKey === toolName); - const mode = toolRule?.mode || AgentPolicyService.modeForCapability(approvalPolicy, capabilityKey); + const capabilityMode = AgentPolicyService.modeForCapability(approvalPolicy, capabilityKey); + + if (toolRule?.mode === 'deny' || (!toolRule && capabilityMode === 'deny')) { + return false; + } + + if (FORCE_APPROVAL_TOOL_CAPABILITIES[toolName]) { + return true; + } + + const mode = toolRule?.mode || capabilityMode; return mode === 'require_approval'; } @@ -657,7 +718,11 @@ export default class ApprovalService { toolName, argumentsSummary: summarizeArguments(input), commandPreview: getCommandPreview(input), - fileChangePreview: summarizeFileChanges(payload.fileChanges), + fileChangePreview: summarizeFileChanges({ + value: payload.fileChanges, + fallbackToolCallId: readString(payload.toolCallId), + fallbackSourceTool: toolName, + }), riskLabels: getRiskLabels(action.capabilityKey), }; } diff --git a/src/server/services/agent/BuildContextChatService.ts b/src/server/services/agent/BuildContextChatService.ts new file mode 100644 index 00000000..c17d6750 --- /dev/null +++ b/src/server/services/agent/BuildContextChatService.ts @@ -0,0 +1,164 @@ +/** + * 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 { RequestUserIdentity } from 'server/lib/get-user'; +import { getLogger } from 'server/lib/logger'; +import AgentSession from 'server/models/AgentSession'; +import AgentThread from 'server/models/AgentThread'; +import Build from 'server/models/Build'; +import { AgentChatStatus, AgentSessionKind } from 'shared/constants'; +import AgentChatSessionService, { type AgentBuildContextChatMetadata } from './ChatSessionService'; +import AgentThreadService from './ThreadService'; + +const ACTIVE_BUILD_CONTEXT_CHAT_UNIQUE_CONSTRAINT = 'agent_sessions_active_build_context_chat_unique'; + +export class BuildContextChatBuildNotFoundError extends Error { + constructor(readonly buildUuid: string) { + super(`Build not found: ${buildUuid}`); + this.name = 'BuildContextChatBuildNotFoundError'; + } +} + +interface LaunchBuildContextChatOptions { + buildUuid: string; + userId: string; + userIdentity?: RequestUserIdentity; + model?: string; +} + +interface LaunchBuildContextChatResult { + session: AgentSession; + thread: AgentThread; + created: boolean; + reused: boolean; + buildContext: AgentBuildContextChatMetadata; +} + +function isUniqueConstraintError(error: unknown, constraintName: string): boolean { + const knexError = error as { code?: string; constraint?: string }; + return knexError?.code === '23505' && knexError?.constraint === constraintName; +} + +function buildLaunchMetadata(build: Build, buildUuid: string): AgentBuildContextChatMetadata { + const pullRequest = build.pullRequest + ? { + fullName: build.pullRequest.fullName || null, + branchName: build.pullRequest.branchName || null, + pullRequestNumber: build.pullRequest.pullRequestNumber || null, + } + : null; + + return { + buildUuid, + buildKind: build.kind || null, + namespace: build.namespace || null, + baseBuildUuid: build.baseBuild?.uuid || null, + revision: build.sha || build.pullRequest?.latestCommit || null, + pullRequest, + contextFreshAt: new Date().toISOString(), + }; +} + +async function findReusableBuildContextChat(buildUuid: string, userId: string): Promise { + return AgentSession.query() + .where({ + userId, + buildUuid, + sessionKind: AgentSessionKind.CHAT, + status: 'active', + chatStatus: AgentChatStatus.READY, + }) + .orderBy('updatedAt', 'desc') + .orderBy('createdAt', 'desc') + .first(); +} + +export default class BuildContextChatService { + static async launchBuildContextChat(opts: LaunchBuildContextChatOptions): Promise { + const build = await Build.query().findOne({ uuid: opts.buildUuid }).withGraphFetched('[pullRequest, baseBuild]'); + if (!build) { + throw new BuildContextChatBuildNotFoundError(opts.buildUuid); + } + + const buildContext = buildLaunchMetadata(build, opts.buildUuid); + const existingSession = await findReusableBuildContextChat(opts.buildUuid, opts.userId); + + if (existingSession) { + const thread = await AgentThreadService.getDefaultThreadForSession(existingSession.uuid, opts.userId); + const reused = true; + + getLogger().info( + `Session: launched build-context chat buildUuid=${opts.buildUuid} sessionId=${existingSession.uuid} reused=${reused}` + ); + + return { + session: existingSession, + thread, + created: false, + reused, + buildContext, + }; + } + + let session: AgentSession; + try { + session = await AgentChatSessionService.createChatSession({ + userId: opts.userId, + userIdentity: opts.userIdentity, + model: opts.model, + buildContext, + }); + } catch (error) { + if (!isUniqueConstraintError(error, ACTIVE_BUILD_CONTEXT_CHAT_UNIQUE_CONSTRAINT)) { + throw error; + } + + const racedSession = await findReusableBuildContextChat(opts.buildUuid, opts.userId); + if (!racedSession) { + throw error; + } + + const thread = await AgentThreadService.getDefaultThreadForSession(racedSession.uuid, opts.userId); + + getLogger().info( + `Session: launched build-context chat buildUuid=${opts.buildUuid} sessionId=${racedSession.uuid} reused=true` + ); + + return { + session: racedSession, + thread, + created: false, + reused: true, + buildContext, + }; + } + + const thread = await AgentThreadService.getDefaultThreadForSession(session.uuid, opts.userId); + const reused = false; + + getLogger().info( + `Session: launched build-context chat buildUuid=${opts.buildUuid} sessionId=${session.uuid} reused=${reused}` + ); + + return { + session, + thread, + created: true, + reused, + buildContext, + }; + } +} diff --git a/src/server/services/agent/CapabilityService.ts b/src/server/services/agent/CapabilityService.ts index 87e99df9..f941ba6b 100644 --- a/src/server/services/agent/CapabilityService.ts +++ b/src/server/services/agent/CapabilityService.ts @@ -18,16 +18,24 @@ 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'; -import { applyMcpDefaultToolArgs } from 'server/services/ai/mcp/runtimeConfig'; -import { usesSessionWorkspaceGatewayExecution } from 'server/services/ai/mcp/sessionPod'; +import { McpConfigService, sanitizeMcpErrorMessage, sanitizeMcpResult } from 'server/services/agentRuntime/mcp/config'; +import { McpClientManager } from 'server/services/agentRuntime/mcp/client'; +import { applyMcpDefaultToolArgs } from 'server/services/agentRuntime/mcp/runtimeConfig'; +import { usesSessionWorkspaceGatewayExecution } from 'server/services/agentRuntime/mcp/sessionPod'; import type { RequestUserIdentity } from 'server/lib/get-user'; import { getLogger } from 'server/lib/logger'; import type { AgentSessionToolRule } from 'server/services/types/agentSessionConfig'; +import type { + CapabilityPolicyConfig, + CustomAgentCreationPolicyConfig, + AgentRuntimeConfig, +} from 'server/services/types/agentRuntimeConfig'; import AgentPolicyService from './PolicyService'; +import type { ResolvedAgentCapabilityAccess } from './PolicyService'; import type { AgentApprovalMode, AgentApprovalPolicy, AgentCapabilityKey, AgentToolAuditRecord } from './types'; -import type { ResolvedMcpServer } from 'server/services/ai/mcp/types'; +import type { AgentCapabilityCatalogId } from './capabilityCatalog'; +import type { ResolvedMcpServer } from 'server/services/agentRuntime/mcp/types'; +import AgentRuntimeConfigService from 'server/services/agentRuntime/config/agentRuntimeConfig'; import { assertSafeWorkspaceMutationCommand, isReadOnlyWorkspaceCommand } from './sandboxExecSafety'; import { buildProposedFileChanges, buildResultFileChanges, didToolResultFail } from './fileChanges'; import type { AgentFileChangeData } from './types'; @@ -46,6 +54,13 @@ import { import { getSessionWorkspaceCatalogEntriesForRuntimeTool } from './sandboxToolCatalog'; import { SessionWorkspaceGatewayUnavailableError } from './errors'; import AgentSandboxService from './SandboxService'; +import { + registerLifecycleDiagnosticFixTools, + registerLifecycleDiagnosticReadTools, + type LifecycleDiagnosticGithubSafety, +} from './diagnosticTools'; +import { YamlConfigParser } from 'server/lib/yamlConfigParser'; +import type { LifecycleConfig } from 'server/models/yaml/Config'; type ToolExecutionHooks = { onToolStarted?: (audit: AgentToolAuditRecord) => Promise; @@ -61,6 +76,7 @@ type SessionWorkspaceGatewayTimeouts = { 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 REDACTED_MCP_DEFAULT_ARG = '******'; const WORKSPACE_EXEC_INPUT_SCHEMA = { type: 'object', required: ['command'], @@ -119,6 +135,7 @@ const WORKSPACE_EDIT_FILE_INPUT_SCHEMA = { }, }, } as const; +const LIFECYCLE_CONFIG_WRITE_PATTERNS = ['lifecycle.yaml', 'lifecycle.yml']; const PUBLISH_HTTP_INPUT_SCHEMA = { type: 'object', required: ['port'], @@ -133,6 +150,14 @@ const PUBLISH_HTTP_INPUT_SCHEMA = { }, } as const; +function toAiJsonSchema(schema: unknown) { + return jsonSchema(schema as any); +} + +function toAiDynamicTool(config: unknown) { + return dynamicTool(config as any); +} + function resolvePrimaryRepo(session: AgentSession): string | undefined { const primaryRepo = (session.workspaceRepos || []).find((repo) => repo.primary)?.repo; if (primaryRepo) { @@ -142,6 +167,84 @@ function resolvePrimaryRepo(session: AgentSession): string | undefined { return session.selectedServices?.[0]?.repo || undefined; } +function resolvePrimaryBranch(session: AgentSession): string | null { + const primaryWorkspaceRepo = + (session.workspaceRepos || []).find((repo) => repo.primary) || session.workspaceRepos?.[0]; + if (primaryWorkspaceRepo?.branch) { + return primaryWorkspaceRepo.branch; + } + + return session.selectedServices?.[0]?.branch || null; +} + +function addReferencedFile(files: Set, value: unknown) { + if (typeof value !== 'string') { + return; + } + + const normalized = value.trim().replace(/^\/+/, '').replace(/^\.\//, ''); + if (normalized) { + files.add(normalized); + } +} + +function collectLifecycleConfigReferencedFiles(config: LifecycleConfig | null | undefined): string[] { + const files = new Set(); + + for (const service of config?.services || []) { + const candidate = service as Record; + addReferencedFile(files, candidate.github?.docker?.app?.dockerfilePath); + addReferencedFile(files, candidate.github?.docker?.init?.dockerfilePath); + addReferencedFile(files, candidate.helm?.docker?.app?.dockerfilePath); + addReferencedFile(files, candidate.helm?.docker?.init?.dockerfilePath); + addReferencedFile(files, candidate.helm?.envMapping?.app?.path); + addReferencedFile(files, candidate.helm?.envMapping?.init?.path); + + for (const valueFile of candidate.helm?.chart?.valueFiles || []) { + addReferencedFile(files, valueFile); + } + } + + return [...files]; +} + +async function resolveLifecycleDiagnosticGithubSafety({ + session, + repoFullName, + config, +}: { + session: AgentSession; + repoFullName?: string; + config?: AgentRuntimeConfig | null; +}): Promise { + const allowedBranch = resolvePrimaryBranch(session); + const allowedWritePatterns = [ + ...new Set([...LIFECYCLE_CONFIG_WRITE_PATTERNS, ...(config?.allowedWritePatterns || [])]), + ]; + const safety: LifecycleDiagnosticGithubSafety = { + allowedBranch, + allowedWritePatterns, + excludedFilePatterns: config?.excludedFilePatterns || [], + referencedFiles: [], + }; + + if (!repoFullName || !allowedBranch) { + return safety; + } + + try { + const lifecycleConfig = await new YamlConfigParser().parseYamlConfigFromBranch(repoFullName, allowedBranch); + safety.referencedFiles = collectLifecycleConfigReferencedFiles(lifecycleConfig); + } catch (error) { + getLogger().warn( + { error, repo: repoFullName, branch: allowedBranch }, + `AgentExec: lifecycle config references unavailable repo=${repoFullName} branch=${allowedBranch}` + ); + } + + return safety; +} + function resolveToolApprovalMode({ toolRules, toolKey, @@ -155,6 +258,43 @@ function resolveToolApprovalMode({ return rule?.mode || capabilityMode; } +function isCatalogCapabilityAllowed( + resolvedCapabilityAccess: ResolvedAgentCapabilityAccess[] | undefined, + capabilityId: AgentCapabilityCatalogId +): boolean { + if (!resolvedCapabilityAccess) { + return false; + } + + return resolvedCapabilityAccess.some((entry) => entry.capabilityId === capabilityId && entry.allowed); +} + +function selectedMcpConnectionRefs(connectionRefs?: string[]): Set | undefined { + if (connectionRefs === undefined) { + return undefined; + } + + return new Set(connectionRefs.map((connectionRef) => connectionRef.trim()).filter(Boolean)); +} + +function redactMcpDefaultArgs( + args: Record, + defaultArgs: Record | undefined +): Record { + if (!defaultArgs || Object.keys(defaultArgs).length === 0) { + return args; + } + + const redacted = { ...args }; + for (const key of Object.keys(defaultArgs)) { + if (key in redacted) { + redacted[key] = REDACTED_MCP_DEFAULT_ARG; + } + } + + return redacted; +} + function resolveSessionWorkspaceGatewayBaseUrl(session: AgentSession): string | null { if (!session.podName || !session.namespace || session.status !== 'active') { return null; @@ -192,6 +332,7 @@ async function resolveSessionWorkspaceGatewayServer( const discoveredTools = await client.listTools(timeouts.discoveryTimeoutMs); return { + scope: 'session', slug: 'sandbox', name: 'Session Workspace', transport: { type: 'http', url }, @@ -360,6 +501,8 @@ function registerChatWorkspaceExecTool({ capabilityKey, description, readOnly, + catalogCapabilityId, + resolvedCapabilityAccess, }: { tools: ToolSet; session: AgentSession; @@ -373,7 +516,13 @@ function registerChatWorkspaceExecTool({ capabilityKey: AgentCapabilityKey; description: string; readOnly: boolean; + catalogCapabilityId: AgentCapabilityCatalogId; + resolvedCapabilityAccess?: ResolvedAgentCapabilityAccess[]; }) { + if (!isCatalogCapabilityAllowed(resolvedCapabilityAccess, catalogCapabilityId)) { + return; + } + const toolKey = buildAgentToolKey(SESSION_WORKSPACE_SERVER_SLUG, toolName); const mode = resolveToolApprovalMode({ toolRules, @@ -385,9 +534,9 @@ function registerChatWorkspaceExecTool({ return; } - tools[toolKey] = dynamicTool({ + tools[toolKey] = toAiDynamicTool({ description, - inputSchema: jsonSchema(WORKSPACE_EXEC_INPUT_SCHEMA), + inputSchema: toAiJsonSchema(WORKSPACE_EXEC_INPUT_SCHEMA), needsApproval: mode === 'require_approval', execute: async (input, context) => { const args = (input as Record) || {}; @@ -467,6 +616,8 @@ function registerChatWorkspaceFileTool({ toolName, inputSchema, description, + catalogCapabilityId, + resolvedCapabilityAccess, }: { tools: ToolSet; session: AgentSession; @@ -479,7 +630,13 @@ function registerChatWorkspaceFileTool({ toolName: string; inputSchema: Record; description: string; + catalogCapabilityId: AgentCapabilityCatalogId; + resolvedCapabilityAccess?: ResolvedAgentCapabilityAccess[]; }) { + if (!isCatalogCapabilityAllowed(resolvedCapabilityAccess, catalogCapabilityId)) { + return; + } + const toolKey = buildAgentToolKey(SESSION_WORKSPACE_SERVER_SLUG, toolName); const capabilityKey: AgentCapabilityKey = 'workspace_write'; const mode = resolveToolApprovalMode({ @@ -492,9 +649,9 @@ function registerChatWorkspaceFileTool({ return; } - tools[toolKey] = dynamicTool({ + tools[toolKey] = toAiDynamicTool({ description, - inputSchema: jsonSchema(inputSchema), + inputSchema: toAiJsonSchema(inputSchema), needsApproval: mode === 'require_approval', onInputAvailable: async ({ input, toolCallId }) => { if (!toolCallId) { @@ -599,6 +756,7 @@ function registerChatPublishHttpTool({ requestGitHubToken, hooks, toolRules, + resolvedCapabilityAccess, }: { tools: ToolSet; session: AgentSession; @@ -607,8 +765,13 @@ function registerChatPublishHttpTool({ requestGitHubToken?: string | null; hooks?: ToolExecutionHooks; toolRules?: AgentSessionToolRule[]; + resolvedCapabilityAccess?: ResolvedAgentCapabilityAccess[]; }) { const toolKey = buildAgentToolKey(LIFECYCLE_BUILTIN_SERVER_SLUG, CHAT_PUBLISH_HTTP_TOOL_NAME); + if (!isCatalogCapabilityAllowed(resolvedCapabilityAccess, 'preview_publish')) { + return; + } + const capabilityKey: AgentCapabilityKey = 'deploy_k8s_mutation'; const mode = resolveToolApprovalMode({ toolRules, @@ -620,10 +783,10 @@ function registerChatPublishHttpTool({ return; } - tools[toolKey] = dynamicTool({ + tools[toolKey] = toAiDynamicTool({ 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), + inputSchema: toAiJsonSchema(PUBLISH_HTTP_INPUT_SCHEMA), needsApproval: mode === 'require_approval', execute: async (input, context) => { const args = (input as Record) || {}; @@ -685,6 +848,7 @@ function registerChatWorkspaceTools({ requestGitHubToken, hooks, toolRules, + resolvedCapabilityAccess, }: { tools: ToolSet; session: AgentSession; @@ -694,6 +858,7 @@ function registerChatWorkspaceTools({ requestGitHubToken?: string | null; hooks?: ToolExecutionHooks; toolRules?: AgentSessionToolRule[]; + resolvedCapabilityAccess?: ResolvedAgentCapabilityAccess[]; }) { registerChatWorkspaceExecTool({ tools, @@ -708,6 +873,8 @@ function registerChatWorkspaceTools({ capabilityKey: 'read', description: buildWorkspaceReadonlyExecDescription(SESSION_WORKSPACE_SERVER_NAME), readOnly: true, + catalogCapabilityId: 'read_context', + resolvedCapabilityAccess, }); registerChatWorkspaceExecTool({ tools, @@ -722,6 +889,8 @@ function registerChatWorkspaceTools({ capabilityKey: 'shell_exec', description: buildWorkspaceMutationExecDescription(SESSION_WORKSPACE_SERVER_NAME), readOnly: false, + catalogCapabilityId: 'workspace_shell', + resolvedCapabilityAccess, }); registerChatWorkspaceFileTool({ tools, @@ -736,6 +905,8 @@ function registerChatWorkspaceTools({ 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.', + catalogCapabilityId: 'workspace_files', + resolvedCapabilityAccess, }); registerChatWorkspaceFileTool({ tools, @@ -750,6 +921,8 @@ function registerChatWorkspaceTools({ 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.', + catalogCapabilityId: 'workspace_files', + resolvedCapabilityAccess, }); } @@ -776,24 +949,25 @@ function registerGenericMcpTool({ }) { const toolKey = buildAgentToolKey(server.slug, exposedToolName); - tools[toolKey] = dynamicTool({ + tools[toolKey] = toAiDynamicTool({ description, - inputSchema: jsonSchema(discoveredTool.inputSchema as Record), + inputSchema: toAiJsonSchema(discoveredTool.inputSchema as Record), needsApproval: mode === 'require_approval', onInputAvailable: async ({ input, toolCallId }) => { if (!toolCallId) { return; } - const args = applyMcpDefaultToolArgs( + const runtimeArgs = applyMcpDefaultToolArgs( discoveredTool.inputSchema as Record, server.defaultArgs, (input as Record) || {} ); + const auditArgs = redactMcpDefaultArgs(runtimeArgs, server.defaultArgs); const changes = buildProposedFileChanges({ toolCallId, sourceTool: exposedToolName, - input: args, + input: auditArgs, previewChars: await getFileChangePreviewChars(), }); @@ -803,32 +977,43 @@ function registerGenericMcpTool({ }, execute: async (input, context) => { const toolCallId = context?.toolCallId; - const args = applyMcpDefaultToolArgs( + const runtimeArgs = applyMcpDefaultToolArgs( discoveredTool.inputSchema as Record, server.defaultArgs, (input as Record) || {} ); + const auditArgs = redactMcpDefaultArgs(runtimeArgs, server.defaultArgs); const audit: AgentToolAuditRecord = { source: 'mcp', serverSlug: server.slug, toolName: exposedToolName, toolCallId, - args, + args: auditArgs, capabilityKey, }; await hooks?.onToolStarted?.(audit); + const mcpSecretSources = [ + { + compiledConfig: { + env: server.env, + defaultArgs: server.defaultArgs, + }, + transport: server.transport, + }, + ]; const client = new McpClientManager(); try { await client.connect(server.transport, server.timeout); - const result = await client.callTool(discoveredTool.name, args, server.timeout); - const failed = result.isError || didToolResultFail(result); + const rawResult = await client.callTool(discoveredTool.name, runtimeArgs, server.timeout); + const failed = rawResult.isError || didToolResultFail(rawResult); + const result = failed ? sanitizeMcpResult(rawResult, mcpSecretSources) : rawResult; if (toolCallId) { const changes = buildResultFileChanges({ toolCallId, sourceTool: exposedToolName, - input: args, + input: auditArgs, result, failed, previewChars: await getFileChangePreviewChars(), @@ -845,17 +1030,18 @@ function registerGenericMcpTool({ }); return result; } catch (error) { + const errorMessage = sanitizeMcpErrorMessage(error, mcpSecretSources); getLogger().warn( - { error }, + { error: errorMessage }, `AgentExec: mcp tool failed sessionId=${session.uuid} server=${server.slug} tool=${exposedToolName}` ); if (toolCallId) { const changes = buildResultFileChanges({ toolCallId, sourceTool: exposedToolName, - input: args, + input: auditArgs, result: { - error: error instanceof Error ? error.message : String(error), + error: errorMessage, }, failed: true, previewChars: await getFileChangePreviewChars(), @@ -868,11 +1054,11 @@ function registerGenericMcpTool({ await hooks?.onToolFinished?.({ ...audit, result: { - error: error instanceof Error ? error.message : String(error), + error: errorMessage, }, status: 'failed', }); - throw error; + throw new Error(errorMessage); } finally { await client.close(); } @@ -897,15 +1083,22 @@ export default class AgentCapabilityService { session: AgentSession; repoFullName?: string; approvalPolicy: AgentApprovalPolicy; + capabilityPolicy?: CapabilityPolicyConfig; + customAgentCreationPolicy?: CustomAgentCreationPolicyConfig; }> { const session = await this.getOwnedSession(sessionUuid, userIdentity.userId); const repoFullName = resolvePrimaryRepo(session); - const approvalPolicy = await AgentPolicyService.getEffectivePolicy(repoFullName); + const [approvalPolicy, effectiveAgentConfig] = await Promise.all([ + AgentPolicyService.getEffectivePolicy(repoFullName), + AgentRuntimeConfigService.getInstance().getEffectiveConfig(repoFullName), + ]); return { session, repoFullName, approvalPolicy, + capabilityPolicy: effectiveAgentConfig.capabilityPolicy, + customAgentCreationPolicy: effectiveAgentConfig.customAgentCreationPolicy, }; } @@ -919,6 +1112,8 @@ export default class AgentCapabilityService { requestGitHubToken, hooks, toolRules, + resolvedCapabilityAccess, + selectedRuntimeMcpConnectionRefs, }: { session: AgentSession; repoFullName?: string; @@ -929,9 +1124,19 @@ export default class AgentCapabilityService { requestGitHubToken?: string | null; hooks?: ToolExecutionHooks; toolRules?: AgentSessionToolRule[]; + resolvedCapabilityAccess?: ResolvedAgentCapabilityAccess[]; + selectedRuntimeMcpConnectionRefs?: string[]; }): Promise { const tools: ToolSet = {}; const chatWorkspaceRuntimeReady = isChatWorkspaceRuntimeReady(session); + const effectiveAgentConfig = await AgentRuntimeConfigService.getInstance().getEffectiveConfig(repoFullName); + const lifecycleDiagnosticGithubSafety = session.buildUuid + ? await resolveLifecycleDiagnosticGithubSafety({ + session, + repoFullName, + config: effectiveAgentConfig, + }) + : undefined; if (session.sessionKind === 'chat') { registerChatWorkspaceTools({ @@ -943,6 +1148,7 @@ export default class AgentCapabilityService { requestGitHubToken, hooks, toolRules, + resolvedCapabilityAccess, }); registerChatPublishHttpTool({ @@ -953,9 +1159,29 @@ export default class AgentCapabilityService { requestGitHubToken, hooks, toolRules, + resolvedCapabilityAccess, }); } + registerLifecycleDiagnosticReadTools({ + tools, + session, + approvalPolicy, + hooks, + toolRules, + resolvedCapabilityAccess, + githubSafety: lifecycleDiagnosticGithubSafety, + }); + registerLifecycleDiagnosticFixTools({ + tools, + session, + approvalPolicy, + hooks, + toolRules, + resolvedCapabilityAccess, + githubSafety: lifecycleDiagnosticGithubSafety, + }); + const mcpConfigService = new McpConfigService(); const [repoServers, workspaceGatewayServer] = await Promise.all([ mcpConfigService.resolveServers(repoFullName, undefined, userIdentity), @@ -966,7 +1192,11 @@ export default class AgentCapabilityService { executionTimeoutMs: workspaceToolExecutionTimeoutMs, }), ]); - const resolvedRepoServers = repoServers.flatMap((server) => { + const selectedRuntimeMcpRefs = selectedMcpConnectionRefs(selectedRuntimeMcpConnectionRefs); + const selectedRepoServers = selectedRuntimeMcpRefs + ? repoServers.filter((server) => selectedRuntimeMcpRefs.has(`${server.scope}:${server.slug}`)) + : repoServers; + const resolvedRepoServers = selectedRepoServers.flatMap((server) => { if (!usesSessionWorkspaceGatewayExecution(server.transport)) { return [server]; } @@ -1000,6 +1230,10 @@ export default class AgentCapabilityService { entry.toolName, entry.annotations || discoveredTool.annotations ); + if (!isCatalogCapabilityAllowed(resolvedCapabilityAccess, entry.catalogCapabilityId)) { + continue; + } + const mode = resolveToolApprovalMode({ toolRules, toolKey: entry.toolKey, @@ -1011,9 +1245,9 @@ export default class AgentCapabilityService { } if (entry.toolName === SESSION_WORKSPACE_READONLY_TOOL_NAME) { - const inputSchema = jsonSchema(WORKSPACE_EXEC_INPUT_SCHEMA); + const inputSchema = toAiJsonSchema(WORKSPACE_EXEC_INPUT_SCHEMA); - tools[entry.toolKey] = dynamicTool({ + tools[entry.toolKey] = toAiDynamicTool({ description: entry.description, inputSchema, needsApproval: mode === 'require_approval', @@ -1072,9 +1306,9 @@ export default class AgentCapabilityService { } if (entry.toolName === SESSION_WORKSPACE_MUTATION_TOOL_NAME) { - const inputSchema = jsonSchema(WORKSPACE_EXEC_INPUT_SCHEMA); + const inputSchema = toAiJsonSchema(WORKSPACE_EXEC_INPUT_SCHEMA); - tools[entry.toolKey] = dynamicTool({ + tools[entry.toolKey] = toAiDynamicTool({ description: entry.description, inputSchema, needsApproval: mode === 'require_approval', @@ -1159,6 +1393,12 @@ export default class AgentCapabilityService { discoveredTool.name, discoveredTool.annotations ); + const catalogCapabilityId: AgentCapabilityCatalogId = + capabilityKey === 'external_mcp_read' ? 'external_mcp_read' : 'external_mcp_write'; + if (!isCatalogCapabilityAllowed(resolvedCapabilityAccess, catalogCapabilityId)) { + continue; + } + const toolName = buildAgentToolKey(server.slug, discoveredTool.name); const mode = resolveToolApprovalMode({ toolRules, diff --git a/src/server/services/agent/ChatSessionService.ts b/src/server/services/agent/ChatSessionService.ts index ee575736..ee115ff5 100644 --- a/src/server/services/agent/ChatSessionService.ts +++ b/src/server/services/agent/ChatSessionService.ts @@ -18,35 +18,102 @@ 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 { normalizeSessionWorkspaceRepo, type AgentSessionWorkspaceRepo } from 'server/lib/agentSession/workspace'; import AgentSession from 'server/models/AgentSession'; import AgentThread from 'server/models/AgentThread'; -import { AgentChatStatus, AgentSessionKind, AgentWorkspaceStatus } from 'shared/constants'; +import { AgentChatStatus, AgentSessionKind, AgentWorkspaceStatus, BuildKind } from 'shared/constants'; import AgentProviderRegistry from './ProviderRegistry'; import AgentSourceService from './SourceService'; +import AgentThreadRuntimeControlsService, { + type AgentThreadRuntimeControlChoiceInput, + type ValidatedEntryRuntimeControlChoices, +} from './ThreadRuntimeControlsService'; import type { ResolvedAgentSessionWorkspaceStorageIntent } from 'server/lib/agentSession/runtimeConfig'; +export interface AgentBuildContextChatMetadata { + buildUuid: string; + buildKind: BuildKind | null; + namespace: string | null; + baseBuildUuid: string | null; + revision: string | null; + pullRequest: { + fullName: string | null; + branchName: string | null; + pullRequestNumber: number | null; + } | null; + contextFreshAt: string; +} + export interface CreateChatSessionOptions { userId: string; userIdentity?: RequestUserIdentity; + provider?: string; model?: string; + runtimeControlChoices?: AgentThreadRuntimeControlChoiceInput; workspaceStorage?: ResolvedAgentSessionWorkspaceStorageIntent; + buildContext?: AgentBuildContextChatMetadata; +} + +function buildContextWorkspaceRepos(buildContext?: AgentBuildContextChatMetadata): AgentSessionWorkspaceRepo[] { + const repo = buildContext?.pullRequest?.fullName?.trim(); + const branch = buildContext?.pullRequest?.branchName?.trim(); + if (!repo || !branch) { + return []; + } + + return [ + normalizeSessionWorkspaceRepo( + { + repo, + repoUrl: `https://github.com/${repo}.git`, + branch, + revision: buildContext.revision, + }, + true + ), + ]; } export default class AgentChatSessionService { static async createChatSession(opts: CreateChatSessionOptions): Promise { const sessionUuid = uuid(); + const requestedProvider = opts.provider?.trim() || undefined; const requestedModelId = opts.model?.trim() || undefined; const providerUserIdentity = { userId: opts.userId, githubUsername: opts.userIdentity?.githubUsername || null, }; + const workspaceRepos = buildContextWorkspaceRepos(opts.buildContext); + const primaryWorkspaceRepo = workspaceRepos.find((repo) => repo.primary) || workspaceRepos[0]; const selection = await AgentProviderRegistry.resolveSelection({ + repoFullName: primaryWorkspaceRepo?.repo, + requestedProvider, requestedModelId, }); await AgentProviderRegistry.getRequiredStoredApiKey({ provider: selection.provider, userIdentity: providerUserIdentity, }); + let validatedRuntimeControlChoices: ValidatedEntryRuntimeControlChoices | null = null; + if (opts.runtimeControlChoices) { + if (!opts.userIdentity) { + throw new Error('userIdentity is required when runtimeControlChoices are provided.'); + } + + validatedRuntimeControlChoices = await AgentThreadRuntimeControlsService.validateEntryChoices({ + userIdentity: opts.userIdentity, + agentId: opts.runtimeControlChoices.agentId, + source: { + adapter: 'blank_workspace', + input: opts.buildContext?.buildUuid ? { buildUuid: opts.buildContext.buildUuid } : {}, + }, + defaults: { + provider: requestedProvider || null, + model: requestedModelId || null, + }, + runtimeControlChoices: opts.runtimeControlChoices, + }); + } const finalizedSession = await AgentSession.transaction(async (trx) => { const session = await AgentSession.query(trx).insertAndFetch({ @@ -54,7 +121,7 @@ export default class AgentChatSessionService { defaultThreadId: null, defaultModel: selection.modelId, defaultHarness: 'lifecycle_ai_sdk', - buildUuid: null, + buildUuid: opts.buildContext?.buildUuid ?? null, buildKind: null, sessionKind: AgentSessionKind.CHAT, userId: opts.userId, @@ -69,7 +136,7 @@ export default class AgentChatSessionService { keepAttachedServicesOnSessionNode: null, devModeSnapshots: {}, forwardedAgentSecretProviders: [], - workspaceRepos: [], + workspaceRepos, selectedServices: [], skillPlan: EMPTY_AGENT_SESSION_SKILL_PLAN, } as unknown as Partial); @@ -80,10 +147,21 @@ export default class AgentChatSessionService { isDefault: true, metadata: { sessionUuid: session.uuid, + ...(validatedRuntimeControlChoices?.selectedAgentMetadataPatch || {}), + ...(validatedRuntimeControlChoices?.runtimeControlChoices + ? { + runtimeControlChoices: validatedRuntimeControlChoices.runtimeControlChoices, + } + : {}), }, } as Partial); - await AgentSourceService.createSessionSource(session, { trx, workspaceStorage: opts.workspaceStorage }); + await AgentSourceService.createSessionSource(session, { + trx, + workspaceStorage: opts.workspaceStorage, + buildContext: opts.buildContext, + defaultProvider: selection.provider, + }); return AgentSession.query(trx).patchAndFetchById(session.id, { defaultThreadId: defaultThread.id, diff --git a/src/server/services/agent/CustomAgentDefinitionService.ts b/src/server/services/agent/CustomAgentDefinitionService.ts new file mode 100644 index 00000000..9a29a523 --- /dev/null +++ b/src/server/services/agent/CustomAgentDefinitionService.ts @@ -0,0 +1,536 @@ +/** + * 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 { RequestUserIdentity } from 'server/lib/get-user'; +import AgentDefinition from 'server/models/AgentDefinition'; +import AgentRuntimeConfigService from 'server/services/agentRuntime/config/agentRuntimeConfig'; +import { v4 as uuid } from 'uuid'; +import type { CapabilityPolicyConfig, CustomAgentCreationPolicyConfig } from 'server/services/types/agentRuntimeConfig'; +import { + listAgentCapabilityCatalogEntries, + type AgentCapabilityCatalogId, + type AgentCapabilityCategory, + type AgentCapabilitySourceKind, +} from './capabilityCatalog'; +import { agentDefinitionRowToContract } from './AgentDefinitionRegistry'; +import type { + AgentDefinitionContract, + AgentDefinitionModelPreference, + AgentDefinitionResourcePolicy, + UserAgentDefinitionListFilters, + UserAgentDefinitionResourceBehavior, + UserAgentDefinitionUpsertInput, +} from './agentDefinitionTypes'; +import AgentPolicyService, { type AgentCapabilityAccessReason } from './PolicyService'; +import AgentProviderRegistry from './ProviderRegistry'; + +type CustomAgentDefinitionErrorCode = + | 'not_found' + | 'invalid_input' + | 'model_unavailable' + | 'creation_unavailable' + | 'creator_capability_reserved' + | AgentCapabilityAccessReason; +type CustomAgentDefinitionUserIdentity = Pick & + Partial>; +export type CustomAgentCreationUnavailableReason = 'creation_disabled' | 'creation_restricted'; +export interface CustomAgentCreationStatus { + canCreate: boolean; + creationUnavailableReason: CustomAgentCreationUnavailableReason | null; +} + +const CAPABILITY_UNAVAILABLE_MESSAGE = + 'Some selected capabilities are no longer available. Review the list and save again.'; +const MODEL_UNAVAILABLE_MESSAGE = 'Selected model is no longer available. Choose another model and save again.'; +const CREATION_UNAVAILABLE_MESSAGE = 'Custom agent creation is not available. Ask an admin for access.'; +const CAPABILITY_DENIAL_REASONS = new Set([ + 'unknown_capability', + 'admin_only', + 'system_only', + 'disabled', + 'source_incompatible', +]); + +export interface UserAgentDefinitionCapabilityDisplaySummary { + name: string; + description: string | null; +} + +export interface UserAgentDefinitionCapability { + capabilityId: AgentCapabilityCatalogId; + label: string; + description: string; + category: AgentCapabilityCategory; + toolCount: number; + resourceCount: number; + requiresWorkspace: boolean; + tools: UserAgentDefinitionCapabilityDisplaySummary[]; + resources: UserAgentDefinitionCapabilityDisplaySummary[]; +} + +export interface UserAgentDefinitionPublicContract { + id: string; + version: number; + name: string; + description: string | null; + instructions: string; + capabilityIds: AgentCapabilityCatalogId[]; + modelPreference: AgentDefinitionModelPreference | null; + resourceBehavior: UserAgentDefinitionResourceBehavior; + status: 'active' | 'archived'; +} + +export class CustomAgentDefinitionServiceError extends Error { + code: CustomAgentDefinitionErrorCode; + + constructor(code: CustomAgentDefinitionErrorCode, message: string) { + super(message); + this.name = 'CustomAgentDefinitionServiceError'; + this.code = code; + } +} + +function trimRequired(value: string, fieldName: string): string { + const trimmed = value.trim(); + if (!trimmed) { + throw new CustomAgentDefinitionServiceError('invalid_input', `${fieldName} is required.`); + } + return trimmed; +} + +function trimNullable(value: string | null | undefined): string | null { + if (typeof value !== 'string') { + return null; + } + const trimmed = value.trim(); + return trimmed || null; +} + +function dedupeCapabilities(capabilityRefs: AgentCapabilityCatalogId[] | undefined): AgentCapabilityCatalogId[] { + return [...new Set(capabilityRefs || [])]; +} + +function normalizeModelPreference( + modelPreference: AgentDefinitionModelPreference | null | undefined +): AgentDefinitionModelPreference | null { + if (!modelPreference) { + return null; + } + + const provider = trimNullable(modelPreference.provider); + const model = trimNullable(modelPreference.model); + + if (!provider && !model) { + return null; + } + + return { provider, model }; +} + +function resourcePolicyForBehavior( + resourceBehavior: UserAgentDefinitionResourceBehavior +): AgentDefinitionResourcePolicy { + if (resourceBehavior === 'current_workspace_when_available') { + return { + sourceKinds: ['freeform_chat', 'workspace_session'], + workspaceRequired: false, + sandboxRequired: false, + }; + } + + return { + sourceKinds: ['freeform_chat'], + workspaceRequired: false, + sandboxRequired: false, + }; +} + +function resourceBehaviorForPolicy(resourcePolicy: AgentDefinitionResourcePolicy): UserAgentDefinitionResourceBehavior { + return resourcePolicy.sourceKinds.includes('workspace_session') ? 'current_workspace_when_available' : 'chat_only'; +} + +function sourceKindsForResourceBehavior( + resourceBehavior: UserAgentDefinitionResourceBehavior +): AgentCapabilitySourceKind[] { + if (resourceBehavior === 'current_workspace_when_available') { + return ['freeform_chat', 'workspace_session']; + } + + return ['freeform_chat']; +} + +function normalizeIdentifierList(values: string[] | undefined, options: { lowercase?: boolean } = {}): Set { + return new Set( + (values || []) + .map((value) => value.trim()) + .filter(Boolean) + .map((value) => (options.lowercase ? value.toLowerCase() : value)) + ); +} + +function isCustomAgentCreatorAllowed( + userIdentity: CustomAgentDefinitionUserIdentity, + policy?: CustomAgentCreationPolicyConfig +): boolean { + const mode = policy?.mode || 'enabled'; + if (mode === 'enabled') { + return true; + } + if (mode === 'disabled') { + return false; + } + if (mode === 'admins_only') { + return Boolean(userIdentity.roles?.includes('admin')); + } + + const allowedUserIds = normalizeIdentifierList(policy?.allowedUserIds); + const allowedGithubUsernames = normalizeIdentifierList(policy?.allowedGithubUsernames, { lowercase: true }); + const githubUsername = userIdentity.githubUsername?.trim().toLowerCase(); + + return ( + allowedUserIds.has(userIdentity.userId) || Boolean(githubUsername && allowedGithubUsernames.has(githubUsername)) + ); +} + +function getCustomAgentCreationUnavailableReason( + userIdentity: CustomAgentDefinitionUserIdentity, + policy?: CustomAgentCreationPolicyConfig +): CustomAgentCreationUnavailableReason | null { + if (isCustomAgentCreatorAllowed(userIdentity, policy)) { + return null; + } + + return policy?.mode === 'disabled' ? 'creation_disabled' : 'creation_restricted'; +} + +function isCreatorCapabilityAvailable( + capabilityId: AgentCapabilityCatalogId, + policy?: CustomAgentCreationPolicyConfig +): boolean { + return policy?.capabilityAvailability?.[capabilityId] !== 'reserved'; +} + +function displayNameFromIdentifier(value: string): string { + return value + .replace(/[_./-]+/g, ' ') + .replace(/\s+/g, ' ') + .trim() + .replace(/\b\w/g, (character) => character.toUpperCase()); +} + +function normalizeUpsertInput(input: UserAgentDefinitionUpsertInput): { + name: string; + description: string | null; + instructionAddendum: string; + capabilityRefs: AgentCapabilityCatalogId[]; + modelPreference: AgentDefinitionModelPreference | null; + resourcePolicy: AgentDefinitionResourcePolicy; +} { + const capabilityRefs = dedupeCapabilities(input.capabilityRefs); + + return { + name: trimRequired(input.name, 'Name'), + description: trimNullable(input.description), + instructionAddendum: trimRequired(input.instructionAddendum, 'Instructions'), + capabilityRefs, + modelPreference: normalizeModelPreference(input.modelPreference), + resourcePolicy: resourcePolicyForBehavior(input.resourceBehavior), + }; +} + +function selectDeniedCapabilityReason( + results: ReturnType +): AgentCapabilityAccessReason { + const policyDenied = results.find((result) => result.reason && result.reason !== 'source_incompatible'); + const reason = policyDenied?.reason || results[0]?.reason; + return reason && CAPABILITY_DENIAL_REASONS.has(reason) ? reason : 'unknown_capability'; +} + +export function serializeUserAgentDefinition(definition: AgentDefinitionContract): UserAgentDefinitionPublicContract { + return { + id: definition.id, + version: definition.version, + name: definition.name, + description: definition.description || null, + instructions: definition.instructionAddendum || '', + capabilityIds: definition.optionalCapabilityRefs?.length + ? definition.optionalCapabilityRefs + : definition.capabilityRefs, + modelPreference: definition.modelPreference || null, + resourceBehavior: resourceBehaviorForPolicy(definition.resourcePolicy), + status: definition.status === 'archived' ? 'archived' : 'active', + }; +} + +export class CustomAgentDefinitionService { + async getUserDefinitionCreationStatus({ + userIdentity, + }: { + userIdentity: RequestUserIdentity; + }): Promise { + const effectiveConfig = await AgentRuntimeConfigService.getInstance().getEffectiveConfig(); + const creationUnavailableReason = getCustomAgentCreationUnavailableReason( + userIdentity, + effectiveConfig.customAgentCreationPolicy + ); + + return { + canCreate: creationUnavailableReason === null, + creationUnavailableReason, + }; + } + + async listUserSelectableCapabilities({ + userIdentity, + resourceBehavior, + }: { + userIdentity: RequestUserIdentity; + resourceBehavior: UserAgentDefinitionResourceBehavior; + }): Promise { + const effectiveConfig = await AgentRuntimeConfigService.getInstance().getEffectiveConfig(); + if (getCustomAgentCreationUnavailableReason(userIdentity, effectiveConfig.customAgentCreationPolicy)) { + return []; + } + + const sourceKinds = sourceKindsForResourceBehavior(resourceBehavior); + + return listAgentCapabilityCatalogEntries() + .filter((entry) => entry.userSelectable) + .filter((entry) => isCreatorCapabilityAvailable(entry.id, effectiveConfig.customAgentCreationPolicy)) + .filter((entry) => + sourceKinds.some((sourceKind) => { + const access = AgentPolicyService.resolveCapabilityAccess({ + capabilityId: entry.id, + capabilityPolicy: effectiveConfig.capabilityPolicy, + definitionOwnerKind: 'user', + sourceKind, + }); + return access.allowed; + }) + ) + .map((entry) => { + const runtimeToolNames = entry.toolKeys || []; + const tools = runtimeToolNames.map((runtimeToolName) => ({ + name: displayNameFromIdentifier(runtimeToolName), + description: null, + })); + const resources = (entry.resourceGrants || []).map((resourceGrant) => ({ + name: displayNameFromIdentifier(resourceGrant), + description: null, + })); + + return { + capabilityId: entry.id, + label: entry.label, + description: entry.description, + category: entry.category, + toolCount: tools.length, + resourceCount: resources.length, + requiresWorkspace: Boolean( + entry.sourceKinds?.includes('workspace_session') && !entry.sourceKinds.includes('freeform_chat') + ), + tools, + resources, + }; + }); + } + + async listUserDefinitions({ + userId, + filters = {}, + }: { + userId: string; + filters?: UserAgentDefinitionListFilters; + }): Promise { + const rows = await AgentDefinition.query() + .where({ + ownerKind: 'user', + ownerUserId: userId, + status: filters.status || 'active', + }) + .orderBy('updatedAt', 'desc'); + + return rows.map(agentDefinitionRowToContract); + } + + async getUserDefinition(definitionId: string, userId: string): Promise { + const row = await this.findActiveUserDefinitionRow(definitionId, userId); + return agentDefinitionRowToContract(row); + } + + async createUserDefinition( + userIdentity: CustomAgentDefinitionUserIdentity, + input: UserAgentDefinitionUpsertInput + ): Promise { + const normalized = normalizeUpsertInput(input); + await this.validateNormalizedInput(normalized, userIdentity); + const row = await AgentDefinition.query().insert({ + definitionId: `custom.${uuid()}`, + version: 1, + ownerKind: 'user', + ownerUserId: userIdentity.userId, + ownerOrganizationId: null, + name: normalized.name, + description: normalized.description, + instructionRefs: [], + instructionAddendum: normalized.instructionAddendum, + capabilityRefs: normalized.capabilityRefs, + requiredCapabilityRefs: [], + optionalCapabilityRefs: normalized.capabilityRefs, + resourcePolicy: normalized.resourcePolicy, + modelPreference: normalized.modelPreference, + status: 'active', + codeOwned: false, + readOnly: false, + }); + + return agentDefinitionRowToContract(row); + } + + async updateUserDefinition( + definitionId: string, + userIdentity: CustomAgentDefinitionUserIdentity, + input: UserAgentDefinitionUpsertInput + ): Promise { + const existing = await this.findActiveUserDefinitionRow(definitionId, userIdentity.userId); + const normalized = normalizeUpsertInput(input); + await this.validateNormalizedInput(normalized, userIdentity); + const row = await AgentDefinition.query().patchAndFetchById(existing.id, { + version: existing.version + 1, + name: normalized.name, + description: normalized.description, + instructionAddendum: normalized.instructionAddendum, + capabilityRefs: normalized.capabilityRefs, + requiredCapabilityRefs: [], + optionalCapabilityRefs: normalized.capabilityRefs, + resourcePolicy: normalized.resourcePolicy, + modelPreference: normalized.modelPreference, + codeOwned: false, + readOnly: false, + }); + + return agentDefinitionRowToContract(row); + } + + async archiveUserDefinition(definitionId: string, userId: string): Promise { + const existing = await this.findActiveUserDefinitionRow(definitionId, userId); + const row = await AgentDefinition.query().patchAndFetchById(existing.id, { status: 'archived' }); + return agentDefinitionRowToContract(row); + } + + private async validateNormalizedInput( + input: ReturnType, + userIdentity: CustomAgentDefinitionUserIdentity + ): Promise { + const effectiveConfig = await AgentRuntimeConfigService.getInstance().getEffectiveConfig(); + this.validateCustomAgentCreator(userIdentity, effectiveConfig.customAgentCreationPolicy); + this.validateCapabilityRefs( + input.capabilityRefs, + input.resourcePolicy, + effectiveConfig.capabilityPolicy, + effectiveConfig.customAgentCreationPolicy + ); + await this.validateModelPreference(input.modelPreference, userIdentity); + } + + private validateCustomAgentCreator( + userIdentity: CustomAgentDefinitionUserIdentity, + customAgentCreationPolicy?: CustomAgentCreationPolicyConfig + ): void { + if (isCustomAgentCreatorAllowed(userIdentity, customAgentCreationPolicy)) { + return; + } + + throw new CustomAgentDefinitionServiceError('creation_unavailable', CREATION_UNAVAILABLE_MESSAGE); + } + + private validateCapabilityRefs( + capabilityRefs: AgentCapabilityCatalogId[], + resourcePolicy: AgentDefinitionResourcePolicy, + capabilityPolicy?: CapabilityPolicyConfig, + customAgentCreationPolicy?: CustomAgentCreationPolicyConfig + ): void { + if (capabilityRefs.length === 0) { + return; + } + + const reservedCapability = capabilityRefs.find( + (capabilityId) => !isCreatorCapabilityAvailable(capabilityId, customAgentCreationPolicy) + ); + if (reservedCapability) { + throw new CustomAgentDefinitionServiceError('creator_capability_reserved', CAPABILITY_UNAVAILABLE_MESSAGE); + } + + const accessBySource = resourcePolicy.sourceKinds.map((sourceKind) => + AgentPolicyService.resolveCapabilitySetAccess(capabilityRefs, { + capabilityPolicy, + definitionOwnerKind: 'user', + sourceKind: sourceKind as AgentCapabilitySourceKind, + }) + ); + + for (const capabilityId of capabilityRefs) { + const results = accessBySource.flatMap((sourceResults) => + sourceResults.filter((result) => result.capabilityId === capabilityId) + ); + if (results.some((result) => result.allowed)) { + continue; + } + + throw new CustomAgentDefinitionServiceError( + selectDeniedCapabilityReason(results), + CAPABILITY_UNAVAILABLE_MESSAGE + ); + } + } + + private async validateModelPreference( + modelPreference: AgentDefinitionModelPreference | null, + userIdentity: CustomAgentDefinitionUserIdentity + ): Promise { + if (!modelPreference) { + return; + } + + const models = await AgentProviderRegistry.listAvailableModelsForUser({ userIdentity }); + const modelAvailable = models.some( + (model) => + (!modelPreference.provider || model.provider === modelPreference.provider) && + (!modelPreference.model || model.modelId === modelPreference.model) + ); + + if (!modelAvailable) { + throw new CustomAgentDefinitionServiceError('model_unavailable', MODEL_UNAVAILABLE_MESSAGE); + } + } + + private async findActiveUserDefinitionRow(definitionId: string, userId: string): Promise { + const row = await AgentDefinition.query().findOne({ + definitionId, + ownerKind: 'user', + ownerUserId: userId, + status: 'active', + }); + + if (!row) { + throw new CustomAgentDefinitionServiceError('not_found', 'Agent not found.'); + } + + return row; + } +} + +export const customAgentDefinitionService = new CustomAgentDefinitionService(); diff --git a/src/server/services/agent/MessageStore.ts b/src/server/services/agent/MessageStore.ts index 4c1f152b..d9e0cd9a 100644 --- a/src/server/services/agent/MessageStore.ts +++ b/src/server/services/agent/MessageStore.ts @@ -33,9 +33,28 @@ import { 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 AGENT_SWITCH_METADATA_KIND = 'agent_switch'; export const DEFAULT_AGENT_MESSAGE_PAGE_LIMIT = 50; export const MAX_AGENT_MESSAGE_PAGE_LIMIT = 100; +export type AgentSwitchEventMetadata = { + kind: typeof AGENT_SWITCH_METADATA_KIND; + actor: { + userId: string; + label: string; + }; + beforeAgent: { + id: string; + label: string; + }; + afterAgent: { + id: string; + label: string; + }; + appliesTo: 'future_runs'; + occurredAt: string; +}; + function toAgentUiMessage(message: AgentMessage): AgentUIMessage { return toUiMessageFromCanonicalInput( { @@ -60,6 +79,10 @@ function normalizeMessageId(value: unknown): string | null { return typeof value === 'string' && value.trim() ? value.trim() : null; } +function isAgentSwitchMessage(message: AgentMessage): boolean { + return message.role === 'system' && message.metadata?.kind === AGENT_SWITCH_METADATA_KIND; +} + function getIncomingMessageId(message: Pick): string | null { return normalizeMessageId(message.id); } @@ -121,6 +144,18 @@ function buildStoredCanonicalMessage( }; } +function resolveStoredRunId( + role: CanonicalAgentInputMessage['role'], + row: AgentMessage | undefined, + runId?: number | null +): number | null { + if (role !== 'assistant') { + return row?.runId ?? null; + } + + return row?.runId ?? runId ?? null; +} + async function loadExistingMessagesForIncomingIds( threadId: number, messages: CanonicalAgentInputMessage[], @@ -161,7 +196,7 @@ function serializeCanonicalAgentMessage( threadUuid: string, runUuid?: string | null ): CanonicalAgentMessage | null { - if (message.role !== 'user' && message.role !== 'assistant') { + if (message.role !== 'user' && message.role !== 'assistant' && !isAgentSwitchMessage(message)) { return null; } @@ -182,8 +217,9 @@ function serializeCanonicalAgentMessage( normalizeMessageId(message.metadata?.[CLIENT_MESSAGE_ID_METADATA_KEY]), threadId: threadUuid, runId: runUuid || normalizeMessageId(enrichedMessage.runUuid), - role: message.role, + role: message.role as CanonicalAgentMessage['role'], parts, + ...(isAgentSwitchMessage(message) ? { metadata: message.metadata || {} } : {}), createdAt: enrichedMessage.createdAt || null, }; } @@ -252,7 +288,13 @@ export default class AgentMessageStore { .alias('message') .leftJoin('agent_runs as run', 'message.runId', 'run.id') .where('message.threadId', thread.id) - .whereIn('message.role', ['user', 'assistant']) + .where((builder) => { + builder.whereIn('message.role', ['user', 'assistant']).orWhere((systemBuilder) => { + systemBuilder + .where('message.role', 'system') + .whereRaw('"message"."metadata"->>? = ?', ['kind', AGENT_SWITCH_METADATA_KIND]); + }); + }) .select('message.*', 'run.uuid as runUuid') .orderBy('message.createdAt', 'desc') .orderBy('message.id', 'desc') @@ -332,6 +374,47 @@ export default class AgentMessageStore { }); } + static async createAgentSwitchEvent({ + thread, + actor, + beforeAgent, + afterAgent, + occurredAt = new Date().toISOString(), + trx, + }: { + thread: Pick; + actor: { userId: string; label?: string | null }; + beforeAgent: { id: string; label: string }; + afterAgent: { id: string; label: string }; + occurredAt?: string; + trx?: Transaction; + }): Promise { + const actorLabel = actor.label?.trim() || 'You'; + const text = `${actorLabel} switched ${beforeAgent.label} -> ${afterAgent.label}. Applies to future runs.`; + const metadata: AgentSwitchEventMetadata = { + kind: AGENT_SWITCH_METADATA_KIND, + actor: { + userId: actor.userId, + label: actorLabel, + }, + beforeAgent, + afterAgent, + appliesTo: 'future_runs', + occurredAt, + }; + + return AgentMessage.query(trx).insertAndFetch({ + uuid: uuid(), + threadId: thread.id, + runId: null, + role: 'system', + parts: [{ type: 'text', text }] as unknown as Record[], + uiMessage: null, + clientMessageId: null, + metadata: metadata as unknown as Record, + }); + } + static async syncCanonicalMessages( threadUuid: string, userId: string, @@ -364,7 +447,7 @@ export default class AgentMessageStore { uiMessage: null, clientMessageId: stored.clientMessageId, metadata, - runId: message.role === 'assistant' && runId ? runId : row?.runId ?? null, + runId: resolveStoredRunId(message.role, row, runId), }; if (!row) { @@ -422,7 +505,7 @@ export default class AgentMessageStore { uiMessage: null, clientMessageId: stored.clientMessageId, metadata: toJsonRecord(stored.metadata), - runId: message.role === 'assistant' && options?.runId ? options.runId : row?.runId ?? null, + runId: resolveStoredRunId(message.role, row, options?.runId), }; if (!row) { @@ -506,7 +589,7 @@ export default class AgentMessageStore { uiMessage: null, clientMessageId: stored.clientMessageId, metadata: toJsonRecord(stored.metadata), - runId: message.role === 'assistant' && options?.runId ? options.runId : row?.runId ?? null, + runId: resolveStoredRunId(message.role, row, options?.runId), }; if (!row) { diff --git a/src/server/services/agent/PolicyService.ts b/src/server/services/agent/PolicyService.ts index 55a23caf..e3b25d2f 100644 --- a/src/server/services/agent/PolicyService.ts +++ b/src/server/services/agent/PolicyService.ts @@ -14,7 +14,16 @@ * limitations under the License. */ -import AIAgentConfigService from 'server/services/aiAgentConfig'; +import AgentRuntimeConfigService from 'server/services/agentRuntime/config/agentRuntimeConfig'; +import type { AgentDefinitionOwnerKind } from './agentDefinitionTypes'; +import type { CapabilityPolicyConfig, CustomAgentCreationPolicyConfig } from 'server/services/types/agentRuntimeConfig'; +import { + getAgentCapabilityCatalogEntry, + isAgentCapabilityCatalogId, + type AgentCapabilityAvailability, + type AgentCapabilityCatalogEntry, + type AgentCapabilitySourceKind, +} from './capabilityCatalog'; import type { AgentApprovalMode, AgentApprovalPolicy, AgentCapabilityKey } from './types'; import { DEFAULT_AGENT_APPROVAL_POLICY } from './types'; @@ -29,9 +38,37 @@ type ApprovalPolicyConfig = Partial & { rules?: Partial>; }; +export type AgentCapabilityAccessReason = + | 'admin_only' + | 'system_only' + | 'disabled' + | 'unknown_capability' + | 'source_incompatible' + | 'creator_capability_reserved'; + +export type ResolveCapabilityAccessInput = { + capabilityId: string; + capabilityPolicy?: CapabilityPolicyConfig; + customAgentCreationPolicy?: CustomAgentCreationPolicyConfig; + approvalPolicy?: AgentApprovalPolicy; + definitionOwnerKind: AgentDefinitionOwnerKind; + requesterIsAdmin?: boolean; + sourceKind?: AgentCapabilitySourceKind; +}; + +export type ResolvedAgentCapabilityAccess = { + capabilityId: string; + entry?: AgentCapabilityCatalogEntry; + configuredAvailability?: AgentCapabilityAvailability; + effectiveAvailability?: AgentCapabilityAvailability; + allowed: boolean; + reason?: AgentCapabilityAccessReason; + approvalMode?: AgentApprovalMode; +}; + export default class AgentPolicyService { static async getEffectivePolicy(repoFullName?: string): Promise { - const config = await AIAgentConfigService.getInstance().getEffectiveConfig(repoFullName); + const config = await AgentRuntimeConfigService.getInstance().getEffectiveConfig(repoFullName); const configured = (config as { approvalPolicy?: ApprovalPolicyConfig }).approvalPolicy; return { @@ -89,4 +126,102 @@ export default class AgentPolicyService { static modeForCapability(policy: AgentApprovalPolicy, capabilityKey: AgentCapabilityKey): AgentApprovalMode { return policy.rules[capabilityKey] || policy.defaultMode; } + + static resolveCapabilityAccess(input: ResolveCapabilityAccessInput): ResolvedAgentCapabilityAccess { + const { capabilityId } = input; + if (!isAgentCapabilityCatalogId(capabilityId)) { + return { + capabilityId, + allowed: false, + reason: 'unknown_capability', + }; + } + + const entry = getAgentCapabilityCatalogEntry(capabilityId); + const configuredAvailability = input.capabilityPolicy?.availability?.[capabilityId]; + const effectiveAvailability = configuredAvailability || entry.defaultAvailability; + const approvalPolicy = input.approvalPolicy || DEFAULT_AGENT_APPROVAL_POLICY; + const approvalMode = entry.runtimeCapabilityKey + ? AgentPolicyService.modeForCapability(approvalPolicy, entry.runtimeCapabilityKey) + : entry.defaultApprovalMode; + + if (input.sourceKind && entry.sourceKinds && !entry.sourceKinds.includes(input.sourceKind)) { + return { + capabilityId, + entry, + configuredAvailability, + effectiveAvailability, + allowed: false, + reason: 'source_incompatible', + approvalMode, + }; + } + + if (effectiveAvailability === 'disabled') { + return { + capabilityId, + entry, + configuredAvailability, + effectiveAvailability, + allowed: false, + reason: 'disabled', + approvalMode, + }; + } + + if (effectiveAvailability === 'system_only' && input.definitionOwnerKind !== 'system') { + return { + capabilityId, + entry, + configuredAvailability, + effectiveAvailability, + allowed: false, + reason: 'system_only', + approvalMode, + }; + } + + if (effectiveAvailability === 'admin_only' && input.definitionOwnerKind === 'user') { + return { + capabilityId, + entry, + configuredAvailability, + effectiveAvailability, + allowed: false, + reason: 'admin_only', + approvalMode, + }; + } + + if ( + input.definitionOwnerKind === 'user' && + input.customAgentCreationPolicy?.capabilityAvailability?.[capabilityId] === 'reserved' + ) { + return { + capabilityId, + entry, + configuredAvailability, + effectiveAvailability, + allowed: false, + reason: 'creator_capability_reserved', + approvalMode, + }; + } + + return { + capabilityId, + entry, + configuredAvailability, + effectiveAvailability, + allowed: true, + approvalMode, + }; + } + + static resolveCapabilitySetAccess( + capabilityIds: readonly string[], + input: Omit + ): ResolvedAgentCapabilityAccess[] { + return capabilityIds.map((capabilityId) => AgentPolicyService.resolveCapabilityAccess({ ...input, capabilityId })); + } } diff --git a/src/server/services/agent/ProviderRegistry.ts b/src/server/services/agent/ProviderRegistry.ts index 18ea471c..86232b9d 100644 --- a/src/server/services/agent/ProviderRegistry.ts +++ b/src/server/services/agent/ProviderRegistry.ts @@ -18,9 +18,9 @@ import type { LanguageModel } from 'ai'; import { createAnthropic } from '@ai-sdk/anthropic'; import { createGoogleGenerativeAI } from '@ai-sdk/google'; import { createOpenAI } from '@ai-sdk/openai'; -import AIAgentConfigService from 'server/services/aiAgentConfig'; +import AgentRuntimeConfigService from 'server/services/agentRuntime/config/agentRuntimeConfig'; import UserApiKeyService from 'server/services/userApiKey'; -import { transformProviderModels } from 'server/services/ai/utils/modelTransformation'; +import { transformProviderModels } from 'server/services/agentRuntime/models/modelTransformation'; import type { RequestUserIdentity } from 'server/lib/get-user'; import { getLogger } from 'server/lib/logger'; import type { AgentModelSummary, AgentResolvedModelSelection } from './types'; @@ -48,6 +48,13 @@ export class MissingAgentProviderApiKeyError extends Error { } } +export class AgentModelSelectionError extends Error { + constructor(message: string) { + super(message); + this.name = 'AgentModelSelectionError'; + } +} + function getProviderInstance(provider: AgentResolvedModelSelection['provider'], apiKey: string) { switch (provider) { case 'anthropic': @@ -72,7 +79,7 @@ export function resolveRequestedModelSelection( : undefined; if (models.length === 0) { - throw new Error('No enabled agent models are configured'); + throw new AgentModelSelectionError('No enabled agent models are configured'); } if (normalizedRequestedProvider && requestedModelId) { @@ -80,7 +87,7 @@ export function resolveRequestedModelSelection( (model) => model.provider === normalizedRequestedProvider && model.modelId === requestedModelId ); if (!matched) { - throw new Error(`Model ${requestedProvider}:${requestedModelId} is not enabled`); + throw new AgentModelSelectionError(`Model ${requestedProvider}:${requestedModelId} is not enabled`); } return { @@ -92,11 +99,11 @@ export function resolveRequestedModelSelection( if (requestedModelId) { const matches = models.filter((model) => model.modelId === requestedModelId); if (matches.length === 0) { - throw new Error(`Model ${requestedModelId} is not enabled`); + throw new AgentModelSelectionError(`Model ${requestedModelId} is not enabled`); } if (matches.length > 1) { - throw new Error(`Model id ${requestedModelId} is ambiguous; provider is required`); + throw new AgentModelSelectionError(`Model id ${requestedModelId} is ambiguous; provider is required`); } return { @@ -108,7 +115,7 @@ export function resolveRequestedModelSelection( if (normalizedRequestedProvider) { const providerModels = models.filter((model) => model.provider === normalizedRequestedProvider); if (providerModels.length === 0) { - throw new Error(`Provider ${requestedProvider} has no enabled models`); + throw new AgentModelSelectionError(`Provider ${requestedProvider} has no enabled models`); } const defaultProviderModel = providerModels.find((model) => model.default) || providerModels[0]; @@ -131,7 +138,7 @@ export default class AgentProviderRegistry { } static async listAvailableModels(repoFullName?: string): Promise { - const config = await AIAgentConfigService.getInstance().getEffectiveConfig(repoFullName); + const config = await AgentRuntimeConfigService.getInstance().getEffectiveConfig(repoFullName); return transformProviderModels(config.providers || []).flatMap((model) => { const provider = normalizeModelProvider(model.provider); if (!provider) { @@ -184,7 +191,7 @@ export default class AgentProviderRegistry { }): Promise> { let config; try { - config = await AIAgentConfigService.getInstance().getEffectiveConfig(repoFullName); + config = await AgentRuntimeConfigService.getInstance().getEffectiveConfig(repoFullName); } catch (error) { getLogger().warn( { error, repoFullName }, diff --git a/src/server/services/agent/RunAdmissionService.ts b/src/server/services/agent/RunAdmissionService.ts index 5d15e3ec..7e9f3dfa 100644 --- a/src/server/services/agent/RunAdmissionService.ts +++ b/src/server/services/agent/RunAdmissionService.ts @@ -24,6 +24,7 @@ import type { AgentRunRuntimeOptions, CanonicalAgentRunMessageInput } from './ca import AgentMessageStore from './MessageStore'; import AgentRunEventService from './RunEventService'; import { ActiveAgentRunError, InvalidAgentRunDefaultsError, TERMINAL_RUN_STATUSES } from './RunService'; +import { isAgentRunPlanSnapshotV1, type AgentRunPlanSnapshotV1 } from './runPlanTypes'; function buildPolicySnapshot( policy: AgentApprovalPolicy, @@ -39,6 +40,12 @@ function buildPolicySnapshot( }; } +export function assertRunPlanSnapshot(value: unknown): asserts value is AgentRunPlanSnapshotV1 { + if (!isAgentRunPlanSnapshotV1(value)) { + throw new InvalidAgentRunDefaultsError('Agent run plan snapshot is required.'); + } +} + export default class AgentRunAdmissionService { static async createQueuedRunWithMessage({ thread, @@ -53,6 +60,7 @@ export default class AgentRunAdmissionService { resolvedModel, sandboxRequirement, runtimeOptions, + runPlanSnapshot, }: { thread: AgentThread; session: AgentSession; @@ -66,6 +74,7 @@ export default class AgentRunAdmissionService { resolvedModel: string; sandboxRequirement?: Record; runtimeOptions?: AgentRunRuntimeOptions; + runPlanSnapshot: AgentRunPlanSnapshotV1; }): Promise<{ run: AgentRun; message: AgentMessage; created: boolean }> { if (!resolvedHarness?.trim()) { throw new InvalidAgentRunDefaultsError('Agent run harness is required.'); @@ -78,26 +87,6 @@ export default class AgentRunAdmissionService { } 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(); @@ -124,6 +113,8 @@ export default class AgentRunAdmissionService { } } + assertRunPlanSnapshot(runPlanSnapshot); + const activeRun = await AgentRun.query(trx) .where({ sessionId: session.id }) .whereNotIn('status', TERMINAL_RUN_STATUSES) @@ -134,6 +125,27 @@ export default class AgentRunAdmissionService { throw new ActiveAgentRunError(); } + 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), + runPlanSnapshot, + error: null, + }; const queuedRun = await AgentRun.query(trx).insertAndFetch(record); const storedMessage = await AgentMessageStore.insertUserMessageForRun(thread, queuedRun, message, trx); await AgentThread.query(trx).patchAndFetchById(thread.id, { diff --git a/src/server/services/agent/RunEventService.ts b/src/server/services/agent/RunEventService.ts index 65570e58..4705368c 100644 --- a/src/server/services/agent/RunEventService.ts +++ b/src/server/services/agent/RunEventService.ts @@ -995,7 +995,7 @@ export default class AgentRunEventService { this.requireExecutionOwner( lockedRun.uuid || String(run.id), options.executionOwner, - lockedRun as Pick + lockedRun as unknown as Pick ); } diff --git a/src/server/services/agent/RunExecutor.ts b/src/server/services/agent/RunExecutor.ts index ff09add2..3cddd398 100644 --- a/src/server/services/agent/RunExecutor.ts +++ b/src/server/services/agent/RunExecutor.ts @@ -33,6 +33,9 @@ import { AgentRunObservabilityTracker, buildMessageObservabilityMetadataPatch } import AgentProviderRegistry from './ProviderRegistry'; import AgentRunQueueService from './RunQueueService'; import AgentRunService from './RunService'; +import AgentRunPlanResolver from './RunPlanResolver'; +import AgentSourceService from './SourceService'; +import { isAgentRunPlanSnapshotV1 } from './runPlanTypes'; import type { AgentFileChangeData, AgentUIMessage } from './types'; import { applyApprovalResponsesToFileChangeParts, buildResultFileChanges } from './fileChanges'; import { AgentRunTerminalFailure, SessionWorkspaceGatewayUnavailableError } from './errors'; @@ -162,7 +165,8 @@ function classifyTerminalRunFailure({ } function readRunMaxIterations(run?: AgentRun): number | null { - const runtimeOptions = run?.policySnapshot?.runtimeOptions; + const snapshot = isAgentRunPlanSnapshotV1(run?.runPlanSnapshot) ? run.runPlanSnapshot : null; + const runtimeOptions = snapshot?.runtime.runtimeOptions || run?.policySnapshot?.runtimeOptions; if (!runtimeOptions || typeof runtimeOptions !== 'object' || Array.isArray(runtimeOptions)) { return null; } @@ -204,14 +208,34 @@ export default class AgentRunExecutor { dispatchAttemptId?: string; onFileChange?: (change: AgentFileChangeData) => Promise | void; }) { - const { repoFullName, approvalPolicy } = await AgentCapabilityService.resolveSessionContext( + const { repoFullName, approvalPolicy: contextApprovalPolicy } = await AgentCapabilityService.resolveSessionContext( session.uuid, userIdentity ); + const existingRunPlan = isAgentRunPlanSnapshotV1(existingRun?.runPlanSnapshot) ? existingRun.runPlanSnapshot : null; + let executionRunPlan = existingRunPlan; + let pendingRunPlan: Awaited> | null = null; + if (!existingRun) { + const source = await AgentSourceService.getSessionSource(session.id); + if (!source || source.status !== 'ready') { + throw new Error('Session source is not ready yet.'); + } + pendingRunPlan = await AgentRunPlanResolver.resolveForRunAdmission({ + thread, + session, + source, + userIdentity, + requestedProvider, + requestedModel: requestedModelId, + runtimeOptions: {}, + }); + executionRunPlan = pendingRunPlan.runPlanSnapshot; + } + const approvalPolicy = executionRunPlan?.runtime.approvalPolicy || contextApprovalPolicy; const selection = await AgentProviderRegistry.resolveSelection({ repoFullName, - requestedProvider, - requestedModelId, + requestedProvider: executionRunPlan?.model.resolvedProvider || requestedProvider, + requestedModelId: executionRunPlan?.model.resolvedModel || requestedModelId, }); const model = await AgentProviderRegistry.createLanguageModel({ repoFullName, @@ -258,11 +282,17 @@ export default class AgentRunExecutor { }; try { + if (existingRun && !existingRunPlan) { + throw new Error('Agent run plan snapshot is required for execution.'); + } + const tools = await AgentCapabilityService.buildToolSet({ session, repoFullName, userIdentity, approvalPolicy, + resolvedCapabilityAccess: executionRunPlan?.capabilities.resolvedCapabilityAccess ?? [], + selectedRuntimeMcpConnectionRefs: executionRunPlan?.capabilities.selectedRuntimeMcpConnectionRefs, workspaceToolDiscoveryTimeoutMs: runControlPlaneConfig.workspaceToolDiscoveryTimeoutMs, workspaceToolExecutionTimeoutMs: runControlPlaneConfig.workspaceToolExecutionTimeoutMs, requestGitHubToken, @@ -342,11 +372,12 @@ export default class AgentRunExecutor { throw new Error('Agent run execution owner is required.'); } + const fallbackResolvedHarness = existingRun.resolvedHarness || session.defaultHarness || 'lifecycle_ai_sdk'; run = await AgentRunService.startRunForExecutionOwner( existingRun.uuid, existingRun.executionOwner, { - resolvedHarness: existingRun.requestedHarness || session.defaultHarness || 'lifecycle_ai_sdk', + resolvedHarness: existingRunPlan?.runtime.resolvedHarness || fallbackResolvedHarness, provider: selection.provider, model: selection.modelId, sandboxGeneration: existingRun.sandboxGeneration, @@ -354,16 +385,23 @@ export default class AgentRunExecutor { { dispatchAttemptId } ); } else { + const runPlan = pendingRunPlan; + if (!runPlan) { + throw new Error('Agent run plan has not been initialized.'); + } + executionRunPlan = runPlan.runPlanSnapshot; 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, + policy: runPlan.approvalPolicy, + requestedHarness: runPlan.requestedHarness, + requestedProvider: runPlan.requestedProvider, + requestedModel: runPlan.requestedModel, + resolvedHarness: runPlan.resolvedHarness, + resolvedProvider: runPlan.resolvedProvider, + resolvedModel: runPlan.resolvedModel, + sandboxRequirement: runPlan.sandboxRequirement, + runPlanSnapshot: runPlan.runPlanSnapshot, }); const executionOwner = buildDirectExecutionOwner(); const claimedRun = await AgentRunService.claimQueuedRunForExecution(queuedRun.uuid, executionOwner); @@ -374,9 +412,9 @@ export default class AgentRunExecutor { claimedRun.uuid, executionOwner, { - resolvedHarness: session.defaultHarness || 'lifecycle_ai_sdk', - provider: selection.provider, - model: selection.modelId, + resolvedHarness: runPlan.resolvedHarness, + provider: runPlan.resolvedProvider, + model: runPlan.resolvedModel, sandboxGeneration: claimedRun.sandboxGeneration, }, { dispatchAttemptId } @@ -412,7 +450,11 @@ export default class AgentRunExecutor { } const agent = new ToolLoopAgent({ model, - instructions: buildSystemPrompt([runControlPlaneConfig.systemPrompt, sessionPrompt]), + instructions: buildSystemPrompt([ + runControlPlaneConfig.systemPrompt, + executionRunPlan?.prompt.instructionAddendum || undefined, + sessionPrompt, + ]), tools, stopWhen: stepCountIs(runControlPlaneConfig.maxIterations), onStepFinish: async (step) => { diff --git a/src/server/services/agent/RunPlanResolver.ts b/src/server/services/agent/RunPlanResolver.ts new file mode 100644 index 00000000..c46e3b67 --- /dev/null +++ b/src/server/services/agent/RunPlanResolver.ts @@ -0,0 +1,442 @@ +/** + * 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 { createHash } from 'crypto'; +import type AgentSession from 'server/models/AgentSession'; +import type AgentSource from 'server/models/AgentSource'; +import type AgentThread from 'server/models/AgentThread'; +import type { RequestUserIdentity } from 'server/lib/get-user'; +import { getLogger } from 'server/lib/logger'; +import AgentCapabilityService from './CapabilityService'; +import AgentPolicyService from './PolicyService'; +import AgentProviderRegistry from './ProviderRegistry'; +import * as AgentDefinitionRegistry from './AgentDefinitionRegistry'; +import { CustomAgentDefinitionServiceError, customAgentDefinitionService } from './CustomAgentDefinitionService'; +import type { AgentRunRuntimeOptions } from './canonicalMessages'; +import AgentThreadService from './ThreadService'; +import AgentThreadRuntimeControlsService from './ThreadRuntimeControlsService'; +import type { AgentDefinitionContract } from './agentDefinitionTypes'; +import { getAgentCapabilityCatalogEntry, type AgentCapabilityCatalogId } from './capabilityCatalog'; +import type { AgentRunPlanSnapshotV1, AgentRunPlanSourceKind, AgentRunPlanWarning } from './runPlanTypes'; +import { + isSystemAgentDefinitionId, + sourceKindForSystemAgentDefinitionId, + type SystemAgentDefinitionId, +} from './systemAgentDefinitions'; + +export class AgentRunPlanCapabilityUnavailableError extends Error { + constructor(public readonly capabilityId: string, public readonly reason: string | undefined) { + super(`Agent capability "${capabilityId}" is unavailable${reason ? `: ${reason}` : ''}.`); + this.name = 'AgentRunPlanCapabilityUnavailableError'; + } +} + +export class AgentRunPlanAgentUnavailableError extends Error { + constructor( + public readonly agentId: string, + public readonly reason: string, + public readonly details?: Record + ) { + super(`Agent "${agentId}" is unavailable: ${reason}.`); + this.name = 'AgentRunPlanAgentUnavailableError'; + } +} + +function readString(value: unknown): string | null { + return typeof value === 'string' && value.trim() ? value.trim() : null; +} + +function readSessionDefaultProvider(source: AgentSource): string | null { + const defaults = source.input?.defaults; + if (!defaults || typeof defaults !== 'object' || Array.isArray(defaults)) { + return null; + } + + return readString((defaults as Record).provider); +} + +function readRecord(value: unknown): Record { + return value && typeof value === 'object' && !Array.isArray(value) ? (value as Record) : {}; +} + +function hashPromptRefs( + definitionId: string, + instructionRefs: string[], + version: number, + instructionAddendum?: string | null +): string { + return createHash('sha256') + .update( + JSON.stringify({ definitionId, instructionRefs, version, instructionAddendum: instructionAddendum || null }) + ) + .digest('hex'); +} + +function compactSource({ + session, + source, + sourceKind, + repoFullName, + capturedAt, +}: { + session: AgentSession; + source: AgentSource; + sourceKind: AgentRunPlanSourceKind; + repoFullName?: string; + capturedAt: string; +}): AgentRunPlanSnapshotV1['source'] { + const workspaceRepos = Array.isArray(session.workspaceRepos) ? session.workspaceRepos : []; + const selectedServices = Array.isArray(session.selectedServices) ? session.selectedServices : []; + const primaryRepo = workspaceRepos.find((repo) => repo.primary) || workspaceRepos[0] || null; + const primaryService = selectedServices[0] || null; + const sourceInput = readRecord(source.input); + + return { + id: source.uuid || null, + adapter: source.adapter || null, + status: source.status || null, + sessionKind: session.sessionKind || null, + buildUuid: readString(sourceInput.buildUuid) || session.buildUuid || null, + repoFullName: primaryRepo?.repo || repoFullName || null, + branch: primaryRepo?.branch || readString(sourceInput.branchName) || null, + namespace: session.namespace || readString(sourceInput.namespace) || null, + workspaceLayout: { + repoCount: Array.isArray(session.workspaceRepos) ? session.workspaceRepos.length : 0, + primaryRepo: primaryRepo?.repo || repoFullName || null, + selectedServiceCount: Array.isArray(session.selectedServices) ? session.selectedServices.length : 0, + primaryService: primaryService?.name || null, + }, + sandboxRequirements: source.sandboxRequirements || {}, + freshness: { + capturedAt, + preparedAt: source.preparedAt || null, + freshnessSource: source.preparedAt ? 'source' : sourceKind === 'workspace_session' ? 'session' : 'request', + }, + }; +} + +function uniqueCapabilityIds(capabilityIds: readonly AgentCapabilityCatalogId[]): AgentCapabilityCatalogId[] { + return Array.from(new Set(capabilityIds)); +} + +function isKnownCapabilityId(capabilityId: string): capabilityId is AgentCapabilityCatalogId { + try { + getAgentCapabilityCatalogEntry(capabilityId as AgentCapabilityCatalogId); + return true; + } catch { + return false; + } +} + +function warningForUnavailableOptionalCapability( + capability: ReturnType +): AgentRunPlanWarning { + const label = capability.entry?.label || 'Optional capability'; + return { + code: 'optional_capability_unavailable', + message: `${label} is unavailable and was skipped.`, + ...(capability.reason + ? { + detail: { + reason: capability.reason, + }, + } + : {}), + }; +} + +async function resolveSelectedDefinition({ + selectedAgentDefinitionId, + defaultAgentDefinitionId, + userId, + warnings, +}: { + selectedAgentDefinitionId: string | null; + defaultAgentDefinitionId: SystemAgentDefinitionId; + userId: string; + warnings: AgentRunPlanWarning[]; +}): Promise<{ selectedDefinitionId: string; definition: AgentDefinitionContract }> { + if (!selectedAgentDefinitionId) { + return { + selectedDefinitionId: defaultAgentDefinitionId, + definition: await AgentDefinitionRegistry.getSystemAgentDefinition(defaultAgentDefinitionId), + }; + } + + if (isSystemAgentDefinitionId(selectedAgentDefinitionId)) { + return { + selectedDefinitionId: selectedAgentDefinitionId, + definition: await AgentDefinitionRegistry.getSystemAgentDefinition(selectedAgentDefinitionId), + }; + } + + try { + return { + selectedDefinitionId: selectedAgentDefinitionId, + definition: await customAgentDefinitionService.getUserDefinition(selectedAgentDefinitionId, userId), + }; + } catch (error) { + if (!(error instanceof CustomAgentDefinitionServiceError) || error.code !== 'not_found') { + throw error; + } + + warnings.push({ + code: 'selected_agent_unavailable', + message: 'Selected agent is unavailable. The default agent will be used for this run.', + }); + + return { + selectedDefinitionId: defaultAgentDefinitionId, + definition: await AgentDefinitionRegistry.getSystemAgentDefinition(defaultAgentDefinitionId), + }; + } +} + +export default class AgentRunPlanResolver { + static async resolveForRunAdmission({ + thread, + session, + source, + userIdentity, + requestedProvider, + requestedModel, + runtimeOptions = {}, + }: { + thread: AgentThread; + session: AgentSession; + source: AgentSource; + userIdentity: RequestUserIdentity; + requestedProvider?: string | null; + requestedModel?: string | null; + runtimeOptions?: AgentRunRuntimeOptions; + }): Promise<{ + approvalPolicy: Awaited>['approvalPolicy']; + requestedHarness: null; + requestedProvider: string | null; + requestedModel: string | null; + resolvedHarness: 'lifecycle_ai_sdk'; + resolvedProvider: string; + resolvedModel: string; + sandboxRequirement: Record; + runtimeOptions: AgentRunRuntimeOptions; + repoFullName?: string; + runPlanSnapshot: AgentRunPlanSnapshotV1; + }> { + const warnings: AgentRunPlanWarning[] = []; + const { repoFullName, approvalPolicy, capabilityPolicy, customAgentCreationPolicy } = + await AgentCapabilityService.resolveSessionContext(session.uuid, userIdentity); + const requestedHarness = readString(session.defaultHarness); + if (requestedHarness && requestedHarness !== 'lifecycle_ai_sdk') { + const warning = { + code: 'unsupported_harness_default', + message: `Unsupported session harness "${requestedHarness}" was replaced with lifecycle_ai_sdk.`, + }; + warnings.push(warning); + getLogger().warn( + { sessionId: session.uuid, requestedHarness }, + `AgentExec: run plan harness fallback sessionId=${session.uuid} harness=${requestedHarness}` + ); + } + + await AgentDefinitionRegistry.ensureSystemAgentDefinitionsSeeded(); + const defaultAgentDefinitionId = AgentDefinitionRegistry.inferDefaultSystemAgentDefinitionId(session, source); + const selectedAgentDefinitionId = AgentThreadService.getSelectedAgentDefinitionId(thread); + const { selectedDefinitionId, definition } = await resolveSelectedDefinition({ + selectedAgentDefinitionId, + defaultAgentDefinitionId, + userId: userIdentity.userId, + warnings, + }); + const sourceKind = sourceKindForSystemAgentDefinitionId(defaultAgentDefinitionId); + const resolvedProviderRequest = + requestedProvider || definition.modelPreference?.provider || readSessionDefaultProvider(source) || undefined; + const resolvedModelRequest = + requestedModel || definition.modelPreference?.model || readString(session.defaultModel); + if (!resolvedModelRequest) { + throw new Error('Agent run model is required'); + } + + const selection = await AgentProviderRegistry.resolveSelection({ + repoFullName, + requestedProvider: resolvedProviderRequest, + requestedModelId: resolvedModelRequest, + }); + if (definition.status !== 'active') { + throw new AgentRunPlanAgentUnavailableError(selectedDefinitionId, 'disabled_agent'); + } + + if (!definition.resourcePolicy.sourceKinds.includes(sourceKind)) { + throw new AgentRunPlanAgentUnavailableError(selectedDefinitionId, 'source_incompatible', { + sourceKind, + }); + } + + if (definition.resourcePolicy.workspaceRequired && sourceKind !== 'workspace_session') { + throw new AgentRunPlanAgentUnavailableError(selectedDefinitionId, 'workspace_required', { + sourceKind, + }); + } + + if (definition.resourcePolicy.sandboxRequired && sourceKind !== 'workspace_session') { + throw new AgentRunPlanAgentUnavailableError(selectedDefinitionId, 'sandbox_required', { + sourceKind, + }); + } + + const requiredCapabilityRefs = definition.requiredCapabilityRefs || definition.capabilityRefs; + const optionalCapabilityRefs = definition.optionalCapabilityRefs || []; + const runtimeChoices = await AgentThreadRuntimeControlsService.resolveRunAdmissionChoices({ + thread, + userIdentity, + definition, + sourceKind, + capabilityPolicy, + customAgentCreationPolicy, + approvalPolicy, + repoFullName, + }); + const selectedOptionalCapabilityRefs = runtimeChoices.metadataPresent + ? uniqueCapabilityIds( + (runtimeChoices.selectedRuntimeCapabilityIds || []).filter((capabilityId) => + optionalCapabilityRefs.includes(capabilityId) + ) + ) + : optionalCapabilityRefs; + const requiredCapabilityAccess = AgentPolicyService.resolveCapabilitySetAccess(requiredCapabilityRefs, { + capabilityPolicy, + customAgentCreationPolicy, + approvalPolicy, + definitionOwnerKind: definition.owner.kind, + sourceKind, + }); + const blockedCapability = requiredCapabilityAccess.find((capability) => !capability.allowed); + if (blockedCapability) { + throw new AgentRunPlanCapabilityUnavailableError(blockedCapability.capabilityId, blockedCapability.reason); + } + const optionalCapabilityAccess = AgentPolicyService.resolveCapabilitySetAccess(selectedOptionalCapabilityRefs, { + capabilityPolicy, + customAgentCreationPolicy, + approvalPolicy, + definitionOwnerKind: definition.owner.kind, + sourceKind, + }); + const allowedOptionalCapabilityAccess = optionalCapabilityAccess.filter((capability) => capability.allowed); + for (const capability of optionalCapabilityAccess) { + if (!capability.allowed) { + warnings.push(warningForUnavailableOptionalCapability(capability)); + } + } + const resolvedCapabilityAccess = [...requiredCapabilityAccess, ...allowedOptionalCapabilityAccess]; + const provisionalCapabilityIds = uniqueCapabilityIds( + [ + ...requiredCapabilityAccess.map((capability) => capability.capabilityId), + ...allowedOptionalCapabilityAccess.map((capability) => capability.capabilityId), + ].filter(isKnownCapabilityId) + ); + const selectedRuntimeCapabilityIds = runtimeChoices.metadataPresent + ? (provisionalCapabilityIds as AgentRunPlanSnapshotV1['capabilities']['selectedRuntimeCapabilityIds']) + : undefined; + const runtimeChoicesMatchAllowedCapabilities = + !runtimeChoices.metadataPresent || + (runtimeChoices.selectedRuntimeCapabilityIds || []).every((capabilityId) => + provisionalCapabilityIds.includes(capabilityId) + ); + + const capturedAt = new Date().toISOString(); + const runPlanSnapshot: AgentRunPlanSnapshotV1 = { + version: 1, + capturedAt, + agent: { + id: selectedDefinitionId, + label: definition.name, + ownerKind: definition.owner.kind, + version: definition.version, + sourceKind, + resourcePolicy: definition.resourcePolicy, + modelPreference: definition.modelPreference || null, + }, + source: compactSource({ session, source, sourceKind, repoFullName, capturedAt }), + model: { + requestedProvider: requestedProvider || null, + requestedModel: requestedModel || null, + resolvedProvider: selection.provider, + resolvedModel: selection.modelId, + }, + runtime: { + requestedHarness, + resolvedHarness: 'lifecycle_ai_sdk', + sandboxRequirement: source.sandboxRequirements || {}, + runtimeOptions, + approvalPolicy, + }, + prompt: { + instructionRefs: definition.instructionRefs, + instructionAddendum: definition.instructionAddendum || null, + renderedSummary: definition.description || definition.name, + renderedHash: hashPromptRefs( + selectedDefinitionId, + definition.instructionRefs, + definition.version, + definition.instructionAddendum + ), + }, + capabilities: { + provisionalCapabilityIds: + provisionalCapabilityIds as AgentRunPlanSnapshotV1['capabilities']['provisionalCapabilityIds'], + resolvedCapabilityAccess: resolvedCapabilityAccess.map((capability) => ({ + capabilityId: + capability.capabilityId as AgentRunPlanSnapshotV1['capabilities']['provisionalCapabilityIds'][number], + availability: capability.effectiveAvailability || capability.entry?.defaultAvailability || 'disabled', + allowed: capability.allowed, + ...(capability.reason ? { reason: capability.reason } : {}), + ...(capability.entry?.runtimeCapabilityKey + ? { runtimeCapabilityKey: capability.entry.runtimeCapabilityKey } + : {}), + ...(capability.approvalMode ? { approvalMode: capability.approvalMode } : {}), + })), + ...(runtimeChoices.metadataPresent + ? { + selectedRuntimeToolChoiceIds: runtimeChoicesMatchAllowedCapabilities + ? runtimeChoices.selectedRuntimeToolChoiceIds || [] + : [], + selectedRuntimeMcpChoiceIds: runtimeChoicesMatchAllowedCapabilities + ? runtimeChoices.selectedRuntimeMcpChoiceIds || [] + : [], + selectedRuntimeCapabilityIds, + selectedRuntimeMcpConnectionRefs: runtimeChoicesMatchAllowedCapabilities + ? runtimeChoices.selectedRuntimeMcpConnectionRefs || [] + : [], + } + : {}), + }, + warnings, + }; + + return { + approvalPolicy, + requestedHarness: null, + requestedProvider: requestedProvider || null, + requestedModel: requestedModel || null, + resolvedHarness: 'lifecycle_ai_sdk', + resolvedProvider: selection.provider, + resolvedModel: selection.modelId, + sandboxRequirement: source.sandboxRequirements || {}, + runtimeOptions, + repoFullName, + runPlanSnapshot, + }; + } +} diff --git a/src/server/services/agent/RunService.ts b/src/server/services/agent/RunService.ts index e7864f67..c2bab599 100644 --- a/src/server/services/agent/RunService.ts +++ b/src/server/services/agent/RunService.ts @@ -22,6 +22,8 @@ import AgentSession from 'server/models/AgentSession'; import type { AgentApprovalPolicy, AgentRunStatus, AgentRunUsageSummary } from './types'; import type { AgentUiMessageChunk } from './streamChunks'; import AgentRunEventService from './RunEventService'; +import { isAgentRunPlanSnapshotV1, type AgentRunPlanSnapshotV1 } from './runPlanTypes'; +import { serializeRunPlanSummary } from './runPlanSummary'; import { AgentRunOwnershipLostError } from './AgentRunOwnershipLostError'; import { DEFAULT_AGENT_SESSION_DISPATCH_RECOVERY_LIMIT, @@ -178,6 +180,7 @@ export default class AgentRunService { resolvedProvider, resolvedModel, sandboxRequirement, + runPlanSnapshot, }: { thread: AgentThread; session: AgentSession; @@ -189,6 +192,7 @@ export default class AgentRunService { resolvedProvider: string; resolvedModel: string; sandboxRequirement?: Record; + runPlanSnapshot: AgentRunPlanSnapshotV1; }): Promise { if (!resolvedHarness?.trim()) { throw new InvalidAgentRunDefaultsError('Agent run harness is required.'); @@ -199,6 +203,9 @@ export default class AgentRunService { if (!resolvedModel?.trim()) { throw new InvalidAgentRunDefaultsError('Agent run model is required.'); } + if (!isAgentRunPlanSnapshotV1(runPlanSnapshot)) { + throw new InvalidAgentRunDefaultsError('Agent run plan snapshot is required.'); + } const now = new Date().toISOString(); const record: PartialModelObject = { @@ -219,6 +226,7 @@ export default class AgentRunService { startedAt: null, usageSummary: {}, policySnapshot: policy as unknown as Record, + runPlanSnapshot: runPlanSnapshot as unknown as Record, error: null, }; @@ -261,12 +269,14 @@ export default class AgentRunService { provider, model, policy, + runPlanSnapshot, }: { thread: AgentThread; session: AgentSession; provider: string; model: string; policy: AgentApprovalPolicy; + runPlanSnapshot: AgentRunPlanSnapshotV1; }): Promise { const run = await this.createQueuedRun({ thread, @@ -278,6 +288,7 @@ export default class AgentRunService { resolvedHarness: session.defaultHarness || 'lifecycle_ai_sdk', resolvedProvider: provider, resolvedModel: model, + runPlanSnapshot, }); return this.startRun(run.uuid, { @@ -304,6 +315,12 @@ export default class AgentRunService { return run || undefined; } + static async hasActiveRun(threadId: number, trx?: Transaction): Promise { + const activeRun = await AgentRun.query(trx).where({ threadId }).whereNotIn('status', TERMINAL_RUN_STATUSES).first(); + + return Boolean(activeRun); + } + static async listRunsNeedingDispatch({ limit, now = new Date(), @@ -937,6 +954,7 @@ export default class AgentRunService { cancelledAt: run.cancelledAt, usageSummary: run.usageSummary || {}, policySnapshot: run.policySnapshot || {}, + runPlan: serializeRunPlanSummary(run.runPlanSnapshot), error: run.error, createdAt: run.createdAt || null, updatedAt: run.updatedAt || null, diff --git a/src/server/services/agent/SessionReadService.ts b/src/server/services/agent/SessionReadService.ts index a21f27ba..375a759b 100644 --- a/src/server/services/agent/SessionReadService.ts +++ b/src/server/services/agent/SessionReadService.ts @@ -20,8 +20,10 @@ 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 { AgentChatStatus, AgentSessionKind } from 'shared/constants'; import AgentThreadService from './ThreadService'; import AgentSandboxService from './SandboxService'; +import AgentUsageService, { type AgentUsageAggregate } from './AgentUsageService'; export const DEFAULT_AGENT_SESSION_LIST_LIMIT = 25; export const MAX_AGENT_SESSION_LIST_LIMIT = 100; @@ -37,6 +39,7 @@ interface SessionRecordRelations { sandbox: AgentSandbox | null; exposures: AgentSandboxExposure[]; defaultThread: AgentThread | null; + usage: AgentUsageAggregate; } function mapSessionStatus(session: AgentSession): 'ready' | 'ended' | 'error' { @@ -44,7 +47,11 @@ function mapSessionStatus(session: AgentSession): 'ready' | 'ended' | 'error' { return 'ended'; } - if (session.status === 'error' || session.workspaceStatus === 'failed') { + if (session.status === 'error' || session.chatStatus === AgentChatStatus.ERROR) { + return 'error'; + } + + if (session.workspaceStatus === 'failed' && session.sessionKind !== AgentSessionKind.CHAT) { return 'error'; } @@ -70,6 +77,16 @@ function buildDerivedSourceInput(session: AgentSession, source: AgentSource) { }; } +function readSourceDefaultProvider(source: AgentSource): string | null { + const defaults = source.input?.defaults; + if (!defaults || typeof defaults !== 'object' || Array.isArray(defaults)) { + return null; + } + + const provider = (defaults as Record).provider; + return typeof provider === 'string' && provider.trim() ? provider.trim() : null; +} + export default class AgentSessionReadService { static async getOwnedSessionRecord(sessionId: string, userId: string) { const session = await AgentSession.query().findOne({ uuid: sessionId, userId }); @@ -127,6 +144,7 @@ export default class AgentSessionReadService { userId: session.userId, ownerGithubUsername: session.ownerGithubUsername, defaults: { + provider: readSourceDefaultProvider(source), model: session.defaultModel || session.model, harness: session.defaultHarness, }, @@ -175,6 +193,7 @@ export default class AgentSessionReadService { createdAt: null, updatedAt: null, }, + usage: relations.usage, }; } @@ -187,7 +206,7 @@ export default class AgentSessionReadService { const defaultThreadIds = sessions .map((session) => session.defaultThreadId) .filter((threadId): threadId is number => Number.isInteger(threadId)); - const [sources, sandboxes, defaultThreads, fallbackThreads] = await Promise.all([ + const [sources, sandboxes, defaultThreads, fallbackThreads, usageBySessionId] = 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([]), @@ -196,6 +215,7 @@ export default class AgentSessionReadService { .where({ isDefault: true }) .whereNull('archivedAt') .orderBy('createdAt', 'asc'), + AgentUsageService.aggregateSessionsUsage(sessionIds), ]); const sourceBySessionId = new Map(); for (const source of sources) { @@ -249,6 +269,7 @@ export default class AgentSessionReadService { sandbox, defaultThread, exposures: sandbox ? exposuresBySandboxId.get(sandbox.id) || [] : [], + usage: usageBySessionId.get(session.id) || AgentUsageService.aggregateRuns([]), }); }); } diff --git a/src/server/services/agent/SettingsService.ts b/src/server/services/agent/SettingsService.ts index e8a3e52f..fcb3d973 100644 --- a/src/server/services/agent/SettingsService.ts +++ b/src/server/services/agent/SettingsService.ts @@ -15,10 +15,10 @@ */ import type { RequestUserIdentity } from 'server/lib/get-user'; -import AIAgentConfigService from 'server/services/aiAgentConfig'; +import AgentRuntimeConfigService from 'server/services/agentRuntime/config/agentRuntimeConfig'; import UserApiKeyService from 'server/services/userApiKey'; -import { McpConfigService } from 'server/services/ai/mcp/config'; -import type { AgentMcpConnection } from 'server/services/ai/mcp/types'; +import { McpConfigService } from 'server/services/agentRuntime/mcp/config'; +import type { AgentMcpConnection } from 'server/services/agentRuntime/mcp/types'; import { normalizeStoredAgentProviderName, type StoredAgentProviderName } from './providerConfig'; export type AgentProviderCredentialState = { @@ -36,7 +36,7 @@ export type AgentSettingsSnapshot = { export default class AgentSettingsService { static async getConfiguredProviders(repoFullName?: string): Promise { try { - const config = await AIAgentConfigService.getInstance().getEffectiveConfig(repoFullName); + const config = await AgentRuntimeConfigService.getInstance().getEffectiveConfig(repoFullName); const configuredProviders = (config.providers || []) .map((provider: { name?: unknown; enabled?: unknown }) => provider.enabled !== false && typeof provider.name === 'string' diff --git a/src/server/services/agent/SourceService.ts b/src/server/services/agent/SourceService.ts index 087766df..7a015dea 100644 --- a/src/server/services/agent/SourceService.ts +++ b/src/server/services/agent/SourceService.ts @@ -20,6 +20,7 @@ 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'; +import type { AgentBuildContextChatMetadata } from './ChatSessionService'; function deriveAdapter(session: AgentSession): string { if (session.sessionKind === 'chat') { @@ -56,7 +57,12 @@ export default class AgentSourceService { static async createSessionSource( session: AgentSession, - options: { trx?: Transaction; workspaceStorage?: ResolvedAgentSessionWorkspaceStorageIntent } = {} + options: { + trx?: Transaction; + workspaceStorage?: ResolvedAgentSessionWorkspaceStorageIntent; + buildContext?: AgentBuildContextChatMetadata; + defaultProvider?: string | null; + } = {} ): Promise { const status = deriveStatus(session); const workspaceRepos = session.workspaceRepos ?? []; @@ -71,15 +77,33 @@ export default class AgentSourceService { repos: workspaceRepos, primaryPath: primaryRepo?.mountPath || SESSION_WORKSPACE_ROOT, }; + const metadata = options.buildContext + ? { + buildUuid: options.buildContext.buildUuid, + buildKind: options.buildContext.buildKind, + sessionKind: session.sessionKind, + namespace: options.buildContext.namespace, + baseBuildUuid: options.buildContext.baseBuildUuid, + revision: options.buildContext.revision, + pullRequest: options.buildContext.pullRequest, + contextFreshAt: options.buildContext.contextFreshAt, + } + : { + buildUuid: session.buildUuid, + buildKind: session.buildKind, + sessionKind: session.sessionKind, + }; return AgentSource.query(options.trx).insertAndFetch({ sessionId: session.id, adapter: deriveAdapter(session), status, input: { - buildUuid: session.buildUuid, - buildKind: session.buildKind, - sessionKind: session.sessionKind, + ...metadata, + defaults: { + provider: options.defaultProvider || null, + model: session.defaultModel || session.model || null, + }, ...(options.workspaceStorage?.requestedSize ? { workspace: { @@ -92,11 +116,7 @@ export default class AgentSourceService { kind: 'workspace_snapshot', workspaceLayout, artifactRefs: [], - metadata: { - buildUuid: session.buildUuid, - buildKind: session.buildKind, - sessionKind: session.sessionKind, - }, + metadata, }, sandboxRequirements: { filesystem: 'persistent', diff --git a/src/server/services/agent/ThreadRuntimeControlsService.ts b/src/server/services/agent/ThreadRuntimeControlsService.ts new file mode 100644 index 00000000..99376c51 --- /dev/null +++ b/src/server/services/agent/ThreadRuntimeControlsService.ts @@ -0,0 +1,703 @@ +/** + * 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 { createHash } from 'crypto'; +import type { RequestUserIdentity } from 'server/lib/get-user'; +import type AgentThread from 'server/models/AgentThread'; +import { McpConfigService } from 'server/services/agentRuntime/mcp/config'; +import type { AgentMcpConnection } from 'server/services/agentRuntime/mcp/types'; +import AgentRuntimeConfigService from 'server/services/agentRuntime/config/agentRuntimeConfig'; +import AgentCapabilityService from './CapabilityService'; +import * as AgentDefinitionRegistry from './AgentDefinitionRegistry'; +import { customAgentDefinitionService } from './CustomAgentDefinitionService'; +import AgentPolicyService from './PolicyService'; +import AgentRunService from './RunService'; +import AgentSourceService from './SourceService'; +import AgentThreadService, { type AgentThreadRuntimeControlChoicesMetadata } from './ThreadService'; +import type { AgentDefinitionContract } from './agentDefinitionTypes'; +import { + getAgentCapabilityCatalogEntry, + type AgentCapabilityCatalogId, + type AgentCapabilitySourceKind, +} from './capabilityCatalog'; +import { + isSystemAgentDefinitionId, + sourceKindForSystemAgentDefinitionId, + type SystemAgentDefinitionId, +} from './systemAgentDefinitions'; + +export type AgentThreadRuntimeControlChoice = { + id: string; + label: string; + description: string | null; + required: boolean; + selected: boolean; + available: boolean; +}; + +export type AgentThreadRuntimeControlsState = { + tools: { + required: AgentThreadRuntimeControlChoice[]; + optional: AgentThreadRuntimeControlChoice[]; + selectedChoiceIds: string[]; + }; + mcp: { + connections: AgentThreadRuntimeControlChoice[]; + selectedChoiceIds: string[]; + }; + canEdit: boolean; + disabledReason: string | null; +}; + +export type AgentThreadRuntimeControlChoiceInput = { + agentId?: string | null; + toolChoiceIds?: string[]; + mcpChoiceIds?: string[]; +}; + +export type AgentRuntimeControlsEntrySourceInput = { + adapter?: string; + input?: Record; +}; + +export type AgentRuntimeControlsEntryDefaultsInput = { + provider?: string | null; + model?: string | null; +}; + +export type ValidatedEntryRuntimeControlChoices = { + selectedAgentMetadataPatch: Record | null; + runtimeControlChoices: AgentThreadRuntimeControlChoicesMetadata | null; +}; + +export type ResolvedRunAdmissionRuntimeChoices = { + metadataPresent: boolean; + selectedRuntimeToolChoiceIds?: string[]; + selectedRuntimeMcpChoiceIds?: string[]; + selectedRuntimeCapabilityIds?: AgentCapabilityCatalogId[]; + selectedRuntimeMcpConnectionRefs?: string[]; +}; + +type RuntimeControlsErrorCode = 'invalid_input' | 'unknown_choice' | 'policy_denied' | 'active_run' | 'not_found'; + +export class AgentThreadRuntimeControlsError extends Error { + constructor(public readonly code: RuntimeControlsErrorCode, message: string) { + super(message); + this.name = 'AgentThreadRuntimeControlsError'; + } +} + +const CHOICE_ID_PREFIX = 'rtc_v1_f48b74d9'; +const ACTIVE_RUN_DISABLED_REASON = 'Change after this response finishes.'; + +type RuntimeChoiceContext = { + selectedAgentId: string; + definition: AgentDefinitionContract; + sourceKind: AgentCapabilitySourceKind; + capabilityPolicy: Awaited>['capabilityPolicy']; + customAgentCreationPolicy: Awaited< + ReturnType + >['customAgentCreationPolicy']; + approvalPolicy: Awaited>['approvalPolicy']; + repoFullName?: string; + activeRun: boolean; + savedChoices: AgentThreadRuntimeControlChoicesMetadata | null; + mcpConnections: AgentMcpConnection[]; +}; + +type ThreadRuntimeChoiceContext = RuntimeChoiceContext & { + threadRecordId: number; +}; + +type ChoiceLookup = { + toolsById: Map; + mcpById: Map; +}; + +type RuntimeChoicePatch = { + toolChoiceIds?: string[]; + mcpChoiceIds?: string[]; +}; + +function opaqueChoiceId(kind: 'mcp' | 'tool', rawId: string): string { + const digest = createHash('sha256') + .update(`${CHOICE_ID_PREFIX}:${kind}:${rawId}`) + .digest('base64') + .replace(/\+/g, '-') + .replace(/\//g, '_') + .replace(/=+$/g, '') + .slice(0, 24); + return `${CHOICE_ID_PREFIX}_${kind}_${digest}`; +} + +function inferEntryDefaultAgentDefinitionId(source?: AgentRuntimeControlsEntrySourceInput): SystemAgentDefinitionId { + if (source?.adapter === 'blank_workspace') { + return typeof source.input?.buildUuid === 'string' && source.input.buildUuid.trim() + ? 'system.debug' + : 'system.freeform'; + } + + return 'system.develop'; +} + +function sourceKindForEntrySelection({ + defaultAgentDefinitionId, + selectedAgentId, + source, +}: { + defaultAgentDefinitionId: SystemAgentDefinitionId; + selectedAgentId: string; + source?: AgentRuntimeControlsEntrySourceInput; +}): AgentCapabilitySourceKind { + const isBlankChat = + source?.adapter === 'blank_workspace' && + !(typeof source.input?.buildUuid === 'string' && source.input.buildUuid.trim()); + + if (isBlankChat && selectedAgentId === 'system.develop') { + return 'workspace_session'; + } + + return sourceKindForSystemAgentDefinitionId(defaultAgentDefinitionId); +} + +function readOptionalStringArray(value: unknown, fieldName: string): string[] | undefined { + if (value === undefined) { + return undefined; + } + + if (!Array.isArray(value)) { + throw new AgentThreadRuntimeControlsError('invalid_input', `${fieldName} must be an array of choice ids.`); + } + + const normalized: string[] = []; + const seen = new Set(); + for (const item of value) { + if (typeof item !== 'string' || !item.trim()) { + throw new AgentThreadRuntimeControlsError('invalid_input', `${fieldName} must contain only choice ids.`); + } + const trimmed = item.trim(); + if (!seen.has(trimmed)) { + normalized.push(trimmed); + seen.add(trimmed); + } + } + + return normalized; +} + +function patchFromChoiceInput( + input: AgentThreadRuntimeControlChoiceInput | null | undefined +): RuntimeChoicePatch | null { + if (!input) { + return null; + } + + const toolChoiceIds = readOptionalStringArray(input.toolChoiceIds, 'toolChoiceIds'); + const mcpChoiceIds = readOptionalStringArray(input.mcpChoiceIds, 'mcpChoiceIds'); + if (toolChoiceIds === undefined && mcpChoiceIds === undefined) { + return null; + } + + return { + ...(toolChoiceIds !== undefined ? { toolChoiceIds } : {}), + ...(mcpChoiceIds !== undefined ? { mcpChoiceIds } : {}), + }; +} + +function isMcpCapability(capabilityId: AgentCapabilityCatalogId): boolean { + return getAgentCapabilityCatalogEntry(capabilityId).category === 'mcp'; +} + +function repoFullNameFromEntrySource(source?: AgentRuntimeControlsEntrySourceInput): string | undefined { + const input = source?.input || {}; + const directRepo = typeof input.repo === 'string' ? input.repo.trim() : ''; + if (directRepo) { + return directRepo; + } + + const repoUrl = typeof input.repoUrl === 'string' ? input.repoUrl.trim() : ''; + if (!repoUrl) { + return undefined; + } + + const normalized = repoUrl.replace(/^https?:\/\/github\.com\//, '').replace(/\.git$/, ''); + return normalized || undefined; +} + +function isConnectionAvailable(connection: AgentMcpConnection): boolean { + if (connection.validationError) { + return false; + } + + if (connection.connectionRequired && (!connection.configured || connection.stale)) { + return false; + } + + return (connection.discoveredTools?.length ?? 0) + (connection.sharedDiscoveredTools?.length ?? 0) > 0; +} + +function capabilityAllowed({ + capabilityId, + definition, + context, +}: { + capabilityId: AgentCapabilityCatalogId; + definition: AgentDefinitionContract; + context: RuntimeChoiceContext; +}) { + return AgentPolicyService.resolveCapabilityAccess({ + capabilityId, + capabilityPolicy: context.capabilityPolicy, + customAgentCreationPolicy: context.customAgentCreationPolicy, + approvalPolicy: context.approvalPolicy, + definitionOwnerKind: definition.owner.kind, + sourceKind: context.sourceKind, + }); +} + +function hasAllowedMcpCapability(context: RuntimeChoiceContext): boolean { + const capabilityRefs = [ + ...(context.definition.requiredCapabilityRefs || []), + ...(context.definition.optionalCapabilityRefs || []), + ...context.definition.capabilityRefs, + ].filter(isMcpCapability); + const uniqueCapabilityRefs = [...new Set(capabilityRefs)]; + + return uniqueCapabilityRefs.some( + (capabilityId) => + capabilityAllowed({ + capabilityId, + definition: context.definition, + context, + }).allowed + ); +} + +function assertDefinitionUsable(definition: AgentDefinitionContract, sourceKind: AgentCapabilitySourceKind): void { + if (definition.status !== 'active') { + throw new AgentThreadRuntimeControlsError('policy_denied', `${definition.name} is unavailable.`); + } + + if (!definition.resourcePolicy.sourceKinds.includes(sourceKind)) { + throw new AgentThreadRuntimeControlsError( + 'policy_denied', + `${definition.name} is unavailable for this conversation.` + ); + } +} + +function buildChoiceState(context: RuntimeChoiceContext): { + state: AgentThreadRuntimeControlsState; + lookup: ChoiceLookup; +} { + const savedChoices = context.savedChoices; + const savedToolChoiceIds = savedChoices ? new Set(savedChoices.toolChoiceIds) : null; + const savedMcpChoiceIds = savedChoices ? new Set(savedChoices.mcpChoiceIds) : null; + const requiredCapabilities = new Set(context.definition.requiredCapabilityRefs || []); + const optionalCapabilities = new Set( + (context.definition.optionalCapabilityRefs?.length + ? context.definition.optionalCapabilityRefs + : context.definition.capabilityRefs + ).filter((capabilityId) => !requiredCapabilities.has(capabilityId)) + ); + + const required: AgentThreadRuntimeControlChoice[] = []; + const optional: AgentThreadRuntimeControlChoice[] = []; + const lookup: ChoiceLookup = { + toolsById: new Map(), + mcpById: new Map(), + }; + + for (const capabilityId of [...requiredCapabilities, ...optionalCapabilities]) { + if (isMcpCapability(capabilityId)) { + continue; + } + + const entry = getAgentCapabilityCatalogEntry(capabilityId); + const access = capabilityAllowed({ + capabilityId, + definition: context.definition, + context, + }); + const id = opaqueChoiceId('tool', capabilityId); + const isRequired = requiredCapabilities.has(capabilityId); + const selected = isRequired || (savedToolChoiceIds ? savedToolChoiceIds.has(id) : true); + const choice = { + id, + label: entry.label, + description: entry.description || null, + required: isRequired, + selected, + available: access.allowed, + }; + lookup.toolsById.set(id, { ...choice, rawCapabilityId: capabilityId }); + if (isRequired) { + required.push(choice); + } else { + optional.push(choice); + } + } + + const mcpConnections: AgentThreadRuntimeControlChoice[] = hasAllowedMcpCapability(context) + ? context.mcpConnections.map((connection) => { + const id = opaqueChoiceId('mcp', `${connection.scope}:${connection.slug}`); + const available = isConnectionAvailable(connection); + const choice = { + id, + label: connection.name, + description: connection.description || null, + required: false, + selected: savedMcpChoiceIds ? savedMcpChoiceIds.has(id) : available, + available, + }; + lookup.mcpById.set(id, { ...choice, rawConnectionId: `${connection.scope}:${connection.slug}` }); + return choice; + }) + : []; + + const selectedToolChoiceIds = [ + ...required.filter((choice) => choice.available).map((choice) => choice.id), + ...optional.filter((choice) => choice.available && choice.selected).map((choice) => choice.id), + ]; + const selectedMcpChoiceIds = mcpConnections + .filter((choice) => choice.available && choice.selected) + .map((choice) => choice.id); + + return { + state: { + tools: { + required, + optional, + selectedChoiceIds: selectedToolChoiceIds, + }, + mcp: { + connections: mcpConnections, + selectedChoiceIds: selectedMcpChoiceIds, + }, + canEdit: !context.activeRun, + disabledReason: context.activeRun ? ACTIVE_RUN_DISABLED_REASON : null, + }, + lookup, + }; +} + +function validateChoiceIds({ + lookup, + toolChoiceIds, + mcpChoiceIds, +}: { + lookup: ChoiceLookup; + toolChoiceIds: string[]; + mcpChoiceIds: string[]; +}): AgentThreadRuntimeControlChoicesMetadata { + for (const choiceId of toolChoiceIds) { + const choice = lookup.toolsById.get(choiceId); + if (!choice) { + throw new AgentThreadRuntimeControlsError('unknown_choice', 'Unknown runtime control choice.'); + } + if (!choice.available) { + throw new AgentThreadRuntimeControlsError('policy_denied', 'Runtime control choice is unavailable.'); + } + } + + for (const choiceId of mcpChoiceIds) { + const choice = lookup.mcpById.get(choiceId); + if (!choice) { + throw new AgentThreadRuntimeControlsError('unknown_choice', 'Unknown runtime control choice.'); + } + if (!choice.available) { + throw new AgentThreadRuntimeControlsError('policy_denied', 'Runtime control choice is unavailable.'); + } + } + + return { + version: 1, + toolChoiceIds: toolChoiceIds.filter((choiceId) => { + const choice = lookup.toolsById.get(choiceId); + return choice && !choice.required; + }), + mcpChoiceIds, + }; +} + +function validateRequestedChoicePatch( + lookup: ChoiceLookup, + state: AgentThreadRuntimeControlsState, + patch: RuntimeChoicePatch +): AgentThreadRuntimeControlChoicesMetadata { + return validateChoiceIds({ + lookup, + toolChoiceIds: patch.toolChoiceIds ?? state.tools.selectedChoiceIds, + mcpChoiceIds: patch.mcpChoiceIds ?? state.mcp.selectedChoiceIds, + }); +} + +async function resolveSystemDefinition(agentId: SystemAgentDefinitionId): Promise { + await AgentDefinitionRegistry.ensureSystemAgentDefinitionsSeeded(); + return AgentDefinitionRegistry.getSystemAgentDefinition(agentId); +} + +async function resolveDefinition(agentId: string, userIdentity: RequestUserIdentity): Promise { + if (isSystemAgentDefinitionId(agentId)) { + return resolveSystemDefinition(agentId); + } + + if (agentId.startsWith('custom.')) { + return customAgentDefinitionService.getUserDefinition(agentId, userIdentity.userId); + } + + throw new AgentThreadRuntimeControlsError('policy_denied', 'Selected agent is unavailable.'); +} + +export default class AgentThreadRuntimeControlsService { + static async resolveRunAdmissionChoices({ + thread, + userIdentity, + definition, + sourceKind, + capabilityPolicy, + customAgentCreationPolicy, + approvalPolicy, + repoFullName, + }: { + thread: AgentThread; + userIdentity: RequestUserIdentity; + definition: AgentDefinitionContract; + sourceKind: AgentCapabilitySourceKind; + capabilityPolicy: RuntimeChoiceContext['capabilityPolicy']; + customAgentCreationPolicy: RuntimeChoiceContext['customAgentCreationPolicy']; + approvalPolicy: RuntimeChoiceContext['approvalPolicy']; + repoFullName?: string; + }): Promise { + const savedChoices = AgentThreadService.getRuntimeControlChoices(thread); + if (!savedChoices) { + return { + metadataPresent: false, + }; + } + + const mcpConnections = await new McpConfigService().listEnabledConnectionsForUser(repoFullName, userIdentity); + const { state, lookup } = buildChoiceState({ + selectedAgentId: definition.id, + definition, + sourceKind, + capabilityPolicy, + customAgentCreationPolicy, + approvalPolicy, + repoFullName, + activeRun: false, + savedChoices, + mcpConnections, + }); + + const selectedRuntimeCapabilityIds = state.tools.selectedChoiceIds + .map((choiceId) => lookup.toolsById.get(choiceId)?.rawCapabilityId) + .filter((capabilityId): capabilityId is AgentCapabilityCatalogId => Boolean(capabilityId)); + const selectedRuntimeMcpConnectionRefs = state.mcp.selectedChoiceIds + .map((choiceId) => lookup.mcpById.get(choiceId)?.rawConnectionId) + .filter((connectionRef): connectionRef is string => Boolean(connectionRef)); + + return { + metadataPresent: true, + selectedRuntimeToolChoiceIds: state.tools.selectedChoiceIds, + selectedRuntimeMcpChoiceIds: state.mcp.selectedChoiceIds, + selectedRuntimeCapabilityIds, + selectedRuntimeMcpConnectionRefs, + }; + } + + static async getState({ + threadId, + userIdentity, + }: { + threadId: string; + userIdentity: RequestUserIdentity; + }): Promise { + const context = await this.buildThreadContext({ threadId, userIdentity }); + return buildChoiceState(context).state; + } + + static async patchChoices({ + threadId, + userIdentity, + toolChoiceIds, + mcpChoiceIds, + }: { + threadId: string; + userIdentity: RequestUserIdentity; + toolChoiceIds?: string[]; + mcpChoiceIds?: string[]; + }): Promise { + const context = await this.buildThreadContext({ threadId, userIdentity }); + if (context.activeRun) { + throw new AgentThreadRuntimeControlsError('active_run', ACTIVE_RUN_DISABLED_REASON); + } + + const patch = patchFromChoiceInput({ toolChoiceIds, mcpChoiceIds }); + if (!patch) { + throw new AgentThreadRuntimeControlsError('invalid_input', 'runtimeControlChoices are required.'); + } + const { state, lookup } = buildChoiceState(context); + const validatedMetadata = validateRequestedChoicePatch(lookup, state, patch); + + await AgentThreadService.patchRuntimeControlChoices(context.threadRecordId, validatedMetadata); + + return buildChoiceState({ + ...context, + savedChoices: validatedMetadata, + }).state; + } + + static async getEntryPreview({ + userIdentity, + agentId, + source, + defaults, + runtimeControlChoices, + }: { + userIdentity: RequestUserIdentity; + agentId?: string | null; + source?: AgentRuntimeControlsEntrySourceInput; + defaults?: AgentRuntimeControlsEntryDefaultsInput; + runtimeControlChoices?: AgentThreadRuntimeControlChoiceInput; + }): Promise { + void defaults; + const context = await this.buildEntryContext({ userIdentity, agentId, source }); + const patch = patchFromChoiceInput(runtimeControlChoices); + const { state, lookup } = buildChoiceState(context); + const savedChoices = patch ? validateRequestedChoicePatch(lookup, state, patch) : null; + return buildChoiceState({ + ...context, + savedChoices, + }).state; + } + + static async validateEntryChoices({ + userIdentity, + agentId, + source, + defaults, + runtimeControlChoices, + }: { + userIdentity: RequestUserIdentity; + agentId?: string | null; + source?: AgentRuntimeControlsEntrySourceInput; + defaults?: AgentRuntimeControlsEntryDefaultsInput; + runtimeControlChoices: AgentThreadRuntimeControlChoiceInput; + }): Promise { + void defaults; + const selectedAgentId = agentId || runtimeControlChoices.agentId || null; + const context = await this.buildEntryContext({ userIdentity, agentId: selectedAgentId, source }); + const patch = patchFromChoiceInput(runtimeControlChoices); + const { state, lookup } = buildChoiceState(context); + const validatedMetadata = patch ? validateRequestedChoicePatch(lookup, state, patch) : null; + + return { + selectedAgentMetadataPatch: selectedAgentId + ? AgentThreadService.buildSelectedAgentDefinitionMetadataPatch(context.selectedAgentId) + : null, + runtimeControlChoices: validatedMetadata, + }; + } + + private static async buildThreadContext({ + threadId, + userIdentity, + }: { + threadId: string; + userIdentity: RequestUserIdentity; + }): Promise { + let owned; + try { + owned = await AgentThreadService.getOwnedThreadWithSession(threadId, userIdentity.userId); + } catch (error) { + if ( + error instanceof Error && + (error.message === 'Agent thread not found' || error.message === 'Agent session not found') + ) { + throw new AgentThreadRuntimeControlsError('not_found', error.message); + } + throw error; + } + + const { thread, session } = owned; + const source = await AgentSourceService.getSessionSource(session.id); + if (!source || source.status !== 'ready') { + throw new AgentThreadRuntimeControlsError('policy_denied', 'Session source is not ready yet.'); + } + + const defaultAgentDefinitionId = AgentDefinitionRegistry.inferDefaultSystemAgentDefinitionId(session, source); + const selectedAgentId = AgentThreadService.getSelectedAgentDefinitionId(thread) || defaultAgentDefinitionId; + const sourceKind = sourceKindForSystemAgentDefinitionId(defaultAgentDefinitionId); + const definition = await resolveDefinition(selectedAgentId, userIdentity); + assertDefinitionUsable(definition, sourceKind); + const { repoFullName, approvalPolicy, capabilityPolicy, customAgentCreationPolicy } = + await AgentCapabilityService.resolveSessionContext(session.uuid, userIdentity); + const [activeRun, mcpConnections] = await Promise.all([ + AgentRunService.hasActiveRun(thread.id), + new McpConfigService().listEnabledConnectionsForUser(repoFullName, userIdentity), + ]); + + return { + threadRecordId: thread.id, + selectedAgentId, + definition, + sourceKind, + capabilityPolicy, + customAgentCreationPolicy, + approvalPolicy, + repoFullName, + activeRun, + savedChoices: AgentThreadService.getRuntimeControlChoices(thread), + mcpConnections, + }; + } + + private static async buildEntryContext({ + userIdentity, + agentId, + source, + }: { + userIdentity: RequestUserIdentity; + agentId?: string | null; + source?: AgentRuntimeControlsEntrySourceInput; + }): Promise { + const defaultAgentDefinitionId = inferEntryDefaultAgentDefinitionId(source); + const selectedAgentId = agentId?.trim() || defaultAgentDefinitionId; + const sourceKind = sourceKindForEntrySelection({ defaultAgentDefinitionId, selectedAgentId, source }); + const definition = await resolveDefinition(selectedAgentId, userIdentity); + assertDefinitionUsable(definition, sourceKind); + const repoFullName = repoFullNameFromEntrySource(source); + const [approvalPolicy, effectiveConfig, mcpConnections] = await Promise.all([ + AgentPolicyService.getEffectivePolicy(repoFullName), + AgentRuntimeConfigService.getInstance().getEffectiveConfig(repoFullName), + new McpConfigService().listEnabledConnectionsForUser(repoFullName, userIdentity), + ]); + + return { + selectedAgentId, + definition, + sourceKind, + capabilityPolicy: effectiveConfig.capabilityPolicy, + customAgentCreationPolicy: effectiveConfig.customAgentCreationPolicy, + approvalPolicy, + repoFullName, + activeRun: false, + savedChoices: null, + mcpConnections, + }; + } +} diff --git a/src/server/services/agent/ThreadService.ts b/src/server/services/agent/ThreadService.ts index f6393abc..3d39b244 100644 --- a/src/server/services/agent/ThreadService.ts +++ b/src/server/services/agent/ThreadService.ts @@ -16,13 +16,91 @@ import AgentSession from 'server/models/AgentSession'; import AgentThread from 'server/models/AgentThread'; +import type { Transaction } from 'objection'; +import { canSessionAcceptMessages, getSessionMessageBlockReason } from './sessionReadiness'; + +export const AGENT_THREAD_SELECTED_AGENT_DEFINITION_METADATA_KEY = 'selectedAgentDefinitionId'; +export const AGENT_THREAD_RUNTIME_CONTROL_CHOICES_METADATA_KEY = 'runtimeControlChoices'; + +export type AgentThreadRuntimeControlChoicesMetadata = { + version: 1; + toolChoiceIds: string[]; + mcpChoiceIds: string[]; +}; function normalizeTitle(title?: string | null): string | null { const trimmed = title?.trim(); return trimmed ? trimmed : null; } +function readRecord(value: unknown): Record { + return value && typeof value === 'object' && !Array.isArray(value) ? (value as Record) : {}; +} + +export function getSelectedAgentDefinitionId(thread: AgentThread): string | null { + const metadata = readRecord(thread.metadata); + const selectedDefinitionId = metadata[AGENT_THREAD_SELECTED_AGENT_DEFINITION_METADATA_KEY]; + if (typeof selectedDefinitionId === 'string' && selectedDefinitionId.trim()) { + return selectedDefinitionId.trim(); + } + + return null; +} + +export function buildSelectedAgentDefinitionMetadataPatch(agentId: string): Record { + return { + [AGENT_THREAD_SELECTED_AGENT_DEFINITION_METADATA_KEY]: agentId, + }; +} + +function normalizeChoiceIds(value: unknown): string[] | null { + if (!Array.isArray(value)) { + return null; + } + + return Array.from( + new Set(value.filter((entry): entry is string => typeof entry === 'string').map((entry) => entry.trim())) + ).filter(Boolean); +} + +export function getRuntimeControlChoices(thread: AgentThread): AgentThreadRuntimeControlChoicesMetadata | null { + const metadata = readRecord(thread.metadata)[AGENT_THREAD_RUNTIME_CONTROL_CHOICES_METADATA_KEY]; + if (!metadata || typeof metadata !== 'object' || Array.isArray(metadata)) { + return null; + } + + const record = metadata as Record; + const toolChoiceIds = normalizeChoiceIds(record.toolChoiceIds); + const mcpChoiceIds = normalizeChoiceIds(record.mcpChoiceIds); + if (record.version !== 1 || !toolChoiceIds || !mcpChoiceIds) { + return null; + } + + return { + version: 1, + toolChoiceIds, + mcpChoiceIds, + }; +} + +export function buildRuntimeControlChoicesMetadataPatch( + choices: AgentThreadRuntimeControlChoicesMetadata +): Record { + return { + [AGENT_THREAD_RUNTIME_CONTROL_CHOICES_METADATA_KEY]: { + version: 1, + toolChoiceIds: [...choices.toolChoiceIds], + mcpChoiceIds: [...choices.mcpChoiceIds], + }, + }; +} + export default class AgentThreadService { + static getSelectedAgentDefinitionId = getSelectedAgentDefinitionId; + static buildSelectedAgentDefinitionMetadataPatch = buildSelectedAgentDefinitionMetadataPatch; + static getRuntimeControlChoices = getRuntimeControlChoices; + static buildRuntimeControlChoicesMetadataPatch = buildRuntimeControlChoicesMetadataPatch; + static async getOwnedSession(sessionUuid: string, userId: string): Promise { const session = await AgentSession.query().findOne({ uuid: sessionUuid, userId }); if (!session) { @@ -120,6 +198,12 @@ export default class AgentThreadService { static async createThread(sessionUuid: string, userId: string, title?: string | null): Promise { const session = await this.getOwnedSession(sessionUuid, userId); + if (session.status === 'ended' || session.status === 'error') { + throw new Error('Cannot create a thread for an inactive session'); + } + if (!canSessionAcceptMessages(session)) { + throw new Error(getSessionMessageBlockReason(session)); + } return AgentThread.query().insertAndFetch({ sessionId: session.id, @@ -131,6 +215,24 @@ export default class AgentThreadService { } as Partial); } + static async patchRuntimeControlChoices( + threadId: number, + choices: AgentThreadRuntimeControlChoicesMetadata, + trx?: Transaction + ): Promise { + const thread = await AgentThread.query(trx).findById(threadId); + if (!thread) { + throw new Error('Agent thread not found'); + } + + return AgentThread.query(trx).patchAndFetchById(threadId, { + metadata: { + ...(thread.metadata || {}), + ...buildRuntimeControlChoicesMetadataPatch(choices), + }, + } as Partial); + } + static serializeThread(thread: AgentThread, sessionUuid?: string) { return { id: thread.uuid, diff --git a/src/server/services/agent/__tests__/AdminService.test.ts b/src/server/services/agent/__tests__/AdminService.test.ts index 286e63f3..3e787eb1 100644 --- a/src/server/services/agent/__tests__/AdminService.test.ts +++ b/src/server/services/agent/__tests__/AdminService.test.ts @@ -22,6 +22,8 @@ const mockMessageQuery = jest.fn(); const mockRunQuery = jest.fn(); const mockRunEventQuery = jest.fn(); const mockToolExecutionQuery = jest.fn(); +const mockMcpServerConfigQuery = jest.fn(); +const mockUserMcpConnectionQuery = jest.fn(); const mockSerializeRun = jest.fn(); const mockSerializeRunEvent = jest.fn(); const mockSerializeThread = jest.fn(); @@ -83,6 +85,20 @@ jest.mock('server/models/AgentToolExecution', () => ({ }, })); +jest.mock('server/models/McpServerConfig', () => ({ + __esModule: true, + default: { + query: (...args: unknown[]) => mockMcpServerConfigQuery(...args), + }, +})); + +jest.mock('server/models/UserMcpConnection', () => ({ + __esModule: true, + default: { + query: (...args: unknown[]) => mockUserMcpConnectionQuery(...args), + }, +})); + jest.mock('../RunService', () => ({ __esModule: true, default: { @@ -222,6 +238,132 @@ describe('AgentAdminService.listSessions', () => { }); }); +describe('AgentAdminService.listMcpServerCoverage', () => { + beforeEach(() => { + jest.clearAllMocks(); + }); + + it('redacts transport and shared MCP secrets in admin coverage rows', async () => { + const configQueryBuilder = { + where: jest.fn().mockReturnThis(), + whereNull: jest.fn().mockReturnThis(), + orderBy: jest.fn().mockResolvedValue([ + { + slug: 'sample-connector', + name: 'Sample Connector', + description: 'A sample MCP connector.', + scope: 'global', + preset: null, + transport: { + type: 'http', + url: 'https://mcp.example.test?api_key=transport-query-secret', + headers: { + Authorization: 'Bearer shared-token', + 'X-Api-Key': 'shared-api-key', + }, + }, + sharedConfig: { + headers: { + 'X-Shared-Token': 'header-secret', + }, + query: { + token: 'query-secret', + }, + env: { + SAMPLE_TOKEN: 'env-secret', + }, + defaultArgs: { + project: 'arg-secret', + }, + }, + authConfig: { mode: 'none' }, + enabled: true, + timeout: 5000, + sharedDiscoveredTools: [], + createdAt: '2026-04-20T00:00:00.000Z', + updatedAt: '2026-04-21T00:00:00.000Z', + }, + { + slug: 'sample-cli', + name: 'Sample CLI', + description: null, + scope: 'global', + preset: null, + transport: { + type: 'stdio', + command: 'sample-mcp', + args: ['--mode', 'stdio'], + env: { + SAMPLE_API_TOKEN: 'stdio-secret', + }, + }, + sharedConfig: {}, + authConfig: { mode: 'none' }, + enabled: true, + timeout: 5000, + sharedDiscoveredTools: [], + createdAt: '2026-04-20T00:00:00.000Z', + updatedAt: '2026-04-21T00:00:00.000Z', + }, + ]), + }; + mockMcpServerConfigQuery.mockReturnValue(configQueryBuilder); + + const connectionQueryBuilder = { + where: jest.fn().mockReturnThis(), + orderBy: jest.fn().mockResolvedValue([{ slug: 'sample-connector', validatedAt: '2026-04-22T00:00:00.000Z' }]), + }; + mockUserMcpConnectionQuery.mockReturnValue(connectionQueryBuilder); + + const result = await AgentAdminService.listMcpServerCoverage(); + + expect(configQueryBuilder.where).toHaveBeenCalledWith({ scope: 'global' }); + expect(result).toEqual([ + expect.objectContaining({ + slug: 'sample-connector', + transport: { + type: 'http', + url: 'https://mcp.example.test?api_key=******', + headers: { + Authorization: '******', + 'X-Api-Key': '******', + }, + }, + sharedConfig: { + headers: { + 'X-Shared-Token': '******', + }, + query: { + token: '******', + }, + env: { + SAMPLE_TOKEN: '******', + }, + defaultArgs: { + project: '******', + }, + }, + userConnectionCount: 1, + latestUserValidatedAt: '2026-04-22T00:00:00.000Z', + }), + expect.objectContaining({ + slug: 'sample-cli', + transport: { + type: 'stdio', + command: 'sample-mcp', + args: ['--mode', 'stdio'], + env: { + SAMPLE_API_TOKEN: '******', + }, + }, + sharedConfig: {}, + userConnectionCount: 0, + latestUserValidatedAt: null, + }), + ]); + }); +}); + describe('AgentAdminService.getThreadConversation', () => { beforeEach(() => { jest.clearAllMocks(); diff --git a/src/server/services/agent/__tests__/AgentDefinitionRegistry.test.ts b/src/server/services/agent/__tests__/AgentDefinitionRegistry.test.ts new file mode 100644 index 00000000..4b84c1dc --- /dev/null +++ b/src/server/services/agent/__tests__/AgentDefinitionRegistry.test.ts @@ -0,0 +1,157 @@ +/** + * 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. + */ + +const mockUpsert = jest.fn(); +const mockFindOne = jest.fn(); +const mockOrderBy = jest.fn(); +const mockWhere = jest.fn(); +const mockWhereIn = jest.fn(); + +jest.mock('server/models/AgentDefinition', () => ({ + __esModule: true, + default: { + upsert: (...args: unknown[]) => mockUpsert(...args), + query: jest.fn(() => ({ + findOne: (...args: unknown[]) => mockFindOne(...args), + whereIn: (...args: unknown[]) => { + mockWhereIn(...args); + return { + where: (...whereArgs: unknown[]) => { + mockWhere(...whereArgs); + return { + orderBy: (...orderArgs: unknown[]) => mockOrderBy(...orderArgs), + }; + }, + }; + }, + })), + }, +})); + +import { + assertAgentDefinitionMutable, + ensureSystemAgentDefinitionsSeeded, + getSystemAgentDefinition, + inferDefaultSystemAgentDefinitionId, + listSystemAgentDefinitions, + serializeAgentDefinitionSummary, +} from '../AgentDefinitionRegistry'; +import { SYSTEM_AGENT_DEFINITIONS } from '../systemAgentDefinitions'; +import { AgentSessionKind, AgentWorkspaceStatus } from 'shared/constants'; + +function buildRow(agentId: keyof typeof SYSTEM_AGENT_DEFINITIONS) { + const definition = SYSTEM_AGENT_DEFINITIONS[agentId]; + return { + definitionId: definition.id, + version: definition.version, + ownerKind: definition.owner.kind, + ownerUserId: null, + ownerOrganizationId: null, + name: definition.name, + description: definition.description || null, + instructionRefs: definition.instructionRefs, + instructionAddendum: definition.instructionAddendum || null, + capabilityRefs: definition.capabilityRefs, + requiredCapabilityRefs: definition.requiredCapabilityRefs, + optionalCapabilityRefs: definition.optionalCapabilityRefs, + resourcePolicy: definition.resourcePolicy, + modelPreference: definition.modelPreference || null, + status: definition.status, + codeOwned: definition.codeOwned, + readOnly: definition.readOnly, + }; +} + +describe('AgentDefinitionRegistry', () => { + beforeEach(() => { + jest.clearAllMocks(); + mockUpsert.mockImplementation(async (row) => row); + mockFindOne.mockResolvedValue(buildRow('system.freeform')); + mockOrderBy.mockResolvedValue([buildRow('system.debug'), buildRow('system.develop'), buildRow('system.freeform')]); + }); + + it('seeds exactly the first-party system agent definition definitions as code-owned read-only rows', async () => { + const rows = await ensureSystemAgentDefinitionsSeeded(); + + expect(mockUpsert).toHaveBeenCalledTimes(3); + expect(mockUpsert.mock.calls.map(([row]) => (row as { definitionId: string }).definitionId).sort()).toEqual([ + 'system.debug', + 'system.develop', + 'system.freeform', + ]); + expect(rows).toEqual( + expect.arrayContaining([ + expect.objectContaining({ + id: 'system.debug', + owner: { kind: 'system', userId: null, organizationId: null }, + codeOwned: true, + readOnly: true, + status: 'active', + requiredCapabilityRefs: expect.arrayContaining(['diagnostics_codefresh', 'diagnostics_kubernetes']), + optionalCapabilityRefs: [], + }), + ]) + ); + }); + + it('loads persisted system agent definitions by public id and lists summaries', async () => { + const definition = await getSystemAgentDefinition('system.freeform'); + const summary = serializeAgentDefinitionSummary(definition); + + expect(mockFindOne).toHaveBeenCalledWith({ + definitionId: 'system.freeform', + ownerKind: 'system', + }); + expect(summary).toEqual( + expect.objectContaining({ + id: 'system.freeform', + ownerKind: 'system', + codeOwned: true, + readOnly: true, + }) + ); + + await expect(listSystemAgentDefinitions()).resolves.toHaveLength(3); + expect(mockWhereIn).toHaveBeenCalledWith('definitionId', ['system.debug', 'system.develop', 'system.freeform']); + }); + + it('rejects mutations for code-owned system definitions', () => { + expect(() => assertAgentDefinitionMutable(SYSTEM_AGENT_DEFINITIONS['system.debug'])).toThrow( + 'code-owned and read-only' + ); + }); + + it('infers default system agent definition ids from launch source', () => { + expect( + inferDefaultSystemAgentDefinitionId( + { sessionKind: AgentSessionKind.CHAT } as any, + { input: { buildUuid: 'build-1' } } as any + ) + ).toBe('system.debug'); + expect( + inferDefaultSystemAgentDefinitionId({ sessionKind: AgentSessionKind.CHAT } as any, { input: {} } as any) + ).toBe('system.freeform'); + expect( + inferDefaultSystemAgentDefinitionId( + { sessionKind: AgentSessionKind.CHAT, workspaceStatus: AgentWorkspaceStatus.READY } as any, + { input: {} } as any + ) + ).toBe('system.develop'); + expect( + inferDefaultSystemAgentDefinitionId({ sessionKind: AgentSessionKind.SANDBOX } as any, { input: {} } as any) + ).toBe('system.develop'); + }); +}); diff --git a/src/server/services/agent/__tests__/AgentSelectionService.test.ts b/src/server/services/agent/__tests__/AgentSelectionService.test.ts new file mode 100644 index 00000000..6459c3c3 --- /dev/null +++ b/src/server/services/agent/__tests__/AgentSelectionService.test.ts @@ -0,0 +1,356 @@ +/** + * 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. + */ + +const mockThreadTransaction = jest.fn(); +const mockThreadQuery = jest.fn(); +const mockRunQuery = jest.fn(); +const mockResolveSessionContext = jest.fn(); +const mockEnsureSeeded = jest.fn(); +const mockListSystemDefinitions = jest.fn(); +const mockInferDefaultAgentDefinitionId = jest.fn(); +const mockCreateAgentSwitchEvent = jest.fn(); +const mockGetSessionSource = jest.fn(); +const mockGetOwnedThreadWithSession = jest.fn(); +const mockListUserDefinitions = jest.fn(); + +jest.mock('server/models/AgentThread', () => ({ + __esModule: true, + default: { + query: (...args: unknown[]) => mockThreadQuery(...args), + transaction: (...args: unknown[]) => mockThreadTransaction(...args), + }, +})); + +jest.mock('server/models/AgentRun', () => ({ + __esModule: true, + default: { + query: (...args: unknown[]) => mockRunQuery(...args), + }, +})); + +jest.mock('../CapabilityService', () => ({ + __esModule: true, + default: { + resolveSessionContext: (...args: unknown[]) => mockResolveSessionContext(...args), + }, +})); + +jest.mock('../AgentDefinitionRegistry', () => { + const actual = jest.requireActual('../AgentDefinitionRegistry'); + return { + __esModule: true, + ...actual, + ensureSystemAgentDefinitionsSeeded: (...args: unknown[]) => mockEnsureSeeded(...args), + listSystemAgentDefinitions: (...args: unknown[]) => mockListSystemDefinitions(...args), + inferDefaultSystemAgentDefinitionId: (...args: unknown[]) => mockInferDefaultAgentDefinitionId(...args), + }; +}); + +jest.mock('../CustomAgentDefinitionService', () => ({ + __esModule: true, + customAgentDefinitionService: { + listUserDefinitions: (...args: unknown[]) => mockListUserDefinitions(...args), + }, +})); + +jest.mock('../MessageStore', () => ({ + __esModule: true, + default: { + createAgentSwitchEvent: (...args: unknown[]) => mockCreateAgentSwitchEvent(...args), + }, +})); + +jest.mock('../SourceService', () => ({ + __esModule: true, + default: { + getSessionSource: (...args: unknown[]) => mockGetSessionSource(...args), + }, +})); + +jest.mock('../ThreadService', () => ({ + __esModule: true, + default: { + getOwnedThreadWithSession: (...args: unknown[]) => mockGetOwnedThreadWithSession(...args), + getSelectedAgentDefinitionId: (thread: { metadata?: Record }) => { + const selectedDefinitionId = thread.metadata?.selectedAgentDefinitionId; + if (typeof selectedDefinitionId === 'string' && selectedDefinitionId.trim()) { + return selectedDefinitionId; + } + return null; + }, + buildSelectedAgentDefinitionMetadataPatch: (agentId: string) => ({ + selectedAgentDefinitionId: agentId, + }), + }, +})); + +import AgentSelectionService, { AgentThreadAgentSwitchError } from '../AgentSelectionService'; +import { SYSTEM_AGENT_DEFINITIONS } from '../systemAgentDefinitions'; + +const userIdentity = { + userId: 'sample-user', + githubUsername: 'sample-user', + preferredUsername: 'sample-user', + email: 'sample-user@example.com', + firstName: 'Sample', + lastName: 'User', + displayName: '', + gitUserName: 'Sample User', + gitUserEmail: 'sample-user@example.com', + roles: [], +}; + +const thread = { + id: 7, + uuid: 'thread-1', + metadata: {}, +}; + +const session = { + id: 17, + uuid: 'session-1', +}; + +const source = { + id: 3, + status: 'ready', + input: {}, +}; + +const customDefinition = { + id: 'custom.sample-agent', + version: 2, + owner: { kind: 'user' as const, userId: 'sample-user' }, + name: 'Sample custom agent', + description: 'Uses allowed capabilities', + instructionRefs: [], + instructionAddendum: 'Focus on concise answers.', + capabilityRefs: ['general_chat'], + requiredCapabilityRefs: [], + optionalCapabilityRefs: ['general_chat'], + resourcePolicy: { + sourceKinds: ['freeform_chat'], + workspaceRequired: false, + sandboxRequired: false, + }, + modelPreference: null, + status: 'active' as const, + codeOwned: false, + readOnly: false, +}; + +function mockNoActiveRun() { + const query = { + where: jest.fn(() => query), + whereNotIn: jest.fn(() => query), + first: jest.fn().mockResolvedValue(null), + }; + mockRunQuery.mockReturnValue(query); +} + +describe('AgentSelectionService', () => { + beforeEach(() => { + jest.clearAllMocks(); + mockThreadTransaction.mockImplementation(async (callback) => callback({ trx: true })); + mockThreadQuery.mockReturnValue({ + patchAndFetchById: jest + .fn() + .mockResolvedValue({ ...thread, metadata: { selectedAgentDefinitionId: 'custom.sample-agent' } }), + }); + mockNoActiveRun(); + mockResolveSessionContext.mockResolvedValue({ + approvalPolicy: { defaultMode: 'require_approval', rules: {} }, + capabilityPolicy: undefined, + }); + mockEnsureSeeded.mockResolvedValue(Object.values(SYSTEM_AGENT_DEFINITIONS)); + mockListSystemDefinitions.mockResolvedValue(Object.values(SYSTEM_AGENT_DEFINITIONS)); + mockInferDefaultAgentDefinitionId.mockReturnValue('system.freeform'); + mockGetOwnedThreadWithSession.mockResolvedValue({ thread, session }); + mockGetSessionSource.mockResolvedValue(source); + mockListUserDefinitions.mockResolvedValue([customDefinition]); + mockCreateAgentSwitchEvent.mockResolvedValue({ uuid: 'message-1' }); + }); + + it('returns built_in and my_agents groups with selected, default, and current ids', async () => { + const state = await AgentSelectionService.getThreadAgentState({ threadId: 'thread-1', userIdentity }); + + expect(mockGetOwnedThreadWithSession).toHaveBeenCalledWith('thread-1', 'sample-user'); + expect(mockListUserDefinitions).toHaveBeenCalledWith({ userId: 'sample-user' }); + expect(state).toEqual( + expect.objectContaining({ + selectedId: null, + defaultId: 'system.freeform', + currentId: 'system.freeform', + }) + ); + expect(state.groups.map((group) => group.id)).toEqual(['built_in', 'my_agents']); + expect(state.groups[0].agents.map((agent) => agent.id)).toEqual([ + 'system.debug', + 'system.develop', + 'system.freeform', + ]); + expect(state.groups[1].agents).toEqual([ + expect.objectContaining({ + id: 'custom.sample-agent', + ownerKind: 'user', + group: 'my_agents', + label: 'Sample custom agent', + }), + ]); + }); + + it('marks user custom agents unavailable when a required capability becomes creator-reserved', async () => { + mockResolveSessionContext.mockResolvedValueOnce({ + approvalPolicy: { defaultMode: 'require_approval', rules: {} }, + capabilityPolicy: undefined, + customAgentCreationPolicy: { + capabilityAvailability: { + read_context: 'reserved', + }, + }, + }); + mockListUserDefinitions.mockResolvedValueOnce([ + { + ...customDefinition, + capabilityRefs: ['read_context'], + requiredCapabilityRefs: ['read_context'], + optionalCapabilityRefs: [], + }, + ]); + + const state = await AgentSelectionService.getThreadAgentState({ threadId: 'thread-1', userIdentity }); + + expect(state.groups[1].agents).toEqual([ + expect.objectContaining({ + id: 'custom.sample-agent', + available: false, + unavailableReason: 'disabled_by_policy', + unavailableMessage: 'Sample custom agent is unavailable because a required capability is disabled.', + }), + ]); + }); + + it('switches to an owned custom agent and clears stale agent metadata', async () => { + const result = await AgentSelectionService.switchThreadAgent({ + threadId: 'thread-1', + userIdentity, + agentId: 'custom.sample-agent', + }); + + expect(result.switched).toBe(true); + expect(mockThreadQuery().patchAndFetchById).toHaveBeenCalledWith( + 7, + expect.objectContaining({ + metadata: { selectedAgentDefinitionId: 'custom.sample-agent' }, + }) + ); + expect(mockCreateAgentSwitchEvent).toHaveBeenCalledWith( + expect.objectContaining({ + beforeAgent: { id: 'system.freeform', label: 'Free-form' }, + afterAgent: { id: 'custom.sample-agent', label: 'Sample custom agent' }, + }) + ); + }); + + it('rejects another user custom id and writes no preference', async () => { + await expect( + AgentSelectionService.switchThreadAgent({ + threadId: 'thread-1', + userIdentity, + agentId: 'custom.another-user-agent', + }) + ).rejects.toMatchObject({ + reason: 'unknown_agent', + }); + + expect(mockThreadQuery().patchAndFetchById).not.toHaveBeenCalled(); + expect(mockCreateAgentSwitchEvent).not.toHaveBeenCalled(); + }); + + it('rejects source-incompatible agents and writes no preference', async () => { + await expect( + AgentSelectionService.switchThreadAgent({ + threadId: 'thread-1', + userIdentity, + agentId: 'system.develop', + }) + ).rejects.toMatchObject({ + reason: 'requires_workspace', + }); + + expect(mockThreadQuery().patchAndFetchById).not.toHaveBeenCalled(); + expect(mockCreateAgentSwitchEvent).not.toHaveBeenCalled(); + }); + + it('rejects active run switches and leaves the old selection visible', async () => { + const query = { + where: jest.fn(() => query), + whereNotIn: jest.fn(() => query), + first: jest.fn().mockResolvedValue({ id: 99, status: 'running' }), + }; + mockRunQuery.mockReturnValue(query); + + await expect( + AgentSelectionService.switchThreadAgent({ + threadId: 'thread-1', + userIdentity, + agentId: 'custom.sample-agent', + }) + ).rejects.toMatchObject({ + reason: 'active_run', + }); + + expect(mockThreadQuery().patchAndFetchById).not.toHaveBeenCalled(); + }); + + it('excludes archived custom definitions from my_agents state', async () => { + mockListUserDefinitions.mockResolvedValueOnce([ + customDefinition, + { ...customDefinition, id: 'custom.archived-agent', name: 'Archived', status: 'archived' }, + ]); + + const state = await AgentSelectionService.getThreadAgentState({ threadId: 'thread-1', userIdentity }); + + expect(JSON.stringify(state.groups)).not.toContain('custom.archived-agent'); + expect(state.groups[1].agents).toHaveLength(1); + }); + + it('does not write an event for no-op switches', async () => { + mockGetOwnedThreadWithSession.mockResolvedValueOnce({ + thread: { ...thread, metadata: { selectedAgentDefinitionId: 'custom.sample-agent' } }, + session, + }); + + const result = await AgentSelectionService.switchThreadAgent({ + threadId: 'thread-1', + userIdentity, + agentId: 'custom.sample-agent', + }); + + expect(result.switched).toBe(false); + expect(mockCreateAgentSwitchEvent).not.toHaveBeenCalled(); + }); + + it('raises typed unknown-agent errors', async () => { + await expect( + AgentSelectionService.switchThreadAgent({ + threadId: 'thread-1', + userIdentity, + agentId: 'custom.missing', + }) + ).rejects.toBeInstanceOf(AgentThreadAgentSwitchError); + }); +}); diff --git a/src/server/services/agent/__tests__/AgentUsageService.test.ts b/src/server/services/agent/__tests__/AgentUsageService.test.ts new file mode 100644 index 00000000..2368da96 --- /dev/null +++ b/src/server/services/agent/__tests__/AgentUsageService.test.ts @@ -0,0 +1,343 @@ +/** + * 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(), + }, +})); + +jest.mock('server/services/agent/ThreadService', () => ({ + __esModule: true, + default: { + getOwnedThreadWithSession: jest.fn(), + getOwnedSession: jest.fn(), + }, +})); + +import AgentRun from 'server/models/AgentRun'; +import AgentThreadService from '../ThreadService'; +import AgentUsageService, { type AgentUsageRunRecord } from '../AgentUsageService'; + +const mockRunQuery = AgentRun.query as jest.Mock; +const mockGetOwnedThreadWithSession = AgentThreadService.getOwnedThreadWithSession as jest.Mock; +const mockGetOwnedSession = AgentThreadService.getOwnedSession as jest.Mock; + +function buildRun(overrides: Partial = {}): AgentUsageRunRecord { + return { + status: 'completed', + resolvedProvider: 'openai', + resolvedModel: 'gpt-5.4', + provider: 'openai', + model: 'gpt-5.4', + usageSummary: {}, + ...overrides, + }; +} + +function buildRunQuery(rows: AgentUsageRunRecord[]) { + const query = { + where: jest.fn(), + orderBy: jest.fn(), + }; + query.where.mockReturnValue(query); + query.orderBy.mockReturnValueOnce(query).mockResolvedValueOnce(rows); + return query; +} + +describe('AgentUsageService', () => { + beforeEach(() => { + jest.clearAllMocks(); + }); + + it('uses exact provider totals and exact input plus output totals only', () => { + const aggregate = AgentUsageService.aggregateRuns([ + buildRun({ + usageSummary: { + totalTokens: 120, + inputTokens: 80, + outputTokens: 30, + reasoningTokens: 10, + cachedInputTokens: 20, + cacheCreationInputTokens: 5, + cacheReadInputTokens: 20, + nonCachedInputTokens: 55, + textOutputTokens: 20, + rawUsage: { provider: 'hidden' }, + providerMetadata: { requestId: 'hidden' }, + }, + }), + buildRun({ + usageSummary: { + inputTokens: 40, + outputTokens: 15, + }, + }), + buildRun({ + usageSummary: { + inputTokens: 25, + }, + }), + ]); + + expect(aggregate.usageSummary).toEqual({ + totalTokens: 175, + inputTokens: 145, + outputTokens: 45, + reasoningTokens: 10, + cachedInputTokens: 20, + cacheCreationInputTokens: 5, + cacheReadInputTokens: 20, + nonCachedInputTokens: 55, + textOutputTokens: 20, + }); + expect(aggregate.usageCompleteness).toEqual({ + runCount: 3, + reportedRunCount: 2, + missingUsageRunCount: 1, + complete: false, + }); + expect(JSON.stringify(aggregate)).not.toContain('rawUsage'); + expect(JSON.stringify(aggregate)).not.toContain('providerMetadata'); + }); + + it('attributes usage by resolved provider and model with legacy fallback', () => { + const aggregate = AgentUsageService.aggregateRuns([ + buildRun({ + resolvedProvider: 'openai', + resolvedModel: 'gpt-5.4', + provider: 'gateway', + model: 'provider-response-id', + usageSummary: { + totalTokens: 10, + inputTokens: 6, + outputTokens: 4, + }, + }), + buildRun({ + resolvedProvider: null, + resolvedModel: null, + provider: 'anthropic', + model: 'claude-sonnet-4.6', + usageSummary: { + totalTokens: 20, + inputTokens: 14, + outputTokens: 6, + }, + }), + ]); + + expect(aggregate.usageByModel).toEqual([ + { + provider: 'openai', + model: 'gpt-5.4', + totalTokens: 10, + inputTokens: 6, + outputTokens: 4, + runCount: 1, + reportedRunCount: 1, + missingUsageRunCount: 0, + }, + { + provider: 'anthropic', + model: 'claude-sonnet-4.6', + totalTokens: 20, + inputTokens: 14, + outputTokens: 6, + runCount: 1, + reportedRunCount: 1, + missingUsageRunCount: 0, + }, + ]); + }); + + it('applies status-specific missing usage rules', () => { + const aggregate = AgentUsageService.aggregateRuns([ + buildRun({ status: 'queued' }), + buildRun({ status: 'starting' }), + buildRun({ status: 'running' }), + buildRun({ status: 'waiting_for_approval' }), + buildRun({ + status: 'waiting_for_input', + usageSummary: { + inputTokens: 12, + }, + }), + buildRun({ + status: 'failed', + usageSummary: { + totalTokens: 9, + }, + }), + buildRun({ + status: 'cancelled', + usageSummary: { + inputTokens: 5, + outputTokens: 3, + }, + }), + buildRun({ status: 'completed' }), + ]); + + expect(aggregate.usageSummary.totalTokens).toBe(17); + expect(aggregate.usageSummary.inputTokens).toBe(17); + expect(aggregate.usageSummary.outputTokens).toBe(3); + expect(aggregate.usageCompleteness).toEqual({ + runCount: 8, + reportedRunCount: 2, + missingUsageRunCount: 3, + complete: false, + }); + expect(aggregate.usageByModel[0]).toEqual( + expect.objectContaining({ + runCount: 8, + reportedRunCount: 2, + missingUsageRunCount: 3, + }) + ); + }); + + it('counts zero exact totals as reported usage for completeness', () => { + const aggregate = AgentUsageService.aggregateRuns([ + buildRun({ + usageSummary: { + totalTokens: 0, + }, + }), + ]); + + expect(aggregate.usageSummary.totalTokens).toBe(0); + expect(aggregate.usageCompleteness).toEqual({ + runCount: 1, + reportedRunCount: 1, + missingUsageRunCount: 0, + complete: true, + }); + }); + + it('verifies thread ownership before aggregating thread usage', async () => { + const runs = [ + buildRun({ + usageSummary: { + totalTokens: 13, + }, + }), + ]; + const query = buildRunQuery(runs); + mockGetOwnedThreadWithSession.mockResolvedValue({ + thread: { id: 7, uuid: 'thread-1' }, + session: { id: 17, uuid: 'session-1' }, + }); + mockRunQuery.mockReturnValueOnce(query); + + const usage = await AgentUsageService.getOwnedThreadUsage('thread-1', 'sample-user'); + + expect(mockGetOwnedThreadWithSession).toHaveBeenCalledWith('thread-1', 'sample-user'); + expect(query.where).toHaveBeenCalledWith({ threadId: 7 }); + expect(usage).toEqual( + expect.objectContaining({ + threadId: 'thread-1', + sessionId: 'session-1', + usageSummary: { totalTokens: 13 }, + }) + ); + }); + + it('exposes a session aggregate without filtering archived threads', async () => { + const runs = [ + buildRun({ + usageSummary: { + totalTokens: 21, + }, + }), + ]; + const query = buildRunQuery(runs); + mockGetOwnedSession.mockResolvedValue({ + id: 17, + uuid: 'session-1', + }); + mockRunQuery.mockReturnValueOnce(query); + + const usage = await AgentUsageService.getOwnedSessionUsage('session-1', 'sample-user'); + + expect(mockGetOwnedSession).toHaveBeenCalledWith('session-1', 'sample-user'); + expect(query.where).toHaveBeenCalledWith({ sessionId: 17 }); + expect(usage).toEqual( + expect.objectContaining({ + sessionId: 'session-1', + usageSummary: { totalTokens: 21 }, + }) + ); + }); + + it('aggregates multiple session usage records with one run query', async () => { + const runs = [ + buildRun({ + sessionId: 17, + usageSummary: { + totalTokens: 21, + }, + }), + buildRun({ + sessionId: 18, + usageSummary: { + inputTokens: 9, + outputTokens: 4, + }, + }), + buildRun({ + sessionId: 18, + status: 'completed', + usageSummary: {}, + }), + ]; + const query = { + whereIn: jest.fn(), + orderBy: jest.fn(), + }; + query.whereIn.mockReturnValue(query); + query.orderBy.mockReturnValueOnce(query).mockReturnValueOnce(query).mockResolvedValueOnce(runs); + mockRunQuery.mockReturnValueOnce(query); + + const usageBySessionId = await AgentUsageService.aggregateSessionsUsage([17, 18, 19]); + + expect(mockRunQuery).toHaveBeenCalledTimes(1); + expect(query.whereIn).toHaveBeenCalledWith('sessionId', [17, 18, 19]); + expect(usageBySessionId.get(17)?.usageSummary).toEqual({ totalTokens: 21 }); + expect(usageBySessionId.get(18)?.usageSummary).toEqual({ + totalTokens: 13, + inputTokens: 9, + outputTokens: 4, + }); + expect(usageBySessionId.get(18)?.usageCompleteness).toEqual({ + runCount: 2, + reportedRunCount: 1, + missingUsageRunCount: 1, + complete: false, + }); + expect(usageBySessionId.get(19)).toEqual({ + usageSummary: { totalTokens: 0 }, + usageByModel: [], + usageCompleteness: { + runCount: 0, + reportedRunCount: 0, + missingUsageRunCount: 0, + complete: true, + }, + }); + }); +}); diff --git a/src/server/services/agent/__tests__/ApprovalService.test.ts b/src/server/services/agent/__tests__/ApprovalService.test.ts index 0b414e42..e772a1f2 100644 --- a/src/server/services/agent/__tests__/ApprovalService.test.ts +++ b/src/server/services/agent/__tests__/ApprovalService.test.ts @@ -165,13 +165,25 @@ describe('ApprovalService', () => { }, fileChanges: [ { + id: 'tool-call-1:/workspace/sample-file.txt', + toolCallId: 'tool-call-1', + sourceTool: 'workspace_edit_file', path: '/workspace/sample-file.txt', displayPath: 'sample-file.txt', kind: 'edited', + stage: 'awaiting-approval', summary: 'Updated sample-file.txt', additions: 1, deletions: 1, truncated: false, + unifiedDiff: 'diff --git a/sample-file.txt b/sample-file.txt', + beforeTextPreview: 'before', + afterTextPreview: 'after', + encoding: 'utf-8', + oldSizeBytes: 6, + newSizeBytes: 5, + oldSha256: 'old-hash', + newSha256: 'new-hash', }, ], }, @@ -202,18 +214,158 @@ describe('ApprovalService', () => { commandPreview: null, fileChangePreview: [ { - path: 'sample-file.txt', - action: 'edited', + id: 'tool-call-1:/workspace/sample-file.txt', + toolCallId: 'tool-call-1', + sourceTool: 'workspace_edit_file', + path: '/workspace/sample-file.txt', + displayPath: 'sample-file.txt', + kind: 'edited', + stage: 'awaiting-approval', summary: 'Updated sample-file.txt', additions: 1, deletions: 1, truncated: false, + unifiedDiff: 'diff --git a/sample-file.txt b/sample-file.txt', + beforeTextPreview: 'before', + afterTextPreview: 'after', + encoding: 'utf-8', + oldSizeBytes: 6, + newSizeBytes: 5, + oldSha256: 'old-hash', + newSha256: 'new-hash', }, ], riskLabels: ['Workspace write'], }); }); + it('serializes Lifecycle fix approval data without exposing raw file content as an argument', () => { + const serialized = ApprovalService.serializePendingAction({ + uuid: 'action-update-file', + threadId: 3, + runId: 4, + kind: 'tool_approval', + status: 'pending', + capabilityKey: 'git_write', + title: 'Approve update_file', + description: 'update_file requires approval before it can run.', + payload: { + approvalId: 'approval-update-file', + toolCallId: 'tool-update-file', + toolName: 'mcp__lifecycle__update_file', + input: { + repository_owner: 'example-org', + repository_name: 'example-repo', + branch: 'feature/sample', + file_path: 'lifecycle.yaml', + new_content: 'services:\n sample-service:\n branch: feature/sample', + commit_message: 'Update sample service config', + }, + fileChanges: [ + { + id: 'tool-update-file:lifecycle.yaml', + toolCallId: 'tool-update-file', + sourceTool: 'update_file', + path: 'lifecycle.yaml', + displayPath: 'lifecycle.yaml', + kind: 'edited', + stage: 'awaiting-approval', + summary: 'Proposed update to lifecycle.yaml', + additions: 3, + deletions: 0, + truncated: false, + unifiedDiff: 'diff --git a/lifecycle.yaml b/lifecycle.yaml', + beforeTextPreview: 'services:\n sample-service:\n branch: main', + afterTextPreview: 'services:\n sample-service:\n branch: feature/sample', + encoding: 'utf-8', + oldSizeBytes: 43, + newSizeBytes: 54, + oldSha256: 'old-lifecycle-hash', + newSha256: 'new-lifecycle-hash', + }, + ], + }, + resolution: null, + resolvedAt: null, + createdAt: '2026-04-11T00:00:00.000Z', + updatedAt: '2026-04-11T00:00:00.000Z', + } as any); + + expect(serialized.argumentsSummary).toEqual( + expect.arrayContaining([ + { name: 'repository_owner', value: 'example-org' }, + { name: 'repository_name', value: 'example-repo' }, + { name: 'branch', value: 'feature/sample' }, + { name: 'file_path', value: 'lifecycle.yaml' }, + ]) + ); + expect(serialized.argumentsSummary).not.toEqual( + expect.arrayContaining([expect.objectContaining({ name: 'new_content' })]) + ); + expect(serialized.fileChangePreview).toEqual([ + { + id: 'tool-update-file:lifecycle.yaml', + toolCallId: 'tool-update-file', + sourceTool: 'update_file', + path: 'lifecycle.yaml', + displayPath: 'lifecycle.yaml', + kind: 'edited', + stage: 'awaiting-approval', + summary: 'Proposed update to lifecycle.yaml', + additions: 3, + deletions: 0, + truncated: false, + unifiedDiff: 'diff --git a/lifecycle.yaml b/lifecycle.yaml', + beforeTextPreview: 'services:\n sample-service:\n branch: main', + afterTextPreview: 'services:\n sample-service:\n branch: feature/sample', + encoding: 'utf-8', + oldSizeBytes: 43, + newSizeBytes: 54, + oldSha256: 'old-lifecycle-hash', + newSha256: 'new-lifecycle-hash', + }, + ]); + expect(serialized.riskLabels).toEqual(['Git write']); + }); + + it('serializes Lifecycle Kubernetes fix approvals with deployment risk labeling', () => { + const serialized = ApprovalService.serializePendingAction({ + uuid: 'action-k8s', + threadId: 3, + runId: 4, + kind: 'tool_approval', + status: 'pending', + capabilityKey: 'deploy_k8s_mutation', + title: 'Approve patch_k8s_resource', + description: 'patch_k8s_resource requires approval before it can run.', + payload: { + approvalId: 'approval-k8s', + toolCallId: 'tool-k8s', + toolName: 'mcp__lifecycle__patch_k8s_resource', + input: { + namespace: 'env-sample', + resource_type: 'deployment', + name: 'sample-service', + operation: 'restart', + }, + }, + resolution: null, + resolvedAt: null, + createdAt: '2026-04-11T00:00:00.000Z', + updatedAt: '2026-04-11T00:00:00.000Z', + } as any); + + expect(serialized.argumentsSummary).toEqual( + expect.arrayContaining([ + { name: 'namespace', value: 'env-sample' }, + { name: 'resource_type', value: 'deployment' }, + { name: 'operation', value: 'restart' }, + ]) + ); + expect(serialized.fileChangePreview).toEqual([]); + expect(serialized.riskLabels).toEqual(['Deployment change']); + }); + it('lists only pending actions for the owned thread', async () => { const query: any = {}; query.alias = jest.fn().mockReturnValue(query); @@ -310,6 +462,64 @@ describe('ApprovalService', () => { expect(mockPendingActionQuery).not.toHaveBeenCalled(); }); + it('persists forced Lifecycle fix approval requests even when the policy allows the capability', async () => { + mockGetToolName.mockReturnValue('mcp__lifecycle__update_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.syncApprovalRequestsFromMessages({ + thread: { id: 7 } as any, + run: { id: 11 } as any, + messages: [ + { + role: 'assistant', + parts: [ + { + approval: { id: 'approval-1' }, + input: { + repository_owner: 'example-org', + repository_name: 'example-repo', + branch: 'feature/sample', + file_path: 'lifecycle.yaml', + new_content: 'services: []', + }, + state: 'approval-requested', + toolCallId: 'tool-call-1', + }, + ], + } as any, + ], + approvalPolicy: { + defaultMode: 'allow', + rules: { + git_write: 'allow', + external_mcp_write: 'allow', + }, + } as any, + toolRules: [], + }); + + expect(insertQuery.insertAndFetch).toHaveBeenCalledWith( + expect.objectContaining({ + capabilityKey: 'git_write', + payload: expect.objectContaining({ + approvalId: 'approval-1', + toolCallId: 'tool-call-1', + toolName: 'mcp__lifecycle__update_file', + }), + }) + ); + }); + it('persists approval requests when a tool rule requires approval', async () => { mockGetToolName.mockReturnValue('mcp__sandbox__workspace_write_file'); @@ -427,6 +637,29 @@ describe('ApprovalService', () => { ); }); + it('does not persist stream approval requests when the current policy allows the tool', async () => { + 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', + }, + approvalPolicy: { + defaultMode: 'allow', + rules: { + workspace_write: 'allow', + }, + } as any, + toolRules: [], + }); + + expect(mockPendingActionQuery).not.toHaveBeenCalled(); + }); + it('does not reset resolved approval requests during final message sync', async () => { mockGetToolName.mockReturnValue('mcp__sandbox__workspace_write_file'); diff --git a/src/server/services/agent/__tests__/BuildContextChatService.test.ts b/src/server/services/agent/__tests__/BuildContextChatService.test.ts new file mode 100644 index 00000000..35b52e34 --- /dev/null +++ b/src/server/services/agent/__tests__/BuildContextChatService.test.ts @@ -0,0 +1,604 @@ +/** + * 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 { AgentChatStatus, AgentSessionKind, AgentWorkspaceStatus, BuildKind } from 'shared/constants'; + +const mockBuildQuery = jest.fn(); +const mockAgentSessionQuery = jest.fn(); +const mockAgentSessionTransaction = jest.fn(); +const mockAgentThreadQuery = jest.fn(); +const mockAgentSourceQuery = jest.fn(); +const mockResolveSelection = jest.fn(); +const mockGetRequiredStoredApiKey = jest.fn(); +const mockLoggerInfo = jest.fn(); + +jest.mock('server/models/Build', () => ({ + __esModule: true, + default: { + query: (...args: unknown[]) => mockBuildQuery(...args), + }, +})); + +jest.mock('server/models/AgentSession', () => ({ + __esModule: true, + default: { + query: (...args: unknown[]) => mockAgentSessionQuery(...args), + transaction: (...args: unknown[]) => mockAgentSessionTransaction(...args), + }, +})); + +jest.mock('server/models/AgentThread', () => ({ + __esModule: true, + default: { + query: (...args: unknown[]) => mockAgentThreadQuery(...args), + }, +})); + +jest.mock('server/models/AgentSource', () => ({ + __esModule: true, + default: { + query: (...args: unknown[]) => mockAgentSourceQuery(...args), + }, +})); + +jest.mock('../ProviderRegistry', () => ({ + __esModule: true, + default: { + resolveSelection: (...args: unknown[]) => mockResolveSelection(...args), + getRequiredStoredApiKey: (...args: unknown[]) => mockGetRequiredStoredApiKey(...args), + }, +})); + +jest.mock('server/lib/logger', () => ({ + getLogger: () => ({ + info: mockLoggerInfo, + }), +})); + +jest.mock('server/lib/dependencies', () => ({})); + +import BuildContextChatService, { BuildContextChatBuildNotFoundError } from '../BuildContextChatService'; + +const NOW = '2026-04-30T12:00:00.000Z'; +const TEST_TRX = { name: 'trx' }; + +function mockBuildLookup(build: Record | null) { + const withGraphFetched = jest.fn().mockResolvedValue(build); + const findOne = jest.fn(() => ({ withGraphFetched })); + mockBuildQuery.mockReturnValueOnce({ findOne }); + return { findOne, withGraphFetched }; +} + +function buildSessionReadQuery({ firstResult, findOneResult }: { firstResult: unknown; findOneResult: unknown }) { + const query = { + where: jest.fn(() => query), + orderBy: jest.fn(() => query), + first: jest.fn().mockResolvedValue(firstResult), + findOne: jest.fn().mockResolvedValue(findOneResult), + }; + return query; +} + +function arrangeCreatePath({ + build, + reuseSession = null, +}: { + build: Record; + reuseSession?: Record | null; +}) { + const buildLookup = mockBuildLookup(build); + const insertedThread = { + id: 23, + uuid: 'thread-uuid-1', + sessionId: 17, + isDefault: true, + archivedAt: null, + }; + let insertedSession: Record | null = null; + let finalizedSession: Record | null = null; + + const sessionInsertAndFetch = jest.fn(async (payload) => { + insertedSession = { + id: 17, + uuid: payload.uuid, + defaultThreadId: null, + ...payload, + updatedAt: NOW, + createdAt: NOW, + }; + return insertedSession; + }); + const sessionPatchAndFetchById = jest.fn(async (_id, patch) => { + finalizedSession = { + ...insertedSession, + ...patch, + }; + return finalizedSession; + }); + const sourceInsertAndFetch = jest.fn(async (payload) => ({ id: 31, ...payload })); + const threadInsertAndFetch = jest.fn(async (payload) => ({ + ...insertedThread, + ...payload, + })); + + const reuseQuery = buildSessionReadQuery({ + firstResult: reuseSession, + findOneResult: () => finalizedSession || insertedSession, + }); + const ownedQuery = buildSessionReadQuery({ + firstResult: null, + findOneResult: () => finalizedSession || insertedSession, + }); + const sessionReadQueries = [reuseQuery, ownedQuery]; + + mockAgentSessionQuery.mockImplementation((trx?: unknown) => { + if (trx) { + return { + insertAndFetch: sessionInsertAndFetch, + patchAndFetchById: sessionPatchAndFetchById, + }; + } + + const nextQuery = + sessionReadQueries.shift() || + buildSessionReadQuery({ firstResult: null, findOneResult: () => finalizedSession || insertedSession }); + + return { + ...nextQuery, + findOne: jest.fn(async (...args: unknown[]) => { + const result = await nextQuery.findOne(...args); + return typeof result === 'function' ? result() : result; + }), + }; + }); + + mockAgentSessionTransaction.mockImplementation(async (callback) => callback(TEST_TRX)); + + mockAgentThreadQuery.mockImplementation((trx?: unknown) => { + if (trx) { + return { + insertAndFetch: threadInsertAndFetch, + }; + } + + return { + findOne: jest.fn().mockResolvedValue(insertedThread), + insertAndFetch: threadInsertAndFetch, + }; + }); + + mockAgentSourceQuery.mockReturnValue({ + insertAndFetch: sourceInsertAndFetch, + }); + + return { + buildLookup, + reuseQuery, + sessionInsertAndFetch, + sessionPatchAndFetchById, + sourceInsertAndFetch, + threadInsertAndFetch, + get insertedThread() { + return insertedThread; + }, + }; +} + +function arrangeReusePath({ + build, + existingSession, + defaultThread, +}: { + build: Record; + existingSession: Record; + defaultThread: Record | null; +}) { + const buildLookup = mockBuildLookup(build); + const reuseQuery = buildSessionReadQuery({ firstResult: existingSession, findOneResult: existingSession }); + const ownedQuery = buildSessionReadQuery({ firstResult: null, findOneResult: existingSession }); + const sessionReadQueries = [reuseQuery, ownedQuery]; + const recreatedThread = { + id: 29, + uuid: 'thread-uuid-recreated', + sessionId: existingSession.id, + isDefault: true, + archivedAt: null, + metadata: { + sessionUuid: existingSession.uuid, + }, + }; + const threadFindOne = jest.fn().mockResolvedValue(defaultThread); + const threadInsertAndFetch = jest.fn().mockResolvedValue(recreatedThread); + const threadQueries = defaultThread + ? [{ findOne: threadFindOne }] + : [{ findOne: threadFindOne }, { insertAndFetch: threadInsertAndFetch }]; + + mockAgentSessionQuery.mockImplementation((trx?: unknown) => { + if (trx) { + throw new Error('create path should not run when an active build-context chat can be reused'); + } + + return sessionReadQueries.shift() || ownedQuery; + }); + mockAgentSessionTransaction.mockImplementation(() => { + throw new Error('transaction should not run when an active build-context chat can be reused'); + }); + mockAgentThreadQuery.mockImplementation(() => threadQueries.shift() || { findOne: threadFindOne }); + + return { + buildLookup, + reuseQuery, + threadFindOne, + threadInsertAndFetch, + recreatedThread, + }; +} + +function sampleBuild(overrides: Record = {}) { + return { + uuid: 'build-uuid-1', + kind: BuildKind.ENVIRONMENT, + namespace: 'env-sample-123', + sha: 'commit-sha-1', + baseBuild: { + uuid: 'base-build-uuid-1', + }, + pullRequest: { + fullName: 'example-org/example-repo', + branchName: 'feature/sample-change', + pullRequestNumber: 42, + latestCommit: 'latest-pr-commit-1', + }, + ...overrides, + }; +} + +function sampleActiveChatSession(overrides: Record = {}) { + return { + id: 17, + uuid: 'session-uuid-existing', + buildUuid: 'build-uuid-1', + buildKind: null, + sessionKind: AgentSessionKind.CHAT, + userId: 'sample-user', + status: 'active', + chatStatus: AgentChatStatus.READY, + workspaceStatus: AgentWorkspaceStatus.NONE, + createdAt: '2026-04-30T11:00:00.000Z', + updatedAt: '2026-04-30T11:10:00.000Z', + ...overrides, + }; +} + +describe('BuildContextChatService', () => { + beforeEach(() => { + jest.clearAllMocks(); + jest.useFakeTimers().setSystemTime(new Date(NOW)); + mockResolveSelection.mockResolvedValue({ + provider: 'gemini', + modelId: 'gemini-3-flash-preview', + }); + mockGetRequiredStoredApiKey.mockResolvedValue('sample-api-key'); + }); + + afterEach(() => { + jest.useRealTimers(); + }); + + it('creates a ready chat session, default thread, and build-context source for a valid build', async () => { + const build = sampleBuild(); + const arranged = arrangeCreatePath({ build }); + + const result = await BuildContextChatService.launchBuildContextChat({ + buildUuid: 'build-uuid-1', + userId: 'sample-user', + userIdentity: { + userId: 'sample-user', + githubUsername: 'sample-user', + preferredUsername: null, + email: 'sample-user@example.com', + firstName: null, + lastName: null, + displayName: 'Sample User', + gitUserName: 'Sample User', + gitUserEmail: 'sample-user@example.com', + }, + model: 'gemini-3-flash-preview', + }); + + expect(arranged.buildLookup.findOne).toHaveBeenCalledWith({ uuid: 'build-uuid-1' }); + expect(arranged.buildLookup.withGraphFetched).toHaveBeenCalledWith('[pullRequest, baseBuild]'); + expect(mockResolveSelection).toHaveBeenCalledWith({ + repoFullName: 'example-org/example-repo', + requestedProvider: undefined, + requestedModelId: 'gemini-3-flash-preview', + }); + expect(mockGetRequiredStoredApiKey).toHaveBeenCalledWith({ + provider: 'gemini', + userIdentity: { + userId: 'sample-user', + githubUsername: 'sample-user', + }, + }); + expect(arranged.sessionInsertAndFetch).toHaveBeenCalledWith( + expect.objectContaining({ + buildUuid: 'build-uuid-1', + buildKind: null, + sessionKind: AgentSessionKind.CHAT, + userId: 'sample-user', + ownerGithubUsername: 'sample-user', + podName: null, + namespace: null, + pvcName: null, + status: 'active', + chatStatus: AgentChatStatus.READY, + workspaceStatus: AgentWorkspaceStatus.NONE, + workspaceRepos: [ + { + repo: 'example-org/example-repo', + repoUrl: 'https://github.com/example-org/example-repo.git', + branch: 'feature/sample-change', + revision: 'commit-sha-1', + mountPath: '/workspace', + primary: true, + }, + ], + selectedServices: [], + }) + ); + + const launchMetadata = { + buildUuid: 'build-uuid-1', + buildKind: BuildKind.ENVIRONMENT, + sessionKind: AgentSessionKind.CHAT, + namespace: 'env-sample-123', + baseBuildUuid: 'base-build-uuid-1', + revision: 'commit-sha-1', + pullRequest: { + fullName: 'example-org/example-repo', + branchName: 'feature/sample-change', + pullRequestNumber: 42, + }, + contextFreshAt: NOW, + }; + expect(arranged.sourceInsertAndFetch).toHaveBeenCalledWith( + expect.objectContaining({ + sessionId: 17, + adapter: 'blank_workspace', + status: 'ready', + input: { + ...launchMetadata, + defaults: { + provider: 'gemini', + model: 'gemini-3-flash-preview', + }, + }, + preparedSource: expect.objectContaining({ + workspaceLayout: expect.objectContaining({ + repos: [], + }), + metadata: launchMetadata, + }), + }) + ); + expect(result).toMatchObject({ + thread: arranged.insertedThread, + created: true, + reused: false, + buildContext: { + buildUuid: 'build-uuid-1', + buildKind: BuildKind.ENVIRONMENT, + namespace: 'env-sample-123', + baseBuildUuid: 'base-build-uuid-1', + revision: 'commit-sha-1', + pullRequest: { + fullName: 'example-org/example-repo', + branchName: 'feature/sample-change', + pullRequestNumber: 42, + }, + contextFreshAt: NOW, + }, + }); + expect(result.session).toMatchObject({ + buildUuid: 'build-uuid-1', + buildKind: null, + sessionKind: AgentSessionKind.CHAT, + workspaceStatus: AgentWorkspaceStatus.NONE, + chatStatus: AgentChatStatus.READY, + podName: null, + namespace: null, + pvcName: null, + }); + expect(mockLoggerInfo).toHaveBeenCalledWith( + expect.stringContaining('Session: launched build-context chat buildUuid=build-uuid-1') + ); + }); + + it('throws a typed not-found error for an invalid or unknown buildUuid', async () => { + mockBuildLookup(null); + + await expect( + BuildContextChatService.launchBuildContextChat({ + buildUuid: 'missing-build-uuid', + userId: 'sample-user', + }) + ).rejects.toBeInstanceOf(BuildContextChatBuildNotFoundError); + + expect(mockAgentSessionTransaction).not.toHaveBeenCalled(); + }); + + it('re-reads and reuses the active chat when a concurrent launch wins the unique constraint race', async () => { + mockBuildLookup(sampleBuild()); + const racedSession = sampleActiveChatSession({ + uuid: 'session-uuid-raced', + }); + const defaultThread = { + id: 23, + uuid: 'thread-uuid-raced', + sessionId: 17, + isDefault: true, + archivedAt: null, + }; + const reuseMissQuery = buildSessionReadQuery({ firstResult: null, findOneResult: null }); + const reuseHitQuery = buildSessionReadQuery({ firstResult: racedSession, findOneResult: racedSession }); + const ownedQuery = buildSessionReadQuery({ firstResult: null, findOneResult: racedSession }); + const sessionReadQueries = [reuseMissQuery, reuseHitQuery, ownedQuery]; + mockAgentSessionQuery.mockImplementation((trx?: unknown) => { + if (trx) { + return { + insertAndFetch: jest.fn(), + patchAndFetchById: jest.fn(), + }; + } + + return sessionReadQueries.shift() || ownedQuery; + }); + mockAgentSessionTransaction.mockRejectedValueOnce({ + code: '23505', + constraint: 'agent_sessions_active_build_context_chat_unique', + }); + mockAgentThreadQuery.mockReturnValue({ + findOne: jest.fn().mockResolvedValue(defaultThread), + }); + + const result = await BuildContextChatService.launchBuildContextChat({ + buildUuid: 'build-uuid-1', + userId: 'sample-user', + }); + + expect(mockAgentSessionTransaction).toHaveBeenCalled(); + expect(reuseMissQuery.first).toHaveBeenCalledTimes(1); + expect(reuseHitQuery.first).toHaveBeenCalledTimes(1); + expect(result).toMatchObject({ + session: racedSession, + thread: defaultThread, + created: false, + reused: true, + }); + }); + + it('reuses the latest active same user build-context chat and active default thread', async () => { + const existingSession = sampleActiveChatSession(); + const defaultThread = { + id: 23, + uuid: 'thread-uuid-existing', + sessionId: 17, + isDefault: true, + archivedAt: null, + }; + const arranged = arrangeReusePath({ + build: sampleBuild(), + existingSession, + defaultThread, + }); + + const result = await BuildContextChatService.launchBuildContextChat({ + buildUuid: 'build-uuid-1', + userId: 'sample-user', + }); + + expect(arranged.reuseQuery.where).toHaveBeenCalledWith({ + userId: 'sample-user', + buildUuid: 'build-uuid-1', + sessionKind: AgentSessionKind.CHAT, + status: 'active', + chatStatus: AgentChatStatus.READY, + }); + expect(arranged.reuseQuery.orderBy).toHaveBeenNthCalledWith(1, 'updatedAt', 'desc'); + expect(arranged.reuseQuery.orderBy).toHaveBeenNthCalledWith(2, 'createdAt', 'desc'); + expect(result).toMatchObject({ + session: existingSession, + thread: defaultThread, + created: false, + reused: true, + }); + expect(mockAgentSessionTransaction).not.toHaveBeenCalled(); + }); + + it('creates a separate chat when a different user launches the same buildUuid', async () => { + const arranged = arrangeCreatePath({ build: sampleBuild() }); + + const result = await BuildContextChatService.launchBuildContextChat({ + buildUuid: 'build-uuid-1', + userId: 'sample-user-2', + }); + + expect(arranged.reuseQuery.where).toHaveBeenCalledWith( + expect.objectContaining({ + userId: 'sample-user-2', + buildUuid: 'build-uuid-1', + }) + ); + expect(result).toMatchObject({ + created: true, + reused: false, + }); + }); + + it('ignores active environment or sandbox sessions and ended or errored chat sessions', async () => { + const arranged = arrangeCreatePath({ build: sampleBuild({ kind: BuildKind.SANDBOX }) }); + + const result = await BuildContextChatService.launchBuildContextChat({ + buildUuid: 'build-uuid-1', + userId: 'sample-user', + }); + + expect(arranged.reuseQuery.where).toHaveBeenCalledWith({ + userId: 'sample-user', + buildUuid: 'build-uuid-1', + sessionKind: AgentSessionKind.CHAT, + status: 'active', + chatStatus: AgentChatStatus.READY, + }); + expect(result).toMatchObject({ + created: true, + reused: false, + }); + }); + + it('recreates a missing or archived default thread when reusing a chat session', async () => { + const existingSession = sampleActiveChatSession(); + const arranged = arrangeReusePath({ + build: sampleBuild(), + existingSession, + defaultThread: null, + }); + + const result = await BuildContextChatService.launchBuildContextChat({ + buildUuid: 'build-uuid-1', + userId: 'sample-user', + }); + + expect(arranged.threadFindOne).toHaveBeenCalledWith({ + sessionId: 17, + isDefault: true, + archivedAt: null, + }); + expect(arranged.threadInsertAndFetch).toHaveBeenCalledWith({ + sessionId: 17, + title: 'Default thread', + isDefault: true, + metadata: { + sessionUuid: 'session-uuid-existing', + }, + }); + expect(result).toMatchObject({ + thread: arranged.recreatedThread, + created: false, + reused: true, + }); + }); +}); diff --git a/src/server/services/agent/__tests__/CapabilityCatalog.test.ts b/src/server/services/agent/__tests__/CapabilityCatalog.test.ts new file mode 100644 index 00000000..addb7d77 --- /dev/null +++ b/src/server/services/agent/__tests__/CapabilityCatalog.test.ts @@ -0,0 +1,60 @@ +/** + * 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 { + AGENT_CAPABILITY_CATALOG, + AGENT_CAPABILITY_CATALOG_IDS, + getAgentCapabilityCatalogEntry, + isAgentCapabilityCatalogId, +} from '../capabilityCatalog'; +import { isAgentCapabilityKey } from '../types'; + +describe('agent capability catalog', () => { + it('has unique ids with non-empty labels and descriptions', () => { + expect(new Set(AGENT_CAPABILITY_CATALOG_IDS).size).toBe(AGENT_CAPABILITY_CATALOG_IDS.length); + expect(new Set(AGENT_CAPABILITY_CATALOG.map((entry) => entry.id)).size).toBe(AGENT_CAPABILITY_CATALOG.length); + + for (const entry of AGENT_CAPABILITY_CATALOG) { + expect(entry.label.trim()).not.toHaveLength(0); + expect(entry.description.trim()).not.toHaveLength(0); + expect(isAgentCapabilityCatalogId(entry.id)).toBe(true); + } + }); + + it('maps runtime capability keys to known approval capabilities', () => { + for (const entry of AGENT_CAPABILITY_CATALOG) { + if (entry.runtimeCapabilityKey) { + expect(isAgentCapabilityKey(entry.runtimeCapabilityKey)).toBe(true); + } + } + }); + + it('contains the required governable capability families', () => { + expect(getAgentCapabilityCatalogEntry('read_context').label).toBe('Read/context'); + expect(getAgentCapabilityCatalogEntry('diagnostics_codefresh').category).toBe('diagnostics'); + expect(getAgentCapabilityCatalogEntry('diagnostics_kubernetes').category).toBe('diagnostics'); + expect(getAgentCapabilityCatalogEntry('diagnostics_database').category).toBe('diagnostics'); + expect(getAgentCapabilityCatalogEntry('github_read').category).toBe('source_control'); + expect(getAgentCapabilityCatalogEntry('github_write').runtimeCapabilityKey).toBe('git_write'); + expect(getAgentCapabilityCatalogEntry('workspace_files').runtimeCapabilityKey).toBe('workspace_write'); + expect(getAgentCapabilityCatalogEntry('workspace_shell').runtimeCapabilityKey).toBe('shell_exec'); + expect(getAgentCapabilityCatalogEntry('workspace_git').runtimeCapabilityKey).toBe('git_write'); + expect(getAgentCapabilityCatalogEntry('external_mcp_read').runtimeCapabilityKey).toBe('external_mcp_read'); + expect(getAgentCapabilityCatalogEntry('external_mcp_write').runtimeCapabilityKey).toBe('external_mcp_write'); + expect(getAgentCapabilityCatalogEntry('preview_publish').category).toBe('preview'); + expect(getAgentCapabilityCatalogEntry('approval_controls').category).toBe('approval'); + }); +}); diff --git a/src/server/services/agent/__tests__/CapabilityService.test.ts b/src/server/services/agent/__tests__/CapabilityService.test.ts index ac0a1940..2f69459f 100644 --- a/src/server/services/agent/__tests__/CapabilityService.test.ts +++ b/src/server/services/agent/__tests__/CapabilityService.test.ts @@ -22,24 +22,56 @@ const mockListTools = jest.fn(); const mockCallTool = jest.fn(); const mockClose = jest.fn(); const mockLoggerWarn = jest.fn(); -const mockModeForCapability = jest.fn(() => 'allow'); +const mockModeForCapability = jest.fn((_policy?: unknown, _capability?: unknown) => 'allow'); +const mockGetEffectivePolicy = jest.fn(); +const mockGetEffectiveAgentConfig = jest.fn(); +const mockCapabilityForExternalMcpTool = jest.fn((_toolName?: string) => 'external_mcp_read'); const mockPublishChatHttpPort = jest.fn(); const mockFindSession = jest.fn(); const mockResolveWorkspaceGatewayBaseUrl = jest.fn(); const mockEnsureChatSandbox = jest.fn(); +const mockDiagnosticToolExecute = jest.fn(); +const mockParseYamlConfigFromBranch = jest.fn(); +const mockGithubOctokitRequest = jest.fn(); +const mockGithubClientInstances: Array<{ + setAllowedBranch: jest.Mock; + setReferencedFiles: jest.Mock; + setExcludedFilePatterns: jest.Mock; + setAllowedWritePatterns: jest.Mock; + isFilePathAllowed: jest.Mock; + validateBranch: jest.Mock; + getOctokit: jest.Mock; +}> = []; let currentTransport: Record | null = null; +function mockMakeDiagnosticToolClass(name: string, description = `${name} description`) { + return jest.fn().mockImplementation(() => ({ + name, + description, + parameters: { + type: 'object', + properties: {}, + }, + execute: (args: Record, signal?: AbortSignal) => mockDiagnosticToolExecute(name, args, signal), + })); +} + jest.mock('ai', () => ({ - dynamicTool: (...args: unknown[]) => mockDynamicTool(...args), - jsonSchema: (...args: unknown[]) => mockJsonSchema(...args), + dynamicTool: (config: unknown) => mockDynamicTool(config), + jsonSchema: (schema: unknown) => mockJsonSchema(schema), })); -jest.mock('server/services/ai/mcp/config', () => ({ - McpConfigService: jest.fn().mockImplementation(() => ({ - resolveServers: (...args: unknown[]) => mockResolveServers(...args), - })), -})); +jest.mock('server/services/agentRuntime/mcp/config', () => { + const actual = jest.requireActual('server/services/agentRuntime/mcp/config'); + return { + __esModule: true, + ...actual, + McpConfigService: jest.fn().mockImplementation(() => ({ + resolveServers: (...args: unknown[]) => mockResolveServers(...args), + })), + }; +}); jest.mock('server/models/AgentSession', () => ({ __esModule: true, @@ -65,7 +97,7 @@ jest.mock('../SandboxService', () => ({ }, })); -jest.mock('server/services/ai/mcp/client', () => ({ +jest.mock('server/services/agentRuntime/mcp/client', () => ({ McpClientManager: jest.fn().mockImplementation(() => ({ connect: (...args: unknown[]) => mockConnect(...args), listTools: (...args: unknown[]) => mockListTools(...args), @@ -74,6 +106,67 @@ jest.mock('server/services/ai/mcp/client', () => ({ })), })); +jest.mock('server/services/agent/tools/shared/k8sClient', () => ({ + K8sClient: jest.fn(), +})); +jest.mock('server/services/agent/tools/shared/githubClient', () => ({ + GitHubClient: jest.fn().mockImplementation(() => { + const client = { + setAllowedBranch: jest.fn(), + setReferencedFiles: jest.fn(), + setExcludedFilePatterns: jest.fn(), + setAllowedWritePatterns: jest.fn(), + isFilePathAllowed: jest.fn(), + validateBranch: jest.fn(), + getOctokit: jest.fn().mockResolvedValue({ request: mockGithubOctokitRequest }), + }; + mockGithubClientInstances.push(client); + return client; + }), +})); + +jest.mock('server/lib/yamlConfigParser', () => ({ + YamlConfigParser: jest.fn().mockImplementation(() => ({ + parseYamlConfigFromBranch: (...args: unknown[]) => mockParseYamlConfigFromBranch(...args), + })), +})); +jest.mock('server/services/agent/tools/shared/databaseClient', () => ({ + DatabaseClient: jest.fn(), +})); +jest.mock('server/services/agent/tools/codefresh/getCodefreshLogs', () => ({ + GetCodefreshLogsTool: mockMakeDiagnosticToolClass('get_codefresh_logs'), +})); +jest.mock('server/services/agent/tools/k8s/getK8sResources', () => ({ + GetK8sResourcesTool: mockMakeDiagnosticToolClass('get_k8s_resources'), +})); +jest.mock('server/services/agent/tools/k8s/getPodLogs', () => ({ + GetPodLogsTool: mockMakeDiagnosticToolClass('get_pod_logs'), +})); +jest.mock('server/services/agent/tools/k8s/getLifecycleLogs', () => ({ + GetLifecycleLogsTool: mockMakeDiagnosticToolClass('get_lifecycle_logs'), +})); +jest.mock('server/services/agent/tools/k8s/queryDatabase', () => ({ + QueryDatabaseTool: mockMakeDiagnosticToolClass('query_database'), +})); +jest.mock('server/services/agent/tools/github/getFile', () => ({ + GetFileTool: mockMakeDiagnosticToolClass('get_file'), +})); +jest.mock('server/services/agent/tools/github/listDirectory', () => ({ + ListDirectoryTool: mockMakeDiagnosticToolClass('list_directory'), +})); +jest.mock('server/services/agent/tools/github/getIssueComment', () => ({ + GetIssueCommentTool: mockMakeDiagnosticToolClass('get_issue_comment'), +})); +jest.mock('server/services/agent/tools/github/updateFile', () => ({ + UpdateFileTool: mockMakeDiagnosticToolClass('update_file'), +})); +jest.mock('server/services/agent/tools/github/updatePrLabels', () => ({ + UpdatePrLabelsTool: mockMakeDiagnosticToolClass('update_pr_labels'), +})); +jest.mock('server/services/agent/tools/k8s/patchK8sResource', () => ({ + PatchK8sResourceTool: mockMakeDiagnosticToolClass('patch_k8s_resource'), +})); + jest.mock('server/lib/logger', () => ({ getLogger: () => ({ warn: (...args: unknown[]) => mockLoggerWarn(...args), @@ -101,15 +194,55 @@ jest.mock('server/lib/agentSession/runtimeConfig', () => { jest.mock('../PolicyService', () => ({ __esModule: true, default: { + getEffectivePolicy: (...args: unknown[]) => (mockGetEffectivePolicy as any)(...args), capabilityForSessionWorkspaceTool: jest.fn(() => 'read'), - capabilityForExternalMcpTool: jest.fn(() => 'external_mcp_read'), - modeForCapability: (...args: unknown[]) => mockModeForCapability(...args), + capabilityForExternalMcpTool: (...args: unknown[]) => (mockCapabilityForExternalMcpTool as any)(...args), + modeForCapability: (...args: unknown[]) => (mockModeForCapability as any)(...args), + }, +})); + +jest.mock('server/services/agentRuntime/config/agentRuntimeConfig', () => ({ + __esModule: true, + default: { + getInstance: jest.fn(() => ({ + getEffectiveConfig: (...args: unknown[]) => mockGetEffectiveAgentConfig(...args), + })), }, })); import AgentCapabilityService from '../CapabilityService'; import { SessionWorkspaceGatewayUnavailableError } from '../errors'; +const defaultResolvedCapabilityAccess = [ + 'read_context', + 'diagnostics_logs', + 'diagnostics_codefresh', + 'diagnostics_kubernetes', + 'diagnostics_database', + 'github_read', + 'github_write', + 'workspace_files', + 'workspace_shell', + 'workspace_git', + 'network_access', + 'preview_publish', + 'external_mcp_read', + 'external_mcp_write', + 'approval_controls', +].map((capabilityId) => ({ + capabilityId, + effectiveAvailability: 'all_users' as const, + allowed: true, + approvalMode: 'allow' as const, +})); + +function buildToolSetForTest(args: Parameters[0]) { + return AgentCapabilityService.buildToolSet({ + resolvedCapabilityAccess: defaultResolvedCapabilityAccess, + ...args, + }); +} + describe('AgentCapabilityService.buildToolSet', () => { const session = { uuid: 'session-123', @@ -124,6 +257,7 @@ describe('AgentCapabilityService.buildToolSet', () => { githubUsername: 'sample-user', } as any; const stdioServer = { + scope: 'global', slug: 'figma', name: 'Figma', transport: { @@ -156,7 +290,18 @@ describe('AgentCapabilityService.buildToolSet', () => { beforeEach(() => { jest.clearAllMocks(); + mockGithubClientInstances.length = 0; + mockGithubOctokitRequest.mockResolvedValue({ + data: { + sha: 'existing-file-sha', + content: Buffer.from('services:\n sample-service:\n branch: main').toString('base64'), + }, + }); mockModeForCapability.mockReturnValue('allow'); + mockGetEffectivePolicy.mockResolvedValue({ rules: {} }); + mockGetEffectiveAgentConfig.mockResolvedValue({}); + mockParseYamlConfigFromBranch.mockResolvedValue({ services: [] }); + mockCapabilityForExternalMcpTool.mockReturnValue('external_mcp_read'); currentTransport = null; mockResolveServers.mockResolvedValue([stdioServer]); mockFindSession.mockResolvedValue(session); @@ -190,6 +335,10 @@ describe('AgentCapabilityService.buildToolSet', () => { mockConnect.mockImplementation(async (transport) => { currentTransport = transport as Record; }); + mockDiagnosticToolExecute.mockResolvedValue({ + success: true, + agentContent: JSON.stringify({ ok: true }), + }); mockListTools.mockImplementation(async () => { if ( currentTransport && @@ -219,8 +368,270 @@ describe('AgentCapabilityService.buildToolSet', () => { mockClose.mockResolvedValue(undefined); }); + it('resolves repo-scoped context from build-context chat workspaceRepos', async () => { + mockFindSession.mockResolvedValueOnce({ + uuid: 'session-build-context', + userId: 'sample-user', + sessionKind: 'chat', + workspaceRepos: [ + { + repo: 'example-org/example-repo', + repoUrl: 'https://github.com/example-org/example-repo.git', + branch: 'feature/sample', + revision: 'commit-sha-1', + mountPath: '/workspace', + primary: true, + }, + ], + selectedServices: [], + }); + const policy = { rules: { deploy_k8s_read: 'allow' } }; + mockGetEffectivePolicy.mockResolvedValueOnce(policy); + + const result = await AgentCapabilityService.resolveSessionContext('session-build-context', userIdentity); + + expect(mockFindSession).toHaveBeenCalledWith({ + uuid: 'session-build-context', + userId: 'sample-user', + }); + expect(mockGetEffectivePolicy).toHaveBeenCalledWith('example-org/example-repo'); + expect(result).toEqual({ + session: expect.objectContaining({ + uuid: 'session-build-context', + }), + repoFullName: 'example-org/example-repo', + approvalPolicy: policy, + capabilityPolicy: undefined, + customAgentCreationPolicy: undefined, + }); + }); + it('routes stdio MCP execution through the session-pod proxy endpoint', async () => { + const tools = await buildToolSetForTest({ + session, + repoFullName: 'example-org/example-repo', + userIdentity, + approvalPolicy: {} as any, + workspaceToolDiscoveryTimeoutMs: 4500, + workspaceToolExecutionTimeoutMs: 22000, + }); + + expect(mockConnect).toHaveBeenCalledWith( + { + type: 'http', + url: 'http://agent-123.env-sample.svc.cluster.local:13338/mcp', + }, + 4500 + ); + expect(mockListTools).toHaveBeenCalledWith(4500); + + const tool = tools.mcp__figma__get_design_context as unknown as { + execute: (input: Record) => Promise; + }; + expect(tool).toBeDefined(); + + await tool.execute({}); + + expect(mockConnect).toHaveBeenCalledWith( + { + type: 'http', + url: 'http://agent-123.env-sample.svc.cluster.local:13338/servers/figma/mcp', + }, + 30000 + ); + expect(mockCallTool).toHaveBeenCalledWith('get_design_context', {}, 30000); + }); + + it('does not expose runtime tools without resolved run-plan capabilities', async () => { const tools = await AgentCapabilityService.buildToolSet({ + session: { + ...session, + sessionKind: 'chat', + workspaceStatus: 'none', + podName: null, + namespace: null, + } as any, + repoFullName: 'example-org/example-repo', + userIdentity, + approvalPolicy: {} as any, + workspaceToolDiscoveryTimeoutMs: 4500, + workspaceToolExecutionTimeoutMs: 22000, + }); + + expect(tools.mcp__figma__get_design_context).toBeUndefined(); + expect(tools.mcp__sandbox__workspace_exec).toBeUndefined(); + expect(tools.mcp__sandbox__workspace_write_file).toBeUndefined(); + expect(tools.mcp__lifecycle__publish_http).toBeUndefined(); + }); + + it('omits external MCP tools whose resolved catalog capability is unavailable', async () => { + mockCapabilityForExternalMcpTool.mockImplementation((toolName: string) => + toolName === 'update_design' ? 'external_mcp_write' : 'external_mcp_read' + ); + mockResolveServers.mockResolvedValue([ + { + ...stdioServer, + discoveredTools: [ + { + name: 'get_design_context', + description: 'Read design context', + inputSchema: { + type: 'object', + properties: {}, + }, + annotations: { + readOnlyHint: true, + }, + }, + { + name: 'update_design', + description: 'Update design context', + inputSchema: { + type: 'object', + properties: {}, + }, + }, + ], + }, + ]); + + const tools = await buildToolSetForTest({ + session, + repoFullName: 'example-org/example-repo', + userIdentity, + approvalPolicy: {} as any, + resolvedCapabilityAccess: [ + { + capabilityId: 'external_mcp_read', + effectiveAvailability: 'all_users', + allowed: true, + approvalMode: 'allow', + }, + { + capabilityId: 'external_mcp_write', + effectiveAvailability: 'admin_only', + allowed: false, + reason: 'admin_only', + }, + ], + workspaceToolDiscoveryTimeoutMs: 4500, + workspaceToolExecutionTimeoutMs: 22000, + }); + + expect(tools.mcp__figma__get_design_context).toBeDefined(); + expect(tools.mcp__figma__update_design).toBeUndefined(); + }); + + it('skips optional external MCP tools for explicit empty runtime MCP selections', async () => { + const tools = await buildToolSetForTest({ + session, + repoFullName: 'example-org/example-repo', + userIdentity, + approvalPolicy: {} as any, + selectedRuntimeMcpConnectionRefs: [], + workspaceToolDiscoveryTimeoutMs: 4500, + workspaceToolExecutionTimeoutMs: 22000, + }); + + expect(tools.mcp__figma__get_design_context).toBeUndefined(); + expect(tools.mcp__sandbox__workspace_read_file).toBeDefined(); + }); + + it('registers only the selected runtime MCP connection tools', async () => { + mockResolveServers.mockResolvedValue([ + stdioServer, + { + ...stdioServer, + slug: 'docs', + name: 'Docs', + discoveredTools: [ + { + name: 'search_docs', + description: 'Search docs', + inputSchema: { + type: 'object', + properties: {}, + }, + annotations: { + readOnlyHint: true, + }, + }, + ], + }, + ]); + + const tools = await buildToolSetForTest({ + session, + repoFullName: 'example-org/example-repo', + userIdentity, + approvalPolicy: {} as any, + selectedRuntimeMcpConnectionRefs: ['global:docs'], + workspaceToolDiscoveryTimeoutMs: 4500, + workspaceToolExecutionTimeoutMs: 22000, + }); + + expect(tools.mcp__figma__get_design_context).toBeUndefined(); + expect(tools.mcp__docs__search_docs).toBeDefined(); + expect(tools.mcp__sandbox__workspace_read_file).toBeDefined(); + }); + + it('filters selected runtime MCP connections by scope and slug', async () => { + mockResolveServers.mockResolvedValue([ + { + ...stdioServer, + slug: 'docs', + name: 'Global Docs', + discoveredTools: [ + { + name: 'search_global_docs', + description: 'Search global docs', + inputSchema: { + type: 'object', + properties: {}, + }, + annotations: { + readOnlyHint: true, + }, + }, + ], + }, + { + ...stdioServer, + scope: 'example-org/example-repo', + slug: 'docs', + name: 'Repo Docs', + discoveredTools: [ + { + name: 'search_repo_docs', + description: 'Search repo docs', + inputSchema: { + type: 'object', + properties: {}, + }, + annotations: { + readOnlyHint: true, + }, + }, + ], + }, + ]); + + const tools = await buildToolSetForTest({ + session, + repoFullName: 'example-org/example-repo', + userIdentity, + approvalPolicy: {} as any, + selectedRuntimeMcpConnectionRefs: ['example-org/example-repo:docs'], + workspaceToolDiscoveryTimeoutMs: 4500, + workspaceToolExecutionTimeoutMs: 22000, + }); + + expect(tools.mcp__docs__search_global_docs).toBeUndefined(); + expect(tools.mcp__docs__search_repo_docs).toBeDefined(); + }); + + it('uses the configured workspace execution timeout for sandbox tools', async () => { + const tools = await buildToolSetForTest({ session, repoFullName: 'example-org/example-repo', userIdentity, @@ -229,132 +640,701 @@ describe('AgentCapabilityService.buildToolSet', () => { workspaceToolExecutionTimeoutMs: 22000, }); - expect(mockConnect).toHaveBeenCalledWith( - { - type: 'http', - url: 'http://agent-123.env-sample.svc.cluster.local:13338/mcp', - }, - 4500 + const tool = tools.mcp__sandbox__workspace_read_file as unknown as { + execute: (input: Record) => Promise; + }; + expect(tool).toBeDefined(); + + await tool.execute({}); + + expect(mockConnect).toHaveBeenLastCalledWith( + { + type: 'http', + url: 'http://agent-123.env-sample.svc.cluster.local:13338/mcp', + }, + 22000 + ); + expect(mockCallTool).toHaveBeenCalledWith('workspace.read_file', {}, 22000); + }); + + it('fails the session tool setup when the sandbox gateway is unavailable', async () => { + mockConnect.mockImplementation(async (transport) => { + currentTransport = transport as Record; + if ( + currentTransport && + currentTransport.type === 'http' && + currentTransport.url === 'http://agent-123.env-sample.svc.cluster.local:13338/mcp' + ) { + throw new Error('sandbox unavailable'); + } + }); + + await expect( + buildToolSetForTest({ + session, + repoFullName: 'example-org/example-repo', + userIdentity, + approvalPolicy: {} as any, + workspaceToolDiscoveryTimeoutMs: 3000, + workspaceToolExecutionTimeoutMs: 15000, + }) + ).rejects.toBeInstanceOf(SessionWorkspaceGatewayUnavailableError); + + expect(mockLoggerWarn).toHaveBeenCalled(); + }); + + it('lets session tool rules override the family approval mode for sandbox tools', async () => { + mockModeForCapability.mockReturnValue('deny'); + + const tools = await buildToolSetForTest({ + session, + repoFullName: 'example-org/example-repo', + userIdentity, + approvalPolicy: {} as any, + toolRules: [ + { + toolKey: 'mcp__sandbox__workspace_read_file', + mode: 'allow', + }, + ], + workspaceToolDiscoveryTimeoutMs: 4500, + workspaceToolExecutionTimeoutMs: 22000, + }); + + expect(tools.mcp__sandbox__workspace_read_file).toEqual( + expect.objectContaining({ + needsApproval: false, + }) + ); + }); + + 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 buildToolSetForTest({ + 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(tools.mcp__lifecycle__get_codefresh_logs).toBeUndefined(); + expect(Object.keys(tools).some((key) => key.includes('__source_'))).toBe(false); + expect(tools.mcp__lifecycle__workspace_provision).toBeUndefined(); + }); + + it('registers Lifecycle diagnostic read tools for build-context chat sessions', async () => { + mockResolveServers.mockResolvedValue([]); + const onToolStarted = jest.fn(); + const onToolFinished = jest.fn(); + + const tools = await buildToolSetForTest({ + session: { + uuid: 'session-build-context', + sessionKind: 'chat', + buildUuid: 'sample-build-1', + workspaceStatus: 'none', + status: 'active', + podName: null, + namespace: null, + } as any, + repoFullName: 'example-org/example-repo', + userIdentity, + approvalPolicy: {} as any, + workspaceToolDiscoveryTimeoutMs: 4500, + workspaceToolExecutionTimeoutMs: 22000, + hooks: { + onToolStarted, + onToolFinished, + }, + }); + + expect(tools.mcp__lifecycle__get_codefresh_logs).toEqual(expect.objectContaining({ needsApproval: false })); + expect(tools.mcp__lifecycle__get_k8s_resources).toBeDefined(); + expect(tools.mcp__lifecycle__get_pod_logs).toBeDefined(); + expect(tools.mcp__lifecycle__get_lifecycle_logs).toBeDefined(); + expect(tools.mcp__lifecycle__query_database).toBeDefined(); + expect(tools.mcp__lifecycle__get_file).toBeDefined(); + expect(tools.mcp__lifecycle__list_directory).toBeDefined(); + expect(tools.mcp__lifecycle__get_issue_comment).toBeDefined(); + expect(typeof (tools.mcp__lifecycle__update_file as { needsApproval?: unknown }).needsApproval).toBe('function'); + expect(tools.mcp__lifecycle__update_pr_labels).toEqual(expect.objectContaining({ needsApproval: true })); + expect(tools.mcp__lifecycle__patch_k8s_resource).toEqual(expect.objectContaining({ needsApproval: true })); + + const tool = tools.mcp__lifecycle__get_codefresh_logs as { + execute: (input: Record, context?: { toolCallId?: string }) => Promise; + }; + await tool.execute({ pipeline_id: 'pipeline-1' }, { toolCallId: 'tool-codefresh' }); + + expect(mockDiagnosticToolExecute).toHaveBeenCalledWith( + 'get_codefresh_logs', + { pipeline_id: 'pipeline-1' }, + undefined + ); + expect(onToolStarted).toHaveBeenCalledWith( + expect.objectContaining({ + source: 'mcp', + serverSlug: 'lifecycle', + toolName: 'get_codefresh_logs', + toolCallId: 'tool-codefresh', + capabilityKey: 'read', + }) + ); + expect(onToolFinished).toHaveBeenCalledWith( + expect.objectContaining({ + toolName: 'get_codefresh_logs', + status: 'completed', + }) + ); + }); + + it('does not register Lifecycle diagnostic read tools for generic no-build chats', async () => { + mockResolveServers.mockResolvedValue([]); + + const tools = await buildToolSetForTest({ + 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, + }); + + expect(tools.mcp__lifecycle__get_codefresh_logs).toBeUndefined(); + expect(tools.mcp__lifecycle__query_database).toBeUndefined(); + expect(tools.mcp__lifecycle__get_file).toBeUndefined(); + expect(tools.mcp__lifecycle__update_file).toBeUndefined(); + expect(tools.mcp__lifecycle__patch_k8s_resource).toBeUndefined(); + }); + + it('lets tool rules deny individual Lifecycle diagnostic read tools', async () => { + mockResolveServers.mockResolvedValue([]); + + const tools = await buildToolSetForTest({ + session: { + uuid: 'session-build-context', + sessionKind: 'chat', + buildUuid: 'sample-build-1', + workspaceStatus: 'none', + status: 'active', + podName: null, + namespace: null, + } as any, + repoFullName: 'example-org/example-repo', + userIdentity, + approvalPolicy: {} as any, + toolRules: [ + { + toolKey: 'mcp__lifecycle__query_database', + mode: 'deny', + }, + ], + workspaceToolDiscoveryTimeoutMs: 4500, + workspaceToolExecutionTimeoutMs: 22000, + }); + + expect(tools.mcp__lifecycle__query_database).toBeUndefined(); + expect(tools.mcp__lifecycle__get_codefresh_logs).toBeDefined(); + }); + + it('omits Lifecycle diagnostic tools whose resolved catalog capability is unavailable', async () => { + mockResolveServers.mockResolvedValue([]); + + const tools = await buildToolSetForTest({ + session: { + uuid: 'session-build-context', + sessionKind: 'chat', + buildUuid: 'sample-build-1', + workspaceStatus: 'none', + status: 'active', + podName: null, + namespace: null, + } as any, + repoFullName: 'example-org/example-repo', + userIdentity, + approvalPolicy: {} as any, + resolvedCapabilityAccess: [ + { + capabilityId: 'diagnostics_codefresh', + effectiveAvailability: 'system_only', + allowed: true, + approvalMode: 'allow', + }, + { + capabilityId: 'diagnostics_database', + effectiveAvailability: 'system_only', + allowed: false, + reason: 'disabled', + }, + ], + workspaceToolDiscoveryTimeoutMs: 4500, + workspaceToolExecutionTimeoutMs: 22000, + }); + + expect(tools.mcp__lifecycle__get_codefresh_logs).toBeDefined(); + expect(tools.mcp__lifecycle__query_database).toBeUndefined(); + }); + + it('requires approval for Lifecycle diagnostic fix tools even when capability policy allows them', async () => { + mockResolveServers.mockResolvedValue([]); + mockModeForCapability.mockReturnValue('allow'); + const onToolStarted = jest.fn(); + const onToolFinished = jest.fn(); + const onFileChange = jest.fn(); + + const tools = await buildToolSetForTest({ + session: { + uuid: 'session-build-context', + sessionKind: 'chat', + buildUuid: 'sample-build-1', + workspaceStatus: 'none', + status: 'active', + podName: null, + namespace: null, + } as any, + repoFullName: 'example-org/example-repo', + userIdentity, + approvalPolicy: {} as any, + workspaceToolDiscoveryTimeoutMs: 4500, + workspaceToolExecutionTimeoutMs: 22000, + hooks: { + onToolStarted, + onToolFinished, + onFileChange, + }, + }); + + expect(typeof (tools.mcp__lifecycle__update_file as { needsApproval?: unknown }).needsApproval).toBe('function'); + expect(tools.mcp__lifecycle__update_pr_labels).toEqual(expect.objectContaining({ needsApproval: true })); + expect(tools.mcp__lifecycle__patch_k8s_resource).toEqual(expect.objectContaining({ needsApproval: true })); + + const updateFileTool = tools.mcp__lifecycle__update_file as unknown as { + onInputAvailable: (input: { input: Record; toolCallId?: string }) => Promise; + execute: (input: Record, context?: { toolCallId?: string }) => Promise; + }; + + await updateFileTool.onInputAvailable({ + input: { + repository_owner: 'example-org', + repository_name: 'example-repo', + branch: 'feature/sample', + file_path: './lifecycle.yaml', + new_content: 'services:\\n sample-service:\\n branch: feature/sample', + }, + toolCallId: 'tool-update-file', + }); + await updateFileTool.execute( + { + repository_owner: 'example-org', + repository_name: 'example-repo', + branch: 'feature/sample', + file_path: './lifecycle.yaml', + new_content: 'services:\\n sample-service:\\n branch: feature/sample', + }, + { toolCallId: 'tool-update-file' } + ); + + expect(onFileChange).toHaveBeenCalledWith( + expect.objectContaining({ + toolCallId: 'tool-update-file', + sourceTool: 'update_file', + displayPath: 'lifecycle.yaml', + stage: 'awaiting-approval', + additions: 1, + deletions: 1, + beforeTextPreview: 'services:\n sample-service:\n branch: main', + afterTextPreview: 'services:\n sample-service:\n branch: feature/sample', + unifiedDiff: expect.stringContaining('- branch: main'), + }) + ); + expect(onFileChange).toHaveBeenCalledWith( + expect.objectContaining({ + unifiedDiff: expect.stringContaining('+ branch: feature/sample'), + }) + ); + const proposedFileChange = onFileChange.mock.calls[0]?.[0] as { unifiedDiff?: string | null }; + expect(proposedFileChange.unifiedDiff?.match(/^\+\+\+ b\/lifecycle\.yaml$/gm)).toHaveLength(1); + expect(mockDiagnosticToolExecute).toHaveBeenCalledWith( + 'update_file', + { + repository_owner: 'example-org', + repository_name: 'example-repo', + branch: 'feature/sample', + file_path: './lifecycle.yaml', + new_content: 'services:\\n sample-service:\\n branch: feature/sample', + }, + undefined + ); + expect(onToolStarted).toHaveBeenCalledWith( + expect.objectContaining({ + toolName: 'update_file', + toolCallId: 'tool-update-file', + capabilityKey: 'git_write', + }) + ); + expect(onToolFinished).toHaveBeenCalledWith( + expect.objectContaining({ + toolName: 'update_file', + status: 'completed', + }) + ); + }); + + it('configures Lifecycle diagnostic GitHub writes with the session branch and allowed config files', async () => { + mockResolveServers.mockResolvedValue([]); + mockGetEffectiveAgentConfig.mockResolvedValue({ + allowedWritePatterns: ['sysops/dockerfiles/**'], + excludedFilePatterns: ['secrets/**'], + }); + mockParseYamlConfigFromBranch.mockResolvedValue({ + services: [ + { + name: 'sample-service', + github: { + docker: { + app: { + dockerfilePath: 'services/sample/Dockerfile', + }, + }, + }, + }, + ], + }); + + const tools = await buildToolSetForTest({ + session: { + uuid: 'session-build-context', + sessionKind: 'chat', + buildUuid: 'sample-build-1', + workspaceStatus: 'none', + status: 'active', + podName: null, + namespace: null, + workspaceRepos: [ + { + repo: 'example-org/example-repo', + repoUrl: 'https://github.com/example-org/example-repo.git', + branch: 'feature/sample', + revision: null, + mountPath: '/workspace', + primary: true, + }, + ], + } as any, + repoFullName: 'example-org/example-repo', + userIdentity, + approvalPolicy: {} as any, + workspaceToolDiscoveryTimeoutMs: 4500, + workspaceToolExecutionTimeoutMs: 22000, + }); + + expect(mockParseYamlConfigFromBranch).toHaveBeenCalledWith('example-org/example-repo', 'feature/sample'); + const fixToolClient = mockGithubClientInstances.at(-1); + expect(fixToolClient?.setAllowedBranch).toHaveBeenCalledWith('feature/sample'); + expect(fixToolClient?.setAllowedWritePatterns).toHaveBeenCalledWith( + expect.arrayContaining(['lifecycle.yaml', 'lifecycle.yml', 'sysops/dockerfiles/**']) ); - expect(mockListTools).toHaveBeenCalledWith(4500); + expect(fixToolClient?.setExcludedFilePatterns).toHaveBeenCalledWith(['secrets/**']); + expect(fixToolClient?.setReferencedFiles).toHaveBeenCalledWith(['services/sample/Dockerfile']); - const tool = tools.mcp__figma__get_design_context as { - execute: (input: Record) => Promise; + const updateFileTool = tools.mcp__lifecycle__update_file as unknown as { + needsApproval: (input: Record) => Promise; }; - expect(tool).toBeDefined(); + fixToolClient?.isFilePathAllowed.mockReturnValue(true); + fixToolClient?.validateBranch.mockReturnValue({ valid: true }); + await expect( + updateFileTool.needsApproval({ + branch: 'feature/sample', + file_path: 'services/sample/Dockerfile', + }) + ).resolves.toBe(true); - await tool.execute({}); + fixToolClient?.isFilePathAllowed.mockReturnValue(false); + await expect( + updateFileTool.needsApproval({ + branch: 'feature/sample', + file_path: 'secrets/token.txt', + }) + ).resolves.toBe(false); + }); - expect(mockConnect).toHaveBeenCalledWith( + it('lets tool rules deny individual Lifecycle diagnostic fix tools', async () => { + mockResolveServers.mockResolvedValue([]); + + const tools = await buildToolSetForTest({ + session: { + uuid: 'session-build-context', + sessionKind: 'chat', + buildUuid: 'sample-build-1', + workspaceStatus: 'none', + status: 'active', + podName: null, + namespace: null, + } as any, + repoFullName: 'example-org/example-repo', + userIdentity, + approvalPolicy: {} as any, + toolRules: [ + { + toolKey: 'mcp__lifecycle__update_file', + mode: 'deny', + }, + ], + workspaceToolDiscoveryTimeoutMs: 4500, + workspaceToolExecutionTimeoutMs: 22000, + }); + + expect(tools.mcp__lifecycle__update_file).toBeUndefined(); + expect(tools.mcp__lifecycle__update_pr_labels).toEqual(expect.objectContaining({ needsApproval: true })); + expect(tools.mcp__lifecycle__patch_k8s_resource).toEqual(expect.objectContaining({ needsApproval: true })); + }); + + it('redacts MCP default args from tool audit hooks while preserving runtime execution args', async () => { + mockResolveServers.mockResolvedValue([ { - type: 'http', - url: 'http://agent-123.env-sample.svc.cluster.local:13338/servers/figma/mcp', + slug: 'docs', + name: 'Docs', + transport: { + type: 'http', + url: 'https://mcp.example.test', + }, + timeout: 30000, + defaultArgs: { + token: 'shared-secret-token', + project: 'secret-project', + }, + env: {}, + discoveredTools: [ + { + name: 'lookup', + description: 'Lookup docs', + inputSchema: { + type: 'object', + properties: { + query: { type: 'string' }, + token: { type: 'string' }, + project: { type: 'string' }, + }, + }, + annotations: { + readOnlyHint: true, + }, + }, + ], }, - 30000 - ); - expect(mockCallTool).toHaveBeenCalledWith('get_design_context', {}, 30000); - }); + ]); + const onToolStarted = jest.fn(); + const onToolFinished = jest.fn(); - it('uses the configured workspace execution timeout for sandbox tools', async () => { - const tools = await AgentCapabilityService.buildToolSet({ + const tools = await buildToolSetForTest({ session, repoFullName: 'example-org/example-repo', userIdentity, approvalPolicy: {} as any, workspaceToolDiscoveryTimeoutMs: 4500, workspaceToolExecutionTimeoutMs: 22000, + hooks: { + onToolStarted, + onToolFinished, + }, }); - const tool = tools.mcp__sandbox__workspace_read_file as { - execute: (input: Record) => Promise; + const tool = tools.mcp__docs__lookup as { + execute: (input: Record, context?: { toolCallId?: string }) => Promise; }; - expect(tool).toBeDefined(); - - await tool.execute({}); + await tool.execute({ query: 'routing' }, { toolCallId: 'tool-defaults' }); - expect(mockConnect).toHaveBeenLastCalledWith( + expect(mockCallTool).toHaveBeenCalledWith( + 'lookup', { - type: 'http', - url: 'http://agent-123.env-sample.svc.cluster.local:13338/mcp', + query: 'routing', + token: 'shared-secret-token', + project: 'secret-project', }, - 22000 + 30000 ); - expect(mockCallTool).toHaveBeenCalledWith('workspace.read_file', {}, 22000); - }); - - it('fails the session tool setup when the sandbox gateway is unavailable', async () => { - mockConnect.mockImplementation(async (transport) => { - currentTransport = transport as Record; - if ( - currentTransport && - currentTransport.type === 'http' && - currentTransport.url === 'http://agent-123.env-sample.svc.cluster.local:13338/mcp' - ) { - throw new Error('sandbox unavailable'); - } - }); - - await expect( - AgentCapabilityService.buildToolSet({ - session, - repoFullName: 'example-org/example-repo', - userIdentity, - approvalPolicy: {} as any, - workspaceToolDiscoveryTimeoutMs: 3000, - workspaceToolExecutionTimeoutMs: 15000, + expect(onToolStarted).toHaveBeenCalledWith( + expect.objectContaining({ + args: { + query: 'routing', + token: '******', + project: '******', + }, }) - ).rejects.toBeInstanceOf(SessionWorkspaceGatewayUnavailableError); - - expect(mockLoggerWarn).toHaveBeenCalled(); + ); + expect(onToolFinished).toHaveBeenCalledWith( + expect.objectContaining({ + args: { + query: 'routing', + token: '******', + project: '******', + }, + }) + ); }); - it('lets session tool rules override the family approval mode for sandbox tools', async () => { - mockModeForCapability.mockReturnValue('deny'); + it('redacts transport and env secrets from MCP tool failure audit output', async () => { + mockResolveServers.mockResolvedValue([ + { + slug: 'docs', + name: 'Docs', + transport: { + type: 'http', + url: 'https://mcp.example.test?api_key=query/secret+value', + headers: { + Authorization: 'Bearer transport-secret', + }, + }, + timeout: 30000, + defaultArgs: { + token: 'shared-secret-token', + }, + env: { + SAMPLE_ENV_TOKEN: 'env-secret', + }, + discoveredTools: [ + { + name: 'lookup', + description: 'Lookup docs', + inputSchema: { + type: 'object', + properties: { + query: { type: 'string' }, + token: { type: 'string' }, + }, + }, + annotations: { + readOnlyHint: true, + }, + }, + ], + }, + ]); + mockCallTool.mockRejectedValueOnce( + new Error( + 'failed Authorization=Bearer transport-secret query=query/secret+value encoded=query%2Fsecret%2Bvalue env=env-secret token=shared-secret-token' + ) + ); + const onToolFinished = jest.fn(); - const tools = await AgentCapabilityService.buildToolSet({ + const tools = await buildToolSetForTest({ session, repoFullName: 'example-org/example-repo', userIdentity, approvalPolicy: {} as any, - toolRules: [ - { - toolKey: 'mcp__sandbox__workspace_read_file', - mode: 'allow', - }, - ], workspaceToolDiscoveryTimeoutMs: 4500, workspaceToolExecutionTimeoutMs: 22000, + hooks: { + onToolFinished, + }, }); - expect(tools.mcp__sandbox__workspace_read_file).toEqual( + const tool = tools.mcp__docs__lookup as { + execute: (input: Record, context?: { toolCallId?: string }) => Promise; + }; + await expect(tool.execute({ query: 'routing' }, { toolCallId: 'tool-failed' })).rejects.toThrow( + 'failed Authorization=****** query=****** encoded=****** env=****** token=******' + ); + + expect(onToolFinished).toHaveBeenCalledWith( expect.objectContaining({ - needsApproval: false, + status: 'failed', + result: { + error: 'failed Authorization=****** query=****** encoded=****** env=****** token=******', + }, }) ); }); - 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' - ); + it('redacts transport and env secrets from non-throwing MCP error results', async () => { mockResolveServers.mockResolvedValue([ { slug: 'docs', name: 'Docs', transport: { type: 'http', - url: 'https://mcp.example.test', + url: 'https://mcp.example.test?api_key=query/secret+value', + headers: { + Authorization: 'Bearer transport-secret', + }, }, timeout: 30000, - defaultArgs: {}, - env: {}, + defaultArgs: { + token: 'shared-secret-token', + }, + env: { + SAMPLE_ENV_TOKEN: 'env-secret', + }, discoveredTools: [ { - name: 'search_docs', - description: 'Search docs', + name: 'lookup', + description: 'Lookup docs', inputSchema: { type: 'object', - properties: {}, + properties: { + query: { type: 'string' }, + token: { type: 'string' }, + }, }, annotations: { readOnlyHint: true, @@ -363,48 +1343,83 @@ describe('AgentCapabilityService.buildToolSet', () => { ], }, ]); + mockCallTool.mockResolvedValueOnce({ + content: [ + { + type: 'text', + text: JSON.stringify({ + ok: false, + error: + 'failed Authorization=Bearer transport-secret query=query/secret+value encoded=query%2Fsecret%2Bvalue env=env-secret token=shared-secret-token', + fileChanges: [ + { + path: 'docs-output.txt', + kind: 'created', + additions: 1, + deletions: 0, + afterTextPreview: 'query=query/secret+value env=env-secret token=shared-secret-token', + summary: 'Created with Authorization=Bearer transport-secret', + }, + ], + }), + }, + ], + isError: true, + }); + const onToolFinished = jest.fn(); + const onFileChange = jest.fn(); - const tools = await AgentCapabilityService.buildToolSet({ - session: { - ...session, - sessionKind: 'chat', - podName: null, - namespace: null, - workspaceStatus: 'none', - } as any, - repoFullName: undefined, + const tools = await buildToolSetForTest({ + session, + repoFullName: 'example-org/example-repo', userIdentity, approvalPolicy: {} as any, workspaceToolDiscoveryTimeoutMs: 4500, workspaceToolExecutionTimeoutMs: 22000, + hooks: { + onToolFinished, + onFileChange, + }, }); - 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( + const tool = tools.mcp__docs__lookup as { + execute: (input: Record, context?: { toolCallId?: string }) => Promise; + }; + const result = await tool.execute({ query: 'routing' }, { toolCallId: 'tool-error-result' }); + const serializedResult = JSON.stringify(result); + + expect(serializedResult).toContain('Authorization=******'); + expect(serializedResult).toContain('query=******'); + expect(serializedResult).toContain('encoded=******'); + expect(serializedResult).toContain('env=******'); + expect(serializedResult).toContain('token=******'); + expect(serializedResult).not.toContain('Bearer transport-secret'); + expect(serializedResult).not.toContain('query/secret+value'); + expect(serializedResult).not.toContain('query%2Fsecret%2Bvalue'); + expect(serializedResult).not.toContain('env-secret'); + expect(serializedResult).not.toContain('shared-secret-token'); + expect(onToolFinished).toHaveBeenCalledWith( expect.objectContaining({ - needsApproval: false, + status: 'failed', + result, }) ); - expect(tools.mcp__lifecycle__publish_http).toEqual( + expect(onFileChange).toHaveBeenCalledWith( expect.objectContaining({ - needsApproval: true, + toolCallId: 'tool-error-result', + sourceTool: 'lookup', + stage: 'failed', + afterTextPreview: 'query=****** env=****** token=******', + summary: 'Created with Authorization=******', }) ); - 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({ + const tools = await buildToolSetForTest({ session: { ...session, sessionKind: 'chat', @@ -440,7 +1455,7 @@ describe('AgentCapabilityService.buildToolSet', () => { namespace: null, }); - const tools = await AgentCapabilityService.buildToolSet({ + const tools = await buildToolSetForTest({ session: { uuid: 'session-chat', sessionKind: 'chat', @@ -522,7 +1537,7 @@ describe('AgentCapabilityService.buildToolSet', () => { it('lets tool rules require approval for lazy chat workspace tools before runtime exists', async () => { mockResolveServers.mockResolvedValue([]); - const tools = await AgentCapabilityService.buildToolSet({ + const tools = await buildToolSetForTest({ session: { uuid: 'session-chat', sessionKind: 'chat', @@ -551,10 +1566,59 @@ describe('AgentCapabilityService.buildToolSet', () => { ); }); + it('omits lazy chat workspace tools whose resolved catalog capability is unavailable', async () => { + mockResolveServers.mockResolvedValue([]); + + const tools = await buildToolSetForTest({ + session: { + uuid: 'session-chat', + sessionKind: 'chat', + workspaceStatus: 'none', + status: 'active', + podName: null, + namespace: null, + } as any, + repoFullName: undefined, + userIdentity, + approvalPolicy: {} as any, + resolvedCapabilityAccess: [ + { + capabilityId: 'read_context', + effectiveAvailability: 'all_users', + allowed: true, + approvalMode: 'allow', + }, + { + capabilityId: 'workspace_files', + effectiveAvailability: 'all_users', + allowed: true, + approvalMode: 'allow', + }, + { + capabilityId: 'workspace_shell', + effectiveAvailability: 'disabled', + allowed: false, + reason: 'disabled', + }, + { + capabilityId: 'preview_publish', + effectiveAvailability: 'all_users', + allowed: true, + approvalMode: 'allow', + }, + ], + workspaceToolDiscoveryTimeoutMs: 4500, + workspaceToolExecutionTimeoutMs: 22000, + }); + + expect(tools.mcp__sandbox__workspace_write_file).toBeDefined(); + expect(tools.mcp__sandbox__workspace_exec_mutation).toBeUndefined(); + }); + it('runs GitHub CLI commands through the generic workspace mutation tool with request GitHub auth', async () => { mockResolveServers.mockResolvedValue([]); - const tools = await AgentCapabilityService.buildToolSet({ + const tools = await buildToolSetForTest({ session: { uuid: 'session-chat', sessionKind: 'chat', @@ -571,7 +1635,7 @@ describe('AgentCapabilityService.buildToolSet', () => { requestGitHubToken: 'sample-gh-token', }); - const tool = tools.mcp__sandbox__workspace_exec_mutation as { + const tool = tools.mcp__sandbox__workspace_exec_mutation as unknown as { execute: (input: Record) => Promise; }; mockFindSession.mockResolvedValueOnce({ @@ -633,7 +1697,7 @@ describe('AgentCapabilityService.buildToolSet', () => { }); const onFileChange = jest.fn(); - const tools = await AgentCapabilityService.buildToolSet({ + const tools = await buildToolSetForTest({ session: { uuid: 'session-chat', sessionKind: 'chat', @@ -723,7 +1787,7 @@ describe('AgentCapabilityService.buildToolSet', () => { }); const onFileChange = jest.fn(); - const tools = await AgentCapabilityService.buildToolSet({ + const tools = await buildToolSetForTest({ session, repoFullName: undefined, userIdentity, @@ -769,7 +1833,7 @@ describe('AgentCapabilityService.buildToolSet', () => { it('blocks unsafe broad process kill commands before they reach the workspace gateway', async () => { mockResolveServers.mockResolvedValue([]); - const tools = await AgentCapabilityService.buildToolSet({ + const tools = await buildToolSetForTest({ session: { uuid: 'session-chat', sessionKind: 'chat', @@ -785,7 +1849,7 @@ describe('AgentCapabilityService.buildToolSet', () => { workspaceToolExecutionTimeoutMs: 22000, }); - const tool = tools.mcp__sandbox__workspace_exec_mutation as { + const tool = tools.mcp__sandbox__workspace_exec_mutation as unknown as { execute: (input: Record) => Promise; }; diff --git a/src/server/services/agent/__tests__/ChatSessionService.test.ts b/src/server/services/agent/__tests__/ChatSessionService.test.ts new file mode 100644 index 00000000..8ca65be2 --- /dev/null +++ b/src/server/services/agent/__tests__/ChatSessionService.test.ts @@ -0,0 +1,324 @@ +/** + * 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. + */ + +const mockAgentSessionQuery = jest.fn(); +const mockAgentSessionTransaction = jest.fn(); +const mockAgentThreadQuery = jest.fn(); +const mockCreateSessionSource = jest.fn(); +const mockResolveSelection = jest.fn(); +const mockGetRequiredStoredApiKey = jest.fn(); +const mockValidateEntryChoices = jest.fn(); +const mockLoggerInfo = jest.fn(); + +jest.mock('server/models/AgentSession', () => ({ + __esModule: true, + default: { + query: (...args: unknown[]) => mockAgentSessionQuery(...args), + transaction: (...args: unknown[]) => mockAgentSessionTransaction(...args), + }, +})); + +jest.mock('server/models/AgentThread', () => ({ + __esModule: true, + default: { + query: (...args: unknown[]) => mockAgentThreadQuery(...args), + }, +})); + +jest.mock('../ProviderRegistry', () => ({ + __esModule: true, + default: { + resolveSelection: (...args: unknown[]) => mockResolveSelection(...args), + getRequiredStoredApiKey: (...args: unknown[]) => mockGetRequiredStoredApiKey(...args), + }, +})); + +jest.mock('../SourceService', () => ({ + __esModule: true, + default: { + createSessionSource: (...args: unknown[]) => mockCreateSessionSource(...args), + }, +})); + +jest.mock('../ThreadRuntimeControlsService', () => ({ + __esModule: true, + default: { + validateEntryChoices: (...args: unknown[]) => mockValidateEntryChoices(...args), + }, +})); + +jest.mock('server/lib/logger', () => ({ + getLogger: () => ({ + info: mockLoggerInfo, + }), +})); + +import { AgentChatStatus, AgentSessionKind, AgentWorkspaceStatus } from 'shared/constants'; +import AgentChatSessionService from '../ChatSessionService'; + +const TEST_TRX = { name: 'trx' }; + +function arrangePersistence() { + let insertedSession: Record | null = null; + let finalizedSession: Record | null = null; + + const sessionInsertAndFetch = jest.fn(async (payload) => { + insertedSession = { + id: 17, + uuid: payload.uuid, + defaultThreadId: null, + ...payload, + }; + return insertedSession; + }); + const sessionPatchAndFetchById = jest.fn(async (_id, patch) => { + finalizedSession = { + ...insertedSession, + ...patch, + }; + return finalizedSession; + }); + const threadInsertAndFetch = jest.fn(async (payload) => ({ + id: 23, + uuid: 'sample-thread', + ...payload, + })); + + mockAgentSessionQuery.mockImplementation((trx?: unknown) => { + if (trx) { + return { + insertAndFetch: sessionInsertAndFetch, + patchAndFetchById: sessionPatchAndFetchById, + }; + } + + return {}; + }); + mockAgentThreadQuery.mockImplementation((trx?: unknown) => { + if (trx) { + return { + insertAndFetch: threadInsertAndFetch, + }; + } + + return {}; + }); + mockAgentSessionTransaction.mockImplementation(async (callback) => callback(TEST_TRX)); + mockCreateSessionSource.mockResolvedValue({ id: 31 }); + + return { + sessionInsertAndFetch, + sessionPatchAndFetchById, + threadInsertAndFetch, + }; +} + +describe('AgentChatSessionService.createChatSession', () => { + beforeEach(() => { + jest.clearAllMocks(); + mockGetRequiredStoredApiKey.mockResolvedValue('sample-key'); + mockValidateEntryChoices.mockResolvedValue(null); + }); + + it('passes provider and model to provider resolution so duplicate model ids are disambiguated', async () => { + const persistence = arrangePersistence(); + mockResolveSelection.mockResolvedValue({ + provider: 'sample-provider', + modelId: 'sample-model', + }); + + await AgentChatSessionService.createChatSession({ + userId: 'sample-user', + provider: 'sample-provider', + model: 'sample-model', + }); + + expect(mockResolveSelection).toHaveBeenCalledWith({ + repoFullName: undefined, + requestedProvider: 'sample-provider', + requestedModelId: 'sample-model', + }); + expect(mockGetRequiredStoredApiKey).toHaveBeenCalledWith({ + provider: 'sample-provider', + userIdentity: { + userId: 'sample-user', + githubUsername: null, + }, + }); + expect(persistence.sessionInsertAndFetch).toHaveBeenCalledWith( + expect.objectContaining({ + defaultModel: 'sample-model', + model: 'sample-model', + defaultHarness: 'lifecycle_ai_sdk', + sessionKind: AgentSessionKind.CHAT, + chatStatus: AgentChatStatus.READY, + workspaceStatus: AgentWorkspaceStatus.NONE, + }) + ); + expect(mockCreateSessionSource).toHaveBeenCalledWith( + expect.objectContaining({ + defaultModel: 'sample-model', + model: 'sample-model', + }), + expect.objectContaining({ + defaultProvider: 'sample-provider', + }) + ); + }); + + it('rejects an invalid provider and model pair through provider registry resolution', async () => { + const invalidProviderModelPairError = new Error('Model sample-provider:sample-model is not enabled'); + mockResolveSelection.mockRejectedValue(invalidProviderModelPairError); + + await expect( + AgentChatSessionService.createChatSession({ + userId: 'sample-user', + provider: 'sample-provider', + model: 'sample-model', + }) + ).rejects.toBe(invalidProviderModelPairError); + + expect(mockResolveSelection).toHaveBeenCalledWith({ + repoFullName: undefined, + requestedProvider: 'sample-provider', + requestedModelId: 'sample-model', + }); + expect(mockAgentSessionTransaction).not.toHaveBeenCalled(); + }); + + it('validates and persists runtime choices to the created default thread before returning', async () => { + const persistence = arrangePersistence(); + mockResolveSelection.mockResolvedValue({ + provider: 'sample-provider', + modelId: 'sample-model', + }); + mockValidateEntryChoices.mockResolvedValue({ + selectedAgentMetadataPatch: { + selectedAgentDefinitionId: 'custom.sample-agent', + }, + runtimeControlChoices: { + version: 1, + toolChoiceIds: ['rtc_tool_choice'], + mcpChoiceIds: [], + }, + }); + + await AgentChatSessionService.createChatSession({ + userId: 'sample-user', + userIdentity: { + userId: 'sample-user', + githubUsername: 'sample-user', + } as any, + provider: 'sample-provider', + model: 'sample-model', + runtimeControlChoices: { + agentId: 'custom.sample-agent', + toolChoiceIds: ['rtc_tool_choice'], + mcpChoiceIds: [], + }, + }); + + expect(mockValidateEntryChoices).toHaveBeenCalledWith({ + userIdentity: expect.objectContaining({ + userId: 'sample-user', + githubUsername: 'sample-user', + }), + agentId: 'custom.sample-agent', + source: { adapter: 'blank_workspace', input: {} }, + defaults: { + provider: 'sample-provider', + model: 'sample-model', + }, + runtimeControlChoices: { + agentId: 'custom.sample-agent', + toolChoiceIds: ['rtc_tool_choice'], + mcpChoiceIds: [], + }, + }); + expect(persistence.threadInsertAndFetch).toHaveBeenCalledWith( + expect.objectContaining({ + metadata: { + sessionUuid: expect.any(String), + selectedAgentDefinitionId: 'custom.sample-agent', + runtimeControlChoices: { + version: 1, + toolChoiceIds: ['rtc_tool_choice'], + mcpChoiceIds: [], + }, + }, + }) + ); + }); + + it('leaves runtime choice metadata absent when bootstrap choices are omitted', async () => { + const persistence = arrangePersistence(); + mockResolveSelection.mockResolvedValue({ + provider: 'sample-provider', + modelId: 'sample-model', + }); + + await AgentChatSessionService.createChatSession({ + userId: 'sample-user', + provider: 'sample-provider', + model: 'sample-model', + }); + + expect(mockValidateEntryChoices).not.toHaveBeenCalled(); + expect(persistence.threadInsertAndFetch).toHaveBeenCalledWith( + expect.objectContaining({ + metadata: { + sessionUuid: expect.any(String), + }, + }) + ); + }); + + it('persists selected agent metadata without runtime choices for agent-only bootstrap input', async () => { + const persistence = arrangePersistence(); + mockResolveSelection.mockResolvedValue({ + provider: 'sample-provider', + modelId: 'sample-model', + }); + mockValidateEntryChoices.mockResolvedValue({ + selectedAgentMetadataPatch: { + selectedAgentDefinitionId: 'custom.sample-agent', + }, + runtimeControlChoices: null, + }); + + await AgentChatSessionService.createChatSession({ + userId: 'sample-user', + userIdentity: { + userId: 'sample-user', + githubUsername: 'sample-user', + } as any, + provider: 'sample-provider', + model: 'sample-model', + runtimeControlChoices: { + agentId: 'custom.sample-agent', + }, + }); + + expect(persistence.threadInsertAndFetch).toHaveBeenCalledWith( + expect.objectContaining({ + metadata: { + sessionUuid: expect.any(String), + selectedAgentDefinitionId: 'custom.sample-agent', + }, + }) + ); + }); +}); diff --git a/src/server/services/agent/__tests__/CustomAgentDefinitionService.test.ts b/src/server/services/agent/__tests__/CustomAgentDefinitionService.test.ts new file mode 100644 index 00000000..448efda3 --- /dev/null +++ b/src/server/services/agent/__tests__/CustomAgentDefinitionService.test.ts @@ -0,0 +1,595 @@ +/** + * 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. + */ + +const mockFindOne = jest.fn(); +const mockInsert = jest.fn(); +const mockOrderBy = jest.fn(); +const mockPatchAndFetchById = jest.fn(); +const mockWhere = jest.fn(); +const mockGetEffectiveConfig = jest.fn(); +const mockListAvailableModelsForUser = jest.fn(); + +jest.mock('server/models/AgentDefinition', () => ({ + __esModule: true, + default: { + query: jest.fn(() => ({ + findOne: (...args: unknown[]) => mockFindOne(...args), + insert: (...args: unknown[]) => mockInsert(...args), + patchAndFetchById: (...args: unknown[]) => mockPatchAndFetchById(...args), + where: (...args: unknown[]) => { + mockWhere(...args); + return { + orderBy: (...orderArgs: unknown[]) => mockOrderBy(...orderArgs), + }; + }, + })), + }, +})); + +jest.mock('server/services/agentRuntime/config/agentRuntimeConfig', () => ({ + __esModule: true, + default: { + getInstance: jest.fn(() => ({ + getEffectiveConfig: (...args: unknown[]) => mockGetEffectiveConfig(...args), + })), + }, +})); + +jest.mock('server/services/agent/ProviderRegistry', () => ({ + __esModule: true, + default: { + listAvailableModelsForUser: (...args: unknown[]) => mockListAvailableModelsForUser(...args), + }, +})); + +import { CustomAgentDefinitionService, CustomAgentDefinitionServiceError } from '../CustomAgentDefinitionService'; + +function buildRow(overrides: Record = {}) { + return { + id: 1, + definitionId: 'custom.sample-agent', + version: 1, + ownerKind: 'user', + ownerUserId: 'sample-user', + ownerOrganizationId: null, + name: 'Sample agent', + description: 'Helps with sample workflows', + instructionRefs: [], + instructionAddendum: 'Answer with concise steps.', + capabilityRefs: ['read_context'], + requiredCapabilityRefs: [], + optionalCapabilityRefs: ['read_context'], + resourcePolicy: { + sourceKinds: ['freeform_chat'], + workspaceRequired: false, + sandboxRequired: false, + }, + modelPreference: null, + status: 'active', + codeOwned: false, + readOnly: false, + updatedAt: '2026-05-01T12:00:00.000Z', + ...overrides, + }; +} + +describe('CustomAgentDefinitionService', () => { + const service = new CustomAgentDefinitionService(); + const userIdentity = { userId: 'sample-user', githubUsername: 'sample-user' }; + + beforeEach(() => { + jest.clearAllMocks(); + mockGetEffectiveConfig.mockResolvedValue({ capabilityPolicy: undefined }); + mockListAvailableModelsForUser.mockResolvedValue([ + { + provider: 'openai', + modelId: 'sample-model', + displayName: 'Sample model', + default: true, + maxTokens: 4096, + }, + ]); + }); + + it('listUserDefinitions returns only active rows for the owner newest first', async () => { + mockOrderBy.mockResolvedValue([ + buildRow({ + id: 2, + definitionId: 'custom.newest', + name: 'Newest', + updatedAt: '2026-05-01T13:00:00.000Z', + }), + buildRow({ + id: 1, + definitionId: 'custom.oldest', + name: 'Oldest', + updatedAt: '2026-05-01T12:00:00.000Z', + }), + ]); + + const definitions = await service.listUserDefinitions({ userId: 'sample-user' }); + + expect(mockWhere).toHaveBeenCalledWith({ + ownerKind: 'user', + ownerUserId: 'sample-user', + status: 'active', + }); + expect(mockOrderBy).toHaveBeenCalledWith('updatedAt', 'desc'); + expect(definitions.map((definition) => definition.id)).toEqual(['custom.newest', 'custom.oldest']); + }); + + it('getUserDefinition returns not found for another user, an archived row, or a system row', async () => { + mockFindOne.mockResolvedValue(null); + + await expect(service.getUserDefinition('custom.other-user', 'sample-user')).rejects.toMatchObject({ + code: 'not_found', + }); + await expect(service.getUserDefinition('custom.archived', 'sample-user')).rejects.toBeInstanceOf( + CustomAgentDefinitionServiceError + ); + await expect(service.getUserDefinition('system.freeform', 'sample-user')).rejects.toMatchObject({ + code: 'not_found', + }); + + expect(mockFindOne).toHaveBeenCalledWith({ + definitionId: 'custom.other-user', + ownerKind: 'user', + ownerUserId: 'sample-user', + status: 'active', + }); + }); + + it('create and update trim fields, dedupe capabilities, increment version, and ignore codeOwned/readOnly input', async () => { + mockInsert.mockImplementation(async (row) => buildRow({ id: 3, ...row })); + mockFindOne.mockResolvedValue(buildRow({ id: 3, version: 2 })); + mockPatchAndFetchById.mockImplementation(async (_id, patch) => buildRow({ id: 3, version: 3, ...patch })); + + const created = await service.createUserDefinition(userIdentity, { + name: ' Release helper ', + description: ' Summarizes release notes. ', + instructionAddendum: ' Keep the response brief. ', + capabilityRefs: ['read_context', 'read_context'], + resourceBehavior: 'chat_only', + codeOwned: true, + readOnly: true, + } as any); + + expect(created).toEqual( + expect.objectContaining({ + owner: { kind: 'user', userId: 'sample-user', organizationId: null }, + name: 'Release helper', + description: 'Summarizes release notes.', + instructionAddendum: 'Keep the response brief.', + capabilityRefs: ['read_context'], + optionalCapabilityRefs: ['read_context'], + codeOwned: false, + readOnly: false, + }) + ); + expect(mockInsert).toHaveBeenCalledWith( + expect.objectContaining({ + definitionId: expect.stringMatching(/^custom\./), + ownerKind: 'user', + ownerUserId: 'sample-user', + ownerOrganizationId: null, + codeOwned: false, + readOnly: false, + }) + ); + + await service.updateUserDefinition('custom.sample-agent', userIdentity, { + name: ' Updated helper ', + description: ' ', + instructionAddendum: ' Prefer bullets. ', + capabilityRefs: ['read_context'], + resourceBehavior: 'current_workspace_when_available', + codeOwned: true, + readOnly: true, + } as any); + + expect(mockPatchAndFetchById).toHaveBeenCalledWith( + 3, + expect.objectContaining({ + version: 3, + name: 'Updated helper', + description: null, + instructionAddendum: 'Prefer bullets.', + codeOwned: false, + readOnly: false, + }) + ); + }); + + it('archiveUserDefinition changes status to archived and keeps the row', async () => { + mockFindOne.mockResolvedValue(buildRow({ id: 4, definitionId: 'custom.to-archive' })); + mockPatchAndFetchById.mockResolvedValue(buildRow({ id: 4, definitionId: 'custom.to-archive', status: 'archived' })); + + const archived = await service.archiveUserDefinition('custom.to-archive', 'sample-user'); + + expect(mockPatchAndFetchById).toHaveBeenCalledWith(4, { status: 'archived' }); + expect(archived.status).toBe('archived'); + }); + + it('rejects unknown capability ids before persistence', async () => { + await expect( + service.createUserDefinition(userIdentity, { + name: 'Sample agent', + instructionAddendum: 'Answer briefly.', + capabilityRefs: ['read_context', 'sample_unknown_capability' as any], + resourceBehavior: 'chat_only', + }) + ).rejects.toMatchObject({ + code: 'unknown_capability', + message: 'Some selected capabilities are no longer available. Review the list and save again.', + }); + + expect(mockInsert).not.toHaveBeenCalled(); + }); + + it.each([ + ['admin_only', 'external_mcp_write', undefined], + ['system_only', 'approval_controls', undefined], + ['disabled', 'read_context', { availability: { read_context: 'disabled' } }], + ])('rejects %s capabilities for user-owned definitions', async (reason, capabilityId, capabilityPolicy) => { + mockGetEffectiveConfig.mockResolvedValueOnce({ capabilityPolicy }); + + await expect( + service.createUserDefinition(userIdentity, { + name: 'Sample agent', + instructionAddendum: 'Answer briefly.', + capabilityRefs: [capabilityId as any], + resourceBehavior: 'chat_only', + }) + ).rejects.toMatchObject({ + code: reason, + message: 'Some selected capabilities are no longer available. Review the list and save again.', + }); + + expect(mockInsert).not.toHaveBeenCalled(); + }); + + it.each([ + ['unknown_capability', 'sample_unknown_capability', undefined], + ['admin_only', 'external_mcp_write', undefined], + ['system_only', 'approval_controls', undefined], + ['disabled', 'read_context', { availability: { read_context: 'disabled' } }], + ['source_incompatible', 'workspace_shell', undefined], + ])( + 'rejects update payloads with %s capabilities without persistence', + async (reason, capabilityId, capabilityPolicy) => { + mockFindOne.mockResolvedValue( + buildRow({ + id: 5, + version: 7, + capabilityRefs: ['external_mcp_write'], + optionalCapabilityRefs: ['external_mcp_write'], + }) + ); + mockGetEffectiveConfig.mockResolvedValueOnce({ capabilityPolicy }); + + await expect( + service.updateUserDefinition('custom.sample-agent', userIdentity, { + name: 'Sample agent', + instructionAddendum: 'Answer briefly.', + capabilityRefs: [capabilityId as any], + resourceBehavior: 'chat_only', + }) + ).rejects.toMatchObject({ + code: reason, + message: 'Some selected capabilities are no longer available. Review the list and save again.', + }); + + expect(mockPatchAndFetchById).not.toHaveBeenCalled(); + } + ); + + it('replaces stale stored restricted capability selections during allowed updates', async () => { + mockFindOne.mockResolvedValue( + buildRow({ + id: 6, + version: 2, + capabilityRefs: ['external_mcp_write'], + optionalCapabilityRefs: ['external_mcp_write'], + }) + ); + mockPatchAndFetchById.mockImplementation(async (_id, patch) => buildRow({ id: 6, version: 3, ...patch })); + + await service.updateUserDefinition('custom.sample-agent', userIdentity, { + name: 'Sample agent', + instructionAddendum: 'Answer briefly.', + capabilityRefs: ['read_context'], + resourceBehavior: 'chat_only', + }); + + expect(mockPatchAndFetchById).toHaveBeenCalledWith( + 6, + expect.objectContaining({ + capabilityRefs: ['read_context'], + optionalCapabilityRefs: ['read_context'], + }) + ); + expect(JSON.stringify(mockPatchAndFetchById.mock.calls[0][1])).not.toContain('external_mcp_write'); + }); + + it('rejects source-incompatible required capabilities for chat_only custom agents', async () => { + await expect( + service.createUserDefinition(userIdentity, { + name: 'Sample agent', + instructionAddendum: 'Answer briefly.', + capabilityRefs: ['workspace_shell'], + resourceBehavior: 'chat_only', + }) + ).rejects.toMatchObject({ + code: 'source_incompatible', + message: 'Some selected capabilities are no longer available. Review the list and save again.', + }); + }); + + it.each([ + ['disabled', { mode: 'disabled' }], + ['admins_only', { mode: 'admins_only' }], + ['allowlist', { mode: 'allowlist', allowedUserIds: ['other-user'] }], + ])( + 'rejects create when custom-agent creation policy is %s for the caller', + async (_label, customAgentCreationPolicy) => { + mockGetEffectiveConfig.mockResolvedValueOnce({ customAgentCreationPolicy }); + + await expect( + service.createUserDefinition(userIdentity, { + name: 'Sample agent', + instructionAddendum: 'Answer briefly.', + capabilityRefs: ['read_context'], + resourceBehavior: 'chat_only', + }) + ).rejects.toMatchObject({ + code: 'creation_unavailable', + message: 'Custom agent creation is not available. Ask an admin for access.', + }); + + expect(mockInsert).not.toHaveBeenCalled(); + } + ); + + it('allows admin-role and allowlisted creators when creation policy is restricted', async () => { + mockGetEffectiveConfig.mockResolvedValueOnce({ + customAgentCreationPolicy: { mode: 'admins_only' }, + }); + mockInsert.mockImplementationOnce(async (row) => buildRow({ id: 7, ...row })); + + await expect( + service.createUserDefinition( + { ...userIdentity, roles: ['admin'] }, + { + name: 'Sample agent', + instructionAddendum: 'Answer briefly.', + capabilityRefs: ['read_context'], + resourceBehavior: 'chat_only', + } + ) + ).resolves.toMatchObject({ id: expect.stringMatching(/^custom\./) }); + + mockGetEffectiveConfig.mockResolvedValueOnce({ + customAgentCreationPolicy: { + mode: 'allowlist', + allowedGithubUsernames: ['SAMPLE-USER'], + }, + }); + mockInsert.mockImplementationOnce(async (row) => buildRow({ id: 8, ...row })); + + await expect( + service.createUserDefinition(userIdentity, { + name: 'Allowlisted agent', + instructionAddendum: 'Answer briefly.', + capabilityRefs: ['read_context'], + resourceBehavior: 'chat_only', + }) + ).resolves.toMatchObject({ name: 'Allowlisted agent' }); + }); + + it('reports current-user custom-agent creation status from policy', async () => { + await expect( + service.getUserDefinitionCreationStatus({ userIdentity: { ...userIdentity, roles: [] } as any }) + ).resolves.toEqual({ + canCreate: true, + creationUnavailableReason: null, + }); + + mockGetEffectiveConfig.mockResolvedValueOnce({ + customAgentCreationPolicy: { mode: 'disabled' }, + }); + + await expect( + service.getUserDefinitionCreationStatus({ userIdentity: { ...userIdentity, roles: [] } as any }) + ).resolves.toEqual({ + canCreate: false, + creationUnavailableReason: 'creation_disabled', + }); + + mockGetEffectiveConfig.mockResolvedValueOnce({ + customAgentCreationPolicy: { mode: 'allowlist', allowedUserIds: ['other-user'] }, + }); + + await expect( + service.getUserDefinitionCreationStatus({ userIdentity: { ...userIdentity, roles: [] } as any }) + ).resolves.toEqual({ + canCreate: false, + creationUnavailableReason: 'creation_restricted', + }); + }); + + it('rejects update when custom-agent creation policy no longer allows the caller', async () => { + mockFindOne.mockResolvedValue(buildRow({ id: 9, version: 3 })); + mockGetEffectiveConfig.mockResolvedValueOnce({ + customAgentCreationPolicy: { mode: 'allowlist', allowedUserIds: ['other-user'] }, + }); + + await expect( + service.updateUserDefinition('custom.sample-agent', userIdentity, { + name: 'Sample agent', + instructionAddendum: 'Answer briefly.', + capabilityRefs: ['read_context'], + resourceBehavior: 'chat_only', + }) + ).rejects.toMatchObject({ + code: 'creation_unavailable', + }); + + expect(mockPatchAndFetchById).not.toHaveBeenCalled(); + }); + + it('hides and rejects creator-reserved capabilities separately from runtime availability', async () => { + mockGetEffectiveConfig.mockResolvedValueOnce({ + capabilityPolicy: { + availability: { + external_mcp_write: 'all_users', + }, + }, + customAgentCreationPolicy: { + capabilityAvailability: { + external_mcp_write: 'reserved', + }, + }, + }); + + const capabilities = await service.listUserSelectableCapabilities({ + userIdentity: { userId: 'sample-user', githubUsername: 'sample-user', roles: [] } as any, + resourceBehavior: 'chat_only', + }); + + expect(capabilities.find((capability) => capability.capabilityId === 'external_mcp_write')).toBeUndefined(); + + mockGetEffectiveConfig.mockResolvedValueOnce({ + capabilityPolicy: { + availability: { + external_mcp_write: 'all_users', + }, + }, + customAgentCreationPolicy: { + capabilityAvailability: { + external_mcp_write: 'reserved', + }, + }, + }); + + await expect( + service.createUserDefinition(userIdentity, { + name: 'Sample agent', + instructionAddendum: 'Answer briefly.', + capabilityRefs: ['external_mcp_write'], + resourceBehavior: 'chat_only', + }) + ).rejects.toMatchObject({ + code: 'creator_capability_reserved', + message: 'Some selected capabilities are no longer available. Review the list and save again.', + }); + + expect(mockInsert).not.toHaveBeenCalled(); + }); + + it('rejects unavailable modelPreference selections with model_unavailable', async () => { + await expect( + service.createUserDefinition(userIdentity, { + name: 'Sample agent', + instructionAddendum: 'Answer briefly.', + capabilityRefs: ['read_context'], + modelPreference: { provider: 'internal-provider', model: 'internal-model' }, + resourceBehavior: 'chat_only', + }) + ).rejects.toMatchObject({ + code: 'model_unavailable', + message: 'Selected model is no longer available. Choose another model and save again.', + }); + expect(mockListAvailableModelsForUser).toHaveBeenCalledWith({ userIdentity }); + expect(mockInsert).not.toHaveBeenCalled(); + }); + + it('sanitizes validation errors without raw tool keys, toolKey fields, serverSlug values, or internal path details', async () => { + try { + await service.createUserDefinition(userIdentity, { + name: 'Sample agent', + instructionAddendum: 'Answer briefly.', + capabilityRefs: ['external_mcp_write'], + resourceBehavior: 'chat_only', + }); + throw new Error('Expected create to fail'); + } catch (error) { + expect(error).toBeInstanceOf(CustomAgentDefinitionServiceError); + expect((error as Error).message).toBe( + 'Some selected capabilities are no longer available. Review the list and save again.' + ); + expect((error as Error).message).not.toContain('external_mcp_write'); + expect((error as Error).message).not.toContain('workspace.exec'); + expect((error as Error).message).not.toContain('toolKey'); + expect((error as Error).message).not.toContain('serverSlug'); + expect((error as Error).message).not.toContain('/internal/path'); + } + }); + + it('listUserSelectableCapabilities returns only user-visible chat capabilities', async () => { + const capabilities = await service.listUserSelectableCapabilities({ + userIdentity: { userId: 'sample-user' } as any, + resourceBehavior: 'chat_only', + }); + + expect(capabilities.map((capability) => capability.capabilityId)).toEqual(['read_context', 'external_mcp_read']); + expect(capabilities.find((capability) => capability.capabilityId === 'external_mcp_write')).toBeUndefined(); + expect(capabilities.find((capability) => capability.capabilityId === 'approval_controls')).toBeUndefined(); + expect(capabilities.find((capability) => capability.capabilityId === 'workspace_shell')).toBeUndefined(); + expect(JSON.stringify(capabilities)).not.toContain('workspace.exec'); + expect(JSON.stringify(capabilities)).not.toContain('toolKey'); + expect(JSON.stringify(capabilities)).not.toContain('serverSlug'); + }); + + it('listUserSelectableCapabilities includes workspace flags and hides disabled capabilities', async () => { + mockGetEffectiveConfig.mockResolvedValueOnce({ + capabilityPolicy: { + availability: { + read_context: 'disabled', + }, + }, + }); + + const capabilities = await service.listUserSelectableCapabilities({ + userIdentity: { userId: 'sample-user' } as any, + resourceBehavior: 'current_workspace_when_available', + }); + + expect(capabilities.find((capability) => capability.capabilityId === 'read_context')).toBeUndefined(); + expect(capabilities).toEqual( + expect.arrayContaining([ + expect.objectContaining({ + capabilityId: 'workspace_shell', + requiresWorkspace: true, + toolCount: 1, + resourceCount: 1, + }), + ]) + ); + }); + + it('listUserSelectableCapabilities returns empty inventory when creation is disabled for the caller', async () => { + mockGetEffectiveConfig.mockResolvedValueOnce({ + customAgentCreationPolicy: { mode: 'disabled' }, + }); + + const capabilities = await service.listUserSelectableCapabilities({ + userIdentity: { userId: 'sample-user', githubUsername: 'sample-user', roles: [] } as any, + resourceBehavior: 'current_workspace_when_available', + }); + + expect(capabilities).toEqual([]); + }); +}); diff --git a/src/server/services/agent/__tests__/FirstPartyAgentDefinitions.integration.test.ts b/src/server/services/agent/__tests__/FirstPartyAgentDefinitions.integration.test.ts new file mode 100644 index 00000000..dd47998d --- /dev/null +++ b/src/server/services/agent/__tests__/FirstPartyAgentDefinitions.integration.test.ts @@ -0,0 +1,483 @@ +/** + * 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. + */ + +const mockDefinitionUpsert = jest.fn(); +const mockDefinitionFindOne = jest.fn(); +const mockDefinitionOrderBy = jest.fn(); +const mockDefinitionWhere = jest.fn(); +const mockDefinitionWhereIn = jest.fn(); +const mockDefinitionQuery = jest.fn(() => ({ + findOne: (...args: unknown[]) => mockDefinitionFindOne(...args), + whereIn: (...args: unknown[]) => { + mockDefinitionWhereIn(...args); + return { + where: (...whereArgs: unknown[]) => { + mockDefinitionWhere(...whereArgs); + return { + orderBy: (...orderArgs: unknown[]) => mockDefinitionOrderBy(...orderArgs), + }; + }, + }; + }, +})); + +const mockResolveSessionContext = jest.fn(); +const mockResolveSelection = jest.fn(); +const mockThreadTransaction = jest.fn(); +const mockThreadQuery = jest.fn(); +const mockPatchAndFetchById = jest.fn(); +const mockRunQuery = jest.fn(); +const mockMessageQuery = jest.fn(); +const mockInsertAndFetch = jest.fn(); +const mockGetSessionSource = jest.fn(); +const mockGetOwnedThreadWithSession = jest.fn(); +const mockWarn = jest.fn(); + +let mockDefinitionRows: any[] = []; +let mockActiveRunResults: any[] = []; + +jest.mock('server/models/AgentDefinition', () => ({ + __esModule: true, + default: { + upsert: (...args: unknown[]) => mockDefinitionUpsert(...args), + query: (...args: unknown[]) => mockDefinitionQuery.apply(null, args), + }, +})); + +jest.mock('server/models/AgentThread', () => ({ + __esModule: true, + default: { + query: (...args: unknown[]) => mockThreadQuery(...args), + transaction: (...args: unknown[]) => mockThreadTransaction(...args), + }, +})); + +jest.mock('server/models/AgentRun', () => ({ + __esModule: true, + default: { + query: (...args: unknown[]) => mockRunQuery(...args), + }, +})); + +jest.mock('server/models/AgentMessage', () => ({ + __esModule: true, + default: { + query: (...args: unknown[]) => mockMessageQuery(...args), + }, +})); + +jest.mock('server/lib/logger', () => ({ + getLogger: jest.fn(() => ({ + warn: mockWarn, + })), +})); + +jest.mock('../CapabilityService', () => ({ + __esModule: true, + default: { + resolveSessionContext: (...args: unknown[]) => mockResolveSessionContext(...args), + }, +})); + +jest.mock('../ProviderRegistry', () => ({ + __esModule: true, + default: { + resolveSelection: (...args: unknown[]) => mockResolveSelection(...args), + }, +})); + +jest.mock('../SourceService', () => ({ + __esModule: true, + default: { + getSessionSource: (...args: unknown[]) => mockGetSessionSource(...args), + }, +})); + +jest.mock('../ThreadService', () => { + const actual = jest.requireActual('../ThreadService'); + return { + __esModule: true, + ...actual, + default: { + getOwnedThreadWithSession: (...args: unknown[]) => mockGetOwnedThreadWithSession(...args), + getSelectedAgentDefinitionId: actual.getSelectedAgentDefinitionId, + getRuntimeControlChoices: actual.getRuntimeControlChoices, + serializeThread: jest.fn(), + }, + }; +}); + +jest.mock('uuid', () => ({ + v4: jest.fn(() => '11111111-1111-4111-8111-111111111111'), +})); + +import { AgentSessionKind } from 'shared/constants'; +import { + ensureSystemAgentDefinitionsSeeded, + getSystemAgentDefinition, + serializeAgentDefinitionSummary, +} from '../AgentDefinitionRegistry'; +import AgentRunPlanResolver, { AgentRunPlanAgentUnavailableError } from '../RunPlanResolver'; +import AgentPolicyService from '../PolicyService'; +import { serializeRunPlanSummary } from '../runPlanSummary'; +import { SYSTEM_AGENT_DEFINITIONS, SYSTEM_AGENT_DEFINITION_IDS } from '../systemAgentDefinitions'; +import type { AgentDefinitionContract } from '../agentDefinitionTypes'; +import type { AgentCapabilityCatalogId } from '../capabilityCatalog'; + +const userIdentity = { + userId: 'sample-user', + githubUsername: 'sample-user', + preferredUsername: 'sample-user', + email: 'sample-user@example.com', + firstName: 'Sample', + lastName: 'User', + displayName: 'Sample User', + gitUserName: 'Sample User', + gitUserEmail: 'sample-user@example.com', + roles: [], +}; + +function buildDefinitionRow(definition: AgentDefinitionContract) { + return { + definitionId: definition.id, + version: definition.version, + ownerKind: definition.owner.kind, + ownerUserId: definition.owner.userId || null, + ownerOrganizationId: definition.owner.organizationId || null, + name: definition.name, + description: definition.description || null, + instructionRefs: definition.instructionRefs, + instructionAddendum: definition.instructionAddendum || null, + capabilityRefs: definition.capabilityRefs, + requiredCapabilityRefs: definition.requiredCapabilityRefs || definition.capabilityRefs, + optionalCapabilityRefs: definition.optionalCapabilityRefs || [], + resourcePolicy: definition.resourcePolicy, + modelPreference: definition.modelPreference || null, + status: definition.status, + codeOwned: Boolean(definition.codeOwned), + readOnly: Boolean(definition.readOnly), + }; +} + +function buildThread(overrides: Record = {}) { + return { + id: 7, + uuid: 'sample-thread', + sessionId: 17, + metadata: {}, + ...overrides, + } as any; +} + +function buildSession(overrides: Record = {}) { + return { + id: 17, + uuid: 'sample-session', + userId: 'sample-user', + sessionKind: AgentSessionKind.CHAT, + defaultHarness: 'lifecycle_ai_sdk', + defaultModel: 'gpt-5.4', + buildUuid: null, + namespace: 'sample-namespace', + workspaceRepos: [ + { + repo: 'example-org/example-repo', + branch: 'main', + primary: true, + }, + ], + selectedServices: [ + { + name: 'sample-service', + repo: 'example-org/example-repo', + branch: 'main', + }, + ], + ...overrides, + } as any; +} + +function buildSource(overrides: Record = {}) { + return { + id: 3, + uuid: 'sample-source', + adapter: 'blank_workspace', + status: 'ready', + input: {}, + preparedSource: {}, + sandboxRequirements: {}, + preparedAt: null, + ...overrides, + } as any; +} + +function mockActiveRuns(...results: any[]) { + mockActiveRunResults = [...results]; +} + +async function resolveRunPlan({ + thread = {}, + session = {}, + source = {}, +}: { + thread?: Record; + session?: Record; + source?: Record; +} = {}) { + return AgentRunPlanResolver.resolveForRunAdmission({ + thread: buildThread(thread), + session: buildSession(session), + source: buildSource(source), + userIdentity, + requestedProvider: null, + requestedModel: null, + runtimeOptions: { maxIterations: 12 }, + }); +} + +function getCapabilityAccess( + result: Awaited>, + capabilityId: AgentCapabilityCatalogId +) { + return result.runPlanSnapshot.capabilities.resolvedCapabilityAccess.find( + (capability) => capability.capabilityId === capabilityId + ); +} + +describe('First-party agent definition integration regressions', () => { + let thread: ReturnType; + let session: ReturnType; + let source: ReturnType; + + beforeEach(() => { + jest.clearAllMocks(); + + mockDefinitionRows = SYSTEM_AGENT_DEFINITION_IDS.map((agentId) => + buildDefinitionRow(SYSTEM_AGENT_DEFINITIONS[agentId]) + ); + thread = buildThread(); + session = buildSession(); + source = buildSource(); + mockActiveRuns(); + + mockDefinitionUpsert.mockImplementation(async (row) => row); + mockDefinitionFindOne.mockImplementation(async ({ definitionId, ownerKind }) => { + return mockDefinitionRows.find((row) => row.definitionId === definitionId && row.ownerKind === ownerKind) || null; + }); + mockDefinitionOrderBy.mockImplementation(async () => mockDefinitionRows); + + mockResolveSessionContext.mockResolvedValue({ + repoFullName: 'example-org/example-repo', + approvalPolicy: { defaultMode: 'require_approval', rules: {} }, + capabilityPolicy: undefined, + }); + mockResolveSelection.mockResolvedValue({ + provider: 'openai', + modelId: 'gpt-5.4', + }); + + mockGetOwnedThreadWithSession.mockImplementation(async () => ({ thread, session })); + mockGetSessionSource.mockImplementation(async () => source); + mockThreadTransaction.mockImplementation(async (callback) => callback({ trx: true })); + mockPatchAndFetchById.mockImplementation(async (_id, patch) => ({ + ...thread, + ...patch, + })); + mockThreadQuery.mockImplementation(() => ({ + patchAndFetchById: mockPatchAndFetchById, + })); + mockRunQuery.mockImplementation(() => { + const query = { + where: jest.fn(() => query), + whereNotIn: jest.fn(() => query), + first: jest.fn(async () => (mockActiveRunResults.length > 0 ? mockActiveRunResults.shift() : null)), + }; + return query; + }); + mockInsertAndFetch.mockResolvedValue({ uuid: 'sample-switch-message' }); + mockMessageQuery.mockImplementation(() => ({ + insertAndFetch: mockInsertAndFetch, + })); + }); + + it('seeds public system agent definitions without reserved capability leakage or compat ids', async () => { + const seeded = await ensureSystemAgentDefinitionsSeeded(); + + expect(mockDefinitionUpsert).toHaveBeenCalledTimes(3); + expect(seeded.map((definition) => definition.id).sort()).toEqual([ + 'system.debug', + 'system.develop', + 'system.freeform', + ]); + expect(seeded).toEqual( + expect.arrayContaining([ + expect.objectContaining({ + id: 'system.debug', + owner: { kind: 'system', userId: null, organizationId: null }, + codeOwned: true, + readOnly: true, + }), + expect.objectContaining({ + id: 'system.develop', + owner: { kind: 'system', userId: null, organizationId: null }, + codeOwned: true, + readOnly: true, + }), + expect.objectContaining({ + id: 'system.freeform', + owner: { kind: 'system', userId: null, organizationId: null }, + codeOwned: true, + readOnly: true, + }), + ]) + ); + + expect(seeded.map((definition) => serializeAgentDefinitionSummary(definition).id).sort()).toEqual([ + 'system.debug', + 'system.develop', + 'system.freeform', + ]); + }); + + it('keeps Debug diagnostic and protected fix capability refs with workspaceRequired false', async () => { + const debug = await getSystemAgentDefinition('system.debug'); + + expect(debug.resourcePolicy.workspaceRequired).toBe(false); + expect(debug.requiredCapabilityRefs).toEqual( + expect.arrayContaining([ + 'diagnostics_logs', + 'diagnostics_codefresh', + 'diagnostics_kubernetes', + 'diagnostics_database', + 'github_write', + 'external_mcp_write', + ]) + ); + + const result = await resolveRunPlan({ + source: { + input: { buildUuid: 'sample-build', branchName: 'main' }, + }, + }); + + expect(result.runPlanSnapshot.agent.id).toBe('system.debug'); + expect(result.runPlanSnapshot.source.repoFullName).toBe('example-org/example-repo'); + expect(getCapabilityAccess(result, 'diagnostics_kubernetes')).toEqual( + expect.objectContaining({ + allowed: true, + availability: 'system_only', + }) + ); + expect(getCapabilityAccess(result, 'github_write')).toEqual( + expect.objectContaining({ + allowed: true, + availability: 'system_only', + approvalMode: 'require_approval', + }) + ); + expect(getCapabilityAccess(result, 'external_mcp_write')).toEqual( + expect.objectContaining({ + allowed: true, + availability: 'admin_only', + approvalMode: 'require_approval', + }) + ); + }); + + it('fails Develop without prepared workspace/source resources and keeps Free-form minimal capabilities', async () => { + const develop = await getSystemAgentDefinition('system.develop'); + const freeform = await getSystemAgentDefinition('system.freeform'); + + expect(develop.resourcePolicy.workspaceRequired).toBe(true); + expect(develop.resourcePolicy.sandboxRequired).toBe(true); + expect(freeform.requiredCapabilityRefs).toEqual(['read_context', 'external_mcp_read']); + + const freeformRun = await resolveRunPlan(); + expect(freeformRun.runPlanSnapshot.agent.id).toBe('system.freeform'); + expect(freeformRun.runPlanSnapshot.capabilities.provisionalCapabilityIds).toEqual([ + 'read_context', + 'external_mcp_read', + ]); + expect(serializeRunPlanSummary(freeformRun.runPlanSnapshot)?.agent.id).toBe('system.freeform'); + + await expect( + resolveRunPlan({ + thread: { + metadata: { selectedAgentDefinitionId: 'system.develop' }, + }, + }) + ).rejects.toMatchObject({ + name: AgentRunPlanAgentUnavailableError.name, + agentId: 'system.develop', + reason: 'source_incompatible', + details: { sourceKind: 'freeform_chat' }, + }); + }); + + it('allows system_only capabilities for system definitions and denies user-definition privilege escalation', () => { + const systemAccess = AgentPolicyService.resolveCapabilitySetAccess( + ['diagnostics_database', 'github_write', 'approval_controls'], + { + definitionOwnerKind: 'system', + sourceKind: 'build_context_chat', + } + ); + expect(systemAccess).toEqual( + expect.arrayContaining([ + expect.objectContaining({ + capabilityId: 'diagnostics_database', + allowed: true, + effectiveAvailability: 'system_only', + }), + expect.objectContaining({ + capabilityId: 'github_write', + allowed: true, + effectiveAvailability: 'system_only', + }), + ]) + ); + + const syntheticUserDefinition: AgentDefinitionContract = { + id: 'user.sample-agent', + version: 1, + owner: { kind: 'user', userId: 'sample-user' }, + name: 'Sample agent', + instructionRefs: [], + capabilityRefs: ['diagnostics_database'], + requiredCapabilityRefs: ['diagnostics_database'], + optionalCapabilityRefs: [], + resourcePolicy: { + sourceKinds: ['build_context_chat'], + workspaceRequired: false, + sandboxRequired: false, + }, + status: 'active', + }; + const userAccess = AgentPolicyService.resolveCapabilitySetAccess(syntheticUserDefinition.requiredCapabilityRefs!, { + definitionOwnerKind: syntheticUserDefinition.owner.kind, + sourceKind: 'build_context_chat', + }); + + expect(userAccess).toEqual([ + expect.objectContaining({ + capabilityId: 'diagnostics_database', + allowed: false, + reason: 'system_only', + effectiveAvailability: 'system_only', + }), + ]); + }); +}); diff --git a/src/server/services/agent/__tests__/MessageStore.test.ts b/src/server/services/agent/__tests__/MessageStore.test.ts index 83601d08..bde747cc 100644 --- a/src/server/services/agent/__tests__/MessageStore.test.ts +++ b/src/server/services/agent/__tests__/MessageStore.test.ts @@ -32,6 +32,7 @@ jest.mock('../ThreadService', () => ({ __esModule: true, default: { getOwnedThread: jest.fn(), + getOwnedThreadWithSession: jest.fn(), }, })); @@ -49,6 +50,8 @@ const mockGetOwnedThread = AgentThreadService.getOwnedThread as jest.Mock; describe('AgentMessageStore', () => { beforeEach(() => { jest.clearAllMocks(); + mockMessageQuery.mockReset(); + mockGetOwnedThread.mockReset(); }); describe('serializeCanonicalMessage', () => { @@ -76,6 +79,99 @@ describe('AgentMessageStore', () => { createdAt: '2026-04-25T00:00:00.000Z', }); }); + + it('returns typed agent switch system messages but rejects unrelated system messages', () => { + const switchMessage = AgentMessageStore.serializeCanonicalMessage( + { + uuid: '22222222-2222-4222-8222-222222222222', + clientMessageId: null, + role: 'system', + parts: [{ type: 'text', text: 'You switched Debug -> Develop. Applies to future runs.' }], + metadata: { + kind: 'agent_switch', + beforeAgent: { id: 'system.debug', label: 'Debug' }, + afterAgent: { id: 'system.develop', label: 'Develop' }, + }, + createdAt: '2026-04-25T00:00:00.000Z', + } as any, + 'thread-uuid' + ); + + expect(switchMessage).toEqual( + expect.objectContaining({ + role: 'system', + metadata: expect.objectContaining({ kind: 'agent_switch' }), + }) + ); + expect(() => + AgentMessageStore.serializeCanonicalMessage( + { + uuid: '33333333-3333-4333-8333-333333333333', + role: 'system', + parts: [{ type: 'text', text: 'Hidden status' }], + metadata: { kind: 'internal_status' }, + } as any, + 'thread-uuid' + ) + ).toThrow('Agent message is not a public canonical message'); + }); + }); + + describe('createAgentSwitchEvent', () => { + it('inserts a server-authored agent_switch system event with metadata fields and no run id', async () => { + const insertAndFetch = jest.fn().mockResolvedValue({ uuid: 'message-1' }); + mockMessageQuery.mockReturnValueOnce({ insertAndFetch }); + + await AgentMessageStore.createAgentSwitchEvent({ + thread: { id: 17 }, + actor: { userId: 'sample-user', label: 'You' }, + beforeAgent: { id: 'system.debug', label: 'Debug' }, + afterAgent: { id: 'system.develop', label: 'Develop' }, + occurredAt: '2026-05-01T00:00:00.000Z', + }); + + expect(insertAndFetch).toHaveBeenCalledWith( + expect.objectContaining({ + threadId: 17, + runId: null, + role: 'system', + parts: [{ type: 'text', text: 'You switched Debug -> Develop. Applies to future runs.' }], + metadata: expect.objectContaining({ + kind: 'agent_switch', + actor: { userId: 'sample-user', label: 'You' }, + beforeAgent: { id: 'system.debug', label: 'Debug' }, + afterAgent: { id: 'system.develop', label: 'Develop' }, + appliesTo: 'future_runs', + occurredAt: '2026-05-01T00:00:00.000Z', + }), + }) + ); + }); + + it('inserts a custom agent switch event with label-only visible copy and backend id metadata', async () => { + const insertAndFetch = jest.fn().mockResolvedValue({ uuid: 'message-1' }); + mockMessageQuery.mockReturnValueOnce({ insertAndFetch }); + + await AgentMessageStore.createAgentSwitchEvent({ + thread: { id: 17 }, + actor: { userId: 'sample-user', label: 'Sample User' }, + beforeAgent: { id: 'system.debug', label: 'Debug' }, + afterAgent: { id: 'custom.sample-agent', label: 'Custom helper' }, + occurredAt: '2026-05-01T00:00:00.000Z', + }); + + expect(insertAndFetch).toHaveBeenCalledWith( + expect.objectContaining({ + parts: [{ type: 'text', text: 'Sample User switched Debug -> Custom helper. Applies to future runs.' }], + metadata: expect.objectContaining({ + kind: 'agent_switch', + beforeAgent: { id: 'system.debug', label: 'Debug' }, + afterAgent: { id: 'custom.sample-agent', label: 'Custom helper' }, + }), + }) + ); + expect(JSON.stringify(insertAndFetch.mock.calls[0][0].parts)).not.toContain('custom.sample-agent'); + }); }); describe('listMessages', () => { @@ -306,6 +402,70 @@ describe('AgentMessageStore', () => { }) ); }); + + it('preserves existing assistant run ownership when saving a later run transcript', async () => { + const existingRow = { + id: 11, + uuid: '22222222-2222-4222-8222-222222222222', + threadId: 17, + runId: 101, + role: 'assistant', + parts: [{ type: 'text', text: 'Previous response' }], + clientMessageId: null, + metadata: { runId: 'run-old' }, + }; + const existingWhere = jest.fn().mockResolvedValue([existingRow]); + const patchAndFetchById = jest.fn().mockResolvedValue(existingRow); + const insert = jest.fn().mockResolvedValue({ + id: 12, + uuid: '33333333-3333-4333-8333-333333333333', + threadId: 17, + runId: 202, + role: 'assistant', + clientMessageId: null, + metadata: { runId: 'run-new' }, + }); + + mockMessageQuery + .mockReturnValueOnce({ where: existingWhere }) + .mockReturnValueOnce({ patchAndFetchById }) + .mockReturnValueOnce({ insert }); + + await AgentMessageStore.upsertCanonicalUiMessagesForThread( + { id: 17 }, + [ + { + id: '22222222-2222-4222-8222-222222222222', + role: 'assistant', + metadata: { runId: 'run-old' }, + parts: [{ type: 'text', text: 'Previous response' }], + } as any, + { + id: '33333333-3333-4333-8333-333333333333', + role: 'assistant', + metadata: { runId: 'run-new' }, + parts: [{ type: 'text', text: 'Current response' }], + } as any, + ], + { runId: 202 } + ); + + expect(patchAndFetchById).toHaveBeenCalledWith( + 11, + expect.objectContaining({ + runId: 101, + parts: [{ type: 'text', text: 'Previous response' }], + }) + ); + expect(insert).toHaveBeenCalledWith( + expect.objectContaining({ + threadId: 17, + role: 'assistant', + runId: 202, + parts: [{ type: 'text', text: 'Current response' }], + }) + ); + }); }); describe('syncCanonicalMessagesFromUiMessages', () => { diff --git a/src/server/services/agent/__tests__/PolicyService.test.ts b/src/server/services/agent/__tests__/PolicyService.test.ts index c3deded9..478b1b08 100644 --- a/src/server/services/agent/__tests__/PolicyService.test.ts +++ b/src/server/services/agent/__tests__/PolicyService.test.ts @@ -47,4 +47,233 @@ describe('AgentPolicyService', () => { 'require_approval' ); }); + + it('allows all user-owned definitions to use all-users capabilities', () => { + const result = AgentPolicyService.resolveCapabilityAccess({ + capabilityId: 'workspace_files', + definitionOwnerKind: 'user', + sourceKind: 'workspace_session', + }); + + expect(result).toEqual( + expect.objectContaining({ + allowed: true, + effectiveAvailability: 'all_users', + approvalMode: 'require_approval', + }) + ); + }); + + it('blocks disabled capabilities for every definition owner', () => { + const result = AgentPolicyService.resolveCapabilityAccess({ + capabilityId: 'read_context', + capabilityPolicy: { + availability: { + read_context: 'disabled', + }, + }, + definitionOwnerKind: 'system', + sourceKind: 'freeform_chat', + }); + + expect(result).toEqual( + expect.objectContaining({ + allowed: false, + reason: 'disabled', + effectiveAvailability: 'disabled', + }) + ); + }); + + it('blocks system-only capabilities for non-system definitions', () => { + const result = AgentPolicyService.resolveCapabilityAccess({ + capabilityId: 'diagnostics_database', + definitionOwnerKind: 'admin', + requesterIsAdmin: true, + sourceKind: 'build_context_chat', + }); + + expect(result).toEqual( + expect.objectContaining({ + allowed: false, + reason: 'system_only', + effectiveAvailability: 'system_only', + }) + ); + }); + + it('allows system-owned definitions to use system-only capabilities', () => { + const result = AgentPolicyService.resolveCapabilityAccess({ + capabilityId: 'diagnostics_kubernetes', + definitionOwnerKind: 'system', + sourceKind: 'build_context_chat', + }); + + expect(result).toEqual( + expect.objectContaining({ + allowed: true, + effectiveAvailability: 'system_only', + approvalMode: 'allow', + }) + ); + }); + + it('blocks user-owned definitions from admin-only capabilities', () => { + const result = AgentPolicyService.resolveCapabilityAccess({ + capabilityId: 'external_mcp_write', + definitionOwnerKind: 'user', + sourceKind: 'freeform_chat', + }); + + expect(result).toEqual( + expect.objectContaining({ + allowed: false, + reason: 'admin_only', + effectiveAvailability: 'admin_only', + }) + ); + }); + + it('allows admin-owned definitions to use admin-only capabilities', () => { + const result = AgentPolicyService.resolveCapabilityAccess({ + capabilityId: 'external_mcp_write', + definitionOwnerKind: 'admin', + requesterIsAdmin: true, + sourceKind: 'workspace_session', + }); + + expect(result).toEqual( + expect.objectContaining({ + allowed: true, + effectiveAvailability: 'admin_only', + }) + ); + }); + + it('uses configured availability over catalog defaults', () => { + const result = AgentPolicyService.resolveCapabilityAccess({ + capabilityId: 'diagnostics_codefresh', + capabilityPolicy: { + availability: { + diagnostics_codefresh: 'all_users', + }, + }, + definitionOwnerKind: 'user', + sourceKind: 'build_context_chat', + }); + + expect(result).toEqual( + expect.objectContaining({ + allowed: true, + configuredAvailability: 'all_users', + effectiveAvailability: 'all_users', + }) + ); + }); + + it('blocks user-owned definitions from creator-reserved capabilities even when runtime policy allows them', () => { + const result = AgentPolicyService.resolveCapabilityAccess({ + capabilityId: 'workspace_files', + capabilityPolicy: { + availability: { + workspace_files: 'all_users', + }, + }, + customAgentCreationPolicy: { + capabilityAvailability: { + workspace_files: 'reserved', + }, + }, + definitionOwnerKind: 'user', + sourceKind: 'workspace_session', + }); + + expect(result).toEqual( + expect.objectContaining({ + allowed: false, + reason: 'creator_capability_reserved', + effectiveAvailability: 'all_users', + }) + ); + }); + + it('does not apply creator-reserved policy to system-owned definitions', () => { + const result = AgentPolicyService.resolveCapabilityAccess({ + capabilityId: 'diagnostics_kubernetes', + customAgentCreationPolicy: { + capabilityAvailability: { + diagnostics_kubernetes: 'reserved', + }, + }, + definitionOwnerKind: 'system', + sourceKind: 'build_context_chat', + }); + + expect(result).toEqual( + expect.objectContaining({ + allowed: true, + effectiveAvailability: 'system_only', + }) + ); + }); + + it('derives approval mode from mapped runtime approval policy', () => { + const result = AgentPolicyService.resolveCapabilityAccess({ + capabilityId: 'workspace_shell', + approvalPolicy: { + defaultMode: 'allow', + rules: { + ...DEFAULT_AGENT_APPROVAL_POLICY.rules, + shell_exec: 'deny', + }, + }, + definitionOwnerKind: 'user', + sourceKind: 'workspace_session', + }); + + expect(result).toEqual( + expect.objectContaining({ + allowed: true, + approvalMode: 'deny', + }) + ); + }); + + it('blocks source-incompatible capabilities', () => { + const result = AgentPolicyService.resolveCapabilityAccess({ + capabilityId: 'workspace_shell', + definitionOwnerKind: 'user', + sourceKind: 'freeform_chat', + }); + + expect(result).toEqual( + expect.objectContaining({ + allowed: false, + reason: 'source_incompatible', + }) + ); + }); + + it('blocks unknown capability ids', () => { + const result = AgentPolicyService.resolveCapabilityAccess({ + capabilityId: 'sample_unknown', + definitionOwnerKind: 'system', + }); + + expect(result).toEqual({ + capabilityId: 'sample_unknown', + allowed: false, + reason: 'unknown_capability', + }); + }); + + it('resolves capability sets in input order', () => { + const result = AgentPolicyService.resolveCapabilitySetAccess(['read_context', 'external_mcp_write'], { + definitionOwnerKind: 'user', + sourceKind: 'freeform_chat', + }); + + expect(result.map((entry) => entry.capabilityId)).toEqual(['read_context', 'external_mcp_write']); + expect(result.map((entry) => entry.allowed)).toEqual([true, false]); + }); }); diff --git a/src/server/services/agent/__tests__/ProviderRegistry.test.ts b/src/server/services/agent/__tests__/ProviderRegistry.test.ts index ce5ad68f..b9a205ef 100644 --- a/src/server/services/agent/__tests__/ProviderRegistry.test.ts +++ b/src/server/services/agent/__tests__/ProviderRegistry.test.ts @@ -18,7 +18,7 @@ import type { AgentModelSummary } from '../types'; const mockGetEffectiveConfig = jest.fn(); -jest.mock('server/services/aiAgentConfig', () => ({ +jest.mock('server/services/agentRuntime/config/agentRuntimeConfig', () => ({ __esModule: true, default: { getInstance: jest.fn(() => ({ diff --git a/src/server/services/agent/__tests__/RunAdmissionService.test.ts b/src/server/services/agent/__tests__/RunAdmissionService.test.ts index 7abc4970..83850e9d 100644 --- a/src/server/services/agent/__tests__/RunAdmissionService.test.ts +++ b/src/server/services/agent/__tests__/RunAdmissionService.test.ts @@ -68,6 +68,112 @@ const mockFindCanonicalMessageByClientMessageId = AgentMessageStore.findCanonica const mockInsertUserMessageForRun = AgentMessageStore.insertUserMessageForRun as jest.Mock; const mockAppendStatusEvent = AgentRunEventService.appendStatusEvent as jest.Mock; +const runPlanSnapshot = { + version: 1, + capturedAt: '2026-05-01T00:00:00.000Z', + agent: { + id: 'system.freeform', + label: 'Free-form', + ownerKind: 'system', + version: 1, + sourceKind: 'freeform_chat', + resourcePolicy: { + sourceKinds: ['build_context_chat', 'workspace_session', 'freeform_chat'], + workspaceRequired: false, + sandboxRequired: false, + }, + modelPreference: null, + }, + source: { + id: 'source-1', + adapter: 'blank_workspace', + status: 'ready', + sessionKind: 'chat', + repoFullName: 'example-org/example-repo', + freshness: { + capturedAt: '2026-05-01T00:00:00.000Z', + freshnessSource: 'source', + }, + }, + model: { + requestedProvider: null, + requestedModel: null, + resolvedProvider: 'openai', + resolvedModel: 'gpt-5.4', + }, + runtime: { + requestedHarness: null, + resolvedHarness: 'lifecycle_ai_sdk', + sandboxRequirement: { filesystem: 'persistent' }, + runtimeOptions: { maxIterations: 12 }, + approvalPolicy: { defaultMode: 'require_approval', rules: {} }, + }, + prompt: { + instructionRefs: [], + renderedSummary: 'Sample prompt summary', + renderedHash: 'sha256:sample-rendered-prompt', + }, + capabilities: { + provisionalCapabilityIds: [], + resolvedCapabilityAccess: [], + }, + warnings: [], +} as const; + +const customRunPlanSnapshot = { + ...runPlanSnapshot, + agent: { + id: 'custom.sample-agent', + label: 'Sample custom agent', + ownerKind: 'user', + version: 3, + sourceKind: 'freeform_chat', + resourcePolicy: { + sourceKinds: ['freeform_chat'], + workspaceRequired: false, + sandboxRequired: false, + }, + modelPreference: { + provider: 'anthropic', + model: 'claude-sonnet-4.6', + }, + }, + model: { + requestedProvider: 'anthropic', + requestedModel: 'claude-sonnet-4.6', + resolvedProvider: 'anthropic', + resolvedModel: 'claude-sonnet-4.6', + }, + runtime: { + requestedHarness: null, + resolvedHarness: 'lifecycle_ai_sdk', + sandboxRequirement: { filesystem: 'persistent' }, + runtimeOptions: { maxIterations: 9 }, + approvalPolicy: { + defaultMode: 'require_approval', + rules: { read: 'allow' }, + }, + }, + prompt: { + instructionRefs: [], + instructionAddendum: 'Use the sample custom instructions.', + renderedSummary: 'Sample custom agent description', + renderedHash: 'sha256:sample-custom-agent-prompt', + }, + capabilities: { + provisionalCapabilityIds: ['read_context'], + resolvedCapabilityAccess: [ + { + capabilityId: 'read_context', + availability: 'all_users', + allowed: true, + runtimeCapabilityKey: 'read', + approvalMode: 'allow', + }, + ], + }, +} as const; + function buildActiveRunQuery(activeRun: unknown = null) { const query = { where: jest.fn(), @@ -123,6 +229,7 @@ describe('AgentRunAdmissionService', () => { resolvedProvider: 'openai', resolvedModel: 'gpt-5.4', runtimeOptions: { maxIterations: 12 }, + runPlanSnapshot, }); expect(admission).toEqual({ @@ -147,6 +254,7 @@ describe('AgentRunAdmissionService', () => { policySnapshot: expect.objectContaining({ runtimeOptions: { maxIterations: 12 }, }), + runPlanSnapshot, }) ); expect(mockAppendStatusEvent).toHaveBeenCalledWith('run-1', 'run.queued', { @@ -172,6 +280,7 @@ describe('AgentRunAdmissionService', () => { resolvedHarness: 'lifecycle_ai_sdk', resolvedProvider: 'openai', resolvedModel: 'gpt-5.4', + runPlanSnapshot, }) ).rejects.toThrow('Wait for the current agent run to finish before starting another run.'); @@ -208,6 +317,7 @@ describe('AgentRunAdmissionService', () => { resolvedHarness: 'lifecycle_ai_sdk', resolvedProvider: 'openai', resolvedModel: 'gpt-5.4', + runPlanSnapshot: null as any, }); expect(admission).toEqual({ @@ -218,4 +328,78 @@ describe('AgentRunAdmissionService', () => { expect(mockInsertUserMessageForRun).not.toHaveBeenCalled(); expect(mockAppendStatusEvent).not.toHaveBeenCalled(); }); + + it('persists custom-agent resolver fields without recomputing the snapshot', 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); + + 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: customRunPlanSnapshot.runtime.approvalPolicy as any, + message: { clientMessageId: 'client-message-1', parts: [{ type: 'text', text: 'Hi' }] }, + requestedHarness: customRunPlanSnapshot.runtime.requestedHarness, + requestedProvider: customRunPlanSnapshot.model.requestedProvider, + requestedModel: customRunPlanSnapshot.model.requestedModel, + resolvedHarness: customRunPlanSnapshot.runtime.resolvedHarness, + resolvedProvider: customRunPlanSnapshot.model.resolvedProvider, + resolvedModel: customRunPlanSnapshot.model.resolvedModel, + sandboxRequirement: customRunPlanSnapshot.runtime.sandboxRequirement, + runtimeOptions: customRunPlanSnapshot.runtime.runtimeOptions, + runPlanSnapshot: customRunPlanSnapshot, + }); + + expect(insertRunQuery.insertAndFetch).toHaveBeenCalledWith( + expect.objectContaining({ + requestedProvider: 'anthropic', + requestedModel: 'claude-sonnet-4.6', + resolvedProvider: 'anthropic', + resolvedModel: 'claude-sonnet-4.6', + sandboxRequirement: { filesystem: 'persistent' }, + policySnapshot: { + defaultMode: 'require_approval', + rules: { read: 'allow' }, + runtimeOptions: { maxIterations: 9 }, + }, + runPlanSnapshot: customRunPlanSnapshot, + }) + ); + }); + + it('rejects missing run plan snapshots before insert', async () => { + const activeRunQuery = buildActiveRunQuery(); + 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', + runPlanSnapshot: null as any, + }) + ).rejects.toThrow('Agent run plan snapshot is required.'); + + expect(mockInsertUserMessageForRun).not.toHaveBeenCalled(); + expect(mockAppendStatusEvent).not.toHaveBeenCalled(); + }); }); diff --git a/src/server/services/agent/__tests__/RunExecutor.test.ts b/src/server/services/agent/__tests__/RunExecutor.test.ts index a6340d75..29a27115 100644 --- a/src/server/services/agent/__tests__/RunExecutor.test.ts +++ b/src/server/services/agent/__tests__/RunExecutor.test.ts @@ -34,6 +34,130 @@ jest.mock('server/services/agent/ProviderRegistry', () => ({ }, })); +const runPlanSnapshot = { + version: 1, + capturedAt: '2026-05-01T00:00:00.000Z', + agent: { + id: 'system.freeform', + label: 'Free-form', + sourceKind: 'freeform_chat', + }, + source: { + id: 'source-1', + adapter: 'blank_workspace', + status: 'ready', + sessionKind: 'chat', + freshness: { + capturedAt: '2026-05-01T00:00:00.000Z', + freshnessSource: 'source', + }, + }, + model: { + requestedProvider: null, + requestedModel: null, + resolvedProvider: 'openai', + resolvedModel: 'gpt-5.4', + }, + runtime: { + requestedHarness: null, + resolvedHarness: 'lifecycle_ai_sdk', + sandboxRequirement: { filesystem: 'persistent' }, + runtimeOptions: {}, + approvalPolicy: 'on-request', + }, + prompt: { + instructionRefs: [], + renderedSummary: 'Sample prompt summary', + renderedHash: 'sha256:sample-rendered-prompt', + }, + capabilities: { + provisionalCapabilityIds: [], + resolvedCapabilityAccess: [], + }, + warnings: [], +} as const; + +const customAgentRunPlanSnapshot = { + ...runPlanSnapshot, + agent: { + id: 'custom.sample-agent', + label: 'Sample custom agent', + ownerKind: 'user', + version: 4, + sourceKind: 'freeform_chat', + modelPreference: { + provider: 'anthropic', + model: 'claude-sonnet-4.6', + }, + }, + model: { + requestedProvider: 'anthropic', + requestedModel: 'claude-sonnet-4.6', + resolvedProvider: 'anthropic', + resolvedModel: 'claude-sonnet-4.6', + }, + runtime: { + ...runPlanSnapshot.runtime, + runtimeOptions: { maxIterations: 6 }, + approvalPolicy: { + defaultMode: 'require_approval', + rules: { read: 'allow' }, + }, + }, + prompt: { + instructionRefs: [], + instructionAddendum: 'Use the sample custom instructions.', + renderedSummary: 'Sample custom agent description', + renderedHash: 'sha256:sample-custom-agent-prompt', + }, + capabilities: { + provisionalCapabilityIds: ['read_context'], + resolvedCapabilityAccess: [ + { + capabilityId: 'read_context', + availability: 'all_users', + allowed: true, + runtimeCapabilityKey: 'read', + approvalMode: 'allow', + }, + ], + }, +} as const; + +const mockResolveForRunAdmission = jest.fn().mockResolvedValue({ + approvalPolicy: 'on-request', + requestedHarness: null, + requestedProvider: null, + requestedModel: null, + resolvedHarness: 'lifecycle_ai_sdk', + resolvedProvider: 'openai', + resolvedModel: 'gpt-5.4', + sandboxRequirement: { filesystem: 'persistent' }, + runtimeOptions: {}, + runPlanSnapshot, +}); + +jest.mock('server/services/agent/RunPlanResolver', () => ({ + __esModule: true, + default: { + resolveForRunAdmission: (...args: unknown[]) => mockResolveForRunAdmission(...args), + }, +})); + +const mockGetSessionSource = jest.fn().mockResolvedValue({ + id: 3, + uuid: 'source-1', + status: 'ready', + sandboxRequirements: { filesystem: 'persistent' }, +}); + +jest.mock('server/services/agent/SourceService', () => ({ + __esModule: true, + default: { + getSessionSource: (...args: unknown[]) => mockGetSessionSource(...args), + }, +})); + const mockResolveSessionContext = jest.fn().mockResolvedValue({ repoFullName: 'example-org/example-repo', approvalPolicy: 'on-request', @@ -213,6 +337,24 @@ describe('AgentRunExecutor', () => { jest.clearAllMocks(); mockResolveSelection.mockResolvedValue({ provider: 'openai', modelId: 'gpt-5.4' }); mockCreateLanguageModel.mockResolvedValue({ id: 'model-instance' }); + mockGetSessionSource.mockResolvedValue({ + id: 3, + uuid: 'source-1', + status: 'ready', + sandboxRequirements: { filesystem: 'persistent' }, + }); + mockResolveForRunAdmission.mockResolvedValue({ + approvalPolicy: 'on-request', + requestedHarness: null, + requestedProvider: null, + requestedModel: null, + resolvedHarness: 'lifecycle_ai_sdk', + resolvedProvider: 'openai', + resolvedModel: 'gpt-5.4', + sandboxRequirement: { filesystem: 'persistent' }, + runtimeOptions: {}, + runPlanSnapshot, + }); mockResolveSessionContext.mockResolvedValue({ repoFullName: 'example-org/example-repo', approvalPolicy: 'on-request', @@ -383,6 +525,367 @@ describe('AgentRunExecutor', () => { ); }); + it('prefers snapshot runtime maxIterations before policySnapshot runtime options', async () => { + await 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', + policySnapshot: { runtimeOptions: { maxIterations: 3 } }, + runPlanSnapshot: { + ...runPlanSnapshot, + runtime: { + ...runPlanSnapshot.runtime, + runtimeOptions: { maxIterations: 21 }, + }, + }, + } as any, + }); + + expect(mockStepCountIs).toHaveBeenCalledWith(21); + }); + + it('prefers snapshot model and approval policy for existing queued runs', async () => { + const snapshotApprovalPolicy = { defaultMode: 'require_approval', rules: { read: 'allow' } }; + + await AgentRunExecutor.execute({ + session: { uuid: 'sess-1' } as any, + thread: { id: 7, uuid: 'thread-1' } as any, + userIdentity: { userId: 'sample-user' } as any, + messages: [], + requestedProvider: 'openai', + requestedModelId: 'gpt-5.4', + existingRun: { + id: 11, + uuid: 'queued-run-1', + status: 'queued', + executionOwner: 'worker-1', + resolvedHarness: 'lifecycle_ai_sdk', + runPlanSnapshot: { + ...runPlanSnapshot, + model: { + requestedProvider: 'openai', + requestedModel: 'gpt-5.4', + resolvedProvider: 'anthropic', + resolvedModel: 'claude-sonnet-4.6', + }, + runtime: { + ...runPlanSnapshot.runtime, + approvalPolicy: snapshotApprovalPolicy, + }, + }, + } as any, + }); + + expect(mockResolveSelection).toHaveBeenCalledWith({ + repoFullName: 'example-org/example-repo', + requestedProvider: 'anthropic', + requestedModelId: 'claude-sonnet-4.6', + }); + expect(mockBuildToolSet).toHaveBeenCalledWith( + expect.objectContaining({ + approvalPolicy: snapshotApprovalPolicy, + }) + ); + }); + + it('passes immutable snapshot MCP filters into tool setup for existing queued runs', async () => { + const snapshotCapabilities = { + ...customAgentRunPlanSnapshot.capabilities, + selectedRuntimeMcpConnectionRefs: ['global:docs'], + }; + + await 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', + resolvedHarness: 'lifecycle_ai_sdk', + runPlanSnapshot: { + ...customAgentRunPlanSnapshot, + capabilities: snapshotCapabilities, + }, + } as any, + }); + + expect(mockBuildToolSet).toHaveBeenCalledWith( + expect.objectContaining({ + resolvedCapabilityAccess: snapshotCapabilities.resolvedCapabilityAccess, + selectedRuntimeMcpConnectionRefs: ['global:docs'], + }) + ); + }); + + it('passes explicit empty snapshot capability access into tool setup for existing queued runs', async () => { + const snapshotApprovalPolicy = { + defaultMode: 'require_approval', + rules: { + read: 'allow', + }, + }; + const snapshotCapabilities = { + ...runPlanSnapshot.capabilities, + resolvedCapabilityAccess: [], + }; + + mockResolveSessionContext.mockResolvedValueOnce({ + repoFullName: 'example-org/example-repo', + approvalPolicy: { + defaultMode: 'allow', + rules: { + workspace_write: 'allow', + shell_exec: 'allow', + }, + }, + binding: null, + }); + + await 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', + resolvedHarness: 'lifecycle_ai_sdk', + runPlanSnapshot: { + ...runPlanSnapshot, + runtime: { + ...runPlanSnapshot.runtime, + approvalPolicy: snapshotApprovalPolicy, + }, + capabilities: snapshotCapabilities, + }, + } as any, + }); + + expect(mockBuildToolSet).toHaveBeenCalledWith( + expect.objectContaining({ + approvalPolicy: snapshotApprovalPolicy, + resolvedCapabilityAccess: [], + }) + ); + }); + + it('executes queued custom-agent snapshots through normal model, approval, message, and tool audit paths', async () => { + mockResolveSelection.mockResolvedValueOnce({ + provider: 'anthropic', + modelId: 'claude-sonnet-4.6', + }); + + const execution = await 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-custom-run-1', + status: 'queued', + executionOwner: 'worker-1', + resolvedHarness: 'lifecycle_ai_sdk', + runPlanSnapshot: customAgentRunPlanSnapshot, + } as any, + requestGitHubToken: 'sample-gh-token', + }); + + expect(mockResolveSelection).toHaveBeenCalledWith({ + repoFullName: 'example-org/example-repo', + requestedProvider: 'anthropic', + requestedModelId: 'claude-sonnet-4.6', + }); + expect(mockCreateLanguageModel).toHaveBeenCalledWith( + expect.objectContaining({ + selection: { provider: 'anthropic', modelId: 'claude-sonnet-4.6' }, + }) + ); + expect(mockStartRunForExecutionOwner).toHaveBeenCalledWith( + 'queued-custom-run-1', + 'worker-1', + expect.objectContaining({ + resolvedHarness: 'lifecycle_ai_sdk', + provider: 'anthropic', + model: 'claude-sonnet-4.6', + }), + { dispatchAttemptId: undefined } + ); + expect(mockStepCountIs).toHaveBeenCalledWith(6); + expect(mockToolLoopAgent).toHaveBeenCalledWith( + expect.objectContaining({ + instructions: 'DB prompt as stored\n\nUse the sample custom instructions.\n\nAppend prompt', + }) + ); + expect(mockBuildToolSet).toHaveBeenCalledWith( + expect.objectContaining({ + approvalPolicy: customAgentRunPlanSnapshot.runtime.approvalPolicy, + resolvedCapabilityAccess: customAgentRunPlanSnapshot.capabilities.resolvedCapabilityAccess, + toolRules: [], + }) + ); + + const toolSetArgs = mockBuildToolSet.mock.calls[0]?.[0]; + mockPendingActionFirst.mockResolvedValue({ + id: 55, + status: 'approved', + }); + + await toolSetArgs.hooks.onToolStarted({ + source: 'mcp', + serverSlug: 'sandbox', + toolName: 'workspace.read_file', + toolCallId: 'tool-call-1', + args: { path: 'sample-file.ts' }, + capabilityKey: 'read', + }); + + expect(mockToolExecutionInsert).toHaveBeenCalledWith( + expect.objectContaining({ + runId: 11, + toolName: 'workspace.read_file', + toolCallId: 'tool-call-1', + pendingActionId: 55, + approved: true, + safetyLevel: 'read', + }) + ); + + mockToolExecutionFirst.mockResolvedValue({ + id: 99, + startedAt: '2026-04-08T00:00:00.000Z', + }); + + await toolSetArgs.hooks.onToolFinished({ + source: 'mcp', + serverSlug: 'sandbox', + toolName: 'workspace.read_file', + toolCallId: 'tool-call-1', + args: { path: 'sample-file.ts' }, + capabilityKey: 'read', + result: { ok: true }, + status: 'completed', + }); + + expect(mockToolExecutionPatchAndFetchById).toHaveBeenCalledWith( + 99, + expect.objectContaining({ + status: 'completed', + durationMs: expect.any(Number), + }) + ); + + await execution.onStreamFinish({ + messages: [ + { + id: 'assistant-1', + role: 'assistant', + parts: [{ type: 'text', text: 'Done' }], + metadata: { runId: 'queued-custom-run-1' }, + } as any, + ], + finishReason: 'stop', + isAborted: false, + }); + + expect(mockUpsertCanonicalUiMessagesForThread).toHaveBeenCalledWith( + { id: 7, uuid: 'thread-1' }, + expect.any(Array), + expect.objectContaining({ + runId: 11, + }) + ); + expect(mockSyncApprovalRequestState).toHaveBeenCalledWith( + expect.objectContaining({ + approvalPolicy: customAgentRunPlanSnapshot.runtime.approvalPolicy, + toolRules: [], + }) + ); + expect(mockLastFinalizeResult).toEqual( + expect.objectContaining({ + status: 'completed', + }) + ); + expect(mockEnqueueRun).not.toHaveBeenCalledWith('queued-custom-run-1', 'approval_resolved', { + githubToken: 'sample-gh-token', + }); + }); + + it('uses resolver-built run plans when creating direct queued runs', async () => { + const directRunPlanSnapshot = { + ...runPlanSnapshot, + capabilities: { + ...customAgentRunPlanSnapshot.capabilities, + selectedRuntimeMcpConnectionRefs: ['global:docs'], + }, + }; + mockResolveForRunAdmission.mockResolvedValueOnce({ + approvalPolicy: directRunPlanSnapshot.runtime.approvalPolicy, + requestedHarness: null, + requestedProvider: null, + requestedModel: null, + resolvedHarness: 'lifecycle_ai_sdk', + resolvedProvider: 'openai', + resolvedModel: 'gpt-5.4', + sandboxRequirement: { filesystem: 'persistent' }, + runtimeOptions: {}, + runPlanSnapshot: directRunPlanSnapshot, + }); + + await AgentRunExecutor.execute({ + session: { id: 17, uuid: 'sess-1' } as any, + thread: { id: 7, uuid: 'thread-1' } as any, + userIdentity: { userId: 'sample-user' } as any, + messages: [], + requestedProvider: 'openai', + requestedModelId: 'gpt-5.4', + }); + + expect(mockGetSessionSource).toHaveBeenCalledWith(17); + expect(mockResolveForRunAdmission).toHaveBeenCalledWith( + expect.objectContaining({ + thread: expect.objectContaining({ id: 7, uuid: 'thread-1' }), + session: expect.objectContaining({ id: 17, uuid: 'sess-1' }), + source: expect.objectContaining({ uuid: 'source-1', status: 'ready' }), + userIdentity: { userId: 'sample-user' }, + requestedProvider: 'openai', + requestedModel: 'gpt-5.4', + runtimeOptions: {}, + }) + ); + expect(mockCreateQueuedRun).toHaveBeenCalledWith( + expect.objectContaining({ + runPlanSnapshot: directRunPlanSnapshot, + sandboxRequirement: { filesystem: 'persistent' }, + }) + ); + expect(mockBuildToolSet).toHaveBeenCalledWith( + expect.objectContaining({ + resolvedCapabilityAccess: directRunPlanSnapshot.capabilities.resolvedCapabilityAccess, + selectedRuntimeMcpConnectionRefs: ['global:docs'], + }) + ); + expect(mockStartRunForExecutionOwner).toHaveBeenCalledWith( + 'run-1', + expect.stringMatching(/^direct:/), + expect.objectContaining({ + resolvedHarness: 'lifecycle_ai_sdk', + }), + { dispatchAttemptId: undefined } + ); + }); + it('does not create a run when tool setup fails before execution starts', async () => { mockBuildToolSet.mockRejectedValueOnce(new Error('tool setup failed')); @@ -408,7 +911,13 @@ describe('AgentRunExecutor', () => { 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, + existingRun: { + id: 11, + uuid: 'queued-run-1', + status: 'queued', + executionOwner: 'worker-1', + runPlanSnapshot, + } as any, }) ).rejects.toThrow('tool setup failed'); @@ -422,6 +931,27 @@ describe('AgentRunExecutor', () => { ); }); + it('rejects existing queued runs that do not have an immutable run plan snapshot', async () => { + 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('Agent run plan snapshot is required for execution.'); + + expect(mockBuildToolSet).not.toHaveBeenCalled(); + expect(mockMarkFailedForExecutionOwner).toHaveBeenCalledWith( + 'queued-run-1', + 'worker-1', + expect.objectContaining({ message: 'Agent run plan snapshot is required for execution.' }), + expect.any(Object), + { dispatchAttemptId: undefined } + ); + }); + it('records a runtime session failure when the workspace gateway is unavailable', async () => { mockBuildToolSet.mockRejectedValueOnce( new SessionWorkspaceGatewayUnavailableError({ diff --git a/src/server/services/agent/__tests__/RunPlanResolver.test.ts b/src/server/services/agent/__tests__/RunPlanResolver.test.ts new file mode 100644 index 00000000..508e06a7 --- /dev/null +++ b/src/server/services/agent/__tests__/RunPlanResolver.test.ts @@ -0,0 +1,937 @@ +/** + * 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. + */ + +const mockWarn = jest.fn(); + +jest.mock('server/lib/logger', () => ({ + getLogger: jest.fn(() => ({ + warn: mockWarn, + })), +})); + +const mockResolveSessionContext = jest.fn(); + +jest.mock('server/services/agent/CapabilityService', () => ({ + __esModule: true, + default: { + resolveSessionContext: (...args: unknown[]) => mockResolveSessionContext(...args), + }, +})); + +const mockResolveSelection = jest.fn(); +const mockEnsureSystemAgentDefinitionsSeeded = jest.fn(); +const mockGetSystemAgentDefinition = jest.fn(); +const mockGetUserDefinition = jest.fn(); +const mockResolveRunAdmissionChoices = jest.fn(); + +jest.mock('server/services/agent/ProviderRegistry', () => ({ + __esModule: true, + default: { + resolveSelection: (...args: unknown[]) => mockResolveSelection(...args), + }, +})); + +jest.mock('../AgentDefinitionRegistry', () => { + const actual = jest.requireActual('../AgentDefinitionRegistry'); + return { + __esModule: true, + ...actual, + ensureSystemAgentDefinitionsSeeded: (...args: unknown[]) => mockEnsureSystemAgentDefinitionsSeeded(...args), + getSystemAgentDefinition: (...args: unknown[]) => mockGetSystemAgentDefinition(...args), + }; +}); + +jest.mock('../CustomAgentDefinitionService', () => { + class MockCustomAgentDefinitionServiceError extends Error { + code: string; + + constructor(code: string, message: string) { + super(message); + this.name = 'CustomAgentDefinitionServiceError'; + this.code = code; + } + } + + return { + __esModule: true, + CustomAgentDefinitionServiceError: MockCustomAgentDefinitionServiceError, + customAgentDefinitionService: { + getUserDefinition: (...args: unknown[]) => mockGetUserDefinition(...args), + }, + }; +}); + +jest.mock('../ThreadRuntimeControlsService', () => ({ + __esModule: true, + default: { + resolveRunAdmissionChoices: (...args: unknown[]) => mockResolveRunAdmissionChoices(...args), + }, +})); + +import AgentRunPlanResolver, { + AgentRunPlanCapabilityUnavailableError, + AgentRunPlanAgentUnavailableError, +} from '../RunPlanResolver'; +import { CustomAgentDefinitionServiceError } from '../CustomAgentDefinitionService'; +import { serializeRunPlanSummary } from '../runPlanSummary'; +import { SYSTEM_AGENT_DEFINITIONS } from '../systemAgentDefinitions'; +import { AgentSessionKind, AgentWorkspaceStatus } from 'shared/constants'; + +const userIdentity = { + userId: 'sample-user', + githubUsername: 'sample-user', + preferredUsername: 'sample-user', + email: 'sample-user@example.com', + firstName: 'Sample', + lastName: 'User', + displayName: 'Sample User', + gitUserName: 'Sample User', + gitUserEmail: 'sample-user@example.com', + roles: [], +}; + +const customDefinition = { + id: 'custom.sample-agent', + version: 3, + owner: { kind: 'user' as const, userId: 'sample-user' }, + name: 'Sample custom agent', + description: 'Custom agent description', + instructionRefs: [], + instructionAddendum: 'Use the sample custom instructions.', + capabilityRefs: ['read_context'], + requiredCapabilityRefs: [], + optionalCapabilityRefs: ['read_context'], + resourcePolicy: { + sourceKinds: ['freeform_chat'], + workspaceRequired: false, + sandboxRequired: false, + }, + modelPreference: { + provider: 'openai', + model: 'gpt-5.4', + }, + status: 'active' as const, + codeOwned: false, + readOnly: false, +}; + +function buildSession(overrides: Record = {}) { + return { + id: 17, + uuid: 'session-1', + sessionKind: AgentSessionKind.CHAT, + defaultHarness: 'lifecycle_ai_sdk', + defaultModel: 'gpt-5.4', + buildUuid: null, + namespace: 'sample-namespace', + workspaceRepos: [ + { + repo: 'example-org/example-repo', + branch: 'main', + primary: true, + }, + ], + selectedServices: [ + { + name: 'sample-service', + repo: 'example-org/example-repo', + branch: 'main', + }, + ], + ...overrides, + } as any; +} + +function buildSource(overrides: Record = {}) { + return { + id: 3, + uuid: 'source-1', + adapter: 'blank_workspace', + status: 'ready', + input: {}, + preparedSource: {}, + sandboxRequirements: { filesystem: 'persistent' }, + preparedAt: '2026-05-01T00:00:00.000Z', + ...overrides, + } as any; +} + +async function resolve( + overrides: { + thread?: Record; + session?: Record; + source?: Record; + } = {} +) { + return AgentRunPlanResolver.resolveForRunAdmission({ + thread: { id: 7, uuid: 'thread-1', metadata: {}, ...overrides.thread } as any, + session: buildSession(overrides.session), + source: buildSource(overrides.source), + userIdentity, + requestedProvider: null, + requestedModel: null, + runtimeOptions: { maxIterations: 12 }, + }); +} + +describe('AgentRunPlanResolver', () => { + beforeEach(() => { + jest.clearAllMocks(); + mockResolveSessionContext.mockResolvedValue({ + repoFullName: 'example-org/example-repo', + approvalPolicy: { defaultMode: 'require_approval', rules: {} }, + capabilityPolicy: undefined, + }); + mockResolveSelection.mockResolvedValue({ + provider: 'openai', + modelId: 'gpt-5.4', + }); + mockEnsureSystemAgentDefinitionsSeeded.mockResolvedValue(Object.values(SYSTEM_AGENT_DEFINITIONS)); + mockGetSystemAgentDefinition.mockImplementation(async (agentId) => SYSTEM_AGENT_DEFINITIONS[agentId]); + mockGetUserDefinition.mockResolvedValue(customDefinition); + mockResolveRunAdmissionChoices.mockResolvedValue({ + metadataPresent: false, + selectedRuntimeToolChoiceIds: undefined, + selectedRuntimeMcpChoiceIds: undefined, + selectedRuntimeCapabilityIds: undefined, + selectedRuntimeMcpConnectionRefs: undefined, + }); + }); + + it('infers Debug for build-context chat before generic chat', async () => { + const result = await resolve({ + source: { + input: { buildUuid: 'build-1', branchName: 'feature-branch' }, + }, + }); + + expect(result.runPlanSnapshot.agent.id).toBe('system.debug'); + expect(result.runPlanSnapshot.agent.label).toBe('Debug'); + expect(result.runPlanSnapshot.agent.sourceKind).toBe('build_context_chat'); + expect(result.runPlanSnapshot.source.buildUuid).toBe('build-1'); + expect(result.runPlanSnapshot.capabilities.provisionalCapabilityIds).toEqual( + expect.arrayContaining(['diagnostics_codefresh', 'diagnostics_kubernetes', 'diagnostics_database']) + ); + expect(result.runPlanSnapshot.capabilities.resolvedCapabilityAccess).toEqual( + expect.arrayContaining([ + expect.objectContaining({ + capabilityId: 'diagnostics_codefresh', + availability: 'system_only', + allowed: true, + runtimeCapabilityKey: 'read', + }), + ]) + ); + }); + + it('infers Free-form for chat sessions without build context', async () => { + const result = await resolve(); + + expect(result.runPlanSnapshot.agent.id).toBe('system.freeform'); + expect(result.runPlanSnapshot.agent.label).toBe('Free-form'); + expect(result.runPlanSnapshot.agent.sourceKind).toBe('freeform_chat'); + expect(result.runPlanSnapshot.capabilities.provisionalCapabilityIds).toEqual(['read_context', 'external_mcp_read']); + expect(serializeRunPlanSummary(result.runPlanSnapshot)?.agent).toEqual( + expect.objectContaining({ + id: 'system.freeform', + label: 'Free-form', + }) + ); + }); + + it('does not apply creator-reserved policy to system agent definitions during run admission', async () => { + mockResolveSessionContext.mockResolvedValueOnce({ + repoFullName: 'example-org/example-repo', + approvalPolicy: { defaultMode: 'require_approval', rules: {} }, + customAgentCreationPolicy: { + capabilityAvailability: { + external_mcp_read: 'reserved', + }, + }, + }); + + const result = await resolve(); + + expect(result.runPlanSnapshot.agent.id).toBe('system.freeform'); + expect(result.runPlanSnapshot.capabilities.provisionalCapabilityIds).toEqual(['read_context', 'external_mcp_read']); + expect(result.runPlanSnapshot.capabilities.resolvedCapabilityAccess).toEqual( + expect.arrayContaining([ + expect.objectContaining({ + capabilityId: 'external_mcp_read', + allowed: true, + }), + ]) + ); + }); + + it('infers Develop for environment and sandbox workspace sessions', async () => { + const environment = await resolve({ + session: { + sessionKind: AgentSessionKind.ENVIRONMENT, + }, + source: { + adapter: 'lifecycle_environment', + }, + }); + const sandbox = await resolve({ + session: { + sessionKind: AgentSessionKind.SANDBOX, + }, + source: { + adapter: 'lifecycle_fork', + }, + }); + + expect(environment.runPlanSnapshot.agent.id).toBe('system.develop'); + expect(environment.runPlanSnapshot.agent.sourceKind).toBe('workspace_session'); + expect(environment.runPlanSnapshot.source).toEqual( + expect.objectContaining({ + adapter: 'lifecycle_environment', + sessionKind: AgentSessionKind.ENVIRONMENT, + repoFullName: 'example-org/example-repo', + namespace: 'sample-namespace', + }) + ); + expect(environment.runPlanSnapshot.source.workspaceLayout).toEqual( + expect.objectContaining({ + repoCount: 1, + selectedServiceCount: 1, + primaryService: 'sample-service', + }) + ); + expect(sandbox.runPlanSnapshot.agent.id).toBe('system.develop'); + expect(sandbox.runPlanSnapshot.agent.sourceKind).toBe('workspace_session'); + expect(sandbox.runPlanSnapshot.source).toEqual( + expect.objectContaining({ + adapter: 'lifecycle_fork', + sessionKind: AgentSessionKind.SANDBOX, + }) + ); + }); + + it('uses a valid selected thread agent preference for future run admission', async () => { + const result = await resolve({ + thread: { + metadata: { selectedAgentDefinitionId: 'system.debug' }, + }, + source: { + input: { buildUuid: 'build-1' }, + }, + }); + + expect(result.runPlanSnapshot.agent.id).toBe('system.debug'); + expect(mockGetSystemAgentDefinition).toHaveBeenCalledWith('system.debug'); + }); + + it('resolves selected owned custom agents and snapshots definition details for run admission', async () => { + const result = await resolve({ + thread: { + metadata: { selectedAgentDefinitionId: 'custom.sample-agent' }, + }, + }); + + expect(mockGetUserDefinition).toHaveBeenCalledWith('custom.sample-agent', 'sample-user'); + expect(result.runPlanSnapshot.agent).toEqual( + expect.objectContaining({ + id: 'custom.sample-agent', + label: 'Sample custom agent', + ownerKind: 'user', + version: 3, + modelPreference: { + provider: 'openai', + model: 'gpt-5.4', + }, + resourcePolicy: expect.objectContaining({ + sourceKinds: ['freeform_chat'], + }), + }) + ); + expect(result.runPlanSnapshot.prompt).toEqual( + expect.objectContaining({ + instructionAddendum: 'Use the sample custom instructions.', + renderedSummary: 'Custom agent description', + }) + ); + expect(result.runPlanSnapshot.capabilities.provisionalCapabilityIds).toEqual(['read_context']); + expect(result.runPlanSnapshot.source).toEqual( + expect.objectContaining({ + adapter: 'blank_workspace', + sessionKind: AgentSessionKind.CHAT, + buildUuid: null, + }) + ); + expect(serializeRunPlanSummary(result.runPlanSnapshot)?.agent).toEqual( + expect.objectContaining({ + id: 'custom.sample-agent', + label: 'Sample custom agent', + }) + ); + }); + + it('never loads another user custom id across ownership and falls back to the source default', async () => { + mockGetUserDefinition.mockRejectedValueOnce(new CustomAgentDefinitionServiceError('not_found', 'Agent not found.')); + + const result = await resolve({ + thread: { + metadata: { selectedAgentDefinitionId: 'custom.another-user-agent' }, + }, + }); + + expect(mockGetUserDefinition).toHaveBeenCalledWith('custom.another-user-agent', 'sample-user'); + expect(result.runPlanSnapshot.agent.id).toBe('system.freeform'); + expect(result.runPlanSnapshot.warnings).toEqual( + expect.arrayContaining([ + expect.objectContaining({ + code: 'selected_agent_unavailable', + }), + ]) + ); + }); + + it('falls back from archived selected custom agents with selected_agent_unavailable warning', async () => { + mockGetUserDefinition.mockRejectedValueOnce(new CustomAgentDefinitionServiceError('not_found', 'Agent not found.')); + + const result = await resolve({ + thread: { + metadata: { selectedAgentDefinitionId: 'custom.archived-agent' }, + }, + }); + + expect(result.runPlanSnapshot.agent.id).toBe('system.freeform'); + expect(result.runPlanSnapshot.warnings).toEqual( + expect.arrayContaining([ + expect.objectContaining({ + code: 'selected_agent_unavailable', + message: expect.not.stringContaining('custom.archived-agent'), + }), + ]) + ); + }); + + it('fails closed when selected custom-agent lookup has an unexpected error', async () => { + mockGetUserDefinition.mockRejectedValueOnce(new Error('database unavailable')); + + await expect( + resolve({ + thread: { + metadata: { selectedAgentDefinitionId: 'custom.sample-agent' }, + }, + }) + ).rejects.toThrow('database unavailable'); + + expect(mockGetSystemAgentDefinition).not.toHaveBeenCalledWith('system.freeform'); + }); + + it('rejects restricted required capability refs on custom definitions before admission fields are returned', async () => { + mockResolveSessionContext.mockResolvedValueOnce({ + repoFullName: 'example-org/example-repo', + approvalPolicy: { defaultMode: 'require_approval', rules: {} }, + capabilityPolicy: { + availability: { + external_mcp_write: 'admin_only', + }, + }, + }); + mockGetUserDefinition.mockResolvedValueOnce({ + ...customDefinition, + capabilityRefs: ['external_mcp_write'], + requiredCapabilityRefs: ['external_mcp_write'], + optionalCapabilityRefs: [], + }); + + await expect( + resolve({ + thread: { + metadata: { selectedAgentDefinitionId: 'custom.sample-agent' }, + }, + }) + ).rejects.toThrow(AgentRunPlanCapabilityUnavailableError); + }); + + it('rejects creator-reserved required capability refs on user custom definitions before admission fields are returned', async () => { + mockResolveSessionContext.mockResolvedValueOnce({ + repoFullName: 'example-org/example-repo', + approvalPolicy: { defaultMode: 'require_approval', rules: {} }, + capabilityPolicy: { + availability: { + read_context: 'all_users', + }, + }, + customAgentCreationPolicy: { + capabilityAvailability: { + read_context: 'reserved', + }, + }, + }); + mockGetUserDefinition.mockResolvedValueOnce({ + ...customDefinition, + capabilityRefs: ['read_context'], + requiredCapabilityRefs: ['read_context'], + optionalCapabilityRefs: [], + }); + + await expect( + resolve({ + thread: { + metadata: { selectedAgentDefinitionId: 'custom.sample-agent' }, + }, + }) + ).rejects.toMatchObject({ + capabilityId: 'read_context', + reason: 'creator_capability_reserved', + }); + }); + + it('skips optional custom capabilities that become unavailable at runtime', async () => { + mockResolveSessionContext.mockResolvedValueOnce({ + repoFullName: 'example-org/example-repo', + approvalPolicy: { defaultMode: 'require_approval', rules: {} }, + capabilityPolicy: { + availability: { + external_mcp_write: 'disabled', + }, + }, + }); + mockGetUserDefinition.mockResolvedValueOnce({ + ...customDefinition, + capabilityRefs: ['read_context', 'external_mcp_write'], + requiredCapabilityRefs: [], + optionalCapabilityRefs: ['read_context', 'external_mcp_write'], + }); + + const result = await resolve({ + thread: { + metadata: { selectedAgentDefinitionId: 'custom.sample-agent' }, + }, + }); + + expect(result.runPlanSnapshot.capabilities.provisionalCapabilityIds).toEqual(['read_context']); + expect(result.runPlanSnapshot.warnings).toEqual( + expect.arrayContaining([ + expect.objectContaining({ + code: 'optional_capability_unavailable', + }), + ]) + ); + }); + + it('skips creator-reserved optional custom capabilities with sanitized warnings', async () => { + mockResolveSessionContext.mockResolvedValueOnce({ + repoFullName: 'example-org/example-repo', + approvalPolicy: { defaultMode: 'require_approval', rules: {} }, + capabilityPolicy: { + availability: { + external_mcp_read: 'all_users', + }, + }, + customAgentCreationPolicy: { + capabilityAvailability: { + external_mcp_read: 'reserved', + }, + }, + }); + mockGetUserDefinition.mockResolvedValueOnce({ + ...customDefinition, + capabilityRefs: ['read_context', 'external_mcp_read'], + requiredCapabilityRefs: ['read_context'], + optionalCapabilityRefs: ['external_mcp_read'], + }); + + const result = await resolve({ + thread: { + metadata: { selectedAgentDefinitionId: 'custom.sample-agent' }, + }, + }); + + expect(result.runPlanSnapshot.capabilities.provisionalCapabilityIds).toEqual(['read_context']); + expect(result.runPlanSnapshot.warnings).toEqual( + expect.arrayContaining([ + expect.objectContaining({ + code: 'optional_capability_unavailable', + message: 'MCP read is unavailable and was skipped.', + detail: { + reason: 'creator_capability_reserved', + }, + }), + ]) + ); + expect(JSON.stringify(serializeRunPlanSummary(result.runPlanSnapshot))).not.toContain('external_mcp_read'); + }); + + it('blocks Develop when the selected agent lacks prepared workspace resources', async () => { + await expect( + resolve({ + thread: { + metadata: { selectedAgentDefinitionId: 'system.develop' }, + }, + }) + ).rejects.toThrow(AgentRunPlanAgentUnavailableError); + }); + + it('allows Develop when a chat session workspace runtime is ready', async () => { + const result = await resolve({ + session: { + workspaceStatus: AgentWorkspaceStatus.READY, + podName: 'agent-session-pod', + pvcName: 'agent-session-pvc', + }, + thread: { + metadata: { selectedAgentDefinitionId: 'system.develop' }, + }, + }); + + expect(result.runPlanSnapshot.agent.id).toBe('system.develop'); + expect(result.runPlanSnapshot.agent.sourceKind).toBe('workspace_session'); + expect(result.runPlanSnapshot.capabilities.provisionalCapabilityIds).toEqual( + expect.arrayContaining(['workspace_files', 'workspace_shell', 'workspace_git']) + ); + }); + + it('stores compact repo and service summaries instead of full arrays', async () => { + const result = await resolve({ + session: { + workspaceRepos: [ + { repo: 'example-org/secondary-repo', branch: 'main' }, + { repo: 'example-org/primary-repo', branch: 'feature-branch', primary: true }, + ], + selectedServices: [ + { name: 'sample-service', repo: 'example-org/primary-repo', branch: 'feature-branch' }, + { name: 'sample-worker', repo: 'example-org/primary-repo', branch: 'feature-branch' }, + ], + }, + }); + + expect(result.runPlanSnapshot.source.workspaceLayout).toEqual( + expect.objectContaining({ + repoCount: 2, + primaryRepo: 'example-org/primary-repo', + selectedServiceCount: 2, + primaryService: 'sample-service', + }) + ); + expect(JSON.stringify(result.runPlanSnapshot.source)).not.toContain('sample-worker'); + }); + + it('warns and continues when the session harness default is not supported', async () => { + const result = await resolve({ + session: { + defaultHarness: 'legacy_harness', + }, + }); + + expect(result.resolvedHarness).toBe('lifecycle_ai_sdk'); + expect(result.runPlanSnapshot.warnings).toEqual([ + expect.objectContaining({ + code: 'unsupported_harness_default', + }), + ]); + expect(mockWarn).toHaveBeenCalled(); + }); + + it('rejects unavailable capability refs before admission fields are returned', async () => { + mockResolveSessionContext.mockResolvedValueOnce({ + repoFullName: 'example-org/example-repo', + approvalPolicy: { defaultMode: 'require_approval', rules: {} }, + capabilityPolicy: { + availability: { + diagnostics_codefresh: 'disabled', + }, + }, + }); + + await expect( + resolve({ + source: { + input: { buildUuid: 'build-1' }, + }, + }) + ).rejects.toThrow(AgentRunPlanCapabilityUnavailableError); + }); + + it('blocks disabled selected agents before admission fields are returned', async () => { + mockGetSystemAgentDefinition.mockResolvedValueOnce({ + ...SYSTEM_AGENT_DEFINITIONS['system.freeform'], + status: 'disabled', + }); + + await expect(resolve()).rejects.toThrow(AgentRunPlanAgentUnavailableError); + }); + + it('resolves policy and model selection from current services', async () => { + const result = await AgentRunPlanResolver.resolveForRunAdmission({ + thread: { id: 7, uuid: 'thread-1' } as any, + session: buildSession(), + source: buildSource(), + userIdentity, + requestedProvider: 'openai', + requestedModel: 'gpt-5.4', + runtimeOptions: {}, + }); + + expect(mockResolveSessionContext).toHaveBeenCalledWith('session-1', userIdentity); + expect(mockResolveSelection).toHaveBeenCalledWith({ + repoFullName: 'example-org/example-repo', + requestedProvider: 'openai', + requestedModelId: 'gpt-5.4', + }); + expect(result.runPlanSnapshot.prompt.renderedHash).toEqual(expect.any(String)); + expect(JSON.stringify(result.runPlanSnapshot.prompt)).not.toContain('DB prompt as stored'); + }); + + it('uses durable session default provider when a run omits provider', async () => { + await AgentRunPlanResolver.resolveForRunAdmission({ + thread: { id: 7, uuid: 'thread-1' } as any, + session: buildSession({ defaultModel: 'shared-model' }), + source: buildSource({ + input: { + defaults: { + provider: 'sample-provider-b', + model: 'shared-model', + }, + }, + }), + userIdentity, + requestedProvider: null, + requestedModel: null, + runtimeOptions: {}, + }); + + expect(mockResolveSelection).toHaveBeenCalledWith({ + repoFullName: 'example-org/example-repo', + requestedProvider: 'sample-provider-b', + requestedModelId: 'shared-model', + }); + }); + + it('snapshots explicit empty runtime choices without removing required capabilities', async () => { + mockGetUserDefinition.mockResolvedValueOnce({ + ...customDefinition, + capabilityRefs: ['read_context', 'external_mcp_read'], + requiredCapabilityRefs: ['read_context'], + optionalCapabilityRefs: ['external_mcp_read'], + }); + mockResolveRunAdmissionChoices.mockResolvedValueOnce({ + metadataPresent: true, + selectedRuntimeToolChoiceIds: ['choice-required-read-context'], + selectedRuntimeMcpChoiceIds: [], + selectedRuntimeCapabilityIds: ['read_context'], + selectedRuntimeMcpConnectionRefs: [], + }); + + const result = await resolve({ + thread: { + metadata: { + selectedAgentDefinitionId: 'custom.sample-agent', + runtimeControlChoices: { + version: 1, + toolChoiceIds: [], + mcpChoiceIds: [], + }, + }, + }, + }); + + expect(mockResolveRunAdmissionChoices).toHaveBeenCalledWith( + expect.objectContaining({ + thread: expect.objectContaining({ uuid: 'thread-1' }), + definition: expect.objectContaining({ id: 'custom.sample-agent' }), + sourceKind: 'freeform_chat', + }) + ); + expect(result.runPlanSnapshot.capabilities.provisionalCapabilityIds).toEqual(['read_context']); + expect(result.runPlanSnapshot.capabilities.selectedRuntimeToolChoiceIds).toEqual(['choice-required-read-context']); + expect(result.runPlanSnapshot.capabilities.selectedRuntimeCapabilityIds).toEqual(['read_context']); + expect(result.runPlanSnapshot.capabilities.selectedRuntimeMcpChoiceIds).toEqual([]); + expect(result.runPlanSnapshot.capabilities.selectedRuntimeMcpConnectionRefs).toEqual([]); + }); + + it('adds selected optional runtime tool capabilities to the next run snapshot', async () => { + mockGetUserDefinition.mockResolvedValueOnce({ + ...customDefinition, + capabilityRefs: ['read_context', 'external_mcp_read'], + requiredCapabilityRefs: ['read_context'], + optionalCapabilityRefs: ['external_mcp_read'], + }); + mockResolveRunAdmissionChoices.mockResolvedValueOnce({ + metadataPresent: true, + selectedRuntimeToolChoiceIds: ['choice-required-read-context', 'choice-mcp-read'], + selectedRuntimeMcpChoiceIds: ['choice-mcp-sample'], + selectedRuntimeCapabilityIds: ['read_context', 'external_mcp_read'], + selectedRuntimeMcpConnectionRefs: ['user:sample-mcp'], + }); + + const result = await resolve({ + thread: { + metadata: { + selectedAgentDefinitionId: 'custom.sample-agent', + runtimeControlChoices: { + version: 1, + toolChoiceIds: ['choice-mcp-read'], + mcpChoiceIds: ['choice-mcp-sample'], + }, + }, + }, + }); + + expect(result.runPlanSnapshot.capabilities.provisionalCapabilityIds).toEqual(['read_context', 'external_mcp_read']); + expect(result.runPlanSnapshot.capabilities.resolvedCapabilityAccess).toEqual( + expect.arrayContaining([ + expect.objectContaining({ + capabilityId: 'external_mcp_read', + allowed: true, + }), + ]) + ); + expect(result.runPlanSnapshot.capabilities.selectedRuntimeCapabilityIds).toEqual([ + 'read_context', + 'external_mcp_read', + ]); + expect(result.runPlanSnapshot.capabilities.selectedRuntimeMcpConnectionRefs).toEqual(['user:sample-mcp']); + }); + + it('preserves current optional capability and MCP behavior when runtime metadata is absent', async () => { + mockGetUserDefinition.mockResolvedValueOnce({ + ...customDefinition, + capabilityRefs: ['read_context', 'external_mcp_read'], + requiredCapabilityRefs: ['read_context'], + optionalCapabilityRefs: ['external_mcp_read'], + }); + + const result = await resolve({ + thread: { + metadata: { selectedAgentDefinitionId: 'custom.sample-agent' }, + }, + }); + + expect(result.runPlanSnapshot.capabilities.provisionalCapabilityIds).toEqual(['read_context', 'external_mcp_read']); + expect(result.runPlanSnapshot.capabilities.selectedRuntimeToolChoiceIds).toBeUndefined(); + expect(result.runPlanSnapshot.capabilities.selectedRuntimeMcpChoiceIds).toBeUndefined(); + expect(result.runPlanSnapshot.capabilities.selectedRuntimeMcpConnectionRefs).toBeUndefined(); + }); + + it('omits policy-denied selected optional choices with sanitized warnings', async () => { + mockResolveSessionContext.mockResolvedValueOnce({ + repoFullName: 'example-org/example-repo', + approvalPolicy: { defaultMode: 'require_approval', rules: {} }, + capabilityPolicy: { + availability: { + external_mcp_write: 'disabled', + }, + }, + }); + mockGetUserDefinition.mockResolvedValueOnce({ + ...customDefinition, + capabilityRefs: ['read_context', 'external_mcp_write'], + requiredCapabilityRefs: ['read_context'], + optionalCapabilityRefs: ['external_mcp_write'], + }); + mockResolveRunAdmissionChoices.mockResolvedValueOnce({ + metadataPresent: true, + selectedRuntimeToolChoiceIds: ['choice-required-read-context', 'choice-mcp-write'], + selectedRuntimeMcpChoiceIds: [], + selectedRuntimeCapabilityIds: ['read_context', 'external_mcp_write'], + selectedRuntimeMcpConnectionRefs: [], + }); + + const result = await resolve({ + thread: { + metadata: { + selectedAgentDefinitionId: 'custom.sample-agent', + runtimeControlChoices: { + version: 1, + toolChoiceIds: ['choice-mcp-write'], + mcpChoiceIds: [], + }, + }, + }, + }); + + expect(result.runPlanSnapshot.capabilities.provisionalCapabilityIds).toEqual(['read_context']); + expect(result.runPlanSnapshot.warnings).toEqual( + expect.arrayContaining([ + expect.objectContaining({ + code: 'optional_capability_unavailable', + message: expect.stringContaining('MCP write'), + }), + ]) + ); + expect(JSON.stringify(serializeRunPlanSummary(result.runPlanSnapshot))).not.toContain('external_mcp_write'); + expect(JSON.stringify(serializeRunPlanSummary(result.runPlanSnapshot))).not.toContain('choice-mcp-write'); + expect(result.runPlanSnapshot.capabilities.resolvedCapabilityAccess).not.toEqual( + expect.arrayContaining([ + expect.objectContaining({ + capabilityId: 'external_mcp_write', + }), + ]) + ); + expect(result.runPlanSnapshot.capabilities.selectedRuntimeToolChoiceIds).toEqual([]); + }); + + it('omits creator-reserved selected optional MCP choices with sanitized warnings', async () => { + mockResolveSessionContext.mockResolvedValueOnce({ + repoFullName: 'example-org/example-repo', + approvalPolicy: { defaultMode: 'require_approval', rules: {} }, + customAgentCreationPolicy: { + capabilityAvailability: { + external_mcp_read: 'reserved', + }, + }, + }); + mockGetUserDefinition.mockResolvedValueOnce({ + ...customDefinition, + capabilityRefs: ['read_context', 'external_mcp_read'], + requiredCapabilityRefs: ['read_context'], + optionalCapabilityRefs: ['external_mcp_read'], + }); + mockResolveRunAdmissionChoices.mockResolvedValueOnce({ + metadataPresent: true, + selectedRuntimeToolChoiceIds: ['choice-required-read-context'], + selectedRuntimeMcpChoiceIds: ['choice-mcp-sample'], + selectedRuntimeCapabilityIds: ['read_context', 'external_mcp_read'], + selectedRuntimeMcpConnectionRefs: ['user:sample-mcp'], + }); + + const result = await resolve({ + thread: { + metadata: { + selectedAgentDefinitionId: 'custom.sample-agent', + runtimeControlChoices: { + version: 1, + toolChoiceIds: [], + mcpChoiceIds: ['choice-mcp-sample'], + }, + }, + }, + }); + + expect(result.runPlanSnapshot.capabilities.provisionalCapabilityIds).toEqual(['read_context']); + expect(result.runPlanSnapshot.warnings).toEqual( + expect.arrayContaining([ + expect.objectContaining({ + code: 'optional_capability_unavailable', + detail: { + reason: 'creator_capability_reserved', + }, + }), + ]) + ); + expect(JSON.stringify(serializeRunPlanSummary(result.runPlanSnapshot))).not.toContain('external_mcp_read'); + expect(JSON.stringify(serializeRunPlanSummary(result.runPlanSnapshot))).not.toContain('choice-mcp-sample'); + expect(result.runPlanSnapshot.capabilities.selectedRuntimeToolChoiceIds).toEqual([]); + expect(result.runPlanSnapshot.capabilities.selectedRuntimeMcpChoiceIds).toEqual([]); + expect(result.runPlanSnapshot.capabilities.selectedRuntimeMcpConnectionRefs).toEqual([]); + }); +}); diff --git a/src/server/services/agent/__tests__/RunService.test.ts b/src/server/services/agent/__tests__/RunService.test.ts index 99f28f46..64de83a0 100644 --- a/src/server/services/agent/__tests__/RunService.test.ts +++ b/src/server/services/agent/__tests__/RunService.test.ts @@ -73,6 +73,72 @@ const mockAppendChunkEventsForRunInTransaction = AgentRunEventService.appendChun const mockNotifyRunEventsInserted = AgentRunEventService.notifyRunEventsInserted as jest.Mock; const mockResolveDurabilityConfig = resolveAgentSessionDurabilityConfig as jest.Mock; const VALID_RUN_UUID = '123e4567-e89b-12d3-a456-426614174000'; +const runPlanSnapshot = { + version: 1, + capturedAt: '2026-05-01T00:00:00.000Z', + agent: { + id: 'system.freeform', + label: 'Free-form', + sourceKind: 'freeform_chat', + }, + source: { + id: 'source-1', + adapter: 'blank_workspace', + status: 'ready', + sessionKind: 'chat', + buildUuid: 'build-1', + repoFullName: 'example-org/example-repo', + branch: 'feature-branch', + namespace: 'sample-namespace', + freshness: { + capturedAt: '2026-05-01T00:00:00.000Z', + freshnessSource: 'source', + }, + }, + model: { + requestedProvider: null, + requestedModel: null, + resolvedProvider: 'openai', + resolvedModel: 'gpt-5.4', + }, + runtime: { + requestedHarness: null, + resolvedHarness: 'lifecycle_ai_sdk', + sandboxRequirement: { filesystem: 'persistent' }, + runtimeOptions: { maxIterations: 12 }, + approvalPolicy: { defaultMode: 'require_approval', rules: {} }, + }, + prompt: { + instructionRefs: [], + renderedSummary: 'Sample prompt summary', + renderedHash: 'sha256:sample-rendered-prompt', + }, + capabilities: { + provisionalCapabilityIds: ['read_context', 'workspace_files'], + resolvedCapabilityAccess: [ + { + capabilityId: 'read_context', + availability: 'all_users', + allowed: true, + runtimeCapabilityKey: 'read', + approvalMode: 'allow', + }, + { + capabilityId: 'workspace_files', + availability: 'admin_only', + allowed: false, + reason: 'sample denied reason', + runtimeCapabilityKey: 'workspace_write', + approvalMode: 'require_approval', + }, + ], + selectedRuntimeToolChoiceIds: ['choice-read-context'], + selectedRuntimeMcpChoiceIds: ['choice-sample-mcp'], + selectedRuntimeCapabilityIds: ['read_context'], + selectedRuntimeMcpConnectionRefs: ['user:sample-mcp'], + }, + warnings: [{ code: 'sample_warning', message: 'Sample warning', detail: { hidden: true } }], +} as const; describe('AgentRunService', () => { beforeEach(() => { @@ -147,6 +213,152 @@ describe('AgentRunService', () => { }); }); + describe('serializeRun', () => { + const baseRun = { + uuid: VALID_RUN_UUID, + threadId: 7, + sessionId: 17, + status: 'queued', + requestedHarness: null, + resolvedHarness: 'lifecycle_ai_sdk', + requestedProvider: null, + requestedModel: null, + resolvedProvider: 'openai', + resolvedModel: 'gpt-5.4', + provider: 'openai', + model: 'gpt-5.4', + sandboxRequirement: {}, + sandboxGeneration: null, + queuedAt: '2026-05-01T00:00:00.000Z', + startedAt: null, + completedAt: null, + cancelledAt: null, + usageSummary: {}, + policySnapshot: { defaultMode: 'require_approval', rules: {} }, + error: null, + createdAt: '2026-05-01T00:00:00.000Z', + updatedAt: '2026-05-01T00:00:00.000Z', + }; + + it('returns runPlan: null for historical runs without snapshots', () => { + expect(AgentRunService.serializeRun({ ...baseRun, runPlanSnapshot: null } as any)).toEqual( + expect.objectContaining({ + runPlan: null, + }) + ); + }); + + it('returns a safe runPlan summary for versioned snapshots', () => { + const serialized = AgentRunService.serializeRun({ ...baseRun, runPlanSnapshot } as any); + + expect(serialized).not.toHaveProperty('runPlanSnapshot'); + expect(serialized.runPlan).toEqual({ + version: 1, + agent: { + id: 'system.freeform', + label: 'Free-form', + sourceKind: 'freeform_chat', + }, + source: { + kind: 'freeform_chat', + repoFullName: 'example-org/example-repo', + branch: 'feature-branch', + buildUuid: 'build-1', + namespace: 'sample-namespace', + }, + model: { + provider: 'openai', + model: 'gpt-5.4', + }, + runtime: { + harness: 'lifecycle_ai_sdk', + maxIterations: 12, + }, + approval: { + defaultMode: 'require_approval', + }, + capabilities: { + effective: [ + { + capabilityId: 'read_context', + availability: 'all_users', + allowed: true, + approvalMode: 'allow', + }, + { + capabilityId: 'workspace_files', + availability: 'admin_only', + allowed: false, + approvalMode: 'require_approval', + }, + ], + selected: { + capabilityIds: ['read_context'], + toolChoiceIds: ['choice-read-context'], + mcpChoiceIds: ['choice-sample-mcp'], + }, + }, + warnings: [{ code: 'sample_warning', message: 'Sample warning' }], + }); + const runPlanJson = JSON.stringify(serialized.runPlan); + for (const forbidden of [ + 'renderedHash', + 'renderedSummary', + 'sha256:sample-rendered-prompt', + 'selectedRuntimeMcpConnectionRefs', + 'runtimeCapabilityKey', + 'workspace.read_file', + 'sample denied reason', + 'rules', + ]) { + expect(runPlanJson).not.toContain(forbidden); + } + }); + + it('defaults missing selected runtime choice arrays to empty arrays', () => { + const snapshotWithoutSelections = { + ...runPlanSnapshot, + capabilities: { + ...runPlanSnapshot.capabilities, + selectedRuntimeToolChoiceIds: undefined, + selectedRuntimeMcpChoiceIds: undefined, + selectedRuntimeCapabilityIds: undefined, + selectedRuntimeMcpConnectionRefs: undefined, + }, + }; + + const serialized = AgentRunService.serializeRun({ + ...baseRun, + runPlanSnapshot: snapshotWithoutSelections, + } as any); + + expect(serialized.runPlan?.capabilities.selected).toEqual({ + capabilityIds: [], + toolChoiceIds: [], + mcpChoiceIds: [], + }); + }); + + it('returns runPlan: null for snapshots missing resolved capability access', () => { + const incompleteSnapshot = { + ...runPlanSnapshot, + capabilities: { + provisionalCapabilityIds: runPlanSnapshot.capabilities.provisionalCapabilityIds, + selectedRuntimeToolChoiceIds: ['choice-read-context'], + selectedRuntimeMcpChoiceIds: ['choice-sample-mcp'], + selectedRuntimeMcpConnectionRefs: ['user:sample-mcp'], + }, + }; + + const serialized = AgentRunService.serializeRun({ + ...baseRun, + runPlanSnapshot: incompleteSnapshot, + } as any); + + expect(serialized.runPlan).toBeNull(); + }); + }); + describe('listRunsNeedingDispatch', () => { it('finds stale queued runs and expired execution leases', async () => { const staleQueuedBuilder: any = { diff --git a/src/server/services/agent/__tests__/SessionReadService.test.ts b/src/server/services/agent/__tests__/SessionReadService.test.ts index a8c4e440..b24e9e61 100644 --- a/src/server/services/agent/__tests__/SessionReadService.test.ts +++ b/src/server/services/agent/__tests__/SessionReadService.test.ts @@ -74,6 +74,23 @@ jest.mock('server/services/agent/ThreadService', () => ({ }, })); +jest.mock('server/services/agent/AgentUsageService', () => ({ + __esModule: true, + default: { + aggregateRuns: jest.fn(() => ({ + usageSummary: { totalTokens: 0 }, + usageByModel: [], + usageCompleteness: { + runCount: 0, + reportedRunCount: 0, + missingUsageRunCount: 0, + complete: true, + }, + })), + aggregateSessionsUsage: jest.fn(), + }, +})); + jest.mock('server/lib/dependencies', () => ({})); import AgentSession from 'server/models/AgentSession'; @@ -81,6 +98,7 @@ 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 AgentUsageService from 'server/services/agent/AgentUsageService'; import AgentSessionReadService from '../SessionReadService'; const mockSessionQuery = AgentSession.query as jest.Mock; @@ -88,6 +106,7 @@ 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; +const mockAggregateSessionsUsage = AgentUsageService.aggregateSessionsUsage as jest.Mock; function buildSession(overrides: Record = {}) { return { @@ -146,6 +165,23 @@ function buildOrderedQuery(rows: T[], orderCalls = 1) { describe('AgentSessionReadService', () => { beforeEach(() => { jest.clearAllMocks(); + mockAggregateSessionsUsage.mockResolvedValue( + new Map([ + [ + 17, + { + usageSummary: { totalTokens: 0 }, + usageByModel: [], + usageCompleteness: { + runCount: 0, + reportedRunCount: 0, + missingUsageRunCount: 0, + complete: true, + }, + }, + ], + ]) + ); }); it('lists owned sessions with capped pagination and batched related reads', async () => { @@ -156,7 +192,12 @@ describe('AgentSessionReadService', () => { sessionId: 17, adapter: 'lifecycle_environment', status: 'ready', - input: {}, + input: { + defaults: { + provider: 'sample-provider', + model: 'gpt-5.4', + }, + }, sandboxRequirements: { filesystem: 'persistent' }, error: null, preparedAt: '2026-04-24T12:00:00.000Z', @@ -219,7 +260,18 @@ describe('AgentSessionReadService', () => { limit: 100, }); expect(result.records).toHaveLength(1); + expect(result.records[0].session.defaults.provider).toBe('sample-provider'); expect(result.records[0].session.defaultThreadId).toBe('thread-1'); + expect(result.records[0].usage).toEqual({ + usageSummary: { totalTokens: 0 }, + usageByModel: [], + usageCompleteness: { + runCount: 0, + reportedRunCount: 0, + missingUsageRunCount: 0, + complete: true, + }, + }); expect(result.records[0].source.id).toBe('source-1'); expect(result.records[0].sandbox.exposures).toEqual([ expect.objectContaining({ @@ -230,5 +282,91 @@ describe('AgentSessionReadService', () => { expect(sourceQuery.whereIn).toHaveBeenCalledWith('sessionId', [17]); expect(sandboxQuery.whereIn).toHaveBeenCalledWith('sessionId', [17]); expect(exposureQuery.whereIn).toHaveBeenCalledWith('sandboxId', [4]); + expect(mockAggregateSessionsUsage).toHaveBeenCalledTimes(1); + expect(mockAggregateSessionsUsage).toHaveBeenCalledWith([17]); + }); + + it('keeps chat sessions ready when an on-demand workspace sandbox failed', async () => { + const session = buildSession({ + sessionKind: 'chat', + workspaceStatus: 'failed', + }); + const source = { + id: 3, + uuid: 'source-1', + sessionId: 17, + adapter: 'blank_workspace', + status: 'ready', + input: { + defaults: { + provider: 'sample-provider', + model: 'gpt-5.4', + }, + sessionKind: 'chat', + }, + 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: 'failed', + capabilitySnapshot: {}, + suspendedAt: null, + endedAt: null, + error: { message: 'Sandbox failed' }, + createdAt: '2026-04-24T12:00:00.000Z', + updatedAt: '2026-04-24T12:00:00.000Z', + }; + const defaultThread = { id: 9, uuid: 'thread-1', sessionId: 17 }; + + mockSourceQuery.mockReturnValueOnce({ whereIn: jest.fn().mockResolvedValue([source]) }); + mockSandboxQuery.mockReturnValueOnce(buildOrderedQuery([sandbox], 2)); + mockThreadQuery.mockReturnValueOnce({ whereIn: jest.fn().mockResolvedValue([defaultThread]) }); + mockThreadQuery.mockReturnValueOnce(buildOrderedQuery([], 1)); + mockSandboxExposureQuery.mockReturnValueOnce(buildOrderedQuery([], 1)); + + const [record] = await AgentSessionReadService.listSessionRecords([session] as any); + + expect(record.session.status).toBe('ready'); + expect(record.sandbox.status).toBe('failed'); + }); + + it('marks non-chat failed workspaces unavailable', async () => { + const session = buildSession({ + sessionKind: 'environment', + workspaceStatus: 'failed', + }); + 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 defaultThread = { id: 9, uuid: 'thread-1', sessionId: 17 }; + + mockSourceQuery.mockReturnValueOnce({ whereIn: jest.fn().mockResolvedValue([source]) }); + mockSandboxQuery.mockReturnValueOnce(buildOrderedQuery([], 2)); + mockThreadQuery.mockReturnValueOnce({ whereIn: jest.fn().mockResolvedValue([defaultThread]) }); + mockThreadQuery.mockReturnValueOnce(buildOrderedQuery([], 1)); + + const [record] = await AgentSessionReadService.listSessionRecords([session] as any); + + expect(record.session.status).toBe('error'); }); }); diff --git a/src/server/services/agent/__tests__/SettingsService.test.ts b/src/server/services/agent/__tests__/SettingsService.test.ts index 6c964920..3c4553f5 100644 --- a/src/server/services/agent/__tests__/SettingsService.test.ts +++ b/src/server/services/agent/__tests__/SettingsService.test.ts @@ -14,7 +14,7 @@ * limitations under the License. */ -jest.mock('server/services/aiAgentConfig', () => ({ +jest.mock('server/services/agentRuntime/config/agentRuntimeConfig', () => ({ __esModule: true, default: { getInstance: jest.fn(() => ({ @@ -40,7 +40,7 @@ jest.mock('server/services/userApiKey', () => ({ const mockListEnabledConnectionsForUser = jest.fn(); -jest.mock('server/services/ai/mcp/config', () => ({ +jest.mock('server/services/agentRuntime/mcp/config', () => ({ __esModule: true, McpConfigService: jest.fn().mockImplementation(() => ({ listEnabledConnectionsForUser: (...args: unknown[]) => mockListEnabledConnectionsForUser(...args), diff --git a/src/server/services/agent/__tests__/SourceService.test.ts b/src/server/services/agent/__tests__/SourceService.test.ts index 283d0fe5..1aa9c9fb 100644 --- a/src/server/services/agent/__tests__/SourceService.test.ts +++ b/src/server/services/agent/__tests__/SourceService.test.ts @@ -37,18 +37,25 @@ describe('AgentSourceService', () => { 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]); + 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, + defaultModel: 'sample-model', + model: 'sample-model', + } as Parameters[0], + { + defaultProvider: 'sample-provider', + } + ); expect(insertAndFetch).toHaveBeenCalledWith( expect.objectContaining({ @@ -59,6 +66,10 @@ describe('AgentSourceService', () => { buildUuid: 'build-1', buildKind: 'environment', sessionKind: 'environment', + defaults: { + provider: 'sample-provider', + model: 'sample-model', + }, }, sandboxRequirements: expect.objectContaining({ filesystem: 'persistent', diff --git a/src/server/services/agent/__tests__/ThreadRuntimeControlsService.test.ts b/src/server/services/agent/__tests__/ThreadRuntimeControlsService.test.ts new file mode 100644 index 00000000..2f41811a --- /dev/null +++ b/src/server/services/agent/__tests__/ThreadRuntimeControlsService.test.ts @@ -0,0 +1,662 @@ +/** + * 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. + */ + +const mockGetOwnedThreadWithSession = jest.fn(); +const mockGetSelectedAgentDefinitionId = jest.fn(); +const mockGetRuntimeControlChoices = jest.fn(); +const mockPatchRuntimeControlChoices = jest.fn(); +const mockBuildSelectedAgentDefinitionMetadataPatch = jest.fn(); +const mockGetSessionSource = jest.fn(); +const mockResolveSessionContext = jest.fn(); +const mockHasActiveRun = jest.fn(); +const mockEnsureSeeded = jest.fn(); +const mockGetSystemAgentDefinition = jest.fn(); +const mockInferDefaultAgentDefinitionId = jest.fn(); +const mockListUserDefinitions = jest.fn(); +const mockGetUserDefinition = jest.fn(); +const mockListEnabledConnectionsForUser = jest.fn(); +const mockGetEffectiveConfig = jest.fn(); + +jest.mock('../ThreadService', () => ({ + __esModule: true, + default: { + getOwnedThreadWithSession: (...args: unknown[]) => mockGetOwnedThreadWithSession(...args), + getSelectedAgentDefinitionId: (...args: unknown[]) => mockGetSelectedAgentDefinitionId(...args), + getRuntimeControlChoices: (...args: unknown[]) => mockGetRuntimeControlChoices(...args), + patchRuntimeControlChoices: (...args: unknown[]) => mockPatchRuntimeControlChoices(...args), + buildSelectedAgentDefinitionMetadataPatch: (...args: unknown[]) => + mockBuildSelectedAgentDefinitionMetadataPatch(...args), + }, +})); + +jest.mock('../SourceService', () => ({ + __esModule: true, + default: { + getSessionSource: (...args: unknown[]) => mockGetSessionSource(...args), + }, +})); + +jest.mock('../CapabilityService', () => ({ + __esModule: true, + default: { + resolveSessionContext: (...args: unknown[]) => mockResolveSessionContext(...args), + }, +})); + +jest.mock('../RunService', () => ({ + __esModule: true, + default: { + hasActiveRun: (...args: unknown[]) => mockHasActiveRun(...args), + }, +})); + +jest.mock('../AgentDefinitionRegistry', () => { + const actual = jest.requireActual('../AgentDefinitionRegistry'); + return { + __esModule: true, + ...actual, + ensureSystemAgentDefinitionsSeeded: (...args: unknown[]) => mockEnsureSeeded(...args), + getSystemAgentDefinition: (...args: unknown[]) => mockGetSystemAgentDefinition(...args), + inferDefaultSystemAgentDefinitionId: (...args: unknown[]) => mockInferDefaultAgentDefinitionId(...args), + }; +}); + +jest.mock('../CustomAgentDefinitionService', () => ({ + __esModule: true, + customAgentDefinitionService: { + listUserDefinitions: (...args: unknown[]) => mockListUserDefinitions(...args), + getUserDefinition: (...args: unknown[]) => mockGetUserDefinition(...args), + }, +})); + +jest.mock('server/services/agentRuntime/mcp/config', () => ({ + McpConfigService: jest.fn().mockImplementation(() => ({ + listEnabledConnectionsForUser: (...args: unknown[]) => mockListEnabledConnectionsForUser(...args), + })), +})); + +jest.mock('server/services/agentRuntime/config/agentRuntimeConfig', () => ({ + __esModule: true, + default: { + getInstance: jest.fn(() => ({ + getEffectiveConfig: (...args: unknown[]) => mockGetEffectiveConfig(...args), + })), + }, +})); + +import AgentThreadRuntimeControlsService, { AgentThreadRuntimeControlsError } from '../ThreadRuntimeControlsService'; +import { SYSTEM_AGENT_DEFINITIONS } from '../systemAgentDefinitions'; + +const userIdentity = { + userId: 'sample-user', + githubUsername: 'sample-user', + preferredUsername: 'sample-user', + email: 'sample-user@example.com', + firstName: 'Sample', + lastName: 'User', + displayName: 'Sample User', + gitUserName: 'Sample User', + gitUserEmail: 'sample-user@example.com', + roles: [], +}; + +const session = { + id: 17, + uuid: 'session-1', + sessionKind: 'chat', + buildUuid: null, + workspaceRepos: [{ repo: 'example-org/example-repo', primary: true }], +}; + +const thread = { + id: 23, + uuid: 'thread-1', + metadata: {}, +}; + +const source = { + id: 7, + status: 'ready', + input: {}, +}; + +const customDefinition = { + id: 'custom.sample-agent', + version: 1, + owner: { kind: 'user' as const, userId: 'sample-user' }, + name: 'Sample agent', + description: 'Helps with sample work.', + instructionRefs: [], + instructionAddendum: 'Answer clearly.', + capabilityRefs: ['read_context', 'workspace_files', 'external_mcp_read'], + requiredCapabilityRefs: ['read_context'], + optionalCapabilityRefs: ['workspace_files', 'external_mcp_read'], + resourcePolicy: { + sourceKinds: ['freeform_chat', 'workspace_session'], + workspaceRequired: false, + sandboxRequired: false, + }, + modelPreference: null, + status: 'active' as const, + codeOwned: false, + readOnly: false, +}; + +function mockBaseContext() { + mockGetOwnedThreadWithSession.mockResolvedValue({ thread, session }); + mockGetSelectedAgentDefinitionId.mockReturnValue('custom.sample-agent'); + mockGetRuntimeControlChoices.mockReturnValue(null); + mockGetSessionSource.mockResolvedValue(source); + mockResolveSessionContext.mockResolvedValue({ + repoFullName: 'example-org/example-repo', + approvalPolicy: { defaultMode: 'allow', rules: {} }, + capabilityPolicy: undefined, + }); + mockHasActiveRun.mockResolvedValue(false); + mockEnsureSeeded.mockResolvedValue([]); + mockInferDefaultAgentDefinitionId.mockReturnValue('system.develop'); + mockGetEffectiveConfig.mockResolvedValue({ + approvalPolicy: { defaultMode: 'allow', rules: {} }, + capabilityPolicy: undefined, + }); + mockListUserDefinitions.mockResolvedValue([customDefinition]); + mockGetUserDefinition.mockResolvedValue(customDefinition); + mockBuildSelectedAgentDefinitionMetadataPatch.mockImplementation((agentId: string) => ({ + selectedAgentDefinitionId: agentId, + })); + mockListEnabledConnectionsForUser.mockResolvedValue([ + { + slug: 'sample-mcp', + name: 'Sample MCP', + description: 'Provides sample context.', + scope: 'global', + connectionRequired: false, + configured: true, + stale: false, + discoveredTools: [{ name: 'readSample', annotations: { readOnlyHint: true } }], + }, + ]); +} + +function getOptionalChoiceId(state: Awaited>) { + const optional = state.tools.optional.find((choice) => choice.label === 'Workspace files'); + if (!optional) { + throw new Error('Expected Workspace files optional choice'); + } + return optional.id; +} + +describe('AgentThreadRuntimeControlsService', () => { + beforeEach(() => { + jest.clearAllMocks(); + mockBaseContext(); + }); + + it('returns sanitized opaque tool and MCP state for an existing thread', async () => { + const state = await AgentThreadRuntimeControlsService.getState({ + threadId: 'thread-1', + userIdentity, + }); + + expect(state.canEdit).toBe(true); + expect(state.disabledReason).toBeNull(); + expect(state.tools.required).toEqual([ + expect.objectContaining({ + label: 'Read/context', + required: true, + selected: true, + available: true, + }), + ]); + expect(state.tools.optional).toEqual([ + expect.objectContaining({ + label: 'Workspace files', + required: false, + selected: true, + available: true, + }), + ]); + expect(state.mcp.connections).toEqual([ + expect.objectContaining({ + label: 'Sample MCP', + selected: true, + available: true, + }), + ]); + expect(JSON.stringify(state)).not.toContain('workspace_files'); + expect(JSON.stringify(state)).not.toContain('external_mcp_read'); + expect(JSON.stringify(state)).not.toContain('sample-mcp'); + expect(state.tools.selectedChoiceIds.every((id) => id.startsWith('rtc_'))).toBe(true); + expect(state.mcp.selectedChoiceIds.every((id) => id.startsWith('rtc_'))).toBe(true); + }); + + it('hides creator-reserved optional capabilities from selected runtime choices', async () => { + mockResolveSessionContext.mockResolvedValueOnce({ + repoFullName: 'example-org/example-repo', + approvalPolicy: { defaultMode: 'allow', rules: {} }, + capabilityPolicy: undefined, + customAgentCreationPolicy: { + capabilityAvailability: { + workspace_files: 'reserved', + external_mcp_read: 'reserved', + }, + }, + }); + + const state = await AgentThreadRuntimeControlsService.getState({ + threadId: 'thread-1', + userIdentity, + }); + + expect(state.tools.optional).toEqual([ + expect.objectContaining({ + label: 'Workspace files', + available: false, + }), + ]); + expect(state.tools.selectedChoiceIds).toEqual([state.tools.required[0].id]); + expect(state.mcp.connections).toEqual([]); + expect(state.mcp.selectedChoiceIds).toEqual([]); + }); + + it('persists valid optional choices as opaque ids and leaves raw ids out of metadata', async () => { + const state = await AgentThreadRuntimeControlsService.getState({ + threadId: 'thread-1', + userIdentity, + }); + const optionalChoiceId = getOptionalChoiceId(state); + const mcpChoiceId = state.mcp.connections[0].id; + mockPatchRuntimeControlChoices.mockResolvedValue({ + ...thread, + metadata: {}, + }); + + await AgentThreadRuntimeControlsService.patchChoices({ + threadId: 'thread-1', + userIdentity, + toolChoiceIds: [optionalChoiceId], + mcpChoiceIds: [mcpChoiceId], + }); + + expect(mockPatchRuntimeControlChoices).toHaveBeenCalledWith(23, { + version: 1, + toolChoiceIds: [optionalChoiceId], + mcpChoiceIds: [mcpChoiceId], + }); + expect(JSON.stringify(mockPatchRuntimeControlChoices.mock.calls[0][1])).not.toContain('workspace_files'); + expect(JSON.stringify(mockPatchRuntimeControlChoices.mock.calls[0][1])).not.toContain('sample-mcp'); + }); + + it('preserves current MCP choices when patching only tool choices', async () => { + const state = await AgentThreadRuntimeControlsService.getState({ + threadId: 'thread-1', + userIdentity, + }); + const optionalChoiceId = getOptionalChoiceId(state); + const mcpChoiceId = state.mcp.connections[0].id; + mockPatchRuntimeControlChoices.mockResolvedValue({ + ...thread, + metadata: {}, + }); + + await AgentThreadRuntimeControlsService.patchChoices({ + threadId: 'thread-1', + userIdentity, + toolChoiceIds: [optionalChoiceId], + }); + + expect(mockPatchRuntimeControlChoices).toHaveBeenCalledWith(23, { + version: 1, + toolChoiceIds: [optionalChoiceId], + mcpChoiceIds: [mcpChoiceId], + }); + }); + + it('preserves current tool choices when patching only MCP choices', async () => { + const state = await AgentThreadRuntimeControlsService.getState({ + threadId: 'thread-1', + userIdentity, + }); + const optionalChoiceId = getOptionalChoiceId(state); + const mcpChoiceId = state.mcp.connections[0].id; + mockPatchRuntimeControlChoices.mockResolvedValue({ + ...thread, + metadata: {}, + }); + + await AgentThreadRuntimeControlsService.patchChoices({ + threadId: 'thread-1', + userIdentity, + mcpChoiceIds: [mcpChoiceId], + }); + + expect(mockPatchRuntimeControlChoices).toHaveBeenCalledWith(23, { + version: 1, + toolChoiceIds: [optionalChoiceId], + mcpChoiceIds: [mcpChoiceId], + }); + }); + + it('keeps required tool choices selected when a patch omits them', async () => { + mockPatchRuntimeControlChoices.mockResolvedValue({ + ...thread, + metadata: {}, + }); + + const updatedState = await AgentThreadRuntimeControlsService.patchChoices({ + threadId: 'thread-1', + userIdentity, + toolChoiceIds: [], + mcpChoiceIds: [], + }); + + expect(mockPatchRuntimeControlChoices).toHaveBeenCalledWith(23, { + version: 1, + toolChoiceIds: [], + mcpChoiceIds: [], + }); + expect(updatedState.tools.required).toEqual([ + expect.objectContaining({ + label: 'Read/context', + required: true, + selected: true, + available: true, + }), + ]); + expect(updatedState.tools.selectedChoiceIds).toEqual([updatedState.tools.required[0].id]); + }); + + it('treats shared discovered MCP tools as available runtime choices', async () => { + mockListEnabledConnectionsForUser.mockResolvedValue([ + { + slug: 'shared-sample-mcp', + name: 'Shared Sample MCP', + description: 'Provides shared sample context.', + scope: 'global', + connectionRequired: false, + configured: false, + stale: false, + validationError: null, + discoveredTools: [], + sharedDiscoveredTools: [{ name: 'readSharedSample', annotations: { readOnlyHint: true } }], + }, + ]); + + const state = await AgentThreadRuntimeControlsService.getState({ + threadId: 'thread-1', + userIdentity, + }); + + expect(state.mcp.connections).toEqual([ + expect.objectContaining({ + label: 'Shared Sample MCP', + selected: true, + available: true, + }), + ]); + expect(state.mcp.selectedChoiceIds).toEqual([state.mcp.connections[0].id]); + }); + + it('marks MCP connections with validation errors unavailable and rejects saved choices', async () => { + mockListEnabledConnectionsForUser.mockResolvedValue([ + { + slug: 'broken-sample-mcp', + name: 'Broken Sample MCP', + description: 'Broken sample context.', + scope: 'global', + connectionRequired: false, + configured: true, + stale: false, + validationError: 'Connection failed', + discoveredTools: [{ name: 'readBrokenSample', annotations: { readOnlyHint: true } }], + sharedDiscoveredTools: [], + }, + ]); + + const state = await AgentThreadRuntimeControlsService.getState({ + threadId: 'thread-1', + userIdentity, + }); + const brokenChoiceId = state.mcp.connections[0].id; + + expect(state.mcp.connections).toEqual([ + expect.objectContaining({ + label: 'Broken Sample MCP', + selected: false, + available: false, + }), + ]); + expect(state.mcp.selectedChoiceIds).toEqual([]); + + await expect( + AgentThreadRuntimeControlsService.patchChoices({ + threadId: 'thread-1', + userIdentity, + toolChoiceIds: [], + mcpChoiceIds: [brokenChoiceId], + }) + ).rejects.toMatchObject({ + code: 'policy_denied', + }); + }); + + it('keeps metadata absent until runtime choices are saved', async () => { + await AgentThreadRuntimeControlsService.getState({ + threadId: 'thread-1', + userIdentity, + }); + + expect(mockPatchRuntimeControlChoices).not.toHaveBeenCalled(); + }); + + it('rejects raw, unknown, and policy-denied choices', async () => { + await expect( + AgentThreadRuntimeControlsService.patchChoices({ + threadId: 'thread-1', + userIdentity, + toolChoiceIds: ['workspace_files'], + mcpChoiceIds: [], + }) + ).rejects.toMatchObject({ + code: 'unknown_choice', + }); + + mockResolveSessionContext.mockResolvedValue({ + repoFullName: 'example-org/example-repo', + approvalPolicy: { defaultMode: 'allow', rules: {} }, + capabilityPolicy: { availability: { workspace_files: 'disabled' } }, + }); + const deniedState = await AgentThreadRuntimeControlsService.getState({ + threadId: 'thread-1', + userIdentity, + }); + const deniedChoiceId = getOptionalChoiceId(deniedState); + + await expect( + AgentThreadRuntimeControlsService.patchChoices({ + threadId: 'thread-1', + userIdentity, + toolChoiceIds: [deniedChoiceId], + mcpChoiceIds: [], + }) + ).rejects.toMatchObject({ + code: 'policy_denied', + }); + }); + + it('blocks existing-thread edits while an active run exists', async () => { + mockHasActiveRun.mockResolvedValue(true); + + const state = await AgentThreadRuntimeControlsService.getState({ + threadId: 'thread-1', + userIdentity, + }); + + expect(state.canEdit).toBe(false); + expect(state.disabledReason).toBe('Change after this response finishes.'); + + await expect( + AgentThreadRuntimeControlsService.patchChoices({ + threadId: 'thread-1', + userIdentity, + toolChoiceIds: [], + mcpChoiceIds: [], + }) + ).rejects.toMatchObject({ + code: 'active_run', + message: 'Change after this response finishes.', + }); + }); + + it('returns a sanitized /new preview without requiring a thread id', async () => { + const state = await AgentThreadRuntimeControlsService.getEntryPreview({ + userIdentity, + agentId: 'custom.sample-agent', + source: { adapter: 'lifecycle_fork', input: {} }, + defaults: {}, + }); + + expect(mockGetOwnedThreadWithSession).not.toHaveBeenCalled(); + expect(state.canEdit).toBe(true); + expect(state.tools.optional.map((choice) => choice.label)).toEqual(['Workspace files']); + expect(state.mcp.connections.map((choice) => choice.label)).toEqual(['Sample MCP']); + expect(JSON.stringify(state)).not.toContain('custom.sample-agent'); + expect(JSON.stringify(state)).not.toContain('sample-mcp'); + }); + + it('previews Develop tools for blank chat entry without a prepared workspace yet', async () => { + mockGetSystemAgentDefinition.mockResolvedValueOnce(SYSTEM_AGENT_DEFINITIONS['system.develop']); + + const state = await AgentThreadRuntimeControlsService.getEntryPreview({ + userIdentity, + agentId: 'system.develop', + source: { adapter: 'blank_workspace', input: {} }, + defaults: {}, + }); + + expect(state.tools.required.map((choice) => choice.label)).toEqual( + expect.arrayContaining(['Workspace files', 'Command tools', 'Source control']) + ); + expect(state.canEdit).toBe(true); + }); + + it('validates create-session bootstrap choices and preserves explicit empty arrays', async () => { + const preview = await AgentThreadRuntimeControlsService.getEntryPreview({ + userIdentity, + agentId: 'custom.sample-agent', + source: { adapter: 'lifecycle_fork', input: {} }, + defaults: {}, + }); + const optionalChoiceId = getOptionalChoiceId(preview); + + const metadata = await AgentThreadRuntimeControlsService.validateEntryChoices({ + userIdentity, + agentId: 'custom.sample-agent', + source: { adapter: 'lifecycle_fork', input: {} }, + defaults: {}, + runtimeControlChoices: { + toolChoiceIds: [optionalChoiceId], + mcpChoiceIds: [], + }, + }); + + expect(metadata).toEqual({ + selectedAgentMetadataPatch: { + selectedAgentDefinitionId: 'custom.sample-agent', + }, + runtimeControlChoices: { + version: 1, + toolChoiceIds: [optionalChoiceId], + mcpChoiceIds: [], + }, + }); + expect(mockBuildSelectedAgentDefinitionMetadataPatch).toHaveBeenCalledWith('custom.sample-agent'); + }); + + it('stores selected agent metadata without runtime-choice metadata for agent-only create-session input', async () => { + const metadata = await AgentThreadRuntimeControlsService.validateEntryChoices({ + userIdentity, + agentId: 'custom.sample-agent', + source: { adapter: 'lifecycle_fork', input: {} }, + defaults: {}, + runtimeControlChoices: { + agentId: 'custom.sample-agent', + }, + }); + + expect(metadata).toEqual({ + selectedAgentMetadataPatch: { + selectedAgentDefinitionId: 'custom.sample-agent', + }, + runtimeControlChoices: null, + }); + }); + + it('stores Develop metadata for blank chat create-session input', async () => { + mockGetSystemAgentDefinition.mockResolvedValueOnce(SYSTEM_AGENT_DEFINITIONS['system.develop']); + + const metadata = await AgentThreadRuntimeControlsService.validateEntryChoices({ + userIdentity, + agentId: 'system.develop', + source: { adapter: 'blank_workspace', input: {} }, + defaults: {}, + runtimeControlChoices: { + agentId: 'system.develop', + }, + }); + + expect(metadata).toEqual({ + selectedAgentMetadataPatch: { + selectedAgentDefinitionId: 'system.develop', + }, + runtimeControlChoices: null, + }); + }); + + it('keeps default tool choices in /new preview when only MCP choices are provided', async () => { + const preview = await AgentThreadRuntimeControlsService.getEntryPreview({ + userIdentity, + agentId: 'custom.sample-agent', + source: { adapter: 'lifecycle_fork', input: {} }, + defaults: {}, + }); + const optionalChoiceId = getOptionalChoiceId(preview); + + const updatedPreview = await AgentThreadRuntimeControlsService.getEntryPreview({ + userIdentity, + agentId: 'custom.sample-agent', + source: { adapter: 'lifecycle_fork', input: {} }, + defaults: {}, + runtimeControlChoices: { + mcpChoiceIds: [], + }, + }); + + expect(updatedPreview.tools.selectedChoiceIds).toContain(optionalChoiceId); + expect(updatedPreview.mcp.selectedChoiceIds).toEqual([]); + }); + + it('throws a typed not_found error for missing threads', async () => { + mockGetOwnedThreadWithSession.mockRejectedValueOnce(new Error('Agent thread not found')); + + await expect( + AgentThreadRuntimeControlsService.getState({ + threadId: 'missing-thread', + userIdentity, + }) + ).rejects.toBeInstanceOf(AgentThreadRuntimeControlsError); + }); +}); diff --git a/src/server/services/agent/__tests__/ThreadService.test.ts b/src/server/services/agent/__tests__/ThreadService.test.ts index 6e724849..37c5394a 100644 --- a/src/server/services/agent/__tests__/ThreadService.test.ts +++ b/src/server/services/agent/__tests__/ThreadService.test.ts @@ -95,4 +95,105 @@ describe('AgentThreadService', () => { await expect(AgentThreadService.listThreadsForSession('session-1', 'user-123')).resolves.toEqual(listedThreads); }); + + it.each(['ended', 'error'])('blocks new threads for %s sessions', async (status) => { + mockAgentSessionQuery.mockReturnValueOnce({ + findOne: jest.fn().mockResolvedValue({ + id: 17, + uuid: 'session-1', + userId: 'user-123', + status, + }), + }); + + await expect(AgentThreadService.createThread('session-1', 'user-123', 'New chat')).rejects.toThrow( + 'Cannot create a thread for an inactive session' + ); + expect(mockAgentThreadQuery).not.toHaveBeenCalled(); + }); + + it('blocks new threads when the session runtime cannot accept messages', async () => { + mockAgentSessionQuery.mockReturnValueOnce({ + findOne: jest.fn().mockResolvedValue({ + id: 17, + uuid: 'session-1', + userId: 'user-123', + status: 'active', + sessionKind: 'environment', + chatStatus: 'ready', + workspaceStatus: 'failed', + }), + }); + + await expect(AgentThreadService.createThread('session-1', 'user-123', 'New chat')).rejects.toThrow( + 'This session is no longer available for new messages.' + ); + expect(mockAgentThreadQuery).not.toHaveBeenCalled(); + }); + + it('creates new threads when the session can accept messages', async () => { + const createdThread = { uuid: 'thread-2', sessionId: 17, isDefault: false }; + const insertAndFetch = jest.fn().mockResolvedValue(createdThread); + + mockAgentSessionQuery.mockReturnValueOnce({ + findOne: jest.fn().mockResolvedValue({ + id: 17, + uuid: 'session-1', + userId: 'user-123', + status: 'active', + sessionKind: 'chat', + chatStatus: 'ready', + workspaceStatus: 'none', + }), + }); + mockAgentThreadQuery.mockReturnValueOnce({ + insertAndFetch, + }); + + await expect(AgentThreadService.createThread('session-1', 'user-123', 'New chat')).resolves.toBe(createdThread); + expect(insertAndFetch).toHaveBeenCalledWith({ + sessionId: 17, + title: 'New chat', + isDefault: false, + metadata: { + sessionUuid: 'session-1', + }, + }); + }); + + it('reads selected agent definition metadata without agent-definition fallback', () => { + expect( + AgentThreadService.getSelectedAgentDefinitionId({ + metadata: { selectedAgentDefinitionId: 'system.debug' }, + } as any) + ).toBe('system.debug'); + expect( + AgentThreadService.getSelectedAgentDefinitionId({ + metadata: {}, + } as any) + ).toBeNull(); + }); + + it('builds a scoped selected agent definition metadata patch', () => { + expect(AgentThreadService.buildSelectedAgentDefinitionMetadataPatch('custom.sample-agent')).toEqual({ + selectedAgentDefinitionId: 'custom.sample-agent', + }); + }); + + it('trims explicit selected agent definition metadata', () => { + expect( + AgentThreadService.getSelectedAgentDefinitionId({ + metadata: { + selectedAgentDefinitionId: ' custom.sample-agent ', + }, + } as any) + ).toBe('custom.sample-agent'); + expect( + AgentThreadService.getSelectedAgentDefinitionId({ + metadata: { + selectedAgentDefinitionId: ' ', + }, + } as any) + ).toBeNull(); + }); }); diff --git a/src/server/services/agent/agentDefinitionTypes.ts b/src/server/services/agent/agentDefinitionTypes.ts new file mode 100644 index 00000000..afff2f17 --- /dev/null +++ b/src/server/services/agent/agentDefinitionTypes.ts @@ -0,0 +1,70 @@ +/** + * 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 { AgentCapabilityCatalogId } from './capabilityCatalog'; + +export type AgentDefinitionOwnerKind = 'system' | 'admin' | 'user'; +export type AgentDefinitionStatus = 'active' | 'disabled' | 'archived'; + +export interface AgentDefinitionOwner { + kind: AgentDefinitionOwnerKind; + userId?: string | null; + organizationId?: string | null; +} + +export interface AgentDefinitionResourcePolicy { + sourceKinds: string[]; + sandboxRequired?: boolean; + workspaceRequired?: boolean; +} + +export interface AgentDefinitionModelPreference { + provider?: string | null; + model?: string | null; +} + +export type UserAgentDefinitionResourceBehavior = 'chat_only' | 'current_workspace_when_available'; + +export interface UserAgentDefinitionUpsertInput { + name: string; + description?: string | null; + instructionAddendum: string; + capabilityRefs?: AgentCapabilityCatalogId[]; + modelPreference?: AgentDefinitionModelPreference | null; + resourceBehavior: UserAgentDefinitionResourceBehavior; +} + +export interface UserAgentDefinitionListFilters { + status?: Extract; +} + +export interface AgentDefinitionContract { + id: string; + version: number; + owner: AgentDefinitionOwner; + name: string; + description?: string | null; + instructionRefs: string[]; + instructionAddendum?: string | null; + capabilityRefs: AgentCapabilityCatalogId[]; + requiredCapabilityRefs?: AgentCapabilityCatalogId[]; + optionalCapabilityRefs?: AgentCapabilityCatalogId[]; + resourcePolicy: AgentDefinitionResourcePolicy; + modelPreference?: AgentDefinitionModelPreference | null; + status: AgentDefinitionStatus; + codeOwned?: boolean; + readOnly?: boolean; +} diff --git a/src/server/services/agent/canonicalMessages.ts b/src/server/services/agent/canonicalMessages.ts index 2011e9da..0c65a6c5 100644 --- a/src/server/services/agent/canonicalMessages.ts +++ b/src/server/services/agent/canonicalMessages.ts @@ -40,8 +40,9 @@ export type CanonicalAgentMessage = { clientMessageId: string | null; threadId: string; runId: string | null; - role: 'user' | 'assistant'; + role: 'user' | 'assistant' | 'system'; parts: CanonicalAgentMessagePart[]; + metadata?: Record; createdAt: string | null; }; diff --git a/src/server/services/agent/capabilityCatalog.ts b/src/server/services/agent/capabilityCatalog.ts new file mode 100644 index 00000000..6c277238 --- /dev/null +++ b/src/server/services/agent/capabilityCatalog.ts @@ -0,0 +1,282 @@ +/** + * 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 { AgentApprovalMode, AgentCapabilityKey } from './types'; + +export const AGENT_CAPABILITY_AVAILABILITIES = ['all_users', 'admin_only', 'system_only', 'disabled'] as const; + +export type AgentCapabilityAvailability = (typeof AGENT_CAPABILITY_AVAILABILITIES)[number]; + +export const AGENT_CAPABILITY_CATALOG_IDS = [ + 'read_context', + 'diagnostics_logs', + 'diagnostics_codefresh', + 'diagnostics_kubernetes', + 'diagnostics_database', + 'github_read', + 'github_write', + 'workspace_files', + 'workspace_shell', + 'workspace_git', + 'network_access', + 'preview_publish', + 'external_mcp_read', + 'external_mcp_write', + 'approval_controls', +] as const; + +export type AgentCapabilityCatalogId = (typeof AGENT_CAPABILITY_CATALOG_IDS)[number]; + +export type AgentCapabilityCategory = + | 'read' + | 'diagnostics' + | 'workspace' + | 'source_control' + | 'mcp' + | 'deployment' + | 'network' + | 'preview' + | 'approval'; + +export type AgentCapabilitySourceKind = 'build_context_chat' | 'workspace_session' | 'freeform_chat'; + +export interface AgentCapabilityCatalogEntry { + id: AgentCapabilityCatalogId; + category: AgentCapabilityCategory; + label: string; + description: string; + defaultAvailability: AgentCapabilityAvailability; + defaultApprovalMode: AgentApprovalMode; + runtimeCapabilityKey?: AgentCapabilityKey; + toolKeys?: readonly string[]; + resourceGrants?: readonly string[]; + sourceKinds?: readonly AgentCapabilitySourceKind[]; + userSelectable: boolean; +} + +export const AGENT_CAPABILITY_CATALOG: readonly AgentCapabilityCatalogEntry[] = [ + { + id: 'read_context', + category: 'read', + label: 'Read/context', + description: 'Read session context, workspace state, logs, and non-mutating reference data.', + defaultAvailability: 'all_users', + defaultApprovalMode: 'allow', + runtimeCapabilityKey: 'read', + resourceGrants: ['session_context'], + sourceKinds: ['build_context_chat', 'workspace_session', 'freeform_chat'], + userSelectable: true, + }, + { + id: 'diagnostics_logs', + category: 'diagnostics', + label: 'Diagnostic logs', + description: 'Inspect Lifecycle and workload logs for troubleshooting.', + defaultAvailability: 'system_only', + defaultApprovalMode: 'allow', + runtimeCapabilityKey: 'read', + toolKeys: ['lifecycle.get_logs', 'k8s.get_pod_logs'], + resourceGrants: ['build_context'], + sourceKinds: ['build_context_chat'], + userSelectable: false, + }, + { + id: 'diagnostics_codefresh', + category: 'diagnostics', + label: 'Codefresh diagnostics', + description: 'Read Codefresh pipeline logs and build details.', + defaultAvailability: 'system_only', + defaultApprovalMode: 'allow', + runtimeCapabilityKey: 'read', + toolKeys: ['codefresh.get_logs'], + resourceGrants: ['build_context', 'codefresh'], + sourceKinds: ['build_context_chat'], + userSelectable: false, + }, + { + id: 'diagnostics_kubernetes', + category: 'diagnostics', + label: 'Kubernetes diagnostics', + description: 'Read Kubernetes resources, events, and pod state.', + defaultAvailability: 'system_only', + defaultApprovalMode: 'allow', + runtimeCapabilityKey: 'read', + toolKeys: ['k8s.get_resources', 'k8s.get_pod_logs', 'lifecycle.get_logs'], + resourceGrants: ['build_context', 'kubernetes_read'], + sourceKinds: ['build_context_chat'], + userSelectable: false, + }, + { + id: 'diagnostics_database', + category: 'diagnostics', + label: 'Database diagnostics', + description: 'Run read-only database inspection needed for environment troubleshooting.', + defaultAvailability: 'system_only', + defaultApprovalMode: 'allow', + runtimeCapabilityKey: 'read', + toolKeys: ['query_database'], + resourceGrants: ['build_context', 'database_read'], + sourceKinds: ['build_context_chat'], + userSelectable: false, + }, + { + id: 'github_read', + category: 'source_control', + label: 'GitHub read', + description: 'Read repository files, pull request context, and issue comments.', + defaultAvailability: 'all_users', + defaultApprovalMode: 'allow', + runtimeCapabilityKey: 'read', + toolKeys: ['github.get_file', 'github.list_directory', 'github.get_issue_comment'], + resourceGrants: ['github_read'], + sourceKinds: ['build_context_chat', 'workspace_session'], + userSelectable: true, + }, + { + id: 'github_write', + category: 'source_control', + label: 'GitHub write', + description: 'Apply repository or pull request fixes through GitHub.', + defaultAvailability: 'system_only', + defaultApprovalMode: 'require_approval', + runtimeCapabilityKey: 'git_write', + toolKeys: ['github.update_file', 'github.update_pr_labels'], + resourceGrants: ['github_write'], + sourceKinds: ['build_context_chat'], + userSelectable: false, + }, + { + id: 'workspace_files', + category: 'workspace', + label: 'Workspace files', + description: 'Create and edit files inside a development workspace.', + defaultAvailability: 'all_users', + defaultApprovalMode: 'require_approval', + runtimeCapabilityKey: 'workspace_write', + toolKeys: ['workspace.write_file', 'workspace.edit_file'], + resourceGrants: ['workspace_write'], + sourceKinds: ['workspace_session'], + userSelectable: true, + }, + { + id: 'workspace_shell', + category: 'workspace', + label: 'Command tools', + description: 'Run shell commands inside a development workspace.', + defaultAvailability: 'all_users', + defaultApprovalMode: 'require_approval', + runtimeCapabilityKey: 'shell_exec', + toolKeys: ['workspace.exec'], + resourceGrants: ['workspace_shell'], + sourceKinds: ['workspace_session'], + userSelectable: true, + }, + { + id: 'workspace_git', + category: 'source_control', + label: 'Source control', + description: 'Stage, commit, and manage repository branches in a workspace.', + defaultAvailability: 'all_users', + defaultApprovalMode: 'require_approval', + runtimeCapabilityKey: 'git_write', + toolKeys: ['git.add', 'git.commit', 'git.branch'], + resourceGrants: ['git_write'], + sourceKinds: ['workspace_session'], + userSelectable: true, + }, + { + id: 'network_access', + category: 'network', + label: 'Network access', + description: 'Use tools that can reach external network resources.', + defaultAvailability: 'all_users', + defaultApprovalMode: 'require_approval', + runtimeCapabilityKey: 'network_access', + resourceGrants: ['network_access'], + sourceKinds: ['workspace_session'], + userSelectable: true, + }, + { + id: 'preview_publish', + category: 'preview', + label: 'Preview/publish', + description: 'Publish or expose workspace preview services.', + defaultAvailability: 'all_users', + defaultApprovalMode: 'require_approval', + runtimeCapabilityKey: 'deploy_k8s_mutation', + toolKeys: ['publish_http'], + resourceGrants: ['preview_publish'], + sourceKinds: ['workspace_session'], + userSelectable: true, + }, + { + id: 'external_mcp_read', + category: 'mcp', + label: 'MCP read', + description: 'Use connected MCP tools that declare read-only behavior.', + defaultAvailability: 'all_users', + defaultApprovalMode: 'allow', + runtimeCapabilityKey: 'external_mcp_read', + resourceGrants: ['mcp_read'], + sourceKinds: ['build_context_chat', 'workspace_session', 'freeform_chat'], + userSelectable: true, + }, + { + id: 'external_mcp_write', + category: 'mcp', + label: 'MCP write', + description: 'Use connected MCP tools that can mutate external systems.', + defaultAvailability: 'admin_only', + defaultApprovalMode: 'require_approval', + runtimeCapabilityKey: 'external_mcp_write', + resourceGrants: ['mcp_write'], + sourceKinds: ['build_context_chat', 'workspace_session', 'freeform_chat'], + userSelectable: true, + }, + { + id: 'approval_controls', + category: 'approval', + label: 'Approval controls', + description: 'Control approval behavior for protected tool execution.', + defaultAvailability: 'system_only', + defaultApprovalMode: 'require_approval', + resourceGrants: ['approval_policy'], + sourceKinds: ['build_context_chat', 'workspace_session', 'freeform_chat'], + userSelectable: false, + }, +] as const; + +const AGENT_CAPABILITY_CATALOG_BY_ID = new Map(AGENT_CAPABILITY_CATALOG.map((entry) => [entry.id, entry])); + +export function isAgentCapabilityCatalogId(value: string): value is AgentCapabilityCatalogId { + return AGENT_CAPABILITY_CATALOG_IDS.includes(value as AgentCapabilityCatalogId); +} + +export function isAgentCapabilityAvailability(value: string): value is AgentCapabilityAvailability { + return AGENT_CAPABILITY_AVAILABILITIES.includes(value as AgentCapabilityAvailability); +} + +export function listAgentCapabilityCatalogEntries(): readonly AgentCapabilityCatalogEntry[] { + return AGENT_CAPABILITY_CATALOG; +} + +export function getAgentCapabilityCatalogEntry(id: AgentCapabilityCatalogId): AgentCapabilityCatalogEntry { + const entry = AGENT_CAPABILITY_CATALOG_BY_ID.get(id); + if (!entry) { + throw new Error(`Unknown agent capability catalog id: ${id}`); + } + return entry; +} diff --git a/src/server/services/agent/diagnosticTools.ts b/src/server/services/agent/diagnosticTools.ts new file mode 100644 index 00000000..dd8cdd28 --- /dev/null +++ b/src/server/services/agent/diagnosticTools.ts @@ -0,0 +1,511 @@ +/** + * 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 { createHash } from 'crypto'; +import { dynamicTool, jsonSchema, type ToolSet } from 'ai'; +import type AgentSession from 'server/models/AgentSession'; +import * as models from 'server/models'; +import { GetCodefreshLogsTool } from 'server/services/agent/tools/codefresh/getCodefreshLogs'; +import { GetFileTool } from 'server/services/agent/tools/github/getFile'; +import { GetIssueCommentTool } from 'server/services/agent/tools/github/getIssueComment'; +import { ListDirectoryTool } from 'server/services/agent/tools/github/listDirectory'; +import { UpdateFileTool } from 'server/services/agent/tools/github/updateFile'; +import { UpdatePrLabelsTool } from 'server/services/agent/tools/github/updatePrLabels'; +import { GetK8sResourcesTool } from 'server/services/agent/tools/k8s/getK8sResources'; +import { GetLifecycleLogsTool } from 'server/services/agent/tools/k8s/getLifecycleLogs'; +import { PatchK8sResourceTool } from 'server/services/agent/tools/k8s/patchK8sResource'; +import { GetPodLogsTool } from 'server/services/agent/tools/k8s/getPodLogs'; +import { QueryDatabaseTool } from 'server/services/agent/tools/k8s/queryDatabase'; +import { DatabaseClient } from 'server/services/agent/tools/shared/databaseClient'; +import { GitHubClient } from 'server/services/agent/tools/shared/githubClient'; +import { K8sClient } from 'server/services/agent/tools/shared/k8sClient'; +import type { Tool } from 'server/services/agent/tools/types'; +import type { + AgentApprovalMode, + AgentApprovalPolicy, + AgentCapabilityKey, + AgentFileChangeData, + AgentToolAuditRecord, +} from './types'; +import AgentPolicyService from './PolicyService'; +import type { ResolvedAgentCapabilityAccess } from './PolicyService'; +import type { AgentCapabilityCatalogId } from './capabilityCatalog'; +import type { AgentSessionToolRule } from 'server/services/types/agentSessionConfig'; +import { buildAgentToolKey, LIFECYCLE_BUILTIN_SERVER_SLUG } from './toolKeys'; +import { getLogger } from 'server/lib/logger'; +import { DEFAULT_AGENT_SESSION_FILE_CHANGE_PREVIEW_CHARS } from 'server/lib/agentSession/runtimeConfig'; + +type ToolExecutionHooks = { + onToolStarted?: (audit: AgentToolAuditRecord) => Promise; + onToolFinished?: (audit: AgentToolAuditRecord & { result: unknown; status: 'completed' | 'failed' }) => Promise; + onFileChange?: (change: AgentFileChangeData) => Promise; +}; + +const LIFECYCLE_DIAGNOSTIC_READ_CAPABILITY: AgentCapabilityKey = 'read'; +const FILE_CHANGE_PREVIEW_CHARS = DEFAULT_AGENT_SESSION_FILE_CHANGE_PREVIEW_CHARS; +const MAX_EXACT_DIFF_MATRIX_CELLS = 1_000_000; + +function toAiJsonSchema(schema: unknown) { + return jsonSchema(schema as any); +} + +function toAiDynamicTool(config: unknown) { + return dynamicTool(config as any); +} + +type LifecycleDiagnosticToolSpec = { + tool: Tool; + capabilityKey: AgentCapabilityKey; + catalogCapabilityId: AgentCapabilityCatalogId; + forceApproval: boolean; + shouldRequestApproval?: (input: Record) => boolean | Promise; + buildProposedFileChanges?: ( + input: Record, + toolCallId: string, + sourceTool: string + ) => AgentFileChangeData[] | Promise; +}; + +export type LifecycleDiagnosticGithubSafety = { + allowedBranch?: string | null; + referencedFiles?: string[]; + excludedFilePatterns?: string[]; + allowedWritePatterns?: string[]; +}; + +function resolveToolMode({ + toolRules, + toolKey, + approvalPolicy, + capabilityKey, + forceApproval, +}: { + toolRules?: AgentSessionToolRule[]; + toolKey: string; + approvalPolicy: AgentApprovalPolicy; + capabilityKey: AgentCapabilityKey; + forceApproval: boolean; +}): AgentApprovalMode { + const toolRule = toolRules?.find((rule) => rule.toolKey === toolKey); + const capabilityMode = AgentPolicyService.modeForCapability(approvalPolicy, capabilityKey); + + if (toolRule?.mode === 'deny' || (!toolRule && capabilityMode === 'deny')) { + return 'deny'; + } + + if (forceApproval) { + return 'require_approval'; + } + + return toolRule?.mode || capabilityMode; +} + +function configureGithubClient(client: GitHubClient, safety?: LifecycleDiagnosticGithubSafety): GitHubClient { + const allowedBranch = safety?.allowedBranch?.trim(); + if (allowedBranch) { + client.setAllowedBranch(allowedBranch); + } + + client.setReferencedFiles(safety?.referencedFiles || []); + client.setExcludedFilePatterns(safety?.excludedFilePatterns || []); + client.setAllowedWritePatterns(safety?.allowedWritePatterns || []); + + return client; +} + +function readString(value: unknown): string | null { + return typeof value === 'string' && value.trim() ? value : null; +} + +function shouldRequestUpdateFileApproval(client: GitHubClient, input: Record): boolean { + const filePath = readString(input.file_path); + const branch = readString(input.branch); + if (!filePath || !branch) { + return false; + } + + return client.isFilePathAllowed(filePath, 'write') && client.validateBranch(branch).valid; +} + +function createLifecycleDiagnosticReadToolSpecs( + safety?: LifecycleDiagnosticGithubSafety +): LifecycleDiagnosticToolSpec[] { + const k8sClient = new K8sClient(); + const githubClient = configureGithubClient(new GitHubClient(), safety); + const databaseClient = new DatabaseClient({ models }); + + const specs: Array<{ tool: Tool; catalogCapabilityId: AgentCapabilityCatalogId }> = [ + { tool: new GetCodefreshLogsTool(), catalogCapabilityId: 'diagnostics_codefresh' }, + { tool: new GetK8sResourcesTool(k8sClient), catalogCapabilityId: 'diagnostics_kubernetes' }, + { tool: new GetPodLogsTool(k8sClient), catalogCapabilityId: 'diagnostics_logs' }, + { tool: new GetLifecycleLogsTool(k8sClient), catalogCapabilityId: 'diagnostics_logs' }, + { tool: new QueryDatabaseTool(databaseClient), catalogCapabilityId: 'diagnostics_database' }, + { tool: new GetFileTool(githubClient), catalogCapabilityId: 'github_read' }, + { tool: new ListDirectoryTool(githubClient), catalogCapabilityId: 'github_read' }, + { tool: new GetIssueCommentTool(githubClient), catalogCapabilityId: 'github_read' }, + ]; + + return specs.map(({ tool, catalogCapabilityId }) => ({ + tool, + catalogCapabilityId, + capabilityKey: LIFECYCLE_DIAGNOSTIC_READ_CAPABILITY, + forceApproval: false, + })); +} + +function trimPreview(value: string): string { + return value.length > FILE_CHANGE_PREVIEW_CHARS + ? `${value.slice(0, FILE_CHANGE_PREVIEW_CHARS)}\n\n[truncated]` + : value; +} + +function countLines(value: string): number { + return value ? value.split('\n').length : 0; +} + +function splitLines(value: string): string[] { + return value.length === 0 ? [] : value.split('\n'); +} + +function normalizeFilePath(path: string): string { + return path.replace(/^\/+/, '').replace(/^\.\//, ''); +} + +function buildSingleHunkUnifiedDiff(path: string, oldContent: string, newContent: string) { + if (oldContent === newContent) { + return { + unifiedDiff: null, + additions: 0, + deletions: 0, + }; + } + + const oldLines = splitLines(oldContent); + const newLines = splitLines(newContent); + + if (oldLines.length * newLines.length > MAX_EXACT_DIFF_MATRIX_CELLS) { + return { + unifiedDiff: [ + `diff --git a/${path} b/${path}`, + `--- a/${path}`, + `+++ b/${path}`, + `@@ -1,${oldLines.length} +1,${newLines.length} @@`, + ...oldLines.map((line) => `-${line}`), + ...newLines.map((line) => `+${line}`), + ].join('\n'), + additions: newLines.length, + deletions: oldLines.length, + }; + } + + const dp = Array.from({ length: oldLines.length + 1 }, () => Array(newLines.length + 1).fill(0)); + + for (let oldIndex = oldLines.length - 1; oldIndex >= 0; oldIndex -= 1) { + for (let newIndex = newLines.length - 1; newIndex >= 0; newIndex -= 1) { + dp[oldIndex][newIndex] = + oldLines[oldIndex] === newLines[newIndex] + ? dp[oldIndex + 1][newIndex + 1] + 1 + : Math.max(dp[oldIndex + 1][newIndex], dp[oldIndex][newIndex + 1]); + } + } + + let oldIndex = 0; + let newIndex = 0; + let additions = 0; + let deletions = 0; + const diffLines: string[] = []; + + while (oldIndex < oldLines.length || newIndex < newLines.length) { + if (oldIndex < oldLines.length && newIndex < newLines.length && oldLines[oldIndex] === newLines[newIndex]) { + diffLines.push(` ${oldLines[oldIndex]}`); + oldIndex += 1; + newIndex += 1; + continue; + } + + if ( + newIndex < newLines.length && + (oldIndex >= oldLines.length || dp[oldIndex][newIndex + 1] >= dp[oldIndex + 1][newIndex]) + ) { + diffLines.push(`+${newLines[newIndex]}`); + additions += 1; + newIndex += 1; + continue; + } + + if (oldIndex < oldLines.length) { + diffLines.push(`-${oldLines[oldIndex]}`); + deletions += 1; + oldIndex += 1; + } + } + + return { + unifiedDiff: [ + `diff --git a/${path} b/${path}`, + `--- a/${path}`, + `+++ b/${path}`, + `@@ -1,${oldLines.length} +1,${newLines.length} @@`, + ...diffLines, + ].join('\n'), + additions, + deletions, + }; +} + +async function readGithubFileContent( + githubClient: GitHubClient, + input: Record, + path: string +): Promise { + const owner = readString(input.repository_owner); + const repo = readString(input.repository_name); + const branch = readString(input.branch); + if (!owner || !repo || !branch) { + return null; + } + + try { + const octokit = await githubClient.getOctokit('agent-runtime-update-file-preview'); + const currentFile = await octokit.request(`GET /repos/${owner}/${repo}/contents/${path}`, { + ref: branch, + }); + const data = currentFile.data; + if (!data || Array.isArray(data) || !('content' in data) || typeof data.content !== 'string') { + return null; + } + + return Buffer.from(data.content, 'base64').toString('utf-8'); + } catch { + return null; + } +} + +async function buildUpdateFilePreview( + githubClient: GitHubClient, + input: Record, + toolCallId: string, + sourceTool: string +): Promise { + if (typeof input.file_path !== 'string' || typeof input.new_content !== 'string') { + return []; + } + + const path = normalizeFilePath(input.file_path); + const content = input.new_content.replace(/\\n/g, '\n').replace(/\\r/g, '\r').replace(/\\t/g, '\t'); + const oldContent = await readGithubFileContent(githubClient, input, path); + const diff = oldContent === null ? null : buildSingleHunkUnifiedDiff(path, oldContent, content); + const beforeTextPreview = oldContent === null ? null : trimPreview(oldContent); + const afterTextPreview = trimPreview(content); + + return [ + { + id: `${toolCallId}:${path}`, + toolCallId, + sourceTool, + path, + displayPath: path, + kind: oldContent === null ? 'created' : 'edited', + stage: 'awaiting-approval', + additions: diff?.additions ?? countLines(content), + deletions: diff?.deletions ?? 0, + truncated: + content.length > FILE_CHANGE_PREVIEW_CHARS || + (oldContent !== null && oldContent.length > FILE_CHANGE_PREVIEW_CHARS), + unifiedDiff: diff?.unifiedDiff ?? null, + beforeTextPreview, + afterTextPreview, + summary: `Proposed update to ${path}`, + encoding: 'utf-8', + oldSizeBytes: oldContent === null ? null : Buffer.byteLength(oldContent, 'utf8'), + newSizeBytes: Buffer.byteLength(content, 'utf8'), + oldSha256: oldContent === null ? null : createHash('sha256').update(oldContent).digest('hex'), + newSha256: createHash('sha256').update(content).digest('hex'), + }, + ]; +} + +function createLifecycleDiagnosticFixToolSpecs( + safety?: LifecycleDiagnosticGithubSafety +): LifecycleDiagnosticToolSpec[] { + const k8sClient = new K8sClient(); + const githubClient = configureGithubClient(new GitHubClient(), safety); + + return [ + { + tool: new UpdateFileTool(githubClient), + capabilityKey: 'git_write', + catalogCapabilityId: 'github_write', + forceApproval: true, + shouldRequestApproval: (input) => shouldRequestUpdateFileApproval(githubClient, input), + buildProposedFileChanges: (input, toolCallId, sourceTool) => + buildUpdateFilePreview(githubClient, input, toolCallId, sourceTool), + }, + { + tool: new UpdatePrLabelsTool(githubClient), + capabilityKey: 'git_write', + catalogCapabilityId: 'github_write', + forceApproval: true, + }, + { + tool: new PatchK8sResourceTool(k8sClient), + capabilityKey: 'deploy_k8s_mutation', + catalogCapabilityId: 'diagnostics_kubernetes', + forceApproval: true, + }, + ]; +} + +function isCatalogCapabilityAllowed( + resolvedCapabilityAccess: ResolvedAgentCapabilityAccess[] | undefined, + capabilityId: AgentCapabilityCatalogId +): boolean { + if (!resolvedCapabilityAccess) { + return false; + } + + return resolvedCapabilityAccess.some((entry) => entry.capabilityId === capabilityId && entry.allowed); +} + +function registerLifecycleDiagnosticToolSpecs({ + tools, + session, + approvalPolicy, + hooks, + toolRules, + specs, + resolvedCapabilityAccess, +}: { + tools: ToolSet; + session: AgentSession; + approvalPolicy: AgentApprovalPolicy; + hooks?: ToolExecutionHooks; + toolRules?: AgentSessionToolRule[]; + specs: LifecycleDiagnosticToolSpec[]; + resolvedCapabilityAccess?: ResolvedAgentCapabilityAccess[]; + githubSafety?: LifecycleDiagnosticGithubSafety; +}) { + if (!session.buildUuid) { + return; + } + + for (const { + tool: diagnosticTool, + capabilityKey, + catalogCapabilityId, + forceApproval, + shouldRequestApproval, + buildProposedFileChanges, + } of specs) { + if (!isCatalogCapabilityAllowed(resolvedCapabilityAccess, catalogCapabilityId)) { + continue; + } + + const toolKey = buildAgentToolKey(LIFECYCLE_BUILTIN_SERVER_SLUG, diagnosticTool.name); + const mode = resolveToolMode({ + toolRules, + toolKey, + approvalPolicy, + capabilityKey, + forceApproval, + }); + + if (mode === 'deny') { + continue; + } + + tools[toolKey] = toAiDynamicTool({ + description: diagnosticTool.description, + inputSchema: toAiJsonSchema(diagnosticTool.parameters as Record), + needsApproval: + mode === 'require_approval' + ? shouldRequestApproval + ? async (input: unknown) => + shouldRequestApproval(((input as Record) || {}) as Record) + : true + : false, + onInputAvailable: buildProposedFileChanges + ? async ({ input, toolCallId }) => { + if (!toolCallId) { + return; + } + + const args = (input as Record) || {}; + for (const change of await buildProposedFileChanges(args, toolCallId, diagnosticTool.name)) { + await hooks?.onFileChange?.(change); + } + } + : undefined, + execute: async (input, context) => { + const args = (input as Record) || {}; + const toolCallId = context?.toolCallId; + const audit: AgentToolAuditRecord = { + source: 'mcp', + serverSlug: LIFECYCLE_BUILTIN_SERVER_SLUG, + toolName: diagnosticTool.name, + toolCallId, + args, + capabilityKey, + }; + + await hooks?.onToolStarted?.(audit); + + try { + const abortSignal = (context as { abortSignal?: AbortSignal } | undefined)?.abortSignal; + const result = await diagnosticTool.execute(args, abortSignal); + await hooks?.onToolFinished?.({ + ...audit, + result, + status: result.success ? 'completed' : 'failed', + }); + return result; + } catch (error) { + const result = { + error: error instanceof Error ? error.message : String(error), + }; + getLogger().warn( + { error }, + `AgentExec: lifecycle diagnostic tool failed sessionId=${session.uuid} tool=${diagnosticTool.name}` + ); + await hooks?.onToolFinished?.({ + ...audit, + result, + status: 'failed', + }); + throw error; + } + }, + }); + } +} + +export function registerLifecycleDiagnosticReadTools( + options: Omit[0], 'specs'> +) { + registerLifecycleDiagnosticToolSpecs({ + ...options, + specs: createLifecycleDiagnosticReadToolSpecs(options.githubSafety), + }); +} + +export function registerLifecycleDiagnosticFixTools( + options: Omit[0], 'specs'> +) { + registerLifecycleDiagnosticToolSpecs({ + ...options, + specs: createLifecycleDiagnosticFixToolSpecs(options.githubSafety), + }); +} diff --git a/src/server/services/agent/fileChanges.ts b/src/server/services/agent/fileChanges.ts index 237fd930..75c21bb0 100644 --- a/src/server/services/agent/fileChanges.ts +++ b/src/server/services/agent/fileChanges.ts @@ -14,7 +14,7 @@ * limitations under the License. */ -import { createHash } from 'node:crypto'; +import { createHash } from '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'; diff --git a/src/server/services/agent/runPlanSummary.ts b/src/server/services/agent/runPlanSummary.ts new file mode 100644 index 00000000..c6fc802d --- /dev/null +++ b/src/server/services/agent/runPlanSummary.ts @@ -0,0 +1,205 @@ +/** + * 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 { AgentRunPlanPublicSummary, AgentRunPlanSourceKind } from './runPlanTypes'; +import { isAgentRunPlanSnapshotV1 } from './runPlanTypes'; +import { isAgentCapabilityAvailability, isAgentCapabilityCatalogId } from './capabilityCatalog'; +import type { AgentApprovalMode } from './types'; + +function readRecord(value: unknown): Record | null { + if (!value || typeof value !== 'object' || Array.isArray(value)) { + return null; + } + + return value as Record; +} + +function readSourceKind(value: unknown): AgentRunPlanSourceKind | null { + if (value === 'build_context_chat' || value === 'workspace_session' || value === 'freeform_chat') { + return value; + } + + return null; +} + +function readNullableString(value: unknown): string | null { + return typeof value === 'string' ? value : null; +} + +function readApprovalMode(value: unknown): AgentApprovalMode | null { + if (value === 'allow' || value === 'require_approval' || value === 'deny') { + return value; + } + + return null; +} + +function readStringArray(value: unknown): string[] { + return Array.isArray(value) ? value.filter((entry): entry is string => typeof entry === 'string') : []; +} + +function readCapabilityIds(value: unknown): AgentRunPlanPublicSummary['capabilities']['selected']['capabilityIds'] { + return readStringArray(value).filter(isAgentCapabilityCatalogId); +} + +function readCapabilitySummary(value: unknown): AgentRunPlanPublicSummary['capabilities']['effective'][number] | null { + const capability = readRecord(value); + + if (!capability) { + return null; + } + + const capabilityId = readNullableString(capability.capabilityId); + const availability = readNullableString(capability.availability); + const allowed = capability.allowed; + const approvalMode = readApprovalMode(capability.approvalMode); + + if ( + !capabilityId || + !isAgentCapabilityCatalogId(capabilityId) || + !availability || + !isAgentCapabilityAvailability(availability) || + typeof allowed !== 'boolean' + ) { + return null; + } + + return { + capabilityId, + availability, + allowed, + ...(approvalMode ? { approvalMode } : {}), + }; +} + +function readCapabilitySummaries(value: unknown): AgentRunPlanPublicSummary['capabilities']['effective'] | null { + if (!Array.isArray(value)) { + return null; + } + + const summaries: AgentRunPlanPublicSummary['capabilities']['effective'] = []; + + for (const entry of value) { + const summary = readCapabilitySummary(entry); + if (!summary) { + return null; + } + summaries.push(summary); + } + + return summaries; +} + +function readWarningSummary(warnings: unknown[]): AgentRunPlanPublicSummary['warnings'] { + return warnings + .map(readRecord) + .filter((warning): warning is Record => Boolean(warning)) + .filter( + (warning): warning is { code: string; message: string } => + typeof warning.code === 'string' && typeof warning.message === 'string' + ) + .map((warning) => ({ + code: warning.code, + message: warning.message, + })); +} + +export function serializeRunPlanSummary(snapshot: unknown): AgentRunPlanPublicSummary | null { + if (!isAgentRunPlanSnapshotV1(snapshot)) { + return null; + } + + const agent = readRecord(snapshot.agent); + const source = readRecord(snapshot.source); + const model = readRecord(snapshot.model); + const runtime = readRecord(snapshot.runtime); + const approvalPolicy = runtime ? readRecord(runtime.approvalPolicy) : null; + const runtimeOptions = runtime ? readRecord(runtime.runtimeOptions) : null; + const capabilities = readRecord(snapshot.capabilities); + + if ( + !agent || + !source || + !model || + !runtime || + !approvalPolicy || + !runtimeOptions || + !capabilities || + !Array.isArray(snapshot.warnings) + ) { + return null; + } + + const label = readNullableString(agent.label); + const agentId = readNullableString(agent.id); + const sourceKind = readSourceKind(agent.sourceKind); + const provider = readNullableString(model.resolvedProvider); + const resolvedModel = readNullableString(model.resolvedModel); + const harness = runtime.resolvedHarness; + const maxIterations = typeof runtimeOptions.maxIterations === 'number' ? runtimeOptions.maxIterations : null; + const defaultMode = readApprovalMode(approvalPolicy.defaultMode); + const effectiveCapabilities = readCapabilitySummaries(capabilities.resolvedCapabilityAccess); + const selectedCapabilityIds = readCapabilityIds(capabilities.selectedRuntimeCapabilityIds); + + if ( + !agentId || + !label || + !sourceKind || + !provider || + !resolvedModel || + harness !== 'lifecycle_ai_sdk' || + !defaultMode || + !effectiveCapabilities + ) { + return null; + } + + return { + version: 1, + agent: { + id: agentId, + label, + sourceKind, + }, + source: { + kind: sourceKind, + repoFullName: readNullableString(source.repoFullName), + branch: readNullableString(source.branch), + buildUuid: readNullableString(source.buildUuid), + namespace: readNullableString(source.namespace), + }, + model: { + provider, + model: resolvedModel, + }, + runtime: { + harness, + maxIterations, + }, + approval: { + defaultMode, + }, + capabilities: { + effective: effectiveCapabilities, + selected: { + capabilityIds: selectedCapabilityIds, + toolChoiceIds: readStringArray(capabilities.selectedRuntimeToolChoiceIds), + mcpChoiceIds: readStringArray(capabilities.selectedRuntimeMcpChoiceIds), + }, + }, + warnings: readWarningSummary(snapshot.warnings), + }; +} diff --git a/src/server/services/agent/runPlanTypes.ts b/src/server/services/agent/runPlanTypes.ts new file mode 100644 index 00000000..6f7a2555 --- /dev/null +++ b/src/server/services/agent/runPlanTypes.ts @@ -0,0 +1,170 @@ +/** + * 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 { AgentApprovalMode, AgentApprovalPolicy, AgentCapabilityKey } from './types'; +import type { AgentCapabilityAvailability, AgentCapabilityCatalogId } from './capabilityCatalog'; +import type { AgentRunRuntimeOptions } from './canonicalMessages'; +import type { + AgentDefinitionModelPreference, + AgentDefinitionOwnerKind, + AgentDefinitionResourcePolicy, +} from './agentDefinitionTypes'; + +export type AgentRunPlanSourceKind = 'build_context_chat' | 'workspace_session' | 'freeform_chat'; + +export interface AgentRunPlanWarning { + code: string; + message: string; + detail?: Record; +} + +export interface AgentRunPlanAgentSnapshot { + id: string; + label: string; + ownerKind?: AgentDefinitionOwnerKind; + version?: number; + sourceKind: AgentRunPlanSourceKind; + resourcePolicy?: AgentDefinitionResourcePolicy; + modelPreference?: AgentDefinitionModelPreference | null; +} + +export interface AgentRunPlanSourceSnapshot { + id?: string | null; + adapter?: string | null; + status?: string | null; + sessionKind?: string | null; + buildUuid?: string | null; + repoFullName?: string | null; + branch?: string | null; + namespace?: string | null; + workspaceLayout?: { + repoCount?: number; + primaryRepo?: string | null; + selectedServiceCount?: number; + primaryService?: string | null; + }; + sandboxRequirements?: Record; + freshness: { + capturedAt: string; + preparedAt?: string | null; + freshnessSource: 'source' | 'session' | 'request'; + }; +} + +export interface AgentRunPlanModelSnapshot { + requestedProvider?: string | null; + requestedModel?: string | null; + resolvedProvider: string; + resolvedModel: string; +} + +export interface AgentRunPlanRuntimeSnapshot { + resolvedHarness: 'lifecycle_ai_sdk'; + requestedHarness?: string | null; + sandboxRequirement: Record; + runtimeOptions: AgentRunRuntimeOptions; + approvalPolicy: AgentApprovalPolicy; +} + +export interface AgentRunPlanPromptSnapshot { + instructionRefs: string[]; + instructionAddendum?: string | null; + renderedSummary: string; + renderedHash: string; +} + +export interface AgentRunPlanCapabilitiesSnapshot { + provisionalCapabilityIds: AgentCapabilityCatalogId[]; + resolvedCapabilityAccess: AgentRunPlanResolvedCapabilityAccess[]; + selectedRuntimeToolChoiceIds?: string[]; + selectedRuntimeMcpChoiceIds?: string[]; + selectedRuntimeCapabilityIds?: AgentCapabilityCatalogId[]; + selectedRuntimeMcpConnectionRefs?: string[]; +} + +export interface AgentRunPlanResolvedCapabilityAccess { + capabilityId: AgentCapabilityCatalogId; + availability: AgentCapabilityAvailability; + allowed: boolean; + reason?: string; + runtimeCapabilityKey?: AgentCapabilityKey; + approvalMode?: AgentApprovalMode; +} + +export interface AgentRunPlanSnapshotV1 { + version: 1; + capturedAt: string; + agent: AgentRunPlanAgentSnapshot; + source: AgentRunPlanSourceSnapshot; + model: AgentRunPlanModelSnapshot; + runtime: AgentRunPlanRuntimeSnapshot; + prompt: AgentRunPlanPromptSnapshot; + capabilities: AgentRunPlanCapabilitiesSnapshot; + warnings: AgentRunPlanWarning[]; +} + +export interface AgentRunPlanPublicSummary { + version: 1; + agent: { + id: string; + label: string; + sourceKind: AgentRunPlanSourceKind; + }; + source: { + kind: AgentRunPlanSourceKind; + repoFullName?: string | null; + branch?: string | null; + buildUuid?: string | null; + namespace?: string | null; + }; + model: { + provider: string; + model: string; + }; + runtime: { + harness: 'lifecycle_ai_sdk'; + maxIterations: number | null; + }; + approval: { + defaultMode: AgentApprovalMode; + }; + capabilities: { + effective: Array<{ + capabilityId: AgentCapabilityCatalogId; + availability: AgentCapabilityAvailability; + allowed: boolean; + approvalMode?: AgentApprovalMode; + }>; + selected: { + capabilityIds: AgentCapabilityCatalogId[]; + toolChoiceIds: string[]; + mcpChoiceIds: string[]; + }; + }; + warnings: Array<{ + code: string; + message: string; + }>; +} + +export function isAgentRunPlanSnapshotV1(value: unknown): value is AgentRunPlanSnapshotV1 { + return ( + Boolean(value) && + typeof value === 'object' && + !Array.isArray(value) && + (value as { version?: unknown }).version === 1 + ); +} diff --git a/src/server/services/agent/sandboxToolCatalog.ts b/src/server/services/agent/sandboxToolCatalog.ts index f23460a0..74ee2238 100644 --- a/src/server/services/agent/sandboxToolCatalog.ts +++ b/src/server/services/agent/sandboxToolCatalog.ts @@ -14,8 +14,9 @@ * limitations under the License. */ -import type { McpDiscoveredTool } from 'server/services/ai/mcp/types'; +import type { McpDiscoveredTool } from 'server/services/agentRuntime/mcp/types'; import type { AgentSessionToolRule } from 'server/services/types/agentSessionConfig'; +import type { AgentCapabilityCatalogId } from './capabilityCatalog'; import type { AgentApprovalPolicy } from './types'; import AgentPolicyService from './PolicyService'; import { @@ -36,6 +37,7 @@ type SessionWorkspaceToolCatalogRecord = { toolName: string; runtimeToolName: string; category: SessionWorkspaceToolCategory; + catalogCapabilityId: AgentCapabilityCatalogId; order: number; adminVisibility: SessionWorkspaceToolAdminVisibility; annotations?: McpDiscoveredTool['annotations']; @@ -52,6 +54,7 @@ const SESSION_WORKSPACE_TOOL_CATALOG: readonly SessionWorkspaceToolCatalogRecord toolName: 'skills.list', runtimeToolName: 'skills.list', category: 'skills', + catalogCapabilityId: 'read_context', order: 10, adminVisibility: 'hidden', annotations: { readOnlyHint: true }, @@ -61,6 +64,7 @@ const SESSION_WORKSPACE_TOOL_CATALOG: readonly SessionWorkspaceToolCatalogRecord toolName: 'skills.learn', runtimeToolName: 'skills.learn', category: 'skills', + catalogCapabilityId: 'read_context', order: 20, adminVisibility: 'hidden', annotations: { readOnlyHint: true }, @@ -70,6 +74,7 @@ const SESSION_WORKSPACE_TOOL_CATALOG: readonly SessionWorkspaceToolCatalogRecord toolName: 'workspace.read_file', runtimeToolName: 'workspace.read_file', category: 'inspect', + catalogCapabilityId: 'read_context', order: 30, adminVisibility: 'visible', annotations: { readOnlyHint: true }, @@ -79,6 +84,7 @@ const SESSION_WORKSPACE_TOOL_CATALOG: readonly SessionWorkspaceToolCatalogRecord toolName: 'workspace.glob', runtimeToolName: 'workspace.glob', category: 'inspect', + catalogCapabilityId: 'read_context', order: 40, adminVisibility: 'visible', annotations: { readOnlyHint: true }, @@ -88,6 +94,7 @@ const SESSION_WORKSPACE_TOOL_CATALOG: readonly SessionWorkspaceToolCatalogRecord toolName: 'workspace.grep', runtimeToolName: 'workspace.grep', category: 'inspect', + catalogCapabilityId: 'read_context', order: 50, adminVisibility: 'visible', annotations: { readOnlyHint: true }, @@ -97,6 +104,7 @@ const SESSION_WORKSPACE_TOOL_CATALOG: readonly SessionWorkspaceToolCatalogRecord toolName: SESSION_WORKSPACE_READONLY_TOOL_NAME, runtimeToolName: 'workspace.exec', category: 'inspect', + catalogCapabilityId: 'read_context', order: 60, adminVisibility: 'visible', annotations: { readOnlyHint: true }, @@ -106,6 +114,7 @@ const SESSION_WORKSPACE_TOOL_CATALOG: readonly SessionWorkspaceToolCatalogRecord toolName: 'session.get_workspace_state', runtimeToolName: 'session.get_workspace_state', category: 'inspect', + catalogCapabilityId: 'read_context', order: 70, adminVisibility: 'hidden', annotations: { readOnlyHint: true }, @@ -115,6 +124,7 @@ const SESSION_WORKSPACE_TOOL_CATALOG: readonly SessionWorkspaceToolCatalogRecord toolName: 'session.list_ports', runtimeToolName: 'session.list_ports', category: 'inspect', + catalogCapabilityId: 'read_context', order: 80, adminVisibility: 'hidden', annotations: { readOnlyHint: true }, @@ -124,6 +134,7 @@ const SESSION_WORKSPACE_TOOL_CATALOG: readonly SessionWorkspaceToolCatalogRecord toolName: 'session.list_processes', runtimeToolName: 'session.list_processes', category: 'inspect', + catalogCapabilityId: 'read_context', order: 90, adminVisibility: 'hidden', annotations: { readOnlyHint: true }, @@ -133,6 +144,7 @@ const SESSION_WORKSPACE_TOOL_CATALOG: readonly SessionWorkspaceToolCatalogRecord toolName: 'session.get_service_status', runtimeToolName: 'session.get_service_status', category: 'inspect', + catalogCapabilityId: 'read_context', order: 100, adminVisibility: 'hidden', annotations: { readOnlyHint: true }, @@ -142,6 +154,7 @@ const SESSION_WORKSPACE_TOOL_CATALOG: readonly SessionWorkspaceToolCatalogRecord toolName: 'git.status', runtimeToolName: 'git.status', category: 'inspect', + catalogCapabilityId: 'workspace_git', order: 110, adminVisibility: 'visible', annotations: { readOnlyHint: true }, @@ -151,6 +164,7 @@ const SESSION_WORKSPACE_TOOL_CATALOG: readonly SessionWorkspaceToolCatalogRecord toolName: 'git.diff', runtimeToolName: 'git.diff', category: 'inspect', + catalogCapabilityId: 'workspace_git', order: 120, adminVisibility: 'visible', annotations: { readOnlyHint: true }, @@ -160,6 +174,7 @@ const SESSION_WORKSPACE_TOOL_CATALOG: readonly SessionWorkspaceToolCatalogRecord toolName: 'workspace.write_file', runtimeToolName: 'workspace.write_file', category: 'file_change', + catalogCapabilityId: 'workspace_files', order: 130, adminVisibility: 'visible', description: 'Write or overwrite a text file within the workspace.', @@ -168,6 +183,7 @@ const SESSION_WORKSPACE_TOOL_CATALOG: readonly SessionWorkspaceToolCatalogRecord toolName: 'workspace.edit_file', runtimeToolName: 'workspace.edit_file', category: 'file_change', + catalogCapabilityId: 'workspace_files', order: 140, adminVisibility: 'visible', description: 'Replace text inside a workspace file using an exact-match edit.', @@ -176,6 +192,7 @@ const SESSION_WORKSPACE_TOOL_CATALOG: readonly SessionWorkspaceToolCatalogRecord toolName: SESSION_WORKSPACE_MUTATION_TOOL_NAME, runtimeToolName: 'workspace.exec', category: 'command', + catalogCapabilityId: 'workspace_shell', order: 150, adminVisibility: 'visible', description: (serverName: string) => buildWorkspaceMutationExecDescription(serverName), @@ -184,6 +201,7 @@ const SESSION_WORKSPACE_TOOL_CATALOG: readonly SessionWorkspaceToolCatalogRecord toolName: 'git.add', runtimeToolName: 'git.add', category: 'git_change', + catalogCapabilityId: 'workspace_git', order: 160, adminVisibility: 'visible', description: 'Stage one or more paths in the workspace repository.', @@ -192,6 +210,7 @@ const SESSION_WORKSPACE_TOOL_CATALOG: readonly SessionWorkspaceToolCatalogRecord toolName: 'git.commit', runtimeToolName: 'git.commit', category: 'git_change', + catalogCapabilityId: 'workspace_git', order: 170, adminVisibility: 'visible', description: 'Create a commit from the current staged changes.', @@ -200,6 +219,7 @@ const SESSION_WORKSPACE_TOOL_CATALOG: readonly SessionWorkspaceToolCatalogRecord toolName: 'git.branch', runtimeToolName: 'git.branch', category: 'git_change', + catalogCapabilityId: 'workspace_git', order: 180, adminVisibility: 'visible', description: 'Inspect branches or create or switch a branch.', diff --git a/src/server/services/agent/sessionReadiness.ts b/src/server/services/agent/sessionReadiness.ts new file mode 100644 index 00000000..41628e19 --- /dev/null +++ b/src/server/services/agent/sessionReadiness.ts @@ -0,0 +1,49 @@ +/** + * 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 AgentSession from 'server/models/AgentSession'; +import { AgentChatStatus, AgentSessionKind, AgentWorkspaceStatus } from 'shared/constants'; + +export function canSessionAcceptMessages( + session: Pick +): boolean { + if (session.chatStatus !== AgentChatStatus.READY) { + return false; + } + + if (session.sessionKind === AgentSessionKind.CHAT) { + return true; + } + + return session.workspaceStatus === AgentWorkspaceStatus.READY; +} + +export 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.'; +} diff --git a/src/server/services/agent/systemAgentDefinitions.ts b/src/server/services/agent/systemAgentDefinitions.ts new file mode 100644 index 00000000..3ae07944 --- /dev/null +++ b/src/server/services/agent/systemAgentDefinitions.ts @@ -0,0 +1,118 @@ +/** + * 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 { AgentDefinitionContract } from './agentDefinitionTypes'; +import type { AgentCapabilitySourceKind } from './capabilityCatalog'; + +export const SYSTEM_AGENT_DEFINITION_IDS = ['system.debug', 'system.develop', 'system.freeform'] as const; + +export type SystemAgentDefinitionId = (typeof SYSTEM_AGENT_DEFINITION_IDS)[number]; + +function defineSystemAgent( + systemId: SystemAgentDefinitionId, + definition: Omit< + AgentDefinitionContract, + | 'id' + | 'version' + | 'owner' + | 'requiredCapabilityRefs' + | 'optionalCapabilityRefs' + | 'status' + | 'codeOwned' + | 'readOnly' + > +): AgentDefinitionContract { + return { + id: systemId, + version: 1, + owner: { kind: 'system' }, + ...definition, + requiredCapabilityRefs: [...definition.capabilityRefs], + optionalCapabilityRefs: [], + status: 'active', + codeOwned: true, + readOnly: true, + }; +} + +export const SYSTEM_AGENT_DEFINITIONS: Record = { + 'system.debug': defineSystemAgent('system.debug', { + name: 'Debug', + description: 'Investigate build and environment context.', + instructionRefs: ['system:debug'], + capabilityRefs: [ + 'read_context', + 'diagnostics_logs', + 'diagnostics_codefresh', + 'diagnostics_kubernetes', + 'diagnostics_database', + 'github_read', + 'github_write', + 'external_mcp_read', + 'external_mcp_write', + ], + resourcePolicy: { + sourceKinds: ['build_context_chat'], + sandboxRequired: false, + workspaceRequired: false, + }, + }), + 'system.develop': defineSystemAgent('system.develop', { + name: 'Develop', + description: 'Work in a prepared Lifecycle workspace.', + instructionRefs: ['system:develop'], + capabilityRefs: [ + 'read_context', + 'workspace_files', + 'workspace_shell', + 'workspace_git', + 'network_access', + 'preview_publish', + 'external_mcp_read', + ], + resourcePolicy: { + sourceKinds: ['workspace_session'], + sandboxRequired: true, + workspaceRequired: true, + }, + }), + 'system.freeform': defineSystemAgent('system.freeform', { + name: 'Free-form', + description: 'Answer general questions without build or workspace requirements.', + instructionRefs: ['system:freeform'], + capabilityRefs: ['read_context', 'external_mcp_read'], + resourcePolicy: { + sourceKinds: ['build_context_chat', 'workspace_session', 'freeform_chat'], + sandboxRequired: false, + workspaceRequired: false, + }, + }), +}; + +export function isSystemAgentDefinitionId(value: unknown): value is SystemAgentDefinitionId { + return typeof value === 'string' && SYSTEM_AGENT_DEFINITION_IDS.includes(value as SystemAgentDefinitionId); +} + +export function sourceKindForSystemAgentDefinitionId(id: SystemAgentDefinitionId): AgentCapabilitySourceKind { + switch (id) { + case 'system.debug': + return 'build_context_chat'; + case 'system.develop': + return 'workspace_session'; + case 'system.freeform': + return 'freeform_chat'; + } +} diff --git a/src/server/services/ai/tools/__tests__/outputLimiter.test.ts b/src/server/services/agent/tools/__tests__/outputLimiter.test.ts similarity index 100% rename from src/server/services/ai/tools/__tests__/outputLimiter.test.ts rename to src/server/services/agent/tools/__tests__/outputLimiter.test.ts diff --git a/src/server/services/ai/tools/__tests__/registry.test.ts b/src/server/services/agent/tools/__tests__/registry.test.ts similarity index 98% rename from src/server/services/ai/tools/__tests__/registry.test.ts rename to src/server/services/agent/tools/__tests__/registry.test.ts index c4173e3a..a03fbf5c 100644 --- a/src/server/services/ai/tools/__tests__/registry.test.ts +++ b/src/server/services/agent/tools/__tests__/registry.test.ts @@ -15,7 +15,7 @@ */ import { ToolRegistry } from '../registry'; -import { Tool, ToolSafetyLevel, ToolCategory } from '../../types/tool'; +import { Tool, ToolSafetyLevel, ToolCategory } from '../types'; function makeTool(overrides: Partial = {}): Tool { return { diff --git a/src/server/services/ai/tools/baseTool.ts b/src/server/services/agent/tools/baseTool.ts similarity index 97% rename from src/server/services/ai/tools/baseTool.ts rename to src/server/services/agent/tools/baseTool.ts index 1fb17f59..31623653 100644 --- a/src/server/services/ai/tools/baseTool.ts +++ b/src/server/services/agent/tools/baseTool.ts @@ -14,7 +14,7 @@ * limitations under the License. */ -import { Tool, ToolResult, ToolSafetyLevel, ToolCategory, JSONSchema, ConfirmationDetails } from '../types/tool'; +import { Tool, ToolResult, ToolSafetyLevel, ToolCategory, JSONSchema, ConfirmationDetails } from './types'; export abstract class BaseTool implements Tool { static readonly Name: string; diff --git a/src/server/services/ai/tools/codefresh/getCodefreshLogs.ts b/src/server/services/agent/tools/codefresh/getCodefreshLogs.ts similarity index 98% rename from src/server/services/ai/tools/codefresh/getCodefreshLogs.ts rename to src/server/services/agent/tools/codefresh/getCodefreshLogs.ts index 9a387697..eaf4be08 100644 --- a/src/server/services/ai/tools/codefresh/getCodefreshLogs.ts +++ b/src/server/services/agent/tools/codefresh/getCodefreshLogs.ts @@ -15,7 +15,7 @@ */ import { BaseTool } from '../baseTool'; -import { ToolResult, ToolSafetyLevel } from '../../types/tool'; +import { ToolResult, ToolSafetyLevel } from '../types'; import { getLogs } from 'server/lib/codefresh'; import { OutputLimiter } from '../outputLimiter'; diff --git a/src/server/services/ai/tools/codefresh/index.ts b/src/server/services/agent/tools/codefresh/index.ts similarity index 100% rename from src/server/services/ai/tools/codefresh/index.ts rename to src/server/services/agent/tools/codefresh/index.ts diff --git a/src/server/services/ai/tools/github/__tests__/getFile.test.ts b/src/server/services/agent/tools/github/__tests__/getFile.test.ts similarity index 100% rename from src/server/services/ai/tools/github/__tests__/getFile.test.ts rename to src/server/services/agent/tools/github/__tests__/getFile.test.ts diff --git a/src/server/services/ai/tools/github/__tests__/listDirectory.test.ts b/src/server/services/agent/tools/github/__tests__/listDirectory.test.ts similarity index 100% rename from src/server/services/ai/tools/github/__tests__/listDirectory.test.ts rename to src/server/services/agent/tools/github/__tests__/listDirectory.test.ts diff --git a/src/server/services/ai/tools/github/__tests__/updateFile.test.ts b/src/server/services/agent/tools/github/__tests__/updateFile.test.ts similarity index 90% rename from src/server/services/ai/tools/github/__tests__/updateFile.test.ts rename to src/server/services/agent/tools/github/__tests__/updateFile.test.ts index 9497d672..bafdae27 100644 --- a/src/server/services/ai/tools/github/__tests__/updateFile.test.ts +++ b/src/server/services/agent/tools/github/__tests__/updateFile.test.ts @@ -15,6 +15,7 @@ */ import { validateDiff, UpdateFileTool, MAX_LINES_CHANGED, MAX_LINES_REMOVED } from '../updateFile'; +import { GitHubClient } from '../../shared/githubClient'; describe('validateDiff', () => { it('allows identical content', () => { @@ -86,6 +87,18 @@ describe('validateDiff', () => { }); }); +describe('GitHubClient write path safety', () => { + it('matches normalized config and referenced file paths', () => { + const client = new GitHubClient(); + client.setAllowedWritePatterns(['lifecycle.yaml']); + client.setReferencedFiles(['grpc-echo/grpc-echo.Dockerfile']); + + expect(client.isFilePathAllowed('./lifecycle.yaml', 'write')).toBe(true); + expect(client.isFilePathAllowed('/grpc-echo/grpc-echo.Dockerfile', 'write')).toBe(true); + expect(client.isFilePathAllowed('secrets/token.txt', 'write')).toBe(false); + }); +}); + describe('UpdateFileTool', () => { const mockOctokit = { request: jest.fn() }; const mockGithubClient = { diff --git a/src/server/services/ai/tools/github/__tests__/updatePrLabels.test.ts b/src/server/services/agent/tools/github/__tests__/updatePrLabels.test.ts similarity index 100% rename from src/server/services/ai/tools/github/__tests__/updatePrLabels.test.ts rename to src/server/services/agent/tools/github/__tests__/updatePrLabels.test.ts diff --git a/src/server/services/ai/tools/github/getFile.ts b/src/server/services/agent/tools/github/getFile.ts similarity index 96% rename from src/server/services/ai/tools/github/getFile.ts rename to src/server/services/agent/tools/github/getFile.ts index 081d19ae..a4b72af7 100644 --- a/src/server/services/ai/tools/github/getFile.ts +++ b/src/server/services/agent/tools/github/getFile.ts @@ -15,7 +15,7 @@ */ import { BaseTool } from '../baseTool'; -import { ToolResult, ToolSafetyLevel } from '../../types/tool'; +import { ToolResult, ToolSafetyLevel } from '../types'; import { GitHubClient } from '../shared/githubClient'; import { OutputLimiter } from '../outputLimiter'; @@ -63,7 +63,7 @@ export class GetFileTool extends BaseTool { ); } - const octokit = await this.githubClient.getOctokit('ai-agent-get-file'); + const octokit = await this.githubClient.getOctokit('agent-runtime-get-file'); const response = await octokit.request(`GET /repos/${owner}/${repo}/contents/${filePath}`, { ref: branch, diff --git a/src/server/services/ai/tools/github/getIssueComment.ts b/src/server/services/agent/tools/github/getIssueComment.ts similarity index 94% rename from src/server/services/ai/tools/github/getIssueComment.ts rename to src/server/services/agent/tools/github/getIssueComment.ts index 17dfd96c..c85c9d98 100644 --- a/src/server/services/ai/tools/github/getIssueComment.ts +++ b/src/server/services/agent/tools/github/getIssueComment.ts @@ -15,7 +15,7 @@ */ import { BaseTool } from '../baseTool'; -import { ToolResult, ToolSafetyLevel } from '../../types/tool'; +import { ToolResult, ToolSafetyLevel } from '../types'; import { GitHubClient } from '../shared/githubClient'; export class GetIssueCommentTool extends BaseTool { @@ -48,7 +48,7 @@ export class GetIssueCommentTool extends BaseTool { const repo = args.repository_name as string; const commentId = args.comment_id as number; - const octokit = await this.githubClient.getOctokit('ai-agent-get-issue-comment'); + const octokit = await this.githubClient.getOctokit('agent-runtime-get-issue-comment'); const response = await octokit.request('GET /repos/{owner}/{repo}/issues/comments/{comment_id}', { owner, diff --git a/src/server/services/ai/tools/github/index.ts b/src/server/services/agent/tools/github/index.ts similarity index 100% rename from src/server/services/ai/tools/github/index.ts rename to src/server/services/agent/tools/github/index.ts diff --git a/src/server/services/ai/tools/github/listDirectory.ts b/src/server/services/agent/tools/github/listDirectory.ts similarity index 95% rename from src/server/services/ai/tools/github/listDirectory.ts rename to src/server/services/agent/tools/github/listDirectory.ts index 4f410ec0..7389af61 100644 --- a/src/server/services/ai/tools/github/listDirectory.ts +++ b/src/server/services/agent/tools/github/listDirectory.ts @@ -15,7 +15,7 @@ */ import { BaseTool } from '../baseTool'; -import { ToolResult, ToolSafetyLevel } from '../../types/tool'; +import { ToolResult, ToolSafetyLevel } from '../types'; import { GitHubClient } from '../shared/githubClient'; export class ListDirectoryTool extends BaseTool { @@ -54,7 +54,7 @@ export class ListDirectoryTool extends BaseTool { const branch = args.branch as string; const directoryPath = args.directory_path as string; - const octokit = await this.githubClient.getOctokit('ai-agent-list-directory'); + const octokit = await this.githubClient.getOctokit('agent-runtime-list-directory'); const response = await octokit.request(`GET /repos/${owner}/${repo}/contents/${directoryPath}`, { ref: branch, diff --git a/src/server/services/ai/tools/github/updateFile.ts b/src/server/services/agent/tools/github/updateFile.ts similarity index 95% rename from src/server/services/ai/tools/github/updateFile.ts rename to src/server/services/agent/tools/github/updateFile.ts index 433a42df..4c9e968e 100644 --- a/src/server/services/ai/tools/github/updateFile.ts +++ b/src/server/services/agent/tools/github/updateFile.ts @@ -15,13 +15,17 @@ */ import { BaseTool } from '../baseTool'; -import { ToolResult, ToolSafetyLevel, ConfirmationDetails } from '../../types/tool'; +import { ToolResult, ToolSafetyLevel, ConfirmationDetails } from '../types'; import { GitHubClient } from '../shared/githubClient'; // TODO: Make this configurable in db export const MAX_LINES_REMOVED = 10; export const MAX_LINES_CHANGED = 150; +function normalizeRepoPath(filePath: string): string { + return filePath.trim().replace(/^\/+/, '').replace(/^\.\//, ''); +} + export function validateDiff( oldContent: string, newContent: string @@ -107,7 +111,7 @@ export class UpdateFileTool extends BaseTool { const owner = args.repository_owner as string; const repo = args.repository_name as string; const branch = args.branch as string; - const filePath = args.file_path as string; + const filePath = normalizeRepoPath(args.file_path as string); const newContent = args.new_content as string; const commitMessage = args.commit_message as string; @@ -116,7 +120,7 @@ export class UpdateFileTool extends BaseTool { `SAFETY ERROR: File path "${filePath}" is not allowed for modification. Allowed files include: 1) Configuration files (lifecycle.yaml, lifecycle.yml) 2) Files explicitly referenced in lifecycle configuration - 3) Additional paths configured via allowedWritePatterns in the AI agent config`, + 3) Additional paths configured via allowedWritePatterns in the agent runtime config`, 'FILE_PATH_NOT_ALLOWED', false ); @@ -127,7 +131,7 @@ export class UpdateFileTool extends BaseTool { return this.createErrorResult(branchValidation.error!, 'BRANCH_VALIDATION_FAILED', false); } - const octokit = await this.githubClient.getOctokit('ai-agent-update-file'); + const octokit = await this.githubClient.getOctokit('agent-runtime-update-file'); let currentFileSha: string | undefined; let currentFileContent: string | undefined; diff --git a/src/server/services/ai/tools/github/updatePrLabels.ts b/src/server/services/agent/tools/github/updatePrLabels.ts similarity index 98% rename from src/server/services/ai/tools/github/updatePrLabels.ts rename to src/server/services/agent/tools/github/updatePrLabels.ts index 86dd4aeb..69fbde97 100644 --- a/src/server/services/ai/tools/github/updatePrLabels.ts +++ b/src/server/services/agent/tools/github/updatePrLabels.ts @@ -15,7 +15,7 @@ */ import { BaseTool } from '../baseTool'; -import { ToolResult, ToolSafetyLevel, ConfirmationDetails } from '../../types/tool'; +import { ToolResult, ToolSafetyLevel, ConfirmationDetails } from '../types'; import { GitHubClient } from '../shared/githubClient'; type LabelAction = 'add' | 'remove' | 'set'; @@ -127,7 +127,7 @@ export class UpdatePrLabelsTool extends BaseTool { return this.createErrorResult('At least one non-empty label is required', 'INVALID_LABELS', false); } - const octokit = await this.githubClient.getOctokit('ai-agent-update-pr-labels'); + const octokit = await this.githubClient.getOctokit('agent-runtime-update-pr-labels'); let currentLabels: string[] = []; if (action !== 'set') { diff --git a/src/server/services/ai/tools/index.ts b/src/server/services/agent/tools/index.ts similarity index 100% rename from src/server/services/ai/tools/index.ts rename to src/server/services/agent/tools/index.ts diff --git a/src/server/services/ai/tools/k8s/__tests__/getK8sResources.test.ts b/src/server/services/agent/tools/k8s/__tests__/getK8sResources.test.ts similarity index 100% rename from src/server/services/ai/tools/k8s/__tests__/getK8sResources.test.ts rename to src/server/services/agent/tools/k8s/__tests__/getK8sResources.test.ts diff --git a/src/server/services/ai/tools/k8s/__tests__/getPodLogs.test.ts b/src/server/services/agent/tools/k8s/__tests__/getPodLogs.test.ts similarity index 100% rename from src/server/services/ai/tools/k8s/__tests__/getPodLogs.test.ts rename to src/server/services/agent/tools/k8s/__tests__/getPodLogs.test.ts diff --git a/src/server/services/ai/tools/k8s/__tests__/patchK8sResource.test.ts b/src/server/services/agent/tools/k8s/__tests__/patchK8sResource.test.ts similarity index 100% rename from src/server/services/ai/tools/k8s/__tests__/patchK8sResource.test.ts rename to src/server/services/agent/tools/k8s/__tests__/patchK8sResource.test.ts diff --git a/src/server/services/ai/tools/k8s/getK8sResources.ts b/src/server/services/agent/tools/k8s/getK8sResources.ts similarity index 99% rename from src/server/services/ai/tools/k8s/getK8sResources.ts rename to src/server/services/agent/tools/k8s/getK8sResources.ts index 2aa77313..4ef722d3 100644 --- a/src/server/services/ai/tools/k8s/getK8sResources.ts +++ b/src/server/services/agent/tools/k8s/getK8sResources.ts @@ -15,7 +15,7 @@ */ import { BaseTool } from '../baseTool'; -import { ToolResult, ToolSafetyLevel } from '../../types/tool'; +import { ToolResult, ToolSafetyLevel } from '../types'; import { K8sClient } from '../shared/k8sClient'; import { OutputLimiter } from '../outputLimiter'; diff --git a/src/server/services/ai/tools/k8s/getLifecycleLogs.ts b/src/server/services/agent/tools/k8s/getLifecycleLogs.ts similarity index 99% rename from src/server/services/ai/tools/k8s/getLifecycleLogs.ts rename to src/server/services/agent/tools/k8s/getLifecycleLogs.ts index e2604646..21467780 100644 --- a/src/server/services/ai/tools/k8s/getLifecycleLogs.ts +++ b/src/server/services/agent/tools/k8s/getLifecycleLogs.ts @@ -15,7 +15,7 @@ */ import { BaseTool } from '../baseTool'; -import { ToolResult, ToolSafetyLevel } from '../../types/tool'; +import { ToolResult, ToolSafetyLevel } from '../types'; import { K8sClient } from '../shared/k8sClient'; import { OutputLimiter } from '../outputLimiter'; diff --git a/src/server/services/ai/tools/k8s/getPodLogs.ts b/src/server/services/agent/tools/k8s/getPodLogs.ts similarity index 98% rename from src/server/services/ai/tools/k8s/getPodLogs.ts rename to src/server/services/agent/tools/k8s/getPodLogs.ts index 8aefe89c..e2977fd3 100644 --- a/src/server/services/ai/tools/k8s/getPodLogs.ts +++ b/src/server/services/agent/tools/k8s/getPodLogs.ts @@ -15,7 +15,7 @@ */ import { BaseTool } from '../baseTool'; -import { ToolResult, ToolSafetyLevel } from '../../types/tool'; +import { ToolResult, ToolSafetyLevel } from '../types'; import { K8sClient } from '../shared/k8sClient'; import { OutputLimiter } from '../outputLimiter'; diff --git a/src/server/services/ai/tools/k8s/index.ts b/src/server/services/agent/tools/k8s/index.ts similarity index 100% rename from src/server/services/ai/tools/k8s/index.ts rename to src/server/services/agent/tools/k8s/index.ts diff --git a/src/server/services/ai/tools/k8s/patchK8sResource.ts b/src/server/services/agent/tools/k8s/patchK8sResource.ts similarity index 99% rename from src/server/services/ai/tools/k8s/patchK8sResource.ts rename to src/server/services/agent/tools/k8s/patchK8sResource.ts index 82cd32f0..dddc0ebf 100644 --- a/src/server/services/ai/tools/k8s/patchK8sResource.ts +++ b/src/server/services/agent/tools/k8s/patchK8sResource.ts @@ -15,7 +15,7 @@ */ import { BaseTool } from '../baseTool'; -import { ToolResult, ToolSafetyLevel, ConfirmationDetails } from '../../types/tool'; +import { ToolResult, ToolSafetyLevel, ConfirmationDetails } from '../types'; import { K8sClient } from '../shared/k8sClient'; export class PatchK8sResourceTool extends BaseTool { diff --git a/src/server/services/ai/tools/k8s/queryDatabase.ts b/src/server/services/agent/tools/k8s/queryDatabase.ts similarity index 98% rename from src/server/services/ai/tools/k8s/queryDatabase.ts rename to src/server/services/agent/tools/k8s/queryDatabase.ts index 5a088d88..c8976e48 100644 --- a/src/server/services/ai/tools/k8s/queryDatabase.ts +++ b/src/server/services/agent/tools/k8s/queryDatabase.ts @@ -15,7 +15,7 @@ */ import { BaseTool } from '../baseTool'; -import { ToolResult, ToolSafetyLevel } from '../../types/tool'; +import { ToolResult, ToolSafetyLevel } from '../types'; import { DatabaseClient } from '../shared/databaseClient'; export class QueryDatabaseTool extends BaseTool { diff --git a/src/server/services/ai/tools/outputLimiter.ts b/src/server/services/agent/tools/outputLimiter.ts similarity index 100% rename from src/server/services/ai/tools/outputLimiter.ts rename to src/server/services/agent/tools/outputLimiter.ts diff --git a/src/server/services/ai/tools/registry.ts b/src/server/services/agent/tools/registry.ts similarity index 97% rename from src/server/services/ai/tools/registry.ts rename to src/server/services/agent/tools/registry.ts index a00e64db..fa271339 100644 --- a/src/server/services/ai/tools/registry.ts +++ b/src/server/services/agent/tools/registry.ts @@ -14,7 +14,7 @@ * limitations under the License. */ -import { Tool, ToolResult, ToolCategory } from '../types/tool'; +import { Tool, ToolResult, ToolCategory } from './types'; export class ToolRegistry { private tools: Map = new Map(); diff --git a/src/server/services/ai/tools/shared/databaseClient.ts b/src/server/services/agent/tools/shared/databaseClient.ts similarity index 100% rename from src/server/services/ai/tools/shared/databaseClient.ts rename to src/server/services/agent/tools/shared/databaseClient.ts diff --git a/src/server/services/ai/tools/shared/githubClient.ts b/src/server/services/agent/tools/shared/githubClient.ts similarity index 84% rename from src/server/services/ai/tools/shared/githubClient.ts rename to src/server/services/agent/tools/shared/githubClient.ts index 5528189b..a2f333fa 100644 --- a/src/server/services/ai/tools/shared/githubClient.ts +++ b/src/server/services/agent/tools/shared/githubClient.ts @@ -23,12 +23,16 @@ export class GitHubClient { private excludedFilePatterns: string[] = []; private allowedWritePatterns: string[] = []; + private normalizeFilePath(filePath: string): string { + return filePath.trim().replace(/^\/+/, '').replace(/^\.\//, ''); + } + setAllowedBranch(branch: string) { this.allowedBranch = branch; } setReferencedFiles(files: string[]) { - this.referencedFiles = new Set(files.map((f) => f.toLowerCase())); + this.referencedFiles = new Set(files.map((file) => this.normalizeFilePath(file).toLowerCase())); } setExcludedFilePatterns(patterns: string[]): void { @@ -41,7 +45,7 @@ export class GitHubClient { isFileExcluded(filePath: string): boolean { if (this.excludedFilePatterns.length === 0) return false; - return picomatch.isMatch(filePath, this.excludedFilePatterns, { dot: true, nocase: true }); + return picomatch.isMatch(this.normalizeFilePath(filePath), this.excludedFilePatterns, { dot: true, nocase: true }); } getAllowedBranch(): string | null { @@ -57,15 +61,16 @@ export class GitHubClient { return true; } - const normalizedPath = filePath.toLowerCase(); + const normalizedPath = this.normalizeFilePath(filePath); + const normalizedLowerPath = normalizedPath.toLowerCase(); - if (this.referencedFiles.has(normalizedPath)) { + if (this.referencedFiles.has(normalizedLowerPath)) { return true; } if ( this.allowedWritePatterns.length > 0 && - picomatch.isMatch(filePath, this.allowedWritePatterns, { dot: true, nocase: true }) + picomatch.isMatch(normalizedPath, this.allowedWritePatterns, { dot: true, nocase: true }) ) { return true; } diff --git a/src/server/services/ai/tools/shared/k8sClient.ts b/src/server/services/agent/tools/shared/k8sClient.ts similarity index 100% rename from src/server/services/ai/tools/shared/k8sClient.ts rename to src/server/services/agent/tools/shared/k8sClient.ts diff --git a/src/server/services/ai/types/tool.ts b/src/server/services/agent/tools/types.ts similarity index 100% rename from src/server/services/ai/types/tool.ts rename to src/server/services/agent/tools/types.ts diff --git a/src/server/services/agentRuntime/config/agentRuntimeConfig.ts b/src/server/services/agentRuntime/config/agentRuntimeConfig.ts new file mode 100644 index 00000000..3e5a2732 --- /dev/null +++ b/src/server/services/agentRuntime/config/agentRuntimeConfig.ts @@ -0,0 +1,504 @@ +/** + * Copyright 2025 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 BaseService from '../../_service'; +import GlobalConfigService from '../../globalConfig'; +import { getLogger } from 'server/lib/logger'; +import type { + AgentRuntimeConfig, + AgentRuntimeRepoOverride, + AgentRuntimeRepoConfigRow, + ApprovalPolicyConfig, + CapabilityPolicyConfig, + CustomAgentCreationPolicyConfig, +} from 'server/services/types/agentRuntimeConfig'; +import { + validateAgentRuntimeConfig, + validateAgentRuntimeRepoOverride, +} from 'server/lib/validation/agentRuntimeConfigValidator'; + +const REDIS_KEY_PREFIX = 'agent_runtime_repo_config:'; + +export class AgentRuntimeConfigService extends BaseService { + private static instance: AgentRuntimeConfigService; + + private memoryCache: Map = new Map(); + private globalCache: { data: AgentRuntimeConfig | null; expiry: number } = { data: null, expiry: 0 }; + private static MEMORY_CACHE_TTL_MS = 30000; + + static getInstance(): AgentRuntimeConfigService { + if (!this.instance) { + this.instance = new AgentRuntimeConfigService(); + } + return this.instance; + } + + async getEffectiveConfig(repoFullName?: string): Promise { + const globalDefaults = await this.getGlobalDefaults(); + + if (!repoFullName) { + return globalDefaults; + } + + const normalized = repoFullName.toLowerCase(); + + const now = Date.now(); + const cached = this.memoryCache.get(normalized); + if (cached && now < cached.expiry) { + return cached.data; + } + + try { + const redisKey = `${REDIS_KEY_PREFIX}${normalized}`; + const redisValue = await this.redis.get(redisKey); + + if (redisValue) { + const repoOverride = JSON.parse(redisValue) as Partial; + const merged = this.mergeConfigs(globalDefaults, repoOverride); + this.memoryCache.set(normalized, { data: merged, expiry: now + AgentRuntimeConfigService.MEMORY_CACHE_TTL_MS }); + return merged; + } + + const row = await this.db + .knex('agent_runtime_repo_config') + .where({ repositoryFullName: normalized }) + .whereNull('deletedAt') + .first(); + + if (row) { + const config = typeof row.config === 'string' ? JSON.parse(row.config) : row.config; + await this.redis.set(redisKey, JSON.stringify(config), 'EX', 300); + const merged = this.mergeConfigs(globalDefaults, config as Partial); + this.memoryCache.set(normalized, { data: merged, expiry: now + AgentRuntimeConfigService.MEMORY_CACHE_TTL_MS }); + return merged; + } + + return globalDefaults; + } catch (error) { + getLogger().warn(`AgentRuntimeConfig: repo config lookup failed repo=${normalized} error=${error}`); + return globalDefaults; + } + } + + private async getGlobalDefaults(): Promise { + const now = Date.now(); + if (this.globalCache.data && now < this.globalCache.expiry) { + return this.globalCache.data; + } + + const config = await GlobalConfigService.getInstance().getConfig('agentRuntime'); + this.globalCache = { data: config, expiry: now + AgentRuntimeConfigService.MEMORY_CACHE_TTL_MS }; + return config; + } + + private mergeApprovalPolicy( + globalPolicy?: ApprovalPolicyConfig, + repoPolicy?: ApprovalPolicyConfig + ): ApprovalPolicyConfig | undefined { + if (!globalPolicy && !repoPolicy) { + return undefined; + } + + return { + defaultMode: repoPolicy?.defaultMode ?? globalPolicy?.defaultMode, + rules: { + ...(globalPolicy?.rules || {}), + ...(repoPolicy?.rules || {}), + }, + }; + } + + private mergeCapabilityPolicy( + globalPolicy?: CapabilityPolicyConfig, + repoPolicy?: CapabilityPolicyConfig + ): CapabilityPolicyConfig | undefined { + if (!globalPolicy && !repoPolicy) { + return undefined; + } + + return { + availability: { + ...(globalPolicy?.availability || {}), + ...(repoPolicy?.availability || {}), + }, + }; + } + + private mergeConfigs( + global: AgentRuntimeConfig, + repoOverride: Partial + ): AgentRuntimeConfig { + const result = { ...global }; + + if (repoOverride.enabled !== undefined) { + result.enabled = repoOverride.enabled; + } + if (repoOverride.maxMessagesPerSession !== undefined) { + result.maxMessagesPerSession = repoOverride.maxMessagesPerSession; + } + if (repoOverride.sessionTTL !== undefined) { + result.sessionTTL = repoOverride.sessionTTL; + } + result.approvalPolicy = this.mergeApprovalPolicy(global.approvalPolicy, repoOverride.approvalPolicy); + result.capabilityPolicy = this.mergeCapabilityPolicy(global.capabilityPolicy, repoOverride.capabilityPolicy); + if (repoOverride.systemPromptOverride !== undefined) { + result.systemPromptOverride = repoOverride.systemPromptOverride; + } + const arrayFields: (keyof AgentRuntimeRepoOverride)[] = [ + 'additiveRules', + 'excludedTools', + 'excludedFilePatterns', + 'allowedWritePatterns', + ]; + + for (const field of arrayFields) { + const globalArray = (global as any)[field] as string[] | undefined; + const repoArray = repoOverride[field] as string[] | undefined; + if (repoArray && repoArray.length > 0) { + const combined = [...(globalArray || []), ...repoArray]; + (result as any)[field] = [...new Set(combined)]; + } + } + + return result; + } + + async getGlobalConfig(): Promise { + const config = await GlobalConfigService.getInstance().getConfig('agentRuntime'); + if (!config) { + return { + enabled: false, + providers: [], + maxMessagesPerSession: 50, + sessionTTL: 3600, + allowedWritePatterns: ['lifecycle.yaml', 'lifecycle.yml'], + }; + } + return config as AgentRuntimeConfig; + } + + async setGlobalConfig(config: AgentRuntimeConfig): Promise { + validateAgentRuntimeConfig(config); + await GlobalConfigService.getInstance().setConfig('agentRuntime', config); + this.invalidateCaches(); + getLogger().info('AgentRuntimeConfig: global config updated via=api'); + } + + async updateGlobalAdditiveRules(additiveRules: string[]): Promise { + validateAgentRuntimeRepoOverride({ additiveRules }); + + const currentConfig = await this.getGlobalConfig(); + const nextConfig: AgentRuntimeConfig = { + ...currentConfig, + additiveRules, + }; + + await GlobalConfigService.getInstance().setConfig('agentRuntime', nextConfig); + this.invalidateCaches(); + getLogger().info(`AgentRuntimeConfig: global additive rules updated count=${additiveRules.length} via=api`); + + return nextConfig; + } + + async updateGlobalApprovalPolicy(approvalPolicy: ApprovalPolicyConfig): Promise { + validateAgentRuntimeRepoOverride({ approvalPolicy }); + + const currentConfig = await this.getGlobalConfig(); + const nextApprovalPolicy = this.normalizeApprovalPolicy(approvalPolicy); + const nextConfig: AgentRuntimeConfig = { + ...currentConfig, + }; + + if (nextApprovalPolicy) { + nextConfig.approvalPolicy = nextApprovalPolicy; + } else { + delete nextConfig.approvalPolicy; + } + + await GlobalConfigService.getInstance().setConfig('agentRuntime', nextConfig); + this.invalidateCaches(); + getLogger().info('AgentRuntimeConfig: global approval policy updated via=api'); + + return nextConfig; + } + + async updateGlobalCapabilityPolicy(capabilityPolicy: CapabilityPolicyConfig): Promise { + validateAgentRuntimeRepoOverride({ capabilityPolicy }); + + const currentConfig = await this.getGlobalConfig(); + const nextCapabilityPolicy = this.normalizeCapabilityPolicy(capabilityPolicy); + const nextConfig: AgentRuntimeConfig = { + ...currentConfig, + }; + + if (nextCapabilityPolicy) { + nextConfig.capabilityPolicy = nextCapabilityPolicy; + } else { + delete nextConfig.capabilityPolicy; + } + + await GlobalConfigService.getInstance().setConfig('agentRuntime', nextConfig); + this.invalidateCaches(); + getLogger().info('AgentRuntimeConfig: global capability policy updated via=api'); + + return nextConfig; + } + + async updateGlobalCustomAgentCreationPolicy( + customAgentCreationPolicy: CustomAgentCreationPolicyConfig + ): Promise { + validateAgentRuntimeConfig({ + enabled: false, + providers: [], + maxMessagesPerSession: 50, + sessionTTL: 3600, + customAgentCreationPolicy, + }); + + const currentConfig = await this.getGlobalConfig(); + const nextCustomAgentCreationPolicy = this.normalizeCustomAgentCreationPolicy(customAgentCreationPolicy); + const nextConfig: AgentRuntimeConfig = { + ...currentConfig, + }; + + if (nextCustomAgentCreationPolicy) { + nextConfig.customAgentCreationPolicy = nextCustomAgentCreationPolicy; + } else { + delete nextConfig.customAgentCreationPolicy; + } + + await GlobalConfigService.getInstance().setConfig('agentRuntime', nextConfig); + this.invalidateCaches(); + getLogger().info('AgentRuntimeConfig: global custom agent creation policy updated via=api'); + + return nextConfig; + } + + async listRepoConfigs(): Promise { + const rows = await this.db + .knex('agent_runtime_repo_config') + .whereNull('deletedAt') + .orderBy('repositoryFullName', 'asc'); + + return rows.map((row: any) => ({ + id: row.id, + repositoryFullName: row.repositoryFullName, + config: typeof row.config === 'string' ? JSON.parse(row.config) : row.config, + createdAt: row.createdAt, + updatedAt: row.updatedAt, + })); + } + + async getRepoConfig(repoFullName: string): Promise | null> { + const normalized = repoFullName.toLowerCase(); + const row = await this.db + .knex('agent_runtime_repo_config') + .where({ repositoryFullName: normalized }) + .whereNull('deletedAt') + .first(); + + if (!row) { + return null; + } + + return typeof row.config === 'string' ? JSON.parse(row.config) : row.config; + } + + async setRepoConfig(repoFullName: string, config: Partial): Promise { + const normalized = repoFullName.toLowerCase(); + validateAgentRuntimeRepoOverride(config); + await this.upsertRepoConfig(normalized, config); + } + + async updateRepoAdditiveRules( + repoFullName: string, + additiveRules: string[] + ): Promise> { + const normalized = repoFullName.toLowerCase(); + validateAgentRuntimeRepoOverride({ additiveRules }); + + const currentConfig = (await this.getRepoConfig(normalized)) ?? {}; + const nextConfig: Partial = { + ...currentConfig, + additiveRules, + }; + + await this.upsertRepoConfig(normalized, nextConfig); + getLogger().info( + `AgentRuntimeConfig: repo additive rules updated repo=${normalized} count=${additiveRules.length} via=api` + ); + + return nextConfig; + } + + async updateRepoCapabilityPolicy( + repoFullName: string, + capabilityPolicy: CapabilityPolicyConfig + ): Promise> { + const normalized = repoFullName.toLowerCase(); + validateAgentRuntimeRepoOverride({ capabilityPolicy }); + + const currentConfig = (await this.getRepoConfig(normalized)) ?? {}; + const nextCapabilityPolicy = this.normalizeCapabilityPolicy(capabilityPolicy); + const nextConfig: Partial = { + ...currentConfig, + }; + + if (nextCapabilityPolicy) { + nextConfig.capabilityPolicy = nextCapabilityPolicy; + } else { + delete nextConfig.capabilityPolicy; + } + + await this.upsertRepoConfig(normalized, nextConfig); + getLogger().info(`AgentRuntimeConfig: repo capability policy updated repo=${normalized} via=api`); + + return nextConfig; + } + + private async upsertRepoConfig( + normalizedRepoFullName: string, + config: Partial + ): Promise { + await this.db + .knex('agent_runtime_repo_config') + .insert({ + repositoryFullName: normalizedRepoFullName, + config: JSON.stringify(config), + createdAt: this.db.knex.fn.now(), + updatedAt: this.db.knex.fn.now(), + }) + .onConflict('repositoryFullName') + .merge({ + config: JSON.stringify(config), + updatedAt: this.db.knex.fn.now(), + deletedAt: null, + }); + + const redisKey = `${REDIS_KEY_PREFIX}${normalizedRepoFullName}`; + await this.redis.del(redisKey); + this.memoryCache.delete(normalizedRepoFullName); + } + + async deleteRepoConfig(repoFullName: string): Promise { + const normalized = repoFullName.toLowerCase(); + await this.db + .knex('agent_runtime_repo_config') + .where({ repositoryFullName: normalized }) + .update({ deletedAt: this.db.knex.fn.now() }); + + const redisKey = `${REDIS_KEY_PREFIX}${normalized}`; + await this.redis.del(redisKey); + this.memoryCache.delete(normalized); + } + + clearCache(repoFullName?: string): void { + if (repoFullName) { + const normalized = repoFullName.toLowerCase(); + this.memoryCache.delete(normalized); + const redisKey = `${REDIS_KEY_PREFIX}${normalized}`; + this.redis.del(redisKey); + } else { + this.invalidateCaches(); + } + } + + private normalizeApprovalPolicy(approvalPolicy?: ApprovalPolicyConfig): ApprovalPolicyConfig | undefined { + if (!approvalPolicy) { + return undefined; + } + + const normalizedRules = + approvalPolicy.rules && Object.keys(approvalPolicy.rules).length > 0 ? approvalPolicy.rules : undefined; + + if (!approvalPolicy.defaultMode && !normalizedRules) { + return undefined; + } + + return { + ...(approvalPolicy.defaultMode ? { defaultMode: approvalPolicy.defaultMode } : {}), + ...(normalizedRules ? { rules: normalizedRules } : {}), + }; + } + + private normalizeCapabilityPolicy(capabilityPolicy?: CapabilityPolicyConfig): CapabilityPolicyConfig | undefined { + if (!capabilityPolicy) { + return undefined; + } + + const normalizedAvailability = + capabilityPolicy.availability && Object.keys(capabilityPolicy.availability).length > 0 + ? capabilityPolicy.availability + : undefined; + + if (!normalizedAvailability) { + return undefined; + } + + return { + availability: normalizedAvailability, + }; + } + + private normalizeCustomAgentCreationPolicy( + customAgentCreationPolicy?: CustomAgentCreationPolicyConfig + ): CustomAgentCreationPolicyConfig | undefined { + if (!customAgentCreationPolicy) { + return undefined; + } + + const allowedUserIds = this.normalizeStringList(customAgentCreationPolicy.allowedUserIds); + const allowedGithubUsernames = this.normalizeStringList(customAgentCreationPolicy.allowedGithubUsernames, { + lowercase: true, + }); + const capabilityAvailability = + customAgentCreationPolicy.capabilityAvailability && + Object.keys(customAgentCreationPolicy.capabilityAvailability).length > 0 + ? customAgentCreationPolicy.capabilityAvailability + : undefined; + + if (!customAgentCreationPolicy.mode && !allowedUserIds && !allowedGithubUsernames && !capabilityAvailability) { + return undefined; + } + + return { + ...(customAgentCreationPolicy.mode ? { mode: customAgentCreationPolicy.mode } : {}), + ...(allowedUserIds ? { allowedUserIds } : {}), + ...(allowedGithubUsernames ? { allowedGithubUsernames } : {}), + ...(capabilityAvailability ? { capabilityAvailability } : {}), + }; + } + + private normalizeStringList(values?: string[], options: { lowercase?: boolean } = {}): string[] | undefined { + const normalized = [ + ...new Set( + (values || []) + .map((value) => value.trim()) + .filter(Boolean) + .map((value) => (options.lowercase ? value.toLowerCase() : value)) + ), + ]; + + return normalized.length > 0 ? normalized : undefined; + } + + private invalidateCaches(): void { + this.globalCache = { data: null, expiry: 0 }; + this.memoryCache.clear(); + } +} + +export default AgentRuntimeConfigService; diff --git a/src/server/services/ai/mcp/__tests__/client.test.ts b/src/server/services/agentRuntime/mcp/__tests__/client.test.ts similarity index 79% rename from src/server/services/ai/mcp/__tests__/client.test.ts rename to src/server/services/agentRuntime/mcp/__tests__/client.test.ts index dbed3083..6200f5fe 100644 --- a/src/server/services/ai/mcp/__tests__/client.test.ts +++ b/src/server/services/agentRuntime/mcp/__tests__/client.test.ts @@ -19,6 +19,7 @@ const mockClose = jest.fn(); const mockExecute = jest.fn(); const mockCreateMCPClient = jest.fn(); const mockExperimentalStdioTransport = jest.fn(); +const mockLoggerWarn = jest.fn(); jest.mock('@ai-sdk/mcp', () => ({ createMCPClient: (...args: unknown[]) => mockCreateMCPClient(...args), @@ -31,7 +32,7 @@ jest.mock('@ai-sdk/mcp/mcp-stdio', () => ({ })); jest.mock('server/lib/logger', () => ({ - getLogger: () => ({ info: jest.fn(), warn: jest.fn(), error: jest.fn() }), + getLogger: () => ({ info: jest.fn(), warn: (...args: unknown[]) => mockLoggerWarn(...args), error: jest.fn() }), })); import { McpClientManager } from '../client'; @@ -79,6 +80,33 @@ describe('McpClientManager', () => { ); }); + it('redacts transport secrets from uncaught MCP client errors', async () => { + mockCreateMCPClient.mockImplementationOnce(async (options: { onUncaughtError?: (error: Error) => void }) => { + options.onUncaughtError?.( + new Error( + 'uncaught Authorization=Bearer transport-secret query=query/secret+value encoded=query%2Fsecret%2Bvalue' + ) + ); + return { + listTools: mockListTools, + close: mockClose, + toolsFromDefinitions: jest.fn(() => ({ + inspectItem: { execute: mockExecute }, + })), + }; + }); + + await manager.connect({ + type: 'http', + url: 'https://mcp.example.com/v1/mcp?api_key=query/secret+value', + headers: { Authorization: 'Bearer transport-secret' }, + }); + + expect(mockLoggerWarn).toHaveBeenCalledWith( + 'MCP client uncaught error: uncaught Authorization=****** query=****** encoded=******' + ); + }); + it('wraps stdio transport with the AI SDK stdio helper', async () => { await manager.connect({ type: 'stdio', diff --git a/src/server/services/ai/mcp/__tests__/config.test.ts b/src/server/services/agentRuntime/mcp/__tests__/config.test.ts similarity index 52% rename from src/server/services/ai/mcp/__tests__/config.test.ts rename to src/server/services/agentRuntime/mcp/__tests__/config.test.ts index b183d200..926f5543 100644 --- a/src/server/services/ai/mcp/__tests__/config.test.ts +++ b/src/server/services/agentRuntime/mcp/__tests__/config.test.ts @@ -49,7 +49,7 @@ jest.mock('server/models/McpServerConfig', () => { }); import McpServerConfig from 'server/models/McpServerConfig'; -import { McpConfigService } from '../config'; +import { McpConfigService, redactMcpConfigSecrets, redactSharedConfigSecrets } from '../config'; const MockModel = McpServerConfig as any; @@ -88,6 +88,62 @@ describe('McpConfigService', () => { mockListDecryptedConnectionsByScopes.mockResolvedValue(new Map()); }); + describe('redactSharedConfigSecrets', () => { + it('redacts all shared config value sections before returning config records', () => { + const result = redactSharedConfigSecrets({ + sharedConfig: { + headers: { Authorization: 'Bearer sample-token' }, + query: { api_key: 'sample-query-token' }, + env: { SAMPLE_API_TOKEN: 'sample-env-token' }, + defaultArgs: { token: 'sample-arg-token' }, + }, + }); + + expect(result.sharedConfig).toEqual({ + headers: { Authorization: '******' }, + query: { api_key: '******' }, + env: { SAMPLE_API_TOKEN: '******' }, + defaultArgs: { token: '******' }, + }); + }); + }); + + describe('redactMcpConfigSecrets', () => { + it('redacts HTTP transport headers before returning config records', () => { + const result = redactMcpConfigSecrets({ + transport: { + type: 'http', + url: 'https://mcp.example.com/v1/mcp?api_key=sample-query-token&workspace=sample-workspace', + headers: { Authorization: 'Bearer sample-token' }, + }, + }); + + expect(result.transport).toEqual({ + type: 'http', + url: 'https://mcp.example.com/v1/mcp?api_key=******&workspace=******', + headers: { Authorization: '******' }, + }); + }); + + it('redacts stdio transport env before returning config records', () => { + const result = redactMcpConfigSecrets({ + transport: { + type: 'stdio', + command: 'npx', + args: ['-y', 'sample-mcp-server'], + env: { SAMPLE_API_TOKEN: 'sample-token' }, + }, + }); + + expect(result.transport).toEqual({ + type: 'stdio', + command: 'npx', + args: ['-y', 'sample-mcp-server'], + env: { SAMPLE_API_TOKEN: '******' }, + }); + }); + }); + describe('create', () => { it('creates a shared connector definition with transport and shared discovered tools', async () => { const { insert } = makeQueryResult(undefined); @@ -111,6 +167,32 @@ describe('McpConfigService', () => { ); }); + it('redacts shared secrets from shared discovery validation errors', async () => { + makeQueryResult(undefined); + mockConnect.mockRejectedValueOnce( + new Error( + 'connect failed Authorization=Bearer header-secret query=query/secret+value encoded=query%2Fsecret%2Bvalue env=env-secret arg=arg-secret' + ) + ); + + await expect( + service.create({ + slug: 'sample-connector', + name: 'Sample connector', + scope: 'global', + transport: { type: 'http', url: 'https://mcp.example.com/v1/mcp' }, + sharedConfig: { + headers: { Authorization: 'Bearer header-secret' }, + query: { api_key: 'query/secret+value' }, + env: { SAMPLE_API_TOKEN: 'env-secret' }, + defaultArgs: { token: 'arg-secret' }, + }, + }) + ).rejects.toThrow( + 'MCP server connectivity validation failed: connect failed Authorization=****** query=****** encoded=****** env=****** arg=******' + ); + }); + it('defers shared discovery for connectors that require user configuration', async () => { const { insert } = makeQueryResult(undefined); const inserted = { id: 1, slug: 'sample-connector' }; @@ -277,29 +359,29 @@ describe('McpConfigService', () => { describe('resolveSessionPodServersForRepo', () => { it('returns only stdio connectors and preserves compiled per-user env bindings', async () => { const stdioConfig = { - slug: 'figma', - name: 'Figma', + slug: 'sample-stdio', + name: 'Sample stdio', scope: 'global', transport: { type: 'stdio', command: 'npx', - args: ['-y', 'figma-developer-mcp', '--stdio'], + args: ['-y', 'sample-mcp-server', '--stdio'], env: {}, }, sharedConfig: {}, authConfig: { mode: 'user-fields', schema: { - fields: [{ key: 'figmaToken', label: 'Figma token', required: true, inputType: 'password' }], - bindings: [{ target: 'env', key: 'FIGMA_API_KEY', fieldKey: 'figmaToken' }], + fields: [{ key: 'apiToken', label: 'API token', required: true, inputType: 'password' }], + bindings: [{ target: 'env', key: 'SAMPLE_API_TOKEN', fieldKey: 'apiToken' }], }, }, timeout: 45000, sharedDiscoveredTools: [], }; const httpConfig = { - slug: 'jira', - name: 'Jira', + slug: 'sample-http', + name: 'Sample HTTP', scope: 'global', transport: { type: 'http', url: 'https://mcp.example.com/v1/mcp', headers: {} }, sharedConfig: {}, @@ -323,13 +405,13 @@ describe('McpConfigService', () => { mockListDecryptedConnectionsByScopes.mockResolvedValue( new Map([ [ - 'global:figma', + 'global:sample-stdio', { state: { type: 'fields', - values: { figmaToken: 'figma-pat-token' }, + values: { apiToken: 'sample-api-token' }, }, - discoveredTools: [{ name: 'get_design_context', inputSchema: {} }], + discoveredTools: [{ name: 'inspectItem', inputSchema: {} }], validationError: null, validatedAt: '2026-04-06T16:00:00.000Z', updatedAt: '2026-04-06T16:00:00.000Z', @@ -345,29 +427,30 @@ describe('McpConfigService', () => { expect(result).toEqual([ { - slug: 'figma', - name: 'Figma', + slug: 'sample-stdio', + name: 'Sample stdio', + scope: 'global', transport: { type: 'stdio', command: 'npx', - args: ['-y', 'figma-developer-mcp', '--stdio'], + args: ['-y', 'sample-mcp-server', '--stdio'], env: { - FIGMA_API_KEY: 'figma-pat-token', + SAMPLE_API_TOKEN: 'sample-api-token', }, }, timeout: 45000, defaultArgs: {}, env: { - FIGMA_API_KEY: 'figma-pat-token', + SAMPLE_API_TOKEN: 'sample-api-token', }, - discoveredTools: [{ name: 'get_design_context', inputSchema: {} }], + discoveredTools: [{ name: 'inspectItem', inputSchema: {} }], }, ]); }); }); describe('update', () => { - it('preserves redacted shared header secrets when updating a connector', async () => { + it('preserves redacted shared config secrets when updating a connector', async () => { const existing = { id: 1, slug: 'sample-connector', @@ -375,11 +458,26 @@ describe('McpConfigService', () => { scope: 'global', description: 'Original description', preset: null, - transport: { type: 'http', url: 'https://mcp.example.com/v1/mcp', headers: {} }, + transport: { + type: 'http', + url: 'https://mcp.example.com/v1/mcp?api_key=transport-query-secret&workspace=sample-workspace', + headers: { + Authorization: 'Bearer transport-token', + }, + }, sharedConfig: { headers: { Authorization: 'Bearer top-secret-token', }, + query: { + api_key: 'top-secret-query-token', + }, + env: { + SAMPLE_API_TOKEN: 'top-secret-env-token', + }, + defaultArgs: { + token: 'top-secret-arg-token', + }, }, authConfig: { mode: 'none' }, enabled: true, @@ -405,10 +503,26 @@ describe('McpConfigService', () => { const result = await service.update('sample-connector', 'global', { description: 'Updated description', + transport: { + type: 'http', + url: 'https://mcp.example.com/v1/mcp?api_key=******&workspace=sample-workspace', + headers: { + Authorization: '******', + }, + }, sharedConfig: { headers: { Authorization: '******', }, + query: { + api_key: '******', + }, + env: { + SAMPLE_API_TOKEN: '******', + }, + defaultArgs: { + token: '******', + }, }, }); @@ -421,12 +535,236 @@ describe('McpConfigService', () => { 1, expect.objectContaining({ description: 'Updated description', + transport: { + type: 'http', + url: 'https://mcp.example.com/v1/mcp?api_key=transport-query-secret&workspace=sample-workspace', + headers: { + Authorization: 'Bearer transport-token', + }, + }, sharedConfig: expect.objectContaining({ headers: { Authorization: 'Bearer top-secret-token', }, + query: { + api_key: 'top-secret-query-token', + }, + env: { + SAMPLE_API_TOKEN: 'top-secret-env-token', + }, + defaultArgs: { + token: 'top-secret-arg-token', + }, + }), + }) + ); + }); + + it('requires transport secrets to be re-entered when the connector target changes', async () => { + const existing = { + id: 1, + slug: 'sample-connector', + name: 'Sample connector', + scope: 'global', + description: 'Original description', + preset: null, + transport: { + type: 'http', + url: 'https://mcp.example.com/v1/mcp?api_key=transport-query-secret', + headers: { + Authorization: 'Bearer transport-token', + }, + }, + sharedConfig: {}, + authConfig: { mode: 'none' }, + enabled: true, + timeout: 30000, + sharedDiscoveredTools: [{ name: 'inspectItem', inputSchema: {} }], + }; + const patchAndFetchById = jest.fn(); + + MockModel.query.mockReturnValueOnce({ + where: jest.fn().mockReturnValue({ + whereNull: jest.fn().mockReturnValue({ + first: jest.fn().mockResolvedValue(existing), + }), + }), + }); + + await expect( + service.update('sample-connector', 'global', { + transport: { + type: 'http', + url: 'https://other.example.com/v1/mcp?api_key=******', + headers: { + Authorization: '******', + }, + }, + }) + ).rejects.toThrow('Re-enter MCP transport secrets when changing the MCP transport target'); + + expect(mockConnect).not.toHaveBeenCalled(); + expect(patchAndFetchById).not.toHaveBeenCalled(); + }); + + it('requires shared secrets to be re-entered when the connector target changes', async () => { + const existing = { + id: 1, + slug: 'sample-connector', + name: 'Sample connector', + scope: 'global', + description: 'Original description', + preset: null, + transport: { + type: 'http', + url: 'https://mcp.example.com/v1/mcp', + headers: {}, + }, + sharedConfig: { + headers: { + Authorization: 'Bearer shared-token', + }, + }, + authConfig: { mode: 'none' }, + enabled: true, + timeout: 30000, + sharedDiscoveredTools: [{ name: 'inspectItem', inputSchema: {} }], + }; + const patchAndFetchById = jest.fn(); + + MockModel.query.mockReturnValueOnce({ + where: jest.fn().mockReturnValue({ + whereNull: jest.fn().mockReturnValue({ + first: jest.fn().mockResolvedValue(existing), + }), + }), + }); + + await expect( + service.update('sample-connector', 'global', { + transport: { + type: 'http', + url: 'https://other.example.com/v1/mcp', + headers: {}, + }, + }) + ).rejects.toThrow('Re-enter MCP shared secrets when changing the MCP transport target'); + + expect(mockConnect).not.toHaveBeenCalled(); + expect(patchAndFetchById).not.toHaveBeenCalled(); + }); + + it('requires transport secrets to be re-entered when the connector transport type changes', async () => { + const existing = { + id: 1, + slug: 'sample-connector', + name: 'Sample connector', + scope: 'global', + description: 'Original description', + preset: null, + transport: { + type: 'http', + url: 'https://mcp.example.com/v1/mcp?api_key=transport-query-secret', + headers: {}, + }, + sharedConfig: {}, + authConfig: { mode: 'none' }, + enabled: true, + timeout: 30000, + sharedDiscoveredTools: [{ name: 'inspectItem', inputSchema: {} }], + }; + const patchAndFetchById = jest.fn(); + + MockModel.query.mockReturnValueOnce({ + where: jest.fn().mockReturnValue({ + whereNull: jest.fn().mockReturnValue({ + first: jest.fn().mockResolvedValue(existing), + }), + }), + }); + + await expect( + service.update('sample-connector', 'global', { + transport: { + type: 'stdio', + command: 'sample-mcp', + args: ['--stdio'], + env: { + SAMPLE_API_TOKEN: '******', + }, + }, + }) + ).rejects.toThrow('Re-enter MCP transport secrets when changing the MCP transport target'); + + expect(mockConnect).not.toHaveBeenCalled(); + expect(patchAndFetchById).not.toHaveBeenCalled(); + }); + + it('preserves redacted stdio transport env secrets when updating a connector', async () => { + const existing = { + id: 1, + slug: 'sample-stdio', + name: 'Sample stdio', + scope: 'global', + description: 'Original description', + preset: null, + transport: { + type: 'stdio', + command: 'npx', + args: ['-y', 'sample-mcp-server'], + env: { + SAMPLE_API_TOKEN: 'top-secret-env-token', + }, + }, + sharedConfig: {}, + authConfig: { mode: 'none' }, + enabled: true, + timeout: 30000, + sharedDiscoveredTools: [{ name: 'inspectItem', inputSchema: {} }], + }; + const patchAndFetchById = jest.fn().mockResolvedValue({ + ...existing, + description: 'Updated description', + }); + + MockModel.query + .mockReturnValueOnce({ + where: jest.fn().mockReturnValue({ + whereNull: jest.fn().mockReturnValue({ + first: jest.fn().mockResolvedValue(existing), + }), }), }) + .mockReturnValueOnce({ + patchAndFetchById, + }); + + await service.update('sample-stdio', 'global', { + description: 'Updated description', + transport: { + type: 'stdio', + command: 'npx', + args: ['-y', 'sample-mcp-server'], + env: { + SAMPLE_API_TOKEN: '******', + }, + }, + }); + + expect(mockConnect).not.toHaveBeenCalled(); + expect(patchAndFetchById).toHaveBeenCalledWith( + 1, + expect.objectContaining({ + description: 'Updated description', + transport: { + type: 'stdio', + command: 'npx', + args: ['-y', 'sample-mcp-server'], + env: { + SAMPLE_API_TOKEN: 'top-secret-env-token', + }, + }, + }) ); }); }); diff --git a/src/server/services/ai/mcp/__tests__/oauthFlow.test.ts b/src/server/services/agentRuntime/mcp/__tests__/oauthFlow.test.ts similarity index 100% rename from src/server/services/ai/mcp/__tests__/oauthFlow.test.ts rename to src/server/services/agentRuntime/mcp/__tests__/oauthFlow.test.ts diff --git a/src/server/services/ai/mcp/__tests__/oauthProvider.test.ts b/src/server/services/agentRuntime/mcp/__tests__/oauthProvider.test.ts similarity index 100% rename from src/server/services/ai/mcp/__tests__/oauthProvider.test.ts rename to src/server/services/agentRuntime/mcp/__tests__/oauthProvider.test.ts diff --git a/src/server/services/ai/mcp/__tests__/runtimeConfig.test.ts b/src/server/services/agentRuntime/mcp/__tests__/runtimeConfig.test.ts similarity index 100% rename from src/server/services/ai/mcp/__tests__/runtimeConfig.test.ts rename to src/server/services/agentRuntime/mcp/__tests__/runtimeConfig.test.ts diff --git a/src/server/services/ai/mcp/client.ts b/src/server/services/agentRuntime/mcp/client.ts similarity index 71% rename from src/server/services/ai/mcp/client.ts rename to src/server/services/agentRuntime/mcp/client.ts index 5cedea42..d48fdd6a 100644 --- a/src/server/services/ai/mcp/client.ts +++ b/src/server/services/agentRuntime/mcp/client.ts @@ -20,6 +20,8 @@ import type { McpDiscoveredTool, McpResolvedTransportConfig, McpToolAnnotations const DEFAULT_HANDSHAKE_TIMEOUT_MS = 5000; const DEFAULT_CALL_TIMEOUT_MS = 30000; +const REDACTED_MCP_SECRET = '******'; +const MIN_SECRET_REDACTION_LENGTH = 4; type ListToolsDefinitions = Awaited>; type ExperimentalStdioMCPModule = typeof import('@ai-sdk/mcp/dist/mcp-stdio'); @@ -87,6 +89,72 @@ function createTransport(transport: McpResolvedTransportConfig) { return transport; } +function addSecretValue(secrets: Set, value: unknown): void { + if (typeof value !== 'string') { + return; + } + + const secret = value.trim(); + if (secret.length < MIN_SECRET_REDACTION_LENGTH || secret === REDACTED_MCP_SECRET) { + return; + } + + secrets.add(secret); + secrets.add(encodeURIComponent(secret)); + secrets.add(new URLSearchParams({ value: secret }).toString().slice('value='.length)); +} + +function collectRawQuerySecretValues(url: string, secrets: Set): void { + const queryStart = url.indexOf('?'); + if (queryStart === -1) { + return; + } + + const hashStart = url.indexOf('#', queryStart); + const query = url.slice(queryStart + 1, hashStart === -1 ? undefined : hashStart); + for (const part of query.split('&')) { + if (!part) { + continue; + } + + const valueStart = part.indexOf('='); + const rawValue = valueStart === -1 ? '' : part.slice(valueStart + 1); + if (!rawValue) { + continue; + } + + addSecretValue(secrets, rawValue); + try { + addSecretValue(secrets, decodeURIComponent(rawValue)); + } catch { + // Ignore malformed percent-encoding; the raw value is still redacted. + } + } +} + +function sanitizeTransportErrorMessage(error: unknown, transport: McpResolvedTransportConfig): string { + const message = error instanceof Error ? error.message : String(error); + const secrets = new Set(); + + if (transport.type === 'http' || transport.type === 'sse') { + Object.values(transport.headers || {}).forEach((value) => addSecretValue(secrets, value)); + collectRawQuerySecretValues(transport.url, secrets); + try { + new URL(transport.url).searchParams.forEach((value) => addSecretValue(secrets, value)); + } catch { + // Ignore invalid URLs here; transport validation happens before connection. + } + } + + if (transport.type === 'stdio') { + Object.values(transport.env || {}).forEach((value) => addSecretValue(secrets, value)); + } + + return Array.from(secrets) + .sort((a, b) => b.length - a.length) + .reduce((current, secret) => current.split(secret).join(REDACTED_MCP_SECRET), message); +} + export class McpClientManager { private client: MCPClient | null = null; private toolDefinitions: ListToolsDefinitions | null = null; @@ -101,7 +169,7 @@ export class McpClientManager { name: 'lifecycle', version: '1.0.0', onUncaughtError: (error) => { - getLogger().warn(`MCP client uncaught error: ${error instanceof Error ? error.message : String(error)}`); + getLogger().warn(`MCP client uncaught error: ${sanitizeTransportErrorMessage(error, transport)}`); }, }), handshakeTimeoutMs, diff --git a/src/server/services/ai/mcp/config.ts b/src/server/services/agentRuntime/mcp/config.ts similarity index 51% rename from src/server/services/ai/mcp/config.ts rename to src/server/services/agentRuntime/mcp/config.ts index 2bd7d335..6d81257f 100644 --- a/src/server/services/ai/mcp/config.ts +++ b/src/server/services/agentRuntime/mcp/config.ts @@ -36,10 +36,12 @@ import { usesSessionWorkspaceGatewayExecution } from './sessionPod'; import type { AgentMcpConnection, CreateMcpServerConfigInput, + McpCompiledConnectionConfig, McpAuthConfig, McpDiscoveredTool, McpResolvedTransportConfig, McpServerConfigRecord, + McpSharedConnectionConfig, McpTransportConfig, ResolvedMcpServer, UpdateMcpServerConfigInput, @@ -49,6 +51,26 @@ const SLUG_REGEX = /^[a-z0-9][a-z0-9-]*[a-z0-9]$|^[a-z0-9]$/; const MAX_SLUG_LENGTH = 100; const VALIDATION_TIMEOUT_MS = 5000; const REDACTED_SHARED_SECRET = '******'; +const MIN_SECRET_REDACTION_LENGTH = 4; +const NON_SECRET_REDACTION_VALUES = new Set([ + 'bearer', + 'basic', + 'false', + 'http', + 'https', + 'none', + 'null', + 'oauth', + 'true', +]); +const SHARED_SECRET_SECTIONS: (keyof McpSharedConnectionConfig)[] = ['headers', 'query', 'env', 'defaultArgs']; + +export type McpErrorRedactionSource = { + values?: Record | null; + compiledConfig?: Partial | null; + transport?: McpResolvedTransportConfig | McpTransportConfig | null; + extraSecrets?: unknown[]; +}; function buildConnectionKey(scope: string, slug: string): string { return `${scope}:${slug}`; @@ -84,40 +106,486 @@ function buildOAuthCallbackUrl(slug: string, scope: string): string { return url.toString(); } -export function redactSharedConfigSecrets } | null }>( - config: T -): T { - if (!config.sharedConfig?.headers || typeof config.sharedConfig.headers !== 'object') { +function addSecretValue(secrets: Set, value: unknown): void { + if (typeof value !== 'string') { + return; + } + + const secret = value.trim(); + if ( + secret.length < MIN_SECRET_REDACTION_LENGTH || + secret === REDACTED_SHARED_SECRET || + NON_SECRET_REDACTION_VALUES.has(secret.toLowerCase()) + ) { + return; + } + + secrets.add(secret); + secrets.add(encodeURIComponent(secret)); + const encoded = new URLSearchParams({ value: secret }).toString().slice('value='.length); + secrets.add(encoded); +} + +function collectUnknownSecretValues(value: unknown, secrets: Set): void { + if (typeof value === 'string') { + addSecretValue(secrets, value); + return; + } + + if (!value || typeof value !== 'object') { + return; + } + + if (Array.isArray(value)) { + for (const item of value) { + collectUnknownSecretValues(item, secrets); + } + return; + } + + for (const item of Object.values(value)) { + collectUnknownSecretValues(item, secrets); + } +} + +function collectRecordSecretValues(values: Record | undefined | null, secrets: Set): void { + if (!values) { + return; + } + + for (const value of Object.values(values)) { + addSecretValue(secrets, value); + } +} + +function collectCompiledConfigSecretValues( + config: Partial | undefined | null, + secrets: Set +): void { + if (!config) { + return; + } + + collectRecordSecretValues(config.headers, secrets); + collectRecordSecretValues(config.query, secrets); + collectRecordSecretValues(config.env, secrets); + collectRecordSecretValues(config.defaultArgs, secrets); +} + +function collectTransportSecretValues( + transport: McpResolvedTransportConfig | McpTransportConfig | undefined | null, + secrets: Set +): void { + if (!transport) { + return; + } + + if (transport.type === 'http' || transport.type === 'sse') { + collectRecordSecretValues(transport.headers, secrets); + collectRawQuerySecretValues(transport.url, secrets); + try { + const parsed = new URL(transport.url); + parsed.searchParams.forEach((value) => { + addSecretValue(secrets, value); + }); + } catch { + // Ignore non-URL transport strings; normalized transports should already be valid URLs. + } + return; + } + + if (transport.type === 'stdio') { + collectRecordSecretValues(transport.env, secrets); + } +} + +function collectRawQuerySecretValues(url: string, secrets: Set): void { + const queryStart = url.indexOf('?'); + if (queryStart === -1) { + return; + } + + const hashStart = url.indexOf('#', queryStart); + const query = url.slice(queryStart + 1, hashStart === -1 ? undefined : hashStart); + for (const part of query.split('&')) { + if (!part) { + continue; + } + + const valueStart = part.indexOf('='); + const rawValue = valueStart === -1 ? '' : part.slice(valueStart + 1); + if (!rawValue) { + continue; + } + + addSecretValue(secrets, rawValue); + try { + addSecretValue(secrets, decodeURIComponent(rawValue)); + } catch { + // Ignore malformed percent-encoding; the raw value is still redacted. + } + } +} + +function collectMcpSecretValues(sources: McpErrorRedactionSource[]): Set { + const secrets = new Set(); + + for (const source of sources) { + collectUnknownSecretValues(source.values, secrets); + collectCompiledConfigSecretValues(source.compiledConfig, secrets); + collectTransportSecretValues(source.transport, secrets); + for (const secret of source.extraSecrets || []) { + collectUnknownSecretValues(secret, secrets); + } + } + + return secrets; +} + +function redactSecretValues(value: string, secrets: Set): string { + return Array.from(secrets) + .sort((a, b) => b.length - a.length) + .reduce((current, secret) => current.split(secret).join(REDACTED_SHARED_SECRET), value); +} + +function sanitizeUnknownValue(value: unknown, secrets: Set): unknown { + if (typeof value === 'string') { + return redactSecretValues(value, secrets); + } + + if (!value || typeof value !== 'object') { + return value; + } + + if (Array.isArray(value)) { + return value.map((item) => sanitizeUnknownValue(item, secrets)); + } + + return Object.fromEntries(Object.entries(value).map(([key, item]) => [key, sanitizeUnknownValue(item, secrets)])); +} + +export function sanitizeMcpErrorMessage(error: unknown, sources: McpErrorRedactionSource[] = []): string { + const message = error instanceof Error ? error.message : String(error); + return redactSecretValues(message, collectMcpSecretValues(sources)); +} + +export function sanitizeMcpResult(result: T, sources: McpErrorRedactionSource[] = []): T { + return sanitizeUnknownValue(result, collectMcpSecretValues(sources)) as T; +} + +export function redactSharedConfigSecrets(config: T): T { + const sharedConfig = config.sharedConfig ? { ...config.sharedConfig } : undefined; + let changed = false; + + if (sharedConfig) { + for (const section of SHARED_SECRET_SECTIONS) { + const values = sharedConfig[section]; + if (!values || typeof values !== 'object') { + continue; + } + + sharedConfig[section] = Object.fromEntries( + Object.keys(values).map((key) => [key, REDACTED_SHARED_SECRET]) + ) as Record; + changed = true; + } + } + + if (!changed) { return config; } return { ...config, - sharedConfig: { - ...config.sharedConfig, - headers: Object.fromEntries(Object.keys(config.sharedConfig.headers).map((key) => [key, REDACTED_SHARED_SECRET])), + ...(sharedConfig ? { sharedConfig } : {}), + }; +} + +function redactSecretRecord(values: Record | undefined): Record | undefined { + return values ? Object.fromEntries(Object.keys(values).map((key) => [key, REDACTED_SHARED_SECRET])) : values; +} + +function redactTransportUrlQuery(url: string): string { + const queryStart = url.indexOf('?'); + if (queryStart === -1) { + return url; + } + + const hashStart = url.indexOf('#', queryStart); + const base = url.slice(0, queryStart); + const query = url.slice(queryStart + 1, hashStart === -1 ? undefined : hashStart); + const hash = hashStart === -1 ? '' : url.slice(hashStart); + if (!query) { + return url; + } + + const redactedQuery = query + .split('&') + .map((part) => { + if (!part) { + return part; + } + + const valueStart = part.indexOf('='); + const key = valueStart === -1 ? part : part.slice(0, valueStart); + return `${key}=${REDACTED_SHARED_SECRET}`; + }) + .join('&'); + + return `${base}?${redactedQuery}${hash}`; +} + +export function redactMcpConfigSecrets< + T extends { sharedConfig?: McpSharedConnectionConfig | null; transport?: McpTransportConfig | null } +>(config: T): T { + const redacted = redactSharedConfigSecrets(config); + if (!redacted.transport) { + return redacted; + } + + if (redacted.transport.type === 'http' || redacted.transport.type === 'sse') { + const url = redactTransportUrlQuery(redacted.transport.url); + if (!redacted.transport.headers && url === redacted.transport.url) { + return redacted; + } + + return { + ...redacted, + transport: { + ...redacted.transport, + url, + ...(redacted.transport.headers ? { headers: redactSecretRecord(redacted.transport.headers) } : {}), + }, + }; + } + + if (!redacted.transport.env) { + return redacted; + } + + return { + ...redacted, + transport: { + ...redacted.transport, + env: redactSecretRecord(redacted.transport.env), }, }; } -function restoreRedactedSharedHeaders( +function restoreRedactedSharedConfig( nextSharedConfig: McpServerConfigRecord['sharedConfig'], existingSharedConfig: McpServerConfigRecord['sharedConfig'] ): McpServerConfigRecord['sharedConfig'] { - if (!nextSharedConfig?.headers || !existingSharedConfig?.headers) { + if (!nextSharedConfig || !existingSharedConfig) { return nextSharedConfig; } let changed = false; - const headers = { ...nextSharedConfig.headers }; - for (const [key, value] of Object.entries(headers)) { - if (value === REDACTED_SHARED_SECRET && existingSharedConfig.headers[key]) { - headers[key] = existingSharedConfig.headers[key]; + const sharedConfig = { ...nextSharedConfig }; + + for (const section of SHARED_SECRET_SECTIONS) { + const nextValues = sharedConfig[section]; + const existingValues = existingSharedConfig[section]; + if (!nextValues || !existingValues) { + continue; + } + + const restoredValues = { ...nextValues }; + for (const [key, value] of Object.entries(restoredValues)) { + if (value === REDACTED_SHARED_SECRET && existingValues[key]) { + restoredValues[key] = existingValues[key]; + changed = true; + } + } + + sharedConfig[section] = restoredValues; + } + + return changed ? sharedConfig : nextSharedConfig; +} + +function restoreRedactedSecretRecord( + nextValues: Record | undefined, + existingValues: Record | undefined +): Record | undefined { + if (!nextValues || !existingValues) { + return nextValues; + } + + let changed = false; + const restoredValues = { ...nextValues }; + for (const [key, value] of Object.entries(restoredValues)) { + if (value === REDACTED_SHARED_SECRET && existingValues[key]) { + restoredValues[key] = existingValues[key]; + changed = true; + } + } + + return changed ? restoredValues : nextValues; +} + +function parseUrlQueryParts(url: string): { base: string; query: string; hash: string } | null { + const queryStart = url.indexOf('?'); + if (queryStart === -1) { + return null; + } + + const hashStart = url.indexOf('#', queryStart); + return { + base: url.slice(0, queryStart), + query: url.slice(queryStart + 1, hashStart === -1 ? undefined : hashStart), + hash: hashStart === -1 ? '' : url.slice(hashStart), + }; +} + +function restoreRedactedTransportUrlQuery(nextUrl: string, existingUrl: string): string { + const nextParts = parseUrlQueryParts(nextUrl); + const existingParts = parseUrlQueryParts(existingUrl); + if (!nextParts || !existingParts || !nextParts.query || !existingParts.query) { + return nextUrl; + } + + const existingValuesByKey = new Map(); + for (const part of existingParts.query.split('&')) { + const valueStart = part.indexOf('='); + if (valueStart === -1) { + continue; + } + + const key = part.slice(0, valueStart); + const value = part.slice(valueStart + 1); + existingValuesByKey.set(key, [...(existingValuesByKey.get(key) || []), value]); + } + + let changed = false; + const restoredQuery = nextParts.query + .split('&') + .map((part) => { + const valueStart = part.indexOf('='); + if (valueStart === -1) { + return part; + } + + const key = part.slice(0, valueStart); + const value = part.slice(valueStart + 1); + if (value !== REDACTED_SHARED_SECRET) { + return part; + } + + const existingValues = existingValuesByKey.get(key); + const existingValue = existingValues?.shift(); + if (!existingValue) { + return part; + } + changed = true; + return `${key}=${existingValue}`; + }) + .join('&'); + + return changed ? `${nextParts.base}?${restoredQuery}${nextParts.hash}` : nextUrl; +} + +function transportTargetChanged(nextTransport: McpTransportConfig, existingTransport: McpTransportConfig): boolean { + if (nextTransport.type !== existingTransport.type) { + return true; + } + + if (nextTransport.type === 'http' || nextTransport.type === 'sse') { + if (existingTransport.type !== 'http' && existingTransport.type !== 'sse') { + return true; + } + + try { + const nextUrl = new URL(nextTransport.url); + const existingUrl = new URL(existingTransport.url); + return ( + nextUrl.protocol !== existingUrl.protocol || + nextUrl.host !== existingUrl.host || + nextUrl.pathname !== existingUrl.pathname + ); + } catch { + return nextTransport.url.split('?')[0] !== existingTransport.url.split('?')[0]; + } + } + + if (existingTransport.type !== 'stdio') { + return true; + } + + return ( + nextTransport.command !== existingTransport.command || + JSON.stringify(nextTransport.args || []) !== JSON.stringify(existingTransport.args || []) + ); +} + +function recordContainsRedactedSecret(values: Record | undefined): boolean { + return !!values && Object.values(values).some((value) => value === REDACTED_SHARED_SECRET); +} + +function transportContainsRedactedSecret(transport: McpTransportConfig): boolean { + if (transport.type === 'http' || transport.type === 'sse') { + let redactedQueryValue = false; + new URLSearchParams(parseUrlQueryParts(transport.url)?.query || '').forEach((value) => { + if (value === REDACTED_SHARED_SECRET) { + redactedQueryValue = true; + } + }); + return recordContainsRedactedSecret(transport.headers) || redactedQueryValue; + } + + return recordContainsRedactedSecret(transport.env); +} + +function sharedConfigContainsRedactedSecret(sharedConfig: McpSharedConnectionConfig): boolean { + return SHARED_SECRET_SECTIONS.some((section) => recordContainsRedactedSecret(sharedConfig[section])); +} + +function sharedConfigContainsSecretValue(sharedConfig: McpSharedConnectionConfig): boolean { + return SHARED_SECRET_SECTIONS.some((section) => { + const values = sharedConfig[section]; + return !!values && Object.values(values).length > 0; + }); +} + +function restoreRedactedTransport( + nextTransport: McpTransportConfig, + existingTransport: McpTransportConfig +): McpTransportConfig { + if (nextTransport.type !== existingTransport.type) { + if (transportContainsRedactedSecret(nextTransport)) { + throw new Error('Re-enter MCP transport secrets when changing the MCP transport target'); + } + + return nextTransport; + } + + if (nextTransport.type === 'http' || nextTransport.type === 'sse') { + if (existingTransport.type !== 'http' && existingTransport.type !== 'sse') { + return nextTransport; + } + + if (transportTargetChanged(nextTransport, existingTransport) && transportContainsRedactedSecret(nextTransport)) { + throw new Error('Re-enter MCP transport secrets when changing the MCP transport target'); } + + const headers = restoreRedactedSecretRecord(nextTransport.headers, existingTransport.headers); + const url = restoreRedactedTransportUrlQuery(nextTransport.url, existingTransport.url); + return headers === nextTransport.headers && url === nextTransport.url + ? nextTransport + : { ...nextTransport, url, headers }; } - return changed ? { ...nextSharedConfig, headers } : nextSharedConfig; + if (existingTransport.type !== 'stdio') { + return nextTransport; + } + + if (transportTargetChanged(nextTransport, existingTransport) && transportContainsRedactedSecret(nextTransport)) { + throw new Error('Re-enter MCP transport secrets when changing the MCP transport target'); + } + + const env = restoreRedactedSecretRecord(nextTransport.env, existingTransport.env); + return env === nextTransport.env ? nextTransport : { ...nextTransport, env }; } export class McpConfigService { @@ -191,10 +659,28 @@ export class McpConfigService { } const nextPreset = input.preset ?? config.preset ?? null; - const nextTransport = input.transport ? normalizeTransportConfig(input.transport) : config.transport; + const currentTransport = normalizeTransportConfig(config.transport); + const normalizedInputTransport = input.transport ? normalizeTransportConfig(input.transport) : undefined; + const targetChanged = normalizedInputTransport + ? transportTargetChanged(normalizedInputTransport, currentTransport) + : false; + const nextTransport = input.transport + ? restoreRedactedTransport(normalizedInputTransport as McpTransportConfig, currentTransport) + : currentTransport; const currentSharedConfig = normalizeSharedConnectionConfig(config.sharedConfig); - const nextSharedConfig = restoreRedactedSharedHeaders( - input.sharedConfig ? normalizeSharedConnectionConfig(input.sharedConfig) : currentSharedConfig, + const normalizedInputSharedConfig = input.sharedConfig + ? normalizeSharedConnectionConfig(input.sharedConfig) + : undefined; + if ( + targetChanged && + (normalizedInputSharedConfig + ? sharedConfigContainsRedactedSecret(normalizedInputSharedConfig) + : sharedConfigContainsSecretValue(currentSharedConfig)) + ) { + throw new Error('Re-enter MCP shared secrets when changing the MCP transport target'); + } + const nextSharedConfig = restoreRedactedSharedConfig( + normalizedInputSharedConfig || currentSharedConfig, currentSharedConfig ); const nextAuthConfig = @@ -202,7 +688,7 @@ export class McpConfigService { ? resolveAuthConfig(input.authConfig, nextPreset) : resolveAuthConfig(config.authConfig, nextPreset); - const transportChanged = JSON.stringify(nextTransport) !== JSON.stringify(config.transport); + const transportChanged = JSON.stringify(nextTransport) !== JSON.stringify(currentTransport); const sharedConfigChanged = JSON.stringify(nextSharedConfig) !== JSON.stringify(currentSharedConfig); const authConfigChanged = JSON.stringify(nextAuthConfig) !== JSON.stringify(resolveAuthConfig(config.authConfig, config.preset)); @@ -258,7 +744,7 @@ export class McpConfigService { description: config.description ?? null, scope: config.scope, preset: config.preset ?? null, - transport: config.transport, + transport: redactMcpConfigSecrets({ transport: config.transport }).transport || config.transport, sharedConfig: redactSharedConfigSecrets({ sharedConfig: config.sharedConfig || {} }).sharedConfig || {}, authConfig, connectionRequired, @@ -319,6 +805,7 @@ export class McpConfigService { return [ { + scope: config.scope, slug: config.slug, name: config.name, transport: resolvedTransport, @@ -362,6 +849,7 @@ export class McpConfigService { const compiledConfig = mergeCompiledConnectionConfig(config.sharedConfig || {}, undefined); return [ { + scope: config.scope, slug: config.slug, name: config.name, transport: resolvedTransport, @@ -383,6 +871,7 @@ export class McpConfigService { return [ { + scope: config.scope, slug: config.slug, name: config.name, transport: resolvedTransport, @@ -486,7 +975,7 @@ export class McpConfigService { try { return await this.discoverTools(resolvedTransport, options.timeoutMs ?? VALIDATION_TIMEOUT_MS); } catch (error) { - const message = error instanceof Error ? error.message : String(error); + const message = sanitizeMcpErrorMessage(error, [{ compiledConfig, transport: resolvedTransport }]); throw new Error(`MCP server connectivity validation failed: ${message}`); } } diff --git a/src/server/services/ai/mcp/connectionConfig.ts b/src/server/services/agentRuntime/mcp/connectionConfig.ts similarity index 100% rename from src/server/services/ai/mcp/connectionConfig.ts rename to src/server/services/agentRuntime/mcp/connectionConfig.ts diff --git a/src/server/services/ai/mcp/oauthFlow.ts b/src/server/services/agentRuntime/mcp/oauthFlow.ts similarity index 100% rename from src/server/services/ai/mcp/oauthFlow.ts rename to src/server/services/agentRuntime/mcp/oauthFlow.ts diff --git a/src/server/services/ai/mcp/oauthProvider.ts b/src/server/services/agentRuntime/mcp/oauthProvider.ts similarity index 100% rename from src/server/services/ai/mcp/oauthProvider.ts rename to src/server/services/agentRuntime/mcp/oauthProvider.ts diff --git a/src/server/services/ai/mcp/presets.ts b/src/server/services/agentRuntime/mcp/presets.ts similarity index 100% rename from src/server/services/ai/mcp/presets.ts rename to src/server/services/agentRuntime/mcp/presets.ts diff --git a/src/server/services/ai/mcp/runtimeConfig.ts b/src/server/services/agentRuntime/mcp/runtimeConfig.ts similarity index 100% rename from src/server/services/ai/mcp/runtimeConfig.ts rename to src/server/services/agentRuntime/mcp/runtimeConfig.ts diff --git a/src/server/services/ai/mcp/sessionPod.ts b/src/server/services/agentRuntime/mcp/sessionPod.ts similarity index 100% rename from src/server/services/ai/mcp/sessionPod.ts rename to src/server/services/agentRuntime/mcp/sessionPod.ts diff --git a/src/server/services/ai/mcp/types.ts b/src/server/services/agentRuntime/mcp/types.ts similarity index 99% rename from src/server/services/ai/mcp/types.ts rename to src/server/services/agentRuntime/mcp/types.ts index a31b54b0..cbc6399e 100644 --- a/src/server/services/ai/mcp/types.ts +++ b/src/server/services/agentRuntime/mcp/types.ts @@ -244,6 +244,7 @@ export const MCP_ERROR_CODES = { } as const; export interface ResolvedMcpServer { + scope: string; slug: string; name: string; transport: McpResolvedTransportConfig; diff --git a/src/server/services/ai/utils/modelTransformation.ts b/src/server/services/agentRuntime/models/modelTransformation.ts similarity index 100% rename from src/server/services/ai/utils/modelTransformation.ts rename to src/server/services/agentRuntime/models/modelTransformation.ts diff --git a/src/server/services/agentSession.ts b/src/server/services/agentSession.ts index 29d53f5b..a80698d8 100644 --- a/src/server/services/agentSession.ts +++ b/src/server/services/agentSession.ts @@ -90,12 +90,12 @@ import { toPublicAgentSessionStartupFailure, } from 'server/lib/agentSession/startupFailureState'; import { BuildEnvironmentVariables } from 'server/lib/buildEnvVariables'; -import { McpConfigService } from 'server/services/ai/mcp/config'; +import { McpConfigService } from 'server/services/agentRuntime/mcp/config'; import GlobalConfigService from './globalConfig'; import { SESSION_POD_MCP_CONFIG_SECRET_KEY, serializeSessionWorkspaceGatewayServers, -} from 'server/services/ai/mcp/sessionPod'; +} from 'server/services/agentRuntime/mcp/sessionPod'; import AgentPrewarmService from './agentPrewarm'; import AgentSessionConfigService from './agentSessionConfig'; import AgentChatSessionService from './agent/ChatSessionService'; @@ -104,6 +104,7 @@ import AgentProviderRegistry from './agent/ProviderRegistry'; import AgentSandboxService from './agent/SandboxService'; import AgentSourceService from './agent/SourceService'; import { buildSessionWorkspacePromptLines } from './agent/sandboxToolCatalog'; +import { canSessionAcceptMessages, getSessionMessageBlockReason } from './agent/sessionReadiness'; import { loadAgentSessionServiceCandidates, resolveRequestedAgentSessionServices, @@ -153,37 +154,6 @@ 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 @@ -760,6 +730,7 @@ export interface CreateSessionOptions { workspacePath?: string; workDir?: string | null; }>; + provider?: string; model?: string; environmentSkillRefs?: AgentSessionSkillRef[]; repoUrl?: string; @@ -1021,6 +992,8 @@ export default class AgentSessionService { const pvcName = `agent-pvc-${session.uuid.slice(0, 8)}`; const apiKeySecretName = `agent-secret-${session.uuid.slice(0, 8)}`; const redis = RedisClient.getInstance().getRedis(); + const workspaceRepos = session.workspaceRepos || []; + const selectedServices = session.selectedServices || []; if (session.namespace && session.namespace !== namespace) { await deleteNamespace(session.namespace).catch(() => {}); @@ -1041,8 +1014,8 @@ export default class AgentSessionService { status: 'active', chatStatus: AgentChatStatus.READY, workspaceStatus: AgentWorkspaceStatus.PROVISIONING, - workspaceRepos: [], - selectedServices: [], + workspaceRepos, + selectedServices, devModeSnapshots: {}, forwardedAgentSecretProviders: [], skillPlan: session.skillPlan || EMPTY_AGENT_SESSION_SKILL_PLAN, @@ -1055,7 +1028,15 @@ export default class AgentSessionService { await AgentSandboxService.recordSessionSandboxState(provisioningSession, { workspaceStorage }); try { - const sessionPodMcpConfigJson = serializeSessionWorkspaceGatewayServers([]); + const primaryWorkspaceRepo = workspaceRepos.find((repo) => repo.primary) || workspaceRepos[0]; + const sessionPodServers = primaryWorkspaceRepo?.repo + ? await new McpConfigService().resolveSessionPodServersForRepo( + primaryWorkspaceRepo.repo, + undefined, + opts.userIdentity || null + ) + : []; + const sessionPodMcpConfigJson = serializeSessionWorkspaceGatewayServers(sessionPodServers); const [, , agentServiceAccountName, useGvisor] = await Promise.all([ createAgentPvc(namespace, pvcName, workspaceStorage.storageSize, undefined, workspaceStorage.accessMode), createAgentApiKeySecret( @@ -1085,7 +1066,7 @@ export default class AgentSessionService { apiKeySecretName, hasGitHubToken: Boolean(opts.githubToken), workspacePath: SESSION_WORKSPACE_ROOT, - workspaceRepos: [], + workspaceRepos, skillPlan: session.skillPlan || EMPTY_AGENT_SESSION_SKILL_PLAN, forwardedAgentEnv: {}, forwardedAgentSecretRefs: [], @@ -1158,8 +1139,6 @@ export default class AgentSessionService { namespace: null, podName: null, pvcName: null, - workspaceRepos: [], - selectedServices: [], devModeSnapshots: {}, } as unknown as Partial); @@ -1274,6 +1253,7 @@ export default class AgentSessionService { const sessionKind = resolveSessionKindFromBuildKind(buildKind); const podName = buildAgentSessionPodName(sessionUuid, opts.buildUuid); const apiKeySecretName = `agent-secret-${sessionUuid.slice(0, 8)}`; + const requestedProvider = opts.provider?.trim() || undefined; const requestedModelId = opts.model?.trim() || undefined; const devModeSnapshots: SessionSnapshotMap = {}; const enabledDevModeDeployIds: number[] = []; @@ -1317,6 +1297,7 @@ export default class AgentSessionService { const preflightStartedAt = Date.now(); const selection = await AgentProviderRegistry.resolveSelection({ repoFullName: primaryWorkspaceRepo?.repo, + requestedProvider, requestedModelId, }); resolvedModelId = selection.modelId; @@ -1392,7 +1373,11 @@ export default class AgentSessionService { skillPlan, } as unknown as Partial); - await AgentSourceService.createSessionSource(createdSession, { trx, workspaceStorage }); + await AgentSourceService.createSessionSource(createdSession, { + trx, + workspaceStorage, + defaultProvider: selection.provider, + }); await AgentSandboxService.recordSessionSandboxState(createdSession, { trx, workspaceStorage }); return createdSession; }); @@ -1674,7 +1659,10 @@ export default class AgentSessionService { .catch(() => null); if (failedSession) { await Promise.all([ - AgentSourceService.createSessionSource(failedSession, { workspaceStorage }).catch(() => {}), + AgentSourceService.createSessionSource(failedSession, { + workspaceStorage, + defaultProvider: selection.provider, + }).catch(() => {}), AgentSandboxService.recordSessionSandboxState(failedSession, { workspaceStorage }).catch(() => {}), ]); } @@ -2083,14 +2071,14 @@ export default class AgentSessionService { return resolvedConfiguredPrompt; } - if (!session.namespace) { + if (!session.namespace && !session.buildUuid) { return resolvedConfiguredPrompt; } try { const context = await resolveAgentSessionPromptContext({ sessionDbId: session.id, - namespace: session.namespace, + namespace: session.namespace || null, buildUuid: session.buildUuid, }); const toolLines = session.namespace diff --git a/src/server/services/agentSessionConfig.ts b/src/server/services/agentSessionConfig.ts index ba8f7e40..f8fc1022 100644 --- a/src/server/services/agentSessionConfig.ts +++ b/src/server/services/agentSessionConfig.ts @@ -25,6 +25,7 @@ import { validateAgentSessionRuntimeSettings, } from 'server/lib/validation/agentSessionConfigValidator'; import type { + AgentCapabilityInventoryEntry, AgentSessionControlPlaneConfigValue, AgentSessionRuntimeSettingsValue, AgentSessionToolInventoryEntry, @@ -32,6 +33,8 @@ import type { AgentSessionToolRuleSelection, EffectiveAgentSessionControlPlaneConfig, } from './types/agentSessionConfig'; +import AgentRuntimeConfigService from 'server/services/agentRuntime/config/agentRuntimeConfig'; +import type { CapabilityPolicyConfig } from './types/agentRuntimeConfig'; import type { GlobalConfig, AgentSessionDefaults } from './types/globalConfig'; import { DEFAULT_AGENT_SESSION_CONTROL_PLANE_APPEND_SYSTEM_PROMPT, @@ -40,9 +43,10 @@ import { DEFAULT_AGENT_SESSION_WORKSPACE_TOOL_DISCOVERY_TIMEOUT_MS, DEFAULT_AGENT_SESSION_WORKSPACE_TOOL_EXECUTION_TIMEOUT_MS, } from 'server/lib/agentSession/runtimeConfig'; -import { McpConfigService } from 'server/services/ai/mcp/config'; -import { normalizeAuthConfig, requiresUserConnection } from 'server/services/ai/mcp/connectionConfig'; +import { McpConfigService } from 'server/services/agentRuntime/mcp/config'; +import { normalizeAuthConfig, requiresUserConnection } from 'server/services/agentRuntime/mcp/connectionConfig'; import AgentPolicyService from './agent/PolicyService'; +import { listAgentCapabilityCatalogEntries, type AgentCapabilityCatalogId } from './agent/capabilityCatalog'; import { buildAgentToolKey, CHAT_PUBLISH_HTTP_TOOL_NAME, @@ -51,7 +55,7 @@ import { SESSION_WORKSPACE_SERVER_NAME, SESSION_WORKSPACE_SERVER_SLUG, } from './agent/toolKeys'; -import type { McpDiscoveredTool } from 'server/services/ai/mcp/types'; +import type { McpDiscoveredTool } from 'server/services/agentRuntime/mcp/types'; import { getSessionWorkspaceToolSortKey, listAdminVisibleSessionWorkspaceToolCatalog, @@ -409,6 +413,30 @@ function toRuleSelection(toolRules: AgentSessionToolRule[], toolKey: string): Ag return toolRules.find((rule) => rule.toolKey === toolKey)?.mode || 'inherit'; } +function catalogCapabilityForTool(entry: AgentSessionToolInventoryEntry): AgentCapabilityCatalogId { + if (entry.sourceType === 'mcp') { + return entry.capabilityKey === 'external_mcp_read' ? 'external_mcp_read' : 'external_mcp_write'; + } + + if (entry.toolName === CHAT_PUBLISH_HTTP_TOOL_NAME) { + return 'preview_publish'; + } + + if (entry.toolName === 'workspace.write_file' || entry.toolName === 'workspace.edit_file') { + return 'workspace_files'; + } + + if (entry.toolName === 'workspace.exec_mutation') { + return 'workspace_shell'; + } + + if (entry.toolName.startsWith('git.')) { + return 'workspace_git'; + } + + return 'read_context'; +} + function hasConfigValues(config: Partial): boolean { return Boolean( normalizeOptionalString(config.systemPrompt) || @@ -731,6 +759,80 @@ export default class AgentSessionConfigService extends BaseService { }); } + async listCapabilityInventory(scope: string): Promise { + const repoFullName = scope === 'global' ? undefined : normalizeRepoFullName(scope); + const agentRuntimeConfigService = AgentRuntimeConfigService.getInstance(); + const [globalAgentConfig, repoAgentConfig, effectiveAgentConfig, approvalPolicy, toolInventory] = await Promise.all( + [ + agentRuntimeConfigService.getGlobalConfig(), + repoFullName ? agentRuntimeConfigService.getRepoConfig(repoFullName) : Promise.resolve(null), + agentRuntimeConfigService.getEffectiveConfig(repoFullName), + AgentPolicyService.getEffectivePolicy(repoFullName), + this.listToolInventory(scope), + ] + ); + const globalPolicy = globalAgentConfig.capabilityPolicy; + const repoPolicy = repoAgentConfig?.capabilityPolicy as CapabilityPolicyConfig | undefined; + const activePolicy = repoFullName ? repoPolicy : globalPolicy; + const effectivePolicy = effectiveAgentConfig.capabilityPolicy; + const toolsByCapability = new Map(); + + for (const tool of toolInventory) { + const capabilityId = catalogCapabilityForTool(tool); + const existing = toolsByCapability.get(capabilityId) || []; + existing.push(tool); + toolsByCapability.set(capabilityId, existing); + } + + return listAgentCapabilityCatalogEntries().map((entry) => { + const configuredAvailability = activePolicy?.availability?.[entry.id]; + const inheritedAvailability = repoFullName + ? globalPolicy?.availability?.[entry.id] || entry.defaultAvailability + : undefined; + const effectiveAvailability = + effectivePolicy?.availability?.[entry.id] || inheritedAvailability || entry.defaultAvailability; + const resolvedAccess = AgentPolicyService.resolveCapabilityAccess({ + capabilityId: entry.id, + capabilityPolicy: { availability: { [entry.id]: effectiveAvailability } }, + approvalPolicy, + definitionOwnerKind: 'system', + sourceKind: entry.sourceKinds?.[0], + }); + const mappedTools = toolsByCapability.get(entry.id) || []; + + return { + capabilityId: entry.id, + label: entry.label, + description: entry.description, + category: entry.category, + defaultAvailability: entry.defaultAvailability, + ...(configuredAvailability ? { configuredAvailability } : {}), + ...(inheritedAvailability ? { inheritedAvailability } : {}), + effectiveAvailability, + approvalMode: resolvedAccess.approvalMode || entry.defaultApprovalMode, + ...(entry.runtimeCapabilityKey ? { runtimeCapabilityKey: entry.runtimeCapabilityKey } : {}), + userSelectable: entry.userSelectable, + toolCount: mappedTools.length || entry.toolKeys?.length || 0, + resourceCount: entry.resourceGrants?.length || 0, + resourceGrants: [...(entry.resourceGrants || [])], + tools: mappedTools.map((tool) => ({ + toolKey: tool.toolKey, + toolName: tool.toolName, + description: tool.description, + serverSlug: tool.serverSlug, + serverName: tool.serverName, + sourceType: tool.sourceType, + sourceScope: tool.sourceScope, + })), + ...(effectiveAvailability === 'disabled' || + effectiveAvailability === 'system_only' || + effectiveAvailability === 'admin_only' + ? { blockedReason: effectiveAvailability } + : {}), + }; + }); + } + private async listDiscoveredToolsForDefinition( config: Pick ): Promise { diff --git a/src/server/services/ai/context/__tests__/contextSummarizer.test.ts b/src/server/services/ai/context/__tests__/contextSummarizer.test.ts deleted file mode 100644 index 0bea0256..00000000 --- a/src/server/services/ai/context/__tests__/contextSummarizer.test.ts +++ /dev/null @@ -1,596 +0,0 @@ -/** - * Copyright 2025 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 { summarizeLifecycleYaml } from '../contextSummarizer'; -import { countTokens } from '../../prompts/tokenCounter'; - -describe('summarizeLifecycleYaml', () => { - describe('parsing and structure', () => { - it('parses valid YAML with environment and services', () => { - const yaml = ` -version: "1.0.0" -environment: - autoDeploy: true - defaultServices: - - name: api - - name: web - optionalServices: - - name: redis -services: - - name: api-service - helm: - repository: org/api - branchName: main - chart: - name: org-chart - - name: db - docker: - dockerImage: postgres - defaultTag: "14" - - name: ext - externalHttp: - defaultInternalHostname: ext.internal - defaultPublicUrl: https://ext.example.com -`; - const result = summarizeLifecycleYaml(yaml); - expect(result.parsed).toBe(true); - expect(result.serviceCount).toBe(3); - expect(result.text).toContain('ENVIRONMENT:'); - expect(result.text).toContain('SERVICES (3):'); - expect(result.text).toContain('api-service'); - expect(result.text).toContain('(helm)'); - expect(result.text).toContain('db'); - expect(result.text).toContain('(docker)'); - expect(result.text).toContain('ext'); - expect(result.text).toContain('(externalHttp)'); - }); - - it('includes version in environment section', () => { - const yaml = ` -version: "1.0.0" -services: [] -`; - const result = summarizeLifecycleYaml(yaml); - expect(result.text).toContain('Version: 1.0.0'); - }); - - it('includes autoDeploy in environment section', () => { - const yaml = ` -environment: - autoDeploy: true -services: [] -`; - const result = summarizeLifecycleYaml(yaml); - expect(result.text).toContain('AutoDeploy: true'); - }); - - it('lists default and optional services', () => { - const yaml = ` -environment: - defaultServices: - - name: api - - name: web - optionalServices: - - name: redis -services: [] -`; - const result = summarizeLifecycleYaml(yaml); - expect(result.text).toContain('Default Services: api, web'); - expect(result.text).toContain('Optional Services: redis'); - }); - }); - - describe('service type detection', () => { - it('detects helm type', () => { - const yaml = ` -services: - - name: svc - helm: - chart: - name: my-chart -`; - const result = summarizeLifecycleYaml(yaml); - expect(result.text).toContain('(helm)'); - }); - - it('detects codefresh type', () => { - const yaml = ` -services: - - name: svc - codefresh: - repository: org/repo - branchName: main -`; - const result = summarizeLifecycleYaml(yaml); - expect(result.text).toContain('(codefresh)'); - }); - - it('detects github type', () => { - const yaml = ` -services: - - name: svc - github: - repository: org/repo - branchName: main - docker: - app: - dockerfilePath: Dockerfile -`; - const result = summarizeLifecycleYaml(yaml); - expect(result.text).toContain('(github)'); - }); - - it('detects docker type', () => { - const yaml = ` -services: - - name: svc - docker: - dockerImage: nginx - defaultTag: latest -`; - const result = summarizeLifecycleYaml(yaml); - expect(result.text).toContain('(docker)'); - }); - - it('detects externalHttp type', () => { - const yaml = ` -services: - - name: svc - externalHttp: - defaultInternalHostname: host.internal - defaultPublicUrl: https://host.example.com -`; - const result = summarizeLifecycleYaml(yaml); - expect(result.text).toContain('(externalHttp)'); - }); - - it('prefers helm over codefresh when both present', () => { - const yaml = ` -services: - - name: svc - helm: - chart: - name: my-chart - codefresh: - repository: org/repo - branchName: main -`; - const result = summarizeLifecycleYaml(yaml); - expect(result.text).toContain('(helm)'); - expect(result.text).not.toContain('(codefresh)'); - }); - }); - - describe('helm service extraction', () => { - it('extracts repo, chart, valueFiles, docker config', () => { - const yaml = ` -services: - - name: api - helm: - repository: org/api - branchName: main - chart: - name: org-chart - repoUrl: https://charts.example.com - valueFiles: - - helm/values.yaml - - helm/overrides.yaml - docker: - builder: - engine: buildkit - app: - dockerfilePath: Dockerfile - ports: - - 8080 -`; - const result = summarizeLifecycleYaml(yaml); - expect(result.text).toContain('Repo: org/api @ main'); - expect(result.text).toContain('Chart: org-chart (repoUrl: https://charts.example.com)'); - expect(result.text).toContain('ValueFiles: helm/values.yaml, helm/overrides.yaml'); - expect(result.text).toContain('Docker: buildkit | dockerfilePath: Dockerfile'); - expect(result.text).toContain('Ports: 8080'); - }); - - it('handles helm service with minimal config', () => { - const yaml = ` -services: - - name: minimal-helm - helm: - chart: - name: my-chart -`; - const result = summarizeLifecycleYaml(yaml); - expect(result.parsed).toBe(true); - expect(result.text).toContain('minimal-helm'); - expect(result.text).toContain('(helm)'); - }); - }); - - describe('docker service extraction', () => { - it('extracts image and ports', () => { - const yaml = ` -services: - - name: db - docker: - dockerImage: postgres - defaultTag: "14" - ports: - - 5432 -`; - const result = summarizeLifecycleYaml(yaml); - expect(result.text).toContain('Image: postgres:14'); - expect(result.text).toContain('Ports: 5432'); - }); - }); - - describe('config pointers', () => { - it('aggregates referenced files from all services', () => { - const yaml = ` -services: - - name: api - helm: - repository: org/api - branchName: main - chart: - name: org-chart - valueFiles: - - helm/api-values.yaml - docker: - app: - dockerfilePath: Dockerfile.api - - name: web - helm: - repository: org/web - branchName: main - chart: - name: org-chart - valueFiles: - - helm/web-values.yaml - docker: - app: - dockerfilePath: Dockerfile.web -`; - const result = summarizeLifecycleYaml(yaml); - expect(result.text).toContain('CONFIG POINTERS:'); - expect(result.text).toContain('Referenced Files:'); - expect(result.text).toContain('helm/api-values.yaml'); - expect(result.text).toContain('Dockerfile.api'); - expect(result.text).toContain('helm/web-values.yaml'); - expect(result.text).toContain('Dockerfile.web'); - }); - - it('deduplicates referenced files', () => { - const yaml = ` -services: - - name: api - helm: - chart: - name: chart - docker: - app: - dockerfilePath: Dockerfile - - name: web - helm: - chart: - name: chart - docker: - app: - dockerfilePath: Dockerfile -`; - const result = summarizeLifecycleYaml(yaml); - const pointerSection = result.text.split('CONFIG POINTERS:')[1]; - const matches = pointerSection.match(/Dockerfile/g); - expect(matches).toHaveLength(1); - }); - - it('omits CONFIG POINTERS when no files referenced', () => { - const yaml = ` -services: - - name: db - docker: - dockerImage: postgres - defaultTag: "14" -`; - const result = summarizeLifecycleYaml(yaml); - expect(result.text).not.toContain('CONFIG POINTERS'); - }); - }); - - describe('dependency extraction', () => { - it('extracts deploymentDependsOn', () => { - const yaml = ` -services: - - name: api - deploymentDependsOn: - - postgres - - redis - helm: - chart: - name: chart -`; - const result = summarizeLifecycleYaml(yaml); - expect(result.text).toContain('DependsOn: postgres, redis'); - }); - - it('extracts requires', () => { - const yaml = ` -services: - - name: api - requires: - - name: cache - helm: - chart: - name: chart -`; - const result = summarizeLifecycleYaml(yaml); - expect(result.text).toContain('DependsOn: cache'); - }); - }); - - describe('fallback behavior', () => { - it('returns raw YAML on empty string', () => { - const result = summarizeLifecycleYaml(''); - expect(result.parsed).toBe(false); - expect(result.text).toBe(''); - expect(result.serviceCount).toBe(0); - }); - - it('returns raw YAML on malformed YAML', () => { - const input = '{{invalid: yaml::'; - const result = summarizeLifecycleYaml(input); - expect(result.parsed).toBe(false); - expect(result.text).toBe(input); - expect(result.serviceCount).toBe(0); - }); - - it('returns raw YAML when parsed result is not an object', () => { - const result = summarizeLifecycleYaml('"just a string"'); - expect(result.parsed).toBe(false); - }); - - it('handles YAML with no services array', () => { - const yaml = ` -environment: - autoDeploy: true -`; - const result = summarizeLifecycleYaml(yaml); - expect(result.parsed).toBe(true); - expect(result.serviceCount).toBe(0); - expect(result.text).toContain('SERVICES (0)'); - }); - - it('handles YAML with empty services array', () => { - const yaml = ` -services: [] -`; - const result = summarizeLifecycleYaml(yaml); - expect(result.parsed).toBe(true); - expect(result.serviceCount).toBe(0); - }); - }); - - describe('compression ratio', () => { - const realisticYaml = ` -version: "1.0.0" -environment: - autoDeploy: true - githubDeployments: true - useGithubStatusComment: true - enabledFeatures: - - native-build - - auto-deploy - defaultServices: - - name: api - repository: org/api-service - branch: main - - name: web - repository: org/web-app - branch: main - - name: worker - repository: org/api-service - branch: main - optionalServices: - - name: redis - - name: postgres - - name: monitoring - webhooks: - onSuccess: - - url: https://hooks.slack.com/services/T00/B00/xxx - method: POST - onFailure: - - url: https://hooks.slack.com/services/T00/B00/yyy - method: POST -services: - - name: api-service - appShort: api - deploymentDependsOn: - - postgres - - redis - requires: - - name: postgres - - name: redis - helm: - type: install - action: install - repository: org/api-service - branchName: feature-branch - envLens: true - grpc: false - disableIngressHost: false - overrideDefaultIpWhitelist: false - chart: - name: org-chart - repoUrl: https://charts.example.com - version: "2.1.0" - valueFiles: - - helm/api-values.yaml - - helm/api-overrides.yaml - - helm/api-secrets.yaml - values: - - "replicaCount=2" - - "resources.requests.cpu=500m" - - "resources.requests.memory=512Mi" - - "resources.limits.cpu=1000m" - - "resources.limits.memory=1024Mi" - - "livenessProbe.httpGet.path=/health" - - "livenessProbe.httpGet.port=8080" - - "readinessProbe.httpGet.path=/ready" - - "readinessProbe.httpGet.port=8080" - docker: - defaultTag: latest - ecr: "123456789012.dkr.ecr.us-west-2.amazonaws.com/org/api" - pipelineId: pipeline-api-build - builder: - engine: buildkit - buildArgs: - NODE_ENV: production - BUILD_DATE: "2024-01-15" - app: - dockerfilePath: sysops/dockerfiles/api.dockerfile - command: "node dist/server.js" - arguments: "--max-old-space-size=4096" - ports: - - 8080 - - 9090 - env: - NODE_ENV: production - LOG_LEVEL: info - DATABASE_URL: postgres://user:pass@db:5432/app - REDIS_URL: redis://redis:6379 - API_SECRET_KEY: vault:secret/api-key - CORS_ORIGIN: https://app.example.com - RATE_LIMIT_MAX: "100" - RATE_LIMIT_WINDOW: "60000" - afterBuildPipelineConfig: - afterBuildPipelineId: pipeline-api-post - detatchAfterBuildPipeline: false - description: "Run migrations after build" - init: - dockerfilePath: sysops/dockerfiles/init.dockerfile - command: "node scripts/migrate.js" - arguments: "--run-seeds" - env: - DATABASE_URL: postgres://user:pass@db:5432/app - MIGRATION_DIR: ./migrations - - name: web-app - appShort: web - deploymentDependsOn: - - api-service - helm: - type: install - action: install - repository: org/web-app - branchName: feature-branch - envLens: true - chart: - name: org-chart - repoUrl: https://charts.example.com - version: "2.1.0" - valueFiles: - - helm/web-values.yaml - - helm/web-ingress.yaml - values: - - "replicaCount=3" - - "resources.requests.cpu=250m" - - "resources.requests.memory=256Mi" - - "resources.limits.cpu=500m" - - "resources.limits.memory=512Mi" - - "ingress.enabled=true" - - "ingress.hosts[0].host=app.example.com" - docker: - defaultTag: latest - ecr: "123456789012.dkr.ecr.us-west-2.amazonaws.com/org/web" - builder: - engine: buildkit - app: - dockerfilePath: Dockerfile - ports: - - 3000 - env: - NEXT_PUBLIC_API_URL: /api - NEXT_PUBLIC_WS_URL: wss://ws.example.com - NODE_ENV: production - - name: worker - appShort: wrk - deploymentDependsOn: - - postgres - - redis - requires: - - name: postgres - helm: - type: install - repository: org/api-service - branchName: feature-branch - chart: - name: org-chart - repoUrl: https://charts.example.com - version: "2.1.0" - valueFiles: - - helm/worker-values.yaml - values: - - "replicaCount=1" - - "resources.requests.cpu=1000m" - - "resources.requests.memory=2048Mi" - docker: - defaultTag: latest - ecr: "123456789012.dkr.ecr.us-west-2.amazonaws.com/org/worker" - builder: - engine: buildkit - app: - dockerfilePath: sysops/dockerfiles/worker.dockerfile - command: "node dist/worker.js" - env: - WORKER_CONCURRENCY: "5" - QUEUE_PREFIX: prod - DATABASE_URL: postgres://user:pass@db:5432/app - REDIS_URL: redis://redis:6379 - - name: postgres - docker: - dockerImage: postgres - defaultTag: "14" - command: "postgres" - arguments: "-c shared_buffers=256MB -c max_connections=200" - ports: - - 5432 - env: - POSTGRES_DB: myapp - POSTGRES_USER: admin - POSTGRES_PASSWORD: secret - PGDATA: /var/lib/postgresql/data/pgdata - - name: external-api - externalHttp: - defaultInternalHostname: api.partner.internal - defaultPublicUrl: https://api.partner.example.com -`; - - it('achieves 2.5-8x compression on a realistic multi-service YAML', () => { - const result = summarizeLifecycleYaml(realisticYaml); - const rawTokens = countTokens(realisticYaml); - const summaryTokens = countTokens(result.text); - const ratio = rawTokens / summaryTokens; - - expect(result.parsed).toBe(true); - expect(result.serviceCount).toBe(5); - expect(ratio).toBeGreaterThanOrEqual(2.5); - expect(ratio).toBeLessThanOrEqual(8); - }); - - it('summary is always smaller than raw YAML for multi-service files', () => { - const result = summarizeLifecycleYaml(realisticYaml); - expect(countTokens(result.text)).toBeLessThan(countTokens(realisticYaml)); - }); - }); -}); diff --git a/src/server/services/ai/context/contextSummarizer.ts b/src/server/services/ai/context/contextSummarizer.ts deleted file mode 100644 index 6929fe39..00000000 --- a/src/server/services/ai/context/contextSummarizer.ts +++ /dev/null @@ -1,265 +0,0 @@ -/** - * Copyright 2025 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 yaml from 'js-yaml'; - -export interface LifecycleYamlSummary { - text: string; - parsed: boolean; - serviceCount: number; -} - -const SERVICE_TYPE_KEYS = [ - 'helm', - 'codefresh', - 'github', - 'docker', - 'externalHttp', - 'auroraRestore', - 'configuration', -] as const; - -function extractServiceType(svc: Record): string { - for (const key of SERVICE_TYPE_KEYS) { - if (svc[key] != null) return key; - } - return 'unknown'; -} - -function formatRepo(block: Record): string | null { - if (!block.repository) return null; - return block.branchName ? `${block.repository} @ ${block.branchName}` : block.repository; -} - -function extractDependencies(svc: Record): string[] { - const deps: string[] = []; - if (Array.isArray(svc.deploymentDependsOn)) { - deps.push(...svc.deploymentDependsOn); - } - if (Array.isArray(svc.requires)) { - for (const req of svc.requires) { - if (req?.name) deps.push(req.name); - } - } - return deps; -} - -function extractHelmLines(block: Record, lines: string[], referencedFiles: string[]): void { - const repo = formatRepo(block); - if (repo) lines.push(` Repo: ${repo}`); - - if (block.chart) { - let chartLine = ` Chart: ${block.chart.name}`; - if (block.chart.repoUrl) chartLine += ` (repoUrl: ${block.chart.repoUrl})`; - if (block.chart.version) chartLine += ` (version: ${block.chart.version})`; - lines.push(chartLine); - - if (Array.isArray(block.chart.valueFiles) && block.chart.valueFiles.length) { - lines.push(` ValueFiles: ${block.chart.valueFiles.join(', ')}`); - referencedFiles.push(...block.chart.valueFiles); - } - } - - const dockerfilePath = block.docker?.app?.dockerfilePath; - const engine = block.docker?.builder?.engine; - if (dockerfilePath) { - const dockerLine = engine - ? ` Docker: ${engine} | dockerfilePath: ${dockerfilePath}` - : ` Docker: dockerfilePath: ${dockerfilePath}`; - lines.push(dockerLine); - referencedFiles.push(dockerfilePath); - } - - const ports = block.docker?.app?.ports; - if (Array.isArray(ports) && ports.length) { - lines.push(` Ports: ${ports.join(', ')}`); - } -} - -function extractCodefreshLines(block: Record, lines: string[]): void { - const repo = formatRepo(block); - if (repo) lines.push(` Repo: ${repo}`); - - if (block.deploy?.pipelineId) { - lines.push(` Pipeline: ${block.deploy.pipelineId}`); - } -} - -function extractGithubLines(block: Record, lines: string[], referencedFiles: string[]): void { - const repo = formatRepo(block); - if (repo) lines.push(` Repo: ${repo}`); - - const dockerfilePath = block.docker?.app?.dockerfilePath; - const engine = block.docker?.builder?.engine; - if (dockerfilePath) { - const dockerLine = engine - ? ` Docker: ${engine} | dockerfilePath: ${dockerfilePath}` - : ` Docker: dockerfilePath: ${dockerfilePath}`; - lines.push(dockerLine); - referencedFiles.push(dockerfilePath); - } -} - -function extractDockerLines(block: Record, lines: string[]): void { - if (block.dockerImage) { - const image = block.defaultTag ? `${block.dockerImage}:${block.defaultTag}` : block.dockerImage; - lines.push(` Image: ${image}`); - } - - if (Array.isArray(block.ports) && block.ports.length) { - lines.push(` Ports: ${block.ports.join(', ')}`); - } -} - -function extractExternalHttpLines(block: Record, lines: string[]): void { - if (block.defaultInternalHostname) { - lines.push(` InternalHost: ${block.defaultInternalHostname}`); - } - if (block.defaultPublicUrl) { - lines.push(` PublicUrl: ${block.defaultPublicUrl}`); - } -} - -function extractAuroraRestoreLines(block: Record, lines: string[]): void { - if (block.command) lines.push(` Command: ${block.command}`); - if (block.arguments) lines.push(` Arguments: ${block.arguments}`); -} - -function extractConfigurationLines(block: Record, lines: string[]): void { - if (block.defaultTag) lines.push(` DefaultTag: ${block.defaultTag}`); - if (block.branchName) lines.push(` Branch: ${block.branchName}`); -} - -export function summarizeLifecycleYaml(rawYaml: string): LifecycleYamlSummary { - try { - const doc = yaml.load(rawYaml) as Record; - if (!doc || typeof doc !== 'object' || Array.isArray(doc)) { - return { text: rawYaml, parsed: false, serviceCount: 0 }; - } - - const lines: string[] = []; - - if (doc.environment || doc.version) { - lines.push('ENVIRONMENT:'); - if (doc.version) lines.push(` Version: ${doc.version}`); - if (doc.environment) { - if (doc.environment.autoDeploy !== undefined) { - lines.push(` AutoDeploy: ${doc.environment.autoDeploy}`); - } - if (Array.isArray(doc.environment.enabledFeatures) && doc.environment.enabledFeatures.length) { - lines.push(` EnabledFeatures: ${doc.environment.enabledFeatures.join(', ')}`); - } - if (Array.isArray(doc.environment.defaultServices) && doc.environment.defaultServices.length) { - lines.push(` Default Services: ${doc.environment.defaultServices.map((s: any) => s.name).join(', ')}`); - } - if (Array.isArray(doc.environment.optionalServices) && doc.environment.optionalServices.length) { - lines.push(` Optional Services: ${doc.environment.optionalServices.map((s: any) => s.name).join(', ')}`); - } - } - lines.push(''); - } - - const services = Array.isArray(doc.services) ? doc.services : []; - const referencedFiles: string[] = []; - - lines.push(`SERVICES (${services.length}):`); - lines.push(''); - - if (services.length > 30) { - const typeCounts = new Map(); - const servicesWithDeps: Array<{ name: string; type: string; deps: string[] }> = []; - - for (const svc of services) { - const type = extractServiceType(svc); - typeCounts.set(type, (typeCounts.get(type) || 0) + 1); - const deps = extractDependencies(svc); - if (deps.length > 0) { - servicesWithDeps.push({ name: svc.name, type, deps }); - } - } - - const typeBreakdown = [...typeCounts.entries()].map(([type, count]) => `${count} ${type}`).join(', '); - lines.push(`Types: ${typeBreakdown}`); - lines.push(''); - - if (servicesWithDeps.length > 0) { - lines.push(`Services with dependencies (${servicesWithDeps.length}):`); - for (const svc of servicesWithDeps) { - lines.push(` ${svc.name} (${svc.type}) → DependsOn: ${svc.deps.join(', ')}`); - } - lines.push(''); - } - - lines.push('[Use get_file("lifecycle.yaml") for full service details]'); - lines.push(''); - } else { - services.forEach((svc: any, i: number) => { - const type = extractServiceType(svc); - const serviceLines: string[] = []; - serviceLines.push(`[${i + 1}] ${svc.name} (${type})`); - - const block = svc[type]; - if (block && typeof block === 'object') { - switch (type) { - case 'helm': - extractHelmLines(block, serviceLines, referencedFiles); - break; - case 'codefresh': - extractCodefreshLines(block, serviceLines); - break; - case 'github': - extractGithubLines(block, serviceLines, referencedFiles); - break; - case 'docker': - extractDockerLines(block, serviceLines); - break; - case 'externalHttp': - extractExternalHttpLines(block, serviceLines); - break; - case 'auroraRestore': - extractAuroraRestoreLines(block, serviceLines); - break; - case 'configuration': - extractConfigurationLines(block, serviceLines); - break; - } - } - - const deps = extractDependencies(svc); - if (deps.length) { - serviceLines.push(` DependsOn: ${deps.join(', ')}`); - } - - lines.push(...serviceLines); - lines.push(''); - }); - } - - const uniqueFiles = [...new Set(referencedFiles)]; - if (uniqueFiles.length) { - lines.push('CONFIG POINTERS:'); - lines.push(` Referenced Files: ${uniqueFiles.join(', ')}`); - } - - return { - text: lines.join('\n'), - parsed: true, - serviceCount: services.length, - }; - } catch { - return { text: rawYaml, parsed: false, serviceCount: 0 }; - } -} diff --git a/src/server/services/ai/context/gatherer.ts b/src/server/services/ai/context/gatherer.ts deleted file mode 100644 index afbc5be2..00000000 --- a/src/server/services/ai/context/gatherer.ts +++ /dev/null @@ -1,628 +0,0 @@ -/** - * Copyright 2025 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 BaseService from '../../_service'; -import * as k8s from '@kubernetes/client-node'; -import { - DebugContext, - LifecycleContext, - ServiceDebugInfo, - PodDebugInfo, - DiagnosedIssue, - ContextWarning, - ContextError, - K8sEvent, -} from '../../types/aiAgent'; -import { GitHubClient } from '../tools/shared/githubClient'; -import { getLogger } from 'server/lib/logger'; - -export default class AIAgentContextService extends BaseService { - private static defaultBranchCache: Map = new Map(); - private static DEFAULT_BRANCH_CACHE_TTL_MS = 86400000; - - async gatherFullContext(buildUuid: string): Promise { - const warnings: ContextWarning[] = []; - const errors: ContextError[] = []; - - const build = await this.db.models.Build.query() - .findOne({ uuid: buildUuid }) - .withGraphFetched( - '[pullRequest.repository, environment, deploys.[service, repository, deployable], deployables]' - ); - - if (!build) { - throw new Error(`Build not found: ${buildUuid}`); - } - - const githubClient = new GitHubClient(); - const octokit = await githubClient.getOctokit('ai-agent-context-gatherer'); - - let lifecycleContext; - let kubernetesServices; - let lifecycleYaml; - - try { - lifecycleContext = await this.gatherLifecycleContext(build, octokit); - } catch (error) { - errors.push({ - source: 'lifecycle', - message: 'Failed to gather Lifecycle database context', - error: error.message, - recoverable: false, - }); - throw error; - } - - try { - kubernetesServices = await this.gatherKubernetesInfo(build, warnings, errors); - } catch (error) { - errors.push({ - source: 'kubernetes', - message: 'Failed to gather Kubernetes context', - error: error.message, - recoverable: true, - }); - kubernetesServices = []; - } - - try { - const fullName = build.pullRequest.fullName; - const branch = build.pullRequest.branchName; - - if (fullName && branch) { - const [owner, repo] = fullName.split('/'); - - const possiblePaths = ['lifecycle.yaml', 'lifecycle.yml']; - let response; - let foundPath: string | null = null; - - for (const path of possiblePaths) { - try { - response = await octokit.request(`GET /repos/${owner}/${repo}/contents/${path}`, { - ref: branch, - }); - foundPath = path; - break; - } catch (error) { - continue; - } - } - - if (foundPath && response?.data && 'content' in response.data && response.data.type === 'file') { - const content = Buffer.from(response.data.content, 'base64').toString('utf-8'); - lifecycleYaml = { - path: foundPath, - content, - }; - } else { - lifecycleYaml = { - path: 'lifecycle.yaml', - content: '', - error: 'lifecycle.yaml or lifecycle.yml not found in repository', - }; - warnings.push({ - source: 'lifecycle', - message: 'Could not fetch lifecycle configuration from repository', - details: 'Neither lifecycle.yaml nor lifecycle.yml found', - }); - } - } - } catch (error) { - warnings.push({ - source: 'lifecycle', - message: 'Failed to fetch lifecycle configuration', - details: error.message, - }); - } - - return { - buildUuid, - namespace: build.namespace, - lifecycleContext, - lifecycleYaml, - services: kubernetesServices, - gatheredAt: new Date(), - warnings: warnings.length > 0 ? warnings : undefined, - errors: errors.length > 0 ? errors : undefined, - }; - } - - private async gatherLifecycleContext(build: any, octokit: any): Promise { - const fullName = build.pullRequest.fullName; - const [owner, repo] = fullName ? fullName.split('/') : []; - const defaultBranch = owner && repo ? await this.getDefaultBranch(owner, repo, octokit) : 'main'; - - return { - build: { - uuid: build.uuid, - status: build.status, - statusMessage: build.statusMessage, - namespace: build.namespace, - sha: build.sha, - trackDefaultBranches: build.trackDefaultBranches, - capacityType: build.capacityType, - enabledFeatures: build.enabledFeatures || [], - dependencyGraph: build.dependencyGraph || {}, - dashboardLinks: build.dashboardLinks || {}, - createdAt: build.createdAt, - updatedAt: build.updatedAt, - }, - pullRequest: { - number: build.pullRequest.number, - title: build.pullRequest.title, - username: build.pullRequest.githubLogin, - branch: build.pullRequest.branchName, - baseBranch: defaultBranch, - status: build.pullRequest.status, - url: `https://github.com/${build.pullRequest.fullName}/pull/${build.pullRequest.pullRequestNumber}`, - latestCommit: build.pullRequest.latestCommit, - fullName: build.pullRequest.fullName, - commentId: build.pullRequest.commentId, - labels: build.pullRequest.labels || [], - }, - environment: { - id: build.environment.id, - name: build.environment.name, - config: build.environment.config || {}, - }, - deploys: build.deploys.map((deploy: any) => ({ - uuid: deploy.uuid, - serviceName: deploy.deployable?.name || deploy.service?.name || deploy.uuid, - status: deploy.status, - statusMessage: deploy.statusMessage, - type: deploy.deployable?.type || deploy.type, - dockerImage: deploy.dockerImage, - branch: deploy.branch, - repoName: deploy.repoName, - buildNumber: deploy.buildNumber, - buildPipelineId: deploy.buildPipelineId, - deployPipelineId: deploy.deployPipelineId, - builderEngine: deploy.deployable?.builder?.engine, - helmChart: deploy.deployable?.helm?.chart, - repositoryId: deploy.deployable?.repositoryId, - })), - repository: { - name: build.pullRequest.repository.name, - githubRepositoryId: build.pullRequest.repository.githubRepositoryId, - url: build.pullRequest.repository.url, - }, - }; - } - - private async getDefaultBranch(owner: string, repo: string, octokit: any): Promise { - const fullName = `${owner}/${repo}`.toLowerCase(); - - try { - const now = Date.now(); - const memoryCached = AIAgentContextService.defaultBranchCache.get(fullName); - if (memoryCached && now < memoryCached.expiry) { - return memoryCached.data; - } - - const redisKey = `default_branch:${fullName}`; - const redisValue = await this.redis.get(redisKey); - if (redisValue) { - AIAgentContextService.defaultBranchCache.set(fullName, { - data: redisValue, - expiry: now + AIAgentContextService.DEFAULT_BRANCH_CACHE_TTL_MS, - }); - return redisValue; - } - - const response = await octokit.request('GET /repos/{owner}/{repo}', { owner, repo }); - const defaultBranch = response.data.default_branch; - - await this.redis.set(redisKey, defaultBranch, 'EX', 86400); - AIAgentContextService.defaultBranchCache.set(fullName, { - data: defaultBranch, - expiry: now + AIAgentContextService.DEFAULT_BRANCH_CACHE_TTL_MS, - }); - - return defaultBranch; - } catch (error) { - getLogger().warn(`AIAgentContext: failed to fetch default branch repo=${fullName} error=${error}`); - return 'main'; - } - } - - private static FAILED_DEPLOY_STATUSES = new Set(['BUILD_FAILED', 'DEPLOY_FAILED', 'ERROR']); - - private async gatherKubernetesInfo( - build: any, - warnings: ContextWarning[], - errors: ContextError[] - ): Promise { - const namespace = build.namespace; - const services: ServiceDebugInfo[] = []; - - for (const deploy of build.deploys || []) { - const serviceName = deploy.deployable?.name || deploy.service?.name || deploy.uuid; - const isFailed = AIAgentContextService.FAILED_DEPLOY_STATUSES.has(deploy.status); - - if (isFailed) { - try { - const serviceInfo = await this.gatherServiceDebugInfo(deploy, namespace, warnings); - services.push(serviceInfo); - } catch (error) { - errors.push({ - source: 'kubernetes', - message: `Failed to gather info for service ${serviceName}`, - error: error.message, - recoverable: true, - }); - } - } else { - services.push({ - name: serviceName, - type: deploy.deployable?.type || deploy.type, - status: this.determineServiceStatus(deploy, []), - deployInfo: { - uuid: deploy.uuid, - serviceName, - status: deploy.status, - statusMessage: deploy.statusMessage, - type: deploy.deployable?.type || deploy.type, - dockerImage: deploy.dockerImage, - branch: deploy.branch, - repoName: deploy.repoName, - buildNumber: deploy.buildNumber, - env: deploy.env, - initEnv: deploy.initEnv, - createdAt: deploy.createdAt, - updatedAt: deploy.updatedAt, - }, - pods: [], - events: [], - issues: [], - }); - } - } - - return services; - } - - private async gatherServiceDebugInfo( - deploy: any, - namespace: string, - warnings: ContextWarning[] - ): Promise { - const kc = new k8s.KubeConfig(); - kc.loadFromDefault(); - const coreApi = kc.makeApiClient(k8s.CoreV1Api); - const appsApi = kc.makeApiClient(k8s.AppsV1Api); - - const serviceName = deploy.deployable?.name || deploy.service?.name || deploy.uuid; - - let pods: PodDebugInfo[] = []; - let events: K8sEvent[] = []; - let deployment: any = undefined; - - try { - pods = await this.gatherPodsInfo(deploy.uuid, serviceName, namespace, coreApi, warnings); - } catch (error) { - warnings.push({ - source: 'kubernetes', - message: `Could not fetch pods for ${serviceName} (${deploy.uuid})`, - details: error.message, - }); - } - - try { - events = await this.getServiceEvents(serviceName, namespace, coreApi); - } catch (error) { - warnings.push({ - source: 'kubernetes', - message: `Could not fetch events for ${serviceName}`, - details: error.message, - }); - } - - try { - deployment = await this.gatherDeploymentInfo(deploy.uuid, namespace, appsApi); - } catch (error) { - warnings.push({ - source: 'kubernetes', - message: `Could not fetch deployment for ${serviceName} (${deploy.uuid})`, - details: error.message, - }); - } - - const issues = this.diagnoseIssues(pods, events); - - return { - name: serviceName, - type: deploy.deployable?.type || deploy.type, - status: this.determineServiceStatus(deploy, pods), - deployInfo: { - uuid: deploy.uuid, - serviceName: serviceName, - status: deploy.status, - statusMessage: deploy.statusMessage, - type: deploy.deployable?.type || deploy.type, - dockerImage: deploy.dockerImage, - branch: deploy.branch, - repoName: deploy.repoName, - buildNumber: deploy.buildNumber, - env: deploy.env, - initEnv: deploy.initEnv, - createdAt: deploy.createdAt, - updatedAt: deploy.updatedAt, - }, - deployment, - pods, - events, - issues, - }; - } - - private async gatherPodsInfo( - deploymentUuid: string, - serviceName: string, - namespace: string, - coreApi: k8s.CoreV1Api, - warnings: ContextWarning[] - ): Promise { - // Try to find pods by deployment uuid first (using label selector) - let podsResponse = await coreApi.listNamespacedPod( - namespace, - undefined, - undefined, - undefined, - undefined, - `deployment=${deploymentUuid}` - ); - - // If no pods found, try by service name as fallback - if (podsResponse.body.items.length === 0) { - podsResponse = await coreApi.listNamespacedPod( - namespace, - undefined, - undefined, - undefined, - undefined, - `app=${serviceName}` - ); - } - - return Promise.all( - podsResponse.body.items.map(async (pod) => { - let recentLogs = ''; - try { - recentLogs = await this.getRecentLogs(pod.metadata.name, namespace, coreApi); - } catch (error) { - warnings.push({ - source: 'logs', - message: `Could not fetch logs for pod ${pod.metadata.name}`, - details: error.message, - }); - } - - return { - name: pod.metadata.name, - phase: pod.status.phase, - conditions: pod.status.conditions || [], - containerStatuses: pod.status.containerStatuses || [], - recentLogs, - events: await this.getPodEvents(pod.metadata.name, namespace, coreApi).catch(() => []), - }; - }) - ); - } - - private async gatherDeploymentInfo(serviceName: string, namespace: string, appsApi: k8s.AppsV1Api): Promise { - try { - const deploymentResponse = await appsApi.readNamespacedDeployment(serviceName, namespace); - const deployment = deploymentResponse.body; - - return { - name: deployment.metadata.name, - replicas: { - desired: deployment.spec.replicas || 0, - current: deployment.status.replicas || 0, - ready: deployment.status.readyReplicas || 0, - available: deployment.status.availableReplicas || 0, - }, - conditions: deployment.status.conditions || [], - strategy: deployment.spec.strategy?.type || 'Unknown', - containers: deployment.spec.template.spec.containers.map((c) => ({ - name: c.name, - image: c.image, - })), - }; - } catch (error) { - // Deployment might not exist yet or service doesn't use deployments - return undefined; - } - } - - private async getRecentLogs( - podName: string, - namespace: string, - coreApi: k8s.CoreV1Api, - tailLines: number = 50 - ): Promise { - const logsResponse = await coreApi.readNamespacedPodLog( - podName, - namespace, - undefined, - undefined, - undefined, - undefined, - undefined, - undefined, - undefined, - tailLines - ); - return logsResponse.body; - } - - private async getPodEvents(podName: string, namespace: string, coreApi: k8s.CoreV1Api): Promise { - try { - const eventsResponse = await coreApi.listNamespacedEvent( - namespace, - undefined, - undefined, - undefined, - `involvedObject.name=${podName}` - ); - - return eventsResponse.body.items.slice(-5).map((event) => ({ - type: event.type, - reason: event.reason, - message: event.message, - count: event.count, - firstTimestamp: event.firstTimestamp, - lastTimestamp: event.lastTimestamp, - })); - } catch (error) { - return []; - } - } - - private async getServiceEvents(serviceName: string, namespace: string, coreApi: k8s.CoreV1Api): Promise { - try { - const eventsResponse = await coreApi.listNamespacedEvent(namespace); - - return eventsResponse.body.items - .filter((event) => { - const name = event.involvedObject?.name || ''; - return name.includes(serviceName); - }) - .slice(-5) - .map((event) => ({ - type: event.type, - reason: event.reason, - message: event.message, - count: event.count, - firstTimestamp: event.firstTimestamp, - lastTimestamp: event.lastTimestamp, - })); - } catch (error) { - return []; - } - } - - private diagnoseIssues(pods: PodDebugInfo[], events: K8sEvent[]): DiagnosedIssue[] { - const issues: DiagnosedIssue[] = []; - - for (const pod of pods) { - if (pod.phase === 'Pending') { - const unschedulable = pod.conditions.find((c: any) => c.type === 'PodScheduled' && c.status === 'False'); - if (unschedulable) { - issues.push({ - severity: 'critical', - category: 'resources', - title: 'Pod cannot be scheduled', - description: unschedulable.message || 'Pod is pending scheduling', - suggestedFix: 'Check resource requests, node capacity, and node selectors', - detectedBy: 'rules', - }); - } - } - - for (const container of pod.containerStatuses) { - if (container.state?.waiting?.reason === 'ImagePullBackOff') { - issues.push({ - severity: 'critical', - category: 'image', - title: 'Cannot pull container image', - description: `Image ${container.image} cannot be pulled from registry`, - suggestedFix: 'Verify image name, registry access, and image pull secrets', - detectedBy: 'rules', - }); - } - - if (container.state?.waiting?.reason === 'CrashLoopBackOff') { - issues.push({ - severity: 'critical', - category: 'configuration', - title: 'Container is crash looping', - description: `Container ${container.name} repeatedly crashes on startup`, - suggestedFix: 'Check application logs for startup errors and verify configuration', - detectedBy: 'rules', - }); - } - - if (container.lastState?.terminated?.reason === 'OOMKilled') { - issues.push({ - severity: 'critical', - category: 'resources', - title: 'Container killed due to out of memory', - description: `Container ${container.name} exceeded memory limit and was killed`, - suggestedFix: 'Increase memory limits or optimize application memory usage', - detectedBy: 'rules', - }); - } - - if (container.restartCount > 5) { - issues.push({ - severity: 'warning', - category: 'configuration', - title: 'High restart count', - description: `Container has restarted ${container.restartCount} times`, - suggestedFix: 'Investigate logs for recurring errors or configuration issues', - detectedBy: 'rules', - }); - } - } - } - - for (const event of events) { - if (event.type === 'Warning' && event.reason === 'FailedScheduling') { - issues.push({ - severity: 'critical', - category: 'resources', - title: 'Failed to schedule pod', - description: event.message, - suggestedFix: 'Check cluster capacity and resource requests', - detectedBy: 'rules', - }); - } - } - - return issues; - } - - private determineServiceStatus( - deploy: any, - pods: PodDebugInfo[] - ): 'pending' | 'building' | 'deploying' | 'running' | 'failed' { - if (deploy.status === 'ERROR' || deploy.status === 'BUILD_FAILED' || deploy.status === 'DEPLOY_FAILED') { - return 'failed'; - } - if (deploy.status === 'BUILDING' || deploy.status === 'CLONING') { - return 'building'; - } - if (deploy.status === 'DEPLOYING' || deploy.status === 'WAITING') { - return 'deploying'; - } - if (deploy.status === 'PENDING' || deploy.status === 'QUEUED') { - return 'pending'; - } - - const runningPods = pods.filter((p) => p.phase === 'Running'); - if (runningPods.length > 0 && runningPods.length === pods.length) { - return 'running'; - } - - const failedPods = pods.filter((p) => p.phase === 'Failed' || p.phase === 'CrashLoopBackOff'); - if (failedPods.length > 0) { - return 'failed'; - } - - return 'deploying'; - } -} diff --git a/src/server/services/ai/conversation/__tests__/manager.test.ts b/src/server/services/ai/conversation/__tests__/manager.test.ts deleted file mode 100644 index b587eac5..00000000 --- a/src/server/services/ai/conversation/__tests__/manager.test.ts +++ /dev/null @@ -1,112 +0,0 @@ -/** - * Copyright 2025 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/logger', () => ({ - getLogger: () => ({ - info: jest.fn(), - warn: jest.fn(), - error: jest.fn(), - }), -})); - -const mockCountTokens = jest.fn(); -jest.mock('../../prompts/tokenCounter', () => ({ - countTokens: (...args: any[]) => mockCountTokens(...args), -})); - -import { ConversationManager, ConversationState } from '../manager'; -import { textMessage } from '../../types/message'; - -describe('ConversationManager', () => { - let manager: ConversationManager; - - beforeEach(() => { - manager = new ConversationManager(); - mockCountTokens.mockReset(); - }); - - describe('shouldCompress', () => { - it('returns false when tokens are below threshold', async () => { - mockCountTokens.mockResolvedValue(50000); - const result = await manager.shouldCompress([textMessage('user', 'hi')]); - expect(result).toBe(false); - }); - - it('returns true when tokens are above threshold', async () => { - mockCountTokens.mockResolvedValue(90000); - const result = await manager.shouldCompress([textMessage('user', 'hi')]); - expect(result).toBe(true); - }); - - it('returns false at exactly the threshold (> not >=)', async () => { - mockCountTokens.mockResolvedValue(80000); - const result = await manager.shouldCompress([textMessage('user', 'hi')]); - expect(result).toBe(false); - }); - }); - - describe('compress', () => { - it('calls LLM provider and returns parsed ConversationState', async () => { - const mockState = { - summary: 'Debugging pod crash', - identifiedIssues: [{ service: 'web', issue: 'OOM', confidence: 'high' }], - investigatedServices: ['web'], - toolsUsed: ['getK8sResources'], - currentTask: 'Check memory limits', - tokenCount: 0, - messageCount: 0, - compressionLevel: 0, - }; - const mockProvider = { - streamCompletion: jest.fn().mockImplementation(async function* () { - yield { type: 'text', content: JSON.stringify(mockState) }; - }), - }; - mockCountTokens.mockResolvedValue(500); - - const result = await manager.compress([textMessage('user', 'check pods')], mockProvider as any); - - expect(result.summary).toBe('Debugging pod crash'); - expect(result.tokenCount).toBe(500); - expect(result.messageCount).toBe(1); - expect(result.compressionLevel).toBe(1); - }); - }); - - describe('buildPromptFromState', () => { - it('produces markdown with issues, services, and current task', () => { - const state: ConversationState = { - summary: 'Investigating crash loop', - identifiedIssues: [{ service: 'api', issue: 'OOMKilled', confidence: 'high' }], - investigatedServices: ['api', 'redis'], - toolsUsed: ['getK8sResources'], - currentTask: 'Check resource limits', - tokenCount: 1000, - messageCount: 5, - compressionLevel: 1, - }; - - const prompt = manager.buildPromptFromState(state); - - expect(prompt).toContain('Summary'); - expect(prompt).toContain('Identified Issues'); - expect(prompt).toContain('Already Investigated'); - expect(prompt).toContain('Current Task'); - expect(prompt).toContain('api'); - expect(prompt).toContain('OOMKilled'); - }); - }); -}); diff --git a/src/server/services/ai/conversation/__tests__/persistence.test.ts b/src/server/services/ai/conversation/__tests__/persistence.test.ts deleted file mode 100644 index 1f5d8118..00000000 --- a/src/server/services/ai/conversation/__tests__/persistence.test.ts +++ /dev/null @@ -1,303 +0,0 @@ -/** - * Copyright 2025 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. - */ - -const mockGetConversation = jest.fn(); - -jest.mock('../storage', () => ({ - __esModule: true, - default: jest.fn().mockImplementation(() => ({ - getConversation: mockGetConversation, - })), -})); - -const mockLogger = { info: jest.fn(), warn: jest.fn(), error: jest.fn() }; -jest.mock('server/lib/logger', () => ({ - getLogger: () => mockLogger, -})); - -jest.mock('server/models/Conversation', () => { - const model: any = {}; - model.query = jest.fn(); - model.transact = jest.fn(); - return { __esModule: true, default: model }; -}); - -jest.mock('server/models/ConversationMessage', () => { - const model: any = {}; - model.query = jest.fn(); - return { __esModule: true, default: model }; -}); - -import ConversationPersistenceService from '../persistence'; -import AIAgentConversationService from '../storage'; -import Conversation from 'server/models/Conversation'; -import ConversationMessage from 'server/models/ConversationMessage'; - -const MockConversation = Conversation as any; -const MockConversationMessage = ConversationMessage as any; - -interface TransactionSetup { - existingConversation?: any; - existingMessages?: Array<{ role: string; timestamp: number }>; - onConversationInsert?: (row: any) => void; - onConversationPatch?: (row: any) => void; - onMessageInsert?: (rows: any[]) => void; -} - -function setupTransaction({ - existingConversation, - existingMessages = [], - onConversationInsert, - onConversationPatch, - onMessageInsert, -}: TransactionSetup) { - const findById = jest.fn().mockResolvedValue(existingConversation); - const insertConversation = jest.fn().mockImplementation(async (row: any) => { - onConversationInsert?.(row); - return row; - }); - const patchWhere = jest.fn().mockResolvedValue(1); - const patchConversation = jest.fn().mockImplementation((row: any) => { - onConversationPatch?.(row); - return { where: patchWhere }; - }); - - const selectMessages = jest.fn().mockResolvedValue(existingMessages); - const whereMessages = jest.fn().mockReturnValue({ select: selectMessages }); - const insertMessages = jest.fn().mockImplementation(async (rows: any[]) => { - onMessageInsert?.(rows); - return rows; - }); - - MockConversation.transact.mockImplementation(async (cb: any) => { - const trx = { id: 'mock-trx' }; - - MockConversation.query.mockImplementation((arg: any) => { - if (arg === trx) { - return { - findById, - insert: insertConversation, - patch: patchConversation, - }; - } - return {}; - }); - - MockConversationMessage.query.mockImplementation((arg: any) => { - if (arg === trx) { - return { - where: whereMessages, - insert: insertMessages, - }; - } - return {}; - }); - - return cb(trx); - }); - - return { - findById, - insertConversation, - patchConversation, - patchWhere, - whereMessages, - selectMessages, - insertMessages, - }; -} - -describe('ConversationPersistenceService', () => { - let service: ConversationPersistenceService; - - beforeEach(() => { - jest.clearAllMocks(); - const conversationService = new AIAgentConversationService({} as any, {} as any, {} as any, {} as any); - service = new ConversationPersistenceService(conversationService); - }); - - describe('persistConversation', () => { - it('returns false when Redis has no conversation', async () => { - mockGetConversation.mockResolvedValue(null); - - const result = await service.persistConversation('uuid-1', 'example-org/example-repo'); - - expect(result).toBe(false); - expect(mockGetConversation).toHaveBeenCalledWith('uuid-1'); - expect(MockConversation.transact).not.toHaveBeenCalled(); - }); - - it('returns false when Redis conversation has empty messages', async () => { - mockGetConversation.mockResolvedValue({ - buildUuid: 'uuid-1', - messages: [], - lastActivity: 1000, - }); - - const result = await service.persistConversation('uuid-1', 'example-org/example-repo'); - - expect(result).toBe(false); - expect(MockConversation.transact).not.toHaveBeenCalled(); - }); - - it('persists a new conversation and all messages', async () => { - const conversation = { - buildUuid: 'uuid-1', - messages: [ - { role: 'user', content: 'hello', timestamp: 1000 }, - { role: 'assistant', content: 'hi there', timestamp: 1001 }, - { role: 'user', content: 'thanks', timestamp: 1002 }, - ], - lastActivity: 1002, - contextSnapshot: { buildUuid: 'uuid-1' }, - }; - mockGetConversation.mockResolvedValue(conversation); - - let insertedConversation: any; - let insertedMessages: any[] = []; - setupTransaction({ - existingConversation: undefined, - existingMessages: [], - onConversationInsert: (row) => { - insertedConversation = row; - }, - onMessageInsert: (rows) => { - insertedMessages = rows; - }, - }); - - const result = await service.persistConversation('uuid-1', 'example-org/example-repo', 'gpt-4'); - - expect(result).toBe(true); - expect(insertedConversation).toMatchObject({ - buildUuid: 'uuid-1', - repo: 'example-org/example-repo', - model: 'gpt-4', - messageCount: 3, - metadata: { - contextSnapshot: { buildUuid: 'uuid-1' }, - lastActivity: 1002, - }, - }); - expect(insertedMessages).toHaveLength(3); - expect(mockLogger.info).toHaveBeenCalledWith('AI: conversation persisted buildUuid=uuid-1 messageCount=3'); - }); - - it('syncs existing conversation and inserts only missing messages', async () => { - const conversation = { - buildUuid: 'uuid-2', - messages: [ - { role: 'user', content: 'first', timestamp: 1000 }, - { role: 'assistant', content: 'second', timestamp: 1001 }, - { role: 'assistant', content: 'third', timestamp: 1002 }, - ], - lastActivity: 1002, - }; - mockGetConversation.mockResolvedValue(conversation); - - let patchedConversation: any; - let insertedMessages: any[] = []; - setupTransaction({ - existingConversation: { buildUuid: 'uuid-2', repo: 'example-org/example-repo', model: null }, - existingMessages: [ - { role: 'user', timestamp: 1000 }, - { role: 'assistant', timestamp: 1001 }, - ], - onConversationPatch: (row) => { - patchedConversation = row; - }, - onMessageInsert: (rows) => { - insertedMessages = rows; - }, - }); - - const result = await service.persistConversation('uuid-2', 'example-org/example-repo'); - - expect(result).toBe(true); - expect(insertedMessages).toHaveLength(1); - expect(insertedMessages[0]).toMatchObject({ - buildUuid: 'uuid-2', - role: 'assistant', - content: 'third', - timestamp: 1002, - }); - expect(patchedConversation).toMatchObject({ - repo: 'example-org/example-repo', - messageCount: 3, - }); - }); - - it('truncates large tool results in metadata', async () => { - const largeResult = 'x'.repeat(20000); - const conversation = { - buildUuid: 'uuid-3', - messages: [ - { - role: 'assistant', - content: 'done', - timestamp: 1000, - debugToolData: [ - { - toolCallId: 'tc-1', - toolName: 'kubectl', - toolArgs: { cmd: 'get pods' }, - toolResult: largeResult, - }, - ], - }, - ], - lastActivity: 1000, - }; - mockGetConversation.mockResolvedValue(conversation); - - let insertedMessages: any[] = []; - setupTransaction({ - existingConversation: undefined, - existingMessages: [], - onMessageInsert: (rows) => { - insertedMessages = rows; - }, - }); - - await service.persistConversation('uuid-3', 'example-org/example-repo'); - - expect(insertedMessages).toHaveLength(1); - const storedToolData = insertedMessages[0].metadata.debugToolData; - expect(storedToolData[0].toolResult).toHaveLength(10000 + '... [truncated]'.length); - expect(storedToolData[0].toolResult).toContain('... [truncated]'); - }); - - it('returns false and logs error on Postgres failure', async () => { - const conversation = { - buildUuid: 'uuid-fail', - messages: [{ role: 'user', content: 'hi', timestamp: 1000 }], - lastActivity: 1000, - }; - mockGetConversation.mockResolvedValue(conversation); - - MockConversation.transact.mockRejectedValue(new Error('persistent failure')); - - const result = await service.persistConversation('uuid-fail', 'example-org/example-repo'); - - expect(result).toBe(false); - expect(MockConversation.transact).toHaveBeenCalledTimes(3); - expect(mockLogger.warn).toHaveBeenCalledTimes(2); - expect(mockLogger.error).toHaveBeenCalledWith( - expect.stringContaining('AI: conversation persistence failed buildUuid=uuid-fail') - ); - }); - }); -}); diff --git a/src/server/services/ai/conversation/__tests__/storage.test.ts b/src/server/services/ai/conversation/__tests__/storage.test.ts deleted file mode 100644 index 07013c52..00000000 --- a/src/server/services/ai/conversation/__tests__/storage.test.ts +++ /dev/null @@ -1,106 +0,0 @@ -/** - * Copyright 2025 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 AIAgentConversationService from '../storage'; - -describe('AIAgentConversationService', () => { - let service: AIAgentConversationService; - let mockRedis: { - get: jest.Mock; - setex: jest.Mock; - del: jest.Mock; - expire: jest.Mock; - }; - - beforeEach(() => { - mockRedis = { - get: jest.fn(), - setex: jest.fn(), - del: jest.fn(), - expire: jest.fn(), - }; - service = new AIAgentConversationService({} as any, mockRedis as any, {} as any, {} as any); - }); - - describe('getConversation', () => { - it('returns parsed conversation when it exists', async () => { - const stored = { buildUuid: 'uuid-1', messages: [{ role: 'user', content: 'hi' }], lastActivity: 123 }; - mockRedis.get.mockResolvedValue(JSON.stringify(stored)); - - const result = await service.getConversation('uuid-1'); - - expect(result).toEqual(stored); - expect(mockRedis.get).toHaveBeenCalledWith('lifecycle:agent:conversation:uuid-1'); - }); - - it('returns null when conversation does not exist', async () => { - mockRedis.get.mockResolvedValue(null); - const result = await service.getConversation('uuid-1'); - expect(result).toBeNull(); - }); - }); - - describe('addMessage', () => { - it('creates new conversation when none exists', async () => { - mockRedis.get.mockResolvedValue(null); - - const result = await service.addMessage('uuid-1', { role: 'user', content: 'hello' } as any); - - expect(result.messages).toHaveLength(1); - expect(result.buildUuid).toBe('uuid-1'); - expect(mockRedis.setex).toHaveBeenCalledWith('lifecycle:agent:conversation:uuid-1', 3600, expect.any(String)); - }); - - it('appends to existing conversation', async () => { - const existing = { buildUuid: 'uuid-1', messages: [{ role: 'user', content: 'hi' }], lastActivity: 100 }; - mockRedis.get.mockResolvedValue(JSON.stringify(existing)); - - const result = await service.addMessage('uuid-1', { role: 'assistant', content: 'hello' } as any); - - expect(result.messages).toHaveLength(2); - expect(mockRedis.setex).toHaveBeenCalledWith('lifecycle:agent:conversation:uuid-1', 3600, expect.any(String)); - }); - }); - - describe('clearConversation', () => { - it('returns message count and deletes key when conversation exists', async () => { - const stored = { - buildUuid: 'uuid-1', - messages: [{ role: 'user' }, { role: 'assistant' }, { role: 'user' }], - lastActivity: 123, - }; - mockRedis.get.mockResolvedValue(JSON.stringify(stored)); - - const count = await service.clearConversation('uuid-1'); - - expect(count).toBe(3); - expect(mockRedis.del).toHaveBeenCalledWith('lifecycle:agent:conversation:uuid-1'); - }); - - it('returns 0 when conversation does not exist', async () => { - mockRedis.get.mockResolvedValue(null); - const count = await service.clearConversation('uuid-1'); - expect(count).toBe(0); - }); - }); - - describe('refreshTTL', () => { - it('calls redis expire with correct key and TTL', async () => { - await service.refreshTTL('uuid-1'); - expect(mockRedis.expire).toHaveBeenCalledWith('lifecycle:agent:conversation:uuid-1', 3600); - }); - }); -}); diff --git a/src/server/services/ai/conversation/manager.ts b/src/server/services/ai/conversation/manager.ts deleted file mode 100644 index e560bec3..00000000 --- a/src/server/services/ai/conversation/manager.ts +++ /dev/null @@ -1,125 +0,0 @@ -/** - * Copyright 2025 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 { LLMProvider, StreamChunk } from '../types/provider'; -import { ConversationMessage, extractTextFromParts, textMessage } from '../types/message'; -import { getLogger } from 'server/lib/logger'; -import { countTokens } from '../prompts/tokenCounter'; - -export interface ConversationState { - summary: string; - identifiedIssues: Array<{ - service: string; - issue: string; - confidence: 'high' | 'medium' | 'low'; - }>; - investigatedServices: string[]; - toolsUsed: string[]; - currentTask: string; - tokenCount: number; - messageCount: number; - compressionLevel: number; -} - -export class ConversationManager { - private readonly COMPRESSION_THRESHOLD: number; - - constructor(compressionThreshold?: number) { - this.COMPRESSION_THRESHOLD = compressionThreshold || 80000; - } - - async shouldCompress(messages: ConversationMessage[]): Promise { - const tokenCount = await this.estimateTokens(messages); - return tokenCount > this.COMPRESSION_THRESHOLD; - } - - async compress( - messages: ConversationMessage[], - llmProvider: LLMProvider, - buildUuid?: string - ): Promise { - getLogger().info(`AI: compression starting messageCount=${messages.length} buildUuid=${buildUuid || 'none'}`); - const compressionPrompt = ` -Analyze this debugging conversation and create a structured summary. - -Extract: -1. What issues have been identified -2. Which services have been investigated -3. What tools were used -4. Current task/focus -5. Key findings - -Return JSON matching ConversationState schema. - -Conversation: -${this.formatMessages(messages)} -`; - - const chunks: StreamChunk[] = []; - for await (const chunk of llmProvider.streamCompletion([textMessage('user', compressionPrompt)], { - systemPrompt: 'You are a conversation summarizer.', - maxTokens: 2000, - })) { - chunks.push(chunk); - } - - const responseText = chunks - .filter((c) => c.type === 'text') - .map((c) => c.content) - .join(''); - - const state: ConversationState = JSON.parse(responseText); - state.tokenCount = await this.estimateTokens([textMessage('user', JSON.stringify(state))]); - state.messageCount = messages.length; - state.compressionLevel = 1; - - getLogger().info( - `AIAgentConversationManager: compression complete messageCount=${messages.length} tokenCount=${state.tokenCount} issueCount=${state.identifiedIssues.length} serviceCount=${state.investigatedServices.length}` - ); - - return state; - } - - buildPromptFromState(state: ConversationState): string { - return ` -# Conversation Context (Compressed) - -## Summary -${state.summary} - -## Identified Issues -${state.identifiedIssues.map((i) => `- **${i.service}**: ${i.issue} (${i.confidence} confidence)`).join('\n')} - -## Already Investigated -Services: ${state.investigatedServices.join(', ')} -Tools used: ${state.toolsUsed.join(', ')} - -## Current Task -${state.currentTask} - -Continue the investigation from this point. -`; - } - - private async estimateTokens(messages: ConversationMessage[]): Promise { - const text = messages.map((m) => extractTextFromParts(m.parts)).join(' '); - return countTokens(text); - } - - private formatMessages(messages: ConversationMessage[]): string { - return messages.map((m) => `${m.role}: ${extractTextFromParts(m.parts)}`).join('\n\n'); - } -} diff --git a/src/server/services/ai/conversation/persistence.ts b/src/server/services/ai/conversation/persistence.ts deleted file mode 100644 index 59808ea8..00000000 --- a/src/server/services/ai/conversation/persistence.ts +++ /dev/null @@ -1,191 +0,0 @@ -/** - * Copyright 2025 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 AIAgentConversationService from './storage'; -import { ConversationState, DebugMessage } from '../../types/aiAgent'; -import Conversation from 'server/models/Conversation'; -import ConversationMessage from 'server/models/ConversationMessage'; - -const TOOL_RESULT_TRUNCATION_LIMIT = 10000; -const MAX_RETRY_ATTEMPTS = 3; -const BASE_RETRY_DELAY_MS = 200; - -export default class ConversationPersistenceService { - private conversationService: AIAgentConversationService; - - constructor(conversationService: AIAgentConversationService) { - this.conversationService = conversationService; - } - - async persistConversation(buildUuid: string, repo: string, model?: string): Promise { - const logger = getLogger(); - try { - const conversation = await this.conversationService.getConversation(buildUuid); - if (!conversation || conversation.messages.length === 0) { - return false; - } - - await this.withRetry(() => this.writeToPostgres(buildUuid, repo, model, conversation)); - logger.info(`AI: conversation persisted buildUuid=${buildUuid} messageCount=${conversation.messages.length}`); - return true; - } catch (err) { - logger.error(`AI: conversation persistence failed buildUuid=${buildUuid} error=${err.message}`); - return false; - } - } - - private async writeToPostgres( - buildUuid: string, - repo: string, - model: string | undefined, - conversation: ConversationState - ): Promise { - await Conversation.transact(async (trx) => { - const existingConversation = await Conversation.query(trx).findById(buildUuid); - const existingMessages = existingConversation - ? await ConversationMessage.query(trx).where({ buildUuid }).select('role', 'timestamp') - : []; - const existingKeys = new Set(existingMessages.map((msg) => `${msg.role}:${msg.timestamp}`)); - - const messageRows = conversation.messages.map((msg) => ({ - buildUuid, - role: msg.role, - content: msg.content, - timestamp: msg.timestamp, - metadata: this.extractMessageMetadata(msg), - })); - - const newMessageRows = messageRows.filter((msg) => !existingKeys.has(`${msg.role}:${msg.timestamp}`)); - const mergedRepo = existingConversation?.repo || repo; - const mergedModel = model || existingConversation?.model || null; - const mergedMessageCount = existingMessages.length + newMessageRows.length; - const metadata = { - contextSnapshot: conversation.contextSnapshot || null, - lastActivity: conversation.lastActivity, - }; - - if (!existingConversation) { - await Conversation.query(trx).insert({ - buildUuid, - repo: mergedRepo, - model: mergedModel, - messageCount: mergedMessageCount, - metadata, - } as any); - } - - if (newMessageRows.length > 0) { - await ConversationMessage.query(trx).insert(newMessageRows as any); - } - - if (existingConversation) { - await Conversation.query(trx) - .patch({ - repo: mergedRepo, - model: mergedModel, - messageCount: mergedMessageCount, - metadata, - } as any) - .where({ buildUuid }); - } - }); - } - - private extractMessageMetadata(msg: DebugMessage): Record { - const metadata: Record = {}; - - if (msg.isSystemAction) { - metadata.isSystemAction = msg.isSystemAction; - } - if (msg.activityHistory) { - metadata.activityHistory = msg.activityHistory; - } - if (msg.evidenceItems) { - metadata.evidenceItems = msg.evidenceItems; - } - if (msg.totalInvestigationTimeMs) { - metadata.totalInvestigationTimeMs = msg.totalInvestigationTimeMs; - } - if (msg.debugContext) { - metadata.debugContext = msg.debugContext; - } - if (msg.debugToolData) { - metadata.debugToolData = this.truncateToolResults(msg.debugToolData); - } - if (msg.debugMetrics) { - metadata.debugMetrics = msg.debugMetrics; - } - - return metadata; - } - - private truncateToolResults(toolData: DebugMessage['debugToolData']): DebugMessage['debugToolData'] { - if (!toolData) { - return undefined; - } - - return toolData.map((entry) => { - if (!entry.toolResult) { - return entry; - } - - if (typeof entry.toolResult === 'string') { - if (entry.toolResult.length > TOOL_RESULT_TRUNCATION_LIMIT) { - return { - ...entry, - toolResult: entry.toolResult.substring(0, TOOL_RESULT_TRUNCATION_LIMIT) + '... [truncated]', - }; - } - return entry; - } - - if (typeof entry.toolResult === 'object') { - const stringified = JSON.stringify(entry.toolResult); - if (stringified.length > TOOL_RESULT_TRUNCATION_LIMIT) { - return { - ...entry, - toolResult: { - truncated: true, - originalLength: stringified.length, - preview: stringified.substring(0, 500), - }, - }; - } - } - - return entry; - }); - } - - private async withRetry(fn: () => Promise): Promise { - const logger = getLogger(); - for (let attempt = 1; attempt <= MAX_RETRY_ATTEMPTS; attempt++) { - try { - return await fn(); - } catch (err) { - if (attempt < MAX_RETRY_ATTEMPTS) { - const delay = BASE_RETRY_DELAY_MS * Math.pow(2, attempt - 1); - logger.warn(`AI: conversation persistence retry attempt=${attempt} error=${err.message}`); - await new Promise((resolve) => setTimeout(resolve, delay)); - } else { - throw err; - } - } - } - throw new Error('withRetry: unreachable'); - } -} diff --git a/src/server/services/ai/conversation/storage.ts b/src/server/services/ai/conversation/storage.ts deleted file mode 100644 index 6292b64f..00000000 --- a/src/server/services/ai/conversation/storage.ts +++ /dev/null @@ -1,66 +0,0 @@ -/** - * Copyright 2025 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 BaseService from '../../_service'; -import { ConversationState, DebugMessage } from '../../types/aiAgent'; - -export default class AIAgentConversationService extends BaseService { - private readonly CONVERSATION_TTL = 3600; - private readonly KEY_PREFIX = 'lifecycle:agent:conversation:'; - - private getConversationKey(buildUuid: string): string { - return `${this.KEY_PREFIX}${buildUuid}`; - } - - async getConversation(buildUuid: string): Promise { - const key = this.getConversationKey(buildUuid); - const data = await this.redis.get(key); - - if (!data) { - return null; - } - - return JSON.parse(data); - } - - async addMessage(buildUuid: string, message: DebugMessage): Promise { - const key = this.getConversationKey(buildUuid); - const conversation = (await this.getConversation(buildUuid)) || { - buildUuid, - messages: [], - lastActivity: Date.now(), - }; - - conversation.messages.push(message); - conversation.lastActivity = Date.now(); - - await this.redis.setex(key, this.CONVERSATION_TTL, JSON.stringify(conversation)); - - return conversation; - } - - async clearConversation(buildUuid: string): Promise { - const key = this.getConversationKey(buildUuid); - const conversation = await this.getConversation(buildUuid); - await this.redis.del(key); - return conversation?.messages.length || 0; - } - - async refreshTTL(buildUuid: string): Promise { - const key = this.getConversationKey(buildUuid); - await this.redis.expire(key, this.CONVERSATION_TTL); - } -} diff --git a/src/server/services/ai/errors/__tests__/classification.test.ts b/src/server/services/ai/errors/__tests__/classification.test.ts deleted file mode 100644 index e149b2e2..00000000 --- a/src/server/services/ai/errors/__tests__/classification.test.ts +++ /dev/null @@ -1,46 +0,0 @@ -/** - * Copyright 2025 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 { ErrorCategory, isRetryable } from '../classification'; - -describe('ErrorCategory', () => { - it('has exactly 4 values', () => { - const values = Object.values(ErrorCategory); - expect(values).toHaveLength(4); - expect(values).toContain('transient'); - expect(values).toContain('rate-limited'); - expect(values).toContain('deterministic'); - expect(values).toContain('ambiguous'); - }); -}); - -describe('isRetryable', () => { - it('returns true for TRANSIENT', () => { - expect(isRetryable(ErrorCategory.TRANSIENT)).toBe(true); - }); - - it('returns true for RATE_LIMITED', () => { - expect(isRetryable(ErrorCategory.RATE_LIMITED)).toBe(true); - }); - - it('returns true for AMBIGUOUS', () => { - expect(isRetryable(ErrorCategory.AMBIGUOUS)).toBe(true); - }); - - it('returns false for DETERMINISTIC', () => { - expect(isRetryable(ErrorCategory.DETERMINISTIC)).toBe(false); - }); -}); diff --git a/src/server/services/ai/errors/__tests__/providerErrors.test.ts b/src/server/services/ai/errors/__tests__/providerErrors.test.ts deleted file mode 100644 index 632d1f1d..00000000 --- a/src/server/services/ai/errors/__tests__/providerErrors.test.ts +++ /dev/null @@ -1,173 +0,0 @@ -/** - * Copyright 2025 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 OpenAI from 'openai'; -import Anthropic from '@anthropic-ai/sdk'; -import { ApiError as GeminiApiError } from '@google/genai'; -import { ErrorCategory } from '../classification'; -import { - classifyOpenAIError, - classifyAnthropicError, - classifyGeminiError, - classifyError, - createClassifiedError, -} from '../providerErrors'; - -const emptyHeaders = new Headers(); - -describe('classifyOpenAIError', () => { - it('maps RateLimitError to RATE_LIMITED', () => { - const err = new OpenAI.RateLimitError(429, undefined, 'rate limited', emptyHeaders); - expect(classifyOpenAIError(err)).toBe(ErrorCategory.RATE_LIMITED); - }); - - it('maps InternalServerError to TRANSIENT', () => { - const err = new OpenAI.InternalServerError(500, undefined, 'internal', emptyHeaders); - expect(classifyOpenAIError(err)).toBe(ErrorCategory.TRANSIENT); - }); - - it('maps APIConnectionError to TRANSIENT', () => { - const err = new OpenAI.APIConnectionError({ message: 'connection error' }); - expect(classifyOpenAIError(err)).toBe(ErrorCategory.TRANSIENT); - }); - - it('maps BadRequestError to DETERMINISTIC', () => { - const err = new OpenAI.BadRequestError(400, undefined, 'bad request', emptyHeaders); - expect(classifyOpenAIError(err)).toBe(ErrorCategory.DETERMINISTIC); - }); - - it('maps AuthenticationError to DETERMINISTIC', () => { - const err = new OpenAI.AuthenticationError(401, undefined, 'auth error', emptyHeaders); - expect(classifyOpenAIError(err)).toBe(ErrorCategory.DETERMINISTIC); - }); - - it('maps unknown Error to AMBIGUOUS', () => { - const err = new Error('something unexpected'); - expect(classifyOpenAIError(err)).toBe(ErrorCategory.AMBIGUOUS); - }); -}); - -describe('classifyAnthropicError', () => { - it('maps RateLimitError to RATE_LIMITED', () => { - const err = new Anthropic.RateLimitError(429, undefined, 'rate limited', emptyHeaders); - expect(classifyAnthropicError(err)).toBe(ErrorCategory.RATE_LIMITED); - }); - - it('maps InternalServerError to TRANSIENT', () => { - const err = new Anthropic.InternalServerError(500, undefined, 'internal', emptyHeaders); - expect(classifyAnthropicError(err)).toBe(ErrorCategory.TRANSIENT); - }); - - it('maps BadRequestError to DETERMINISTIC', () => { - const err = new Anthropic.BadRequestError(400, undefined, 'bad request', emptyHeaders); - expect(classifyAnthropicError(err)).toBe(ErrorCategory.DETERMINISTIC); - }); - - it('maps unknown Error to AMBIGUOUS', () => { - const err = new Error('something unexpected'); - expect(classifyAnthropicError(err)).toBe(ErrorCategory.AMBIGUOUS); - }); -}); - -describe('classifyGeminiError', () => { - it('maps GeminiApiError with status 429 to RATE_LIMITED', () => { - const err = new GeminiApiError({ message: 'rate limited', status: 429 }); - expect(classifyGeminiError(err)).toBe(ErrorCategory.RATE_LIMITED); - }); - - it('maps GeminiApiError with status >= 500 to TRANSIENT', () => { - const err = new GeminiApiError({ message: 'server error', status: 500 }); - expect(classifyGeminiError(err)).toBe(ErrorCategory.TRANSIENT); - }); - - it('maps GeminiApiError with status 400 to DETERMINISTIC', () => { - const err = new GeminiApiError({ message: 'bad request', status: 400 }); - expect(classifyGeminiError(err)).toBe(ErrorCategory.DETERMINISTIC); - }); - - it('maps error with MALFORMED_FUNCTION_CALL to TRANSIENT', () => { - const err = new Error('Gemini MALFORMED_FUNCTION_CALL detected'); - expect(classifyGeminiError(err)).toBe(ErrorCategory.TRANSIENT); - }); - - it('maps error with empty response and STOP to AMBIGUOUS', () => { - const err = new Error('Gemini returned an empty response. finishReason: STOP'); - expect(classifyGeminiError(err)).toBe(ErrorCategory.AMBIGUOUS); - }); - - it('maps unknown Error to AMBIGUOUS', () => { - const err = new Error('something unexpected'); - expect(classifyGeminiError(err)).toBe(ErrorCategory.AMBIGUOUS); - }); -}); - -describe('classifyError', () => { - it('dispatches to classifyOpenAIError for openai', () => { - const err = new OpenAI.RateLimitError(429, undefined, 'rate limited', emptyHeaders); - expect(classifyError('openai', err)).toBe(ErrorCategory.RATE_LIMITED); - }); - - it('dispatches to classifyAnthropicError for anthropic', () => { - const err = new Anthropic.InternalServerError(500, undefined, 'internal', emptyHeaders); - expect(classifyError('anthropic', err)).toBe(ErrorCategory.TRANSIENT); - }); - - it('dispatches to classifyGeminiError for gemini', () => { - const err = new GeminiApiError({ message: 'rate limited', status: 429 }); - expect(classifyError('gemini', err)).toBe(ErrorCategory.RATE_LIMITED); - }); - - it('returns AMBIGUOUS for unknown provider', () => { - const err = new Error('some error'); - expect(classifyError('unknown-provider', err)).toBe(ErrorCategory.AMBIGUOUS); - }); -}); - -describe('createClassifiedError', () => { - it('returns a ClassifiedError with all fields populated', () => { - const err = new OpenAI.InternalServerError(500, undefined, 'internal', emptyHeaders); - const classified = createClassifiedError('openai', err); - expect(classified.category).toBe(ErrorCategory.TRANSIENT); - expect(classified.original).toBe(err); - expect(classified.retryable).toBe(true); - expect(classified.providerName).toBe('openai'); - }); - - it('sets retryable=true for transient errors', () => { - const err = new OpenAI.InternalServerError(500, undefined, 'internal', emptyHeaders); - const classified = createClassifiedError('openai', err); - expect(classified.retryable).toBe(true); - }); - - it('sets retryable=false for deterministic errors', () => { - const err = new OpenAI.BadRequestError(400, undefined, 'bad', emptyHeaders); - const classified = createClassifiedError('openai', err); - expect(classified.retryable).toBe(false); - }); - - it('extracts httpStatus from error.status', () => { - const err = new OpenAI.RateLimitError(429, undefined, 'rate limited', emptyHeaders); - const classified = createClassifiedError('openai', err); - expect(classified.httpStatus).toBe(429); - }); - - it('wraps non-Error values in new Error', () => { - const classified = createClassifiedError('openai', 'string error'); - expect(classified.original).toBeInstanceOf(Error); - expect(classified.original.message).toBe('string error'); - expect(classified.category).toBe(ErrorCategory.AMBIGUOUS); - }); -}); diff --git a/src/server/services/ai/errors/__tests__/retryBudget.test.ts b/src/server/services/ai/errors/__tests__/retryBudget.test.ts deleted file mode 100644 index 87e6158b..00000000 --- a/src/server/services/ai/errors/__tests__/retryBudget.test.ts +++ /dev/null @@ -1,95 +0,0 @@ -/** - * Copyright 2025 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 { RetryBudget } from '../retryBudget'; - -describe('RetryBudget', () => { - it('defaults to maxRetries=10', () => { - const budget = new RetryBudget(); - expect(budget.used).toBe(0); - expect(budget.exhausted).toBe(false); - for (let i = 0; i < 10; i++) { - expect(budget.canRetry()).toBe(true); - budget.consume(); - } - expect(budget.canRetry()).toBe(false); - }); - - it('canRetry returns true initially', () => { - const budget = new RetryBudget(); - expect(budget.canRetry()).toBe(true); - }); - - it('consume decrements remaining', () => { - const budget = new RetryBudget(3); - expect(budget.used).toBe(0); - budget.consume(); - expect(budget.used).toBe(1); - budget.consume(); - expect(budget.used).toBe(2); - }); - - it('exhausted returns true after consuming all retries', () => { - const budget = new RetryBudget(2); - budget.consume(); - budget.consume(); - expect(budget.exhausted).toBe(true); - }); - - it('canRetry returns false when exhausted', () => { - const budget = new RetryBudget(1); - budget.consume(); - expect(budget.canRetry()).toBe(false); - }); - - it('used returns correct count', () => { - const budget = new RetryBudget(5); - expect(budget.used).toBe(0); - budget.consume(); - budget.consume(); - budget.consume(); - expect(budget.used).toBe(3); - }); - - it('supports custom maxRetries via constructor', () => { - const budget = new RetryBudget(3); - budget.consume(); - budget.consume(); - budget.consume(); - expect(budget.exhausted).toBe(true); - expect(budget.used).toBe(3); - }); - - it('reset restores budget to full', () => { - const budget = new RetryBudget(5); - budget.consume(); - budget.consume(); - expect(budget.used).toBe(2); - budget.reset(); - expect(budget.used).toBe(0); - expect(budget.canRetry()).toBe(true); - expect(budget.exhausted).toBe(false); - }); - - it('consume does not go below 0', () => { - const budget = new RetryBudget(1); - budget.consume(); - budget.consume(); - budget.consume(); - expect(budget.used).toBe(1); - expect(budget.exhausted).toBe(true); - }); -}); diff --git a/src/server/services/ai/errors/__tests__/userMessages.test.ts b/src/server/services/ai/errors/__tests__/userMessages.test.ts deleted file mode 100644 index e7afae3c..00000000 --- a/src/server/services/ai/errors/__tests__/userMessages.test.ts +++ /dev/null @@ -1,127 +0,0 @@ -/** - * Copyright 2025 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 { ErrorCategory } from '../classification'; -import { getUserErrorMessage, getSuggestedAction, isAuthError, ErrorContext } from '../userMessages'; - -describe('getUserErrorMessage', () => { - const baseCtx: ErrorContext = { - modelName: 'GPT-4o', - providerName: 'OpenAI', - }; - - it('returns rate limit message with retry seconds', () => { - const msg = getUserErrorMessage(ErrorCategory.RATE_LIMITED, { - ...baseCtx, - retryAfter: 30, - }); - expect(msg).toBe('GPT-4o is rate limited. Retrying in 30s...'); - }); - - it('returns rate limit message without retry seconds', () => { - const msg = getUserErrorMessage(ErrorCategory.RATE_LIMITED, { - ...baseCtx, - retryAfter: null, - }); - expect(msg).toBe('GPT-4o is rate limited. Please wait and try again.'); - }); - - it('returns transient message', () => { - const msg = getUserErrorMessage(ErrorCategory.TRANSIENT, { - ...baseCtx, - modelName: 'Claude 3.5', - }); - expect(msg).toBe('Claude 3.5 is temporarily unavailable. Retrying...'); - }); - - it('returns deterministic auth error message', () => { - const msg = getUserErrorMessage(ErrorCategory.DETERMINISTIC, { - ...baseCtx, - providerName: 'Anthropic', - isAuthError: true, - }); - expect(msg).toBe('Anthropic API key is invalid. Check AI agent configuration in admin settings.'); - }); - - it('returns deterministic non-auth error message', () => { - const msg = getUserErrorMessage(ErrorCategory.DETERMINISTIC, { - ...baseCtx, - isAuthError: false, - }); - expect(msg).toBe('Request failed. Please try a different approach.'); - }); - - it('returns ambiguous fallback message', () => { - const msg = getUserErrorMessage(ErrorCategory.AMBIGUOUS, baseCtx); - expect(msg).toBe('Something went wrong. Please try again.'); - }); -}); - -describe('getSuggestedAction', () => { - it('returns retry for RATE_LIMITED', () => { - expect(getSuggestedAction(ErrorCategory.RATE_LIMITED)).toBe('retry'); - }); - - it('returns switch-model for TRANSIENT', () => { - expect(getSuggestedAction(ErrorCategory.TRANSIENT)).toBe('switch-model'); - }); - - it('returns check-config for DETERMINISTIC auth error', () => { - expect(getSuggestedAction(ErrorCategory.DETERMINISTIC, true)).toBe('check-config'); - }); - - it('returns null for DETERMINISTIC non-auth error', () => { - expect(getSuggestedAction(ErrorCategory.DETERMINISTIC, false)).toBeNull(); - }); - - it('returns retry for AMBIGUOUS', () => { - expect(getSuggestedAction(ErrorCategory.AMBIGUOUS)).toBe('retry'); - }); -}); - -describe('isAuthError', () => { - it('detects 401 status', () => { - const err = Object.assign(new Error('Unauthorized'), { status: 401 }); - expect(isAuthError(err)).toBe(true); - }); - - it('detects 403 status', () => { - const err = Object.assign(new Error('Forbidden'), { status: 403 }); - expect(isAuthError(err)).toBe(true); - }); - - it('detects AuthenticationError by name', () => { - const err = new Error('auth failed'); - err.name = 'AuthenticationError'; - expect(isAuthError(err)).toBe(true); - }); - - it('detects PermissionDeniedError by name', () => { - const err = new Error('denied'); - err.name = 'PermissionDeniedError'; - expect(isAuthError(err)).toBe(true); - }); - - it('returns false for non-auth errors', () => { - expect(isAuthError(new Error('something else'))).toBe(false); - }); - - it('returns false for non-error values', () => { - expect(isAuthError(null)).toBe(false); - expect(isAuthError(undefined)).toBe(false); - expect(isAuthError('string')).toBe(false); - }); -}); diff --git a/src/server/services/ai/errors/classification.ts b/src/server/services/ai/errors/classification.ts deleted file mode 100644 index 16e3d905..00000000 --- a/src/server/services/ai/errors/classification.ts +++ /dev/null @@ -1,49 +0,0 @@ -/** - * Copyright 2025 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. - */ - -export enum ErrorCategory { - TRANSIENT = 'transient', - RATE_LIMITED = 'rate-limited', - DETERMINISTIC = 'deterministic', - AMBIGUOUS = 'ambiguous', -} - -export interface ClassifiedError { - category: ErrorCategory; - original: Error; - retryable: boolean; - providerName: string; - httpStatus?: number; - finishReason?: string; - retryAfter?: number | null; -} - -export function isRetryable(category: ErrorCategory): boolean { - return ( - category === ErrorCategory.TRANSIENT || - category === ErrorCategory.RATE_LIMITED || - category === ErrorCategory.AMBIGUOUS - ); -} - -export function isRateLimitError(error: any): boolean { - return ( - error?.status === 429 || - error?.error?.error?.type === 'rate_limit_error' || - error?.message?.includes('RATE_LIMIT_EXCEEDED') || - error?.message?.includes('quota exceeded') - ); -} diff --git a/src/server/services/ai/errors/index.ts b/src/server/services/ai/errors/index.ts deleted file mode 100644 index 8eb9dfaf..00000000 --- a/src/server/services/ai/errors/index.ts +++ /dev/null @@ -1,29 +0,0 @@ -/** - * Copyright 2025 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. - */ - -export { ErrorCategory, isRetryable } from './classification'; -export type { ClassifiedError } from './classification'; -export { - classifyOpenAIError, - classifyAnthropicError, - classifyGeminiError, - classifyError, - createClassifiedError, - extractRetryAfter, -} from './providerErrors'; -export { RetryBudget } from './retryBudget'; -export { getUserErrorMessage, getSuggestedAction, isAuthError } from './userMessages'; -export type { ErrorContext } from './userMessages'; diff --git a/src/server/services/ai/errors/providerErrors.ts b/src/server/services/ai/errors/providerErrors.ts deleted file mode 100644 index 67f0190e..00000000 --- a/src/server/services/ai/errors/providerErrors.ts +++ /dev/null @@ -1,116 +0,0 @@ -/** - * Copyright 2025 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 OpenAI from 'openai'; -import Anthropic from '@anthropic-ai/sdk'; -import { ApiError as GeminiApiError } from '@google/genai'; -import { ErrorCategory, ClassifiedError, isRetryable } from './classification'; - -const MAX_RETRY_AFTER_SECONDS = 300; - -export function extractRetryAfter(error: unknown): number | null { - const err = error as any; - let raw: string | undefined; - - if (err?.headers?.['retry-after'] != null) { - raw = String(err.headers['retry-after']); - } else if (typeof err?.headers?.get === 'function') { - const val = err.headers.get('retry-after'); - if (val != null) raw = String(val); - } - - if (raw == null) return null; - - const seconds = Number(raw); - if (!isNaN(seconds) && seconds >= 0) { - return Math.min(seconds, MAX_RETRY_AFTER_SECONDS); - } - - const date = Date.parse(raw); - if (!isNaN(date)) { - const delta = Math.ceil((date - Date.now()) / 1000); - if (delta > 0) return Math.min(delta, MAX_RETRY_AFTER_SECONDS); - return 0; - } - - return null; -} - -export function classifyOpenAIError(error: unknown): ErrorCategory { - if (error instanceof OpenAI.RateLimitError) return ErrorCategory.RATE_LIMITED; - if (error instanceof OpenAI.InternalServerError) return ErrorCategory.TRANSIENT; - if (error instanceof OpenAI.APIConnectionError) return ErrorCategory.TRANSIENT; - if (error instanceof OpenAI.ConflictError) return ErrorCategory.TRANSIENT; - if (error instanceof OpenAI.BadRequestError) return ErrorCategory.DETERMINISTIC; - if (error instanceof OpenAI.AuthenticationError) return ErrorCategory.DETERMINISTIC; - if (error instanceof OpenAI.PermissionDeniedError) return ErrorCategory.DETERMINISTIC; - if (error instanceof OpenAI.NotFoundError) return ErrorCategory.DETERMINISTIC; - if (error instanceof OpenAI.UnprocessableEntityError) return ErrorCategory.DETERMINISTIC; - return ErrorCategory.AMBIGUOUS; -} - -export function classifyAnthropicError(error: unknown): ErrorCategory { - if (error instanceof Anthropic.RateLimitError) return ErrorCategory.RATE_LIMITED; - if (error instanceof Anthropic.InternalServerError) return ErrorCategory.TRANSIENT; - if (error instanceof Anthropic.APIConnectionError) return ErrorCategory.TRANSIENT; - if (error instanceof Anthropic.ConflictError) return ErrorCategory.TRANSIENT; - if (error instanceof Anthropic.BadRequestError) return ErrorCategory.DETERMINISTIC; - if (error instanceof Anthropic.AuthenticationError) return ErrorCategory.DETERMINISTIC; - if (error instanceof Anthropic.PermissionDeniedError) return ErrorCategory.DETERMINISTIC; - if (error instanceof Anthropic.NotFoundError) return ErrorCategory.DETERMINISTIC; - if (error instanceof Anthropic.UnprocessableEntityError) return ErrorCategory.DETERMINISTIC; - return ErrorCategory.AMBIGUOUS; -} - -export function classifyGeminiError(error: unknown): ErrorCategory { - if (error instanceof GeminiApiError) { - if (error.status === 429) return ErrorCategory.RATE_LIMITED; - if (error.status !== undefined && error.status >= 500) return ErrorCategory.TRANSIENT; - if (error.status === 400 || error.status === 401 || error.status === 403) return ErrorCategory.DETERMINISTIC; - } - if (error instanceof Error) { - if (error.message.includes('MALFORMED_FUNCTION_CALL')) return ErrorCategory.TRANSIENT; - if (error.message.includes('empty response') && error.message.includes('STOP')) return ErrorCategory.AMBIGUOUS; - } - return ErrorCategory.AMBIGUOUS; -} - -export function classifyError(providerName: string, error: unknown): ErrorCategory { - switch (providerName) { - case 'openai': - return classifyOpenAIError(error); - case 'anthropic': - return classifyAnthropicError(error); - case 'gemini': - return classifyGeminiError(error); - default: - return ErrorCategory.AMBIGUOUS; - } -} - -export function createClassifiedError(providerName: string, error: unknown): ClassifiedError { - const category = classifyError(providerName, error); - const original = error instanceof Error ? error : new Error(String(error)); - return { - category, - original, - retryable: isRetryable(category), - providerName, - httpStatus: (error as any)?.status ?? (error as any)?.statusCode, - finishReason: (error as any)?.finishReason, - retryAfter: extractRetryAfter(error), - }; -} diff --git a/src/server/services/ai/errors/retryBudget.ts b/src/server/services/ai/errors/retryBudget.ts deleted file mode 100644 index a61a66c4..00000000 --- a/src/server/services/ai/errors/retryBudget.ts +++ /dev/null @@ -1,47 +0,0 @@ -/** - * Copyright 2025 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. - */ - -export class RetryBudget { - private readonly maxRetries: number; - private remaining: number; - - constructor(maxRetries: number = 10) { - this.maxRetries = maxRetries; - this.remaining = maxRetries; - } - - canRetry(): boolean { - return this.remaining > 0; - } - - consume(): void { - if (this.remaining > 0) { - this.remaining--; - } - } - - get exhausted(): boolean { - return this.remaining <= 0; - } - - get used(): number { - return this.maxRetries - this.remaining; - } - - reset(): void { - this.remaining = this.maxRetries; - } -} diff --git a/src/server/services/ai/errors/userMessages.ts b/src/server/services/ai/errors/userMessages.ts deleted file mode 100644 index f3d69e55..00000000 --- a/src/server/services/ai/errors/userMessages.ts +++ /dev/null @@ -1,68 +0,0 @@ -/** - * Copyright 2025 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 { ErrorCategory } from './classification'; - -export interface ErrorContext { - modelName: string; - providerName: string; - retryAfter?: number | null; - isAuthError?: boolean; -} - -const messageBuilders: Record string> = { - [ErrorCategory.RATE_LIMITED]: (ctx) => - ctx.retryAfter - ? `${ctx.modelName} is rate limited. Retrying in ${ctx.retryAfter}s...` - : `${ctx.modelName} is rate limited. Please wait and try again.`, - [ErrorCategory.TRANSIENT]: (ctx) => `${ctx.modelName} is temporarily unavailable. Retrying...`, - [ErrorCategory.DETERMINISTIC]: (ctx) => - ctx.isAuthError - ? `${ctx.providerName} API key is invalid. Check AI agent configuration in admin settings.` - : 'Request failed. Please try a different approach.', - [ErrorCategory.AMBIGUOUS]: () => 'Something went wrong. Please try again.', -}; - -export function getUserErrorMessage(category: ErrorCategory, ctx: ErrorContext): string { - return messageBuilders[category](ctx); -} - -export function getSuggestedAction( - category: ErrorCategory, - authError?: boolean -): 'retry' | 'switch-model' | 'check-config' | null { - switch (category) { - case ErrorCategory.RATE_LIMITED: - return 'retry'; - case ErrorCategory.TRANSIENT: - return 'switch-model'; - case ErrorCategory.DETERMINISTIC: - return authError ? 'check-config' : null; - case ErrorCategory.AMBIGUOUS: - return 'retry'; - } -} - -const AUTH_ERROR_NAMES = new Set(['AuthenticationError', 'PermissionDeniedError']); -const AUTH_STATUS_CODES = new Set([401, 403]); - -export function isAuthError(error: unknown): boolean { - if (!error || typeof error !== 'object') return false; - const err = error as Record; - if (AUTH_ERROR_NAMES.has(err.name)) return true; - if (AUTH_STATUS_CODES.has(err.status)) return true; - return false; -} diff --git a/src/server/services/ai/evidence/__tests__/extractor.test.ts b/src/server/services/ai/evidence/__tests__/extractor.test.ts deleted file mode 100644 index 845ef268..00000000 --- a/src/server/services/ai/evidence/__tests__/extractor.test.ts +++ /dev/null @@ -1,212 +0,0 @@ -/** - * Copyright 2025 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 { extractEvidence, generateResultPreview } from '../extractor'; -import type { ToolResult } from '../../types/tool'; - -const ctx = { toolCallId: 'tc-1' }; - -function ok(agentContent: string): ToolResult { - return { success: true, agentContent }; -} - -function fail(): ToolResult { - return { success: false }; -} - -describe('extractEvidence', () => { - it('returns evidence_file for get_file', () => { - const result = extractEvidence( - 'get_file', - { file_path: 'src/index.ts', repository_owner: 'org', repository_name: 'repo' }, - ok(JSON.stringify({ path: 'src/index.ts', content: 'hello' })), - ctx - ); - expect(result).toHaveLength(1); - expect(result[0]).toMatchObject({ - type: 'evidence_file', - toolCallId: 'tc-1', - filePath: 'src/index.ts', - repository: 'org/repo', - }); - }); - - it('infers typescript language for .ts file', () => { - const result = extractEvidence( - 'get_file', - { file_path: 'src/app.ts' }, - ok(JSON.stringify({ path: 'src/app.ts' })), - ctx - ); - expect(result[0]).toHaveProperty('language', 'typescript'); - }); - - it('returns undefined language for unknown extension', () => { - const result = extractEvidence( - 'get_file', - { file_path: 'data.xyz' }, - ok(JSON.stringify({ path: 'data.xyz' })), - ctx - ); - expect(result[0]).toHaveProperty('language', undefined); - }); - - it('returns evidence_commit + evidence_file for update_file with commit info', () => { - const result = extractEvidence( - 'update_file', - { file_path: 'src/a.ts', commit_message: 'fix bug', repository_owner: 'o', repository_name: 'r' }, - ok(JSON.stringify({ commit_sha: 'abc123', commit_url: 'https://github.com/o/r/commit/abc123' })), - ctx - ); - expect(result).toHaveLength(2); - expect(result[0]).toMatchObject({ type: 'evidence_commit', commitUrl: 'https://github.com/o/r/commit/abc123' }); - expect(result[1]).toMatchObject({ type: 'evidence_file', filePath: 'src/a.ts' }); - }); - - it('returns only evidence_file for update_file without commit info', () => { - const result = extractEvidence( - 'update_file', - { file_path: 'src/a.ts' }, - ok(JSON.stringify({ message: 'ok' })), - ctx - ); - expect(result).toHaveLength(1); - expect(result[0]).toMatchObject({ type: 'evidence_file' }); - }); - - it('returns evidence_resource for get_k8s_resources', () => { - const result = extractEvidence( - 'get_k8s_resources', - { resource_type: 'deployment', name: 'web', namespace: 'ns' }, - ok(JSON.stringify({})), - ctx - ); - expect(result).toHaveLength(1); - expect(result[0]).toMatchObject({ - type: 'evidence_resource', - resourceType: 'deployment', - resourceName: 'web', - namespace: 'ns', - }); - }); - - it('returns evidence_resource with type=pod for get_pod_logs', () => { - const result = extractEvidence( - 'get_pod_logs', - { pod_name: 'web-abc', namespace: 'ns' }, - ok(JSON.stringify({ logs: 'line1' })), - ctx - ); - expect(result[0]).toMatchObject({ type: 'evidence_resource', resourceType: 'pod', resourceName: 'web-abc' }); - }); - - it('returns evidence_resource for patch_k8s_resource', () => { - const result = extractEvidence( - 'patch_k8s_resource', - { resource_type: 'deployment', name: 'api', namespace: 'ns' }, - ok(JSON.stringify({})), - ctx - ); - expect(result[0]).toMatchObject({ type: 'evidence_resource', resourceType: 'deployment', resourceName: 'api' }); - }); - - it('returns evidence_resource for get_lifecycle_logs', () => { - const result = extractEvidence( - 'get_lifecycle_logs', - { resource_type: 'service', name: 'worker', namespace: 'ns' }, - ok(JSON.stringify({})), - ctx - ); - expect(result[0]).toMatchObject({ type: 'evidence_resource' }); - }); - - it('returns empty array for unknown tool', () => { - expect(extractEvidence('unknown_tool', {}, ok('{}'), ctx)).toEqual([]); - }); - - it('returns empty array for failed result', () => { - expect(extractEvidence('get_file', { file_path: 'a.ts' }, fail(), ctx)).toEqual([]); - }); - - it('returns empty array on extraction error', () => { - expect(extractEvidence('get_file', null as any, ok('not-json'), ctx)).toEqual([]); - }); -}); - -describe('generateResultPreview', () => { - it('returns path and line count for get_file', () => { - const preview = generateResultPreview( - 'get_file', - { file_path: 'src/a.ts' }, - ok(JSON.stringify({ path: 'src/a.ts', content: 'line1\nline2\nline3' })) - ); - expect(preview).toBe('src/a.ts (3 lines)'); - }); - - it('returns committed message for update_file', () => { - const preview = generateResultPreview('update_file', { commit_message: 'fix typo' }, ok(JSON.stringify({}))); - expect(preview).toBe('Committed: fix typo'); - }); - - it('returns pod phase summary for get_k8s_resources with pods', () => { - const preview = generateResultPreview( - 'get_k8s_resources', - {}, - ok(JSON.stringify({ pods: [{ phase: 'Running' }, { phase: 'Running' }, { phase: 'Pending' }] })) - ); - expect(preview).toBe('3 pods: 2 Running, 1 Pending'); - }); - - it('returns item count for get_k8s_resources with items', () => { - const preview = generateResultPreview( - 'get_k8s_resources', - { resource_type: 'service' }, - ok(JSON.stringify({ items: [{}, {}, {}] })) - ); - expect(preview).toBe('3 service found'); - }); - - it('returns log line count for get_pod_logs', () => { - const preview = generateResultPreview( - 'get_pod_logs', - { pod_name: 'web-1' }, - ok(JSON.stringify({ logs: 'a\nb\nc' })) - ); - expect(preview).toBe('3 log lines from web-1'); - }); - - it('returns row count for query_database', () => { - const preview = generateResultPreview('query_database', {}, ok(JSON.stringify({ rows: [{ id: 1 }, { id: 2 }] }))); - expect(preview).toBe('2 rows returned'); - }); - - it('returns patched info for patch_k8s_resource', () => { - const preview = generateResultPreview( - 'patch_k8s_resource', - { resource_type: 'deployment', name: 'api' }, - ok(JSON.stringify({})) - ); - expect(preview).toBe('Patched deployment/api'); - }); - - it('returns undefined for unknown tool', () => { - expect(generateResultPreview('unknown', {}, ok('{}'))).toBeUndefined(); - }); - - it('returns undefined for failed result', () => { - expect(generateResultPreview('get_file', {}, fail())).toBeUndefined(); - }); -}); diff --git a/src/server/services/ai/evidence/extractor.ts b/src/server/services/ai/evidence/extractor.ts deleted file mode 100644 index 4b251b11..00000000 --- a/src/server/services/ai/evidence/extractor.ts +++ /dev/null @@ -1,226 +0,0 @@ -/** - * Copyright 2025 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 { AIChatEvidenceEvent } from 'shared/types/aiChat'; -import type { ToolResult } from '../types/tool'; - -export interface EvidenceExtractorContext { - toolCallId: string; - repositoryOwner?: string; - repositoryName?: string; - commitSha?: string; -} - -const EXTENSION_LANGUAGE_MAP: Record = { - ts: 'typescript', - tsx: 'typescriptreact', - js: 'javascript', - jsx: 'javascriptreact', - py: 'python', - go: 'go', - yaml: 'yaml', - yml: 'yaml', - json: 'json', - md: 'markdown', - rs: 'rust', - rb: 'ruby', - java: 'java', - sh: 'shell', - css: 'css', - html: 'html', -}; - -function inferLanguage(filePath: string): string | undefined { - const ext = filePath.split('.').pop()?.toLowerCase(); - if (!ext) return undefined; - return EXTENSION_LANGUAGE_MAP[ext]; -} - -function truncate(str: string, maxLen: number): string { - if (str.length <= maxLen) return str; - return str.slice(0, maxLen - 3) + '...'; -} - -function safeParse(content: unknown): Record | null { - if (typeof content === 'string') { - try { - return JSON.parse(content); - } catch { - return null; - } - } - if (content && typeof content === 'object') { - return content as Record; - } - return null; -} - -export function extractEvidence( - toolName: string, - toolArgs: Record, - result: ToolResult, - context: EvidenceExtractorContext -): AIChatEvidenceEvent[] { - try { - if (!result.success) return []; - - switch (toolName) { - case 'get_file': { - const parsed = safeParse(result.agentContent); - const filePath = String((parsed && parsed.path) || toolArgs.file_path || ''); - const repository = `${toolArgs.repository_owner || ''}/${toolArgs.repository_name || ''}`; - const branch = toolArgs.branch ? String(toolArgs.branch) : undefined; - return [ - { - type: 'evidence_file', - toolCallId: context.toolCallId, - filePath, - repository, - branch, - language: inferLanguage(filePath), - }, - ]; - } - - case 'update_file': { - const parsed = safeParse(result.agentContent); - const events: AIChatEvidenceEvent[] = []; - - if (parsed && parsed.commit_sha && parsed.commit_url) { - events.push({ - type: 'evidence_commit', - toolCallId: context.toolCallId, - commitUrl: String(parsed.commit_url), - commitMessage: String(toolArgs.commit_message || ''), - filePaths: [String(toolArgs.file_path || '')], - }); - } - - events.push({ - type: 'evidence_file', - toolCallId: context.toolCallId, - filePath: String(toolArgs.file_path || ''), - repository: `${toolArgs.repository_owner || ''}/${toolArgs.repository_name || ''}`, - branch: toolArgs.branch ? String(toolArgs.branch) : undefined, - }); - - return events; - } - - case 'get_k8s_resources': - case 'get_pod_logs': - case 'patch_k8s_resource': - case 'get_lifecycle_logs': { - let resourceType: string; - let resourceName: string; - - if (toolName === 'get_pod_logs') { - resourceType = 'pod'; - resourceName = String(toolArgs.pod_name || ''); - } else { - resourceType = String(toolArgs.resource_type || 'unknown'); - resourceName = String(toolArgs.name || ''); - } - - return [ - { - type: 'evidence_resource', - toolCallId: context.toolCallId, - resourceType, - resourceName, - namespace: String(toolArgs.namespace || ''), - }, - ]; - } - - default: - return []; - } - } catch { - return []; - } -} - -export function generateResultPreview( - toolName: string, - toolArgs: Record, - result: ToolResult -): string | undefined { - try { - if (!result.success) return undefined; - - const parsed = safeParse(result.agentContent); - - switch (toolName) { - case 'get_file': { - if (!parsed) return undefined; - const path = String(parsed.path || toolArgs.file_path || ''); - const content = parsed.content; - const lineCount = typeof content === 'string' ? content.split('\n').length : 0; - return truncate(`${path} (${lineCount} lines)`, 100); - } - - case 'update_file': { - const message = String((parsed && parsed.message) || toolArgs.commit_message || ''); - return truncate(`Committed: ${message}`, 100); - } - - case 'get_k8s_resources': { - if (!parsed) return undefined; - if (Array.isArray(parsed.pods)) { - const total = parsed.pods.length; - const phases: Record = {}; - for (const pod of parsed.pods as Array>) { - const phase = String(pod.phase || 'Unknown'); - phases[phase] = (phases[phase] || 0) + 1; - } - const phaseSummary = Object.entries(phases) - .map(([p, c]) => `${c} ${p}`) - .join(', '); - return truncate(`${total} pods: ${phaseSummary}`, 100); - } - if (Array.isArray(parsed.items)) { - return truncate(`${parsed.items.length} ${toolArgs.resource_type || 'resources'} found`, 100); - } - return undefined; - } - - case 'get_pod_logs': { - if (!parsed) return undefined; - const logContent = String(parsed.logs || parsed.content || ''); - const lineCount = logContent ? logContent.split('\n').length : 0; - return truncate(`${lineCount} log lines from ${toolArgs.pod_name || 'pod'}`, 100); - } - - case 'query_database': { - if (!parsed) return undefined; - if (Array.isArray(parsed.rows)) { - return truncate(`${parsed.rows.length} rows returned`, 100); - } - return undefined; - } - - case 'patch_k8s_resource': { - return truncate(`Patched ${toolArgs.resource_type || 'resource'}/${toolArgs.name || ''}`, 100); - } - - default: - return undefined; - } - } catch { - return undefined; - } -} diff --git a/src/server/services/ai/feedback/FeedbackService.ts b/src/server/services/ai/feedback/FeedbackService.ts deleted file mode 100644 index 613554b5..00000000 --- a/src/server/services/ai/feedback/FeedbackService.ts +++ /dev/null @@ -1,276 +0,0 @@ -/** - * Copyright 2025 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 ConversationPersistenceService from '../conversation/persistence'; -import Conversation from 'server/models/Conversation'; -import ConversationMessage from 'server/models/ConversationMessage'; -import MessageFeedback from 'server/models/MessageFeedback'; -import ConversationFeedback from 'server/models/ConversationFeedback'; - -interface MessageFeedbackParams { - buildUuid: string; - messageId?: number; - messageTimestamp?: number; - rating: 'up' | 'down'; - text?: string; - userIdentifier?: string; - repo: string; - prNumber?: number; -} - -interface ConversationFeedbackParams { - buildUuid: string; - rating: 'up' | 'down'; - text?: string; - userIdentifier?: string; - repo: string; - prNumber?: number; -} - -export default class FeedbackService { - private static readonly TIMESTAMP_FALLBACK_BACKWARD_WINDOW_MS = 60_000; - private static readonly TIMESTAMP_FALLBACK_FORWARD_WINDOW_MS = 10 * 60_000; - - private persistenceService: ConversationPersistenceService; - - constructor(persistenceService: ConversationPersistenceService) { - this.persistenceService = persistenceService; - } - - private isUniqueViolation(error: unknown): boolean { - if (!error || typeof error !== 'object') { - return false; - } - - const code = (error as { code?: unknown }).code; - return code === '23505'; - } - - private async resolveAssistantMessageByTimestamp( - buildUuid: string, - messageTimestamp: number - ): Promise { - // Prefer exact timestamp matches when available. - const exactMatch = await ConversationMessage.query() - .where({ buildUuid, role: 'assistant', timestamp: messageTimestamp }) - .orderBy('id', 'desc') - .first(); - - if (exactMatch) { - return exactMatch; - } - - // Use a bounded window around the client timestamp to reduce mismatches - // for long-running or stale sessions. - const nearbyCandidates = await ConversationMessage.query() - .where({ buildUuid, role: 'assistant' }) - .andWhere('timestamp', '>=', messageTimestamp - FeedbackService.TIMESTAMP_FALLBACK_BACKWARD_WINDOW_MS) - .andWhere('timestamp', '<=', messageTimestamp + FeedbackService.TIMESTAMP_FALLBACK_FORWARD_WINDOW_MS) - .orderBy('timestamp', 'asc') - .orderBy('id', 'desc'); - - if (nearbyCandidates.length === 1) { - return nearbyCandidates[0]; - } - - if (nearbyCandidates.length > 1) { - const ranked = nearbyCandidates - .map((candidate) => ({ - candidate, - delta: Math.abs(Number(candidate.timestamp) - messageTimestamp), - })) - .sort((a, b) => { - if (a.delta !== b.delta) return a.delta - b.delta; - if (Number(a.candidate.timestamp) !== Number(b.candidate.timestamp)) { - return Number(a.candidate.timestamp) - Number(b.candidate.timestamp); - } - return b.candidate.id - a.candidate.id; - }); - - // Do not guess when two candidates are equally close. - if (ranked[1] && ranked[0].delta === ranked[1].delta) { - return undefined; - } - - return ranked[0].candidate; - } - - // Conservative fallback: if the conversation has exactly one assistant - // message total, use it. - const assistantMessages = await ConversationMessage.query() - .where({ buildUuid, role: 'assistant' }) - .orderBy('timestamp', 'desc') - .limit(2); - - if (assistantMessages.length === 1) { - return assistantMessages[0]; - } - - return undefined; - } - - async submitMessageFeedback(params: MessageFeedbackParams): Promise { - const { buildUuid, messageId, messageTimestamp, rating, text, userIdentifier, repo, prNumber } = params; - const logger = getLogger(); - - try { - await this.persistenceService.persistConversation(buildUuid, repo); - } catch (err) { - logger.warn(`AI: feedback persistence trigger failed buildUuid=${buildUuid} error=${err.message}`); - } - - let message = messageId ? await ConversationMessage.query().findOne({ id: messageId, buildUuid }) : undefined; - - if (!message && messageTimestamp != null) { - message = await this.resolveAssistantMessageByTimestamp(buildUuid, messageTimestamp); - if (message && Number(message.timestamp) !== messageTimestamp) { - logger.info( - `AI: feedback timestamp fallback matched buildUuid=${buildUuid} requestedTimestamp=${messageTimestamp} resolvedTimestamp=${message.timestamp} resolvedMessageId=${message.id}` - ); - } - } - - if (!message) { - throw new Error( - `Message not found: messageId=${messageId || 'n/a'} messageTimestamp=${ - messageTimestamp || 'n/a' - } buildUuid=${buildUuid}` - ); - } - - const patch: Record = { - rating, - repo, - prNumber: prNumber || null, - }; - if (text !== undefined) { - patch.text = text || null; - } - if (userIdentifier !== undefined) { - patch.userIdentifier = userIdentifier || null; - } - - const existingRecord = await MessageFeedback.query() - .where({ buildUuid, messageId: message.id }) - .orderBy('id', 'desc') - .first(); - - if (existingRecord) { - const record = await MessageFeedback.query().patchAndFetchById(existingRecord.id, patch as any); - logger.info(`AI: feedback updated type=message buildUuid=${buildUuid} messageId=${message.id} rating=${rating}`); - return record; - } - - try { - const record = await MessageFeedback.query().insertAndFetch({ - buildUuid, - messageId: message.id, - rating, - text: text || null, - userIdentifier: userIdentifier || null, - repo, - prNumber: prNumber || null, - } as any); - - logger.info( - `AI: feedback submitted type=message buildUuid=${buildUuid} messageId=${message.id} rating=${rating}` - ); - return record; - } catch (error) { - if (!this.isUniqueViolation(error)) { - throw error; - } - - const conflictRecord = await MessageFeedback.query() - .where({ buildUuid, messageId: message.id }) - .orderBy('id', 'desc') - .first(); - - if (!conflictRecord) { - throw error; - } - - const record = await MessageFeedback.query().patchAndFetchById(conflictRecord.id, patch as any); - logger.info(`AI: feedback updated type=message buildUuid=${buildUuid} messageId=${message.id} rating=${rating}`); - return record; - } - } - - async submitConversationFeedback(params: ConversationFeedbackParams): Promise { - const { buildUuid, rating, text, userIdentifier, repo, prNumber } = params; - const logger = getLogger(); - - try { - await this.persistenceService.persistConversation(buildUuid, repo); - } catch (err) { - logger.warn(`AI: feedback persistence trigger failed buildUuid=${buildUuid} error=${err.message}`); - } - - const conversation = await Conversation.query().findById(buildUuid); - if (!conversation) { - throw new Error(`Conversation not found: buildUuid=${buildUuid}`); - } - - const patch: Record = { - rating, - repo, - prNumber: prNumber || null, - }; - if (text !== undefined) { - patch.text = text || null; - } - if (userIdentifier !== undefined) { - patch.userIdentifier = userIdentifier || null; - } - - const existingRecord = await ConversationFeedback.query().where({ buildUuid }).orderBy('id', 'desc').first(); - - if (existingRecord) { - const record = await ConversationFeedback.query().patchAndFetchById(existingRecord.id, patch as any); - logger.info(`AI: feedback updated type=conversation buildUuid=${buildUuid} rating=${rating}`); - return record; - } - - try { - const record = await ConversationFeedback.query().insertAndFetch({ - buildUuid, - rating, - text: text || null, - userIdentifier: userIdentifier || null, - repo, - prNumber: prNumber || null, - } as any); - - logger.info(`AI: feedback submitted type=conversation buildUuid=${buildUuid} rating=${rating}`); - return record; - } catch (error) { - if (!this.isUniqueViolation(error)) { - throw error; - } - - const conflictRecord = await ConversationFeedback.query().where({ buildUuid }).orderBy('id', 'desc').first(); - - if (!conflictRecord) { - throw error; - } - - const record = await ConversationFeedback.query().patchAndFetchById(conflictRecord.id, patch as any); - logger.info(`AI: feedback updated type=conversation buildUuid=${buildUuid} rating=${rating}`); - return record; - } - } -} diff --git a/src/server/services/ai/feedback/__tests__/FeedbackService.test.ts b/src/server/services/ai/feedback/__tests__/FeedbackService.test.ts deleted file mode 100644 index 6dc52b75..00000000 --- a/src/server/services/ai/feedback/__tests__/FeedbackService.test.ts +++ /dev/null @@ -1,288 +0,0 @@ -/** - * Copyright 2025 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. - */ - -const mockLogger = { info: jest.fn(), warn: jest.fn(), error: jest.fn() }; -jest.mock('server/lib/logger', () => ({ - getLogger: () => mockLogger, -})); - -jest.mock('server/models/Conversation', () => { - const model: any = {}; - model.query = jest.fn(); - return { __esModule: true, default: model }; -}); - -jest.mock('server/models/ConversationMessage', () => { - const model: any = {}; - model.query = jest.fn(); - return { __esModule: true, default: model }; -}); - -jest.mock('server/models/MessageFeedback', () => { - const model: any = {}; - model.query = jest.fn(); - return { __esModule: true, default: model }; -}); - -jest.mock('server/models/ConversationFeedback', () => { - const model: any = {}; - model.query = jest.fn(); - return { __esModule: true, default: model }; -}); - -import FeedbackService from '../FeedbackService'; -import Conversation from 'server/models/Conversation'; -import ConversationFeedback from 'server/models/ConversationFeedback'; -import ConversationMessage from 'server/models/ConversationMessage'; -import MessageFeedback from 'server/models/MessageFeedback'; - -const MockConversation = Conversation as any; -const MockConversationFeedback = ConversationFeedback as any; -const MockConversationMessage = ConversationMessage as any; -const MockMessageFeedback = MessageFeedback as any; - -describe('FeedbackService', () => { - const persistenceService = { persistConversation: jest.fn() }; - let service: FeedbackService; - - beforeEach(() => { - jest.clearAllMocks(); - persistenceService.persistConversation.mockResolvedValue(true); - service = new FeedbackService(persistenceService as any); - }); - - it('submits message feedback using direct persisted messageId', async () => { - const findOne = jest.fn().mockResolvedValue({ id: 11, buildUuid: 'uuid-1' }); - MockConversationMessage.query.mockReturnValue({ findOne }); - - const existingFirst = jest.fn().mockResolvedValue(undefined); - const existingOrderBy = jest.fn().mockReturnValue({ first: existingFirst }); - const existingWhere = jest.fn().mockReturnValue({ orderBy: existingOrderBy }); - const insertAndFetch = jest.fn().mockResolvedValue({ id: 1, messageId: 11 }); - MockMessageFeedback.query.mockReturnValueOnce({ where: existingWhere }).mockReturnValueOnce({ insertAndFetch }); - - await service.submitMessageFeedback({ - buildUuid: 'uuid-1', - messageId: 11, - rating: 'up', - repo: 'org/repo', - }); - - expect(findOne).toHaveBeenCalledWith({ id: 11, buildUuid: 'uuid-1' }); - expect(existingWhere).toHaveBeenCalledWith({ buildUuid: 'uuid-1', messageId: 11 }); - expect(insertAndFetch).toHaveBeenCalledWith( - expect.objectContaining({ messageId: 11, buildUuid: 'uuid-1', rating: 'up' }) - ); - }); - - it('submits message feedback by resolving assistant message from timestamp', async () => { - const first = jest.fn().mockResolvedValue({ id: 27, buildUuid: 'uuid-2', timestamp: 1234 }); - const orderBy = jest.fn().mockReturnValue({ first }); - const where = jest.fn().mockReturnValue({ orderBy }); - MockConversationMessage.query.mockReturnValue({ where }); - - const existingFirst = jest.fn().mockResolvedValue(undefined); - const existingOrderBy = jest.fn().mockReturnValue({ first: existingFirst }); - const existingWhere = jest.fn().mockReturnValue({ orderBy: existingOrderBy }); - const insertAndFetch = jest.fn().mockResolvedValue({ id: 2, messageId: 27 }); - MockMessageFeedback.query.mockReturnValueOnce({ where: existingWhere }).mockReturnValueOnce({ insertAndFetch }); - - await service.submitMessageFeedback({ - buildUuid: 'uuid-2', - messageTimestamp: 1234, - rating: 'down', - repo: 'org/repo', - text: 'not helpful', - }); - - expect(where).toHaveBeenCalledWith({ buildUuid: 'uuid-2', role: 'assistant', timestamp: 1234 }); - expect(existingWhere).toHaveBeenCalledWith({ buildUuid: 'uuid-2', messageId: 27 }); - expect(insertAndFetch).toHaveBeenCalledWith( - expect.objectContaining({ messageId: 27, rating: 'down', text: 'not helpful' }) - ); - }); - - it('submits message feedback by falling back to nearest assistant message when exact timestamp misses', async () => { - const exactFirst = jest.fn().mockResolvedValue(undefined); - const exactOrderBy = jest.fn().mockReturnValue({ first: exactFirst }); - const exactWhere = jest.fn().mockReturnValue({ orderBy: exactOrderBy }); - - const nearbySecondOrderBy = jest.fn().mockResolvedValue([ - { id: 35, buildUuid: 'uuid-2', timestamp: 1300 }, - { id: 36, buildUuid: 'uuid-2', timestamp: 1500 }, - ]); - const nearbyFirstOrderBy = jest.fn().mockReturnValue({ orderBy: nearbySecondOrderBy }); - const nearbyUpperBoundWhere = jest.fn().mockReturnValue({ orderBy: nearbyFirstOrderBy }); - const nearbyLowerBoundWhere = jest.fn().mockReturnValue({ andWhere: nearbyUpperBoundWhere }); - const nearbyWhere = jest.fn().mockReturnValue({ andWhere: nearbyLowerBoundWhere }); - - MockConversationMessage.query - .mockReturnValueOnce({ where: exactWhere }) - .mockReturnValueOnce({ where: nearbyWhere }); - - const existingFirst = jest.fn().mockResolvedValue(undefined); - const existingOrderBy = jest.fn().mockReturnValue({ first: existingFirst }); - const existingWhere = jest.fn().mockReturnValue({ orderBy: existingOrderBy }); - const insertAndFetch = jest.fn().mockResolvedValue({ id: 3, messageId: 35 }); - MockMessageFeedback.query.mockReturnValueOnce({ where: existingWhere }).mockReturnValueOnce({ insertAndFetch }); - - await service.submitMessageFeedback({ - buildUuid: 'uuid-2', - messageTimestamp: 1234, - rating: 'down', - repo: 'org/repo', - }); - - expect(exactWhere).toHaveBeenCalledWith({ buildUuid: 'uuid-2', role: 'assistant', timestamp: 1234 }); - expect(nearbyWhere).toHaveBeenCalledWith({ buildUuid: 'uuid-2', role: 'assistant' }); - expect(nearbyLowerBoundWhere).toHaveBeenCalledWith('timestamp', '>=', 1234 - 60_000); - expect(nearbyUpperBoundWhere).toHaveBeenCalledWith('timestamp', '<=', 1234 + 600_000); - expect(existingWhere).toHaveBeenCalledWith({ buildUuid: 'uuid-2', messageId: 35 }); - expect(insertAndFetch).toHaveBeenCalledWith(expect.objectContaining({ messageId: 35, rating: 'down' })); - expect(mockLogger.info).toHaveBeenCalledWith(expect.stringContaining('AI: feedback timestamp fallback matched')); - }); - - it('falls back to the only assistant message when bounded timestamp lookup misses', async () => { - const exactFirst = jest.fn().mockResolvedValue(undefined); - const exactOrderBy = jest.fn().mockReturnValue({ first: exactFirst }); - const exactWhere = jest.fn().mockReturnValue({ orderBy: exactOrderBy }); - - const nearbySecondOrderBy = jest.fn().mockResolvedValue([]); - const nearbyFirstOrderBy = jest.fn().mockReturnValue({ orderBy: nearbySecondOrderBy }); - const nearbyUpperBoundWhere = jest.fn().mockReturnValue({ orderBy: nearbyFirstOrderBy }); - const nearbyLowerBoundWhere = jest.fn().mockReturnValue({ andWhere: nearbyUpperBoundWhere }); - const nearbyWhere = jest.fn().mockReturnValue({ andWhere: nearbyLowerBoundWhere }); - - const fallbackLimit = jest.fn().mockResolvedValue([{ id: 77, buildUuid: 'uuid-2', timestamp: 9000 }]); - const fallbackOrderBy = jest.fn().mockReturnValue({ limit: fallbackLimit }); - const fallbackWhere = jest.fn().mockReturnValue({ orderBy: fallbackOrderBy }); - - MockConversationMessage.query - .mockReturnValueOnce({ where: exactWhere }) - .mockReturnValueOnce({ where: nearbyWhere }) - .mockReturnValueOnce({ where: fallbackWhere }); - - const existingFirst = jest.fn().mockResolvedValue(undefined); - const existingOrderBy = jest.fn().mockReturnValue({ first: existingFirst }); - const existingWhere = jest.fn().mockReturnValue({ orderBy: existingOrderBy }); - const insertAndFetch = jest.fn().mockResolvedValue({ id: 4, messageId: 77 }); - MockMessageFeedback.query.mockReturnValueOnce({ where: existingWhere }).mockReturnValueOnce({ insertAndFetch }); - - await service.submitMessageFeedback({ - buildUuid: 'uuid-2', - messageTimestamp: 1234, - rating: 'up', - repo: 'org/repo', - }); - - expect(fallbackWhere).toHaveBeenCalledWith({ buildUuid: 'uuid-2', role: 'assistant' }); - expect(fallbackLimit).toHaveBeenCalledWith(2); - expect(existingWhere).toHaveBeenCalledWith({ buildUuid: 'uuid-2', messageId: 77 }); - }); - - it('overwrites existing feedback for the same message instead of inserting a new row', async () => { - const findOne = jest.fn().mockResolvedValue({ id: 44, buildUuid: 'uuid-4' }); - MockConversationMessage.query.mockReturnValue({ findOne }); - - const existingFirst = jest.fn().mockResolvedValue({ id: 9, buildUuid: 'uuid-4', messageId: 44, rating: 'up' }); - const existingOrderBy = jest.fn().mockReturnValue({ first: existingFirst }); - const existingWhere = jest.fn().mockReturnValue({ orderBy: existingOrderBy }); - const patchAndFetchById = jest - .fn() - .mockResolvedValue({ id: 9, buildUuid: 'uuid-4', messageId: 44, rating: 'down' }); - - MockMessageFeedback.query.mockReturnValueOnce({ where: existingWhere }).mockReturnValueOnce({ patchAndFetchById }); - - await service.submitMessageFeedback({ - buildUuid: 'uuid-4', - messageId: 44, - rating: 'down', - repo: 'org/repo', - }); - - expect(existingWhere).toHaveBeenCalledWith({ buildUuid: 'uuid-4', messageId: 44 }); - expect(patchAndFetchById).toHaveBeenCalledWith(9, expect.objectContaining({ rating: 'down', repo: 'org/repo' })); - expect(mockLogger.info).toHaveBeenCalledWith(expect.stringContaining('AI: feedback updated type=message')); - }); - - it('throws when no message can be resolved', async () => { - const findOne = jest.fn().mockResolvedValue(undefined); - const exactFirst = jest.fn().mockResolvedValue(undefined); - const exactOrderBy = jest.fn().mockReturnValue({ first: exactFirst }); - const exactWhere = jest.fn().mockReturnValue({ orderBy: exactOrderBy }); - - const nearbySecondOrderBy = jest.fn().mockResolvedValue([]); - const nearbyFirstOrderBy = jest.fn().mockReturnValue({ orderBy: nearbySecondOrderBy }); - const nearbyUpperBoundWhere = jest.fn().mockReturnValue({ orderBy: nearbyFirstOrderBy }); - const nearbyLowerBoundWhere = jest.fn().mockReturnValue({ andWhere: nearbyUpperBoundWhere }); - const nearbyWhere = jest.fn().mockReturnValue({ andWhere: nearbyLowerBoundWhere }); - - const fallbackLimit = jest.fn().mockResolvedValue([ - { id: 81, buildUuid: 'uuid-3', timestamp: 2000 }, - { id: 80, buildUuid: 'uuid-3', timestamp: 1900 }, - ]); - const fallbackOrderBy = jest.fn().mockReturnValue({ limit: fallbackLimit }); - const fallbackWhere = jest.fn().mockReturnValue({ orderBy: fallbackOrderBy }); - - MockConversationMessage.query - .mockReturnValueOnce({ findOne }) - .mockReturnValueOnce({ where: exactWhere }) - .mockReturnValueOnce({ where: nearbyWhere }) - .mockReturnValueOnce({ where: fallbackWhere }); - - await expect( - service.submitMessageFeedback({ - buildUuid: 'uuid-3', - messageId: 99, - messageTimestamp: 1234, - rating: 'up', - repo: 'org/repo', - }) - ).rejects.toThrow('Message not found'); - }); - - it('overwrites existing conversation feedback for the same build', async () => { - const findById = jest.fn().mockResolvedValue({ buildUuid: 'uuid-5' }); - MockConversation.query.mockReturnValue({ findById }); - - const existingFirst = jest.fn().mockResolvedValue({ id: 13, buildUuid: 'uuid-5', rating: 'up' }); - const existingOrderBy = jest.fn().mockReturnValue({ first: existingFirst }); - const existingWhere = jest.fn().mockReturnValue({ orderBy: existingOrderBy }); - const patchAndFetchById = jest.fn().mockResolvedValue({ id: 13, buildUuid: 'uuid-5', rating: 'down' }); - - MockConversationFeedback.query - .mockReturnValueOnce({ where: existingWhere }) - .mockReturnValueOnce({ patchAndFetchById }); - - await service.submitConversationFeedback({ - buildUuid: 'uuid-5', - rating: 'down', - repo: 'org/repo', - text: 'This improved', - }); - - expect(existingWhere).toHaveBeenCalledWith({ buildUuid: 'uuid-5' }); - expect(patchAndFetchById).toHaveBeenCalledWith( - 13, - expect.objectContaining({ - rating: 'down', - repo: 'org/repo', - text: 'This improved', - }) - ); - expect(mockLogger.info).toHaveBeenCalledWith(expect.stringContaining('AI: feedback updated type=conversation')); - }); -}); diff --git a/src/server/services/ai/feedback/__tests__/resolveFeedbackContext.test.ts b/src/server/services/ai/feedback/__tests__/resolveFeedbackContext.test.ts deleted file mode 100644 index 23e5faf5..00000000 --- a/src/server/services/ai/feedback/__tests__/resolveFeedbackContext.test.ts +++ /dev/null @@ -1,73 +0,0 @@ -/** - * Copyright 2025 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/Build', () => { - const model: { query: jest.Mock } = { query: jest.fn() }; - return { __esModule: true, default: model }; -}); - -jest.mock('server/models/Conversation', () => { - const model: { query: jest.Mock } = { query: jest.fn() }; - return { __esModule: true, default: model }; -}); - -import Build from 'server/models/Build'; -import Conversation from 'server/models/Conversation'; -import { resolveFeedbackContext } from '../resolveFeedbackContext'; - -const MockBuild = Build as unknown as { query: jest.Mock }; -const MockConversation = Conversation as unknown as { query: jest.Mock }; - -describe('resolveFeedbackContext', () => { - beforeEach(() => { - jest.clearAllMocks(); - }); - - it('prefers build pull request context when available', async () => { - const findById = jest.fn().mockReturnValue({ select: jest.fn().mockResolvedValue({ repo: 'conversation/repo' }) }); - MockConversation.query.mockReturnValue({ findById }); - - const modifyGraph = jest.fn().mockResolvedValue({ - pullRequest: { - fullName: 'build/repo', - pullRequestNumber: 44, - }, - }); - const withGraphFetched = jest.fn().mockReturnValue({ modifyGraph }); - const findOne = jest.fn().mockReturnValue({ withGraphFetched }); - MockBuild.query.mockReturnValue({ findOne }); - - await expect(resolveFeedbackContext('uuid-1')).resolves.toEqual({ - repo: 'build/repo', - prNumber: 44, - }); - }); - - it('returns conversation repo without prNumber when build pull request data is missing', async () => { - const select = jest.fn().mockResolvedValue({ repo: 'conversation/repo' }); - const findById = jest.fn().mockReturnValue({ select }); - MockConversation.query.mockReturnValue({ findById }); - - const modifyGraph = jest.fn().mockResolvedValue({ pullRequest: null }); - const withGraphFetched = jest.fn().mockReturnValue({ modifyGraph }); - const findOne = jest.fn().mockReturnValue({ withGraphFetched }); - MockBuild.query.mockReturnValue({ findOne }); - - await expect(resolveFeedbackContext('uuid-2')).resolves.toEqual({ - repo: 'conversation/repo', - }); - }); -}); diff --git a/src/server/services/ai/feedback/__tests__/userIdentifier.test.ts b/src/server/services/ai/feedback/__tests__/userIdentifier.test.ts deleted file mode 100644 index e31c750e..00000000 --- a/src/server/services/ai/feedback/__tests__/userIdentifier.test.ts +++ /dev/null @@ -1,54 +0,0 @@ -/** - * Copyright 2025 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 { resolveUserIdentifierFromPayload } from '../userIdentifier'; - -describe('resolveUserIdentifierFromPayload', () => { - it('returns undefined when payload is missing', () => { - expect(resolveUserIdentifierFromPayload(null)).toBeUndefined(); - }); - - it('uses the highest priority claim and trims whitespace', () => { - expect( - resolveUserIdentifierFromPayload({ - preferred_username: ' preferred-user ', - github_username: ' github-user ', - }) - ).toBe('github-user'); - }); - - it('clamps long identifiers to 255 characters', () => { - const longIdentifier = 'a'.repeat(300); - const resolved = resolveUserIdentifierFromPayload({ - preferred_username: longIdentifier, - }); - - expect(resolved).toBeDefined(); - expect(resolved!.length).toBe(255); - expect(resolved).toBe('a'.repeat(255)); - }); - - it('clamps Unicode identifiers without splitting surrogate pairs', () => { - const longUnicodeIdentifier = '😀'.repeat(300); - const resolved = resolveUserIdentifierFromPayload({ - preferred_username: longUnicodeIdentifier, - }); - - expect(resolved).toBeDefined(); - expect(Array.from(resolved!).length).toBe(255); - expect(/[\uD800-\uDBFF]$/.test(resolved!)).toBe(false); - }); -}); diff --git a/src/server/services/ai/feedback/constants.ts b/src/server/services/ai/feedback/constants.ts deleted file mode 100644 index 4946a924..00000000 --- a/src/server/services/ai/feedback/constants.ts +++ /dev/null @@ -1,17 +0,0 @@ -/** - * Copyright 2025 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. - */ - -export const MAX_FEEDBACK_TEXT_LENGTH = 10_000; diff --git a/src/server/services/ai/feedback/resolveFeedbackContext.ts b/src/server/services/ai/feedback/resolveFeedbackContext.ts deleted file mode 100644 index bd72fc70..00000000 --- a/src/server/services/ai/feedback/resolveFeedbackContext.ts +++ /dev/null @@ -1,53 +0,0 @@ -/** - * Copyright 2025 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 Build from 'server/models/Build'; -import Conversation from 'server/models/Conversation'; - -export interface FeedbackContext { - repo: string; - prNumber?: number; -} - -export async function resolveFeedbackContext(buildUuid: string): Promise { - const existingConversation = await Conversation.query().findById(buildUuid).select('repo'); - const build = await Build.query() - .findOne({ uuid: buildUuid }) - .withGraphFetched('pullRequest') - .modifyGraph('pullRequest', (builder) => { - builder.select('id', 'fullName', 'pullRequestNumber'); - }); - - const buildRepo = build?.pullRequest?.fullName?.trim(); - const buildPrNumber = build?.pullRequest?.pullRequestNumber; - - if (buildRepo) { - return { - repo: buildRepo, - ...(buildPrNumber != null ? { prNumber: buildPrNumber } : {}), - }; - } - - if (existingConversation?.repo) { - return { - repo: existingConversation.repo, - }; - } - - return { - repo: '', - }; -} diff --git a/src/server/services/ai/feedback/userIdentifier.ts b/src/server/services/ai/feedback/userIdentifier.ts deleted file mode 100644 index 20d0f4a3..00000000 --- a/src/server/services/ai/feedback/userIdentifier.ts +++ /dev/null @@ -1,64 +0,0 @@ -/** - * Copyright 2025 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 { JWTPayload } from 'jose'; - -const MAX_USER_IDENTIFIER_LENGTH = 255; - -const CANDIDATE_KEYS = [ - 'github_username', - 'githubUsername', - 'preferred_username', - 'preferredUsername', - 'email', - 'upn', - 'name', - 'sub', -] as const; - -function clampIdentifier(value: string): string { - const graphemes = Array.from(value); - if (graphemes.length <= MAX_USER_IDENTIFIER_LENGTH) { - return value; - } - return graphemes.slice(0, MAX_USER_IDENTIFIER_LENGTH).join(''); -} - -function normalizeClaim(value: unknown): string | undefined { - if (typeof value !== 'string') { - return undefined; - } - const normalized = value.trim(); - if (!normalized) { - return undefined; - } - return clampIdentifier(normalized); -} - -export function resolveUserIdentifierFromPayload(payload: JWTPayload | null): string | undefined { - if (!payload) { - return undefined; - } - - for (const key of CANDIDATE_KEYS) { - const candidate = normalizeClaim((payload as Record)[key]); - if (candidate) { - return candidate; - } - } - - return undefined; -} diff --git a/src/server/services/ai/mcp/__tests__/toolAdapter.test.ts b/src/server/services/ai/mcp/__tests__/toolAdapter.test.ts deleted file mode 100644 index f5fdb0e7..00000000 --- a/src/server/services/ai/mcp/__tests__/toolAdapter.test.ts +++ /dev/null @@ -1,182 +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. - */ - -const mockConnect = jest.fn().mockResolvedValue(undefined); -const mockCallTool = jest.fn(); -const mockClose = jest.fn().mockResolvedValue(undefined); - -jest.mock('../client', () => ({ - McpClientManager: jest.fn().mockImplementation(() => ({ - connect: mockConnect, - callTool: mockCallTool, - close: mockClose, - })), -})); - -jest.mock('server/lib/logger', () => ({ - getLogger: () => ({ info: jest.fn(), warn: jest.fn(), error: jest.fn() }), -})); - -import { MCPToolAdapter, createMcpTools } from '../toolAdapter'; -import { ToolSafetyLevel } from '../../types/tool'; -import { MCP_ERROR_CODES, type ResolvedMcpServer } from '../types'; - -describe('MCPToolAdapter', () => { - beforeEach(() => { - jest.clearAllMocks(); - }); - - describe('constructor', () => { - it('produces a namespaced tool name', () => { - const adapter = new MCPToolAdapter( - 'sample-server', - { name: 'inspectItem', inputSchema: { type: 'object' } }, - { type: 'http', url: 'https://mcp.example.com/v1/mcp' }, - {}, - 30000 - ); - - expect(adapter.name).toBe('mcp__sample-server__inspectItem'); - }); - - it('maps destructive annotations to DANGEROUS', () => { - const adapter = new MCPToolAdapter( - 'sample-server', - { name: 'removeItem', inputSchema: { type: 'object' }, annotations: { destructiveHint: true } }, - { type: 'http', url: 'https://mcp.example.com/v1/mcp' }, - {}, - 30000 - ); - - expect(adapter.safetyLevel).toBe(ToolSafetyLevel.DANGEROUS); - expect(adapter.description).toMatch(/^\[DANGEROUS\]/); - }); - - it('maps read-only annotations to SAFE', () => { - const adapter = new MCPToolAdapter( - 'sample-server', - { name: 'readItem', inputSchema: { type: 'object' }, annotations: { readOnlyHint: true } }, - { type: 'http', url: 'https://mcp.example.com/v1/mcp' }, - {}, - 30000 - ); - - expect(adapter.safetyLevel).toBe(ToolSafetyLevel.SAFE); - }); - }); - - describe('execute', () => { - it('returns text content on success', async () => { - mockCallTool.mockResolvedValue({ - content: [{ type: 'text', text: 'result data' }], - isError: false, - }); - - const adapter = new MCPToolAdapter( - 'sample-server', - { name: 'inspectItem', inputSchema: {} }, - { type: 'http', url: 'https://mcp.example.com/v1/mcp' }, - {}, - 30000 - ); - - const result = await adapter.execute({}); - - expect(result.success).toBe(true); - expect(result.agentContent).toBe('result data'); - }); - - it('returns a connection error when connect fails', async () => { - mockConnect.mockRejectedValueOnce(new Error('ECONNREFUSED')); - - const adapter = new MCPToolAdapter( - 'sample-server', - { name: 'inspectItem', inputSchema: {} }, - { type: 'http', url: 'https://mcp.example.com/v1/mcp' }, - {}, - 30000 - ); - - const result = await adapter.execute({}); - - expect(result.success).toBe(false); - expect(result.error?.code).toBe(MCP_ERROR_CODES.CONNECTION); - }); - - it('applies default args before calling the tool', async () => { - mockCallTool.mockResolvedValue({ - content: [{ type: 'text', text: 'ok' }], - isError: false, - }); - - const adapter = new MCPToolAdapter( - 'sample-server', - { - name: 'inspectItem', - inputSchema: { - type: 'object', - properties: { - siteUrl: { type: 'string' }, - }, - }, - }, - { type: 'http', url: 'https://mcp.example.com/v1/mcp' }, - { siteUrl: 'https://sample-site.example.com' }, - 30000 - ); - - await adapter.execute({}); - - expect(mockCallTool).toHaveBeenCalledWith( - 'inspectItem', - { siteUrl: 'https://sample-site.example.com' }, - 30000, - undefined - ); - }); - }); - - describe('createMcpTools', () => { - it('creates tools from resolved servers', () => { - const servers: ResolvedMcpServer[] = [ - { - slug: 'sample-server-a', - name: 'Sample Server A', - transport: { type: 'http', url: 'https://mcp-a.example.com/v1/mcp' }, - timeout: 30000, - defaultArgs: {}, - env: {}, - discoveredTools: [{ name: 'inspectItem', inputSchema: {} }], - }, - { - slug: 'sample-server-b', - name: 'Sample Server B', - transport: { type: 'sse', url: 'https://mcp-b.example.com/sse' }, - timeout: 30000, - defaultArgs: {}, - env: {}, - discoveredTools: [{ name: 'createItem', inputSchema: {} }], - }, - ]; - - const tools = createMcpTools(servers); - - expect(tools).toHaveLength(2); - expect(tools[0].name).toBe('mcp__sample-server-a__inspectItem'); - expect(tools[1].name).toBe('mcp__sample-server-b__createItem'); - }); - }); -}); diff --git a/src/server/services/ai/mcp/toolAdapter.ts b/src/server/services/ai/mcp/toolAdapter.ts deleted file mode 100644 index e814b985..00000000 --- a/src/server/services/ai/mcp/toolAdapter.ts +++ /dev/null @@ -1,176 +0,0 @@ -/** - * Copyright 2025 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 { Tool, ToolResult, ToolSafetyLevel, ToolCategory, JSONSchema } from '../types/tool'; -import { McpToolAnnotations, ResolvedMcpServer, MCP_ERROR_CODES, McpDiscoveredTool } from './types'; -import { McpClientManager } from './client'; -import { applyMcpDefaultToolArgs } from './runtimeConfig'; -import { getLogger } from 'server/lib/logger'; -import { OutputLimiter } from '../tools/outputLimiter'; - -function mapAnnotationsToSafetyLevel(annotations?: McpToolAnnotations): ToolSafetyLevel { - if (!annotations) return ToolSafetyLevel.CAUTIOUS; - if (annotations.destructiveHint === true) return ToolSafetyLevel.DANGEROUS; - if (annotations.readOnlyHint === true) return ToolSafetyLevel.SAFE; - if (annotations.openWorldHint === true) return ToolSafetyLevel.CAUTIOUS; - return ToolSafetyLevel.CAUTIOUS; -} - -function prefixDescription(description: string, level: ToolSafetyLevel): string { - if (level === ToolSafetyLevel.DANGEROUS) return `[DANGEROUS] ${description}`; - if (level === ToolSafetyLevel.SAFE) return `[SAFE] ${description}`; - return description; -} - -function extractTextContent(content: unknown): string { - if (!Array.isArray(content)) { - return typeof content === 'string' ? content : JSON.stringify(content); - } - return content - .filter((block: any) => block?.type === 'text' && typeof block.text === 'string') - .map((block: any) => block.text) - .join('\n'); -} - -function classifyMcpError(error: unknown): { code: string; message: string } { - const msg = error instanceof Error ? error.message : String(error); - const lower = msg.toLowerCase(); - - if ( - lower.includes('econnrefused') || - lower.includes('timeout') || - lower.includes('abort') || - lower.includes('connection failed') || - lower.includes('fetch failed') - ) { - return { code: MCP_ERROR_CODES.CONNECTION, message: msg }; - } - - if (lower.includes('protocol') || lower.includes('json-rpc') || lower.includes('jsonrpc')) { - return { code: MCP_ERROR_CODES.PROTOCOL, message: msg }; - } - - return { code: MCP_ERROR_CODES.TOOL, message: msg }; -} - -export class MCPToolAdapter implements Tool { - name: string; - description: string; - parameters: JSONSchema; - safetyLevel: ToolSafetyLevel; - category: ToolCategory; - - private originalName: string; - private serverTransport: ResolvedMcpServer['transport']; - private serverDefaultArgs: Record; - private serverTimeout: number; - - constructor( - serverId: string, - discoveredTool: McpDiscoveredTool, - serverTransport: ResolvedMcpServer['transport'], - serverDefaultArgs: Record, - serverTimeout: number - ) { - this.name = `mcp__${serverId}__${discoveredTool.name}`; - this.safetyLevel = mapAnnotationsToSafetyLevel(discoveredTool.annotations); - this.description = prefixDescription( - discoveredTool.description || `MCP tool ${discoveredTool.name} from ${serverId}`, - this.safetyLevel - ); - this.parameters = discoveredTool.inputSchema as JSONSchema; - this.category = 'mcp'; - this.originalName = discoveredTool.name; - this.serverTransport = serverTransport; - this.serverDefaultArgs = serverDefaultArgs; - this.serverTimeout = serverTimeout; - } - - async execute(args: Record, signal?: AbortSignal): Promise { - const resolvedArgs = applyMcpDefaultToolArgs( - this.parameters as Record, - this.serverDefaultArgs, - args - ); - const client = new McpClientManager(); - try { - try { - await client.connect(this.serverTransport, this.serverTimeout); - } catch (error) { - getLogger().warn( - `MCP connection failed for tool ${this.name}: ${error instanceof Error ? error.message : String(error)}` - ); - return { - success: false, - error: { - message: `Failed to connect to MCP server: ${error instanceof Error ? error.message : String(error)}`, - code: MCP_ERROR_CODES.CONNECTION, - recoverable: true, - suggestedAction: 'MCP server may be temporarily unavailable. Try again or skip this tool.', - }, - }; - } - - const result = await client.callTool(this.originalName, resolvedArgs, this.serverTimeout, signal); - const textContent = extractTextContent(result.content); - - if (result.isError) { - return { - success: false, - agentContent: textContent, - error: { - message: `MCP tool returned error: ${textContent}`, - code: MCP_ERROR_CODES.TOOL, - recoverable: true, - suggestedAction: 'Check tool arguments and try again.', - }, - }; - } - - return { - success: true, - agentContent: OutputLimiter.truncate(textContent), - }; - } catch (error) { - const classified = classifyMcpError(error); - getLogger().warn(`MCP tool error for ${this.name}: code=${classified.code} ${classified.message}`); - return { - success: false, - error: { - message: classified.message, - code: classified.code, - recoverable: classified.code !== MCP_ERROR_CODES.PROTOCOL, - suggestedAction: - classified.code === MCP_ERROR_CODES.PROTOCOL - ? 'MCP server may have a compatibility issue.' - : 'Check tool arguments and try again.', - }, - }; - } finally { - await client.close(); - } - } -} - -export function createMcpTools(servers: ResolvedMcpServer[]): Tool[] { - const tools: Tool[] = []; - for (const server of servers) { - for (const discoveredTool of server.discoveredTools) { - tools.push(new MCPToolAdapter(server.slug, discoveredTool, server.transport, server.defaultArgs, server.timeout)); - } - } - return tools; -} diff --git a/src/server/services/ai/orchestration/__tests__/loopProtection.test.ts b/src/server/services/ai/orchestration/__tests__/loopProtection.test.ts deleted file mode 100644 index e1f8e9ad..00000000 --- a/src/server/services/ai/orchestration/__tests__/loopProtection.test.ts +++ /dev/null @@ -1,110 +0,0 @@ -/** - * Copyright 2025 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 { LoopDetector } from '../loopProtection'; - -describe('LoopDetector', () => { - it('uses default values', () => { - const d = new LoopDetector(); - const p = d.getProtection(); - expect(p.maxIterations).toBe(20); - expect(p.maxToolCalls).toBe(50); - expect(p.maxRepeatedCalls).toBe(1); - }); - - it('accepts custom options', () => { - const d = new LoopDetector({ maxIterations: 10, maxToolCalls: 25, maxRepeatedCalls: 5 }); - const p = d.getProtection(); - expect(p.maxIterations).toBe(10); - expect(p.maxToolCalls).toBe(25); - expect(p.maxRepeatedCalls).toBe(5); - }); - - it('recordCall adds to history', () => { - const d = new LoopDetector(); - d.recordCall('get_file', { path: 'a.ts' }, 1); - d.recordCall('get_file', { path: 'b.ts' }, 2); - expect(d.getProtection().toolCallHistory).toHaveLength(2); - }); - - it('countRepeatedCalls returns 0 for no history', () => { - const d = new LoopDetector(); - expect(d.countRepeatedCalls('get_file', { path: 'a.ts' }, 1)).toBe(0); - }); - - it('counts exact matches (same tool + same args)', () => { - const d = new LoopDetector(); - d.recordCall('get_file', { path: 'a.ts' }, 1); - d.recordCall('get_file', { path: 'a.ts' }, 2); - d.recordCall('get_file', { path: 'a.ts' }, 3); - expect(d.countRepeatedCalls('get_file', { path: 'a.ts' }, 3)).toBe(3); - }); - - it('ignores calls older than 5 iterations', () => { - const d = new LoopDetector(); - d.recordCall('get_file', { path: 'a.ts' }, 1); - expect(d.countRepeatedCalls('get_file', { path: 'a.ts' }, 7)).toBe(0); - }); - - it('includes calls within 5-iteration window', () => { - const d = new LoopDetector(); - d.recordCall('get_file', { path: 'a.ts' }, 2); - expect(d.countRepeatedCalls('get_file', { path: 'a.ts' }, 7)).toBe(1); - }); - - it('ignores different tool names', () => { - const d = new LoopDetector(); - d.recordCall('get_pod_logs', { path: 'a.ts' }, 1); - expect(d.countRepeatedCalls('get_file', { path: 'a.ts' }, 1)).toBe(0); - }); - - it('ignores different args', () => { - const d = new LoopDetector(); - d.recordCall('get_file', { path: 'b.ts' }, 1); - expect(d.countRepeatedCalls('get_file', { path: 'a.ts' }, 1)).toBe(0); - }); - - it('returns hint about searching resources for get_k8s_resources without name', () => { - const d = new LoopDetector(); - const hint = d.getLoopHint('get_k8s_resources', { namespace: 'ns' }); - expect(hint).toContain('searching for resources'); - }); - - it('returns default hint for get_k8s_resources with name', () => { - const d = new LoopDetector(); - const hint = d.getLoopHint('get_k8s_resources', { name: 'pod1' }); - expect(hint).toBe('Consider trying a different tool or different arguments.'); - }); - - it('returns hint about fetching logs for get_pod_logs', () => { - const d = new LoopDetector(); - const hint = d.getLoopHint('get_pod_logs', {}); - expect(hint).toContain('fetching logs'); - }); - - it('returns default hint for other tools', () => { - const d = new LoopDetector(); - const hint = d.getLoopHint('other_tool', {}); - expect(hint).toBe('Consider trying a different tool or different arguments.'); - }); - - it('reset clears toolCallHistory', () => { - const d = new LoopDetector(); - d.recordCall('get_file', { path: 'a.ts' }, 1); - d.reset(); - expect(d.getProtection().toolCallHistory).toHaveLength(0); - }); -}); diff --git a/src/server/services/ai/orchestration/__tests__/observationMasker.test.ts b/src/server/services/ai/orchestration/__tests__/observationMasker.test.ts deleted file mode 100644 index 36c43448..00000000 --- a/src/server/services/ai/orchestration/__tests__/observationMasker.test.ts +++ /dev/null @@ -1,325 +0,0 @@ -/** - * Copyright 2025 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('../../prompts/tokenCounter', () => ({ - countTokens: (text: string) => text.length, -})); - -import { ConversationMessage } from '../../types/message'; -import { ToolResult } from '../../types/tool'; -import { maskObservations } from '../observationMasker'; - -function buildToolResultMessage(toolName: string, agentContent: string, success: boolean = true): ConversationMessage { - const result: ToolResult = { success, agentContent }; - return { - role: 'user', - parts: [ - { - type: 'tool_result', - toolCallId: `call_${toolName}_${Math.random().toString(36).slice(2, 8)}`, - name: toolName, - result, - }, - ], - }; -} - -function buildToolCallMessage(toolName: string, args: Record): ConversationMessage { - return { - role: 'assistant', - parts: [ - { - type: 'tool_call', - toolCallId: `call_${toolName}_${Math.random().toString(36).slice(2, 8)}`, - name: toolName, - arguments: args, - }, - ], - }; -} - -function buildTextMessage(role: 'user' | 'assistant', content: string): ConversationMessage { - return { role, parts: [{ type: 'text', content }] }; -} - -describe('maskObservations', () => { - it('should not mask when under token threshold', () => { - const messages: ConversationMessage[] = [ - buildTextMessage('user', 'hello'), - buildToolCallMessage('get_k8s_resources', { kind: 'pods' }), - buildToolResultMessage('get_k8s_resources', 'pod-1 Running'), - ]; - - const result = maskObservations(messages, { tokenThreshold: 999999 }); - - expect(result.masked).toBe(false); - expect(result.messages).toBe(messages); - }); - - it('should mask old tool results with placeholder', () => { - const messages: ConversationMessage[] = [ - buildTextMessage('user', 'check pods'), - buildToolCallMessage('get_k8s_resources', { kind: 'pods' }), - buildToolResultMessage('get_k8s_resources', 'pod-1 Running\npod-2 Running\npod-3 CrashLoopBackOff'), - buildTextMessage('assistant', 'I see pod-3 is crashing'), - buildTextMessage('user', 'check logs'), - buildToolCallMessage('get_pod_logs', { pod: 'pod-3' }), - buildToolResultMessage('get_pod_logs', 'Error: OOMKilled'), - ]; - - const result = maskObservations(messages, { tokenThreshold: 1, recencyWindow: 1 }); - - expect(result.masked).toBe(true); - const oldToolResult = result.messages[2]; - expect(oldToolResult.parts[0].type).toBe('tool_result'); - const oldPart = oldToolResult.parts[0] as any; - expect(oldPart.result.agentContent).toBe('[get_k8s_resources output omitted — re-call tool if needed]'); - - const recentToolResult = result.messages[6]; - expect(recentToolResult.parts[0].type).toBe('tool_result'); - const recentPart = recentToolResult.parts[0] as any; - expect(recentPart.result.agentContent).toBe('Error: OOMKilled'); - }); - - it('should protect recent turn-pairs from masking', () => { - const messages: ConversationMessage[] = []; - for (let i = 0; i < 12; i++) { - messages.push(buildToolCallMessage(`tool_${i}`, { i })); - messages.push(buildToolResultMessage(`tool_${i}`, `result-content-for-tool-${i}`)); - } - - const result = maskObservations(messages, { tokenThreshold: 1, recencyWindow: 2 }); - - expect(result.masked).toBe(true); - - const lastToolResultMsg = result.messages[23]; - const lastPart = lastToolResultMsg.parts[0] as any; - expect(lastPart.result.agentContent).toBe('result-content-for-tool-11'); - - const secondLastToolResultMsg = result.messages[21]; - const secondLastPart = secondLastToolResultMsg.parts[0] as any; - expect(secondLastPart.result.agentContent).toBe('result-content-for-tool-10'); - - const oldToolResultMsg = result.messages[1]; - const oldPart = oldToolResultMsg.parts[0] as any; - expect(oldPart.result.agentContent).toBe('[tool_0 output omitted — re-call tool if needed]'); - }); - - it('should never mask error results', () => { - const messages: ConversationMessage[] = [ - buildToolCallMessage('query_database', { sql: 'SELECT 1' }), - buildToolResultMessage('query_database', 'Connection refused: ECONNREFUSED', false), - buildToolCallMessage('get_k8s_resources', { kind: 'pods' }), - buildToolResultMessage('get_k8s_resources', 'pod-1 Running'), - buildToolCallMessage('get_pod_logs', { pod: 'pod-1' }), - buildToolResultMessage('get_pod_logs', 'latest log line'), - ]; - - const result = maskObservations(messages, { tokenThreshold: 1, recencyWindow: 1 }); - - expect(result.masked).toBe(true); - - const errorResult = result.messages[1]; - const errorPart = errorResult.parts[0] as any; - expect(errorPart.result.agentContent).toBe('Connection refused: ECONNREFUSED'); - expect(errorPart.result.success).toBe(false); - - const oldSuccessResult = result.messages[3]; - const oldSuccessPart = oldSuccessResult.parts[0] as any; - expect(oldSuccessPart.result.agentContent).toBe('[get_k8s_resources output omitted — re-call tool if needed]'); - }); - - it('should not mutate the original messages array', () => { - const messages: ConversationMessage[] = [ - buildToolCallMessage('get_k8s_resources', { kind: 'pods' }), - buildToolResultMessage('get_k8s_resources', 'pod-1 Running'), - buildToolCallMessage('get_pod_logs', { pod: 'pod-1' }), - buildToolResultMessage('get_pod_logs', 'latest log line'), - ]; - - const snapshot = JSON.parse(JSON.stringify(messages)); - - maskObservations(messages, { tokenThreshold: 1, recencyWindow: 1 }); - - expect(JSON.stringify(messages)).toBe(JSON.stringify(snapshot)); - }); - - it('should return accurate stats', () => { - const messages: ConversationMessage[] = [ - buildToolCallMessage('get_k8s_resources', { kind: 'pods' }), - buildToolResultMessage( - 'get_k8s_resources', - 'pod-1 Running\npod-2 Running\npod-3 CrashLoopBackOff\npod-4 Pending\npod-5 Running\npod-6 ImagePullBackOff' - ), - buildToolCallMessage('get_pod_logs', { pod: 'pod-1' }), - buildToolResultMessage('get_pod_logs', 'recent log'), - ]; - - const result = maskObservations(messages, { tokenThreshold: 1, recencyWindow: 1 }); - - expect(result.masked).toBe(true); - expect(result.stats.maskedParts).toBe(1); - expect(result.stats.totalTokensBefore).toBeGreaterThan(0); - expect(result.stats.totalTokensAfter).toBeGreaterThan(0); - expect(result.stats.totalTokensAfter).toBeLessThan(result.stats.totalTokensBefore); - expect(result.stats.savedTokens).toBe(result.stats.totalTokensBefore - result.stats.totalTokensAfter); - }); - - it('does not mask when total tokens are just below threshold', () => { - const content = 'x'.repeat(99); - const messages: ConversationMessage[] = [ - buildToolCallMessage('tool_a', { q: 'a' }), - buildToolResultMessage('tool_a', content), - buildToolCallMessage('tool_b', { q: 'b' }), - buildToolResultMessage('tool_b', content), - ]; - - const totalChars = messages.reduce((sum, m) => { - return ( - sum + - m.parts.reduce((ps, p) => { - if (p.type === 'text') return ps + p.content.length; - if (p.type === 'tool_call') return ps + JSON.stringify(p.arguments).length + p.name.length; - if (p.type === 'tool_result') return ps + (p.result.agentContent || JSON.stringify(p.result)).length; - return ps; - }, 0) - ); - }, 0); - - const result = maskObservations(messages, { tokenThreshold: totalChars + 1 }); - expect(result.masked).toBe(false); - expect(result.stats.maskedParts).toBe(0); - }); - - it('masks oldest tool results when over threshold, preserving recent N', () => { - const messages: ConversationMessage[] = []; - for (let i = 0; i < 6; i++) { - messages.push(buildToolCallMessage(`tool_${i}`, { i })); - messages.push(buildToolResultMessage(`tool_${i}`, `result-for-tool-${i}-with-padding`)); - } - - const result = maskObservations(messages, { tokenThreshold: 1, recencyWindow: 3 }); - - expect(result.masked).toBe(true); - - for (let i = 0; i < 3; i++) { - const msg = result.messages[i * 2 + 1]; - const part = msg.parts[0] as any; - expect(part.result.agentContent).toBe(`[tool_${i} output omitted — re-call tool if needed]`); - } - - for (let i = 3; i < 6; i++) { - const msg = result.messages[i * 2 + 1]; - const part = msg.parts[0] as any; - expect(part.result.agentContent).toBe(`result-for-tool-${i}-with-padding`); - } - }); - - it('never masks error results regardless of age', () => { - const messages: ConversationMessage[] = [ - buildToolCallMessage('failing_tool', { x: 1 }), - buildToolResultMessage('failing_tool', 'CRITICAL: connection timeout at db.host:5432', false), - buildToolCallMessage('tool_2', { x: 2 }), - buildToolResultMessage('tool_2', 'success result 2'), - buildToolCallMessage('tool_3', { x: 3 }), - buildToolResultMessage('tool_3', 'success result 3'), - buildToolCallMessage('tool_4', { x: 4 }), - buildToolResultMessage('tool_4', 'success result 4'), - buildToolCallMessage('tool_5', { x: 5 }), - buildToolResultMessage('tool_5', 'success result 5'), - ]; - - const result = maskObservations(messages, { tokenThreshold: 1, recencyWindow: 2 }); - - expect(result.masked).toBe(true); - - const errorPart = result.messages[1].parts[0] as any; - expect(errorPart.result.agentContent).toBe('CRITICAL: connection timeout at db.host:5432'); - expect(errorPart.result.success).toBe(false); - - const oldSuccessPart = result.messages[3].parts[0] as any; - expect(oldSuccessPart.result.agentContent).toBe('[tool_2 output omitted — re-call tool if needed]'); - }); - - it('masks at exact threshold boundary (uses strict less-than)', () => { - const content = 'a'.repeat(50); - const messages: ConversationMessage[] = [ - buildToolCallMessage('tool_old', { q: 'old' }), - buildToolResultMessage('tool_old', content), - buildToolCallMessage('tool_new', { q: 'new' }), - buildToolResultMessage('tool_new', content), - ]; - - const totalChars = messages.reduce((sum, m) => { - return ( - sum + - m.parts.reduce((ps, p) => { - if (p.type === 'text') return ps + p.content.length; - if (p.type === 'tool_call') return ps + JSON.stringify(p.arguments).length + p.name.length; - if (p.type === 'tool_result') return ps + (p.result.agentContent || JSON.stringify(p.result)).length; - return ps; - }, 0) - ); - }, 0); - - const atThreshold = maskObservations(messages, { tokenThreshold: totalChars, recencyWindow: 1 }); - expect(atThreshold.masked).toBe(true); - - const belowThreshold = maskObservations(messages, { tokenThreshold: totalChars + 1, recencyWindow: 1 }); - expect(belowThreshold.masked).toBe(false); - }); - - it('should preserve assistant reasoning text while masking tool results', () => { - const messages: ConversationMessage[] = [ - buildTextMessage('user', 'What pods are running?'), - { - role: 'assistant', - parts: [ - { type: 'text', content: 'Let me check the pods for you.' }, - { - type: 'tool_call', - toolCallId: 'call_1', - name: 'get_k8s_resources', - arguments: { kind: 'pods' }, - }, - ], - }, - buildToolResultMessage('get_k8s_resources', 'pod-1 Running\npod-2 CrashLoop'), - buildTextMessage('assistant', 'I found 2 pods. Let me check the logs.'), - buildToolCallMessage('get_pod_logs', { pod: 'pod-2' }), - buildToolResultMessage('get_pod_logs', 'OOMKilled at 12:00'), - buildTextMessage('assistant', 'Pod-2 is OOMKilled.'), - buildTextMessage('user', 'fix it'), - ]; - - const result = maskObservations(messages, { tokenThreshold: 1, recencyWindow: 1 }); - - expect(result.masked).toBe(true); - - const assistantWithReasoning = result.messages[1]; - const textPart = assistantWithReasoning.parts[0] as any; - expect(textPart.type).toBe('text'); - expect(textPart.content).toBe('Let me check the pods for you.'); - - const oldToolResult = result.messages[2]; - const toolResultPart = oldToolResult.parts[0] as any; - expect(toolResultPart.result.agentContent).toBe('[get_k8s_resources output omitted — re-call tool if needed]'); - - const preservedReasoning = result.messages[3]; - expect(preservedReasoning.parts[0].type).toBe('text'); - expect((preservedReasoning.parts[0] as any).content).toBe('I found 2 pods. Let me check the logs.'); - }); -}); diff --git a/src/server/services/ai/orchestration/__tests__/orchestrator.test.ts b/src/server/services/ai/orchestration/__tests__/orchestrator.test.ts deleted file mode 100644 index 4f08b03f..00000000 --- a/src/server/services/ai/orchestration/__tests__/orchestrator.test.ts +++ /dev/null @@ -1,445 +0,0 @@ -/** - * Copyright 2025 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/logger', () => ({ - getLogger: () => ({ - info: jest.fn(), - warn: jest.fn(), - error: jest.fn(), - }), -})); - -const mockExecute = jest.fn(); -const mockCreateProviderPolicy = jest.fn(() => ({ execute: mockExecute })); - -jest.mock('../../resilience', () => ({ - createProviderPolicy: (...args: unknown[]) => mockCreateProviderPolicy(...args), -})); - -jest.mock('../../errors', () => ({ - RetryBudget: jest.requireActual('../../errors/retryBudget').RetryBudget, - ErrorCategory: jest.requireActual('../../errors/classification').ErrorCategory, - createClassifiedError: jest.requireActual('../../errors/providerErrors').createClassifiedError, -})); - -import { ToolOrchestrator } from '../orchestrator'; -import { LLMProvider, StreamChunk, CompletionOptions } from '../../types/provider'; -import { ConversationMessage } from '../../types/message'; -import { Tool, ToolSafetyLevel } from '../../types/tool'; -import { StreamCallbacks } from '../../types/stream'; -import { ToolRegistry } from '../../tools/registry'; -import { ToolSafetyManager } from '../safety'; -import { RetryBudget } from '../../errors/retryBudget'; - -function createMockCallbacks(): StreamCallbacks { - return { - onTextChunk: jest.fn(), - onThinking: jest.fn(), - onToolCall: jest.fn(), - onToolResult: jest.fn(), - onError: jest.fn(), - onActivity: jest.fn(), - }; -} - -function createMockProvider(chunkSets: StreamChunk[][]): LLMProvider { - let callIndex = 0; - return { - name: 'test-provider', - streamCompletion(_messages: ConversationMessage[], _options: CompletionOptions, _signal?: AbortSignal) { - const chunks = chunkSets[callIndex] || []; - callIndex++; - let i = 0; - return { - next() { - if (i < chunks.length) { - const chunk = chunks[i]; - i++; - if ((chunk as any)._throw) { - return Promise.reject(new Error((chunk as any)._throw)); - } - return Promise.resolve({ value: chunk, done: false }); - } - return Promise.resolve({ value: undefined as any, done: true }); - }, - [Symbol.asyncIterator]() { - return this; - }, - } as any; - }, - supportsTools: () => true, - getModelInfo: () => ({ model: 'test', maxTokens: 4096 }), - formatToolDefinition: (tool: Tool) => tool, - parseToolCall: () => [], - estimateTokens: () => 0, - formatHistory: (messages: ConversationMessage[]) => messages.map((m) => ({ role: m.role, content: '' })), - }; -} - -function createErrorProvider(error: Error): LLMProvider { - return { - name: 'test-provider', - streamCompletion() { - return { - next() { - return Promise.reject(error); - }, - [Symbol.asyncIterator]() { - return this; - }, - } as any; - }, - supportsTools: () => true, - getModelInfo: () => ({ model: 'test', maxTokens: 4096 }), - formatToolDefinition: (tool: Tool) => tool, - parseToolCall: () => [], - estimateTokens: () => 0, - formatHistory: (messages: ConversationMessage[]) => messages.map((m) => ({ role: m.role, content: '' })), - }; -} - -function createPartialErrorProvider(textChunks: StreamChunk[], error: Error): LLMProvider { - return { - name: 'test-provider', - streamCompletion() { - let i = 0; - return { - next() { - if (i < textChunks.length) { - const chunk = textChunks[i]; - i++; - return Promise.resolve({ value: chunk, done: false }); - } - return Promise.reject(error); - }, - [Symbol.asyncIterator]() { - return this; - }, - } as any; - }, - supportsTools: () => true, - getModelInfo: () => ({ model: 'test', maxTokens: 4096 }), - formatToolDefinition: (tool: Tool) => tool, - parseToolCall: () => [], - estimateTokens: () => 0, - formatHistory: (messages: ConversationMessage[]) => messages.map((m) => ({ role: m.role, content: '' })), - }; -} - -function createMockToolRegistry(): ToolRegistry { - const registry = { - get: jest.fn().mockReturnValue({ - name: 'test_tool', - description: 'A test tool', - parameters: { type: 'object', properties: {} }, - safetyLevel: ToolSafetyLevel.SAFE, - category: 'k8s' as const, - execute: jest.fn().mockResolvedValue({ success: true, agentContent: 'result' }), - }), - getAll: jest.fn().mockReturnValue([]), - register: jest.fn(), - registerMultiple: jest.fn(), - unregister: jest.fn(), - getByCategory: jest.fn().mockReturnValue([]), - getFiltered: jest.fn().mockReturnValue([]), - execute: jest.fn(), - } as unknown as ToolRegistry; - return registry; -} - -function createMockSafetyManager(): ToolSafetyManager { - return { - safeExecute: jest.fn().mockResolvedValue({ success: true, agentContent: 'result' }), - } as unknown as ToolSafetyManager; -} - -beforeEach(() => { - mockExecute.mockReset(); - mockCreateProviderPolicy.mockReset(); - mockCreateProviderPolicy.mockReturnValue({ execute: mockExecute }); - mockExecute.mockImplementation(async (fn: Function) => fn()); -}); - -describe('ToolOrchestrator', () => { - let orchestrator: ToolOrchestrator; - let registry: ToolRegistry; - let safetyManager: ToolSafetyManager; - - beforeEach(() => { - registry = createMockToolRegistry(); - safetyManager = createMockSafetyManager(); - orchestrator = new ToolOrchestrator(registry, safetyManager); - }); - - it('returns successful response when provider streams text without errors', async () => { - const provider = createMockProvider([ - [ - { type: 'text', content: 'Hello ' }, - { type: 'text', content: 'world' }, - ], - ]); - const callbacks = createMockCallbacks(); - const controller = new AbortController(); - - const result = await orchestrator.executeToolLoop(provider, 'system prompt', [], [], callbacks, controller.signal); - - expect(result.success).toBe(true); - expect(result.response).toBe('Hello world'); - expect(callbacks.onTextChunk).toHaveBeenCalledWith('Hello '); - expect(callbacks.onTextChunk).toHaveBeenCalledWith('world'); - }); - - it('preserves partial results when stream error occurs after text accumulation', async () => { - const partialChunks: StreamChunk[] = [ - { type: 'text', content: 'Partial ' }, - { type: 'text', content: 'response' }, - ]; - const provider = createPartialErrorProvider(partialChunks, new Error('stream died')); - const callbacks = createMockCallbacks(); - const controller = new AbortController(); - - const result = await orchestrator.executeToolLoop(provider, 'system prompt', [], [], callbacks, controller.signal); - - expect(result.success).toBe(true); - expect(result.response).toContain('Partial response'); - expect(result.error).toContain('Stream interrupted'); - }); - - it('returns failure when stream error occurs with no accumulated content', async () => { - const provider = createErrorProvider(new Error('total failure')); - const callbacks = createMockCallbacks(); - const controller = new AbortController(); - - const result = await orchestrator.executeToolLoop(provider, 'system prompt', [], [], callbacks, controller.signal); - - expect(result.success).toBe(false); - expect(result.error).toBe('total failure'); - }); - - it('passes RetryBudget to createProviderPolicy', async () => { - const provider = createMockProvider([[{ type: 'text', content: 'ok' }]]); - const callbacks = createMockCallbacks(); - const controller = new AbortController(); - - await orchestrator.executeToolLoop(provider, 'system prompt', [], [], callbacks, controller.signal); - - expect(mockCreateProviderPolicy).toHaveBeenCalledWith('test-provider', expect.any(RetryBudget)); - }); - - it('wraps stream consumption in policy.execute', async () => { - const provider = createMockProvider([[{ type: 'text', content: 'ok' }]]); - const callbacks = createMockCallbacks(); - const controller = new AbortController(); - - await orchestrator.executeToolLoop(provider, 'system prompt', [], [], callbacks, controller.signal); - - expect(mockExecute).toHaveBeenCalled(); - expect(typeof mockExecute.mock.calls[0][0]).toBe('function'); - }); - - it('handles tool calls correctly after successful stream', async () => { - const provider = createMockProvider([ - [ - { - type: 'tool_call', - toolCalls: [{ name: 'test_tool', arguments: { query: 'pods' } }], - }, - ], - [{ type: 'text', content: 'Done with tools' }], - ]); - const callbacks = createMockCallbacks(); - const controller = new AbortController(); - - const result = await orchestrator.executeToolLoop(provider, 'system prompt', [], [], callbacks, controller.signal); - - expect(result.success).toBe(true); - expect(result.response).toBe('Done with tools'); - expect(safetyManager.safeExecute).toHaveBeenCalled(); - expect(callbacks.onToolCall).toHaveBeenCalledWith('test_tool', { query: 'pods' }, expect.any(String)); - }); - - it('respects abort signal', async () => { - const provider = createMockProvider([[{ type: 'text', content: 'should not appear' }]]); - const callbacks = createMockCallbacks(); - const controller = new AbortController(); - controller.abort(); - - const result = await orchestrator.executeToolLoop(provider, 'system prompt', [], [], callbacks, controller.signal); - - expect(result.success).toBe(false); - expect(result.cancelled).toBe(true); - expect(result.error).toContain('cancelled'); - }); - - describe('parallel tool execution', () => { - it('executes multiple tool calls in parallel', async () => { - const delayedSafeExecute = jest - .fn() - .mockImplementation( - () => new Promise((resolve) => setTimeout(() => resolve({ success: true, agentContent: 'result' }), 100)) - ); - (safetyManager.safeExecute as jest.Mock) = delayedSafeExecute; - safetyManager.safeExecute = delayedSafeExecute; - - const provider = createMockProvider([ - [ - { - type: 'tool_call', - toolCalls: [ - { name: 'test_tool', arguments: { q: '1' } }, - { name: 'test_tool', arguments: { q: '2' } }, - { name: 'test_tool', arguments: { q: '3' } }, - ], - }, - ], - [{ type: 'text', content: 'Done' }], - ]); - const callbacks = createMockCallbacks(); - const controller = new AbortController(); - - const before = Date.now(); - const result = await orchestrator.executeToolLoop( - provider, - 'system prompt', - [], - [], - callbacks, - controller.signal - ); - const elapsed = Date.now() - before; - - expect(result.success).toBe(true); - expect(delayedSafeExecute).toHaveBeenCalledTimes(3); - expect(elapsed).toBeLessThan(200); - expect(callbacks.onToolResult).toHaveBeenCalledTimes(3); - }); - - it('handles partial failure in parallel execution', async () => { - let callCount = 0; - const mixedSafeExecute = jest.fn().mockImplementation(() => { - callCount++; - if (callCount === 2) { - return Promise.resolve({ - success: false, - error: { message: 'tool 2 failed', code: 'FAIL', recoverable: true }, - }); - } - return Promise.resolve({ success: true, agentContent: 'ok' }); - }); - safetyManager.safeExecute = mixedSafeExecute; - - const provider = createMockProvider([ - [ - { - type: 'tool_call', - toolCalls: [ - { name: 'test_tool', arguments: { q: '1' } }, - { name: 'test_tool', arguments: { q: '2' } }, - { name: 'test_tool', arguments: { q: '3' } }, - ], - }, - ], - [{ type: 'text', content: 'Done' }], - ]); - const callbacks = createMockCallbacks(); - const controller = new AbortController(); - - const result = await orchestrator.executeToolLoop( - provider, - 'system prompt', - [], - [], - callbacks, - controller.signal - ); - - expect(result.success).toBe(true); - expect(callbacks.onToolResult).toHaveBeenCalledTimes(3); - expect(callbacks.onToolCall).toHaveBeenCalledTimes(3); - - const secondResult = (callbacks.onToolResult as jest.Mock).mock.calls[1][0]; - expect(secondResult.success).toBe(false); - }); - - it('preserves abort behavior during parallel execution', async () => { - const controller = new AbortController(); - const abortSafeExecute = jest.fn().mockImplementation(() => { - controller.abort(); - return Promise.resolve({ success: true, agentContent: 'first done' }); - }); - safetyManager.safeExecute = abortSafeExecute; - - const provider = createMockProvider([ - [ - { - type: 'tool_call', - toolCalls: [ - { name: 'test_tool', arguments: { q: '1' } }, - { name: 'test_tool', arguments: { q: '2' } }, - ], - }, - ], - [{ type: 'text', content: 'Done' }], - ]); - const callbacks = createMockCallbacks(); - - await orchestrator.executeToolLoop(provider, 'system prompt', [], [], callbacks, controller.signal); - - expect(callbacks.onToolResult).toHaveBeenCalled(); - }); - - it('assigns llmThinkTime to first tool by index', async () => { - let callIdx = 0; - const orderedSafeExecute = jest.fn().mockImplementation(() => { - callIdx++; - const delay = callIdx === 1 ? 80 : 10; - return new Promise((resolve) => - setTimeout(() => resolve({ success: true, agentContent: `result-${callIdx}` }), delay) - ); - }); - safetyManager.safeExecute = orderedSafeExecute; - - const provider = createMockProvider([ - [ - { - type: 'tool_call', - toolCalls: [ - { name: 'test_tool', arguments: { q: 'slow' } }, - { name: 'test_tool', arguments: { q: 'fast' } }, - ], - }, - ], - [{ type: 'text', content: 'Done' }], - ]); - const callbacks = createMockCallbacks(); - const controller = new AbortController(); - - await orchestrator.executeToolLoop(provider, 'system prompt', [], [], callbacks, controller.signal); - - const firstToolCall = (callbacks.onToolResult as jest.Mock).mock.calls.find( - (call: any[]) => call[1] === 'test_tool' && call[2].q === 'slow' - ); - const secondToolCall = (callbacks.onToolResult as jest.Mock).mock.calls.find( - (call: any[]) => call[1] === 'test_tool' && call[2].q === 'fast' - ); - - expect(firstToolCall).toBeDefined(); - expect(secondToolCall).toBeDefined(); - expect(firstToolCall![4]).toBeGreaterThanOrEqual(firstToolCall![3]); - expect(secondToolCall![4]).toEqual(secondToolCall![3]); - }); - }); -}); diff --git a/src/server/services/ai/orchestration/__tests__/safety.test.ts b/src/server/services/ai/orchestration/__tests__/safety.test.ts deleted file mode 100644 index 0ab752b1..00000000 --- a/src/server/services/ai/orchestration/__tests__/safety.test.ts +++ /dev/null @@ -1,193 +0,0 @@ -/** - * Copyright 2025 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. - */ - -const mockLogger = { - info: jest.fn(), - warn: jest.fn(), - error: jest.fn(), - debug: jest.fn(), -}; - -jest.mock('server/lib/logger', () => ({ - getLogger: () => mockLogger, -})); - -import { ToolSafetyManager } from '../safety'; -import { Tool, ToolSafetyLevel, ToolCategory, ConfirmationDetails } from '../../types/tool'; -import { StreamCallbacks } from '../../types/stream'; - -function makeTool(overrides: Partial = {}): Tool { - return { - name: 'test_tool', - description: 'test', - parameters: { type: 'object' }, - safetyLevel: ToolSafetyLevel.SAFE, - category: 'k8s' as ToolCategory, - execute: jest.fn().mockResolvedValue({ success: true }), - ...overrides, - }; -} - -const confirmDetails: ConfirmationDetails = { - title: 'Confirm', - description: 'Are you sure?', - impact: 'high', - confirmButtonText: 'Yes', -}; - -function makeCallbacks(overrides: Partial = {}): StreamCallbacks { - return { - onToolConfirmation: jest.fn().mockResolvedValue(true), - ...overrides, - } as StreamCallbacks; -} - -describe('ToolSafetyManager', () => { - let manager: ToolSafetyManager; - - beforeEach(() => { - manager = new ToolSafetyManager(true); - jest.clearAllMocks(); - }); - - describe('argument validation', () => { - it('returns INVALID_ARGUMENTS when args do not match schema', async () => { - const tool = makeTool({ - parameters: { type: 'object', properties: { name: { type: 'string' } }, required: ['name'] }, - }); - const result = await manager.safeExecute(tool, {}, makeCallbacks()); - expect(result.success).toBe(false); - expect(result.error?.code).toBe('INVALID_ARGUMENTS'); - }); - - it('passes validation when args match schema', async () => { - const tool = makeTool({ - parameters: { type: 'object', properties: { name: { type: 'string' } }, required: ['name'] }, - }); - const result = await manager.safeExecute(tool, { name: 'hello' }, makeCallbacks()); - expect(result.success).toBe(true); - }); - }); - - describe('confirmation gating', () => { - it('SAFE tool executes without confirmation', async () => { - const callbacks = makeCallbacks(); - const tool = makeTool({ safetyLevel: ToolSafetyLevel.SAFE }); - await manager.safeExecute(tool, {}, callbacks); - expect(callbacks.onToolConfirmation).not.toHaveBeenCalled(); - }); - - it('DANGEROUS tool triggers confirmation callback', async () => { - const callbacks = makeCallbacks(); - const tool = makeTool({ - safetyLevel: ToolSafetyLevel.DANGEROUS, - shouldConfirmExecution: jest.fn().mockResolvedValue(confirmDetails), - }); - await manager.safeExecute(tool, {}, callbacks); - expect(callbacks.onToolConfirmation).toHaveBeenCalledWith(confirmDetails); - }); - - it('CAUTIOUS tool triggers confirmation callback', async () => { - const callbacks = makeCallbacks(); - const tool = makeTool({ - safetyLevel: ToolSafetyLevel.CAUTIOUS, - shouldConfirmExecution: jest.fn().mockResolvedValue(confirmDetails), - }); - await manager.safeExecute(tool, {}, callbacks); - expect(callbacks.onToolConfirmation).toHaveBeenCalledWith(confirmDetails); - }); - - it('unconfirmed request returns AWAITING_APPROVAL error', async () => { - const callbacks = makeCallbacks({ onToolConfirmation: jest.fn().mockResolvedValue(false) }); - const tool = makeTool({ - safetyLevel: ToolSafetyLevel.DANGEROUS, - shouldConfirmExecution: jest.fn().mockResolvedValue(confirmDetails), - }); - const result = await manager.safeExecute(tool, {}, callbacks); - expect(result.error?.code).toBe('AWAITING_APPROVAL'); - }); - - it('missing onToolConfirmation returns NO_CONFIRMATION_HANDLER', async () => { - const tool = makeTool({ - safetyLevel: ToolSafetyLevel.DANGEROUS, - shouldConfirmExecution: jest.fn().mockResolvedValue(confirmDetails), - }); - const result = await manager.safeExecute(tool, {}, {} as StreamCallbacks); - expect(result.error?.code).toBe('NO_CONFIRMATION_HANDLER'); - }); - - it('requireConfirmation=false still requires confirmation for DANGEROUS tool', async () => { - const noConfirmManager = new ToolSafetyManager(false); - const callbacks = makeCallbacks(); - const tool = makeTool({ - safetyLevel: ToolSafetyLevel.DANGEROUS, - shouldConfirmExecution: jest.fn().mockResolvedValue(confirmDetails), - }); - await noConfirmManager.safeExecute(tool, {}, callbacks); - expect(callbacks.onToolConfirmation).toHaveBeenCalled(); - }); - - it('shouldConfirmExecution returning false skips confirmation', async () => { - const callbacks = makeCallbacks(); - const tool = makeTool({ - safetyLevel: ToolSafetyLevel.DANGEROUS, - shouldConfirmExecution: jest.fn().mockResolvedValue(false), - }); - await manager.safeExecute(tool, {}, callbacks); - expect(callbacks.onToolConfirmation).not.toHaveBeenCalled(); - }); - }); - - describe('execution', () => { - it('successful execution returns tool result', async () => { - const tool = makeTool({ execute: jest.fn().mockResolvedValue({ success: true, agentContent: 'done' }) }); - const result = await manager.safeExecute(tool, {}, makeCallbacks()); - expect(result.success).toBe(true); - expect(result.agentContent).toBe('done'); - }); - - it('tool throwing error returns EXECUTION_ERROR with recoverable:true', async () => { - const tool = makeTool({ execute: jest.fn().mockRejectedValue(new Error('boom')) }); - const result = await manager.safeExecute(tool, {}, makeCallbacks()); - expect(result.success).toBe(false); - expect(result.error?.code).toBe('EXECUTION_ERROR'); - expect(result.error?.recoverable).toBe(true); - }); - - it('timeout returns TIMEOUT error', async () => { - const tool = makeTool({ - execute: jest.fn().mockRejectedValue(new Error('Tool execution timeout')), - }); - const result = await manager.safeExecute(tool, {}, makeCallbacks()); - expect(result.success).toBe(false); - expect(result.error?.code).toBe('TIMEOUT'); - expect(result.error?.recoverable).toBe(true); - }); - }); - - describe('logging', () => { - it('non-recoverable error logs via logger.error', async () => { - const tool = makeTool({ - execute: jest.fn().mockResolvedValue({ - success: false, - error: { message: 'fatal', code: 'FATAL', recoverable: false }, - }), - }); - await manager.safeExecute(tool, {}, makeCallbacks()); - expect(mockLogger.error).toHaveBeenCalled(); - }); - }); -}); diff --git a/src/server/services/ai/orchestration/loopProtection.ts b/src/server/services/ai/orchestration/loopProtection.ts deleted file mode 100644 index 20eb1876..00000000 --- a/src/server/services/ai/orchestration/loopProtection.ts +++ /dev/null @@ -1,106 +0,0 @@ -/** - * Copyright 2025 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. - */ - -export interface ToolCallRecord { - tool: string; - args: Record; - iteration: number; - timestamp: number; -} - -export interface LoopProtection { - maxIterations: number; - maxToolCalls: number; - maxRepeatedCalls: number; - toolCallHistory: ToolCallRecord[]; -} - -export class LoopDetector { - private protection: LoopProtection; - - constructor(options?: Partial) { - this.protection = { - maxIterations: options?.maxIterations || 20, - maxToolCalls: options?.maxToolCalls || 50, - maxRepeatedCalls: options?.maxRepeatedCalls || 1, - toolCallHistory: [], - }; - } - - recordCall(toolName: string, args: Record, iteration: number): void { - this.protection.toolCallHistory.push({ - tool: toolName, - args, - iteration, - timestamp: Date.now(), - }); - } - - countRepeatedCalls(toolName: string, args: Record, currentIteration: number): number { - return this.protection.toolCallHistory.filter((record) => { - if (currentIteration - record.iteration > 5) { - return false; - } - - if (record.tool !== toolName) { - return false; - } - - if (JSON.stringify(record.args) === JSON.stringify(args)) { - return true; - } - - if (toolName === 'get_file' && args.file_path && record.args.file_path === args.file_path) { - return true; - } - - return false; - }).length; - } - - getLoopHint(toolName: string, args: Record): string { - if (toolName === 'get_file') { - return ( - `You already read ${args.file_path || 'this file'}. ` + - 'Use the content from the previous result instead of re-fetching.' - ); - } - - if (toolName === 'get_k8s_resources' && !args.name) { - return ( - 'You keep searching for resources with the same criteria. ' + - "If resources don't exist, check deployment status instead." - ); - } - - if (toolName === 'get_pod_logs') { - return ( - "Repeatedly fetching logs suggests the pattern isn't found. " + - 'Try a different search term or check a different service.' - ); - } - - return 'Consider trying a different tool or different arguments.'; - } - - getProtection(): LoopProtection { - return this.protection; - } - - reset(): void { - this.protection.toolCallHistory = []; - } -} diff --git a/src/server/services/ai/orchestration/observationMasker.ts b/src/server/services/ai/orchestration/observationMasker.ts deleted file mode 100644 index 2480b06b..00000000 --- a/src/server/services/ai/orchestration/observationMasker.ts +++ /dev/null @@ -1,164 +0,0 @@ -/** - * Copyright 2025 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 { ConversationMessage, MessagePart, ToolResultPart } from '../types/message'; -import { countTokens } from '../prompts/tokenCounter'; - -export interface MaskingOptions { - recencyWindow: number; - tokenThreshold: number; -} - -export interface MaskingStats { - totalTokensBefore: number; - totalTokensAfter: number; - maskedParts: number; - savedTokens: number; -} - -export interface MaskingResult { - messages: ConversationMessage[]; - masked: boolean; - stats: MaskingStats; -} - -// Validated against eval scenarios and provider limits (anthropic=180K, openai=110K, gemini=900K). -// System prompt uses ~8K tokens. With 10 tool calls averaging ~3K tokens each, conversations -// reach ~38-50K tokens by step 7-8. Threshold of 40000 activates masking before the conversation -// consumes more than ~36% of the smallest context window (openai 110K), leaving headroom for -// continued reasoning. recencyWindow=3 preserves the immediate investigation context (current -// tool result plus the two preceding results that inform the agent's next action). -const DEFAULT_OPTIONS: MaskingOptions = { - recencyWindow: 3, - tokenThreshold: 25000, -}; - -function estimatePartTokens(part: MessagePart): number { - switch (part.type) { - case 'text': - return countTokens(part.content); - case 'tool_call': - return countTokens(JSON.stringify(part.arguments)) + countTokens(part.name); - case 'tool_result': - return countTokens(part.result.agentContent || JSON.stringify(part.result)); - } -} - -function estimateConversationTokens(messages: ConversationMessage[]): number { - let total = 0; - for (const msg of messages) { - for (const part of msg.parts) { - total += estimatePartTokens(part); - } - } - return total; -} - -function hasToolResultPart(msg: ConversationMessage): boolean { - return msg.parts.some((p) => p.type === 'tool_result'); -} - -function findProtectedBoundary(messages: ConversationMessage[], windowSize: number): number { - let toolResultCount = 0; - for (let i = messages.length - 1; i >= 0; i--) { - if (hasToolResultPart(messages[i])) { - toolResultCount++; - if (toolResultCount >= windowSize) { - return i; - } - } - } - return 0; -} - -function buildPlaceholder(part: ToolResultPart): string { - return `[${part.name} output omitted — re-call tool if needed]`; -} - -function maskToolResultsInMessage(msg: ConversationMessage): ConversationMessage { - const newParts = msg.parts.map((part): MessagePart => { - if (part.type === 'tool_result' && part.result.success !== false) { - return { - ...part, - result: { - ...part.result, - agentContent: buildPlaceholder(part), - }, - } as ToolResultPart; - } - return part; - }); - return { ...msg, parts: newParts }; -} - -export function maskObservations(messages: ConversationMessage[], options?: Partial): MaskingResult { - const opts: MaskingOptions = { ...DEFAULT_OPTIONS, ...options }; - - const totalTokensBefore = estimateConversationTokens(messages); - - if (totalTokensBefore < opts.tokenThreshold) { - return { - messages, - masked: false, - stats: { - totalTokensBefore, - totalTokensAfter: totalTokensBefore, - maskedParts: 0, - savedTokens: 0, - }, - }; - } - - const boundary = findProtectedBoundary(messages, opts.recencyWindow); - - let maskedParts = 0; - const newMessages = messages.map((msg, index) => { - if (index >= boundary) { - return msg; - } - - if (!hasToolResultPart(msg)) { - return msg; - } - - const masked = maskToolResultsInMessage(msg); - for (let j = 0; j < msg.parts.length; j++) { - const origPart = msg.parts[j]; - const newPart = masked.parts[j]; - if ( - origPart.type === 'tool_result' && - newPart.type === 'tool_result' && - origPart.result.agentContent !== newPart.result.agentContent - ) { - maskedParts++; - } - } - return masked; - }); - - const totalTokensAfter = estimateConversationTokens(newMessages); - - return { - messages: newMessages, - masked: true, - stats: { - totalTokensBefore, - totalTokensAfter, - maskedParts, - savedTokens: totalTokensBefore - totalTokensAfter, - }, - }; -} diff --git a/src/server/services/ai/orchestration/orchestrator.ts b/src/server/services/ai/orchestration/orchestrator.ts deleted file mode 100644 index 7e0689bf..00000000 --- a/src/server/services/ai/orchestration/orchestrator.ts +++ /dev/null @@ -1,400 +0,0 @@ -/** - * Copyright 2025 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 crypto from 'crypto'; -import { LLMProvider, StreamChunk } from '../types/provider'; -import { Tool, ToolCall } from '../types/tool'; -import { StreamCallbacks } from '../types/stream'; -import { ToolRegistry } from '../tools/registry'; -import { ToolSafetyManager } from './safety'; -import { LoopDetector, LoopProtection } from './loopProtection'; -import { getLogger } from 'server/lib/logger'; -import { RetryBudget, createClassifiedError, ErrorCategory } from '../errors'; -import type { ClassifiedError } from '../errors'; -import { createProviderPolicy } from '../resilience'; -import { isBrokenCircuitError } from 'cockatiel'; -import { ConversationMessage, ToolCallPart, ToolResultPart } from '../types/message'; - -export interface OrchestrationResult { - success: boolean; - response?: string; - error?: string; - cancelled?: boolean; - classifiedError?: ClassifiedError; - metrics: { - iterations: number; - toolCalls: number; - duration: number; - inputTokens: number; - outputTokens: number; - }; -} - -export interface OrchestratorOptions { - loopProtection?: Partial; - retryBudget?: number; -} - -export class ToolOrchestrator { - private loopDetector: LoopDetector; - private retryBudgetMax: number; - - constructor( - private toolRegistry: ToolRegistry, - private safetyManager: ToolSafetyManager, - options?: OrchestratorOptions - ) { - this.loopDetector = new LoopDetector(options?.loopProtection); - this.retryBudgetMax = options?.retryBudget || 10; - } - - async executeToolLoop( - provider: LLMProvider, - systemPrompt: string, - messages: ConversationMessage[], - tools: Tool[], - callbacks: StreamCallbacks, - signal: AbortSignal, - buildUuid?: string - ): Promise { - const startTime = Date.now(); - let iteration = 0; - let totalToolCalls = 0; - let totalInputTokens = 0; - let totalOutputTokens = 0; - let fullResponse = ''; - const protection = this.loopDetector.getProtection(); - - this.loopDetector.reset(); - - const retryBudget = new RetryBudget(this.retryBudgetMax); - const policy = createProviderPolicy(provider.name, retryBudget); - - while (iteration < protection.maxIterations) { - if (signal.aborted) { - return { - success: false, - error: 'Operation cancelled by user', - cancelled: true, - metrics: this.buildMetrics(iteration, totalToolCalls, startTime, totalInputTokens, totalOutputTokens), - }; - } - - iteration++; - - const iterationStartTime = Date.now(); - const chunks: StreamChunk[] = []; - try { - await policy.execute(async () => { - chunks.length = 0; - for await (const chunk of provider.streamCompletion(messages, { systemPrompt, tools, callbacks }, signal)) { - chunks.push(chunk); - - if (chunk.type === 'text' && chunk.content) { - fullResponse += chunk.content; - callbacks.onTextChunk(chunk.content); - } - } - }); - } catch (error: any) { - if (isBrokenCircuitError(error)) { - getLogger().warn( - `AI: circuit breaker rejected request provider=${provider.name} buildUuid=${buildUuid || 'none'}` - ); - const classified: ClassifiedError = { - category: ErrorCategory.TRANSIENT, - original: error instanceof Error ? error : new Error(String(error)), - retryable: true, - providerName: provider.name, - retryAfter: null, - }; - return { - success: false, - error: 'Provider circuit breaker is open', - classifiedError: classified, - metrics: this.buildMetrics(iteration, totalToolCalls, startTime, totalInputTokens, totalOutputTokens), - }; - } - - const hasPartialResults = fullResponse.length > 0 || chunks.length > 0; - - if (hasPartialResults) { - getLogger().warn( - `AI: stream error with partial results, preserving partialTextLen=${fullResponse.length} chunkCount=${ - chunks.length - } buildUuid=${buildUuid || 'none'} error=${error.message}` - ); - return { - success: true, - response: fullResponse || 'The response was interrupted. Here is what was generated before the error.', - error: `Stream interrupted: ${error.message}`, - classifiedError: createClassifiedError(provider.name, error), - metrics: this.buildMetrics(iteration, totalToolCalls, startTime, totalInputTokens, totalOutputTokens), - }; - } - - getLogger().error( - `AI: stream error buildUuid=${buildUuid || 'none'} error=${error.message} budgetUsed=${retryBudget.used}` - ); - return { - success: false, - error: error.message || 'Provider error', - classifiedError: createClassifiedError(provider.name, error), - metrics: this.buildMetrics(iteration, totalToolCalls, startTime, totalInputTokens, totalOutputTokens), - }; - } - - const usageChunk = chunks.find((c) => c.usage); - if (usageChunk?.usage) { - totalInputTokens += usageChunk.usage.inputTokens; - totalOutputTokens += usageChunk.usage.outputTokens; - } - - const llmThinkTime = Date.now() - iterationStartTime; - const toolCalls = this.extractToolCalls(chunks); - - if (toolCalls.length === 0) { - return { - success: true, - response: fullResponse, - metrics: this.buildMetrics(iteration, totalToolCalls, startTime, totalInputTokens, totalOutputTokens), - }; - } - - totalToolCalls += toolCalls.length; - if (totalToolCalls > protection.maxToolCalls) { - getLogger().warn( - `AI: tool call limit exceeded totalToolCalls=${totalToolCalls} maxToolCalls=${ - protection.maxToolCalls - } buildUuid=${buildUuid || 'none'}` - ); - return { - success: false, - error: - `Tool call limit exceeded (${protection.maxToolCalls}). ` + - `The investigation is too broad. Try asking about specific services.`, - metrics: this.buildMetrics(iteration, totalToolCalls, startTime, totalInputTokens, totalOutputTokens), - }; - } - - const toolCallParts: ToolCallPart[] = toolCalls.map((tc) => ({ - type: 'tool_call' as const, - toolCallId: tc.id || crypto.randomBytes(16).toString('hex'), - name: tc.name, - arguments: tc.arguments, - ...(tc.metadata ? { metadata: tc.metadata } : {}), - })); - - messages.push({ - role: 'assistant', - parts: toolCallParts, - }); - - const toolResultParts: ToolResultPart[] = new Array(toolCalls.length); - - const loopDetectedIndices = new Set(); - for (let i = 0; i < toolCalls.length; i++) { - const toolCall = toolCalls[i]; - const repeatCount = this.loopDetector.countRepeatedCalls(toolCall.name, toolCall.arguments, iteration); - - if (repeatCount >= protection.maxRepeatedCalls) { - const loopError = { - success: false, - error: { - message: - `This tool has been called ${repeatCount} times with the same arguments. ` + - `This suggests a loop. Please try a different approach.`, - code: 'LOOP_DETECTED', - recoverable: false, - suggestedAction: this.loopDetector.getLoopHint(toolCall.name, toolCall.arguments), - }, - }; - - toolResultParts[i] = { - type: 'tool_result' as const, - toolCallId: toolCallParts[i].toolCallId, - name: toolCall.name, - result: loopError, - }; - - const totalDuration = i === 0 ? llmThinkTime : 0; - callbacks.onToolResult( - loopError, - toolCall.name, - toolCall.arguments, - 0, - totalDuration, - toolCallParts[i].toolCallId - ); - - console.warn(`[Loop Detection] ${toolCall.name} called ${repeatCount} times`, { - args: toolCall.arguments, - iteration, - }); - - loopDetectedIndices.add(i); - } - } - - const executeIndices = toolCalls.map((_, i) => i).filter((i) => !loopDetectedIndices.has(i)); - - for (const i of executeIndices) { - this.loopDetector.recordCall(toolCalls[i].name, toolCalls[i].arguments, iteration); - callbacks.onToolCall(toolCalls[i].name, toolCalls[i].arguments, toolCallParts[i].toolCallId); - } - - getLogger().info( - `AI: executing tools=[${executeIndices - .map((i) => toolCalls[i].name) - .join(',')}] iteration=${iteration} buildUuid=${buildUuid || 'none'}` - ); - - const settled = await Promise.allSettled( - executeIndices.map(async (i) => { - if (signal.aborted) { - return { - index: i, - result: { - success: false, - error: { message: 'Operation cancelled during tool execution', code: 'CANCELLED', recoverable: false }, - }, - toolDuration: 0, - }; - } - const toolStartTime = Date.now(); - const tool = this.toolRegistry.get(toolCalls[i].name); - if (!tool) { - const registeredNames = this.toolRegistry.getAll().map((t) => t.name); - return { - index: i, - result: { - success: false, - error: { - message: `Tool "${toolCalls[i].name}" is not available. Available tools: ${registeredNames.join( - ', ' - )}`, - code: 'TOOL_NOT_FOUND', - recoverable: true, - }, - }, - toolDuration: 0, - }; - } - const result = await this.safetyManager.safeExecute( - tool, - toolCalls[i].arguments, - callbacks, - signal, - buildUuid - ); - const toolDuration = Date.now() - toolStartTime; - return { index: i, result, toolDuration }; - }) - ); - - for (const entry of settled) { - if (entry.status === 'fulfilled') { - const { index, result, toolDuration } = entry.value; - const totalDuration = index === 0 && !loopDetectedIndices.has(0) ? llmThinkTime + toolDuration : toolDuration; - toolResultParts[index] = { - type: 'tool_result' as const, - toolCallId: toolCallParts[index].toolCallId, - name: toolCalls[index].name, - result, - }; - callbacks.onToolResult( - result, - toolCalls[index].name, - toolCalls[index].arguments, - toolDuration, - totalDuration, - toolCallParts[index].toolCallId - ); - if (!result.success) { - getLogger().warn( - `AI: tool failed tool=${toolCalls[index].name} error=${result.error?.message} code=${ - result.error?.code - } duration=${toolDuration}ms buildUuid=${buildUuid || 'none'}` - ); - } else { - getLogger().info( - `AI: tool completed tool=${toolCalls[index].name} success=true duration=${toolDuration}ms buildUuid=${ - buildUuid || 'none' - }` - ); - } - } else { - const idx = executeIndices[settled.indexOf(entry)]; - const errorResult = { - success: false, - error: { message: entry.reason?.message || 'Unknown error', code: 'EXECUTION_ERROR', recoverable: true }, - }; - toolResultParts[idx] = { - type: 'tool_result' as const, - toolCallId: toolCallParts[idx].toolCallId, - name: toolCalls[idx].name, - result: errorResult, - }; - callbacks.onToolResult( - errorResult, - toolCalls[idx].name, - toolCalls[idx].arguments, - 0, - idx === 0 ? llmThinkTime : 0, - toolCallParts[idx].toolCallId - ); - getLogger().error( - `AI: tool crashed tool=${toolCalls[idx].name} error=${entry.reason?.message} buildUuid=${ - buildUuid || 'none' - }` - ); - } - } - - messages.push({ - role: 'user', - parts: toolResultParts, - }); - } - - getLogger().warn( - `AI: iteration limit reached iteration=${iteration} maxIterations=${ - protection.maxIterations - } totalToolCalls=${totalToolCalls} buildUuid=${buildUuid || 'none'}` - ); - return { - success: false, - error: - `Investigation incomplete - hit iteration limit (${protection.maxIterations}). ` + - `This may indicate the issue is complex or unclear from available data.`, - metrics: this.buildMetrics(iteration, totalToolCalls, startTime, totalInputTokens, totalOutputTokens), - }; - } - - private extractToolCalls(chunks: StreamChunk[]): ToolCall[] { - return chunks.filter((c) => c.type === 'tool_call' && c.toolCalls).flatMap((c) => c.toolCalls || []); - } - - private buildMetrics(iterations: number, toolCalls: number, startTime: number, inputTokens = 0, outputTokens = 0) { - return { - iterations, - toolCalls, - duration: Date.now() - startTime, - inputTokens, - outputTokens, - }; - } -} diff --git a/src/server/services/ai/orchestration/safety.ts b/src/server/services/ai/orchestration/safety.ts deleted file mode 100644 index 4e87a9d3..00000000 --- a/src/server/services/ai/orchestration/safety.ts +++ /dev/null @@ -1,197 +0,0 @@ -/** - * Copyright 2025 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 JsonSchema from 'jsonschema'; -import { Tool, ToolResult, ToolSafetyLevel } from '../types/tool'; -import { StreamCallbacks } from '../types/stream'; -import { getLogger } from 'server/lib/logger'; -import { OutputLimiter } from '../tools/outputLimiter'; - -export class ToolSafetyManager { - private requireConfirmation: boolean; - private validator: JsonSchema.Validator; - private toolExecutionTimeout: number; - private toolOutputMaxChars: number; - - constructor(requireConfirmation: boolean = true, toolExecutionTimeout?: number, toolOutputMaxChars?: number) { - this.requireConfirmation = requireConfirmation; - this.validator = new JsonSchema.Validator(); - this.toolExecutionTimeout = toolExecutionTimeout || 30000; - this.toolOutputMaxChars = toolOutputMaxChars || 30000; - } - - async safeExecute( - tool: Tool, - args: Record, - callbacks: StreamCallbacks, - signal?: AbortSignal, - buildUuid?: string - ): Promise { - const validation = this.validateArgs(tool.parameters, args); - if (!validation.valid) { - getLogger().warn( - `AI: validation failed tool=${tool.name} errors=${validation.errors.join(', ')} buildUuid=${ - buildUuid || 'none' - }` - ); - return { - success: false, - error: { - message: `Invalid arguments: ${validation.errors.join(', ')}`, - code: 'INVALID_ARGUMENTS', - recoverable: true, - }, - }; - } - - if (callbacks.onToolAuthorization) { - const decision = await callbacks.onToolAuthorization( - { - name: tool.name, - description: tool.description, - category: tool.category, - safetyLevel: tool.safetyLevel, - }, - args - ); - if (!decision.allowed) { - return { - success: false, - error: { - message: decision.reason || `Tool ${tool.name} is not authorized for this operation`, - code: 'TOOL_NOT_AUTHORIZED', - recoverable: false, - }, - }; - } - } - - if (this.needsConfirmation(tool)) { - const confirmDetails = await tool.shouldConfirmExecution?.(args); - - if (confirmDetails) { - if (!callbacks.onToolConfirmation) { - getLogger().error(`AI: confirmation callback missing tool=${tool.name} buildUuid=${buildUuid || 'none'}`); - return { - success: false, - error: { - message: `This operation requires user confirmation, but the confirmation system is not available. Please implement onToolConfirmation callback.`, - code: 'NO_CONFIRMATION_HANDLER', - recoverable: false, - }, - }; - } - - const confirmed = await callbacks.onToolConfirmation(confirmDetails); - - if (!confirmed) { - return { - success: false, - error: { - message: `${tool.name} requires user approval. The request has been sent to the user for confirmation. Do not attempt alternative approaches or workarounds — wait for the user to approve or deny this operation.`, - code: 'AWAITING_APPROVAL', - recoverable: false, - }, - }; - } - } - } - - const effectiveTimeout = tool.executionTimeout ?? this.toolExecutionTimeout; - try { - const result = await this.withTimeout(tool.execute(args, signal), effectiveTimeout); - - if (result.success && result.agentContent) { - result.agentContent = OutputLimiter.truncate(result.agentContent, this.toolOutputMaxChars); - } - - this.logToolExecution(tool.name, args, result, buildUuid); - - return result; - } catch (error: any) { - if (error.message === 'Tool execution timeout') { - getLogger().warn( - `AI: tool timeout tool=${tool.name} timeout=${effectiveTimeout}ms buildUuid=${buildUuid || 'none'}` - ); - return { - success: false, - error: { - message: `${tool.name} timed out after ${effectiveTimeout / 1000} seconds`, - code: 'TIMEOUT', - recoverable: true, - suggestedAction: 'The operation took too long. Try narrowing your query.', - }, - }; - } - - getLogger().error( - `AI: tool execution failed tool=${tool.name} error=${error?.message} buildUuid=${buildUuid || 'none'}` - ); - return { - success: false, - error: { - message: error.message || 'Unknown error', - code: error.code || 'EXECUTION_ERROR', - details: error, - recoverable: true, - }, - }; - } - } - - private needsConfirmation(tool: Tool): boolean { - if (tool.safetyLevel === ToolSafetyLevel.DANGEROUS) { - return true; - } - - if (!this.requireConfirmation) { - return false; - } - - if (tool.safetyLevel === ToolSafetyLevel.CAUTIOUS) { - return true; - } - - return false; - } - - private async withTimeout(promise: Promise, timeoutMs: number): Promise { - return Promise.race([ - promise, - new Promise((_, reject) => setTimeout(() => reject(new Error('Tool execution timeout')), timeoutMs)), - ]); - } - - private validateArgs(schema: any, args: Record): { valid: boolean; errors: string[] } { - const validationResult = this.validator.validate(args, schema); - - return { - valid: validationResult.valid, - errors: validationResult.valid ? [] : validationResult.errors.map((e) => e.message || 'Validation error'), - }; - } - - private logToolExecution(name: string, args: Record, result: ToolResult, buildUuid?: string): void { - if (!result.success) { - const level = result.error?.recoverable ? 'warn' : 'error'; - getLogger()[level]( - `AI: tool error tool=${name} error=${result.error?.message} code=${result.error?.code} recoverable=${ - result.error?.recoverable - } buildUuid=${buildUuid || 'none'}` - ); - } - } -} diff --git a/src/server/services/ai/prompts/__tests__/builder.test.ts b/src/server/services/ai/prompts/__tests__/builder.test.ts deleted file mode 100644 index 27aedd9c..00000000 --- a/src/server/services/ai/prompts/__tests__/builder.test.ts +++ /dev/null @@ -1,210 +0,0 @@ -/** - * Copyright 2025 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 { AIAgentPromptBuilder, PromptContext } from '../builder'; -import { DebugContext } from '../../../types/aiAgent'; - -jest.mock('../sectionRegistry', () => ({ - assembleBasePrompt: () => 'base-prompt', - PROMPT_SECTIONS: [{ id: 'safety', content: '# Safety' }], -})); - -jest.mock('../../context/contextSummarizer', () => ({ - summarizeLifecycleYaml: () => ({ parsed: true, text: 'yaml summary' }), -})); - -function makeDeploy(overrides: Record = {}) { - return { - uuid: 'deploy-1', - serviceName: 'svc-a', - status: 'RUNNING', - statusMessage: '', - type: 'service', - dockerImage: 'img:latest', - branch: 'main', - repoName: 'org/repo-a', - buildNumber: 1, - env: {}, - initEnv: {}, - createdAt: new Date(), - updatedAt: new Date(), - ...overrides, - }; -} - -function makeServiceDebug(name: string, status = 'running' as any) { - return { - name, - type: 'service', - status, - deployInfo: makeDeploy({ serviceName: name }), - pods: [], - events: [], - issues: [], - }; -} - -function makeContext(deploys: any[], services: any[] = [], repoFullName = 'org/repo-a'): PromptContext { - const debugContext: DebugContext = { - buildUuid: 'build-123', - namespace: 'ns-test', - gatheredAt: new Date('2026-01-15T10:00:00Z'), - lifecycleContext: { - build: { - uuid: 'build-123', - status: 'RUNNING' as any, - statusMessage: '', - namespace: 'ns-test', - sha: 'abc123', - trackDefaultBranches: false, - capacityType: 'spot', - enabledFeatures: [], - dependencyGraph: {}, - dashboardLinks: {}, - createdAt: new Date(), - updatedAt: new Date(), - }, - pullRequest: { - number: 42, - title: 'Test PR', - username: 'dev', - branch: 'feature', - baseBranch: 'main', - status: 'open' as any, - url: 'https://github.com/org/repo-a/pull/42', - latestCommit: 'abc123', - fullName: repoFullName, - }, - environment: { id: 1, name: 'test', config: {} }, - deploys, - repository: { name: repoFullName, githubRepositoryId: 1, url: '' }, - }, - services, - lifecycleYaml: { path: 'lifecycle.yaml', content: 'version: 2' }, - }; - - return { - provider: 'anthropic', - debugContext, - conversationHistory: [], - userMessage: 'help', - }; -} - -describe('AIAgentPromptBuilder', () => { - let builder: AIAgentPromptBuilder; - - beforeEach(() => { - builder = new AIAgentPromptBuilder(); - }); - - describe('buildEnvironmentContext via build()', () => { - it('includes gathered-at timestamp', () => { - const ctx = makeContext([makeDeploy()], [makeServiceDebug('svc-a')]); - const result = builder.build(ctx); - const userMsg = result.messages[result.messages.length - 1].content; - expect(userMsg).toContain('Environment State (gathered: 2026-01-15T10:00:00.000Z)'); - }); - - it('uses structured header format', () => { - const ctx = makeContext([makeDeploy()], [makeServiceDebug('svc-a')]); - const result = builder.build(ctx); - const userMsg = result.messages[result.messages.length - 1].content; - expect(userMsg).toContain('Build: build-123 | Status: RUNNING | Namespace: ns-test'); - expect(userMsg).toContain('PR: #42 "Test PR" by dev'); - expect(userMsg).toContain('Repo: org/repo-a @ feature (base: main)'); - }); - - it('does not include duplicate K8s section', () => { - const failingSvc = makeServiceDebug('svc-a', 'failed'); - failingSvc.issues = [ - { - severity: 'critical', - category: 'image', - title: 'CrashLoop', - description: '', - suggestedFix: '', - detectedBy: 'rules', - }, - ]; - const ctx = makeContext([makeDeploy({ serviceName: 'svc-a', status: 'DEPLOY_FAILED' })], [failingSvc]); - const result = builder.build(ctx); - const userMsg = result.messages[result.messages.length - 1].content; - expect(userMsg).not.toContain('INITIAL K8S STATE'); - }); - - it('uses markdown header for lifecycle yaml section', () => { - const ctx = makeContext([makeDeploy()], [makeServiceDebug('svc-a')]); - const result = builder.build(ctx); - const userMsg = result.messages[result.messages.length - 1].content; - expect(userMsg).toContain('## Configuration (lifecycle.yaml)'); - expect(userMsg).not.toContain('===== LIFECYCLE.YAML SUMMARY ====='); - }); - - describe('single-repo services', () => { - it('renders flat without repo sub-headers', () => { - const deploys = [ - makeDeploy({ serviceName: 'svc-a', repoName: 'org/repo-a' }), - makeDeploy({ serviceName: 'svc-b', repoName: 'org/repo-a', status: 'BUILD_FAILED' }), - ]; - const services = [makeServiceDebug('svc-a'), makeServiceDebug('svc-b', 'failed')]; - const ctx = makeContext(deploys, services); - const result = builder.build(ctx); - const userMsg = result.messages[result.messages.length - 1].content; - expect(userMsg).toContain('## Services (2 total, 1 failing)'); - expect(userMsg).not.toContain('### org/repo-a'); - expect(userMsg).toContain('FAILING:'); - expect(userMsg).toContain('HEALTHY (1): svc-a'); - }); - }); - - describe('multi-repo services', () => { - it('groups services by repository with sub-headers', () => { - const deploys = [ - makeDeploy({ serviceName: 'svc-a', repoName: 'org/repo-a', status: 'BUILD_FAILED' }), - makeDeploy({ serviceName: 'svc-b', repoName: 'org/repo-a' }), - makeDeploy({ serviceName: 'svc-c', repoName: 'org/repo-b' }), - ]; - const services = [makeServiceDebug('svc-a', 'failed'), makeServiceDebug('svc-b'), makeServiceDebug('svc-c')]; - const ctx = makeContext(deploys, services); - const result = builder.build(ctx); - const userMsg = result.messages[result.messages.length - 1].content; - expect(userMsg).toContain('## Services (3 total, 1 failing)'); - expect(userMsg).toContain('### org/repo-a'); - expect(userMsg).toContain('### org/repo-b'); - - const repoAIdx = userMsg.indexOf('### org/repo-a'); - const repoBIdx = userMsg.indexOf('### org/repo-b'); - const failingIdx = userMsg.indexOf('FAILING:', repoAIdx); - expect(failingIdx).toBeGreaterThan(repoAIdx); - expect(failingIdx).toBeLessThan(repoBIdx); - }); - - it('uses primary repo fullName when repoName is missing', () => { - const deploys = [ - makeDeploy({ serviceName: 'svc-a', repoName: '' }), - makeDeploy({ serviceName: 'svc-b', repoName: 'org/repo-b' }), - ]; - const services = [makeServiceDebug('svc-a'), makeServiceDebug('svc-b')]; - const ctx = makeContext(deploys, services, 'org/repo-a'); - const result = builder.build(ctx); - const userMsg = result.messages[result.messages.length - 1].content; - expect(userMsg).toContain('### org/repo-a'); - expect(userMsg).toContain('### org/repo-b'); - }); - }); - }); -}); diff --git a/src/server/services/ai/prompts/__tests__/sectionRegistry.test.ts b/src/server/services/ai/prompts/__tests__/sectionRegistry.test.ts deleted file mode 100644 index 4aa918e3..00000000 --- a/src/server/services/ai/prompts/__tests__/sectionRegistry.test.ts +++ /dev/null @@ -1,288 +0,0 @@ -/** - * Copyright 2025 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 { assembleBasePrompt, PROMPT_SECTIONS } from '../sectionRegistry'; -import { SAFETY_SECTION } from '../sections/safety'; - -describe('Prompt Section Registry', () => { - describe('assembleBasePrompt', () => { - it('excludes safety section', () => { - const prompt = assembleBasePrompt(); - expect(prompt).not.toContain('Violation Examples'); - expect(prompt).not.toContain('# Final Reminder'); - }); - - it('includes all non-safety sections', () => { - const prompt = assembleBasePrompt(); - expect(prompt).toContain('developers who are blocked'); - expect(prompt).toContain('# Investigation Principles'); - expect(prompt).toContain('# Output Rules'); - expect(prompt).toContain('# Lifecycle Architecture'); - expect(prompt).toContain('# Configuration Architecture'); - }); - - it('excludes sections by id when excludeIds provided', () => { - const prompt = assembleBasePrompt(['reference']); - expect(prompt).not.toContain('# Lifecycle Architecture'); - expect(prompt).not.toContain('# Configuration Architecture'); - }); - - it('still includes non-excluded sections when excluding reference', () => { - const prompt = assembleBasePrompt(['reference']); - expect(prompt).toContain('developers who are blocked'); - expect(prompt).toContain('# Investigation Principles'); - }); - - it('supports excluding multiple sections', () => { - const prompt = assembleBasePrompt(['investigation', 'reference']); - expect(prompt).not.toContain(''); - expect(prompt).not.toContain('# Lifecycle Architecture'); - expect(prompt).toContain('developers who are blocked'); - }); - - it('returns same result with empty excludeIds as no argument', () => { - expect(assembleBasePrompt([])).toBe(assembleBasePrompt()); - }); - - it('assembles sections in correct order (reference before foundations before investigation)', () => { - const prompt = assembleBasePrompt(); - const markers = ['# Configuration Architecture', 'developers who are blocked', '# Investigation Principles']; - - const positions = markers.map((m) => prompt.indexOf(m)); - for (let i = 1; i < positions.length; i++) { - expect(positions[i]).toBeGreaterThan(positions[i - 1]); - } - }); - - it('assembled prompt contains all critical sections', () => { - const assembled = assembleBasePrompt(); - expect(assembled).toContain('Investigation Principles'); - expect(assembled).toContain('Fix Application Workflow'); - expect(assembled).toContain('Hypothesis-Driven'); - expect(assembled).toContain('investigation_complete'); - expect(assembled).toContain('Verification Protocol'); - expect(assembled).toContain('Two-Step Verification'); - expect(assembled).toContain('Lifecycle Architecture'); - expect(assembled).toContain(''); - expect(assembled).toContain('# Multi-Turn Conversation'); - }); - - it('contains XML section tags in assembled prompt', () => { - const assembled = assembleBasePrompt(); - expect(assembled).toContain(''); - expect(assembled).toContain(''); - expect(assembled).toContain(''); - expect(assembled).toContain(''); - expect(assembled).toContain(''); - expect(assembled).toContain(''); - expect(assembled).toContain(''); - expect(assembled).toContain(''); - expect(assembled).toContain(''); - expect(assembled).toContain(''); - expect(assembled).toContain(''); - expect(assembled).toContain(''); - }); - - it('contains output schema and field rules separation', () => { - const assembled = assembleBasePrompt(); - expect(assembled).toContain(''); - expect(assembled).toContain(''); - expect(assembled).toContain('## Field Rules'); - }); - - it('examples contain full JSON with example_output tags', () => { - const assembled = assembleBasePrompt(); - expect(assembled).toContain(''); - expect(assembled).toContain('"type": "investigation_complete"'); - expect(assembled).toContain('"fixesApplied": false'); - expect(assembled).toContain('"fixesApplied": true'); - expect(assembled).toContain('"oldContent"'); - expect(assembled).toContain('"newContent"'); - }); - - it('schema uses lowercase status values', () => { - const assembled = assembleBasePrompt(); - expect(assembled).toContain('build_failed'); - expect(assembled).toContain('deploy_failed'); - expect(assembled).not.toContain('"status": "BUILD_FAILED'); - expect(assembled).not.toContain('"status": "DEPLOY_FAILED'); - }); - }); - - describe('PROMPT_SECTIONS', () => { - it('has exactly 4 entries', () => { - expect(PROMPT_SECTIONS).toHaveLength(4); - }); - - it('has unique ids', () => { - const ids = PROMPT_SECTIONS.map((s) => s.id); - expect(new Set(ids).size).toBe(ids.length); - }); - - it('has unique order values', () => { - const orders = PROMPT_SECTIONS.map((s) => s.order); - expect(new Set(orders).size).toBe(orders.length); - }); - - it('positions safety section last', () => { - const sorted = PROMPT_SECTIONS.slice().sort((a, b) => a.order - b.order); - expect(sorted[sorted.length - 1].id).toBe('safety'); - }); - - it('positions reference section first', () => { - const sorted = PROMPT_SECTIONS.slice().sort((a, b) => a.order - b.order); - expect(sorted[0].id).toBe('reference'); - }); - - it('has non-empty content for each section', () => { - for (const section of PROMPT_SECTIONS) { - expect(typeof section.content).toBe('string'); - expect(section.content.length).toBeGreaterThan(0); - } - }); - - it('has a rationale for each section', () => { - for (const section of PROMPT_SECTIONS) { - expect(typeof section.rationale).toBe('string'); - expect(section.rationale.length).toBeGreaterThan(0); - } - }); - }); - - describe('content fidelity', () => { - it('assembled prompt contains all eval-critical phrases', () => { - const assembled = assembleBasePrompt(); - const requiredPhrases = [ - 'developers who are blocked', - 'reason about what you expect', - '# Output Rules', - 'Greeting, unclear question, need clarification', - 'Investigation Principles', - 'Hypothesis-Driven', - 'Evidence-Based Stopping', - 'Fix Application Workflow', - 'investigation_complete', - 'fixesApplied', - 'canAutoFix', - 'Verification Protocol', - 'Compare States', - 'Multi-Repo Architecture', - 'Two-Step Verification', - 'Multi-Turn Conversation', - 'Challenge Responses', - 'Confidence Levels', - 'Staleness Detection', - ]; - for (const phrase of requiredPhrases) { - expect(assembled).toContain(phrase); - } - }); - - it('assembled prompt does NOT contain V1 CoT suppression', () => { - const assembled = assembleBasePrompt(); - expect(assembled).not.toContain('Execute tools immediately without announcing intent'); - expect(assembled).not.toContain('Analysis AFTER results, not before'); - }); - - it('assembled prompt removes hard tool limit', () => { - const assembled = assembleBasePrompt(); - expect(assembled).not.toContain('Hard limit: 20 tool calls'); - }); - - it('safety section contains required safety rules', () => { - const safetySection = PROMPT_SECTIONS.find((s) => s.id === 'safety'); - expect(safetySection).toBeDefined(); - expect(safetySection!.content).toBe(SAFETY_SECTION); - expect(safetySection!.content).toContain('User Consent'); - expect(safetySection!.content).toContain('Surgical Changes'); - expect(safetySection!.content).toContain('Scope Boundaries'); - expect(safetySection!.content).toContain('Path Verification'); - expect(safetySection!.content).toContain('Compare States'); - expect(safetySection!.content).toContain('Violation Examples'); - }); - - it('safety section contains Content Integrity rule', () => { - const safetySection = PROMPT_SECTIONS.find((s) => s.id === 'safety'); - expect(safetySection!.content).toContain('Content Integrity'); - expect(safetySection!.content).toContain('EXACT content returned by get_file'); - }); - - it('safety section contains XML tags', () => { - const safetySection = PROMPT_SECTIONS.find((s) => s.id === 'safety'); - expect(safetySection!.content).toContain(''); - expect(safetySection!.content).toContain(''); - }); - - it('investigation section fix workflow contains verification step', () => { - const assembled = assembleBasePrompt(); - expect(assembled).toContain('does your new_content differ from the original in ONLY the intended lines'); - }); - - it('investigation section contains unrelated-changes negative example', () => { - const assembled = assembleBasePrompt(); - expect(assembled).toContain('commented-out service configuration'); - }); - - it('safety section contains No Fabrication rule', () => { - const safetySection = PROMPT_SECTIONS.find((s) => s.id === 'safety'); - expect(safetySection!.content).toContain('No Fabrication'); - }); - - it('investigation section contains Insufficient Evidence rule', () => { - const assembled = assembleBasePrompt(); - expect(assembled).toContain('Insufficient Evidence'); - expect(assembled).toContain('do NOT fabricate a root cause from config analysis alone'); - }); - - it('investigation section contains cite-then-conclude evidence pattern', () => { - const assembled = assembleBasePrompt(); - expect(assembled).toContain('cite-then-conclude'); - expect(assembled).toContain('first cite the specific error message from your tool results'); - }); - - it('investigation section contains external knowledge restriction', () => { - const assembled = assembleBasePrompt(); - expect(assembled).toContain('External Knowledge Restriction'); - expect(assembled).toContain('Do not apply general knowledge'); - }); - - it('investigation section explicitly permits saying I dont know', () => { - const assembled = assembleBasePrompt(); - expect(assembled).toContain("I don't have enough information to determine the root cause"); - }); - - it('investigation section contains fabricated-diagnosis negative example', () => { - const assembled = assembleBasePrompt(); - expect(assembled).toContain('finds 0 pods running and no error messages'); - }); - - it('canAutoFix rule requires actual error message', () => { - const assembled = assembleBasePrompt(); - expect(assembled).toContain('a specific error message from logs/K8s/build output points to the problem'); - }); - - it('has compound failure example', () => { - const assembled = assembleBasePrompt(); - expect(assembled).toContain('cascading effect'); - }); - - it('has everything-healthy example', () => { - const assembled = assembleBasePrompt(); - expect(assembled).toContain('All 3 services'); - expect(assembled).toContain('What specific issue are you seeing'); - }); - }); -}); diff --git a/src/server/services/ai/prompts/__tests__/tokenCounter.test.ts b/src/server/services/ai/prompts/__tests__/tokenCounter.test.ts deleted file mode 100644 index 67c49409..00000000 --- a/src/server/services/ai/prompts/__tests__/tokenCounter.test.ts +++ /dev/null @@ -1,132 +0,0 @@ -/** - * Copyright 2025 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 { - countTokens, - countSectionTokens, - getTokenBreakdown, - checkBudget, - PROVIDER_TOKEN_LIMITS, -} from '../tokenCounter'; - -describe('Token Counter', () => { - describe('countTokens', () => { - it('returns a positive number for "hello world"', () => { - const count = countTokens('hello world'); - expect(count).toBeGreaterThan(0); - }); - - it('returns a different result than text.length / 4', () => { - const text = 'hello world'; - const count = countTokens(text); - const naive = Math.ceil(text.length / 4); - expect(count).not.toBe(naive); - }); - - it('returns 0 for empty string', () => { - expect(countTokens('')).toBe(0); - }); - - it('returns more tokens for code with special characters than naive estimate', () => { - const code = 'const x = { foo: "bar", baz: [1, 2, 3] };\nconsole.log(x);'; - const count = countTokens(code); - expect(count).toBeGreaterThan(0); - expect(typeof count).toBe('number'); - }); - - it('returns a whole number', () => { - const count = countTokens('The quick brown fox jumps over the lazy dog'); - expect(Number.isInteger(count)).toBe(true); - }); - }); - - describe('countSectionTokens', () => { - it('returns an object with keys matching all 4 section IDs', () => { - const result = countSectionTokens(); - const expectedIds = ['foundations', 'investigation', 'reference', 'safety']; - for (const id of expectedIds) { - expect(result).toHaveProperty(id); - } - expect(Object.keys(result)).toHaveLength(4); - }); - - it('has a positive number for every section', () => { - const result = countSectionTokens(); - for (const [, value] of Object.entries(result)) { - expect(value).toBeGreaterThan(0); - } - }); - - it('has a total greater than 0', () => { - const result = countSectionTokens(); - const total = Object.values(result).reduce((sum, v) => sum + v, 0); - expect(total).toBeGreaterThan(0); - }); - }); - - describe('getTokenBreakdown', () => { - it('returns object with sections, providerAugmentation, environmentContext, and total', () => { - const breakdown = getTokenBreakdown('test prompt'); - expect(breakdown).toHaveProperty('sections'); - expect(breakdown).toHaveProperty('providerAugmentation'); - expect(breakdown).toHaveProperty('environmentContext'); - expect(breakdown).toHaveProperty('total'); - }); - - it('total equals sum of section values + providerAugmentation + environmentContext', () => { - const breakdown = getTokenBreakdown('some system prompt text for testing'); - const sectionSum = Object.values(breakdown.sections).reduce((sum, v) => sum + v, 0); - expect(breakdown.total).toBe(sectionSum + breakdown.providerAugmentation + breakdown.environmentContext); - }); - - it('uses pre-computed sections when provided', () => { - const precomputed = { identity: 100, communication: 50 }; - const breakdown = getTokenBreakdown('test prompt', precomputed); - expect(breakdown.sections).toEqual(precomputed); - }); - }); - - describe('checkBudget', () => { - it('returns overBudget=false for a small prompt', () => { - const result = checkBudget('hello', 'anthropic'); - expect(result.overBudget).toBe(false); - }); - - it('returns remaining = limit - used', () => { - const result = checkBudget('hello', 'anthropic'); - expect(result.remaining).toBe(result.limit - result.used); - }); - - it('has correct provider limits', () => { - expect(PROVIDER_TOKEN_LIMITS.anthropic).toBe(180000); - expect(PROVIDER_TOKEN_LIMITS.openai).toBe(110000); - expect(PROVIDER_TOKEN_LIMITS.gemini).toBe(900000); - }); - - it('returns overBudget=true when used exceeds limit', () => { - const result = checkBudget('any-text', 'openai', 200000); - expect(result.overBudget).toBe(true); - expect(result.used).toBe(200000); - expect(result.used).toBeGreaterThan(result.limit); - expect(result.remaining).toBeLessThan(0); - }); - - it('returns correct provider name', () => { - const result = checkBudget('hello', 'gemini'); - expect(result.provider).toBe('gemini'); - }); - }); -}); diff --git a/src/server/services/ai/prompts/builder.ts b/src/server/services/ai/prompts/builder.ts deleted file mode 100644 index 657f49a6..00000000 --- a/src/server/services/ai/prompts/builder.ts +++ /dev/null @@ -1,401 +0,0 @@ -/** - * Copyright 2025 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 { assembleBasePrompt, PROMPT_SECTIONS } from './sectionRegistry'; - -import { DebugContext, DebugMessage, ServiceDebugInfo } from '../../types/aiAgent'; -import { summarizeLifecycleYaml } from '../context/contextSummarizer'; - -export interface McpToolInfo { - serverName: string; - serverSlug: string; - toolName: string; - qualifiedName: string; - description: string; -} - -export interface PromptContext { - provider: 'anthropic' | 'openai' | 'gemini'; - debugContext: DebugContext; - conversationHistory: DebugMessage[]; - userMessage: string; - additiveRules?: string[]; - systemPromptOverride?: string; - excludedTools?: string[]; - excludedFilePatterns?: string[]; - mcpTools?: McpToolInfo[]; - excludeSections?: string[]; -} - -export interface BuiltPrompt { - systemPrompt: string; - messages: Array<{ role: string; content: string }>; -} - -export class AIAgentPromptBuilder { - private basePrompt = assembleBasePrompt(); - - public build(context: PromptContext): BuiltPrompt { - const excluded = context.excludeSections || []; - const registryExcludes = [...excluded.filter((id) => id !== 'safety')]; - const base = - context.systemPromptOverride || (excluded.length > 0 ? assembleBasePrompt(registryExcludes) : this.basePrompt); - - const layers = [ - base, - this.buildMcpToolsLayer(context.mcpTools), - this.buildCustomRulesLayer(context), - this.buildAccessRestrictionsNotice(context), - excluded.includes('safety') ? '' : this.buildSafetyRulesLayer(), - ]; - - const systemPrompt = layers.filter(Boolean).join('\n'); - const messages = this.buildMessages(context); - - return { systemPrompt, messages }; - } - - private buildCustomRulesLayer(context: PromptContext): string { - if (!context.additiveRules || context.additiveRules.length === 0) { - return ''; - } - - const rules = context.additiveRules.map((rule) => `- ${rule}`).join('\n'); - return `\n\n---\n\n# Custom Rules\n\n${rules}`; - } - - private buildAccessRestrictionsNotice(context: PromptContext): string { - const parts: string[] = []; - if (context.excludedTools && context.excludedTools.length > 0) { - parts.push( - `The following tools are NOT available to you: ${context.excludedTools.join(', ')}. Do not attempt to use them.` - ); - } - if (context.excludedFilePatterns && context.excludedFilePatterns.length > 0) { - parts.push( - `The following file patterns are restricted and you cannot access them: ${context.excludedFilePatterns.join( - ', ' - )}. If a user asks about these files, explain they are restricted.` - ); - } - if (parts.length === 0) return ''; - return `\n\n---\n\n# Access Restrictions\n\n${parts.join('\n\n')}`; - } - - private buildMcpToolsLayer(mcpTools?: McpToolInfo[]): string { - if (!mcpTools || mcpTools.length === 0) return ''; - - const byServer = new Map(); - for (const tool of mcpTools) { - const key = tool.serverName; - if (!byServer.has(key)) byServer.set(key, []); - byServer.get(key)!.push(tool); - } - - const serverSections = Array.from(byServer.entries()) - .map(([serverName, tools]) => { - const toolList = tools.map((t) => `- **${t.qualifiedName}**: ${t.description}`).join('\n'); - return `### ${serverName}\n${toolList}`; - }) - .join('\n\n'); - - return ` - ---- - -# External Tools (MCP) - -You have access to external tools from connected MCP servers. Use these tools when they can provide better or additional information beyond your built-in tools. - -${serverSections} - -**When to use MCP tools:** -- Use them alongside built-in tools during investigation — they provide complementary data -- If an MCP tool can answer the user's question directly, prefer it over manual investigation -- MCP tool names are prefixed with \`mcp____\` — call them like any other tool -- If an MCP tool fails, fall back to built-in tools`; - } - - private buildSafetyRulesLayer(): string { - return '\n\n---\n\n' + PROMPT_SECTIONS.find((s) => s.id === 'safety')!.content; - } - - private static FAILED_DEPLOY_STATUSES = new Set(['BUILD_FAILED', 'DEPLOY_FAILED', 'ERROR']); - - private isFailingService(deploy: any, serviceDebug?: ServiceDebugInfo): boolean { - if (AIAgentPromptBuilder.FAILED_DEPLOY_STATUSES.has(deploy.status)) return true; - if (serviceDebug?.status === 'failed') return true; - if (serviceDebug?.issues && serviceDebug.issues.length > 0) return true; - return false; - } - - private renderFailingService(d: any, serviceDebug?: ServiceDebugInfo): string { - let info = `- ${d.serviceName}: ${d.status}${d.statusMessage ? ` - ${d.statusMessage}` : ''}`; - info += `\n Type: ${d.type}`; - if (d.builderEngine) info += ` | Builder: ${d.builderEngine}`; - if (d.helmChart) info += ` | Chart: ${d.helmChart}`; - - if (d.buildPipelineId) { - info += `\n Build: Codefresh (buildPipelineId: ${d.buildPipelineId})`; - } else if (d.builderEngine) { - info += `\n Build: Native/${d.builderEngine} (label_selector="lc-service=${d.serviceName}")`; - } - - if (d.deployPipelineId) { - info += `\n Deploy: Codefresh (deployPipelineId: ${d.deployPipelineId})`; - } else { - info += `\n Deploy: Native/Helm (label_selector="lc-service=${d.serviceName}")`; - } - - info += `\n Image: ${d.dockerImage || 'N/A'}`; - - if (serviceDebug) { - const podCount = serviceDebug.pods.length; - const readyPods = serviceDebug.pods.filter( - (p) => p.phase === 'Running' && p.containerStatuses.every((c: any) => c.ready) - ).length; - if (podCount > 0) { - info += `\n K8s: ${readyPods}/${podCount} pods ready`; - } - - for (const pod of serviceDebug.pods) { - for (const cs of pod.containerStatuses) { - const waitReason = cs.state?.waiting?.reason; - const termReason = cs.lastState?.terminated?.reason; - const parts: string[] = []; - if (waitReason) parts.push(waitReason); - if (termReason) parts.push(`last: ${termReason}`); - if (cs.restartCount > 0) parts.push(`restarts: ${cs.restartCount}`); - if (parts.length > 0) { - info += `\n Container ${cs.name}: ${parts.join(', ')}`; - } - } - const unschedulable = pod.conditions.find((c: any) => c.type === 'PodScheduled' && c.status === 'False'); - if (unschedulable) { - info += `\n Scheduling: ${unschedulable.message || 'cannot be scheduled'}`; - } - } - - if (serviceDebug.issues.length > 0) { - info += `\n Issues: ${serviceDebug.issues.map((i) => i.title).join('; ')}`; - } - if (serviceDebug.events.length > 0) { - const warningEvents = serviceDebug.events.filter((e) => e.type === 'Warning'); - if (warningEvents.length > 0) { - info += `\n Events: ${warningEvents - .slice(0, 3) - .map((e) => `${e.reason}: ${e.message}`) - .join('; ')}`; - } - } - - if (serviceDebug.pods.length > 0) { - const logsSource = serviceDebug.pods.find((p) => p.recentLogs) || serviceDebug.pods[0]; - if (logsSource?.recentLogs) { - const logLines = logsSource.recentLogs.split('\n').filter(Boolean); - const tail = logLines.slice(-15).join('\n'); - if (tail) { - info += `\n Recent logs (${logsSource.name}, last ${Math.min(logLines.length, 15)} lines):\n${tail}`; - } - } - } - } - - return info; - } - - private renderRepoGroup(failing: any[], healthy: any[], servicesByName: Map): string { - let section = ''; - if (failing.length > 0) { - section += '\nFAILING:'; - for (const d of failing) { - section += '\n' + this.renderFailingService(d, servicesByName.get(d.serviceName)); - } - } - if (healthy.length > 30) { - section += `\n\nHEALTHY: ${healthy.length} services (use query_database on deploys table to list specific services)`; - } else if (healthy.length > 0) { - section += `\n\nHEALTHY (${healthy.length}): ${healthy.map((d) => d.serviceName).join(', ')}`; - } - return section; - } - - private buildDependencyGraphSummary(dependencyGraph: Record, failingServiceNames: Set): string { - if (!dependencyGraph || failingServiceNames.size === 0) return ''; - - const edges: Array<{ source: string; target: string }> = dependencyGraph.edges; - if (!Array.isArray(edges) || edges.length === 0) return ''; - - const dependsOn = new Map>(); - const requiredBy = new Map>(); - - for (const edge of edges) { - if (edge.source === edge.target) continue; - if (!dependsOn.has(edge.source)) dependsOn.set(edge.source, new Set()); - dependsOn.get(edge.source)!.add(edge.target); - if (!requiredBy.has(edge.target)) requiredBy.set(edge.target, new Set()); - requiredBy.get(edge.target)!.add(edge.source); - } - - const lines: string[] = []; - - for (const name of failingServiceNames) { - const deps = dependsOn.get(name); - const consumers = requiredBy.get(name); - if (!deps && !consumers) continue; - - if (deps && deps.size > 0) { - const failingDeps = [...deps].filter((d) => failingServiceNames.has(d)); - if (failingDeps.length > 0) { - lines.push(`${name} depends on (ALSO FAILING): ${failingDeps.join(', ')}`); - const healthyDeps = [...deps].filter((d) => !failingServiceNames.has(d)); - if (healthyDeps.length > 0) { - lines.push(`${name} depends on (healthy): ${healthyDeps.join(', ')}`); - } - } else { - lines.push(`${name} depends on: ${[...deps].join(', ')}`); - } - } - if (consumers && consumers.size > 0) { - lines.push(`${name} is required by: ${[...consumers].join(', ')}`); - } - } - - if (lines.length === 0) return ''; - - return `\n\n## Service Dependencies (${ - failingServiceNames.size - } failing services)\n\nUse this dependency information to identify root causes. If a failing service depends on another failing service, investigate the dependency first.\n\n${lines.join( - '\n' - )}`; - } - - private buildEnvironmentContext(debugContext: DebugContext): string { - const lc = debugContext.lifecycleContext; - const servicesByName = new Map(); - for (const s of debugContext.services) { - servicesByName.set(s.name, s); - } - - const failingDeploys: any[] = []; - const healthyDeploys: any[] = []; - for (const d of lc.deploys) { - const serviceDebug = servicesByName.get(d.serviceName); - if (this.isFailingService(d, serviceDebug)) { - failingDeploys.push(d); - } else { - healthyDeploys.push(d); - } - } - - const repoNames = new Set(); - for (const d of lc.deploys) { - repoNames.add(d.repoName || lc.pullRequest.fullName); - } - const isMultiRepo = repoNames.size > 1; - - let servicesSection = `## Services (${lc.deploys.length} total, ${failingDeploys.length} failing)`; - - if (isMultiRepo) { - const repoGroups = new Map(); - for (const d of [...failingDeploys, ...healthyDeploys]) { - const repo = d.repoName || lc.pullRequest.fullName; - if (!repoGroups.has(repo)) repoGroups.set(repo, { failing: [], healthy: [] }); - const group = repoGroups.get(repo)!; - if (failingDeploys.includes(d)) { - group.failing.push(d); - } else { - group.healthy.push(d); - } - } - for (const [repo, group] of repoGroups) { - servicesSection += `\n\n### ${repo}`; - servicesSection += this.renderRepoGroup(group.failing, group.healthy, servicesByName); - } - } else { - servicesSection += this.renderRepoGroup(failingDeploys, healthyDeploys, servicesByName); - } - - const failingServiceNames = new Set(failingDeploys.map((d: any) => d.serviceName)); - const dependencySummary = this.buildDependencyGraphSummary(lc.build.dependencyGraph, failingServiceNames); - - let lifecycleYamlSection: string; - if (!debugContext.lifecycleYaml) { - lifecycleYamlSection = 'lifecycle.yaml not available'; - } else if (debugContext.lifecycleYaml.error) { - lifecycleYamlSection = `Could not fetch lifecycle.yaml: ${debugContext.lifecycleYaml.error}`; - } else if (debugContext.lifecycleYaml.content) { - const summary = summarizeLifecycleYaml(debugContext.lifecycleYaml.content); - if (summary.parsed) { - lifecycleYamlSection = `${summary.text}\n[Use get_file("lifecycle.yaml") for full configuration]`; - } else { - const lines = debugContext.lifecycleYaml.content.split('\n'); - const truncated = lines.slice(0, 200).join('\n'); - lifecycleYamlSection = `${truncated}${ - lines.length > 200 - ? `\n... (${ - lines.length - 200 - } more lines truncated)\n[Use get_file("lifecycle.yaml") for full configuration]` - : '' - }`; - } - } else { - lifecycleYamlSection = 'lifecycle.yaml is empty'; - } - - const gatheredAt = - debugContext.gatheredAt instanceof Date ? debugContext.gatheredAt.toISOString() : String(debugContext.gatheredAt); - - return ` - ---- - -# Environment State (gathered: ${gatheredAt}) - -Build: ${debugContext.buildUuid} | Status: ${lc.build.status} | Namespace: ${lc.build.namespace} -PR: #${lc.pullRequest.number || 'N/A'} "${lc.pullRequest.title || 'N/A'}" by ${lc.pullRequest.username || 'N/A'} -Labels: ${lc.pullRequest.labels?.length ? lc.pullRequest.labels.join(', ') : 'none'} -Repo: ${lc.pullRequest.fullName} @ ${lc.pullRequest.branch} (base: ${lc.pullRequest.baseBranch || 'N/A'}) -SHA: ${lc.pullRequest.latestCommit || lc.build.sha || 'N/A'} -${ - lc.pullRequest.commentId - ? `PR Comment ID: ${lc.pullRequest.commentId} (use get_pr_comment to see enabled services)` - : 'PR Comment ID: Not available' -} - -${servicesSection}${dependencySummary} - -## Configuration (lifecycle.yaml) -${lifecycleYamlSection}`; - } - - private buildMessages(context: PromptContext): Array<{ role: string; content: string }> { - const messages = context.conversationHistory.map((m) => ({ - role: m.role, - content: m.content, - })); - - const envContext = this.buildEnvironmentContext(context.debugContext); - const finalMessage = `${envContext}\n\n${context.userMessage}`; - - messages.push({ - role: 'user', - content: finalMessage, - }); - - return messages; - } -} diff --git a/src/server/services/ai/prompts/sectionRegistry.ts b/src/server/services/ai/prompts/sectionRegistry.ts deleted file mode 100644 index c210b35a..00000000 --- a/src/server/services/ai/prompts/sectionRegistry.ts +++ /dev/null @@ -1,60 +0,0 @@ -/** - * Copyright 2025 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 { FOUNDATIONS_SECTION, INVESTIGATION_SECTION, REFERENCE_SECTION, SAFETY_SECTION } from './sections'; - -export interface PromptSection { - id: string; - content: string; - order: number; - rationale: string; -} - -export const PROMPT_SECTIONS: PromptSection[] = [ - { - id: 'reference', - content: REFERENCE_SECTION, - order: 1, - rationale: - 'Domain knowledge first — long-context research shows 30% improvement when reference data precedes instructions', - }, - { - id: 'foundations', - content: FOUNDATIONS_SECTION, - order: 2, - rationale: 'Agent identity, communication style, and constraints after domain context is established', - }, - { - id: 'investigation', - content: INVESTIGATION_SECTION, - order: 3, - rationale: 'Investigation methodology, output format, and examples define operational patterns', - }, - { - id: 'safety', - content: SAFETY_SECTION, - order: 4, - rationale: 'Safety rules positioned LAST for recency bias advantage', - }, -]; - -export function assembleBasePrompt(excludeIds: string[] = []): string { - return PROMPT_SECTIONS.slice() - .sort((a, b) => a.order - b.order) - .filter((s) => s.id !== 'safety' && !excludeIds.includes(s.id)) - .map((s) => s.content) - .join('\n\n'); -} diff --git a/src/server/services/ai/prompts/sections/foundations.ts b/src/server/services/ai/prompts/sections/foundations.ts deleted file mode 100644 index 891f5a12..00000000 --- a/src/server/services/ai/prompts/sections/foundations.ts +++ /dev/null @@ -1,83 +0,0 @@ -/** - * Copyright 2025 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. - */ - -export const FOUNDATIONS_SECTION = ` -You are an SRE debugging agent for Lifecycle, a platform that creates ephemeral Kubernetes environments from pull requests. Your users are developers who are blocked — their environment isn't working and they need to get back to testing. Your job is to identify why and give them a clear path to resolution. - -# Primary Objective - -Identify root causes by comparing desired config state vs actual runtime state, then provide specific fixes. - -# Capabilities & Instructions - -## Data Reuse - -The injected environment context contains deployment statuses, service health, and K8s state gathered at the timestamp shown. Use this data directly. Only call tools for information not already in context or when the user indicates state has changed. - -## Context Freshness - -The injected context reflects current DB state at the \`gathered at\` timestamp. Trust it unless the user indicates something changed ("I just pushed", "I redeployed"). When in doubt about staleness, verify with a single targeted query. - -## Verification - -If the injected summary lacks detail for a specific service, read that service's config via get_file. Reference data from injected context and prior tool results directly rather than re-fetching. - -## Root Cause Focus - -Compare DESIRED (config files) vs ACTUAL (runtime) for root cause identification. - -## Parallel Execution - -Execute independent tool calls in parallel. Targeted over exhaustive — investigate specific failing services, not all. - -## Reasoning - -Before calling tools, briefly reason about what you expect to find. Use evidence from tool results to confirm or refute your hypothesis. - -# Communication Style - -- Get to the point. Lead with findings, not process descriptions. -- Professional and concise — under 5 lines when practical. -- GitHub-flavored Markdown. Clarity over brevity when they conflict. -- Tools for actions, text only for user communication. - -## Handling Truncated Data - -Tool results may be truncated with \`[Truncated: showing X of Y chars]\`. When you see this: -1. Focus on what IS visible — errors typically appear at the end -2. If critical info is missing, re-query with tighter filters -3. Note truncation in your analysis - -## Managing Data Volume - -When results are large or truncated: -1. Note the truncation and tell the user what was cut -2. Re-query with narrower filters rather than requesting larger limits -3. Prefer targeted queries over broad ones -4. If results show "X of Y total", ask the user which subset they care about - -## Two-Step Verification - -For uncertain diagnoses (ambiguous logs, partial data, multiple possible causes): state your confidence level, identify what additional evidence would confirm or refute, and gather that evidence before concluding. - -# Efficiency - -Be efficient with tool calls. The system enforces a maximum, but most investigations should complete in 5-15 calls. Prefer targeted queries over broad scans. -Each tool should be called at most once with the same arguments per conversation. If a call errors or returns not found, move on. -K8s: ONE call per resource type in the namespace. Use label_selector=lc-service={serviceName} to scope when investigating a specific service. -Logs: ONE call per pod. If get_pod_logs fails, move on. -When stuck, output what you know and ask the user. -`; diff --git a/src/server/services/ai/prompts/sections/index.ts b/src/server/services/ai/prompts/sections/index.ts deleted file mode 100644 index 5f161abb..00000000 --- a/src/server/services/ai/prompts/sections/index.ts +++ /dev/null @@ -1,20 +0,0 @@ -/** - * Copyright 2025 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. - */ - -export { FOUNDATIONS_SECTION } from './foundations'; -export { INVESTIGATION_SECTION } from './investigation'; -export { REFERENCE_SECTION } from './reference'; -export { SAFETY_SECTION } from './safety'; diff --git a/src/server/services/ai/prompts/sections/investigation.ts b/src/server/services/ai/prompts/sections/investigation.ts deleted file mode 100644 index 82848e0d..00000000 --- a/src/server/services/ai/prompts/sections/investigation.ts +++ /dev/null @@ -1,355 +0,0 @@ -/** - * Copyright 2025 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. - */ - -export const INVESTIGATION_SECTION = ` -# Investigation Principles - -Vague messages ("something seems wrong", "help") — ask for clarification before investigating. - -**Hypothesis-Driven Flow:** Read injected context (build status, service statuses, K8s state) → form a hypothesis from symptoms → test with targeted tool calls → confirm with evidence or pivot. When a hypothesis is wrong, explain why you're pivoting. - -**Evidence-Based Stopping:** Once you confirm the primary root cause, briefly check if other failing services share the same cause or have independent issues. Report the primary failure in depth and others briefly. - -**Priority:** Build failures before deploy failures. Investigate the most critical failing service first; briefly mention others. BUILD_FAILED with clear error in logs = sufficient root cause — no need to also check K8s. - -**Data Reuse:** Injected context contains environment state — read it before calling ANY tools. Start with recent log lines (tail); request more only if error not found. See Multi-Turn Conversation section for follow-up behavior. - -**Evidence Chains (cite-then-conclude):** Before stating any diagnosis, first cite the specific error message from your tool results that supports it. If you cannot point to a specific error from logs, K8s events, or build output, you do not have evidence and must not make the diagnosis. Reading a config file and deciding it "looks wrong" based on your general knowledge is NOT evidence — only an actual error message from runtime or build output counts. - -**External Knowledge Restriction:** Base your diagnosis ONLY on information from tool results (logs, K8s events, build output, database queries) and injected context. Do not apply general knowledge about how configuration files "should" be structured. A config is only wrong when a tool result contains an error pointing to it. - -**Uncertainty — you are allowed to say "I don't know":** When root cause can't be determined from available evidence, you are explicitly permitted to say "I don't have enough information to determine the root cause." This is always preferable to a wrong diagnosis. Describe what you checked and what remains unknown. - -**Insufficient Evidence:** When your investigation finds no error messages — no build logs with errors, no pod crash reasons, no K8s events — do NOT fabricate a root cause from config analysis alone. Instead, report what you checked, what you found (or didn't find), and state that you couldn't determine the root cause. Common causes of missing runtime artifacts (0 pods, no builds) include missing PR labels, webhook delivery issues, and build triggers — mention these as possibilities without diagnosing one as the cause. - -**You are the investigator.** Read logs, configs, and files yourself. Report findings directly — never suggest the user investigate. Use ACTUAL paths and values, not placeholders. - -## Fix Application Workflow - -"Fix it for me" = EXPLICIT consent for ONLY that specific issue. -1. Call get_file — fetch current content. This is your ground truth. -2. Copy the file content EXACTLY. Change ONLY the lines needed for the fix. Do not remove comments, delete unused sections, reformat, or make any other modifications. -3. Before calling update_file, verify: does your new_content differ from the original in ONLY the intended lines? If you changed anything else, redo step 2. -4. Execute ONLY the mutating tool needed for the selected issue: - - File content fixes: update_file - - PR label fixes: update_pr_labels - - Runtime K8s patch fixes: patch_k8s_resource -5. Verify success response, extract commit_url. -6. Output JSON with fixesApplied: true + commitUrl. - -Two-step fix pattern: (1) patch K8s resource for immediate relief, (2) persist in config and commit. -lifecycle.yaml fixes: offer to apply and commit. Service repo fixes: describe what to change (cannot commit to other repos). -Agent can perform any operation within the service namespace, never outside. - -## Verification Protocol - -Compare States: DESIRED (config) vs ACTUAL (runtime). Match=intentional, mismatch=bug. Always verify config before diagnosing. - - - -# Output Rules - -Decide your output format using this decision tree: -- Greeting, unclear question, need clarification → plain text (markdown, 1-5 lines) -- Everything healthy, no issues found → plain text confirmation, ask what they're seeing -- Any issue found OR config problem even if pods running → JSON schema below -- Fix applied successfully → JSON with fixesApplied: true and commitUrl - -JSON — output ONLY the raw JSON object. Do not include any markdown preamble, conversational text, or code fences before or after it. If you want to provide a conversational summary, put it in the "summary" field of the JSON. The response must start with \`{\` and end with \`}\`: - - -{ - "type": "investigation_complete", - "summary": "string — if fixesApplied=false say 'needs to be fixed' NOT 'has been fixed'", - "fixesApplied": false, - "services": [ - { - "serviceName": "string", - "status": "build_failed | deploy_failed | error | ready", - "issue": "string — ROOT CAUSE with WHY, not just WHAT", - "keyError": "string | undefined — exact error message from logs", - "errorSource": "string | undefined — e.g. 'build_logs', 'pod_logs', 'k8s_events'", - "errorSourceDetail": "string | undefined — e.g. pod name, build job id", - "suggestedFix": "string — actionable fix description", - "canAutoFix": "boolean | undefined", - "filePath": "string | undefined — primary file to fix", - "lineNumber": "number | undefined", - "lineNumberEnd": "number | undefined", - "files": [ - { - "path": "string — file path", - "lineNumber": "number | undefined", - "lineNumberEnd": "number | undefined", - "description": "string | undefined — what the change does", - "oldContent": "string | undefined — current content to replace", - "newContent": "string | undefined — corrected content" - } - ], - "commitUrl": "string | undefined — present only when fixesApplied=true" - } - ], - "repository": { - "owner": "string | undefined", - "name": "string | undefined", - "branch": "string | undefined", - "sha": "string | undefined" - } -} - - -## Field Rules - -**status:** Use lowercase values: \`build_failed\`, \`deploy_failed\`, \`error\`, \`ready\`. These must match the TypeScript union type exactly. - -**canAutoFix=true** ONLY when: (1) a specific error message from logs/K8s/build output points to the problem, (2) a concrete fix is prepared and verified, (3) 100% certain, (4) the required mutating tool is actually available in this run (e.g., update_file for file edits, update_pr_labels for PR labels, patch_k8s_resource for K8s patches). False when: no error messages found, config-only analysis, user decision needed, uncertainty, missing files, or required tool unavailable. - -**fixesApplied=true** ONLY when: you actually executed the intended fix tool for the selected issue and it succeeded. Otherwise false. - -**keyError:** Extract the exact error message from logs. Helps the UI display the root cause prominently. - -**errorSource / errorSourceDetail:** Identify where the error was found (e.g., errorSource="pod_logs", errorSourceDetail="web-build-abc-xyz"). - -**Line numbers:** Count lines from the raw file content returned by get_file (the result includes totalLines). Include lineNumber + lineNumberEnd. If can't determine, omit both. - -**files array:** Include oldContent and newContent when canAutoFix=true. The UI renders these as a diff view. Each entry represents one file change. - -**suggestedFix patterns:** For single-line changes, use the pattern: \`Change from '' to '' in \`. The UI parses this pattern to render inline diffs. For multi-line changes, use the files array with oldContent/newContent instead. - -**commitUrl:** Include only when the successful fix action produced a commit URL (typically update_file). - -**repository:** Include when available from the injected context. Used by the UI to build GitHub links for file paths. - - - -# Examples - - -Why did my web service build fail? - -Hypothesis: Injected context shows web BUILD_FAILED. Likely a dockerfile or dependency issue. -Action: [get_pod_logs: pod="web-build-xyz"] → logs show "COPY failed: file not found ./src/index.ts" -Action: [get_file: lifecycle.yaml] → dockerfile references ./src/index.ts but package.json main is ./dist/index.js -Root cause confirmed: dockerfile references wrong path. Stop investigating. - - -{ - "type": "investigation_complete", - "summary": "Web service build failed because the Dockerfile copies ./src/index.ts which does not exist — the correct entrypoint is ./dist/index.js. This needs to be fixed in lifecycle.yaml.", - "fixesApplied": false, - "services": [ - { - "serviceName": "web", - "status": "build_failed", - "issue": "Dockerfile COPY references ./src/index.ts but the actual entrypoint is ./dist/index.js. The build fails at the COPY step because the source file does not exist.", - "keyError": "COPY failed: file not found ./src/index.ts", - "errorSource": "build_logs", - "errorSourceDetail": "web-build-xyz", - "suggestedFix": "Change dockerfilePath entrypoint from './src/index.ts' to './dist/index.js' in lifecycle.yaml", - "canAutoFix": true, - "filePath": "lifecycle.yaml", - "lineNumber": 12, - "lineNumberEnd": 12, - "files": [ - { - "path": "lifecycle.yaml", - "lineNumber": 12, - "lineNumberEnd": 12, - "description": "Fix Dockerfile entrypoint path", - "oldContent": " entrypoint: ./src/index.ts", - "newContent": " entrypoint: ./dist/index.js" - } - ] - } - ] -} - - - - -My environment looks broken, what's going on? - -Hypothesis: Injected context shows api DEPLOY_FAILED, possible resource or image issue. -Action: [get_k8s_resources: pods, namespace="env-abc"] → api pod shows OOMKilled (256Mi limit) -Root cause confirmed: memory limit too low. Stop — don't also check configs or other healthy services. - - -{ - "type": "investigation_complete", - "summary": "The api service is crashing due to OOMKilled — the memory limit of 256Mi is too low. This needs to be fixed in the Helm values file.", - "fixesApplied": false, - "services": [ - { - "serviceName": "api", - "status": "deploy_failed", - "issue": "Pod is being OOMKilled because the memory limit (256Mi) is insufficient for this service. Kubernetes terminates the container when it exceeds the limit.", - "keyError": "OOMKilled", - "errorSource": "k8s_events", - "errorSourceDetail": "api-deployment-abc-xyz", - "suggestedFix": "Change resources.limits.memory from '256Mi' to '512Mi' in helm/values.yaml", - "canAutoFix": true, - "filePath": "helm/values.yaml", - "lineNumber": 28, - "lineNumberEnd": 28, - "files": [ - { - "path": "helm/values.yaml", - "lineNumber": 28, - "lineNumberEnd": 28, - "description": "Increase memory limit to prevent OOMKill", - "oldContent": " memory: 256Mi", - "newContent": " memory: 512Mi" - } - ] - } - ] -} - - - - -Fix the dockerfile path in lifecycle.yaml - -Action: [get_file: lifecycle.yaml] → read current content (147 lines), identified wrong path on line 42 -Action: Copied original content, changed ONLY line 42. Verified: 146 of 147 lines unchanged. No comments removed, no reformatting. -Action: [update_file: path="lifecycle.yaml", new_content="..."] → success, commit URL returned -Fix applied successfully. - - -{ - "type": "investigation_complete", - "summary": "Fixed the dockerfile path in lifecycle.yaml. The change has been committed.", - "fixesApplied": true, - "services": [ - { - "serviceName": "web", - "status": "build_failed", - "issue": "Dockerfile path referenced a non-existent file ./src/index.ts instead of ./dist/index.js.", - "suggestedFix": "Change dockerfile entrypoint from './src/index.ts' to './dist/index.js' in lifecycle.yaml", - "canAutoFix": true, - "filePath": "lifecycle.yaml", - "lineNumber": 12, - "lineNumberEnd": 12, - "files": [ - { - "path": "lifecycle.yaml", - "lineNumber": 12, - "lineNumberEnd": 12, - "description": "Fix Dockerfile entrypoint path", - "oldContent": " entrypoint: ./src/index.ts", - "newContent": " entrypoint: ./dist/index.js" - } - ], - "commitUrl": "https://github.com/org/repo/commit/abc123" - } - ] -} - - - - -My environment is completely broken, nothing works - -Hypothesis: Injected context shows api DEPLOY_FAILED and worker BUILD_FAILED. Two services failing — could be related or independent. -Action: [get_pod_logs: pod="worker-build-xyz"] → logs show "npm ERR! Missing dependency: @org/shared-lib@2.0" -Root cause for worker: missing dependency in package.json. -Action: [get_k8s_resources: pods, namespace="env-abc"] → api pod shows CrashLoopBackOff, logs show "Error: connect ECONNREFUSED worker:3000" -api failure is downstream — it depends on worker, which failed to build. -Primary root cause: worker build failure (missing dependency). api failure is a cascading effect. - - -{ - "type": "investigation_complete", - "summary": "Worker build failed due to a missing dependency (@org/shared-lib@2.0), which caused a cascading failure in the api service. Fix the worker dependency first — api will recover once worker is running.", - "fixesApplied": false, - "services": [ - { - "serviceName": "worker", - "status": "build_failed", - "issue": "Build fails because @org/shared-lib@2.0 is not in package.json dependencies. This is the primary root cause — the api service depends on worker and fails downstream.", - "keyError": "npm ERR! Missing dependency: @org/shared-lib@2.0", - "errorSource": "build_logs", - "errorSourceDetail": "worker-build-xyz", - "suggestedFix": "Add @org/shared-lib@2.0 to dependencies in package.json", - "canAutoFix": false, - "filePath": "package.json" - }, - { - "serviceName": "api", - "status": "deploy_failed", - "issue": "CrashLoopBackOff because it cannot connect to worker:3000. This is a cascading effect of the worker build failure — not an independent issue.", - "keyError": "Error: connect ECONNREFUSED worker:3000", - "errorSource": "pod_logs", - "errorSourceDetail": "api-deployment-abc-xyz", - "suggestedFix": "No action needed — api will recover once the worker build is fixed and deployed", - "canAutoFix": false - } - ] -} - - - - -Something seems wrong with my environment - -Hypothesis: Check injected context for failures. -Injected context shows: all 3 services READY, all pods Running with ready=desired. -Verify with one fresh K8s check: [get_k8s_resources: pods, namespace="env-abc"] → confirms all healthy. - - -All 3 services (api, web, worker) are running and healthy. What specific issue are you seeing? - - - - -ANTI-PATTERN: Agent finds OOMKilled in pod logs, confirms root cause, then ALSO reads lifecycle.yaml, values.yaml, and checks events. This wastes tool calls after root cause is already confirmed. -CORRECT: Stop after OOMKilled confirmation. Report the fix (increase memory limit) immediately. - - - -ANTI-PATTERN: Agent fixes dockerfilePath on line 42 of lifecycle.yaml, but the committed diff also deletes 28 lines of commented-out service configuration at the bottom of the file. The user only approved the single-line fix. -CORRECT: Copy the entire original file content from get_file verbatim. Change ONLY line 42. Every other line — including comments, blank lines, and disabled sections — must remain byte-for-byte identical. - - - -ANTI-PATTERN: Agent finds 0 pods running and no error messages in any logs. Reads lifecycle.yaml and decides the YAML structure "looks wrong" — proposes config changes with canAutoFix=true. No actual error message pointed to the config. The real cause was outside the agent's investigation (e.g., a missing PR label). -CORRECT: Report that no pods were found and no build/deploy errors were found in logs. State that the root cause could not be determined from available evidence. Suggest possible external causes (PR labels, webhook delivery, build triggers) without diagnosing a specific one. - - - - -# Multi-Turn Conversation - -**How conversation context works:** The LLM provider (Gemini/Anthropic/OpenAI) manages the full message thread natively. Your tool calls and their results from the current turn are always visible to you. Tool results from recent prior turns are preserved in the conversation history by the observation masker; older tool results may be masked. Additionally, fresh environment context (deployment status, service health) is injected at the start of each new user message via Tier 2. - -**Context Reuse:** Given the above, when the user asks a follow-up ("what about the other service?", "and the logs?"), reference findings from your prior tool calls visible in this conversation thread. Only call tools when the question targets data you haven't gathered yet or data that may have gone stale. - -**Staleness Detection:** If the user indicates state has changed ("I just deployed", "I pushed a fix", "try again"), re-gather the relevant data. Otherwise, trust your prior findings. The environment context block injected each turn includes a "gathered at" timestamp -- use it to gauge freshness. - -**Challenge Responses:** When the user disputes your findings ("that's not right", "are you sure?"): -1. Re-examine your evidence for the specific claim -2. If evidence is strong, defend your conclusion with specific citations -3. If evidence is weak or the user provides new information, re-verify with targeted tool calls -4. Acknowledge errors explicitly when corrected - -**Confidence Levels:** Qualify conclusions based on evidence strength: -- Confirmed: direct evidence from logs/K8s state (e.g., "Build failed due to X — confirmed by build logs") -- Likely: strong circumstantial evidence (e.g., "This is likely caused by Y — consistent with symptoms but no direct log entry") -- Uncertain: multiple possible causes (e.g., "Could be X or Y — need more data to confirm") - -**Long Conversations:** In conversations exceeding many turns, prioritize the most recent context. If asked about something discussed many turns ago, acknowledge the data may be stale and offer to re-check. - -**Proactive Observations:** If during investigation you notice related issues affecting other services, briefly mention them after addressing the user's primary question. -`; diff --git a/src/server/services/ai/prompts/sections/reference.ts b/src/server/services/ai/prompts/sections/reference.ts deleted file mode 100644 index 81dbf0ca..00000000 --- a/src/server/services/ai/prompts/sections/reference.ts +++ /dev/null @@ -1,59 +0,0 @@ -/** - * Copyright 2025 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. - */ - -export const REFERENCE_SECTION = ` -# Configuration Architecture - -**Hierarchy:** lifecycle.yaml = SITEMAP referencing other files. Configuration is DISTRIBUTED — follow references for actual values. -- dockerfile → Dockerfiles | helm.valueFiles → Helm values (replicaCount, resources, ports) | helm.chart → Local charts -- valueFiles override inline helm.values - -## Configuration Drift - -Compare config (files) vs actual (K8s): -- Config=0 AND K8s=0 → intentional (working as configured) -- Config>0 AND K8s=0 → MANUAL SCALE DOWN (drift detected) -- Config=X AND K8s=Y → manual override (drift detected) -Report both values, let user decide. - -**Connection failures:** Check both sides. If target service uses port X everywhere but caller's env var says Y → fix caller, not target. Truth = what the target actually uses. - -# Lifecycle Architecture - -Lifecycle creates ephemeral environments from Pull Requests. - -**Source of Truth (ranked):** 1. DB status 2. Config files (Helm values, Dockerfiles) 3. lifecycle.yaml 4. PR comment 5. K8s state 6. Events 7. Logs - -**PR Labels:** Lifecycle requires a deploy label (default: \`lifecycle-deploy!\`) on the PR before it will build and deploy an environment. Without this label, no builds or deploys are created — the environment will have 0 pods and no activity. The \`lifecycle-disabled!\` label explicitly prevents deployment. PR labels are shown in the injected context. If you see 0 pods and no builds/deploys, check whether the deploy label is present. - -**Build vs Deploy:** Builds first, deploys only if ALL builds succeed. Check builds before deploy issues. - -**Build system:** buildPipelineId → Codefresh (get_codefresh_logs) | builderEngine → Native K8s (get_k8s_resources + get_pod_logs) | GITHUB type → no build - -**K8s:** Deployments AND StatefulSets (check both). Labels: lc-service={serviceName}. Namespace: env-{buildUuid}. - -# Multi-Repo Architecture - -Each service's code lives in its OWN repository, shown as "Repo: Owner/name @ branch" in injected context. - -**Single-repo environments:** The PR repo IS the service repo — use the PR repo directly. -**Multi-repo environments:** Query the deploys table to find the correct repo and branch for each service. - -Use \`get_file\` with the service's repo and branch params to read service files. -If a service repo is inaccessible, tell the user explicitly — do not guess at file contents. -Always identify which repo a file is from when referencing it. -Never commit fixes to repos other than the PR repo. -`; diff --git a/src/server/services/ai/prompts/sections/safety.ts b/src/server/services/ai/prompts/sections/safety.ts deleted file mode 100644 index fb8fff9b..00000000 --- a/src/server/services/ai/prompts/sections/safety.ts +++ /dev/null @@ -1,44 +0,0 @@ -/** - * Copyright 2025 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. - */ - -export const SAFETY_SECTION = ` -# Security & Safety - -- **User Consent:** Present findings as JSON first. The user clicks "Fix it for me" before any changes are applied. After applying, show the commit URL. -- **Selected-Issue Scope:** "Fix it for me" approval applies only to the selected issue/target. Do not fix additional services or unrelated findings unless the user explicitly requests each one. -- **Surgical Changes:** Change ONLY what was suggested. No formatting, no cleanup, no unrelated fixes. -- **Content Integrity:** When modifying a file via update_file, use the EXACT content returned by get_file as your starting point. Apply only the specific fix — do not remove comments, reformat whitespace, delete unused sections, or make any other modifications. Your new_content must be identical to the original except for the targeted fix lines. -- **Scope Boundaries:** You can modify files in the PR repo only. For issues in other service repos, describe the fix but do not attempt to commit. -- **Tool-Capability Gate:** Only mark an issue auto-fixable or execute a fix when the matching mutating tool is available (file edits → update_file, PR labels → update_pr_labels, runtime K8s patches → patch_k8s_resource). -- **Path Verification:** Before including a filePath in your JSON response, you must have either (a) read it via get_file, (b) seen it in the injected context, or (c) confirmed it exists via list_directory. -- **No Fabrication:** Never diagnose a root cause you cannot support with a specific error message cited from your tool results. If no tool output contains an error pointing to the problem, say "I don't have enough information to determine the root cause." General knowledge about config structure is not evidence. -- **Ambiguity:** When multiple valid options exist, ask the user to choose. - -## Violation Examples - -WRONG: Agent finds a typo in lifecycle.yaml and also reformats indentation while fixing it. -RIGHT: Agent fixes only the typo. No other changes. - -WRONG: Agent sees canAutoFix=true for a resource limit issue and immediately calls update_file. -RIGHT: Agent outputs JSON with the suggestion. User clicks "Fix it for me". Then agent applies. - -WRONG: Agent fixes a typo on line 42 of lifecycle.yaml but the committed diff also deletes 28 lines of commented-out service configuration at the bottom of the file. -RIGHT: Agent fixes only line 42. All other content — including comments, blank lines, and unused sections — remains byte-for-byte identical to the get_file output. - -# Final Reminder - -Compare States: DESIRED (config) vs ACTUAL (runtime). Intentional or bug? When in doubt, ASK. -`; diff --git a/src/server/services/ai/prompts/systemPrompt.ts b/src/server/services/ai/prompts/systemPrompt.ts deleted file mode 100644 index 0e8accad..00000000 --- a/src/server/services/ai/prompts/systemPrompt.ts +++ /dev/null @@ -1,19 +0,0 @@ -/** - * Copyright 2025 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 { PROMPT_SECTIONS } from './sectionRegistry'; - -export const AI_AGENT_SAFETY_RULES = PROMPT_SECTIONS.find((s) => s.id === 'safety')!.content; diff --git a/src/server/services/ai/prompts/tokenCounter.ts b/src/server/services/ai/prompts/tokenCounter.ts deleted file mode 100644 index 55734b4e..00000000 --- a/src/server/services/ai/prompts/tokenCounter.ts +++ /dev/null @@ -1,94 +0,0 @@ -/** - * Copyright 2025 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 { getEncoding, Tiktoken } from 'js-tiktoken'; -import { PROMPT_SECTIONS } from './sectionRegistry'; - -let _encoder: Tiktoken | null = null; - -function getEncoder(): Tiktoken { - if (!_encoder) { - _encoder = getEncoding('cl100k_base'); - } - return _encoder; -} - -export interface TokenBreakdown { - sections: Record; - providerAugmentation: number; - environmentContext: number; - total: number; -} - -export interface TokenBudget { - provider: string; - limit: number; - used: number; - remaining: number; - overBudget: boolean; -} - -export const PROVIDER_TOKEN_LIMITS: Record = { - anthropic: 180000, - openai: 110000, - gemini: 900000, -}; - -export function countTokens(text: string): number { - if (!text) return 0; - return getEncoder().encode(text).length; -} - -export function countSectionTokens(): Record { - const result: Record = {}; - for (const section of PROMPT_SECTIONS) { - result[section.id] = countTokens(section.content); - } - return result; -} - -export function getTokenBreakdown(systemPrompt: string, sections?: Record): TokenBreakdown { - const sectionCounts = sections || countSectionTokens(); - const sectionTotal = Object.values(sectionCounts).reduce((sum, v) => sum + v, 0); - const totalTokens = countTokens(systemPrompt); - const overhead = Math.max(0, totalTokens - sectionTotal); - - const providerAugmentation = Math.floor(overhead / 2); - const environmentContext = overhead - providerAugmentation; - - return { - sections: sectionCounts, - providerAugmentation, - environmentContext, - total: sectionTotal + providerAugmentation + environmentContext, - }; -} - -export function checkBudget( - systemPrompt: string, - provider: 'anthropic' | 'openai' | 'gemini', - tokenCount?: number -): TokenBudget { - const used = tokenCount ?? countTokens(systemPrompt); - const limit = PROVIDER_TOKEN_LIMITS[provider]; - return { - provider, - limit, - used, - remaining: limit - used, - overBudget: used > limit, - }; -} diff --git a/src/server/services/ai/providers/__tests__/anthropic.test.ts b/src/server/services/ai/providers/__tests__/anthropic.test.ts deleted file mode 100644 index a7fb884f..00000000 --- a/src/server/services/ai/providers/__tests__/anthropic.test.ts +++ /dev/null @@ -1,141 +0,0 @@ -/** - * Copyright 2025 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. - */ - -const mockCreate = jest.fn(); -jest.mock('@anthropic-ai/sdk', () => - jest.fn().mockImplementation(() => ({ - messages: { create: mockCreate }, - })) -); -jest.mock('server/lib/logger', () => ({ getLogger: () => ({ info: jest.fn(), error: jest.fn(), warn: jest.fn() }) })); - -import { AnthropicProvider } from '../anthropic'; -import { ConversationMessage } from '../../types/message'; - -describe('AnthropicProvider.formatHistory', () => { - let provider: AnthropicProvider; - - const TEXT_USER: ConversationMessage = { role: 'user', parts: [{ type: 'text', content: 'What is wrong?' }] }; - const TEXT_ASSISTANT: ConversationMessage = { - role: 'assistant', - parts: [{ type: 'text', content: 'Let me check.' }], - }; - const SYSTEM: ConversationMessage = { role: 'system', parts: [{ type: 'text', content: 'You are a helper.' }] }; - const TOOL_CALL: ConversationMessage = { - role: 'assistant', - parts: [{ type: 'tool_call', toolCallId: 'call-1', name: 'getK8sResources', arguments: { namespace: 'default' } }], - }; - const TOOL_RESULT: ConversationMessage = { - role: 'user', - parts: [ - { - type: 'tool_result', - toolCallId: 'call-1', - name: 'getK8sResources', - result: { success: true, agentContent: '{"pods": []}' }, - }, - ], - }; - - beforeEach(() => { - provider = new AnthropicProvider('test-model', 'test-key'); - }); - - it('formats text messages', () => { - const result = provider.formatHistory([TEXT_USER, TEXT_ASSISTANT]); - expect(result).toEqual([ - { role: 'user', content: 'What is wrong?' }, - { role: 'assistant', content: 'Let me check.' }, - ]); - }); - - it('skips system messages', () => { - const result = provider.formatHistory([SYSTEM, TEXT_USER]); - expect(result).toHaveLength(1); - expect(result[0]).toEqual({ role: 'user', content: 'What is wrong?' }); - }); - - it('maps tool_call to tool_use', () => { - const result = provider.formatHistory([TOOL_CALL]); - expect(result).toEqual([ - { - role: 'assistant', - content: [{ type: 'tool_use', id: 'call-1', name: 'getK8sResources', input: { namespace: 'default' } }], - }, - ]); - }); - - it('maps tool_result with tool_use_id', () => { - const result = provider.formatHistory([TOOL_RESULT]); - expect(result).toEqual([ - { - role: 'user', - content: [{ type: 'tool_result', tool_use_id: 'call-1', content: '{"pods": []}' }], - }, - ]); - }); - - it('formats full conversation skipping system', () => { - const result = provider.formatHistory([SYSTEM, TEXT_USER, TOOL_CALL, TOOL_RESULT, TEXT_ASSISTANT]); - expect(result).toHaveLength(4); - }); - - it('joins multi-part text with space', () => { - const msg: ConversationMessage = { - role: 'user', - parts: [ - { type: 'text', content: 'hello' }, - { type: 'text', content: 'world' }, - ], - }; - const result = provider.formatHistory([msg]); - expect(result).toEqual([{ role: 'user', content: 'hello world' }]); - }); -}); - -describe('AnthropicProvider.streamCompletion cache_control', () => { - let provider: AnthropicProvider; - - beforeEach(() => { - mockCreate.mockReset(); - mockCreate.mockResolvedValue({ content: [{ type: 'text', text: 'response' }] }); - provider = new AnthropicProvider('test-model', 'test-key'); - }); - - it('sends system prompt as content block array with cache_control ephemeral', async () => { - const messages: ConversationMessage[] = [{ role: 'user', parts: [{ type: 'text', content: 'hi' }] }]; - const iter = provider.streamCompletion(messages, { systemPrompt: 'You are helpful.' }); - await iter.next(); - - expect(mockCreate).toHaveBeenCalledWith( - expect.objectContaining({ - system: [{ type: 'text', text: 'You are helpful.', cache_control: { type: 'ephemeral' } }], - }) - ); - }); - - it('sends system as undefined when no systemPrompt provided', async () => { - const messages: ConversationMessage[] = [{ role: 'user', parts: [{ type: 'text', content: 'hi' }] }]; - const iter = provider.streamCompletion(messages, {}); - await iter.next(); - - expect(mockCreate).toHaveBeenCalledWith( - expect.objectContaining({ - system: undefined, - }) - ); - }); -}); diff --git a/src/server/services/ai/providers/__tests__/gemini.test.ts b/src/server/services/ai/providers/__tests__/gemini.test.ts deleted file mode 100644 index 66caac1a..00000000 --- a/src/server/services/ai/providers/__tests__/gemini.test.ts +++ /dev/null @@ -1,129 +0,0 @@ -/** - * Copyright 2025 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('@google/genai', () => ({ GoogleGenAI: jest.fn().mockImplementation(() => ({})) })); -jest.mock('server/lib/logger', () => ({ getLogger: () => ({ info: jest.fn(), error: jest.fn(), warn: jest.fn() }) })); - -import { GeminiProvider } from '../gemini'; -import { ConversationMessage } from '../../types/message'; - -describe('GeminiProvider.formatHistory', () => { - let provider: GeminiProvider; - - const TEXT_USER: ConversationMessage = { role: 'user', parts: [{ type: 'text', content: 'What is wrong?' }] }; - const TEXT_ASSISTANT: ConversationMessage = { - role: 'assistant', - parts: [{ type: 'text', content: 'Let me check.' }], - }; - const SYSTEM: ConversationMessage = { role: 'system', parts: [{ type: 'text', content: 'You are a helper.' }] }; - const TOOL_CALL: ConversationMessage = { - role: 'assistant', - parts: [{ type: 'tool_call', toolCallId: 'call-1', name: 'getK8sResources', arguments: { namespace: 'default' } }], - }; - const TOOL_RESULT: ConversationMessage = { - role: 'user', - parts: [ - { - type: 'tool_result', - toolCallId: 'call-1', - name: 'getK8sResources', - result: { success: true, agentContent: '{"pods": []}' }, - }, - ], - }; - - beforeEach(() => { - provider = new GeminiProvider('test-model', 'test-key'); - }); - - it('formats text messages with assistant mapped to model', () => { - const result = provider.formatHistory([TEXT_USER, TEXT_ASSISTANT]); - expect(result).toEqual([ - { role: 'user', parts: [{ text: 'What is wrong?' }] }, - { role: 'model', parts: [{ text: 'Let me check.' }] }, - ]); - }); - - it('skips system messages', () => { - const result = provider.formatHistory([SYSTEM, TEXT_USER]); - expect(result).toHaveLength(1); - expect(result[0]).toEqual({ role: 'user', parts: [{ text: 'What is wrong?' }] }); - }); - - it('maps tool_call to functionCall', () => { - const result = provider.formatHistory([TOOL_CALL]); - expect(result).toEqual([ - { role: 'model', parts: [{ functionCall: { name: 'getK8sResources', args: { namespace: 'default' } } }] }, - ]); - }); - - it('maps tool_result to functionResponse with JSON.parse', () => { - const result = provider.formatHistory([TOOL_RESULT]); - expect(result).toEqual([ - { - role: 'user', - parts: [{ functionResponse: { name: 'getK8sResources', response: { pods: [] } } }], - }, - ]); - }); - - it('wraps non-JSON agentContent in content key', () => { - const msg: ConversationMessage = { - role: 'user', - parts: [ - { - type: 'tool_result', - toolCallId: 'call-1', - name: 'getK8sResources', - result: { success: true, agentContent: 'plain text' }, - }, - ], - }; - const result = provider.formatHistory([msg]); - expect(result).toEqual([ - { role: 'user', parts: [{ functionResponse: { name: 'getK8sResources', response: { content: 'plain text' } } }] }, - ]); - }); - - it('maps error tool_result with error message', () => { - const msg: ConversationMessage = { - role: 'user', - parts: [ - { - type: 'tool_result', - toolCallId: 'call-1', - name: 'getK8sResources', - result: { - success: false, - error: { message: 'not found', code: 'NOT_FOUND', recoverable: false }, - }, - }, - ], - }; - const result = provider.formatHistory([msg]); - expect(result).toEqual([ - { - role: 'user', - parts: [{ functionResponse: { name: 'getK8sResources', response: { error: 'not found', success: false } } }], - }, - ]); - }); - - it('formats full conversation skipping system', () => { - const result = provider.formatHistory([SYSTEM, TEXT_USER, TOOL_CALL, TOOL_RESULT, TEXT_ASSISTANT]); - expect(result).toHaveLength(4); - }); -}); diff --git a/src/server/services/ai/providers/__tests__/openai.test.ts b/src/server/services/ai/providers/__tests__/openai.test.ts deleted file mode 100644 index 3ad5aeff..00000000 --- a/src/server/services/ai/providers/__tests__/openai.test.ts +++ /dev/null @@ -1,111 +0,0 @@ -/** - * Copyright 2025 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('openai', () => jest.fn().mockImplementation(() => ({}))); -jest.mock('server/lib/logger', () => ({ getLogger: () => ({ info: jest.fn(), error: jest.fn(), warn: jest.fn() }) })); - -import { OpenAIProvider } from '../openai'; -import { ConversationMessage } from '../../types/message'; - -describe('OpenAIProvider.formatHistory', () => { - let provider: OpenAIProvider; - - const TEXT_USER: ConversationMessage = { role: 'user', parts: [{ type: 'text', content: 'What is wrong?' }] }; - const TEXT_ASSISTANT: ConversationMessage = { - role: 'assistant', - parts: [{ type: 'text', content: 'Let me check.' }], - }; - const SYSTEM: ConversationMessage = { role: 'system', parts: [{ type: 'text', content: 'You are a helper.' }] }; - const TOOL_CALL: ConversationMessage = { - role: 'assistant', - parts: [{ type: 'tool_call', toolCallId: 'call-1', name: 'getK8sResources', arguments: { namespace: 'default' } }], - }; - const TOOL_RESULT: ConversationMessage = { - role: 'user', - parts: [ - { - type: 'tool_result', - toolCallId: 'call-1', - name: 'getK8sResources', - result: { success: true, agentContent: '{"pods": []}' }, - }, - ], - }; - - beforeEach(() => { - provider = new OpenAIProvider('test-model', 'test-key'); - }); - - it('formats text messages', () => { - const result = provider.formatHistory([TEXT_USER, TEXT_ASSISTANT]); - expect(result).toEqual([ - { role: 'user', content: 'What is wrong?' }, - { role: 'assistant', content: 'Let me check.' }, - ]); - }); - - it('includes system messages', () => { - const result = provider.formatHistory([SYSTEM, TEXT_USER]); - expect(result).toHaveLength(2); - expect(result[0]).toEqual({ role: 'system', content: 'You are a helper.' }); - }); - - it('formats tool_call with JSON.stringify arguments', () => { - const result = provider.formatHistory([TOOL_CALL]); - expect(result).toEqual([ - { - role: 'assistant', - content: null, - tool_calls: [ - { - id: 'call-1', - type: 'function', - function: { name: 'getK8sResources', arguments: '{"namespace":"default"}' }, - }, - ], - }, - ]); - }); - - it('emits separate tool messages per tool_result part', () => { - const msg: ConversationMessage = { - role: 'user', - parts: [ - { - type: 'tool_result', - toolCallId: 'call-1', - name: 'getK8sResources', - result: { success: true, agentContent: '{"pods": []}' }, - }, - { - type: 'tool_result', - toolCallId: 'call-2', - name: 'queryDatabase', - result: { success: true, agentContent: '{"rows": []}' }, - }, - ], - }; - const result = provider.formatHistory([msg]); - expect(result).toHaveLength(2); - expect(result[0]).toEqual({ role: 'tool', tool_call_id: 'call-1', content: '{"pods": []}' }); - expect(result[1]).toEqual({ role: 'tool', tool_call_id: 'call-2', content: '{"rows": []}' }); - }); - - it('formats full conversation including system', () => { - const result = provider.formatHistory([SYSTEM, TEXT_USER, TOOL_CALL, TOOL_RESULT, TEXT_ASSISTANT]); - expect(result).toHaveLength(5); - }); -}); diff --git a/src/server/services/ai/providers/anthropic.ts b/src/server/services/ai/providers/anthropic.ts deleted file mode 100644 index 235db7d6..00000000 --- a/src/server/services/ai/providers/anthropic.ts +++ /dev/null @@ -1,166 +0,0 @@ -/** - * Copyright 2025 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 Anthropic from '@anthropic-ai/sdk'; -import { - type MessageParam, - type Tool as AnthropicTool, - type ToolUseBlock, -} from '@anthropic-ai/sdk/resources/messages/messages'; -import { BaseLLMProvider } from './base'; -import { ModelInfo, CompletionOptions, StreamChunk } from '../types/provider'; -import { ConversationMessage, TextPart, ToolCallPart, ToolResultPart } from '../types/message'; -import { Tool, ToolCall } from '../types/tool'; - -export class AnthropicProvider extends BaseLLMProvider { - name = 'anthropic'; - private client: Anthropic; - private modelId: string; - - constructor(modelId?: string, apiKey?: string) { - super(); - this.modelId = modelId || 'claude-sonnet-4-5-20250929'; - const key = this.validateApiKey(apiKey, 'Anthropic'); - this.client = new Anthropic({ apiKey: key }); - } - - async *streamCompletion( - messages: ConversationMessage[], - options: CompletionOptions, - signal?: AbortSignal - ): AsyncIterator { - if (signal?.aborted) { - throw new Error('Request aborted'); - } - - const anthropicMessages = this.formatHistory(messages) as MessageParam[]; - - const tools = options.tools?.map((t) => this.formatToolDefinition(t)) as AnthropicTool[] | undefined; - - const response = await this.client.messages.create({ - model: this.modelId, - max_tokens: options.maxTokens || 4096, - system: options.systemPrompt - ? [{ type: 'text' as const, text: options.systemPrompt, cache_control: { type: 'ephemeral' as const } }] - : undefined, - messages: anthropicMessages, - tools: tools || [], - }); - - const toolUseBlocks = response.content.filter((block) => block.type === 'tool_use'); - const textBlocks = response.content.filter((block) => block.type === 'text'); - - for (const block of textBlocks) { - if ('text' in block && block.text) { - yield { - type: 'text', - content: block.text, - }; - } - } - - if (toolUseBlocks.length > 0) { - const toolCalls = this.parseToolCall(toolUseBlocks); - yield { - type: 'tool_call', - toolCalls, - }; - } - - if (response.usage) { - yield { - type: 'text', - usage: { - inputTokens: response.usage.input_tokens, - outputTokens: response.usage.output_tokens, - }, - }; - } - } - - formatHistory(messages: ConversationMessage[]): unknown[] { - const result: unknown[] = []; - - for (const msg of messages) { - if (msg.role === 'system') { - continue; - } - - const toolCallParts = msg.parts.filter((p): p is ToolCallPart => p.type === 'tool_call'); - const toolResultParts = msg.parts.filter((p): p is ToolResultPart => p.type === 'tool_result'); - const textParts = msg.parts.filter((p): p is TextPart => p.type === 'text'); - - if (toolCallParts.length > 0) { - result.push({ - role: 'assistant' as const, - content: toolCallParts.map((p) => ({ - type: 'tool_use', - id: p.toolCallId, - name: p.name, - input: p.arguments, - })), - }); - } else if (toolResultParts.length > 0) { - result.push({ - role: 'user' as const, - content: toolResultParts.map((p) => ({ - type: 'tool_result', - tool_use_id: p.toolCallId, - content: p.result.agentContent || JSON.stringify(p.result), - })), - }); - } else { - const textContent = textParts.map((p) => p.content).join(' '); - result.push({ role: msg.role as 'user' | 'assistant', content: textContent }); - } - } - - return result; - } - - supportsTools(): boolean { - return true; - } - - getModelInfo(): ModelInfo { - return { - model: this.modelId, - maxTokens: 200000, - }; - } - - formatToolDefinition(tool: Tool): unknown { - return { - name: tool.name, - description: tool.description, - input_schema: tool.parameters, - }; - } - - parseToolCall(content: unknown): ToolCall[] { - if (!Array.isArray(content)) { - return []; - } - - return content - .filter((c: ToolUseBlock) => c.type === 'tool_use') - .map((c: ToolUseBlock) => ({ - name: c.name, - arguments: c.input as Record, - id: c.id, - })); - } -} diff --git a/src/server/services/ai/providers/base.ts b/src/server/services/ai/providers/base.ts deleted file mode 100644 index a3538a8b..00000000 --- a/src/server/services/ai/providers/base.ts +++ /dev/null @@ -1,50 +0,0 @@ -/** - * Copyright 2025 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 { LLMProvider, ModelInfo, CompletionOptions, StreamChunk } from '../types/provider'; -import { ConversationMessage } from '../types/message'; -import { Tool, ToolCall } from '../types/tool'; -import { countTokens } from '../prompts/tokenCounter'; - -export abstract class BaseLLMProvider implements LLMProvider { - abstract name: string; - - abstract streamCompletion( - messages: ConversationMessage[], - options: CompletionOptions, - signal?: AbortSignal - ): AsyncIterator; - - abstract supportsTools(): boolean; - abstract getModelInfo(): ModelInfo; - abstract formatToolDefinition(tool: Tool): unknown; - abstract parseToolCall(response: unknown): ToolCall[]; - abstract formatHistory(messages: ConversationMessage[]): unknown[]; - - estimateTokens(text: string): number { - return countTokens(text); - } - - protected validateApiKey(apiKey: string | undefined, providerName: string): string { - if (!apiKey) { - throw new Error( - `${providerName} API key is required. ` + - `Please set the appropriate environment variable (ANTHROPIC_API_KEY, OPENAI_API_KEY, GEMINI_API_KEY, or AI_API_KEY).` - ); - } - return apiKey; - } -} diff --git a/src/server/services/ai/providers/factory.ts b/src/server/services/ai/providers/factory.ts deleted file mode 100644 index 2d81be82..00000000 --- a/src/server/services/ai/providers/factory.ts +++ /dev/null @@ -1,58 +0,0 @@ -/** - * Copyright 2025 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 { LLMProvider } from '../types/provider'; -import { AnthropicProvider } from './anthropic'; -import { OpenAIProvider } from './openai'; -import { GeminiProvider } from './gemini'; - -export type ProviderType = 'anthropic' | 'openai' | 'gemini'; - -export interface ProviderConfig { - provider: ProviderType; - modelId?: string; - apiKey?: string; -} - -export class ProviderFactory { - static create(config: ProviderConfig): LLMProvider { - const apiKey = config.apiKey || this.getDefaultApiKey(config.provider); - - switch (config.provider) { - case 'anthropic': - return new AnthropicProvider(config.modelId, apiKey); - case 'openai': - return new OpenAIProvider(config.modelId, apiKey); - case 'gemini': - return new GeminiProvider(config.modelId, apiKey); - default: - throw new Error(`Unknown provider: ${config.provider}`); - } - } - - private static getDefaultApiKey(provider: ProviderType): string | undefined { - switch (provider) { - case 'anthropic': - return process.env.ANTHROPIC_API_KEY || process.env.AI_API_KEY; - case 'openai': - return process.env.OPENAI_API_KEY || process.env.AI_API_KEY; - case 'gemini': - return process.env.GEMINI_API_KEY || process.env.AI_API_KEY; - default: - return undefined; - } - } -} diff --git a/src/server/services/ai/providers/gemini.ts b/src/server/services/ai/providers/gemini.ts deleted file mode 100644 index 27055473..00000000 --- a/src/server/services/ai/providers/gemini.ts +++ /dev/null @@ -1,301 +0,0 @@ -/** - * Copyright 2025 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 { GoogleGenAI, type Candidate, type Content, type FunctionCall, type FunctionDeclaration } from '@google/genai'; -import { BaseLLMProvider } from './base'; -import { ModelInfo, CompletionOptions, StreamChunk } from '../types/provider'; -import { ConversationMessage, TextPart, ToolCallPart, ToolResultPart } from '../types/message'; -import { Tool, ToolCall } from '../types/tool'; -import { getLogger } from 'server/lib/logger'; -import { ErrorCategory } from '../errors'; - -export class GeminiProvider extends BaseLLMProvider { - name = 'gemini'; - private client: GoogleGenAI; - private modelId: string; - - constructor(modelId?: string, apiKey?: string) { - super(); - this.modelId = modelId || 'gemini-2.5-flash'; - const key = this.validateApiKey(apiKey, 'Gemini'); - this.client = new GoogleGenAI({ apiKey: key }); - } - - async *streamCompletion( - messages: ConversationMessage[], - options: CompletionOptions, - signal?: AbortSignal - ): AsyncIterator { - if (signal?.aborted) { - throw new Error('Request aborted'); - } - - const tools = options.tools?.map((t) => this.formatToolDefinition(t)) as FunctionDeclaration[] | undefined; - - if (tools) { - getLogger().info(`GeminiProvider: sending toolCount=${tools.length} tools=${tools.map((t) => t.name).join(',')}`); - } - - const history = this.formatHistory(messages.slice(0, -1)) as Content[]; - - const isThinkingModel = this.modelId.includes('2.5') || this.modelId.includes('3.'); - - const chat = this.client.chats.create({ - model: this.modelId, - config: { - systemInstruction: options.systemPrompt, - tools: tools ? [{ functionDeclarations: tools }] : undefined, - temperature: options.temperature || 0.1, - topP: 0.95, - ...(isThinkingModel ? {} : { topK: 40 }), - maxOutputTokens: options.maxTokens || 65536, - }, - history, - }); - - const lastMsg = messages[messages.length - 1]; - const toolResultParts = lastMsg?.parts.filter((p): p is ToolResultPart => p.type === 'tool_result') || []; - - let message: string | Array<{ functionResponse: { name: string; response: Record } }>; - if (toolResultParts.length > 0) { - message = toolResultParts.map((part) => { - let responseObj: Record; - if (part.result.success) { - const raw = part.result.agentContent || JSON.stringify(part.result); - try { - responseObj = JSON.parse(raw); - if (Array.isArray(responseObj)) { - responseObj = { items: responseObj }; - } - } catch { - responseObj = { content: raw }; - } - } else { - responseObj = { - error: part.result.error?.message || 'Tool execution failed', - success: false, - }; - } - return { - functionResponse: { - name: part.name, - response: responseObj, - }, - }; - }); - } else { - message = lastMsg - ? lastMsg.parts - .filter((p): p is TextPart => p.type === 'text') - .map((p) => p.content) - .join(' ') - : ''; - } - const stream = await chat.sendMessageStream({ message }); - - let accumulatedText = ''; - const functionCalls: Array = []; - let lastCandidate: Candidate | null = null; - let lastRawChunk: any = null; - - for await (const chunk of stream) { - if (signal?.aborted) { - throw new Error('Request aborted'); - } - - const candidate = chunk.candidates?.[0]; - if (!candidate) { - continue; - } - - lastCandidate = candidate; - lastRawChunk = chunk; - - if (candidate.finishReason === 'STOP' && (!candidate.content?.parts || candidate.content.parts.length === 0)) { - getLogger().error( - `GeminiProvider: returned STOP with no content safetyRatings=${JSON.stringify(candidate.safetyRatings)}` - ); - } - - if (candidate.content?.parts) { - for (const part of candidate.content.parts) { - if ('text' in part && part.text) { - accumulatedText += part.text; - yield { - type: 'text', - content: part.text, - }; - } - - if ('functionCall' in part && part.functionCall) { - functionCalls.push({ - ...part.functionCall, - ...(part.thoughtSignature ? { thoughtSignature: part.thoughtSignature } : {}), - }); - } - } - } - } - - if (lastCandidate?.finishReason === 'MALFORMED_FUNCTION_CALL') { - const error: Error & { category?: ErrorCategory; retryable?: boolean } = new Error( - 'Gemini generated a malformed function call. This is a transient model error.' - ); - error.category = ErrorCategory.TRANSIENT; - error.retryable = true; - throw error; - } - - if (accumulatedText.length === 0 && functionCalls.length === 0) { - getLogger().error(`GeminiProvider: empty response finishReason=${lastCandidate?.finishReason}`); - - const error = new Error( - `Gemini returned an empty response. This may be due to: ` + - `(1) The system prompt being too large (${options.systemPrompt?.length || 0} chars), ` + - `(2) Too many tools (${options.tools?.length || 0}), or ` + - `(3) Incompatible tool definitions. ` + - `finishReason: ${lastCandidate?.finishReason}` - ); - const categorizedError: Error & { category?: ErrorCategory; retryable?: boolean } = error; - if (lastCandidate?.finishReason === 'STOP') { - categorizedError.category = ErrorCategory.AMBIGUOUS; - } else { - categorizedError.category = ErrorCategory.TRANSIENT; - } - categorizedError.retryable = true; - throw categorizedError; - } - - if (functionCalls.length > 0) { - const toolCalls = this.parseToolCall(functionCalls); - yield { - type: 'tool_call', - toolCalls, - }; - } - - if (lastRawChunk?.usageMetadata) { - yield { - type: 'text', - usage: { - inputTokens: lastRawChunk.usageMetadata.promptTokenCount || 0, - outputTokens: lastRawChunk.usageMetadata.candidatesTokenCount || 0, - }, - }; - } - } - - formatHistory(messages: ConversationMessage[]): unknown[] { - const history: unknown[] = []; - - for (const msg of messages) { - if (msg.role === 'system') { - continue; - } - - const toolCallParts = msg.parts.filter((p): p is ToolCallPart => p.type === 'tool_call'); - const toolResultParts = msg.parts.filter((p): p is ToolResultPart => p.type === 'tool_result'); - - if (toolCallParts.length > 0) { - history.push({ - role: 'model' as const, - parts: toolCallParts.map((part) => ({ - functionCall: { name: part.name, args: part.arguments }, - ...(part.metadata?.thoughtSignature ? { thoughtSignature: part.metadata.thoughtSignature as string } : {}), - })), - }); - } else if (toolResultParts.length > 0) { - const responseParts = toolResultParts.map((part) => { - let responseObj: Record; - if (part.result.success) { - const raw = part.result.agentContent || JSON.stringify(part.result); - try { - responseObj = JSON.parse(raw); - if (Array.isArray(responseObj)) { - responseObj = { items: responseObj }; - } - } catch { - responseObj = { content: raw }; - } - } else { - responseObj = { - error: part.result.error?.message || 'Tool execution failed', - success: false, - }; - } - return { functionResponse: { name: part.name, response: responseObj } }; - }); - history.push({ - role: 'user' as const, - parts: responseParts, - }); - } else { - const textContent = msg.parts - .filter((p): p is TextPart => p.type === 'text') - .map((p) => p.content) - .join(' '); - - history.push({ - role: msg.role === 'assistant' ? ('model' as const) : ('user' as const), - parts: [{ text: textContent }], - }); - } - } - - return history; - } - - supportsTools(): boolean { - return true; - } - - getModelInfo(): ModelInfo { - return { - model: this.modelId, - maxTokens: 1000000, - }; - } - - formatToolDefinition(tool: Tool): unknown { - return { - name: tool.name, - description: tool.description, - parametersJsonSchema: { - type: 'object', - properties: tool.parameters.properties || {}, - required: tool.parameters.required || [], - }, - }; - } - - parseToolCall(content: unknown): ToolCall[] { - if (!Array.isArray(content)) { - return []; - } - - return content.map((fc: FunctionCall & { thoughtSignature?: string }) => { - let name = fc.name || ''; - if (name.startsWith('default_api:')) { - name = name.substring('default_api:'.length); - } - return { - name, - arguments: fc.args || {}, - ...(fc.thoughtSignature ? { metadata: { thoughtSignature: fc.thoughtSignature } } : {}), - }; - }); - } -} diff --git a/src/server/services/ai/providers/index.ts b/src/server/services/ai/providers/index.ts deleted file mode 100644 index 48120bd8..00000000 --- a/src/server/services/ai/providers/index.ts +++ /dev/null @@ -1,21 +0,0 @@ -/** - * Copyright 2025 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. - */ - -export * from './base'; -export * from './anthropic'; -export * from './openai'; -export * from './gemini'; -export * from './factory'; diff --git a/src/server/services/ai/providers/openai.ts b/src/server/services/ai/providers/openai.ts deleted file mode 100644 index c64066c9..00000000 --- a/src/server/services/ai/providers/openai.ts +++ /dev/null @@ -1,221 +0,0 @@ -/** - * Copyright 2025 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 OpenAI from 'openai'; -import { type ChatCompletionMessageParam, type ChatCompletionTool } from 'openai/resources/chat/completions'; -import { BaseLLMProvider } from './base'; -import { ModelInfo, CompletionOptions, StreamChunk } from '../types/provider'; -import { ConversationMessage, TextPart, ToolCallPart, ToolResultPart } from '../types/message'; -import { Tool, ToolCall } from '../types/tool'; -import { ErrorCategory } from '../errors'; - -interface AccumulatedToolCall { - id?: string; - type?: string; - function: { - name: string; - arguments: string; - }; -} - -export class OpenAIProvider extends BaseLLMProvider { - name = 'openai'; - private client: OpenAI; - private modelId: string; - - constructor(modelId?: string, apiKey?: string) { - super(); - this.modelId = modelId || 'gpt-4o'; - const key = this.validateApiKey(apiKey, 'OpenAI'); - this.client = new OpenAI({ apiKey: key }); - } - - async *streamCompletion( - messages: ConversationMessage[], - options: CompletionOptions, - signal?: AbortSignal - ): AsyncIterator { - if (signal?.aborted) { - throw new Error('Request aborted'); - } - - const openaiMessages = [ - { role: 'system' as const, content: options.systemPrompt }, - ...(this.formatHistory(messages) as ChatCompletionMessageParam[]), - ]; - - const tools = options.tools?.map((t) => this.formatToolDefinition(t)) as ChatCompletionTool[] | undefined; - - const stream = await this.client.chat.completions.create({ - model: this.modelId, - max_tokens: options.maxTokens || 250000, - stream: true, - stream_options: { include_usage: true }, - messages: openaiMessages, - tools: tools || undefined, - }); - - let accumulatedToolCalls: AccumulatedToolCall[] = []; - let finalUsage: { prompt_tokens: number; completion_tokens: number } | null = null; - - for await (const chunk of stream) { - if (signal?.aborted) { - throw new Error('Request aborted'); - } - - if (chunk.usage) { - finalUsage = { - prompt_tokens: chunk.usage.prompt_tokens, - completion_tokens: chunk.usage.completion_tokens, - }; - } - - const delta = chunk.choices[0]?.delta; - - if (delta?.content) { - yield { - type: 'text', - content: delta.content, - }; - } - - if (delta?.tool_calls) { - for (const toolCallDelta of delta.tool_calls) { - if (!accumulatedToolCalls[toolCallDelta.index]) { - accumulatedToolCalls[toolCallDelta.index] = { - id: toolCallDelta.id, - type: toolCallDelta.type, - function: { - name: toolCallDelta.function?.name || '', - arguments: toolCallDelta.function?.arguments || '', - }, - }; - } else { - if (toolCallDelta.function?.name) { - accumulatedToolCalls[toolCallDelta.index].function.name += toolCallDelta.function.name; - } - if (toolCallDelta.function?.arguments) { - accumulatedToolCalls[toolCallDelta.index].function.arguments += toolCallDelta.function.arguments; - } - } - } - } - } - - if (accumulatedToolCalls.length > 0) { - const toolCalls = this.parseToolCall(accumulatedToolCalls); - yield { - type: 'tool_call', - toolCalls, - }; - } - - if (finalUsage) { - yield { - type: 'text', - usage: { - inputTokens: finalUsage.prompt_tokens, - outputTokens: finalUsage.completion_tokens, - }, - }; - } - } - - formatHistory(messages: ConversationMessage[]): unknown[] { - const result: unknown[] = []; - - for (const msg of messages) { - const toolCallParts = msg.parts.filter((p): p is ToolCallPart => p.type === 'tool_call'); - const toolResultParts = msg.parts.filter((p): p is ToolResultPart => p.type === 'tool_result'); - const textParts = msg.parts.filter((p): p is TextPart => p.type === 'text'); - - if (msg.role === 'system') { - const textContent = textParts.map((p) => p.content).join(' '); - result.push({ role: 'system' as const, content: textContent }); - } else if (toolCallParts.length > 0) { - result.push({ - role: 'assistant' as const, - content: null, - tool_calls: toolCallParts.map((p) => ({ - id: p.toolCallId, - type: 'function', - function: { name: p.name, arguments: JSON.stringify(p.arguments) }, - })), - }); - } else if (toolResultParts.length > 0) { - for (const p of toolResultParts) { - result.push({ - role: 'tool' as const, - tool_call_id: p.toolCallId, - content: p.result.agentContent || JSON.stringify(p.result), - }); - } - } else { - const textContent = textParts.map((p) => p.content).join(' '); - result.push({ role: msg.role as 'user' | 'assistant', content: textContent }); - } - } - - return result; - } - - supportsTools(): boolean { - return true; - } - - getModelInfo(): ModelInfo { - return { - model: this.modelId, - maxTokens: 128000, - }; - } - - formatToolDefinition(tool: Tool): unknown { - return { - type: 'function', - function: { - name: tool.name, - description: tool.description, - parameters: tool.parameters, - }, - }; - } - - parseToolCall(content: unknown): ToolCall[] { - if (!Array.isArray(content)) { - return []; - } - - return content.map((tc: AccumulatedToolCall) => { - let args: Record; - try { - args = JSON.parse(tc.function.arguments); - } catch (parseError) { - const error: Error & { category?: ErrorCategory; retryable?: boolean } = new Error( - `OpenAI returned malformed tool call arguments for ${tc.function.name}: ${parseError}` - ); - error.category = ErrorCategory.TRANSIENT; - error.retryable = true; - throw error; - } - return { - name: tc.function.name, - arguments: args, - id: tc.id, - }; - }); - } -} diff --git a/src/server/services/ai/resilience/__tests__/policies.test.ts b/src/server/services/ai/resilience/__tests__/policies.test.ts deleted file mode 100644 index 257c86de..00000000 --- a/src/server/services/ai/resilience/__tests__/policies.test.ts +++ /dev/null @@ -1,150 +0,0 @@ -/** - * Copyright 2025 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/logger', () => ({ - getLogger: () => ({ - info: jest.fn(), - warn: jest.fn(), - error: jest.fn(), - }), -})); - -import { ErrorCategory } from '../../errors/classification'; -import { RetryBudget } from '../../errors/retryBudget'; - -const mockClassifyError = jest.fn(); -jest.mock('../../errors', () => ({ - ...jest.requireActual('../../errors/classification'), - ...jest.requireActual('../../errors/retryBudget'), - classifyError: (...args: unknown[]) => mockClassifyError(...args), - isRetryable: jest.requireActual('../../errors/classification').isRetryable, - RetryBudget: jest.requireActual('../../errors/retryBudget').RetryBudget, - ErrorCategory: jest.requireActual('../../errors/classification').ErrorCategory, - extractRetryAfter: jest.requireActual('../../errors/providerErrors').extractRetryAfter, -})); - -import { createProviderPolicy } from '../policies'; -import { getProviderCircuitBreaker, resetAllCircuitBreakers } from '../circuitState'; - -beforeEach(() => { - mockClassifyError.mockReset(); - resetAllCircuitBreakers(); -}); - -describe('createProviderPolicy', () => { - it('returns the result of a succeeding function without retry', async () => { - const budget = new RetryBudget(3); - const policy = createProviderPolicy('openai', budget); - const result = await policy.execute(() => 'success'); - expect(result).toBe('success'); - expect(budget.used).toBe(0); - }); - - it('retries on transient error and succeeds on second attempt', async () => { - mockClassifyError.mockReturnValue(ErrorCategory.TRANSIENT); - const budget = new RetryBudget(3); - const policy = createProviderPolicy('openai', budget); - - let callCount = 0; - const result = await policy.execute(() => { - callCount++; - if (callCount === 1) { - throw new Error('transient failure'); - } - return 'recovered'; - }); - - expect(result).toBe('recovered'); - expect(callCount).toBe(2); - expect(budget.used).toBe(1); - }); - - it('does not retry deterministic errors', async () => { - mockClassifyError.mockReturnValue(ErrorCategory.DETERMINISTIC); - const budget = new RetryBudget(3); - const policy = createProviderPolicy('openai', budget); - - let callCount = 0; - await expect( - policy.execute(() => { - callCount++; - throw new Error('bad request'); - }) - ).rejects.toThrow('bad request'); - - expect(callCount).toBe(1); - expect(budget.used).toBe(0); - }); - - it('consumes retry budget on each retry', async () => { - mockClassifyError.mockReturnValue(ErrorCategory.TRANSIENT); - const budget = new RetryBudget(2); - const policy = createProviderPolicy('openai', budget); - - let callCount = 0; - const result = await policy.execute(() => { - callCount++; - if (callCount <= 2) { - throw new Error('transient'); - } - return 'ok'; - }); - - expect(result).toBe('ok'); - expect(budget.used).toBe(2); - expect(budget.exhausted).toBe(true); - }); - - it('stops retrying when budget is exhausted', async () => { - mockClassifyError.mockReturnValue(ErrorCategory.TRANSIENT); - const budget = new RetryBudget(1); - budget.consume(); - expect(budget.exhausted).toBe(true); - - const policy = createProviderPolicy('openai', budget); - - let callCount = 0; - await expect( - policy.execute(() => { - callCount++; - throw new Error('transient'); - }) - ).rejects.toThrow('transient'); - - expect(callCount).toBe(1); - }); -}); - -describe('getProviderCircuitBreaker', () => { - it('returns the same instance for the same provider', () => { - const a = getProviderCircuitBreaker('openai'); - const b = getProviderCircuitBreaker('openai'); - expect(a).toBe(b); - }); - - it('returns different instances for different providers', () => { - const a = getProviderCircuitBreaker('openai'); - const b = getProviderCircuitBreaker('gemini'); - expect(a).not.toBe(b); - }); - - it('clears all instances on resetAllCircuitBreakers', () => { - const before = getProviderCircuitBreaker('openai'); - resetAllCircuitBreakers(); - const after = getProviderCircuitBreaker('openai'); - expect(before).not.toBe(after); - }); -}); diff --git a/src/server/services/ai/resilience/circuitState.ts b/src/server/services/ai/resilience/circuitState.ts deleted file mode 100644 index 687275f6..00000000 --- a/src/server/services/ai/resilience/circuitState.ts +++ /dev/null @@ -1,48 +0,0 @@ -/** - * Copyright 2025 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 { circuitBreaker, handleWhen, ConsecutiveBreaker, CircuitBreakerPolicy } from 'cockatiel'; -import { isRetryable, classifyError } from '../errors'; -import { getLogger } from 'server/lib/logger'; - -const circuitBreakers = new Map(); - -export function getProviderCircuitBreaker(providerName: string): CircuitBreakerPolicy { - const existing = circuitBreakers.get(providerName); - if (existing) return existing; - - const shouldHandle = handleWhen((err) => isRetryable(classifyError(providerName, err))); - - const breaker = circuitBreaker(shouldHandle, { - halfOpenAfter: 30_000, - breaker: new ConsecutiveBreaker(5), - }); - - breaker.onBreak(() => { - getLogger().error(`AI: circuit breaker OPEN provider=${providerName}`); - }); - - breaker.onReset(() => { - getLogger().info(`AI: circuit breaker CLOSED provider=${providerName}`); - }); - - circuitBreakers.set(providerName, breaker); - return breaker; -} - -export function resetAllCircuitBreakers(): void { - circuitBreakers.clear(); -} diff --git a/src/server/services/ai/resilience/index.ts b/src/server/services/ai/resilience/index.ts deleted file mode 100644 index a9df494e..00000000 --- a/src/server/services/ai/resilience/index.ts +++ /dev/null @@ -1,18 +0,0 @@ -/** - * Copyright 2025 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. - */ - -export { createProviderPolicy } from './policies'; -export { getProviderCircuitBreaker, resetAllCircuitBreakers } from './circuitState'; diff --git a/src/server/services/ai/resilience/policies.ts b/src/server/services/ai/resilience/policies.ts deleted file mode 100644 index add512b8..00000000 --- a/src/server/services/ai/resilience/policies.ts +++ /dev/null @@ -1,96 +0,0 @@ -/** - * Copyright 2025 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 { - retry, - handleWhen, - wrap, - ExponentialBackoff, - type IBackoffFactory, - type IBackoff, - type IRetryBackoffContext, -} from 'cockatiel'; -import { isRetryable, classifyError, RetryBudget, extractRetryAfter } from '../errors'; -import { getProviderCircuitBreaker } from './circuitState'; -import { getLogger } from 'server/lib/logger'; - -class RetryAfterBackoffInstance implements IBackoff> { - readonly duration: number; - private readonly fallback: IBackoff; - - constructor(duration: number, fallback: IBackoff) { - this.duration = duration; - this.fallback = fallback; - } - - next(context: IRetryBackoffContext): IBackoff> { - const error = 'error' in context.result ? context.result.error : undefined; - const retryAfterSeconds = error != null ? extractRetryAfter(error) : null; - const nextFallback = this.fallback.next(context as any); - if (retryAfterSeconds != null && retryAfterSeconds > 0) { - return new RetryAfterBackoffInstance(retryAfterSeconds * 1000, nextFallback); - } - return new RetryAfterBackoffInstance(nextFallback.duration, nextFallback); - } -} - -export class RetryAfterBackoff implements IBackoffFactory> { - private readonly fallbackFactory: ExponentialBackoff; - - constructor(fallbackOptions?: { initialDelay?: number; maxDelay?: number; exponent?: number }) { - this.fallbackFactory = new ExponentialBackoff({ - initialDelay: fallbackOptions?.initialDelay ?? 500, - maxDelay: fallbackOptions?.maxDelay ?? 10_000, - exponent: fallbackOptions?.exponent ?? 2, - }); - } - - next(context: IRetryBackoffContext): IBackoff> { - const error = 'error' in context.result ? context.result.error : undefined; - const retryAfterSeconds = error != null ? extractRetryAfter(error) : null; - const fallbackBackoff = this.fallbackFactory.next(context as any); - if (retryAfterSeconds != null && retryAfterSeconds > 0) { - return new RetryAfterBackoffInstance(retryAfterSeconds * 1000, fallbackBackoff); - } - return new RetryAfterBackoffInstance(fallbackBackoff.duration, fallbackBackoff); - } -} - -export function createProviderPolicy(providerName: string, retryBudget: RetryBudget) { - const shouldHandle = handleWhen((err) => { - if (!retryBudget.canRetry()) return false; - return isRetryable(classifyError(providerName, err)); - }); - - const retryPolicy = retry(shouldHandle, { - maxAttempts: 3, - backoff: new RetryAfterBackoff({ initialDelay: 500, maxDelay: 10_000, exponent: 2 }), - }); - - retryPolicy.onRetry((reason) => { - retryBudget.consume(); - const errorMessage = 'error' in reason ? reason.error.message : 'unknown'; - getLogger().warn( - `AI: retrying provider=${providerName} error=${errorMessage} budgetRemaining=${ - retryBudget.canRetry() ? 'yes' : 'exhausted' - }` - ); - }); - - const breakerPolicy = getProviderCircuitBreaker(providerName); - - return wrap(retryPolicy, breakerPolicy); -} diff --git a/src/server/services/ai/service.ts b/src/server/services/ai/service.ts deleted file mode 100644 index e901d4b2..00000000 --- a/src/server/services/ai/service.ts +++ /dev/null @@ -1,450 +0,0 @@ -/** - * Copyright 2025 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 { LLMProvider } from './types/provider'; -import { ConversationMessage, textMessage, extractTextFromParts } from './types/message'; -import { ProviderFactory, ProviderType } from './providers/factory'; -import { ToolRegistry } from './tools/registry'; -import { ToolOrchestrator } from './orchestration/orchestrator'; -import { maskObservations } from './orchestration/observationMasker'; -import { ToolSafetyManager } from './orchestration/safety'; -import { ConversationManager } from './conversation/manager'; -import { StreamCallbacks } from './types/stream'; -import { ResponseHandler } from './streaming/responseHandler'; -import { AIAgentPromptBuilder } from './prompts/builder'; -import { - GetK8sResourcesTool, - GetPodLogsTool, - GetLifecycleLogsTool, - PatchK8sResourceTool, - QueryDatabaseTool, - GetFileTool, - UpdateFileTool, - UpdatePrLabelsTool, - ListDirectoryTool, - GetIssueCommentTool, - GetCodefreshLogsTool, - K8sClient, - DatabaseClient, - GitHubClient, -} from './tools'; -import { DebugContext, DebugMessage } from '../types/aiAgent'; -import { getLogger } from 'server/lib/logger'; -import { createMcpTools } from './mcp/toolAdapter'; -import { McpConfigService } from './mcp/config'; -import { McpToolInfo } from './prompts/builder'; -import { ResolvedMcpServer } from './mcp/types'; - -export interface AIAgentConfig { - provider: ProviderType; - modelId?: string; - db: any; - redis: any; - requireToolConfirmation?: boolean; - mode?: 'investigate' | 'fix'; - additiveRules?: string[]; - systemPromptOverride?: string; - excludedTools?: string[]; - excludedFilePatterns?: string[]; - allowedWritePatterns?: string[]; - modelPricing?: { - inputCostPerMillion: number; - outputCostPerMillion: number; - }; - maxIterations?: number; - maxToolCalls?: number; - maxRepeatedCalls?: number; - compressionThreshold?: number; - observationMaskingRecencyWindow?: number; - observationMaskingTokenThreshold?: number; - toolExecutionTimeout?: number; - toolOutputMaxChars?: number; - retryBudget?: number; -} - -export interface ProcessQueryResult { - response: string; - isJson: boolean; - preamble?: string; - availableTools: Array<{ - name: string; - description: string; - category: string; - safetyLevel: string; - }>; - metrics: { - iterations: number; - toolCalls: number; - duration: number; - inputTokens: number; - outputTokens: number; - }; -} - -export class AIAgentCore { - private provider: LLMProvider; - private toolRegistry: ToolRegistry; - private orchestrator: ToolOrchestrator; - private promptBuilder: AIAgentPromptBuilder; - private conversationManager: ConversationManager; - private mode: 'investigate' | 'fix'; - private additiveRules?: string[]; - private systemPromptOverride?: string; - private excludedTools?: string[]; - private excludedFilePatterns?: string[]; - private allowedWritePatterns?: string[]; - private modelPricing?: { inputCostPerMillion: number; outputCostPerMillion: number }; - private observationMaskingOptions?: { recencyWindow?: number; tokenThreshold?: number }; - - private mcpToolsLoaded = false; - private mcpToolInfos: McpToolInfo[] = []; - - private k8sClient: K8sClient; - private databaseClient: DatabaseClient; - private githubClient: GitHubClient; - - constructor(config: AIAgentConfig) { - this.k8sClient = new K8sClient(); - this.databaseClient = new DatabaseClient(config.db); - this.githubClient = new GitHubClient(); - - this.provider = ProviderFactory.create({ - provider: config.provider, - modelId: config.modelId, - }); - this.mode = config.mode || 'investigate'; - this.additiveRules = config.additiveRules; - this.systemPromptOverride = config.systemPromptOverride; - this.excludedTools = config.excludedTools; - this.excludedFilePatterns = config.excludedFilePatterns; - this.allowedWritePatterns = config.allowedWritePatterns; - this.modelPricing = config.modelPricing; - if (config.observationMaskingRecencyWindow || config.observationMaskingTokenThreshold) { - this.observationMaskingOptions = { - recencyWindow: config.observationMaskingRecencyWindow, - tokenThreshold: config.observationMaskingTokenThreshold, - }; - } - - this.toolRegistry = new ToolRegistry(); - this.registerAllTools(); - - const safetyManager = new ToolSafetyManager( - config.requireToolConfirmation ?? true, - config.toolExecutionTimeout, - config.toolOutputMaxChars - ); - this.orchestrator = new ToolOrchestrator(this.toolRegistry, safetyManager, { - loopProtection: { - maxIterations: config.maxIterations, - maxToolCalls: config.maxToolCalls, - maxRepeatedCalls: config.maxRepeatedCalls, - }, - retryBudget: config.retryBudget, - }); - - this.promptBuilder = new AIAgentPromptBuilder(); - this.conversationManager = new ConversationManager(config.compressionThreshold); - } - - async processQuery( - userMessage: string, - context: DebugContext, - conversationHistory: DebugMessage[], - callbacks: StreamCallbacks, - signal: AbortSignal, - onDebugContext?: (event: any) => void - ): Promise { - const startTime = Date.now(); - - if (!this.mcpToolsLoaded) { - this.mcpToolsLoaded = true; - try { - const repoFullName = context.lifecycleContext?.pullRequest?.fullName; - if (repoFullName) { - const mcpConfig = new McpConfigService(); - const servers = await mcpConfig.resolveServersForRepo(repoFullName); - let mcpTools = createMcpTools(servers); - if (this.excludedTools && this.excludedTools.length > 0) { - mcpTools = mcpTools.filter((tool) => !this.excludedTools!.includes(tool.name)); - } - for (const tool of mcpTools) { - this.toolRegistry.register(tool); - } - this.mcpToolInfos = this.buildMcpToolInfos(servers); - if (mcpTools.length > 0) { - getLogger().info(`AIAgentCore: registered MCP tools count=${mcpTools.length} repo=${repoFullName}`); - } - } - } catch (err) { - getLogger().warn( - 'AIAgentCore: MCP tool resolution failed, continuing with built-in tools only: error=' + - (err instanceof Error ? err.message : String(err)) - ); - } - } - - try { - if (context.lifecycleContext.pullRequest.branch) { - this.githubClient.setAllowedBranch(context.lifecycleContext.pullRequest.branch); - } - - if (context.lifecycleYaml?.content) { - const referencedFiles = this.githubClient.extractReferencedFilesFromYaml(context.lifecycleYaml.content); - referencedFiles.push('lifecycle.yaml'); - referencedFiles.push('lifecycle.yml'); - this.githubClient.setReferencedFiles(referencedFiles); - } - - if (this.excludedFilePatterns && this.excludedFilePatterns.length > 0) { - this.githubClient.setExcludedFilePatterns(this.excludedFilePatterns); - } - - if (this.allowedWritePatterns && this.allowedWritePatterns.length > 0) { - this.githubClient.setAllowedWritePatterns(this.allowedWritePatterns); - } - - const messages: ConversationMessage[] = conversationHistory.map((m) => - textMessage(m.role as 'user' | 'assistant', m.content) - ); - - const maskResult = maskObservations(messages, this.observationMaskingOptions); - if (maskResult.masked) { - getLogger().info( - `AIAgentCore: observation masking applied maskedParts=${maskResult.stats.maskedParts} savedTokens=${maskResult.stats.savedTokens} buildUuid=${context.buildUuid}` - ); - messages.splice(0, messages.length, ...maskResult.messages); - } - - if (await this.conversationManager.shouldCompress(messages)) { - getLogger().info( - `AIAgentCore: compressing conversation fromMessageCount=${messages.length} buildUuid=${context.buildUuid}` - ); - const state = await this.conversationManager.compress(messages, this.provider, context.buildUuid); - messages.splice(0, messages.length - 1); - messages.unshift(textMessage('user', this.conversationManager.buildPromptFromState(state))); - getLogger().info( - `AIAgentCore: conversation compressed toMessageCount=${messages.length} buildUuid=${context.buildUuid}` - ); - } - - const conversationHistoryForBuilder: DebugMessage[] = messages.map((m) => ({ - role: m.role as 'user' | 'assistant', - content: extractTextFromParts(m.parts), - timestamp: Date.now(), - })); - - const prompt = this.promptBuilder.build({ - provider: this.provider.name as ProviderType, - debugContext: context, - conversationHistory: conversationHistoryForBuilder, - userMessage, - additiveRules: this.additiveRules, - systemPromptOverride: this.systemPromptOverride, - excludedTools: this.excludedTools, - excludedFilePatterns: this.excludedFilePatterns, - mcpTools: this.mcpToolInfos, - }); - - onDebugContext?.({ - type: 'debug_context', - systemPrompt: prompt.systemPrompt, - maskingStats: maskResult.masked ? maskResult.stats : null, - provider: this.provider.name, - modelId: this.provider.getModelInfo()?.model || '', - }); - - const responseHandler = new ResponseHandler(callbacks, context.buildUuid); - - const enhancedCallbacks: StreamCallbacks = { - ...callbacks, - onTextChunk: (text) => { - responseHandler.handleChunk(text); - }, - }; - - const messagesForOrchestrator: ConversationMessage[] = prompt.messages.map((m) => - textMessage(m.role as 'user' | 'assistant', m.content) - ); - - const result = await this.orchestrator.executeToolLoop( - this.provider, - prompt.systemPrompt, - messagesForOrchestrator, - this.toolRegistry.getAll(), - enhancedCallbacks, - signal, - context.buildUuid - ); - - const finalResult = responseHandler.getResult(); - - const duration = Date.now() - startTime; - - onDebugContext?.({ - type: 'debug_metrics', - iterations: result.metrics.iterations, - totalToolCalls: result.metrics.toolCalls, - totalDurationMs: duration, - inputTokens: result.metrics.inputTokens, - outputTokens: result.metrics.outputTokens, - ...(this.modelPricing - ? { - inputCostPerMillion: this.modelPricing.inputCostPerMillion, - outputCostPerMillion: this.modelPricing.outputCostPerMillion, - } - : {}), - }); - - getLogger().info( - `AIAgentCore: query processing ${result.success ? 'completed' : 'failed'} iterations=${ - result.metrics.iterations - } toolCalls=${result.metrics.toolCalls} duration=${duration}ms isJson=${finalResult.isJson} buildUuid=${ - context.buildUuid - }` - ); - - if (!result.success && result.classifiedError) { - throw result.classifiedError.original; - } - - return { - response: result.response || result.error || finalResult.response, - isJson: finalResult.isJson, - preamble: finalResult.preamble, - availableTools: this.getRegisteredTools(), - metrics: result.metrics, - }; - } catch (error: any) { - const duration = Date.now() - startTime; - - getLogger().error( - `AIAgentCore: query processing error duration=${duration}ms error=${error?.message} buildUuid=${context.buildUuid}` - ); - - throw error; - } - } - - private registerAllTools(): void { - const k8sTools = [ - new GetK8sResourcesTool(this.k8sClient), - new GetPodLogsTool(this.k8sClient), - new GetLifecycleLogsTool(this.k8sClient), - new PatchK8sResourceTool(this.k8sClient), - new QueryDatabaseTool(this.databaseClient), - ]; - - const githubTools = [ - new GetFileTool(this.githubClient), - new UpdateFileTool(this.githubClient), - new UpdatePrLabelsTool(this.githubClient), - new ListDirectoryTool(this.githubClient), - new GetIssueCommentTool(this.githubClient), - ]; - - const codefreshTools = [new GetCodefreshLogsTool()]; - - let allTools = [...k8sTools, ...githubTools, ...codefreshTools]; - - if (this.mode === 'investigate') { - const writingToolNames = [UpdateFileTool.Name, UpdatePrLabelsTool.Name, PatchK8sResourceTool.Name]; - allTools = allTools.filter((tool) => !writingToolNames.includes(tool.name)); - } else if (this.mode === 'fix') { - const directWriteTools = [UpdateFileTool.Name, PatchK8sResourceTool.Name]; - allTools = allTools.filter((tool) => !directWriteTools.includes(tool.name)); - } - - if (this.excludedTools && this.excludedTools.length > 0) { - allTools = allTools.filter((tool) => !this.excludedTools!.includes(tool.name)); - } - - this.toolRegistry.registerMultiple(allTools); - } - - getProviderName(): string { - return this.provider.name; - } - - getModelInfo() { - return this.provider.getModelInfo(); - } - - getRegisteredTools(): Array<{ - name: string; - description: string; - category: string; - safetyLevel: string; - }> { - const tools = this.toolRegistry.getAll().map((tool) => ({ - name: tool.name, - description: tool.description, - category: tool.category, - safetyLevel: tool.safetyLevel, - })); - - // In investigate mode, write tools are intentionally hidden from execution. - // Add them back as "fix-mode available" capabilities for canAutoFix gating. - if (this.mode === 'investigate') { - const existing = new Set(tools.map((tool) => tool.name)); - const excluded = new Set(this.excludedTools || []); - const fixModeToolHints: typeof tools = [ - { - name: UpdateFileTool.Name, - description: 'Update repository files (available in fix mode)', - category: 'github', - safetyLevel: 'dangerous', - }, - { - name: PatchK8sResourceTool.Name, - description: 'Patch Kubernetes resources (available in fix mode)', - category: 'k8s', - safetyLevel: 'dangerous', - }, - { - name: UpdatePrLabelsTool.Name, - description: 'Update pull request labels in GitHub (available in fix mode)', - category: 'github', - safetyLevel: 'dangerous', - }, - ]; - - for (const tool of fixModeToolHints) { - if (!existing.has(tool.name) && !excluded.has(tool.name)) { - tools.push(tool); - } - } - } - - return tools; - } - - private buildMcpToolInfos(servers: ResolvedMcpServer[]): McpToolInfo[] { - const infos: McpToolInfo[] = []; - for (const server of servers) { - for (const tool of server.discoveredTools) { - infos.push({ - serverName: server.name, - serverSlug: server.slug, - toolName: tool.name, - qualifiedName: `mcp__${server.slug}__${tool.name}`, - description: tool.description || tool.name, - }); - } - } - return infos; - } -} diff --git a/src/server/services/ai/streaming/__tests__/completeJsonEmission.test.ts b/src/server/services/ai/streaming/__tests__/completeJsonEmission.test.ts deleted file mode 100644 index 619b8499..00000000 --- a/src/server/services/ai/streaming/__tests__/completeJsonEmission.test.ts +++ /dev/null @@ -1,173 +0,0 @@ -/** - * Copyright 2025 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/logger', () => ({ - getLogger: () => ({ - info: jest.fn(), - warn: jest.fn(), - error: jest.fn(), - }), -})); - -import { ResponseHandler } from '../responseHandler'; -import { extractJsonFromResponse } from '../../utils/jsonExtraction'; - -function simulateRoutePostProcessing(result: { response: string; isJson: boolean; preamble?: string }) { - const events: Array<{ type: string; content?: string; preamble?: string }> = []; - let aiResponse = result.response; - let isJsonResponse = result.isJson; - let preambleText: string | undefined = result.preamble; - let completeJsonEmitted = false; - - if (aiResponse.includes('"investigation_complete"')) { - const extracted = extractJsonFromResponse(aiResponse, 'test-uuid'); - if (extracted.isJson) { - aiResponse = extracted.response; - isJsonResponse = true; - if (extracted.preamble && !preambleText) { - preambleText = extracted.preamble; - } - } - } - - if (isJsonResponse) { - try { - JSON.parse(aiResponse); - if (!completeJsonEmitted) { - completeJsonEmitted = true; - events.push({ - type: 'complete_json', - content: aiResponse, - ...(preambleText ? { preamble: preambleText } : {}), - }); - } - } catch { - isJsonResponse = false; - } - } - - events.push({ type: 'complete' }); - - return { events, aiResponse, isJsonResponse }; -} - -function createHandler() { - return new ResponseHandler( - { - onThinking: jest.fn(), - onTextChunk: jest.fn(), - onToolCall: jest.fn(), - onToolResult: jest.fn(), - onActivity: jest.fn(), - onError: jest.fn(), - }, - 'test-uuid' - ); -} - -describe('complete_json emission', () => { - it('emits complete_json for direct JSON investigation response', () => { - const handler = createHandler(); - const json = '{"type": "investigation_complete", "summary": "done", "services": []}'; - handler.handleChunk(json); - const { events, isJsonResponse } = simulateRoutePostProcessing(handler.getResult()); - - expect(isJsonResponse).toBe(true); - expect(events[0].type).toBe('complete_json'); - expect(events[1].type).toBe('complete'); - expect(JSON.parse(events[0].content!).type).toBe('investigation_complete'); - }); - - it('emits complete_json for fenced JSON investigation response', () => { - const handler = createHandler(); - handler.handleChunk('```json\n{"type": "investigation_complete", "data": []}\n```'); - const { events, isJsonResponse } = simulateRoutePostProcessing(handler.getResult()); - - expect(isJsonResponse).toBe(true); - expect(events[0].type).toBe('complete_json'); - expect(events[1].type).toBe('complete'); - expect(JSON.parse(events[0].content!).type).toBe('investigation_complete'); - }); - - it('emits complete_json for preamble + fenced JSON', () => { - const handler = createHandler(); - handler.handleChunk('Here are the findings:\n\n```json\n{"type": "investigation_complete", "items": []}\n```'); - const { events, isJsonResponse } = simulateRoutePostProcessing(handler.getResult()); - - expect(isJsonResponse).toBe(true); - expect(events[0].type).toBe('complete_json'); - expect(events[1].type).toBe('complete'); - expect(JSON.parse(events[0].content!).type).toBe('investigation_complete'); - }); - - it('includes preamble in complete_json when present', () => { - const handler = createHandler(); - handler.handleChunk('Here are the findings:\n\n```json\n{"type": "investigation_complete", "items": []}\n```'); - const { events } = simulateRoutePostProcessing(handler.getResult()); - - expect(events[0].type).toBe('complete_json'); - expect(events[0].preamble).toBe('Here are the findings:'); - }); - - it('does not include preamble for pure JSON responses', () => { - const handler = createHandler(); - handler.handleChunk('{"type": "investigation_complete", "summary": "done"}'); - const { events } = simulateRoutePostProcessing(handler.getResult()); - - expect(events[0].type).toBe('complete_json'); - expect(events[0].preamble).toBeUndefined(); - }); - - it('content field is valid JSON in complete_json event', () => { - const handler = createHandler(); - const json = '{"type": "investigation_complete", "nested": {"key": [1, 2]}}'; - handler.handleChunk(json); - const { events } = simulateRoutePostProcessing(handler.getResult()); - - const parsed = JSON.parse(events[0].content!); - expect(parsed.type).toBe('investigation_complete'); - expect(parsed.nested.key).toEqual([1, 2]); - }); - - it('complete_json is emitted before complete', () => { - const handler = createHandler(); - handler.handleChunk('{"type": "investigation_complete"}'); - const { events } = simulateRoutePostProcessing(handler.getResult()); - - const jsonIdx = events.findIndex((e) => e.type === 'complete_json'); - const completeIdx = events.findIndex((e) => e.type === 'complete'); - expect(jsonIdx).toBeLessThan(completeIdx); - }); - - it('emits complete_json at most once', () => { - const handler = createHandler(); - handler.handleChunk('{"type": "investigation_complete", "summary": "done"}'); - const { events } = simulateRoutePostProcessing(handler.getResult()); - - const jsonEvents = events.filter((e) => e.type === 'complete_json'); - expect(jsonEvents).toHaveLength(1); - }); - - it('does not emit complete_json for plain text responses', () => { - const handler = createHandler(); - handler.handleChunk('Just a regular chat message'); - const { events, isJsonResponse } = simulateRoutePostProcessing(handler.getResult()); - - expect(isJsonResponse).toBe(false); - expect(events).toHaveLength(1); - expect(events[0].type).toBe('complete'); - }); -}); diff --git a/src/server/services/ai/streaming/__tests__/responseHandler.test.ts b/src/server/services/ai/streaming/__tests__/responseHandler.test.ts deleted file mode 100644 index b30c2f76..00000000 --- a/src/server/services/ai/streaming/__tests__/responseHandler.test.ts +++ /dev/null @@ -1,192 +0,0 @@ -/** - * Copyright 2025 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/logger', () => ({ - getLogger: () => ({ - info: jest.fn(), - warn: jest.fn(), - error: jest.fn(), - }), -})); - -import { ResponseHandler } from '../responseHandler'; - -describe('ResponseHandler', () => { - let handler: ResponseHandler; - let onThinking: jest.Mock; - let onTextChunk: jest.Mock; - - beforeEach(() => { - onThinking = jest.fn(); - onTextChunk = jest.fn(); - handler = new ResponseHandler( - { - onThinking, - onTextChunk, - onToolCall: jest.fn(), - onToolResult: jest.fn(), - onActivity: jest.fn(), - onError: jest.fn(), - }, - 'test-uuid' - ); - }); - - it('returns plain text for non-JSON responses', () => { - handler.handleChunk('Hello world'); - const result = handler.getResult(); - expect(result).toEqual({ response: 'Hello world', isJson: false }); - expect(onTextChunk).toHaveBeenCalledWith('Hello world'); - }); - - it('detects JSON response when text starts with { and contains type', () => { - handler.handleChunk('{"type": "investigation_complete", "summary": "done"}'); - const result = handler.getResult(); - expect(result.isJson).toBe(true); - expect(onThinking).toHaveBeenCalledWith('Generating structured report...'); - }); - - it('handles multi-chunk JSON', () => { - handler.handleChunk('{"type": "invest'); - handler.handleChunk('igation_complete"}'); - const result = handler.getResult(); - expect(result.isJson).toBe(true); - expect(result.response).toContain('{"type": "invest'); - expect(result.response).toContain('igation_complete"}'); - }); - - it('detects JSON when first chunk is buffered then second triggers detection', () => { - handler.handleChunk('{'); - handler.handleChunk('"type": "x"}'); - const result = handler.getResult(); - expect(result.isJson).toBe(true); - }); - - it('keeps plain text as non-JSON', () => { - handler.handleChunk('Just a regular message'); - const result = handler.getResult(); - expect(result.isJson).toBe(false); - }); - - it('calls onTextChunk for plain text chunks', () => { - handler.handleChunk('one'); - handler.handleChunk('two'); - handler.handleChunk('three'); - expect(onTextChunk).toHaveBeenCalledTimes(3); - }); - - it('detects markdown-fenced JSON', () => { - handler.handleChunk('```json\n{"type": "investigation_complete", "summary": "done"}\n```'); - const result = handler.getResult(); - expect(result.isJson).toBe(true); - expect(JSON.parse(result.response)).toEqual({ type: 'investigation_complete', summary: 'done' }); - }); - - it('detects preamble text + fenced JSON', () => { - handler.handleChunk('Here are the findings:\n\n```json\n{"type": "investigation_complete", "data": []}\n```'); - const result = handler.getResult(); - expect(result.isJson).toBe(true); - expect(JSON.parse(result.response)).toEqual({ type: 'investigation_complete', data: [] }); - }); - - it('detects fenced JSON without json language tag', () => { - handler.handleChunk('```\n{"type": "investigation_complete"}\n```'); - const result = handler.getResult(); - expect(result.isJson).toBe(true); - expect(JSON.parse(result.response)).toEqual({ type: 'investigation_complete' }); - }); - - it('strips trailing fence markers from getResult', () => { - handler.handleChunk('```json\n{"type": "x"}'); - handler.handleChunk('\n```'); - const result = handler.getResult(); - expect(result.isJson).toBe(true); - expect(result.response).not.toContain('```'); - }); - - describe('JSON chunk suppression', () => { - it('does not call onTextChunk for JSON content chunks', () => { - handler.handleChunk('{"type": "investigation_complete", "summary": "done"}'); - expect(onTextChunk).not.toHaveBeenCalled(); - }); - - it('does not leak partial JSON prefix when type arrives in later chunk', () => { - handler.handleChunk('{'); - handler.handleChunk('"type": "investigation_complete"}'); - expect(onTextChunk).not.toHaveBeenCalled(); - }); - - it('does not call onTextChunk for subsequent JSON chunks', () => { - handler.handleChunk('{"type": "invest'); - handler.handleChunk('igation_complete"}'); - expect(onTextChunk).not.toHaveBeenCalled(); - }); - - it('does not call onTextChunk for fenced JSON', () => { - handler.handleChunk('```json\n{"type": "investigation_complete"}\n```'); - expect(onTextChunk).not.toHaveBeenCalled(); - }); - - it('sends preamble via onTextChunk but suppresses JSON', () => { - handler.handleChunk('Here is my analysis:\n\n```json\n{"type": "investigation_complete"}\n```'); - const calls = onTextChunk.mock.calls.map((c: any[]) => c[0]); - const allText = calls.join(''); - expect(allText).not.toContain('"investigation_complete"'); - expect(allText).not.toContain('{'); - }); - - it('emits plain-text preamble and suppresses split raw JSON tail', () => { - handler.handleChunk('Analysis complete.\n{'); - handler.handleChunk('"type": "investigation_complete", "services": []}'); - - const calls = onTextChunk.mock.calls.map((c: any[]) => c[0]); - const allText = calls.join(''); - expect(allText).toContain('Analysis complete.'); - expect(allText).not.toContain('"investigation_complete"'); - expect(allText).not.toContain('{'); - }); - }); - - describe('preamble tracking', () => { - it('returns preamble for mixed text+JSON responses', () => { - handler.handleChunk('Here are the findings:\n\n```json\n{"type": "investigation_complete", "data": []}\n```'); - const result = handler.getResult(); - expect(result.isJson).toBe(true); - expect(result.preamble).toBe('Here are the findings:'); - }); - - it('does not return preamble for pure JSON responses', () => { - handler.handleChunk('{"type": "investigation_complete", "summary": "done"}'); - const result = handler.getResult(); - expect(result.isJson).toBe(true); - expect(result.preamble).toBeUndefined(); - }); - - it('does not return preamble for plain text', () => { - handler.handleChunk('Just a regular message'); - const result = handler.getResult(); - expect(result.isJson).toBe(false); - expect(result.preamble).toBeUndefined(); - }); - - it('does not return preamble for fenced JSON without preamble text', () => { - handler.handleChunk('```json\n{"type": "investigation_complete"}\n```'); - const result = handler.getResult(); - expect(result.isJson).toBe(true); - expect(result.preamble).toBeUndefined(); - }); - }); -}); diff --git a/src/server/services/ai/streaming/jsonBuffer.ts b/src/server/services/ai/streaming/jsonBuffer.ts deleted file mode 100644 index 71bfa937..00000000 --- a/src/server/services/ai/streaming/jsonBuffer.ts +++ /dev/null @@ -1,56 +0,0 @@ -/** - * Copyright 2025 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'; - -export class JSONBuffer { - private buffer: string = ''; - private complete: boolean = false; - - append(text: string): void { - this.buffer += text; - - if (this.buffer.trim().endsWith('}')) { - const openBraces = (this.buffer.match(/{/g) || []).length; - const closeBraces = (this.buffer.match(/}/g) || []).length; - - if (openBraces === closeBraces) { - this.complete = true; - } - } - } - - isComplete(): boolean { - return this.complete; - } - - parse(): any | null { - if (!this.complete) { - return null; - } - - try { - return JSON.parse(this.buffer); - } catch (error: any) { - getLogger().error({ error }, `JSONBuffer: parse failed bufferLength=${this.buffer.length}`); - return null; - } - } - - getContent(): string { - return this.buffer; - } -} diff --git a/src/server/services/ai/streaming/responseHandler.ts b/src/server/services/ai/streaming/responseHandler.ts deleted file mode 100644 index 65430719..00000000 --- a/src/server/services/ai/streaming/responseHandler.ts +++ /dev/null @@ -1,173 +0,0 @@ -/** - * Copyright 2025 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 { StreamCallbacks } from '../types/stream'; -import { JSONBuffer } from './jsonBuffer'; -import { getLogger } from 'server/lib/logger'; - -function stripCodeFence(text: string): string { - return text.replace(/```(?:json)?\s*\n?/g, '').replace(/\n?\s*```/g, ''); -} - -function cleanPreamble(text: string): string { - return stripCodeFence(text).trim(); -} - -export class ResponseHandler { - private jsonBuffer: JSONBuffer; - private isJsonResponse: boolean = false; - private textBuffer: string = ''; - private plainTextResponse: string = ''; - private preambleText: string = ''; - private buildUuid?: string; - - constructor(private callbacks: StreamCallbacks, buildUuid?: string) { - this.jsonBuffer = new JSONBuffer(); - this.buildUuid = buildUuid; - } - - private appendPreamble(text: string): void { - const cleaned = cleanPreamble(text); - if (!cleaned) return; - this.preambleText = this.preambleText ? `${this.preambleText}\n${cleaned}` : cleaned; - } - - private emitTextChunk(text: string): void { - if (!text) return; - this.plainTextResponse += text; - this.callbacks.onTextChunk(text); - } - - private isPotentialJsonPrefix(text: string): boolean { - const trimmed = text.trimStart(); - if (!trimmed) return true; - - if (trimmed.startsWith('```')) { - if (/```(?:json)?\s*\n?\s*\{/.test(trimmed)) return true; - return !trimmed.includes('```', 3); - } - - if (!trimmed.startsWith('{')) return false; - if (trimmed.includes('"type"')) return false; - if (trimmed.includes('}') && !trimmed.includes('"type"')) return false; - - const afterBrace = trimmed.slice(1); - if (afterBrace.length === 0) return true; - if (/^\s*$/.test(afterBrace)) return true; - if (/^\s*[\r\n]/.test(afterBrace)) return true; - if (/^\s*"/.test(afterBrace)) return true; - return false; - } - - private findJsonBoundary(text: string): number { - const fenceMatch = text.match(/```(?:json)?\s*\n?\s*\{/); - if (fenceMatch && fenceMatch.index !== undefined) { - return fenceMatch.index; - } - const braceIdx = text.indexOf('{'); - if (braceIdx >= 0) { - const tail = text.substring(braceIdx); - if (text.includes('"type"') || this.isPotentialJsonPrefix(tail)) { - return braceIdx; - } - } - return -1; - } - - handleChunk(text: string): void { - if (this.isJsonResponse) { - this.jsonBuffer.append(text); - - if (this.jsonBuffer.isComplete()) { - getLogger().info(`AI: JSON response complete buildUuid=${this.buildUuid || 'none'}`); - const parsed = this.jsonBuffer.parse(); - if (parsed) { - getLogger().info(`AI: structured output parsed type=${parsed.type} buildUuid=${this.buildUuid || 'none'}`); - } else { - getLogger().warn(`AI: JSON parse failed buildUuid=${this.buildUuid || 'none'}`); - } - } - return; - } - - this.textBuffer += text; - const combined = this.textBuffer; - - if (this.isJsonStart(combined)) { - this.isJsonResponse = true; - getLogger().info(`AI: JSON response detected buildUuid=${this.buildUuid || 'none'}`); - this.callbacks.onThinking('Generating structured report...'); - - const boundary = this.findJsonBoundary(combined); - const preamble = boundary > 0 ? combined.substring(0, boundary) : ''; - if (preamble.trim()) { - this.appendPreamble(preamble); - this.emitTextChunk(preamble); - } - - const stripped = stripCodeFence(combined); - const jsonIdx = stripped.indexOf('{'); - const jsonContent = jsonIdx >= 0 ? stripped.substring(jsonIdx) : stripped; - this.jsonBuffer.append(jsonContent); - this.textBuffer = ''; - return; - } - - const boundary = this.findJsonBoundary(combined); - if (boundary > 0) { - const preamble = combined.substring(0, boundary); - if (preamble.trim()) { - this.appendPreamble(preamble); - this.emitTextChunk(preamble); - } - this.textBuffer = combined.substring(boundary); - return; - } - - if (this.isPotentialJsonPrefix(combined)) { - return; - } - - this.emitTextChunk(combined); - this.textBuffer = ''; - } - - private isJsonStart(text: string): boolean { - const trimmed = text.trim(); - if (trimmed.startsWith('{') && trimmed.includes('"type"')) return true; - if (/```(?:json)?\s*\n?\s*\{/.test(trimmed) && trimmed.includes('"type"')) return true; - const stripped = stripCodeFence(trimmed); - const braceIdx = stripped.indexOf('{'); - if (braceIdx >= 0 && stripped.includes('"type"')) return true; - return false; - } - - getResult(): { response: string; isJson: boolean; preamble?: string } { - if (this.isJsonResponse) { - const content = stripCodeFence(this.jsonBuffer.getContent()).trim(); - return { - response: content, - isJson: true, - ...(this.preambleText ? { preamble: this.preambleText } : {}), - }; - } - - return { - response: this.plainTextResponse + this.textBuffer, - isJson: false, - }; - } -} diff --git a/src/server/services/ai/types/index.ts b/src/server/services/ai/types/index.ts deleted file mode 100644 index 78961aa3..00000000 --- a/src/server/services/ai/types/index.ts +++ /dev/null @@ -1,20 +0,0 @@ -/** - * Copyright 2025 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. - */ - -export * from './provider'; -export * from './tool'; -export * from './stream'; -export * from './message'; diff --git a/src/server/services/ai/types/message.ts b/src/server/services/ai/types/message.ts deleted file mode 100644 index 09971134..00000000 --- a/src/server/services/ai/types/message.ts +++ /dev/null @@ -1,63 +0,0 @@ -/** - * Copyright 2025 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 { ToolResult } from './tool'; - -export interface TextPart { - type: 'text'; - content: string; -} - -export interface ToolCallPart { - type: 'tool_call'; - toolCallId: string; - name: string; - arguments: Record; - metadata?: Record; -} - -export interface ToolResultPart { - type: 'tool_result'; - toolCallId: string; - name: string; - result: ToolResult; -} - -export type MessagePart = TextPart | ToolCallPart | ToolResultPart; - -export interface ConversationMessage { - role: 'user' | 'assistant' | 'system'; - parts: MessagePart[]; -} - -export function extractTextFromParts(parts: MessagePart[]): string { - return parts - .map((part) => { - switch (part.type) { - case 'text': - return part.content; - case 'tool_call': - return `[Tool: ${part.name}(${JSON.stringify(part.arguments)})]`; - case 'tool_result': - return `[Result: ${part.name} -> ${part.result.agentContent || JSON.stringify(part.result)}]`; - } - }) - .join(' '); -} - -export function textMessage(role: ConversationMessage['role'], content: string): ConversationMessage { - return { role, parts: [{ type: 'text', content }] }; -} diff --git a/src/server/services/ai/types/provider.ts b/src/server/services/ai/types/provider.ts deleted file mode 100644 index 154fee53..00000000 --- a/src/server/services/ai/types/provider.ts +++ /dev/null @@ -1,64 +0,0 @@ -/** - * Copyright 2025 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 { Tool, ToolCall } from './tool'; -import { StreamCallbacks } from './stream'; -import { ConversationMessage } from './message'; - -export interface CompletionOptions { - systemPrompt: string; - tools?: Tool[]; - callbacks?: StreamCallbacks; - maxTokens?: number; - temperature?: number; -} - -export interface TokenUsage { - inputTokens: number; - outputTokens: number; -} - -export type StreamChunkType = 'text' | 'tool_call' | 'thinking' | 'error'; - -export interface StreamChunk { - type: StreamChunkType; - content?: string; - toolCalls?: ToolCall[]; - metadata?: Record; - usage?: TokenUsage; -} - -export interface ModelInfo { - model: string; - maxTokens: number; -} - -export interface LLMProvider { - name: string; - - streamCompletion( - messages: ConversationMessage[], - options: CompletionOptions, - signal?: AbortSignal - ): AsyncIterator; - - supportsTools(): boolean; - getModelInfo(): ModelInfo; - formatToolDefinition(tool: Tool): unknown; - parseToolCall(response: unknown): ToolCall[]; - estimateTokens(text: string): number; - formatHistory(messages: ConversationMessage[]): unknown[]; -} diff --git a/src/server/services/ai/types/stream.ts b/src/server/services/ai/types/stream.ts deleted file mode 100644 index 8804440b..00000000 --- a/src/server/services/ai/types/stream.ts +++ /dev/null @@ -1,51 +0,0 @@ -/** - * Copyright 2025 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 { ToolResult, ConfirmationDetails } from './tool'; - -export interface ActivityEvent { - type: 'tool_call' | 'thinking' | 'processing' | 'error'; - message: string; - details?: unknown; - toolCallId?: string; - timestamp: number; -} - -export interface StreamCallbacks { - onTextChunk(text: string): void; - onThinking(message: string, details?: unknown): void; - onToolCall(tool: string, args: unknown, toolCallId?: string): void; - onToolResult( - result: ToolResult, - toolName: string, - toolArgs: unknown, - toolDurationMs?: number, - totalDurationMs?: number, - toolCallId?: string - ): void; - onError(error: ToolResult['error']): void; - onActivity(activity: ActivityEvent): void; - onToolConfirmation?(details: ConfirmationDetails): Promise; - onToolAuthorization?( - tool: { - name: string; - description: string; - category: string; - safetyLevel: string; - }, - args: Record - ): Promise<{ allowed: boolean; reason?: string }>; -} diff --git a/src/server/services/ai/utils/__tests__/fixTargetAuthorization.test.ts b/src/server/services/ai/utils/__tests__/fixTargetAuthorization.test.ts deleted file mode 100644 index 2a696113..00000000 --- a/src/server/services/ai/utils/__tests__/fixTargetAuthorization.test.ts +++ /dev/null @@ -1,98 +0,0 @@ -/** - * Copyright 2025 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 { authorizeToolForFixTarget } from '../fixTargetAuthorization'; - -describe('authorizeToolForFixTarget', () => { - it('blocks file writes when selected fix target is PR label-only', () => { - const decision = authorizeToolForFixTarget( - { - serviceName: 'Environment', - suggestedFix: 'Add the lifecycle-deploy! label to the PR in GitHub.', - }, - { - name: 'update_file', - description: 'Update repository file', - category: 'github', - safetyLevel: 'dangerous', - args: { - file_path: 'lifecycle.yaml', - }, - } - ); - - expect(decision.allowed).toBe(false); - }); - - it('allows file writes only for selected target files', () => { - const allowDecision = authorizeToolForFixTarget( - { - serviceName: 'sample-helm-service', - suggestedFix: "Change dockerfilePath from 'a' to 'b' in lifecycle.yaml", - filePath: 'lifecycle.yaml', - }, - { - name: 'update_file', - description: 'Update repository file', - category: 'github', - safetyLevel: 'dangerous', - args: { - file_path: 'lifecycle.yaml', - }, - } - ); - - const denyDecision = authorizeToolForFixTarget( - { - serviceName: 'sample-helm-service', - suggestedFix: "Change dockerfilePath from 'a' to 'b' in lifecycle.yaml", - filePath: 'lifecycle.yaml', - }, - { - name: 'update_file', - description: 'Update repository file', - category: 'github', - safetyLevel: 'dangerous', - args: { - file_path: 'sysops/dockerfiles/app.dockerfile', - }, - } - ); - - expect(allowDecision.allowed).toBe(true); - expect(denyDecision.allowed).toBe(false); - }); - - it('allows a PR label tool when target is label fix', () => { - const decision = authorizeToolForFixTarget( - { - serviceName: 'Environment', - suggestedFix: 'Add the lifecycle-deploy! label to the PR in GitHub.', - }, - { - name: 'mcp__github__add_pr_label', - description: 'Add label to pull request', - category: 'mcp', - safetyLevel: 'cautious', - args: { - label: 'lifecycle-deploy!', - }, - } - ); - - expect(decision.allowed).toBe(true); - }); -}); diff --git a/src/server/services/ai/utils/__tests__/jsonExtraction.test.ts b/src/server/services/ai/utils/__tests__/jsonExtraction.test.ts deleted file mode 100644 index b16956ef..00000000 --- a/src/server/services/ai/utils/__tests__/jsonExtraction.test.ts +++ /dev/null @@ -1,148 +0,0 @@ -/** - * Copyright 2025 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/logger', () => ({ - getLogger: () => ({ - info: jest.fn(), - warn: jest.fn(), - error: jest.fn(), - }), -})); - -import { extractJsonFromResponse, extractBalancedJson } from '../jsonExtraction'; - -describe('extractBalancedJson', () => { - it('extracts a simple JSON object', () => { - const result = extractBalancedJson('{"a": 1}', 0); - expect(result).toBe('{"a": 1}'); - }); - - it('extracts nested JSON objects', () => { - const input = '{"a": {"b": {"c": 1}}}'; - expect(extractBalancedJson(input, 0)).toBe(input); - }); - - it('extracts JSON starting at an offset', () => { - const input = 'prefix {"type": "x"} suffix'; - expect(extractBalancedJson(input, 7)).toBe('{"type": "x"}'); - }); - - it('handles braces inside strings', () => { - const input = '{"text": "a { b } c"}'; - expect(extractBalancedJson(input, 0)).toBe(input); - }); - - it('handles escaped quotes in strings', () => { - const input = '{"text": "say \\"hello\\""}'; - expect(extractBalancedJson(input, 0)).toBe(input); - }); - - it('returns null for non-brace start', () => { - expect(extractBalancedJson('abc', 0)).toBeNull(); - }); - - it('returns null for unbalanced braces', () => { - expect(extractBalancedJson('{incomplete', 0)).toBeNull(); - }); -}); - -describe('extractJsonFromResponse', () => { - const buildUuid = 'test-uuid'; - - it('returns isJson: false for non-investigation text', () => { - const result = extractJsonFromResponse('Hello world', buildUuid); - expect(result).toEqual({ response: 'Hello world', isJson: false }); - }); - - it('extracts pure JSON string', () => { - const json = '{"type": "investigation_complete", "summary": "done"}'; - const result = extractJsonFromResponse(json, buildUuid); - expect(result.isJson).toBe(true); - expect(JSON.parse(result.response)).toEqual({ type: 'investigation_complete', summary: 'done' }); - }); - - it('extracts fenced JSON (```json)', () => { - const input = '```json\n{"type": "investigation_complete", "data": []}\n```'; - const result = extractJsonFromResponse(input, buildUuid); - expect(result.isJson).toBe(true); - expect(JSON.parse(result.response)).toEqual({ type: 'investigation_complete', data: [] }); - }); - - it('extracts preamble text + fenced JSON', () => { - const input = 'Here are the findings:\n\n```json\n{"type": "investigation_complete", "summary": "ok"}\n```'; - const result = extractJsonFromResponse(input, buildUuid); - expect(result.isJson).toBe(true); - expect(JSON.parse(result.response)).toEqual({ type: 'investigation_complete', summary: 'ok' }); - }); - - it('extracts preamble text + raw JSON (no fences)', () => { - const input = 'Analysis complete.\n{"type": "investigation_complete", "items": [1, 2]}'; - const result = extractJsonFromResponse(input, buildUuid); - expect(result.isJson).toBe(true); - expect(JSON.parse(result.response)).toEqual({ type: 'investigation_complete', items: [1, 2] }); - }); - - it('content field parses with JSON.parse()', () => { - const json = '{"type": "investigation_complete", "nested": {"a": [1, 2, 3]}}'; - const result = extractJsonFromResponse(json, buildUuid); - expect(result.isJson).toBe(true); - const parsed = JSON.parse(result.response); - expect(parsed.nested.a).toEqual([1, 2, 3]); - }); - - it('handles JSON with trailing text after fence', () => { - const input = '```json\n{"type": "investigation_complete"}\n```\n\nLet me know if you need more details.'; - const result = extractJsonFromResponse(input, buildUuid); - expect(result.isJson).toBe(true); - expect(JSON.parse(result.response)).toEqual({ type: 'investigation_complete' }); - }); - - it('returns isJson: false for text mentioning investigation_complete without valid JSON', () => { - const input = 'The "investigation_complete" status was set but no data found.'; - const result = extractJsonFromResponse(input, buildUuid); - expect(result.isJson).toBe(false); - }); - - describe('preamble extraction', () => { - it('returns preamble for preamble + fenced JSON', () => { - const input = 'Here are the findings:\n\n```json\n{"type": "investigation_complete", "data": []}\n```'; - const result = extractJsonFromResponse(input, buildUuid); - expect(result.isJson).toBe(true); - expect(result.preamble).toBe('Here are the findings:'); - }); - - it('returns preamble for preamble + raw JSON', () => { - const input = 'Analysis complete.\n{"type": "investigation_complete", "items": []}'; - const result = extractJsonFromResponse(input, buildUuid); - expect(result.isJson).toBe(true); - expect(result.preamble).toBe('Analysis complete.'); - }); - - it('does not return preamble for pure JSON', () => { - const json = '{"type": "investigation_complete", "summary": "done"}'; - const result = extractJsonFromResponse(json, buildUuid); - expect(result.isJson).toBe(true); - expect(result.preamble).toBeUndefined(); - }); - - it('does not return preamble for fenced JSON without preamble text', () => { - const input = '```json\n{"type": "investigation_complete"}\n```'; - const result = extractJsonFromResponse(input, buildUuid); - expect(result.isJson).toBe(true); - expect(result.preamble).toBeUndefined(); - }); - }); -}); diff --git a/src/server/services/ai/utils/__tests__/normalizePayload.test.ts b/src/server/services/ai/utils/__tests__/normalizePayload.test.ts deleted file mode 100644 index fbc2d2b0..00000000 --- a/src/server/services/ai/utils/__tests__/normalizePayload.test.ts +++ /dev/null @@ -1,374 +0,0 @@ -/** - * Copyright 2025 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/logger', () => ({ - getLogger: () => ({ - info: jest.fn(), - warn: jest.fn(), - error: jest.fn(), - }), -})); - -import { normalizeInvestigationPayload } from '../normalizePayload'; - -describe('normalizeInvestigationPayload', () => { - it('passes through a valid payload unchanged', () => { - const payload = { - type: 'investigation_complete', - summary: 'Build failed due to missing dependency', - fixesApplied: false, - services: [ - { - serviceName: 'web', - status: 'build_failed', - issue: 'Missing dependency', - suggestedFix: 'Add the dependency', - fixesApplied: false, - }, - ], - }; - const result = normalizeInvestigationPayload(payload); - expect(result).toEqual(payload); - }); - - it('defaults missing summary to empty string', () => { - const payload = { type: 'investigation_complete', services: [] }; - const result = normalizeInvestigationPayload(payload) as any; - expect(result.summary).toBe(''); - }); - - it('defaults missing services to empty array', () => { - const payload = { type: 'investigation_complete', summary: 'test' }; - const result = normalizeInvestigationPayload(payload) as any; - expect(result.services).toEqual([]); - }); - - it('defaults non-array services to empty array', () => { - const payload = { type: 'investigation_complete', summary: 'test', services: 'invalid' }; - const result = normalizeInvestigationPayload(payload) as any; - expect(result.services).toEqual([]); - }); - - it('defaults missing serviceName to unknown', () => { - const payload = { - type: 'investigation_complete', - summary: 'test', - services: [{ status: 'error', issue: 'broken', suggestedFix: 'fix it' }], - }; - const result = normalizeInvestigationPayload(payload) as any; - expect(result.services[0].serviceName).toBe('unknown'); - }); - - it('defaults missing status to error', () => { - const payload = { - type: 'investigation_complete', - summary: 'test', - services: [{ serviceName: 'web', issue: 'broken', suggestedFix: 'fix it' }], - }; - const result = normalizeInvestigationPayload(payload) as any; - expect(result.services[0].status).toBe('error'); - }); - - it('preserves invalid status values but logs a warning', () => { - const payload = { - type: 'investigation_complete', - summary: 'test', - services: [{ serviceName: 'web', status: 'unknown_status', issue: 'broken', suggestedFix: 'fix' }], - }; - const result = normalizeInvestigationPayload(payload) as any; - expect(result.services[0].status).toBe('unknown_status'); - }); - - it('defaults missing issue to empty string', () => { - const payload = { - type: 'investigation_complete', - summary: 'test', - services: [{ serviceName: 'web', status: 'error', suggestedFix: 'fix it' }], - }; - const result = normalizeInvestigationPayload(payload) as any; - expect(result.services[0].issue).toBe(''); - }); - - it('defaults missing suggestedFix to empty string', () => { - const payload = { - type: 'investigation_complete', - summary: 'test', - services: [{ serviceName: 'web', status: 'error', issue: 'broken' }], - }; - const result = normalizeInvestigationPayload(payload) as any; - expect(result.services[0].suggestedFix).toBe(''); - }); - - it('forces canAutoFix=false when missing or non-boolean', () => { - const payload = { - type: 'investigation_complete', - summary: 'test', - services: [ - { serviceName: 'web', status: 'error', issue: 'broken', suggestedFix: 'fix it' }, - { serviceName: 'api', status: 'error', issue: 'broken', suggestedFix: 'fix it', canAutoFix: 'yes' }, - ], - }; - const result = normalizeInvestigationPayload(payload) as any; - expect(result.services[0].canAutoFix).toBe(false); - expect(result.services[1].canAutoFix).toBe(false); - }); - - it('downgrades canAutoFix when specific error evidence is missing', () => { - const payload = { - type: 'investigation_complete', - summary: 'test', - services: [ - { - serviceName: 'web', - status: 'deploy_failed', - issue: 'Schema grant failed', - suggestedFix: "Change objs from 'a' to 'b' in lifecycle.yaml", - canAutoFix: true, - filePath: 'lifecycle.yaml', - lineNumber: 42, - }, - ], - }; - const result = normalizeInvestigationPayload(payload) as any; - expect(result.services[0].canAutoFix).toBe(false); - }); - - it('downgrades canAutoFix when file target is missing', () => { - const payload = { - type: 'investigation_complete', - summary: 'test', - services: [ - { - serviceName: 'web', - status: 'deploy_failed', - issue: 'Schema grant failed', - keyError: 'ERROR: relation "subscriber" does not exist', - errorSource: 'build_logs', - suggestedFix: 'Run migrations first', - canAutoFix: true, - }, - ], - }; - const result = normalizeInvestigationPayload(payload) as any; - expect(result.services[0].canAutoFix).toBe(false); - }); - - it('downgrades canAutoFix for uncertain recommendations', () => { - const payload = { - type: 'investigation_complete', - summary: 'test', - services: [ - { - serviceName: 'web', - status: 'deploy_failed', - issue: 'This might be related to schema order', - keyError: 'ERROR: relation "subscriber" does not exist', - errorSource: 'build_logs', - suggestedFix: "Change objs from 'a' to 'b' in lifecycle.yaml", - canAutoFix: true, - filePath: 'lifecycle.yaml', - }, - ], - }; - const result = normalizeInvestigationPayload(payload) as any; - expect(result.services[0].canAutoFix).toBe(false); - }); - - it('keeps canAutoFix=true for actionable single-line fixes with evidence', () => { - const payload = { - type: 'investigation_complete', - summary: 'test', - services: [ - { - serviceName: 'web', - status: 'deploy_failed', - issue: 'Grant targets wrong table', - keyError: 'ERROR: relation "subscriber" does not exist', - errorSource: 'build_logs', - suggestedFix: "Change objs from 'spatial_ref_sys' to 'subscriber,promo,stripe_customer' in lifecycle.yaml", - canAutoFix: true, - filePath: 'lifecycle.yaml', - lineNumber: 105, - }, - ], - }; - const result = normalizeInvestigationPayload(payload) as any; - expect(result.services[0].canAutoFix).toBe(true); - }); - - it('derives filePath from suggestedFix single-line pattern', () => { - const payload = { - type: 'investigation_complete', - summary: 'test', - services: [ - { - serviceName: 'web', - status: 'build_failed', - issue: 'Dockerfile path points to missing file', - keyError: 'open sysops/dockerfiles/appaasd.dockerfile: no such file or directory', - errorSource: 'build_logs', - suggestedFix: - "Change dockerfilePath from 'sysops/dockerfiles/appaasd.dockerfile' to 'sysops/dockerfiles/app.dockerfile' in lifecycle.yaml", - canAutoFix: true, - }, - ], - }; - - const result = normalizeInvestigationPayload(payload) as any; - expect(result.services[0].filePath).toBe('lifecycle.yaml'); - expect(result.services[0].canAutoFix).toBe(true); - }); - - it('keeps canAutoFix=true for actionable multi-line file diffs with evidence', () => { - const payload = { - type: 'investigation_complete', - summary: 'test', - services: [ - { - serviceName: 'web', - status: 'deploy_failed', - issue: 'Migrations must run before grants', - keyError: 'ERROR: relation "subscriber" does not exist', - errorSource: 'build_logs', - suggestedFix: 'Move migrations block before grants in lifecycle.yaml', - canAutoFix: true, - files: [ - { - path: 'sysops/ansible/playbooks/lifecycle.yaml', - oldContent: '- name: Add grants', - newContent: '- name: Run Migrations', - }, - ], - }, - ], - }; - const result = normalizeInvestigationPayload(payload) as any; - expect(result.services[0].canAutoFix).toBe(true); - }); - - it('defaults non-boolean fixesApplied to false on services', () => { - const payload = { - type: 'investigation_complete', - summary: 'test', - services: [{ serviceName: 'web', status: 'error', issue: 'broken', suggestedFix: 'fix', fixesApplied: 'yes' }], - }; - const result = normalizeInvestigationPayload(payload) as any; - expect(result.services[0].fixesApplied).toBe(false); - }); - - it('defaults non-boolean top-level fixesApplied to false', () => { - const payload = { - type: 'investigation_complete', - summary: 'test', - services: [], - }; - const result = normalizeInvestigationPayload(payload) as any; - expect(result.fixesApplied).toBe(false); - }); - - it('preserves true fixesApplied when valid', () => { - const payload = { - type: 'investigation_complete', - summary: 'Fixed it', - fixesApplied: true, - services: [ - { - serviceName: 'web', - status: 'ready', - issue: 'Was broken', - suggestedFix: 'Fixed', - fixesApplied: true, - }, - ], - }; - const result = normalizeInvestigationPayload(payload) as any; - expect(result.fixesApplied).toBe(true); - expect(result.services[0].fixesApplied).toBe(true); - }); - - it('downgrades canAutoFix for file edits when no file-write tool is available', () => { - const payload = { - type: 'investigation_complete', - summary: 'test', - services: [ - { - serviceName: 'web', - status: 'deploy_failed', - issue: 'Grant targets wrong table', - keyError: 'ERROR: relation "subscriber" does not exist', - errorSource: 'build_logs', - suggestedFix: "Change objs from 'a' to 'b' in lifecycle.yaml", - canAutoFix: true, - filePath: 'lifecycle.yaml', - }, - ], - }; - - const result = normalizeInvestigationPayload(payload, { - availableTools: [{ name: 'get_file', description: 'Read repository files' }], - }) as any; - - expect(result.services[0].canAutoFix).toBe(false); - }); - - it('allows canAutoFix for PR label fixes when a label tool is available', () => { - const payload = { - type: 'investigation_complete', - summary: 'test', - services: [ - { - serviceName: 'environment', - status: 'error', - issue: 'Missing deploy label blocks environment creation', - keyError: 'Deploy label not present', - errorSource: 'lifecycle', - suggestedFix: 'Add the lifecycle-deploy! label to the PR to start deployment.', - canAutoFix: true, - }, - ], - }; - - const result = normalizeInvestigationPayload(payload, { - availableTools: [{ name: 'mcp__github__add_pr_label', description: 'Add a label to a pull request' }], - }) as any; - - expect(result.services[0].canAutoFix).toBe(true); - }); - - it('downgrades PR label fixes when label mutation tool is not available', () => { - const payload = { - type: 'investigation_complete', - summary: 'test', - services: [ - { - serviceName: 'environment', - status: 'error', - issue: 'Missing deploy label blocks environment creation', - keyError: 'Deploy label not present', - errorSource: 'lifecycle', - suggestedFix: 'Add the lifecycle-deploy! label to the PR to start deployment.', - canAutoFix: true, - }, - ], - }; - - const result = normalizeInvestigationPayload(payload, { - availableTools: [{ name: 'get_issue_comment', description: 'Read pull request comment text' }], - }) as any; - - expect(result.services[0].canAutoFix).toBe(false); - }); -}); diff --git a/src/server/services/ai/utils/fixTargetAuthorization.ts b/src/server/services/ai/utils/fixTargetAuthorization.ts deleted file mode 100644 index 53691743..00000000 --- a/src/server/services/ai/utils/fixTargetAuthorization.ts +++ /dev/null @@ -1,246 +0,0 @@ -/** - * Copyright 2025 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. - */ - -type Capability = 'file_write' | 'pr_label_write' | 'k8s_patch'; - -export interface FixTargetScope { - serviceName?: string; - suggestedFix?: string; - filePath?: string; - files?: Array<{ path?: string }>; - autoFixAction?: string; - actionType?: string; - fixType?: string; - tool?: string; - toolName?: string; -} - -export interface ToolAuthorizationInput { - name: string; - description: string; - category: string; - safetyLevel: string; - args: Record; -} - -export interface ToolAuthorizationDecision { - allowed: boolean; - reason?: string; -} - -const SINGLE_LINE_FIX_PATTERN = /from ['"]([^'"]+)['"] to ['"]([^'"]+)['"] in ([\w/.+-]+\.\w+)/i; -const PR_LABEL_MUTATION_PATTERN = /\b(add|apply|set|remove|update|edit)\b/i; -const PR_LABEL_CONTEXT_PATTERN = /\b(pr|pull[\s-]?request)\b/i; -const LABEL_PATTERN = /\blabels?\b/i; - -function isPrLabelFix(target: FixTargetScope): boolean { - const fix = target.suggestedFix || ''; - return ( - fix.length > 0 && - PR_LABEL_MUTATION_PATTERN.test(fix) && - LABEL_PATTERN.test(fix) && - PR_LABEL_CONTEXT_PATTERN.test(fix) - ); -} - -function detectCapabilities(target: FixTargetScope): Set { - const capabilities = new Set(); - const explicitHints = [target.autoFixAction, target.actionType, target.fixType, target.tool, target.toolName] - .filter((v): v is string => typeof v === 'string' && v.trim().length > 0) - .join(' ') - .toLowerCase(); - - if (explicitHints.includes('label') && (explicitHints.includes('pr') || explicitHints.includes('pull'))) { - capabilities.add('pr_label_write'); - } - if ( - explicitHints.includes('patch_k8s') || - (explicitHints.includes('k8s') && explicitHints.includes('patch')) || - (explicitHints.includes('kubernetes') && explicitHints.includes('patch')) - ) { - capabilities.add('k8s_patch'); - } - if ( - explicitHints.includes('update_file') || - explicitHints.includes('commit_lifecycle_fix') || - explicitHints.includes('file') || - explicitHints.includes('config') - ) { - capabilities.add('file_write'); - } - - if (target.filePath?.trim()) { - capabilities.add('file_write'); - } - if (target.files?.some((f) => typeof f.path === 'string' && f.path.trim().length > 0)) { - capabilities.add('file_write'); - } - if (typeof target.suggestedFix === 'string' && SINGLE_LINE_FIX_PATTERN.test(target.suggestedFix)) { - capabilities.add('file_write'); - } - if (isPrLabelFix(target)) { - capabilities.add('pr_label_write'); - } - - return capabilities; -} - -function getAllowedPaths(target: FixTargetScope): Set { - const normalizePath = (value: string): string => - value - .trim() - .replace(/^\.\/+/, '') - .replace(/\/{2,}/g, '/') - .toLowerCase(); - - const paths = new Set(); - if (typeof target.filePath === 'string' && target.filePath.trim().length > 0) { - paths.add(normalizePath(target.filePath)); - } - if (Array.isArray(target.files)) { - for (const file of target.files) { - if (typeof file?.path === 'string' && file.path.trim().length > 0) { - paths.add(normalizePath(file.path)); - } - } - } - return paths; -} - -function toolText(tool: ToolAuthorizationInput): string { - return `${tool.name} ${tool.description}`.toLowerCase(); -} - -function isFileWriteTool(tool: ToolAuthorizationInput): boolean { - const text = toolText(tool); - if ( - text.includes('update_file') || - text.includes('commit_lifecycle_fix') || - text.includes('write_file') || - text.includes('edit_file') || - text.includes('modify_file') - ) { - return true; - } - return /\b(update|edit|modify|write|create|commit|patch|replace)\b/.test(text) && /\bfile\b/.test(text); -} - -function isPrLabelWriteTool(tool: ToolAuthorizationInput): boolean { - const text = toolText(tool); - return ( - /\blabels?\b/.test(text) && - /\b(pr|pull[\s_-]?request|issue)\b/.test(text) && - /\b(add|set|remove|update|edit|patch|apply)\b/.test(text) - ); -} - -function isK8sPatchTool(tool: ToolAuthorizationInput): boolean { - const text = toolText(tool); - if (text.includes('patch_k8s_resource')) return true; - return /\b(k8s|kubernetes)\b/.test(text) && /\b(patch|apply|update|edit)\b/.test(text); -} - -function isReadOnlyTool(tool: ToolAuthorizationInput): boolean { - const text = toolText(tool); - if (tool.safetyLevel === 'safe') return true; - return ( - /\b(get|list|read|query|fetch|describe)\b/.test(text) && - !/\b(update|patch|write|create|delete|add|set|remove|apply|commit|edit)\b/.test(text) - ); -} - -function getToolFilePath(args: Record): string | undefined { - const normalizePath = (value: string): string => - value - .trim() - .replace(/^\.\/+/, '') - .replace(/\/{2,}/g, '/') - .toLowerCase(); - - const candidates = [args.file_path, args.path]; - for (const candidate of candidates) { - if (typeof candidate === 'string' && candidate.trim().length > 0) { - return normalizePath(candidate); - } - } - return undefined; -} - -export function authorizeToolForFixTarget( - target: FixTargetScope, - tool: ToolAuthorizationInput -): ToolAuthorizationDecision { - if (!target || typeof target !== 'object') { - return { allowed: true }; - } - - if (isReadOnlyTool(tool)) { - return { allowed: true }; - } - - const capabilities = detectCapabilities(target); - const allowedPaths = getAllowedPaths(target); - - if (isFileWriteTool(tool)) { - if (!capabilities.has('file_write')) { - return { - allowed: false, - reason: `Blocked ${tool.name}: fix target "${ - target.serviceName || 'selected service' - }" does not allow file edits`, - }; - } - if (allowedPaths.size > 0) { - const requestedPath = getToolFilePath(tool.args); - if (!requestedPath || !allowedPaths.has(requestedPath)) { - return { - allowed: false, - reason: `Blocked ${tool.name}: file path is outside the selected fix target scope`, - }; - } - } - return { allowed: true }; - } - - if (isPrLabelWriteTool(tool)) { - if (!capabilities.has('pr_label_write')) { - return { - allowed: false, - reason: `Blocked ${tool.name}: selected fix target does not allow PR label changes`, - }; - } - return { allowed: true }; - } - - if (isK8sPatchTool(tool)) { - if (!capabilities.has('k8s_patch')) { - return { - allowed: false, - reason: `Blocked ${tool.name}: selected fix target does not allow Kubernetes patching`, - }; - } - return { allowed: true }; - } - - if (tool.safetyLevel === 'dangerous' || tool.safetyLevel === 'cautious') { - return { - allowed: false, - reason: `Blocked ${tool.name}: mutating tool is outside the selected fix target scope`, - }; - } - - return { allowed: true }; -} diff --git a/src/server/services/ai/utils/jsonExtraction.ts b/src/server/services/ai/utils/jsonExtraction.ts deleted file mode 100644 index bd4d1e24..00000000 --- a/src/server/services/ai/utils/jsonExtraction.ts +++ /dev/null @@ -1,120 +0,0 @@ -/** - * Copyright 2025 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'; - -export function extractBalancedJson(str: string, startIndex: number): string | null { - if (str[startIndex] !== '{') return null; - - let depth = 0; - let inString = false; - let escape = false; - - for (let i = startIndex; i < str.length; i++) { - const ch = str[i]; - - if (escape) { - escape = false; - continue; - } - - if (ch === '\\' && inString) { - escape = true; - continue; - } - - if (ch === '"') { - inString = !inString; - continue; - } - - if (inString) continue; - - if (ch === '{') depth++; - else if (ch === '}') { - depth--; - if (depth === 0) { - return str.substring(startIndex, i + 1); - } - } - } - - return null; -} - -export function extractJsonFromResponse( - aiResponse: string, - buildUuid: string -): { response: string; isJson: boolean; preamble?: string } { - if (!aiResponse.includes('"investigation_complete"')) { - return { response: aiResponse, isJson: false }; - } - - const original = aiResponse.trim(); - - let candidate = original; - - const fenceStart = candidate.indexOf('```'); - const hasFence = fenceStart >= 0; - const preambleBeforeFence = hasFence ? candidate.substring(0, fenceStart).trim() : ''; - - candidate = candidate - .replace(/```(?:json)?\s*\n?/g, '') - .replace(/\n?\s*```/g, '') - .trim(); - - const jsonIdx = candidate.indexOf('{"type"'); - if (jsonIdx < 0) { - const braceIdx = candidate.indexOf('{'); - if (braceIdx < 0) return { response: aiResponse, isJson: false }; - - const preambleRaw = candidate.substring(0, braceIdx).trim(); - const balanced = extractBalancedJson(candidate, braceIdx); - if (balanced) { - try { - JSON.parse(balanced); - getLogger().info(`AI: late JSON detection - extracted valid JSON buildUuid=${buildUuid}`); - return { - response: balanced, - isJson: true, - ...(preambleRaw ? { preamble: preambleRaw } : {}), - }; - } catch { - return { response: aiResponse, isJson: false }; - } - } - return { response: aiResponse, isJson: false }; - } - - const preambleRaw = hasFence ? preambleBeforeFence : candidate.substring(0, jsonIdx).trim(); - - const balanced = extractBalancedJson(candidate, jsonIdx); - if (balanced) { - try { - JSON.parse(balanced); - getLogger().info(`AI: late JSON detection - extracted valid JSON buildUuid=${buildUuid}`); - return { - response: balanced, - isJson: true, - ...(preambleRaw ? { preamble: preambleRaw } : {}), - }; - } catch { - return { response: aiResponse, isJson: false }; - } - } - - return { response: aiResponse, isJson: false }; -} diff --git a/src/server/services/ai/utils/normalizePayload.ts b/src/server/services/ai/utils/normalizePayload.ts deleted file mode 100644 index 12c4e648..00000000 --- a/src/server/services/ai/utils/normalizePayload.ts +++ /dev/null @@ -1,254 +0,0 @@ -/** - * Copyright 2025 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'; - -const VALID_STATUSES = new Set(['build_failed', 'deploy_failed', 'error', 'ready']); -const UNCERTAINTY_PATTERN = /\b(maybe|might|could|likely|possibly|uncertain|not sure|probably)\b/i; -const NON_ACTIONABLE_PATTERN = /\b(no action needed|manual fix|choose|decide|depends on)\b/i; -const SINGLE_LINE_FIX_PATTERN = /from ['"]([^'"]+)['"] to ['"]([^'"]+)['"] in ([\w/.+-]+\.\w+)/i; -const PR_LABEL_MUTATION_PATTERN = /\b(add|apply|set|remove|update|edit)\b/i; -const PR_LABEL_CONTEXT_PATTERN = /\b(pr|pull[\s-]?request)\b/i; -const LABEL_PATTERN = /\blabels?\b/i; - -type AutoFixCapability = 'file_write' | 'pr_label_write' | 'k8s_patch'; - -export interface AvailableToolInfo { - name: string; - description?: string; - category?: string; - safetyLevel?: string; -} - -export interface NormalizePayloadOptions { - availableTools?: AvailableToolInfo[]; -} - -function hasSpecificErrorEvidence(service: Record): boolean { - const keyError = typeof service.keyError === 'string' ? service.keyError.trim() : ''; - const errorSource = typeof service.errorSource === 'string' ? service.errorSource.trim() : ''; - return keyError.length > 0 && errorSource.length > 0; -} - -function hasFileDiffPayload(service: Record): boolean { - if (!Array.isArray(service.files)) return false; - return service.files.some((file: Record) => { - if (!file || typeof file !== 'object') return false; - const path = typeof file.path === 'string' ? file.path.trim() : ''; - return ( - path.length > 0 && - typeof file.oldContent === 'string' && - typeof file.newContent === 'string' && - file.oldContent !== file.newContent - ); - }); -} - -function extractSingleLineFixFilePath(service: Record): string | undefined { - const explicitPath = typeof service.filePath === 'string' ? service.filePath.trim() : ''; - if (explicitPath.length > 0) return explicitPath; - - const suggestedFix = typeof service.suggestedFix === 'string' ? service.suggestedFix : ''; - const match = suggestedFix.match(SINGLE_LINE_FIX_PATTERN); - if (!match || typeof match[3] !== 'string') return undefined; - const derivedPath = match[3].trim(); - return derivedPath.length > 0 ? derivedPath : undefined; -} - -function hasSingleLineFileTarget(service: Record): boolean { - const suggestedFix = typeof service.suggestedFix === 'string' ? service.suggestedFix : ''; - return Boolean(extractSingleLineFixFilePath(service)) && SINGLE_LINE_FIX_PATTERN.test(suggestedFix); -} - -function isPrLabelFix(service: Record): boolean { - const suggestedFix = typeof service.suggestedFix === 'string' ? service.suggestedFix : ''; - return ( - suggestedFix.length > 0 && - PR_LABEL_MUTATION_PATTERN.test(suggestedFix) && - LABEL_PATTERN.test(suggestedFix) && - PR_LABEL_CONTEXT_PATTERN.test(suggestedFix) - ); -} - -function isConfidentlyActionable(service: Record): boolean { - const issue = typeof service.issue === 'string' ? service.issue : ''; - const suggestedFix = typeof service.suggestedFix === 'string' ? service.suggestedFix : ''; - return ( - !UNCERTAINTY_PATTERN.test(issue) && - !UNCERTAINTY_PATTERN.test(suggestedFix) && - !NON_ACTIONABLE_PATTERN.test(suggestedFix) - ); -} - -function toolText(tool: AvailableToolInfo): string { - return `${tool.name || ''} ${tool.description || ''}`.toLowerCase(); -} - -function supportsFileWrite(tool: AvailableToolInfo): boolean { - const text = toolText(tool); - if ( - text.includes('update_file') || - text.includes('commit_lifecycle_fix') || - text.includes('edit_file') || - text.includes('write_file') || - text.includes('modify_file') - ) { - return true; - } - const hasMutation = /\b(update|edit|modify|write|create|commit|patch|replace)\b/.test(text); - const hasFileScope = /\b(file|yaml|yml|config|manifest|dockerfile|helm|values)\b/.test(text); - return hasMutation && hasFileScope; -} - -function supportsPrLabelWrite(tool: AvailableToolInfo): boolean { - const text = toolText(tool); - const hasLabel = /\blabels?\b/.test(text); - const hasPrContext = /\b(pr|pull[\s_-]?request|issue)\b/.test(text); - const hasMutation = /\b(add|set|remove|update|edit|patch|apply)\b/.test(text); - return hasLabel && hasPrContext && hasMutation; -} - -function supportsK8sPatch(tool: AvailableToolInfo): boolean { - const text = toolText(tool); - if (text.includes('patch_k8s_resource')) return true; - const hasK8s = /\b(k8s|kubernetes)\b/.test(text); - const hasMutation = /\b(patch|apply|update|edit)\b/.test(text); - const hasResource = /\b(resource|deployment|statefulset|service|configmap|secret|ingress)\b/.test(text); - return hasK8s && hasMutation && hasResource; -} - -function hasCapability(capability: AutoFixCapability, availableTools: AvailableToolInfo[]): boolean { - if (capability === 'file_write') { - return availableTools.some((tool) => supportsFileWrite(tool)); - } - if (capability === 'pr_label_write') { - return availableTools.some((tool) => supportsPrLabelWrite(tool)); - } - return availableTools.some((tool) => supportsK8sPatch(tool)); -} - -function getRequiredCapabilities(service: Record): AutoFixCapability[] { - const explicitHintCandidates = ['autoFixAction', 'actionType', 'fixType', 'tool', 'toolName']; - const explicitHint = explicitHintCandidates - .map((key) => service[key]) - .find((value) => typeof value === 'string' && value.trim().length > 0) - ?.toLowerCase(); - - if (explicitHint) { - if ( - explicitHint.includes('label') && - (explicitHint.includes('pr') || explicitHint.includes('pull') || explicitHint.includes('issue')) - ) { - return ['pr_label_write']; - } - if ( - explicitHint.includes('patch_k8s') || - (explicitHint.includes('k8s') && explicitHint.includes('patch')) || - (explicitHint.includes('kubernetes') && explicitHint.includes('patch')) - ) { - return ['k8s_patch']; - } - if ( - explicitHint.includes('update_file') || - explicitHint.includes('commit_lifecycle_fix') || - explicitHint.includes('file') || - explicitHint.includes('config') - ) { - return ['file_write']; - } - } - - const capabilities: AutoFixCapability[] = []; - if (hasFileDiffPayload(service) || hasSingleLineFileTarget(service)) { - capabilities.push('file_write'); - } - if (isPrLabelFix(service)) { - capabilities.push('pr_label_write'); - } - return [...new Set(capabilities)]; -} - -function shouldAllowAutoFix(service: Record, options: NormalizePayloadOptions): boolean { - if (service.canAutoFix !== true) return false; - if (!hasSpecificErrorEvidence(service)) return false; - if (!isConfidentlyActionable(service)) return false; - - const requiredCapabilities = getRequiredCapabilities(service); - if (requiredCapabilities.length === 0) return false; - - const availableTools = options.availableTools; - if (!availableTools || availableTools.length === 0) { - // Preserve legacy behavior when tool inventory is unavailable. - return requiredCapabilities.every((capability) => capability === 'file_write'); - } - - return requiredCapabilities.every((capability) => hasCapability(capability, availableTools)); -} - -export function normalizeInvestigationPayload(parsed: any, options: NormalizePayloadOptions = {}): object { - if (!parsed || typeof parsed !== 'object') { - return parsed; - } - - if (parsed.summary === undefined || parsed.summary === null) { - parsed.summary = ''; - } - - if (!Array.isArray(parsed.services)) { - parsed.services = []; - } - - for (const service of parsed.services) { - if (typeof service !== 'object' || service === null) continue; - - if (!service.serviceName) { - service.serviceName = 'unknown'; - } - - if (!service.status) { - service.status = 'error'; - } else if (!VALID_STATUSES.has(service.status)) { - getLogger().warn(`AI: invalid service status="${service.status}" serviceName=${service.serviceName}`); - } - - if (!service.issue) { - service.issue = ''; - } - - if (!service.suggestedFix) { - service.suggestedFix = ''; - } - - if (!service.filePath) { - const derivedPath = extractSingleLineFixFilePath(service); - if (derivedPath) { - service.filePath = derivedPath; - } - } - - service.canAutoFix = shouldAllowAutoFix(service, options); - - if (typeof service.fixesApplied !== 'boolean') { - service.fixesApplied = false; - } - } - - if (typeof parsed.fixesApplied !== 'boolean') { - parsed.fixesApplied = false; - } - - return parsed; -} diff --git a/src/server/services/ai/utils/sanitize.ts b/src/server/services/ai/utils/sanitize.ts deleted file mode 100644 index ef4c3a92..00000000 --- a/src/server/services/ai/utils/sanitize.ts +++ /dev/null @@ -1,32 +0,0 @@ -/** - * Copyright 2025 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. - */ - -export function sanitizeForJson(obj: any): any { - if (typeof obj === 'string') { - return obj; - } else if (Array.isArray(obj)) { - return obj.map((item) => sanitizeForJson(item)); - } else if (obj && typeof obj === 'object') { - const sanitized: any = {}; - for (const key in obj) { - if (Object.prototype.hasOwnProperty.call(obj, key)) { - sanitized[key] = sanitizeForJson(obj[key]); - } - } - return sanitized; - } - return obj; -} diff --git a/src/server/services/aiAgent.ts b/src/server/services/aiAgent.ts deleted file mode 100644 index 183f56fc..00000000 --- a/src/server/services/aiAgent.ts +++ /dev/null @@ -1,378 +0,0 @@ -/** - * Copyright 2025 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 BaseService from './_service'; -import AIAgentConfigService from './aiAgentConfig'; -import { AIAgentCore } from './ai/service'; -import { DebugContext, DebugMessage, StructuredDebugResponse } from './types/aiAgent'; -import { StreamCallbacks } from './ai/types/stream'; -import { ProviderType } from './ai/providers/factory'; -import { - GetK8sResourcesTool, - GetPodLogsTool, - GetFileTool, - ListDirectoryTool, - QueryDatabaseTool, - GetCodefreshLogsTool, - UpdateFileTool, - UpdatePrLabelsTool, - PatchK8sResourceTool, - GetIssueCommentTool, -} from './ai/tools'; -import { getLogger } from 'server/lib/logger'; -import type { AIChatEvidenceEvent } from 'shared/types/aiChat'; -import { extractEvidence, generateResultPreview } from './ai/evidence/extractor'; - -export default class AIAgentService extends BaseService { - private service: AIAgentCore | null = null; - private provider: ProviderType = 'anthropic'; - private modelId?: string; - private modelPricing?: { inputCostPerMillion: number; outputCostPerMillion: number }; - private currentMode: 'investigate' | 'fix' = 'investigate'; - private repoFullName?: string; - - async initialize(repoFullName?: string): Promise { - await this.initializeWithMode('investigate', undefined, undefined, repoFullName); - } - - async initializeWithMode( - mode: 'investigate' | 'fix', - provider?: ProviderType, - modelId?: string, - repoFullName?: string - ): Promise { - getLogger().info(`AI: initializeWithMode called mode=${mode} provider=${provider} modelId=${modelId}`); - - if (repoFullName) { - this.repoFullName = repoFullName; - } - const aiAgentConfigService = AIAgentConfigService.getInstance(); - const aiAgentConfig = await aiAgentConfigService.getEffectiveConfig(this.repoFullName); - - if (!aiAgentConfig?.enabled) { - throw new Error('AI Agent feature is not enabled in global_config'); - } - - if (provider && modelId) { - getLogger().info(`AI: explicit provider selection provider=${provider} modelId=${modelId}`); - const providerConfig = aiAgentConfig.providers.find((p: any) => p.name === provider); - if (!providerConfig || !providerConfig.enabled) { - throw new Error(`Provider ${provider} is not enabled`); - } - - const modelConfig = providerConfig.models.find((m: any) => m.id === modelId); - if (!modelConfig || !modelConfig.enabled) { - throw new Error(`Model ${modelId} is not enabled for provider ${provider}`); - } - - this.provider = provider; - this.modelId = modelId; - this.modelPricing = - modelConfig.inputCostPerMillion != null && modelConfig.outputCostPerMillion != null - ? { - inputCostPerMillion: modelConfig.inputCostPerMillion, - outputCostPerMillion: modelConfig.outputCostPerMillion, - } - : undefined; - getLogger().info(`AI: provider set to ${this.provider} modelId=${this.modelId}`); - } else { - getLogger().info(`AI: no provider/modelId provided, using defaults`); - const enabledProvider = aiAgentConfig.providers.find((p: any) => p.enabled); - if (!enabledProvider) { - throw new Error('No enabled providers found in configuration'); - } - const defaultModel = enabledProvider.models.find((m: any) => m.default && m.enabled); - if (!defaultModel) { - throw new Error(`No default model found for provider ${enabledProvider.name}`); - } - this.provider = enabledProvider.name as ProviderType; - this.modelId = defaultModel.id; - this.modelPricing = - defaultModel.inputCostPerMillion != null && defaultModel.outputCostPerMillion != null - ? { - inputCostPerMillion: defaultModel.inputCostPerMillion, - outputCostPerMillion: defaultModel.outputCostPerMillion, - } - : undefined; - getLogger().info(`AI: using default provider=${this.provider} modelId=${this.modelId}`); - } - - this.currentMode = mode; - - getLogger().info(`AI: creating AIAgentCore with provider=${this.provider} modelId=${this.modelId}`); - this.service = new AIAgentCore({ - provider: this.provider, - modelId: this.modelId, - db: this.db, - redis: this.redis, - requireToolConfirmation: mode === 'investigate', - mode: mode, - additiveRules: aiAgentConfig.additiveRules, - systemPromptOverride: aiAgentConfig.systemPromptOverride, - excludedTools: aiAgentConfig.excludedTools, - excludedFilePatterns: aiAgentConfig.excludedFilePatterns, - allowedWritePatterns: aiAgentConfig.allowedWritePatterns, - modelPricing: this.modelPricing, - maxIterations: aiAgentConfig.maxIterations, - maxToolCalls: aiAgentConfig.maxToolCalls, - maxRepeatedCalls: aiAgentConfig.maxRepeatedCalls, - compressionThreshold: aiAgentConfig.compressionThreshold, - observationMaskingRecencyWindow: aiAgentConfig.observationMaskingRecencyWindow, - observationMaskingTokenThreshold: aiAgentConfig.observationMaskingTokenThreshold, - toolExecutionTimeout: aiAgentConfig.toolExecutionTimeout, - toolOutputMaxChars: aiAgentConfig.toolOutputMaxChars, - retryBudget: aiAgentConfig.retryBudget, - }); - } - - getMode(): 'investigate' | 'fix' { - return this.currentMode; - } - - private generateToolActivityMessage(toolName: string, toolArgs: any): string { - const parts: string[] = []; - - switch (toolName) { - case GetK8sResourcesTool.Name: - parts.push(`Checking ${toolArgs.resource_type || 'resources'}`); - if (toolArgs.namespace) parts.push(`in namespace ${toolArgs.namespace}`); - break; - - case GetPodLogsTool.Name: - parts.push(`Getting logs from pod ${toolArgs.pod_name}`); - if (toolArgs.namespace) parts.push(`in ${toolArgs.namespace}`); - break; - - case GetFileTool.Name: - parts.push(`Reading file ${toolArgs.file_path || 'from repository'}`); - if (toolArgs.repository_name) parts.push(`from ${toolArgs.repository_name}`); - break; - - case ListDirectoryTool.Name: - parts.push(`Listing files in ${toolArgs.directory_path || 'directory'}`); - if (toolArgs.repository_name) parts.push(`from ${toolArgs.repository_name}`); - break; - - case QueryDatabaseTool.Name: - parts.push(`Querying ${toolArgs.table} table`); - if (toolArgs.filters) parts.push(`with filters`); - break; - - case GetCodefreshLogsTool.Name: - if (toolArgs.service_name) { - parts.push(`Getting build logs for ${toolArgs.service_name}`); - } else { - parts.push(`Getting pipeline logs`); - } - break; - - case UpdateFileTool.Name: - parts.push(`Updating ${toolArgs.file_path || 'file'}`); - if (toolArgs.commit_message) parts.push(`- ${toolArgs.commit_message}`); - break; - - case UpdatePrLabelsTool.Name: - parts.push(`Updating PR #${toolArgs.pull_request_number} labels`); - if (toolArgs.action) parts.push(`(${toolArgs.action})`); - break; - - case PatchK8sResourceTool.Name: - parts.push(`Patching ${toolArgs.resource_type}/${toolArgs.name}`); - if (toolArgs.namespace) parts.push(`in ${toolArgs.namespace}`); - break; - - case GetIssueCommentTool.Name: - parts.push(`Reading issue/PR comments`); - if (toolArgs.repository_name) parts.push(`from ${toolArgs.repository_name}`); - break; - - default: { - parts.push(toolName.replace(/_/g, ' ')); - const importantArgs = Object.entries(toolArgs) - .filter(([key, val]) => val && !key.includes('token') && !key.includes('key')) - .slice(0, 2); - if (importantArgs.length > 0) { - const argStr = importantArgs.map(([k, v]) => `${k}: ${v}`).join(', '); - parts.push(`(${argStr})`); - } - break; - } - } - - return parts.join(' '); - } - - private generateToolResultMessage(toolName: string, toolArgs: any, result: any): string { - if (!result.success) { - return `Failed to ${this.generateToolActivityMessage(toolName, toolArgs).toLowerCase()}`; - } - - if (result.displayContent?.content) { - const content = result.displayContent.content; - if (typeof content === 'string' && content.length <= 100) { - return `✓ ${content}`; - } - } - - return `✓ ${this.generateToolActivityMessage(toolName, toolArgs)}`; - } - - async processQueryStream( - userMessage: string, - context: DebugContext, - conversationHistory: DebugMessage[], - onChunk: (chunk: string) => void, - onActivity?: (activity: { - type: string; - message: string; - details?: any; - toolCallId?: string; - resultPreview?: string; - }) => void, - onEvidence?: (event: AIChatEvidenceEvent) => void, - onToolConfirmation?: (details: { - title: string; - description: string; - impact: string; - confirmButtonText: string; - }) => Promise, - onToolAuthorization?: ( - tool: { - name: string; - description: string; - category: string; - safetyLevel: string; - }, - args: Record - ) => Promise<{ allowed: boolean; reason?: string }>, - mode?: 'investigate' | 'fix', - onDebugEvent?: (event: any) => void - ): Promise<{ - response: string; - isJson: boolean; - totalInvestigationTimeMs: number; - preamble?: string; - availableTools: Array<{ - name: string; - description: string; - category: string; - safetyLevel: string; - }>; - }> { - const effectiveMode = mode || 'investigate'; - - if (!this.service) { - getLogger().warn('AI: processQueryStream called without initialized service, using defaults'); - await this.initialize(this.repoFullName); - } - - if (this.service && effectiveMode !== this.getMode()) { - getLogger().info( - `AI: mode changed from ${this.getMode()} to ${effectiveMode}, reinitializing with same provider/model` - ); - await this.initializeWithMode(effectiveMode, this.provider, this.modelId, this.repoFullName); - } - - const abortController = new AbortController(); - - const callbacks: StreamCallbacks = { - onTextChunk: (text) => onChunk(text), - onThinking: (message) => onActivity?.({ type: 'thinking', message }), - onToolCall: (tool, args, toolCallId) => { - onActivity?.({ - type: 'tool_call', - message: this.generateToolActivityMessage(tool, args), - toolCallId, - }); - onDebugEvent?.({ type: 'debug_tool_call', toolCallId, toolName: tool, toolArgs: args }); - }, - onToolResult: (result, toolName, toolArgs, toolDurationMs, totalDurationMs, toolCallId) => { - const argsRecord = toolArgs as Record; - onActivity?.({ - type: 'processing', - message: this.generateToolResultMessage(toolName, toolArgs, result), - details: - toolDurationMs !== undefined || totalDurationMs !== undefined - ? { - toolDurationMs: toolDurationMs, - totalDurationMs: totalDurationMs, - } - : undefined, - toolCallId, - resultPreview: generateResultPreview(toolName, argsRecord, result), - }); - onDebugEvent?.({ type: 'debug_tool_result', toolCallId, toolName, toolResult: result, toolDurationMs }); - if (onEvidence) { - try { - const evidenceEvents = extractEvidence(toolName, argsRecord, result, { - toolCallId: toolCallId || '', - repositoryOwner: context.lifecycleContext?.pullRequest?.fullName?.split('/')[0], - repositoryName: context.lifecycleContext?.pullRequest?.fullName?.split('/')[1], - commitSha: context.lifecycleContext?.pullRequest?.latestCommit, - }); - for (const ev of evidenceEvents) { - onEvidence(ev); - } - } catch { - // evidence extraction must never disrupt the tool loop - } - } - }, - onError: (error) => onActivity?.({ type: 'error', message: error?.message || 'Error' }), - onActivity: (activity) => onActivity?.(activity), - onToolConfirmation: onToolConfirmation, - onToolAuthorization: onToolAuthorization, - }; - - const result = await this.service!.processQuery( - userMessage, - context, - conversationHistory, - callbacks, - abortController.signal, - onDebugEvent - ); - - return { - response: result.response, - isJson: result.isJson, - totalInvestigationTimeMs: result.metrics.duration, - preamble: result.preamble, - availableTools: result.availableTools, - }; - } - - parseStructuredResponse(response: string): StructuredDebugResponse | null { - try { - let cleaned = response.trim(); - - if (cleaned.startsWith('```')) { - cleaned = cleaned.replace(/^```(?:json)?\n?/, '').replace(/\n?```\s*$/, ''); - } - - const parsed = JSON.parse(cleaned); - - if (parsed.type === 'investigation_complete' && Array.isArray(parsed.services)) { - const fixesApplied = parsed.fixesApplied ?? false; - return { ...parsed, fixesApplied } as StructuredDebugResponse; - } - - return null; - } catch (error) { - return null; - } - } -} diff --git a/src/server/services/aiAgentConfig.ts b/src/server/services/aiAgentConfig.ts deleted file mode 100644 index da350cc0..00000000 --- a/src/server/services/aiAgentConfig.ts +++ /dev/null @@ -1,330 +0,0 @@ -/** - * Copyright 2025 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 BaseService from './_service'; -import GlobalConfigService from './globalConfig'; -import { getLogger } from 'server/lib/logger'; -import type { - AIAgentConfig, - AIAgentRepoOverride, - AIAgentRepoConfigRow, - ApprovalPolicyConfig, -} from './types/aiAgentConfig'; -import { validateAIAgentConfig, validateAIAgentRepoOverride } from 'server/lib/validation/aiAgentConfigValidator'; - -const REDIS_KEY_PREFIX = 'ai_agent_repo_config:'; - -export default class AIAgentConfigService extends BaseService { - private static instance: AIAgentConfigService; - - private memoryCache: Map = new Map(); - private globalCache: { data: AIAgentConfig | null; expiry: number } = { data: null, expiry: 0 }; - private static MEMORY_CACHE_TTL_MS = 30000; - - static getInstance(): AIAgentConfigService { - if (!this.instance) { - this.instance = new AIAgentConfigService(); - } - return this.instance; - } - - async getEffectiveConfig(repoFullName?: string): Promise { - const globalDefaults = await this.getGlobalDefaults(); - - if (!repoFullName) { - return globalDefaults; - } - - const normalized = repoFullName.toLowerCase(); - - const now = Date.now(); - const cached = this.memoryCache.get(normalized); - if (cached && now < cached.expiry) { - return cached.data; - } - - try { - const redisKey = `${REDIS_KEY_PREFIX}${normalized}`; - const redisValue = await this.redis.get(redisKey); - - if (redisValue) { - const repoOverride = JSON.parse(redisValue) as Partial; - const merged = this.mergeConfigs(globalDefaults, repoOverride); - this.memoryCache.set(normalized, { data: merged, expiry: now + AIAgentConfigService.MEMORY_CACHE_TTL_MS }); - return merged; - } - - const row = await this.db - .knex('ai_agent_repo_config') - .where({ repositoryFullName: normalized }) - .whereNull('deletedAt') - .first(); - - if (row) { - const config = typeof row.config === 'string' ? JSON.parse(row.config) : row.config; - await this.redis.set(redisKey, JSON.stringify(config), 'EX', 300); - const merged = this.mergeConfigs(globalDefaults, config as Partial); - this.memoryCache.set(normalized, { data: merged, expiry: now + AIAgentConfigService.MEMORY_CACHE_TTL_MS }); - return merged; - } - - return globalDefaults; - } catch (error) { - getLogger().warn(`AIAgentConfig: repo config lookup failed repo=${normalized} error=${error}`); - return globalDefaults; - } - } - - private async getGlobalDefaults(): Promise { - const now = Date.now(); - if (this.globalCache.data && now < this.globalCache.expiry) { - return this.globalCache.data; - } - - const config = await GlobalConfigService.getInstance().getConfig('aiAgent'); - this.globalCache = { data: config, expiry: now + AIAgentConfigService.MEMORY_CACHE_TTL_MS }; - return config; - } - - private mergeApprovalPolicy( - globalPolicy?: ApprovalPolicyConfig, - repoPolicy?: ApprovalPolicyConfig - ): ApprovalPolicyConfig | undefined { - if (!globalPolicy && !repoPolicy) { - return undefined; - } - - return { - defaultMode: repoPolicy?.defaultMode ?? globalPolicy?.defaultMode, - rules: { - ...(globalPolicy?.rules || {}), - ...(repoPolicy?.rules || {}), - }, - }; - } - - private mergeConfigs(global: AIAgentConfig, repoOverride: Partial): AIAgentConfig { - const result = { ...global }; - - if (repoOverride.enabled !== undefined) { - result.enabled = repoOverride.enabled; - } - if (repoOverride.maxMessagesPerSession !== undefined) { - result.maxMessagesPerSession = repoOverride.maxMessagesPerSession; - } - if (repoOverride.sessionTTL !== undefined) { - result.sessionTTL = repoOverride.sessionTTL; - } - result.approvalPolicy = this.mergeApprovalPolicy(global.approvalPolicy, repoOverride.approvalPolicy); - if (repoOverride.systemPromptOverride !== undefined) { - result.systemPromptOverride = repoOverride.systemPromptOverride; - } - const arrayFields: (keyof AIAgentRepoOverride)[] = [ - 'additiveRules', - 'excludedTools', - 'excludedFilePatterns', - 'allowedWritePatterns', - ]; - - for (const field of arrayFields) { - const globalArray = (global as any)[field] as string[] | undefined; - const repoArray = repoOverride[field] as string[] | undefined; - if (repoArray && repoArray.length > 0) { - const combined = [...(globalArray || []), ...repoArray]; - (result as any)[field] = [...new Set(combined)]; - } - } - - return result; - } - - async getGlobalConfig(): Promise { - const config = await GlobalConfigService.getInstance().getConfig('aiAgent'); - if (!config) { - return { - enabled: false, - providers: [], - maxMessagesPerSession: 50, - sessionTTL: 3600, - allowedWritePatterns: ['lifecycle.yaml', 'lifecycle.yml'], - }; - } - return config as AIAgentConfig; - } - - async setGlobalConfig(config: AIAgentConfig): Promise { - validateAIAgentConfig(config); - await GlobalConfigService.getInstance().setConfig('aiAgent', config); - this.invalidateCaches(); - getLogger().info('AIAgentConfig: global config updated via=api'); - } - - async updateGlobalAdditiveRules(additiveRules: string[]): Promise { - validateAIAgentRepoOverride({ additiveRules }); - - const currentConfig = await this.getGlobalConfig(); - const nextConfig: AIAgentConfig = { - ...currentConfig, - additiveRules, - }; - - await GlobalConfigService.getInstance().setConfig('aiAgent', nextConfig); - this.invalidateCaches(); - getLogger().info(`AIAgentConfig: global additive rules updated count=${additiveRules.length} via=api`); - - return nextConfig; - } - - async updateGlobalApprovalPolicy(approvalPolicy: ApprovalPolicyConfig): Promise { - validateAIAgentRepoOverride({ approvalPolicy }); - - const currentConfig = await this.getGlobalConfig(); - const nextApprovalPolicy = this.normalizeApprovalPolicy(approvalPolicy); - const nextConfig: AIAgentConfig = { - ...currentConfig, - }; - - if (nextApprovalPolicy) { - nextConfig.approvalPolicy = nextApprovalPolicy; - } else { - delete nextConfig.approvalPolicy; - } - - await GlobalConfigService.getInstance().setConfig('aiAgent', nextConfig); - this.invalidateCaches(); - getLogger().info('AIAgentConfig: global approval policy updated via=api'); - - return nextConfig; - } - - async listRepoConfigs(): Promise { - const rows = await this.db.knex('ai_agent_repo_config').whereNull('deletedAt').orderBy('repositoryFullName', 'asc'); - - return rows.map((row: any) => ({ - id: row.id, - repositoryFullName: row.repositoryFullName, - config: typeof row.config === 'string' ? JSON.parse(row.config) : row.config, - createdAt: row.createdAt, - updatedAt: row.updatedAt, - })); - } - - async getRepoConfig(repoFullName: string): Promise | null> { - const normalized = repoFullName.toLowerCase(); - const row = await this.db - .knex('ai_agent_repo_config') - .where({ repositoryFullName: normalized }) - .whereNull('deletedAt') - .first(); - - if (!row) { - return null; - } - - return typeof row.config === 'string' ? JSON.parse(row.config) : row.config; - } - - async setRepoConfig(repoFullName: string, config: Partial): Promise { - const normalized = repoFullName.toLowerCase(); - validateAIAgentRepoOverride(config); - await this.upsertRepoConfig(normalized, config); - } - - async updateRepoAdditiveRules(repoFullName: string, additiveRules: string[]): Promise> { - const normalized = repoFullName.toLowerCase(); - validateAIAgentRepoOverride({ additiveRules }); - - const currentConfig = (await this.getRepoConfig(normalized)) ?? {}; - const nextConfig: Partial = { - ...currentConfig, - additiveRules, - }; - - await this.upsertRepoConfig(normalized, nextConfig); - getLogger().info( - `AIAgentConfig: repo additive rules updated repo=${normalized} count=${additiveRules.length} via=api` - ); - - return nextConfig; - } - - private async upsertRepoConfig(normalizedRepoFullName: string, config: Partial): Promise { - await this.db - .knex('ai_agent_repo_config') - .insert({ - repositoryFullName: normalizedRepoFullName, - config: JSON.stringify(config), - createdAt: this.db.knex.fn.now(), - updatedAt: this.db.knex.fn.now(), - }) - .onConflict('repositoryFullName') - .merge({ - config: JSON.stringify(config), - updatedAt: this.db.knex.fn.now(), - deletedAt: null, - }); - - const redisKey = `${REDIS_KEY_PREFIX}${normalizedRepoFullName}`; - await this.redis.del(redisKey); - this.memoryCache.delete(normalizedRepoFullName); - } - - async deleteRepoConfig(repoFullName: string): Promise { - const normalized = repoFullName.toLowerCase(); - await this.db - .knex('ai_agent_repo_config') - .where({ repositoryFullName: normalized }) - .update({ deletedAt: this.db.knex.fn.now() }); - - const redisKey = `${REDIS_KEY_PREFIX}${normalized}`; - await this.redis.del(redisKey); - this.memoryCache.delete(normalized); - } - - clearCache(repoFullName?: string): void { - if (repoFullName) { - const normalized = repoFullName.toLowerCase(); - this.memoryCache.delete(normalized); - const redisKey = `${REDIS_KEY_PREFIX}${normalized}`; - this.redis.del(redisKey); - } else { - this.invalidateCaches(); - } - } - - private normalizeApprovalPolicy(approvalPolicy?: ApprovalPolicyConfig): ApprovalPolicyConfig | undefined { - if (!approvalPolicy) { - return undefined; - } - - const normalizedRules = - approvalPolicy.rules && Object.keys(approvalPolicy.rules).length > 0 ? approvalPolicy.rules : undefined; - - if (!approvalPolicy.defaultMode && !normalizedRules) { - return undefined; - } - - return { - ...(approvalPolicy.defaultMode ? { defaultMode: approvalPolicy.defaultMode } : {}), - ...(normalizedRules ? { rules: normalizedRules } : {}), - }; - } - - private invalidateCaches(): void { - this.globalCache = { data: null, expiry: 0 }; - this.memoryCache.clear(); - } -} diff --git a/src/server/services/types/agentRuntimeConfig.ts b/src/server/services/types/agentRuntimeConfig.ts new file mode 100644 index 00000000..64481b9f --- /dev/null +++ b/src/server/services/types/agentRuntimeConfig.ts @@ -0,0 +1,113 @@ +/** + * Copyright 2025 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 { AgentCapabilityAvailability, AgentCapabilityCatalogId } from 'server/services/agent/capabilityCatalog'; + +export interface AgentRuntimeModelConfig { + id: string; + displayName: string; + enabled: boolean; + default: boolean; + maxTokens: number; + inputCostPerMillion?: number; + outputCostPerMillion?: number; +} + +export interface AgentRuntimeProviderConfig { + name: string; + enabled: boolean; + apiKeyEnvVar: string; + models: AgentRuntimeModelConfig[]; +} + +export type ApprovalModeConfig = 'allow' | 'require_approval' | 'deny'; + +export interface ApprovalPolicyConfig { + defaultMode?: ApprovalModeConfig; + rules?: Partial< + Record< + | 'read' + | 'external_mcp_read' + | 'workspace_write' + | 'shell_exec' + | 'git_write' + | 'network_access' + | 'deploy_k8s_mutation' + | 'external_mcp_write', + ApprovalModeConfig + > + >; +} + +export interface CapabilityPolicyConfig { + availability?: Partial>; +} + +export type CustomAgentCreationMode = 'enabled' | 'disabled' | 'admins_only' | 'allowlist'; +export type CreatorCapabilityAvailability = 'available' | 'reserved'; + +export interface CustomAgentCreationPolicyConfig { + mode?: CustomAgentCreationMode; + allowedUserIds?: string[]; + allowedGithubUsernames?: string[]; + capabilityAvailability?: Partial>; +} + +export interface AgentRuntimeConfig { + enabled: boolean; + providers: AgentRuntimeProviderConfig[]; + maxMessagesPerSession: number; + sessionTTL: number; + approvalPolicy?: ApprovalPolicyConfig; + capabilityPolicy?: CapabilityPolicyConfig; + customAgentCreationPolicy?: CustomAgentCreationPolicyConfig; + additiveRules?: string[]; + systemPromptOverride?: string; + excludedTools?: string[]; + excludedFilePatterns?: string[]; + allowedWritePatterns?: string[]; + maxIterations?: number; + maxToolCalls?: number; + maxRepeatedCalls?: number; + compressionThreshold?: number; + observationMaskingRecencyWindow?: number; + observationMaskingTokenThreshold?: number; + toolExecutionTimeout?: number; + toolOutputMaxChars?: number; + retryBudget?: number; +} + +export interface AgentRuntimeRepoConfigRow { + id: number; + repositoryFullName: string; + config: Partial; + createdAt?: string; + updatedAt?: string; + deletedAt?: string | null; +} + +export interface AgentRuntimeRepoOverride { + enabled?: boolean; + maxMessagesPerSession?: number; + sessionTTL?: number; + approvalPolicy?: ApprovalPolicyConfig; + capabilityPolicy?: CapabilityPolicyConfig; + additiveRules?: string[]; + systemPromptOverride?: string; + excludedTools?: string[]; + excludedFilePatterns?: string[]; + allowedWritePatterns?: string[]; +} diff --git a/src/server/services/types/agentSessionConfig.ts b/src/server/services/types/agentSessionConfig.ts index e67de88a..a5c0d5d8 100644 --- a/src/server/services/types/agentSessionConfig.ts +++ b/src/server/services/types/agentSessionConfig.ts @@ -15,6 +15,11 @@ */ import type { AgentApprovalMode, AgentCapabilityKey } from 'server/services/agent/types'; +import type { + AgentCapabilityAvailability, + AgentCapabilityCatalogId, + AgentCapabilityCategory, +} from 'server/services/agent/capabilityCatalog'; import type { AgentSessionWorkspaceStorageAccessMode } from './globalConfig'; export type AgentSessionToolRuleMode = AgentApprovalMode; @@ -110,3 +115,32 @@ export interface AgentSessionToolInventoryEntry { effectiveRuleMode: AgentSessionToolRuleSelection; availability: 'available' | 'blocked_by_tool_rule' | 'blocked_by_policy'; } + +export interface AgentCapabilityInventoryToolEntry { + toolKey: string; + toolName: string; + description: string | null; + serverSlug: string; + serverName: string; + sourceType: 'builtin' | 'mcp'; + sourceScope: string; +} + +export interface AgentCapabilityInventoryEntry { + capabilityId: AgentCapabilityCatalogId; + label: string; + description: string; + category: AgentCapabilityCategory; + defaultAvailability: AgentCapabilityAvailability; + configuredAvailability?: AgentCapabilityAvailability; + inheritedAvailability?: AgentCapabilityAvailability; + effectiveAvailability: AgentCapabilityAvailability; + approvalMode: AgentApprovalMode; + runtimeCapabilityKey?: AgentCapabilityKey; + userSelectable: boolean; + toolCount: number; + resourceCount: number; + resourceGrants: string[]; + tools: AgentCapabilityInventoryToolEntry[]; + blockedReason?: 'admin_only' | 'system_only' | 'disabled'; +} diff --git a/src/server/services/types/aiAgent.ts b/src/server/services/types/aiAgent.ts deleted file mode 100644 index 623f98eb..00000000 --- a/src/server/services/types/aiAgent.ts +++ /dev/null @@ -1,366 +0,0 @@ -/** - * Copyright 2025 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 { BuildStatus, DeployStatus, DeployTypes, PullRequestStatus } from 'shared/constants'; - -export interface ConversationState { - buildUuid: string; - messages: DebugMessage[]; - lastActivity: number; - contextSnapshot?: Partial; -} - -export interface DebugMessage { - role: 'user' | 'assistant' | 'system'; - content: string; - timestamp: number; - isSystemAction?: boolean; - activityHistory?: Array<{ - type: string; - message: string; - status?: 'pending' | 'completed' | 'failed'; - details?: { toolDurationMs?: number; totalDurationMs?: number }; - toolCallId?: string; - resultPreview?: string; - }>; - evidenceItems?: Array>; - totalInvestigationTimeMs?: number; - debugContext?: { - systemPrompt: string; - maskingStats: { - totalTokensBefore: number; - totalTokensAfter: number; - maskedParts: number; - savedTokens: number; - } | null; - provider: string; - modelId: string; - }; - debugToolData?: Array<{ - toolCallId: string; - toolName: string; - toolArgs: Record; - toolResult?: unknown; - toolDurationMs?: number; - }>; - debugMetrics?: { - iterations: number; - totalToolCalls: number; - totalDurationMs: number; - inputTokens?: number; - outputTokens?: number; - inputCostPerMillion?: number; - outputCostPerMillion?: number; - }; -} - -export interface DebugContext { - buildUuid: string; - namespace: string; - lifecycleContext: LifecycleContext; - lifecycleYaml?: { - path: string; - content: string; - error?: string; - }; - services: ServiceDebugInfo[]; - gatheredAt: Date; - warnings?: ContextWarning[]; - errors?: ContextError[]; -} - -export interface ContextWarning { - source: 'kubernetes' | 'lifecycle' | 'logs'; - message: string; - details?: string; -} - -export interface ContextError { - source: 'kubernetes' | 'lifecycle' | 'logs'; - message: string; - error: string; - recoverable: boolean; -} - -export interface LifecycleContext { - build: BuildInfo; - pullRequest: PullRequestInfo; - environment: EnvironmentInfo; - deploys: DeployInfo[]; - repository: RepositoryInfo; -} - -export interface BuildInfo { - uuid: string; - status: BuildStatus; - statusMessage: string; - namespace: string; - sha: string; - trackDefaultBranches: boolean; - capacityType: string; - enabledFeatures: string[]; - dependencyGraph: Record; - dashboardLinks: Record; - createdAt: Date; - updatedAt: Date; -} - -export interface PullRequestInfo { - number: number; - title: string; - username: string; - branch: string; - baseBranch: string; - status: PullRequestStatus; - url: string; - latestCommit: string; - fullName: string; - commentId?: number; - labels: string[]; -} - -export interface EnvironmentInfo { - id: number; - name: string; - config: Record; -} - -export interface DeployInfo { - uuid: string; - serviceName: string; - status: DeployStatus; - statusMessage: string; - type: DeployTypes; - dockerImage: string; - branch: string; - repoName: string; - buildNumber: number; - buildPipelineId?: string; - deployPipelineId?: string; - builderEngine?: string; - helmChart?: string; - repositoryId?: number; - env: Record; - initEnv: Record; - createdAt: Date; - updatedAt: Date; -} - -export interface DeployableInfo { - serviceName: string; - type: DeployTypes; - repositoryId: number; - defaultBranchName: string; - commentBranchName: string; - helm: any; - deploymentDependsOn: string[]; - builder: any; -} - -export interface RepositoryInfo { - name: string; - githubRepositoryId: number; - url: string; -} - -export interface ServiceDebugInfo { - name: string; - type: string; - status: 'pending' | 'building' | 'deploying' | 'running' | 'failed'; - deployInfo: DeployInfo; - deployment?: K8sDeployment; - pods: PodDebugInfo[]; - events: K8sEvent[]; - issues: DiagnosedIssue[]; -} - -export interface K8sDeployment { - name: string; - replicas: { - desired: number; - current: number; - ready: number; - available: number; - }; - conditions: any[]; - strategy: string; - containers: Array<{ - name: string; - image: string; - }>; -} - -export interface PodDebugInfo { - name: string; - phase: string; - conditions: any[]; - containerStatuses: any[]; - recentLogs: string; - events: K8sEvent[]; -} - -export interface K8sEvent { - type: string; - reason: string; - message: string; - count: number; - firstTimestamp: Date; - lastTimestamp: Date; -} - -export interface DiagnosedIssue { - severity: 'critical' | 'warning' | 'info'; - category: 'image' | 'resources' | 'configuration' | 'network' | 'other'; - title: string; - description: string; - suggestedFix: string; - detectedBy: 'rules' | 'llm'; - suggestedActions?: SuggestedAction[]; -} - -export interface SuggestedAction { - id: string; - label: string; - description: string; - action: 'kubectl_patch' | 'kubectl_delete' | 'kubectl_restart' | 'kubectl_scale' | 'view_logs' | 'view_events'; - params: { - resourceType?: string; - resourceName?: string; - namespace?: string; - patch?: any; - replicas?: number; - containerName?: string; - }; - confirmation?: string; - dangerous?: boolean; -} - -export interface InvestigationStep { - stepNumber: number; - tool: string; - args: Record; - completed: boolean; - dependsOn?: number; - result?: any; -} - -export interface ServiceInvestigationPlan { - serviceName: string; - status: DeployStatus.BUILD_FAILED | DeployStatus.DEPLOY_FAILED; - type: 'codefresh' | 'native'; - steps: InvestigationStep[]; - completed: boolean; -} - -export interface InvestigationPlan { - services: ServiceInvestigationPlan[]; - allCompleted: boolean; -} - -export interface FileChange { - path: string; // The exact file path - lineNumber?: number; - lineNumberEnd?: number; - description?: string; // What's being changed in this specific file - oldContent?: string; // For multi-line changes - the current content - newContent?: string; // For multi-line changes - what it should be -} - -export interface ServiceInvestigationResult { - serviceName: string; - status: DeployStatus.BUILD_FAILED | DeployStatus.DEPLOY_FAILED; - issue: string; - keyError?: string; - errorSource?: string; - errorSourceDetail?: string; - suggestedFix: string; - canAutoFix?: boolean; - filePath?: string; // The primary file path (kept for backward compatibility) - lineNumber?: number; - lineNumberEnd?: number; - files?: FileChange[]; // For multi-file fixes - if present, use this instead of single file fields - commitUrl?: string; -} - -export interface StructuredDebugResponse { - type: 'investigation_complete'; - summary: string; - fixesApplied: boolean; - services: ServiceInvestigationResult[]; - repository?: { - owner: string; - name: string; - branch: string; - }; -} - -export interface ModelConfig { - id: string; - displayName: string; - enabled: boolean; - default: boolean; - maxTokens: number; - inputCostPerMillion?: number; - outputCostPerMillion?: number; -} - -export interface ProviderConfig { - name: string; - enabled: boolean; - apiKeyEnvVar: string; - models: ModelConfig[]; -} - -export type ApprovalModeConfig = 'allow' | 'require_approval' | 'deny'; - -export interface ApprovalPolicyConfig { - defaultMode?: ApprovalModeConfig; - rules?: Partial< - Record< - | 'read' - | 'external_mcp_read' - | 'workspace_write' - | 'shell_exec' - | 'git_write' - | 'network_access' - | 'deploy_k8s_mutation' - | 'external_mcp_write', - ApprovalModeConfig - > - >; -} - -export interface AIAgentConfig { - enabled: boolean; - providers: ProviderConfig[]; - maxMessagesPerSession: number; - sessionTTL: number; - approvalPolicy?: ApprovalPolicyConfig; - additiveRules?: string[]; - systemPromptOverride?: string; - excludedTools?: string[]; - excludedFilePatterns?: string[]; - allowedWritePatterns?: string[]; - maxIterations?: number; - maxToolCalls?: number; - maxRepeatedCalls?: number; - compressionThreshold?: number; - observationMaskingRecencyWindow?: number; - observationMaskingTokenThreshold?: number; - toolExecutionTimeout?: number; - toolOutputMaxChars?: number; - retryBudget?: number; -} diff --git a/src/server/services/types/aiAgentConfig.ts b/src/server/services/types/aiAgentConfig.ts deleted file mode 100644 index 04f09710..00000000 --- a/src/server/services/types/aiAgentConfig.ts +++ /dev/null @@ -1,40 +0,0 @@ -/** - * Copyright 2025 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. - */ - -export type { AIAgentConfig, ApprovalPolicyConfig, ProviderConfig, ModelConfig } from './aiAgent'; - -export interface AIAgentRepoConfigRow { - id: number; - repositoryFullName: string; - config: Partial; - createdAt?: string; - updatedAt?: string; - deletedAt?: string | null; -} - -export interface AIAgentRepoOverride { - enabled?: boolean; - maxMessagesPerSession?: number; - sessionTTL?: number; - approvalPolicy?: import('./aiAgent').ApprovalPolicyConfig; - additiveRules?: string[]; - systemPromptOverride?: string; - excludedTools?: string[]; - excludedFilePatterns?: string[]; - allowedWritePatterns?: string[]; -} - -export type EffectiveAIAgentConfig = import('./aiAgent').AIAgentConfig; diff --git a/src/server/services/userMcpConnection.ts b/src/server/services/userMcpConnection.ts index 158fd52a..94850d1e 100644 --- a/src/server/services/userMcpConnection.ts +++ b/src/server/services/userMcpConnection.ts @@ -17,13 +17,13 @@ import 'server/lib/dependencies'; import UserMcpConnection from 'server/models/UserMcpConnection'; import { decrypt, encrypt } from 'server/lib/encryption'; -import { normalizeUserConnectionValues } from 'server/services/ai/mcp/connectionConfig'; +import { normalizeUserConnectionValues } from 'server/services/agentRuntime/mcp/connectionConfig'; import type { McpDiscoveredTool, McpStoredUserConnectionState, UserMcpConnectionMaskedUser, UserMcpConnectionState, -} from 'server/services/ai/mcp/types'; +} from 'server/services/agentRuntime/mcp/types'; type DecryptedUserMcpConnection = { state: McpStoredUserConnectionState | null; diff --git a/src/shared/openApiSpec.test.ts b/src/shared/openApiSpec.test.ts index 416b4eba..11d90929 100644 --- a/src/shared/openApiSpec.test.ts +++ b/src/shared/openApiSpec.test.ts @@ -8,6 +8,10 @@ function getJsonErrorSchema(path: string, method: string, status: string) { return swaggerSpec.paths[path][method].responses[status].content['application/json'].schema; } +function getOperation(path: string, method: string) { + return swaggerSpec.paths[path]?.[method]; +} + describe('OpenAPI v2 agent session contract', () => { it('documents canonical run events with public context and a version', () => { const eventSchema = schemas.AgentRunMessagePartEvent; @@ -20,6 +24,26 @@ describe('OpenAPI v2 agent session contract', () => { expect(eventSchema.properties.version).toEqual({ type: 'integer', enum: [1] }); }); + it('uses one canonical file-change artifact for run events and pending actions', () => { + expect(schemas.AgentFileChangeData.required).toEqual( + expect.arrayContaining(['id', 'toolCallId', 'sourceTool', 'path', 'displayPath', 'kind', 'stage']) + ); + expect(schemas.AgentFileChangeData.properties.kind.enum).toEqual(['created', 'edited', 'deleted']); + expect(schemas.AgentFileChangeData.properties.stage.enum).toEqual([ + 'awaiting-approval', + 'approved', + 'applied', + 'denied', + 'failed', + ]); + expect(schemas.AgentRunToolFileChangeEvent.properties.payload.oneOf[0].properties.data).toEqual({ + $ref: '#/components/schemas/AgentFileChangeData', + }); + expect(schemas.AgentPendingAction.properties.fileChangePreview.items).toEqual({ + $ref: '#/components/schemas/AgentFileChangeData', + }); + }); + 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'); @@ -41,14 +65,250 @@ describe('OpenAPI v2 agent session contract', () => { ); }); + it('documents the public run-plan summary without internal snapshot fields', () => { + expect(schemas.AgentRunPlanSummary.required).toEqual( + expect.arrayContaining(['runtime', 'approval', 'capabilities']) + ); + expect(schemas.AgentRunPlanRuntimeSummary).toBeDefined(); + expect(schemas.AgentRunPlanApprovalSummary).toBeDefined(); + expect(schemas.AgentRunPlanCapabilitySummary).toBeDefined(); + expect(schemas.AgentRunPlanCapabilitiesSummary).toBeDefined(); + expect(schemas.AgentRunPlanSelectedRuntimeChoicesSummary).toBeDefined(); + expect(schemas.AgentRunPlanSummary.properties.runtime).toEqual({ + $ref: '#/components/schemas/AgentRunPlanRuntimeSummary', + }); + expect(schemas.AgentRunPlanSummary.properties.approval).toEqual({ + $ref: '#/components/schemas/AgentRunPlanApprovalSummary', + }); + expect(schemas.AgentRunPlanSummary.properties.capabilities).toEqual({ + $ref: '#/components/schemas/AgentRunPlanCapabilitiesSummary', + }); + + const runPlanSchemas = JSON.stringify({ + AgentRunPlanSummary: schemas.AgentRunPlanSummary, + AgentRunPlanRuntimeSummary: schemas.AgentRunPlanRuntimeSummary, + AgentRunPlanApprovalSummary: schemas.AgentRunPlanApprovalSummary, + AgentRunPlanCapabilitySummary: schemas.AgentRunPlanCapabilitySummary, + AgentRunPlanCapabilitiesSummary: schemas.AgentRunPlanCapabilitiesSummary, + AgentRunPlanSelectedRuntimeChoicesSummary: schemas.AgentRunPlanSelectedRuntimeChoicesSummary, + }); + + for (const forbidden of [ + 'renderedHash', + 'renderedSummary', + `selectedRuntime${'Mcp'}ConnectionRefs`, + 'runtimeCapabilityKey', + 'approvalPolicy', + ]) { + expect(runPlanSchemas).not.toContain(forbidden); + } + }); + + it('uses unified agent platform tags for user-facing agent routes', () => { + for (const [path, method] of [ + ['/api/v2/ai/agent/definitions', 'get'], + ['/api/v2/ai/agent/definitions', 'post'], + ['/api/v2/ai/agent/definitions/{definitionId}', 'get'], + ['/api/v2/ai/agent/definitions/{definitionId}', 'patch'], + ['/api/v2/ai/agent/definitions/{definitionId}', 'delete'], + ['/api/v2/ai/agent/definition-capabilities', 'get'], + ['/api/v2/ai/agent/threads/{threadId}/agent', 'get'], + ['/api/v2/ai/agent/threads/{threadId}/agent', 'patch'], + ['/api/v2/ai/agent/threads/{threadId}/runtime-controls', 'get'], + ['/api/v2/ai/agent/threads/{threadId}/runtime-controls', 'patch'], + ['/api/v2/ai/agent/runtime-controls/preview', 'post'], + ['/api/v2/ai/agent/sessions', 'post'], + ['/api/v2/ai/agent/threads/{threadId}/runs', 'post'], + ['/api/v2/ai/agent/runs/{runId}', 'get'], + ]) { + expect(getOperation(path, method)?.tags).toEqual(['Agent Platform']); + } + }); + + it('keeps create-run and run-detail schemas aligned to public run-plan contracts', () => { + expect(getOperation('/api/v2/ai/agent/threads/{threadId}/runs', 'post')?.description).toContain( + "resolves its run plan server-side from the thread's selected agent" + ); + expect(Object.keys(schemas.CreateAgentThreadRunRequest.properties).sort()).toEqual([ + 'message', + 'model', + 'runtimeOptions', + ]); + expect(schemas.CreateAgentThreadRunRequest.additionalProperties).toBe(false); + expect(schemas.AgentRun.properties.runPlan).toEqual({ + $ref: '#/components/schemas/AgentRunPlanSummary', + }); + }); + + it('documents exact thread usage without raw provider internals', () => { + expect(getOperation('/api/v2/ai/agent/threads/{threadId}/usage', 'get')?.tags).toEqual(['Agent Platform']); + expect(schemas.AgentUsageSummary.required).toEqual(['totalTokens']); + expect(Object.keys(schemas.AgentUsageSummary.properties).sort()).toEqual([ + 'cacheCreationInputTokens', + 'cacheReadInputTokens', + 'cachedInputTokens', + 'inputTokens', + 'nonCachedInputTokens', + 'outputTokens', + 'reasoningTokens', + 'textOutputTokens', + 'totalTokens', + ]); + expect(schemas.AgentUsageByModel.required).toEqual([ + 'provider', + 'model', + 'totalTokens', + 'runCount', + 'reportedRunCount', + 'missingUsageRunCount', + ]); + expect(schemas.AgentUsageAggregate.required).toEqual(['usageSummary', 'usageByModel', 'usageCompleteness']); + expect(schemas.AgentThreadUsageResponse.required).toEqual([ + 'threadId', + 'sessionId', + 'usageSummary', + 'usageByModel', + 'usageCompleteness', + ]); + expect(getJsonErrorSchema('/api/v2/ai/agent/threads/{threadId}/usage', 'get', '404')).toEqual({ + $ref: '#/components/schemas/ApiErrorResponse', + }); + + const usageSchemas = JSON.stringify({ + AgentUsageSummary: schemas.AgentUsageSummary, + AgentUsageByModel: schemas.AgentUsageByModel, + AgentUsageCompleteness: schemas.AgentUsageCompleteness, + AgentUsageAggregate: schemas.AgentUsageAggregate, + AgentThreadUsageResponse: schemas.AgentThreadUsageResponse, + }); + expect(usageSchemas).not.toContain('rawUsage'); + expect(usageSchemas).not.toContain('providerMetadata'); + }); + + it('documents session summaries with lifetime usage', () => { + expect(schemas.AgentSessionSummary.required).toEqual(['session', 'source', 'sandbox', 'usage']); + expect(schemas.AgentSessionSummary.properties.usage).toEqual({ + $ref: '#/components/schemas/AgentUsageAggregate', + }); + }); + + it('keeps admin capability policy docs admin-scoped and payload-compatible', () => { + const adminPath = '/api/v2/ai/admin/agent/capabilities'; + + expect(getOperation(adminPath, 'get')?.tags).toEqual(['Agent Admin']); + expect(getOperation(adminPath, 'put')?.tags).toEqual(['Agent Admin']); + expect(getOperation(adminPath, 'get')?.summary).toBe('Get agent capability policy inventory'); + expect(getOperation(adminPath, 'put')?.summary).toBe('Update agent capability policy'); + expect(schemas.UpdateAdminAgentCapabilitiesRequest.required).toEqual(['capabilityPolicy']); + expect(schemas.AgentCapabilityInventoryEntry.required).toContain('resourceGrants'); + expect(schemas.AgentCapabilityInventoryEntry.properties.resourceGrants).toEqual({ + type: 'array', + items: { type: 'string' }, + }); + }); + + it('documents custom-agent creation policy as an admin-only management contract', () => { + const adminPath = '/api/v2/ai/admin/agent/creation-policy'; + + expect(getOperation(adminPath, 'get')?.tags).toEqual(['Agent Admin']); + expect(getOperation(adminPath, 'put')?.tags).toEqual(['Agent Admin']); + expect(getOperation(adminPath, 'get')?.summary).toBe('Get custom-agent creation policy'); + expect(getOperation(adminPath, 'put')?.summary).toBe('Update custom-agent creation policy'); + expect(schemas.UpdateAdminCustomAgentCreationPolicyRequest.required).toEqual(['customAgentCreationPolicy']); + expect(schemas.CustomAgentCreationPolicy.properties.mode).toEqual({ + $ref: '#/components/schemas/CustomAgentCreationMode', + }); + expect(schemas.CreatorCapabilityAvailability.enum).toEqual(['available', 'reserved']); + expect(schemas.AgentRuntimeConfig.properties.customAgentCreationPolicy).toEqual({ + $ref: '#/components/schemas/CustomAgentCreationPolicy', + }); + }); + + it('documents user custom-agent creator eligibility on the public capabilities contract', () => { + expect(schemas.UserAgentDefinitionCapabilitiesResponse.required).toEqual([ + 'resourceBehavior', + 'canCreate', + 'creationUnavailableReason', + 'capabilities', + ]); + expect(schemas.UserAgentDefinitionCapabilitiesResponse.properties.canCreate).toEqual({ type: 'boolean' }); + expect(schemas.UserAgentDefinitionCapabilitiesResponse.properties.creationUnavailableReason).toEqual({ + $ref: '#/components/schemas/CustomAgentCreationUnavailableReason', + }); + expect(schemas.CustomAgentCreationUnavailableReason.enum).toEqual([ + 'creation_disabled', + 'creation_restricted', + null, + ]); + expect(schemas.CustomAgentCreationUnavailableReason.nullable).toBe(true); + }); + + it('documents custom-agent create and update policy denials as safe JSON 403 responses', () => { + expect(getJsonErrorSchema('/api/v2/ai/agent/definitions', 'post', '403')).toEqual({ + $ref: '#/components/schemas/ApiErrorResponse', + }); + expect(getJsonErrorSchema('/api/v2/ai/agent/definitions/{definitionId}', 'patch', '403')).toEqual({ + $ref: '#/components/schemas/ApiErrorResponse', + }); + }); + 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/threads/{threadId}/preset']).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(); + expect(schemas.AgentThreadPresetState).toBeUndefined(); + expect(schemas.SwitchAgentThreadPresetRequest).toBeUndefined(); + }); + + it('does not advertise legacy AI chat execution, history, session, or model contracts', () => { + expect(swaggerSpec.paths['/api/v2/ai/chat/{buildUuid}']).toBeUndefined(); + expect(swaggerSpec.paths['/api/v2/ai/chat/{buildUuid}/messages']).toBeUndefined(); + expect(swaggerSpec.paths['/api/v2/ai/chat/{buildUuid}/session']).toBeUndefined(); + expect(swaggerSpec.paths['/api/v2/ai/chat/{buildUuid}/feedback']).toBeUndefined(); + expect(swaggerSpec.paths['/api/v2/ai/chat/{buildUuid}/messages/{messageId}/feedback']).toBeUndefined(); + expect(swaggerSpec.paths['/api/v2/ai/models']).toBeUndefined(); + expect(swaggerSpec.paths['/api/v2/ai/admin/feedback']).toBeUndefined(); + expect(swaggerSpec.paths['/api/v2/ai/admin/feedback/{id}/conversation']).toBeUndefined(); + expect(schemas.AIModel).toBeUndefined(); + expect(schemas.GetAIModelsSuccessResponse).toBeUndefined(); + expect(schemas.ConversationMessage).toBeUndefined(); + expect(schemas.ActivityHistoryEntry).toBeUndefined(); + expect(schemas.DebugContext).toBeUndefined(); + expect(schemas.DebugToolData).toBeUndefined(); + expect(schemas.DebugMetrics).toBeUndefined(); + expect(schemas.GetAIMessagesSuccessResponse).toBeUndefined(); + expect(schemas.DeleteAISessionSuccessResponse).toBeUndefined(); + expect(schemas.SSEChunkEvent).toBeUndefined(); + expect(schemas.SSEErrorEvent).toBeUndefined(); + expect(schemas.FeedbackEntry).toBeUndefined(); + expect(schemas.FeedbackListPaginationMetadata).toBeUndefined(); + expect(schemas.FeedbackListResponseMetadata).toBeUndefined(); + expect(schemas.GetAdminFeedbackListSuccessResponse).toBeUndefined(); + expect(schemas.ConversationReplayMessage).toBeUndefined(); + expect(schemas.FeedbackConversationReplayData).toBeUndefined(); + expect(schemas.GetAdminFeedbackConversationSuccessResponse).toBeUndefined(); + }); + + it('keeps replacement build-context chat, agent model, and runtime config contracts', () => { + expect(swaggerSpec.paths['/api/v2/ai/agent/build-context-chats']?.post).toBeDefined(); + expect(swaggerSpec.paths['/api/v2/ai/agent/models']?.get).toBeDefined(); + expect(swaggerSpec.paths['/api/v2/ai/config']?.get).toBeDefined(); + expect(swaggerSpec.paths['/api/v2/ai/agent/runtime-config']?.get).toBeDefined(); + expect(swaggerSpec.paths['/api/v2/ai/agent/runtime-config/repos']?.get).toBeDefined(); + expect(swaggerSpec.paths['/api/v2/ai/agent/runtime-config/repos/{owner}/{repo}']?.put).toBeDefined(); + expect(swaggerSpec.paths['/api/v2/ai/agent-config']).toBeUndefined(); + expect(swaggerSpec.paths['/api/v2/ai/agent-config/repos']).toBeUndefined(); + expect(schemas.BuildContextAgentChatResponse).toBeDefined(); + expect(schemas.AgentModel).toBeDefined(); + expect(schemas.AIConfigStatus).toBeDefined(); + expect(schemas.GetAIConfigSuccessResponse).toBeDefined(); + expect(schemas.AgentRuntimeConfig).toBeDefined(); + expect(schemas.AgentRuntimeRepoOverride).toBeDefined(); }); it('documents JSON error responses for changed canonical endpoints', () => { diff --git a/src/shared/openApiSpec.ts b/src/shared/openApiSpec.ts index 0b37e0c4..f18735b7 100644 --- a/src/shared/openApiSpec.ts +++ b/src/shared/openApiSpec.ts @@ -83,6 +83,18 @@ const agentRunEventPayloadMetadata = { additionalProperties: true, }; +const agentUsageSummaryProperties = { + totalTokens: { type: 'number' }, + inputTokens: { type: 'number' }, + outputTokens: { type: 'number' }, + reasoningTokens: { type: 'number' }, + cachedInputTokens: { type: 'number' }, + cacheCreationInputTokens: { type: 'number' }, + cacheReadInputTokens: { type: 'number' }, + nonCachedInputTokens: { type: 'number' }, + textOutputTokens: { type: 'number' }, +}; + export const openApiSpecificationForV2Api: OAS3Options = { definition: { openapi: '3.0.0', @@ -516,12 +528,15 @@ export const openApiSpecificationForV2Api: OAS3Options = { clientMessageId: { type: 'string', nullable: true }, threadId: { type: 'string' }, runId: { type: 'string', nullable: true }, - role: { type: 'string', enum: ['user', 'assistant'] }, + role: { type: 'string', enum: ['user', 'assistant', 'system'] }, parts: { type: 'array', items: { $ref: '#/components/schemas/CanonicalAgentMessagePart' }, minItems: 1, }, + metadata: { + oneOf: [{ $ref: '#/components/schemas/AgentSwitchEventMetadata' }, { type: 'object' }], + }, createdAt: { type: 'string', format: 'date-time', nullable: true }, }, required: ['id', 'clientMessageId', 'threadId', 'runId', 'role', 'parts', 'createdAt'], @@ -537,242 +552,901 @@ export const openApiSpecificationForV2Api: OAS3Options = { }, }, - AgentThreadMessagesResponse: { + SystemAgentDefinitionId: { + type: 'string', + enum: ['system.debug', 'system.develop', 'system.freeform'], + }, + + AgentSelectionSummary: { type: 'object', properties: { - 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, + id: { type: 'string' }, + ownerKind: { type: 'string', enum: ['system', 'admin', 'user'] }, + label: { type: 'string' }, + description: { type: 'string', nullable: true }, + available: { type: 'boolean' }, + unavailableReason: { + type: 'string', + nullable: true, + enum: [ + 'unknown_agent', + 'active_run', + 'disabled_agent', + 'requires_workspace', + 'source_incompatible', + 'disabled_by_policy', + null, + ], }, + unavailableMessage: { type: 'string', nullable: true }, + group: { type: 'string', enum: ['built_in', 'my_agents'] }, }, - required: ['thread', 'messages', 'pagination'], + required: [ + 'id', + 'ownerKind', + 'label', + 'description', + 'available', + 'unavailableReason', + 'unavailableMessage', + 'group', + ], additionalProperties: false, }, - AgentRunRuntimeOptions: { + AgentSelectionGroup: { type: 'object', properties: { - maxIterations: { - type: 'integer', - minimum: 1, - maximum: 100, + id: { type: 'string', enum: ['built_in', 'my_agents'] }, + label: { type: 'string' }, + agents: { + type: 'array', + items: { $ref: '#/components/schemas/AgentSelectionSummary' }, }, }, + required: ['id', 'label', 'agents'], 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: { + AgentSelectionState: { type: 'object', properties: { - clientMessageId: { type: 'string' }, - parts: { + selectedId: { type: 'string', nullable: true }, + defaultId: { $ref: '#/components/schemas/SystemAgentDefinitionId' }, + currentId: { type: 'string' }, + groups: { type: 'array', - items: { $ref: '#/components/schemas/CanonicalAgentMessagePart' }, - minItems: 1, + items: { $ref: '#/components/schemas/AgentSelectionGroup' }, + minItems: 2, }, }, - required: ['parts'], + required: ['selectedId', 'defaultId', 'currentId', 'groups'], additionalProperties: false, }, - CreateAgentThreadRunRequest: { + SwitchAgentSelectionRequest: { 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' }, + agentId: { type: 'string' }, }, - required: ['message'], + required: ['agentId'], 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, - }, - }, }, - AgentThread: { + SwitchAgentSelectionResponse: { type: 'object', properties: { - id: { type: 'string' }, - sessionId: { type: 'string', nullable: true }, - title: { type: 'string', nullable: true }, - isDefault: { type: 'boolean' }, - archivedAt: { type: 'string', format: 'date-time', nullable: true }, - lastRunAt: { type: 'string', format: 'date-time', nullable: true }, - metadata: { type: 'object', additionalProperties: true }, - createdAt: { type: 'string', format: 'date-time', nullable: true }, - updatedAt: { type: 'string', format: 'date-time', nullable: true }, + previousAgent: { $ref: '#/components/schemas/AgentSelectionSummary' }, + nextAgent: { $ref: '#/components/schemas/AgentSelectionSummary' }, + switched: { type: 'boolean' }, + state: { $ref: '#/components/schemas/AgentSelectionState' }, }, - required: ['id', 'isDefault', 'metadata'], + required: ['previousAgent', 'nextAgent', 'switched', 'state'], + additionalProperties: false, }, - AgentSessionDefaults: { + AgentThreadRuntimeControlChoice: { type: 'object', properties: { - model: { type: 'string' }, - harness: { type: 'string', nullable: true }, + id: { type: 'string' }, + label: { type: 'string' }, + description: { type: 'string', nullable: true }, + required: { type: 'boolean' }, + selected: { type: 'boolean' }, + available: { type: 'boolean' }, }, - required: ['model', 'harness'], + required: ['id', 'label', 'description', 'required', 'selected', 'available'], + additionalProperties: false, }, - AgentSource: { + AgentThreadRuntimeControlsState: { 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 }, + tools: { + type: 'object', + properties: { + required: { + type: 'array', + items: { $ref: '#/components/schemas/AgentThreadRuntimeControlChoice' }, + }, + optional: { + type: 'array', + items: { $ref: '#/components/schemas/AgentThreadRuntimeControlChoice' }, + }, + selectedChoiceIds: { + type: 'array', + items: { type: 'string' }, + }, + }, + required: ['required', 'optional', 'selectedChoiceIds'], + additionalProperties: false, + }, + mcp: { + type: 'object', + properties: { + connections: { + type: 'array', + items: { $ref: '#/components/schemas/AgentThreadRuntimeControlChoice' }, + }, + selectedChoiceIds: { + type: 'array', + items: { type: 'string' }, + }, + }, + required: ['connections', 'selectedChoiceIds'], + additionalProperties: false, + }, + canEdit: { type: 'boolean' }, + disabledReason: { type: 'string', nullable: true }, }, - required: ['id', 'adapter', 'status', 'input', 'sandboxRequirements', 'error'], + required: ['tools', 'mcp', 'canEdit', 'disabledReason'], + additionalProperties: false, }, - AgentSandboxExposure: { + AgentThreadRuntimeControlsPatchRequest: { 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 }, + toolChoiceIds: { + type: 'array', + items: { type: 'string' }, + }, + mcpChoiceIds: { + type: 'array', + items: { type: 'string' }, + }, }, - required: ['id', 'kind', 'status', 'metadata'], + additionalProperties: false, }, - AgentSandbox: { + AgentRuntimeControlChoicesInput: { 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'], + agentId: { type: 'string', nullable: true }, + toolChoiceIds: { + type: 'array', + items: { type: 'string' }, }, - capabilitySnapshot: { type: 'object', additionalProperties: true }, - exposures: { + mcpChoiceIds: { type: 'array', - items: { $ref: '#/components/schemas/AgentSandboxExposure' }, + items: { type: 'string' }, }, - 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'], + additionalProperties: false, }, - AgentSessionSummary: { + AgentRuntimeControlsPreviewRequest: { type: 'object', properties: { - session: { + agentId: { type: 'string', nullable: true }, + source: { 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 }, + adapter: { type: 'string' }, + input: { + type: 'object', + additionalProperties: true, + }, }, - required: ['id', 'status', 'userId', 'ownerGithubUsername', 'defaults', 'defaultThreadId'], + additionalProperties: true, }, - source: { $ref: '#/components/schemas/AgentSource' }, - sandbox: { $ref: '#/components/schemas/AgentSandbox' }, + defaults: { + type: 'object', + properties: { + provider: { type: 'string', nullable: true }, + model: { type: 'string', nullable: true }, + }, + additionalProperties: false, + }, + runtimeControlChoices: { $ref: '#/components/schemas/AgentRuntimeControlChoicesInput' }, }, - required: ['session', 'source', 'sandbox'], + additionalProperties: false, }, - AgentAdminSessionSummary: { + GetAgentThreadRuntimeControlsSuccessResponse: { + allOf: [ + { $ref: '#/components/schemas/SuccessApiResponse' }, + { + type: 'object', + properties: { + data: { $ref: '#/components/schemas/AgentThreadRuntimeControlsState' }, + }, + required: ['data'], + }, + ], + }, + + PatchAgentThreadRuntimeControlsSuccessResponse: { + allOf: [ + { $ref: '#/components/schemas/SuccessApiResponse' }, + { + type: 'object', + properties: { + data: { $ref: '#/components/schemas/AgentThreadRuntimeControlsState' }, + }, + required: ['data'], + }, + ], + }, + + AgentRuntimeControlsPreviewSuccessResponse: { + allOf: [ + { $ref: '#/components/schemas/SuccessApiResponse' }, + { + type: 'object', + properties: { + data: { $ref: '#/components/schemas/AgentThreadRuntimeControlsState' }, + }, + required: ['data'], + }, + ], + }, + + AgentSwitchEventMetadata: { type: 'object', properties: { - id: { type: 'string' }, - sessionKind: { - type: 'string', - enum: Object.values(AgentSessionKind), + kind: { type: 'string', enum: ['agent_switch'] }, + actor: { + type: 'object', + properties: { + userId: { type: 'string' }, + label: { type: 'string' }, + }, + required: ['userId', 'label'], + additionalProperties: false, }, - buildUuid: { type: 'string', nullable: true }, - baseBuildUuid: { type: 'string', nullable: true }, - buildKind: { - type: 'string', - enum: Object.values(BuildKind), - nullable: true, + beforeAgent: { + type: 'object', + properties: { + id: { type: 'string' }, + label: { type: 'string' }, + }, + required: ['id', 'label'], + additionalProperties: false, }, - userId: { type: 'string' }, + afterAgent: { + type: 'object', + properties: { + id: { type: 'string' }, + label: { type: 'string' }, + }, + required: ['id', 'label'], + additionalProperties: false, + }, + appliesTo: { type: 'string', enum: ['future_runs'] }, + occurredAt: { type: 'string', format: 'date-time' }, + }, + required: ['kind', 'actor', 'beforeAgent', 'afterAgent', 'appliesTo', 'occurredAt'], + additionalProperties: false, + }, + + UserAgentDefinitionResourceBehavior: { + type: 'string', + enum: ['chat_only', 'current_workspace_when_available'], + }, + + UserAgentDefinitionModelPreference: { + type: 'object', + nullable: true, + properties: { + provider: { type: 'string', nullable: true }, + model: { type: 'string', nullable: true }, + }, + additionalProperties: false, + }, + + UserAgentDefinition: { + type: 'object', + properties: { + id: { type: 'string' }, + version: { type: 'integer', minimum: 1 }, + name: { type: 'string' }, + description: { type: 'string', nullable: true }, + instructions: { type: 'string' }, + capabilityIds: { + type: 'array', + items: { type: 'string' }, + }, + modelPreference: { $ref: '#/components/schemas/UserAgentDefinitionModelPreference' }, + resourceBehavior: { $ref: '#/components/schemas/UserAgentDefinitionResourceBehavior' }, + status: { type: 'string', enum: ['active', 'archived'] }, + }, + required: [ + 'id', + 'version', + 'name', + 'description', + 'instructions', + 'capabilityIds', + 'modelPreference', + 'resourceBehavior', + 'status', + ], + additionalProperties: false, + }, + + UserAgentDefinitionSummary: { + type: 'object', + properties: { + id: { type: 'string' }, + version: { type: 'integer', minimum: 1 }, + name: { type: 'string' }, + description: { type: 'string', nullable: true }, + capabilityIds: { + type: 'array', + items: { type: 'string' }, + }, + modelPreference: { $ref: '#/components/schemas/UserAgentDefinitionModelPreference' }, + resourceBehavior: { $ref: '#/components/schemas/UserAgentDefinitionResourceBehavior' }, + status: { type: 'string', enum: ['active', 'archived'] }, + }, + required: [ + 'id', + 'version', + 'name', + 'description', + 'capabilityIds', + 'modelPreference', + 'resourceBehavior', + 'status', + ], + additionalProperties: false, + }, + + UserAgentDefinitionUpsertRequest: { + type: 'object', + properties: { + name: { type: 'string', minLength: 1 }, + description: { type: 'string', nullable: true }, + instructions: { type: 'string', minLength: 1 }, + capabilityIds: { + type: 'array', + items: { type: 'string' }, + }, + modelPreference: { $ref: '#/components/schemas/UserAgentDefinitionModelPreference' }, + resourceBehavior: { $ref: '#/components/schemas/UserAgentDefinitionResourceBehavior' }, + }, + required: ['name', 'instructions', 'resourceBehavior'], + additionalProperties: false, + }, + + UserAgentDefinitionDisplaySummary: { + type: 'object', + properties: { + name: { type: 'string' }, + description: { type: 'string', nullable: true }, + }, + required: ['name', 'description'], + additionalProperties: false, + }, + + UserAgentDefinitionCapability: { + type: 'object', + properties: { + capabilityId: { type: 'string' }, + label: { type: 'string' }, + description: { type: 'string' }, + category: { + type: 'string', + enum: [ + 'read', + 'diagnostics', + 'workspace', + 'source_control', + 'mcp', + 'deployment', + 'network', + 'preview', + 'approval', + ], + }, + toolCount: { type: 'integer', minimum: 0 }, + resourceCount: { type: 'integer', minimum: 0 }, + requiresWorkspace: { type: 'boolean' }, + tools: { + type: 'array', + items: { $ref: '#/components/schemas/UserAgentDefinitionDisplaySummary' }, + }, + resources: { + type: 'array', + items: { $ref: '#/components/schemas/UserAgentDefinitionDisplaySummary' }, + }, + }, + required: [ + 'capabilityId', + 'label', + 'description', + 'category', + 'toolCount', + 'resourceCount', + 'requiresWorkspace', + 'tools', + 'resources', + ], + additionalProperties: false, + }, + + ListUserAgentDefinitionsResponse: { + type: 'object', + properties: { + definitions: { + type: 'array', + items: { $ref: '#/components/schemas/UserAgentDefinition' }, + }, + }, + required: ['definitions'], + additionalProperties: false, + }, + + UserAgentDefinitionResponse: { + type: 'object', + properties: { + definition: { $ref: '#/components/schemas/UserAgentDefinition' }, + }, + required: ['definition'], + additionalProperties: false, + }, + + DeleteUserAgentDefinitionResponse: { + type: 'object', + properties: { + archived: { type: 'boolean' }, + definition: { $ref: '#/components/schemas/UserAgentDefinition' }, + }, + required: ['archived', 'definition'], + additionalProperties: false, + }, + + UserAgentDefinitionCapabilitiesResponse: { + type: 'object', + properties: { + resourceBehavior: { $ref: '#/components/schemas/UserAgentDefinitionResourceBehavior' }, + canCreate: { type: 'boolean' }, + creationUnavailableReason: { + $ref: '#/components/schemas/CustomAgentCreationUnavailableReason', + }, + capabilities: { + type: 'array', + items: { $ref: '#/components/schemas/UserAgentDefinitionCapability' }, + }, + }, + required: ['resourceBehavior', 'canCreate', 'creationUnavailableReason', 'capabilities'], + additionalProperties: false, + }, + + ListUserAgentDefinitionsSuccessResponse: { + allOf: [ + { $ref: '#/components/schemas/SuccessApiResponse' }, + { + type: 'object', + required: ['data'], + properties: { + data: { $ref: '#/components/schemas/ListUserAgentDefinitionsResponse' }, + }, + }, + ], + }, + + UserAgentDefinitionSuccessResponse: { + allOf: [ + { $ref: '#/components/schemas/SuccessApiResponse' }, + { + type: 'object', + required: ['data'], + properties: { + data: { $ref: '#/components/schemas/UserAgentDefinitionResponse' }, + }, + }, + ], + }, + + DeleteUserAgentDefinitionSuccessResponse: { + allOf: [ + { $ref: '#/components/schemas/SuccessApiResponse' }, + { + type: 'object', + required: ['data'], + properties: { + data: { $ref: '#/components/schemas/DeleteUserAgentDefinitionResponse' }, + }, + }, + ], + }, + + UserAgentDefinitionCapabilitiesSuccessResponse: { + allOf: [ + { $ref: '#/components/schemas/SuccessApiResponse' }, + { + type: 'object', + required: ['data'], + properties: { + data: { $ref: '#/components/schemas/UserAgentDefinitionCapabilitiesResponse' }, + }, + }, + ], + }, + + AgentThreadMessagesResponse: { + type: 'object', + properties: { + 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: ['thread', 'messages', 'pagination'], + additionalProperties: false, + }, + + AgentUsageSummary: { + type: 'object', + properties: agentUsageSummaryProperties, + required: ['totalTokens'], + additionalProperties: false, + }, + + AgentUsageByModel: { + type: 'object', + properties: { + provider: { type: 'string' }, + model: { type: 'string' }, + ...agentUsageSummaryProperties, + runCount: { type: 'integer' }, + reportedRunCount: { type: 'integer' }, + missingUsageRunCount: { type: 'integer' }, + }, + required: ['provider', 'model', 'totalTokens', 'runCount', 'reportedRunCount', 'missingUsageRunCount'], + additionalProperties: false, + }, + + AgentUsageCompleteness: { + type: 'object', + properties: { + runCount: { type: 'integer' }, + reportedRunCount: { type: 'integer' }, + missingUsageRunCount: { type: 'integer' }, + complete: { type: 'boolean' }, + }, + required: ['runCount', 'reportedRunCount', 'missingUsageRunCount', 'complete'], + additionalProperties: false, + }, + + AgentUsageAggregate: { + type: 'object', + properties: { + usageSummary: { $ref: '#/components/schemas/AgentUsageSummary' }, + usageByModel: { + type: 'array', + items: { $ref: '#/components/schemas/AgentUsageByModel' }, + }, + usageCompleteness: { $ref: '#/components/schemas/AgentUsageCompleteness' }, + }, + required: ['usageSummary', 'usageByModel', 'usageCompleteness'], + additionalProperties: false, + }, + + AgentThreadUsageResponse: { + type: 'object', + properties: { + threadId: { type: 'string' }, + sessionId: { type: 'string' }, + usageSummary: { $ref: '#/components/schemas/AgentUsageSummary' }, + usageByModel: { + type: 'array', + items: { $ref: '#/components/schemas/AgentUsageByModel' }, + }, + usageCompleteness: { $ref: '#/components/schemas/AgentUsageCompleteness' }, + }, + required: ['threadId', 'sessionId', 'usageSummary', 'usageByModel', 'usageCompleteness'], + additionalProperties: false, + }, + + GetAgentThreadUsageSuccessResponse: { + allOf: [ + { $ref: '#/components/schemas/SuccessApiResponse' }, + { + type: 'object', + required: ['data'], + properties: { + data: { $ref: '#/components/schemas/AgentThreadUsageResponse' }, + }, + }, + ], + }, + + AgentRunRuntimeOptions: { + type: 'object', + properties: { + maxIterations: { + type: 'integer', + minimum: 1, + maximum: 100, + }, + }, + additionalProperties: false, + example: { + maxIterations: 12, + }, + }, + + CreateAgentThreadRunMessage: { + type: 'object', + properties: { + clientMessageId: { type: 'string' }, + parts: { + type: 'array', + 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, + }, + }, + }, + + CreateBuildContextAgentChatRequest: { + type: 'object', + properties: { + buildUuid: { type: 'string' }, + defaults: { + type: 'object', + properties: { + model: { type: 'string' }, + }, + additionalProperties: false, + }, + }, + required: ['buildUuid'], + additionalProperties: false, + example: { + buildUuid: '00000000-0000-0000-0000-000000000000', + defaults: { + model: 'gpt-5.4', + }, + }, + }, + + AgentThread: { + type: 'object', + properties: { + id: { type: 'string' }, + sessionId: { type: 'string', nullable: true }, + title: { type: 'string', nullable: true }, + isDefault: { type: 'boolean' }, + archivedAt: { type: 'string', format: 'date-time', nullable: true }, + lastRunAt: { type: 'string', format: 'date-time', nullable: true }, + metadata: { type: 'object', additionalProperties: true }, + createdAt: { type: 'string', format: 'date-time', nullable: true }, + updatedAt: { type: 'string', format: 'date-time', nullable: true }, + }, + required: ['id', 'isDefault', 'metadata'], + }, + + AgentSessionDefaults: { + type: 'object', + properties: { + provider: { type: 'string', nullable: true }, + 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' }, + usage: { $ref: '#/components/schemas/AgentUsageAggregate' }, + }, + required: ['session', 'source', 'sandbox', 'usage'], + }, + + BuildContextAgentChatContext: { + type: 'object', + properties: { + buildUuid: { type: 'string' }, + buildKind: { + type: 'string', + enum: Object.values(BuildKind), + nullable: true, + }, + namespace: { type: 'string', nullable: true }, + baseBuildUuid: { type: 'string', nullable: true }, + repo: { type: 'string', nullable: true }, + branch: { type: 'string', nullable: true }, + pullRequestNumber: { type: 'integer', nullable: true }, + contextFreshAt: { type: 'string', format: 'date-time' }, + }, + required: [ + 'buildUuid', + 'buildKind', + 'namespace', + 'baseBuildUuid', + 'repo', + 'branch', + 'pullRequestNumber', + 'contextFreshAt', + ], + additionalProperties: false, + }, + + BuildContextAgentChatLinks: { + type: 'object', + properties: { + messages: { type: 'string' }, + runs: { type: 'string' }, + events: { type: 'string' }, + eventStream: { type: 'string' }, + pendingActions: { type: 'string' }, + }, + required: ['messages', 'runs', 'events', 'eventStream', 'pendingActions'], + additionalProperties: false, + }, + + BuildContextAgentChatResponse: { + type: 'object', + properties: { + session: { $ref: '#/components/schemas/AgentSessionSummary' }, + thread: { $ref: '#/components/schemas/AgentThread' }, + created: { type: 'boolean' }, + reused: { type: 'boolean' }, + buildContext: { $ref: '#/components/schemas/BuildContextAgentChatContext' }, + links: { $ref: '#/components/schemas/BuildContextAgentChatLinks' }, + }, + required: ['session', 'thread', 'created', 'reused', 'buildContext', 'links'], + additionalProperties: false, + }, + + 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: { + type: 'string', + enum: Object.values(BuildKind), + nullable: true, + }, + userId: { type: 'string' }, ownerGithubUsername: { type: 'string', nullable: true }, podName: { type: 'string', nullable: true }, namespace: { type: 'string', nullable: true }, @@ -868,6 +1542,135 @@ export const openApiSpecificationForV2Api: OAS3Options = { ], }, + AgentRunPlanRuntimeSummary: { + type: 'object', + properties: { + harness: { type: 'string', enum: ['lifecycle_ai_sdk'] }, + maxIterations: { type: 'integer', nullable: true }, + }, + required: ['harness', 'maxIterations'], + additionalProperties: false, + }, + + AgentRunPlanApprovalSummary: { + type: 'object', + properties: { + defaultMode: { $ref: '#/components/schemas/AgentApprovalMode' }, + }, + required: ['defaultMode'], + additionalProperties: false, + }, + + AgentRunPlanCapabilitySummary: { + type: 'object', + properties: { + capabilityId: { type: 'string' }, + availability: { $ref: '#/components/schemas/AgentCapabilityAvailability' }, + allowed: { type: 'boolean' }, + approvalMode: { + allOf: [{ $ref: '#/components/schemas/AgentApprovalMode' }], + nullable: true, + }, + }, + required: ['capabilityId', 'availability', 'allowed'], + additionalProperties: false, + }, + + AgentRunPlanSelectedRuntimeChoicesSummary: { + type: 'object', + properties: { + capabilityIds: { + type: 'array', + items: { type: 'string' }, + }, + toolChoiceIds: { + type: 'array', + items: { type: 'string' }, + }, + mcpChoiceIds: { + type: 'array', + items: { type: 'string' }, + }, + }, + required: ['capabilityIds', 'toolChoiceIds', 'mcpChoiceIds'], + additionalProperties: false, + }, + + AgentRunPlanCapabilitiesSummary: { + type: 'object', + properties: { + effective: { + type: 'array', + items: { $ref: '#/components/schemas/AgentRunPlanCapabilitySummary' }, + }, + selected: { $ref: '#/components/schemas/AgentRunPlanSelectedRuntimeChoicesSummary' }, + }, + required: ['effective', 'selected'], + additionalProperties: false, + }, + + AgentRunPlanSummary: { + type: 'object', + nullable: true, + properties: { + version: { type: 'integer', enum: [1] }, + agent: { + type: 'object', + properties: { + id: { type: 'string' }, + label: { type: 'string' }, + sourceKind: { + type: 'string', + enum: ['build_context_chat', 'workspace_session', 'freeform_chat'], + }, + }, + required: ['id', 'label', 'sourceKind'], + additionalProperties: false, + }, + source: { + type: 'object', + properties: { + kind: { + type: 'string', + enum: ['build_context_chat', 'workspace_session', 'freeform_chat'], + }, + repoFullName: { type: 'string', nullable: true }, + branch: { type: 'string', nullable: true }, + buildUuid: { type: 'string', nullable: true }, + namespace: { type: 'string', nullable: true }, + }, + required: ['kind'], + additionalProperties: false, + }, + model: { + type: 'object', + properties: { + provider: { type: 'string' }, + model: { type: 'string' }, + }, + required: ['provider', 'model'], + additionalProperties: false, + }, + runtime: { $ref: '#/components/schemas/AgentRunPlanRuntimeSummary' }, + approval: { $ref: '#/components/schemas/AgentRunPlanApprovalSummary' }, + capabilities: { $ref: '#/components/schemas/AgentRunPlanCapabilitiesSummary' }, + warnings: { + type: 'array', + items: { + type: 'object', + properties: { + code: { type: 'string' }, + message: { type: 'string' }, + }, + required: ['code', 'message'], + additionalProperties: false, + }, + }, + }, + required: ['version', 'agent', 'source', 'model', 'runtime', 'approval', 'capabilities', 'warnings'], + additionalProperties: false, + }, + AgentRun: { type: 'object', properties: { @@ -891,6 +1694,7 @@ export const openApiSpecificationForV2Api: OAS3Options = { cancelledAt: { type: 'string', format: 'date-time', nullable: true }, usageSummary: { type: 'object', additionalProperties: true }, policySnapshot: { type: 'object', additionalProperties: true }, + runPlan: { $ref: '#/components/schemas/AgentRunPlanSummary' }, error: { allOf: [{ $ref: '#/components/schemas/AgentRunError' }], nullable: true, @@ -913,6 +1717,7 @@ export const openApiSpecificationForV2Api: OAS3Options = { 'sandboxGeneration', 'usageSummary', 'policySnapshot', + 'runPlan', ], }, @@ -1109,7 +1914,7 @@ export const openApiSpecificationForV2Api: OAS3Options = { type: 'object', properties: { id: { type: 'string' }, - data: { type: 'object', additionalProperties: true }, + data: { $ref: '#/components/schemas/AgentFileChangeData' }, transient: { type: 'boolean' }, }, required: ['data'], @@ -1282,17 +2087,50 @@ export const openApiSpecificationForV2Api: OAS3Options = { additionalProperties: false, }, - AgentPendingActionFileChangePreview: { + AgentFileChangeData: { type: 'object', properties: { + id: { type: 'string' }, + toolCallId: { type: 'string' }, + sourceTool: { type: 'string' }, path: { type: 'string' }, - action: { type: 'string' }, - summary: { type: 'string' }, - additions: { type: 'integer', nullable: true }, - deletions: { type: 'integer', nullable: true }, + displayPath: { type: 'string' }, + kind: { type: 'string', enum: ['created', 'edited', 'deleted'] }, + stage: { type: 'string', enum: ['awaiting-approval', 'approved', 'applied', 'denied', 'failed'] }, + additions: { type: 'integer' }, + deletions: { type: 'integer' }, truncated: { type: 'boolean' }, + unifiedDiff: { type: 'string', nullable: true }, + beforeTextPreview: { type: 'string', nullable: true }, + afterTextPreview: { type: 'string', nullable: true }, + summary: { type: 'string', nullable: true }, + encoding: { type: 'string', nullable: true }, + oldSizeBytes: { type: 'integer', nullable: true }, + newSizeBytes: { type: 'integer', nullable: true }, + oldSha256: { type: 'string', nullable: true }, + newSha256: { type: 'string', nullable: true }, }, - required: ['path', 'action', 'summary', 'additions', 'deletions', 'truncated'], + required: [ + 'id', + 'toolCallId', + 'sourceTool', + 'path', + 'displayPath', + 'kind', + 'stage', + 'additions', + 'deletions', + 'truncated', + 'unifiedDiff', + 'beforeTextPreview', + 'afterTextPreview', + 'summary', + 'encoding', + 'oldSizeBytes', + 'newSizeBytes', + 'oldSha256', + 'newSha256', + ], additionalProperties: false, }, @@ -1316,7 +2154,7 @@ export const openApiSpecificationForV2Api: OAS3Options = { commandPreview: { type: 'string', nullable: true }, fileChangePreview: { type: 'array', - items: { $ref: '#/components/schemas/AgentPendingActionFileChangePreview' }, + items: { $ref: '#/components/schemas/AgentFileChangeData' }, }, riskLabels: { type: 'array', @@ -1355,12 +2193,25 @@ export const openApiSpecificationForV2Api: OAS3Options = { commandPreview: null, fileChangePreview: [ { + id: 'tool-call-1:sample-file.txt', + toolCallId: 'tool-call-1', + sourceTool: 'workspace_edit_file', path: 'sample-file.txt', - action: 'edited', + displayPath: 'sample-file.txt', + kind: 'edited', + stage: 'awaiting-approval', summary: 'edited sample-file.txt', additions: 1, deletions: 0, truncated: false, + unifiedDiff: null, + beforeTextPreview: null, + afterTextPreview: null, + encoding: null, + oldSizeBytes: null, + newSizeBytes: null, + oldSha256: null, + newSha256: null, }, ], riskLabels: ['Workspace write'], @@ -2327,10 +3178,10 @@ export const openApiSpecificationForV2Api: OAS3Options = { }, // =================================================================== - // AI Agent Config Schemas + // Agent Runtime Config Schemas // =================================================================== - AIAgentModelConfig: { + AgentRuntimeModelConfig: { type: 'object', properties: { id: { @@ -2354,7 +3205,7 @@ export const openApiSpecificationForV2Api: OAS3Options = { required: ['id', 'displayName', 'enabled', 'default', 'maxTokens'], }, - AIAgentProviderConfig: { + AgentRuntimeProviderConfig: { type: 'object', properties: { name: { type: 'string' }, @@ -2366,20 +3217,106 @@ export const openApiSpecificationForV2Api: OAS3Options = { }, models: { type: 'array', - items: { $ref: '#/components/schemas/AIAgentModelConfig' }, + items: { $ref: '#/components/schemas/AgentRuntimeModelConfig' }, }, }, required: ['name', 'enabled', 'apiKeyEnvVar', 'models'], }, - AIAgentConfig: { + AgentCapabilityAvailability: { + type: 'string', + enum: ['all_users', 'admin_only', 'system_only', 'disabled'], + }, + + AgentCapabilityPolicyAvailability: { + type: 'object', + properties: { + read_context: { $ref: '#/components/schemas/AgentCapabilityAvailability' }, + diagnostics_logs: { $ref: '#/components/schemas/AgentCapabilityAvailability' }, + diagnostics_codefresh: { $ref: '#/components/schemas/AgentCapabilityAvailability' }, + diagnostics_kubernetes: { $ref: '#/components/schemas/AgentCapabilityAvailability' }, + diagnostics_database: { $ref: '#/components/schemas/AgentCapabilityAvailability' }, + github_read: { $ref: '#/components/schemas/AgentCapabilityAvailability' }, + github_write: { $ref: '#/components/schemas/AgentCapabilityAvailability' }, + workspace_files: { $ref: '#/components/schemas/AgentCapabilityAvailability' }, + workspace_shell: { $ref: '#/components/schemas/AgentCapabilityAvailability' }, + workspace_git: { $ref: '#/components/schemas/AgentCapabilityAvailability' }, + network_access: { $ref: '#/components/schemas/AgentCapabilityAvailability' }, + preview_publish: { $ref: '#/components/schemas/AgentCapabilityAvailability' }, + external_mcp_read: { $ref: '#/components/schemas/AgentCapabilityAvailability' }, + external_mcp_write: { $ref: '#/components/schemas/AgentCapabilityAvailability' }, + approval_controls: { $ref: '#/components/schemas/AgentCapabilityAvailability' }, + }, + additionalProperties: false, + }, + + AgentCapabilityPolicy: { + type: 'object', + properties: { + availability: { $ref: '#/components/schemas/AgentCapabilityPolicyAvailability' }, + }, + additionalProperties: false, + }, + + CustomAgentCreationMode: { + type: 'string', + enum: ['enabled', 'disabled', 'admins_only', 'allowlist'], + }, + + CustomAgentCreationUnavailableReason: { + type: 'string', + nullable: true, + enum: ['creation_disabled', 'creation_restricted', null], + }, + + CreatorCapabilityAvailability: { + type: 'string', + enum: ['available', 'reserved'], + }, + + CreatorCapabilityAvailabilityMap: { + type: 'object', + properties: { + read_context: { $ref: '#/components/schemas/CreatorCapabilityAvailability' }, + diagnostics_logs: { $ref: '#/components/schemas/CreatorCapabilityAvailability' }, + diagnostics_codefresh: { $ref: '#/components/schemas/CreatorCapabilityAvailability' }, + diagnostics_kubernetes: { $ref: '#/components/schemas/CreatorCapabilityAvailability' }, + diagnostics_database: { $ref: '#/components/schemas/CreatorCapabilityAvailability' }, + github_read: { $ref: '#/components/schemas/CreatorCapabilityAvailability' }, + github_write: { $ref: '#/components/schemas/CreatorCapabilityAvailability' }, + workspace_files: { $ref: '#/components/schemas/CreatorCapabilityAvailability' }, + workspace_shell: { $ref: '#/components/schemas/CreatorCapabilityAvailability' }, + workspace_git: { $ref: '#/components/schemas/CreatorCapabilityAvailability' }, + network_access: { $ref: '#/components/schemas/CreatorCapabilityAvailability' }, + preview_publish: { $ref: '#/components/schemas/CreatorCapabilityAvailability' }, + external_mcp_read: { $ref: '#/components/schemas/CreatorCapabilityAvailability' }, + external_mcp_write: { $ref: '#/components/schemas/CreatorCapabilityAvailability' }, + approval_controls: { $ref: '#/components/schemas/CreatorCapabilityAvailability' }, + }, + additionalProperties: false, + }, + + CustomAgentCreationPolicy: { + type: 'object', + properties: { + mode: { $ref: '#/components/schemas/CustomAgentCreationMode' }, + allowedUserIds: { type: 'array', items: { type: 'string' } }, + allowedGithubUsernames: { type: 'array', items: { type: 'string' } }, + capabilityAvailability: { $ref: '#/components/schemas/CreatorCapabilityAvailabilityMap' }, + }, + additionalProperties: false, + }, + + AgentRuntimeConfig: { type: 'object', properties: { enabled: { type: 'boolean' }, approvalPolicy: { $ref: '#/components/schemas/AgentApprovalPolicy' }, + capabilityPolicy: { $ref: '#/components/schemas/AgentCapabilityPolicy' }, + customAgentCreationPolicy: { $ref: '#/components/schemas/CustomAgentCreationPolicy' }, providers: { type: 'array', - items: { $ref: '#/components/schemas/AIAgentProviderConfig' }, + items: { $ref: '#/components/schemas/AgentRuntimeProviderConfig' }, }, maxMessagesPerSession: { type: 'integer' }, sessionTTL: { type: 'integer' }, @@ -2413,13 +3350,14 @@ export const openApiSpecificationForV2Api: OAS3Options = { required: ['enabled', 'providers', 'maxMessagesPerSession', 'sessionTTL'], }, - AIAgentRepoOverride: { + AgentRuntimeRepoOverride: { type: 'object', properties: { enabled: { type: 'boolean' }, maxMessagesPerSession: { type: 'integer' }, sessionTTL: { type: 'integer' }, approvalPolicy: { $ref: '#/components/schemas/AgentApprovalPolicy' }, + capabilityPolicy: { $ref: '#/components/schemas/AgentCapabilityPolicy' }, additiveRules: { type: 'array', items: { type: 'string' } }, systemPromptOverride: { type: 'string', maxLength: 50000 }, excludedTools: { type: 'array', items: { type: 'string' } }, @@ -2428,7 +3366,7 @@ export const openApiSpecificationForV2Api: OAS3Options = { }, }, - AIAgentAdditiveRulesUpdateRequest: { + AgentRuntimeAdditiveRulesUpdateRequest: { type: 'object', properties: { additiveRules: { type: 'array', items: { type: 'string' } }, @@ -2437,7 +3375,7 @@ export const openApiSpecificationForV2Api: OAS3Options = { additionalProperties: false, }, - AIAgentApprovalPolicyUpdateRequest: { + AgentRuntimeApprovalPolicyUpdateRequest: { type: 'object', properties: { approvalPolicy: { $ref: '#/components/schemas/AgentApprovalPolicy' }, @@ -2446,7 +3384,7 @@ export const openApiSpecificationForV2Api: OAS3Options = { additionalProperties: false, }, - AIAgentConfigPatchRequest: { + AgentRuntimeConfigPatchRequest: { type: 'object', properties: { additiveRules: { type: 'array', items: { type: 'string' } }, @@ -2458,29 +3396,29 @@ export const openApiSpecificationForV2Api: OAS3Options = { description: 'Provide exactly one patch target: additiveRules or approvalPolicy.', }, - AIAgentRepoConfigEntry: { + AgentRuntimeRepoConfigEntry: { type: 'object', properties: { repositoryFullName: { type: 'string' }, - config: { $ref: '#/components/schemas/AIAgentRepoOverride' }, + config: { $ref: '#/components/schemas/AgentRuntimeRepoOverride' }, createdAt: { type: 'string', format: 'date-time' }, updatedAt: { type: 'string', format: 'date-time' }, }, }, - GetGlobalAIAgentConfigSuccessResponse: { + GetGlobalAgentRuntimeConfigSuccessResponse: { allOf: [ { $ref: '#/components/schemas/SuccessApiResponse' }, { type: 'object', properties: { - data: { $ref: '#/components/schemas/AIAgentConfig' }, + data: { $ref: '#/components/schemas/AgentRuntimeConfig' }, }, }, ], }, - GetRepoAIAgentConfigSuccessResponse: { + GetRepoAgentRuntimeConfigSuccessResponse: { allOf: [ { $ref: '#/components/schemas/SuccessApiResponse' }, { @@ -2490,7 +3428,7 @@ export const openApiSpecificationForV2Api: OAS3Options = { type: 'object', properties: { repoFullName: { type: 'string' }, - config: { $ref: '#/components/schemas/AIAgentRepoOverride' }, + config: { $ref: '#/components/schemas/AgentRuntimeRepoOverride' }, }, }, }, @@ -2498,7 +3436,7 @@ export const openApiSpecificationForV2Api: OAS3Options = { ], }, - GetEffectiveAIAgentConfigSuccessResponse: { + GetEffectiveAgentRuntimeConfigSuccessResponse: { allOf: [ { $ref: '#/components/schemas/SuccessApiResponse' }, { @@ -2508,7 +3446,7 @@ export const openApiSpecificationForV2Api: OAS3Options = { type: 'object', properties: { repoFullName: { type: 'string' }, - effectiveConfig: { $ref: '#/components/schemas/AIAgentConfig' }, + effectiveConfig: { $ref: '#/components/schemas/AgentRuntimeConfig' }, }, }, }, @@ -2516,7 +3454,7 @@ export const openApiSpecificationForV2Api: OAS3Options = { ], }, - ListRepoAIAgentConfigsSuccessResponse: { + ListRepoAgentRuntimeConfigsSuccessResponse: { allOf: [ { $ref: '#/components/schemas/SuccessApiResponse' }, { @@ -2524,7 +3462,7 @@ export const openApiSpecificationForV2Api: OAS3Options = { properties: { data: { type: 'array', - items: { $ref: '#/components/schemas/AIAgentRepoConfigEntry' }, + items: { $ref: '#/components/schemas/AgentRuntimeRepoConfigEntry' }, }, }, }, @@ -2731,252 +3669,162 @@ export const openApiSpecificationForV2Api: OAS3Options = { ], }, - GetGlobalAgentSessionConfigSuccessResponse: { - allOf: [ - { $ref: '#/components/schemas/SuccessApiResponse' }, - { - type: 'object', - properties: { - data: { $ref: '#/components/schemas/AgentSessionControlPlaneConfig' }, - }, - }, - ], - }, - - GetEffectiveGlobalAgentSessionConfigSuccessResponse: { - allOf: [ - { $ref: '#/components/schemas/SuccessApiResponse' }, - { - type: 'object', - properties: { - data: { - $ref: '#/components/schemas/EffectiveAgentSessionControlPlaneConfig', - }, - }, - }, - ], - }, - - GetGlobalAgentSessionRuntimeConfigSuccessResponse: { - allOf: [ - { $ref: '#/components/schemas/SuccessApiResponse' }, - { - type: 'object', - properties: { - data: { $ref: '#/components/schemas/AgentSessionRuntimeSettings' }, - }, - }, - ], - }, - - GetRepoAgentSessionConfigSuccessResponse: { - allOf: [ - { $ref: '#/components/schemas/SuccessApiResponse' }, - { - type: 'object', - properties: { - data: { - type: 'object', - properties: { - repoFullName: { type: 'string' }, - config: { $ref: '#/components/schemas/AgentSessionControlPlaneConfig' }, - }, - required: ['repoFullName', 'config'], - }, - }, - }, - ], - }, - - GetAdminAgentToolInventorySuccessResponse: { - allOf: [ - { $ref: '#/components/schemas/SuccessApiResponse' }, - { - type: 'object', - properties: { - data: { - type: 'array', - items: { $ref: '#/components/schemas/AgentSessionToolInventoryEntry' }, - }, - }, - }, - ], - }, - - FeedbackEntry: { + AgentCapabilityCatalogEntry: { type: 'object', properties: { id: { type: 'string' }, - feedbackType: { + label: { type: 'string' }, + description: { type: 'string' }, + category: { type: 'string', - enum: ['message', 'conversation'], + enum: [ + 'read', + 'diagnostics', + 'workspace', + 'source_control', + 'mcp', + 'deployment', + 'network', + 'preview', + 'approval', + ], }, - buildUuid: { type: 'string' }, - rating: { - type: 'string', - enum: ['up', 'down'], - }, - text: { type: 'string', nullable: true }, - userIdentifier: { type: 'string', nullable: true }, - repo: { type: 'string' }, - prNumber: { type: 'integer', nullable: true }, - messageId: { type: 'integer', nullable: true }, - messagePreview: { type: 'string', nullable: true }, - costUsd: { type: 'number', nullable: true }, - createdAt: { type: 'string', format: 'date-time' }, + defaultAvailability: { $ref: '#/components/schemas/AgentCapabilityAvailability' }, + defaultApprovalMode: { $ref: '#/components/schemas/AgentApprovalMode' }, + runtimeCapabilityKey: { type: 'string' }, + toolKeys: { + type: 'array', + items: { type: 'string' }, + }, + resourceGrants: { + type: 'array', + items: { type: 'string' }, + }, + sourceKinds: { + type: 'array', + items: { + type: 'string', + enum: ['build_context_chat', 'workspace_session', 'freeform_chat'], + }, + }, + userSelectable: { type: 'boolean' }, }, required: [ 'id', - 'feedbackType', - 'buildUuid', - 'rating', - 'text', - 'userIdentifier', - 'repo', - 'prNumber', - 'messageId', - 'messagePreview', - 'costUsd', - 'createdAt', - ], - }, - - FeedbackListPaginationMetadata: { - type: 'object', - properties: { - page: { type: 'integer' }, - limit: { type: 'integer' }, - totalCount: { type: 'integer' }, - totalPages: { type: 'integer' }, - }, - required: ['page', 'limit', 'totalCount', 'totalPages'], - }, - - FeedbackListResponseMetadata: { - type: 'object', - properties: { - pagination: { $ref: '#/components/schemas/FeedbackListPaginationMetadata' }, - }, - required: ['pagination'], - }, - - GetAdminFeedbackListSuccessResponse: { - allOf: [ - { $ref: '#/components/schemas/SuccessApiResponse' }, - { - type: 'object', - properties: { - data: { - type: 'array', - items: { $ref: '#/components/schemas/FeedbackEntry' }, - }, - metadata: { $ref: '#/components/schemas/FeedbackListResponseMetadata' }, - }, - required: ['data', 'metadata'], - }, + 'label', + 'description', + 'category', + 'defaultAvailability', + 'defaultApprovalMode', + 'userSelectable', ], }, - ConversationReplayMessage: { + AgentCapabilityInventoryToolEntry: { type: 'object', properties: { - id: { type: 'integer' }, - role: { + toolKey: { type: 'string' }, + toolName: { type: 'string' }, + description: { type: 'string', nullable: true }, + serverSlug: { type: 'string' }, + serverName: { type: 'string' }, + sourceType: { type: 'string', - enum: ['user', 'assistant', 'system'], - }, - content: { type: 'string' }, - timestamp: { type: 'integer' }, - metadata: { - type: 'object', - additionalProperties: true, + enum: ['builtin', 'mcp'], }, + sourceScope: { type: 'string' }, }, - required: ['id', 'role', 'content', 'timestamp', 'metadata'], + required: ['toolKey', 'toolName', 'serverSlug', 'serverName', 'sourceType', 'sourceScope'], }, - FeedbackConversationReplayData: { + AgentCapabilityInventoryEntry: { type: 'object', properties: { - feedbackType: { + capabilityId: { type: 'string' }, + label: { type: 'string' }, + description: { type: 'string' }, + category: { type: 'string', - enum: ['message', 'conversation'], + enum: [ + 'read', + 'diagnostics', + 'workspace', + 'source_control', + 'mcp', + 'deployment', + 'network', + 'preview', + 'approval', + ], }, - feedbackId: { type: 'integer' }, - buildUuid: { type: 'string' }, - repo: { type: 'string' }, - ratedMessageId: { type: 'integer', nullable: true }, - feedbackRating: { - type: 'string', - enum: ['up', 'down'], + defaultAvailability: { $ref: '#/components/schemas/AgentCapabilityAvailability' }, + configuredAvailability: { $ref: '#/components/schemas/AgentCapabilityAvailability' }, + inheritedAvailability: { $ref: '#/components/schemas/AgentCapabilityAvailability' }, + effectiveAvailability: { $ref: '#/components/schemas/AgentCapabilityAvailability' }, + approvalMode: { $ref: '#/components/schemas/AgentApprovalMode' }, + runtimeCapabilityKey: { type: 'string' }, + userSelectable: { type: 'boolean' }, + toolCount: { type: 'integer', minimum: 0 }, + resourceCount: { type: 'integer', minimum: 0 }, + resourceGrants: { + type: 'array', + items: { type: 'string' }, }, - feedbackText: { type: 'string', nullable: true }, - feedbackUserIdentifier: { type: 'string', nullable: true }, - feedbackCreatedAt: { type: 'string', format: 'date-time' }, - conversation: { - type: 'object', - properties: { - messageCount: { type: 'integer' }, - model: { type: 'string', nullable: true }, - messages: { - type: 'array', - items: { $ref: '#/components/schemas/ConversationReplayMessage' }, - }, - }, - required: ['messageCount', 'model', 'messages'], + tools: { + type: 'array', + items: { $ref: '#/components/schemas/AgentCapabilityInventoryToolEntry' }, + }, + blockedReason: { + type: 'string', + enum: ['admin_only', 'system_only', 'disabled'], }, }, required: [ - 'feedbackType', - 'feedbackId', - 'buildUuid', - 'repo', - 'ratedMessageId', - 'feedbackRating', - 'feedbackText', - 'feedbackUserIdentifier', - 'feedbackCreatedAt', - 'conversation', + 'capabilityId', + 'label', + 'description', + 'category', + 'defaultAvailability', + 'effectiveAvailability', + 'approvalMode', + 'userSelectable', + 'toolCount', + 'resourceCount', + 'resourceGrants', + 'tools', ], }, - GetAdminFeedbackConversationSuccessResponse: { - allOf: [ - { $ref: '#/components/schemas/SuccessApiResponse' }, - { - type: 'object', - properties: { - data: { $ref: '#/components/schemas/FeedbackConversationReplayData' }, - }, - required: ['data'], + AgentCapabilityGovernanceResponse: { + type: 'object', + properties: { + scope: { type: 'string' }, + scopeType: { + type: 'string', + enum: ['global', 'repo'], }, - ], + repoFullName: { type: 'string' }, + capabilityPolicy: { $ref: '#/components/schemas/AgentCapabilityPolicy' }, + inheritedCapabilityPolicy: { $ref: '#/components/schemas/AgentCapabilityPolicy' }, + effectiveCapabilityPolicy: { $ref: '#/components/schemas/AgentCapabilityPolicy' }, + capabilities: { + type: 'array', + items: { $ref: '#/components/schemas/AgentCapabilityInventoryEntry' }, + }, + }, + required: ['scope', 'scopeType', 'capabilityPolicy', 'effectiveCapabilityPolicy', 'capabilities'], }, - // =================================================================== - // AI Chat Schemas - // =================================================================== - - AIModel: { + UpdateAdminAgentCapabilitiesRequest: { type: 'object', - description: 'An available AI model returned by the models endpoint.', properties: { - provider: { type: 'string', description: 'The LLM provider name.', example: 'anthropic' }, - modelId: { - type: 'string', - description: 'The model ID to pass to the chat endpoint.', - example: 'claude-sonnet-4-20250514', - }, - displayName: { type: 'string', example: 'Claude Sonnet' }, - default: { type: 'boolean', description: 'Whether this is the default model.' }, - maxTokens: { type: 'integer', example: 8192 }, - inputCostPerMillion: { type: 'number', description: 'Cost per million input tokens (USD).' }, - outputCostPerMillion: { type: 'number', description: 'Cost per million output tokens (USD).' }, + capabilityPolicy: { $ref: '#/components/schemas/AgentCapabilityPolicy' }, }, - required: ['provider', 'modelId', 'displayName', 'default', 'maxTokens'], + required: ['capabilityPolicy'], + additionalProperties: false, }, - GetAIModelsSuccessResponse: { + GetAdminCustomAgentCreationPolicySuccessResponse: { allOf: [ { $ref: '#/components/schemas/SuccessApiResponse' }, { @@ -2985,146 +3833,63 @@ export const openApiSpecificationForV2Api: OAS3Options = { data: { type: 'object', properties: { - models: { - type: 'array', - items: { $ref: '#/components/schemas/AIModel' }, - }, + customAgentCreationPolicy: { $ref: '#/components/schemas/CustomAgentCreationPolicy' }, }, - required: ['models'], + required: ['customAgentCreationPolicy'], }, }, - required: ['data'], }, ], }, - AIConfigStatus: { + UpdateAdminCustomAgentCreationPolicyRequest: { type: 'object', properties: { - enabled: { type: 'boolean' }, - provider: { type: 'string' }, - configured: { type: 'boolean' }, + customAgentCreationPolicy: { $ref: '#/components/schemas/CustomAgentCreationPolicy' }, }, - required: ['enabled'], + required: ['customAgentCreationPolicy'], + additionalProperties: false, }, - GetAIConfigSuccessResponse: { + GetGlobalAgentSessionConfigSuccessResponse: { allOf: [ { $ref: '#/components/schemas/SuccessApiResponse' }, { type: 'object', properties: { - data: { $ref: '#/components/schemas/AIConfigStatus' }, + data: { $ref: '#/components/schemas/AgentSessionControlPlaneConfig' }, }, - required: ['data'], }, ], }, - ConversationMessage: { - type: 'object', - description: 'A single message in the conversation history.', - properties: { - role: { type: 'string', enum: ['user', 'assistant', 'system'] }, - content: { type: 'string', description: 'The message text or JSON string for structured responses.' }, - timestamp: { type: 'integer', description: 'Unix timestamp in milliseconds.' }, - isSystemAction: { - type: 'boolean', - description: 'True when the message was initiated by a system action rather than a user.', - }, - activityHistory: { - type: 'array', - description: 'Tool call activity recorded during the assistant response.', - items: { $ref: '#/components/schemas/ActivityHistoryEntry' }, - }, - evidenceItems: { - type: 'array', - description: 'Evidence references (files, commits, resources) found during investigation.', - items: { type: 'object' }, - }, - totalInvestigationTimeMs: { - type: 'number', - description: 'Total wall-clock time spent generating this response.', - }, - debugContext: { $ref: '#/components/schemas/DebugContext' }, - debugToolData: { - type: 'array', - description: 'Detailed tool call/result data for debugging.', - items: { $ref: '#/components/schemas/DebugToolData' }, - }, - debugMetrics: { $ref: '#/components/schemas/DebugMetrics' }, - }, - required: ['role', 'content', 'timestamp'], - }, - - ActivityHistoryEntry: { - type: 'object', - properties: { - type: { type: 'string', description: 'The activity type (tool_call, processing, thinking, error).' }, - message: { type: 'string' }, - status: { type: 'string', enum: ['pending', 'completed', 'failed'] }, - details: { + GetEffectiveGlobalAgentSessionConfigSuccessResponse: { + allOf: [ + { $ref: '#/components/schemas/SuccessApiResponse' }, + { type: 'object', properties: { - toolDurationMs: { type: 'number' }, - totalDurationMs: { type: 'number' }, + data: { + $ref: '#/components/schemas/EffectiveAgentSessionControlPlaneConfig', + }, }, }, - toolCallId: { type: 'string' }, - resultPreview: { type: 'string', description: 'Truncated preview of the tool result.' }, - }, - required: ['type', 'message'], + ], }, - DebugContext: { - type: 'object', - description: 'Debug information about the system prompt and model used for a response.', - properties: { - systemPrompt: { type: 'string' }, - maskingStats: { + GetGlobalAgentSessionRuntimeConfigSuccessResponse: { + allOf: [ + { $ref: '#/components/schemas/SuccessApiResponse' }, + { type: 'object', - nullable: true, properties: { - totalTokensBefore: { type: 'integer' }, - totalTokensAfter: { type: 'integer' }, - maskedParts: { type: 'integer' }, - savedTokens: { type: 'integer' }, + data: { $ref: '#/components/schemas/AgentSessionRuntimeSettings' }, }, }, - provider: { type: 'string', description: 'LLM provider name (e.g. anthropic, openai).' }, - modelId: { type: 'string' }, - }, - required: ['systemPrompt', 'provider', 'modelId'], - }, - - DebugToolData: { - type: 'object', - properties: { - toolCallId: { type: 'string' }, - toolName: { type: 'string' }, - toolArgs: { type: 'object', description: 'Arguments passed to the tool. May be truncated for storage.' }, - toolResult: { description: 'Result returned by the tool. May be truncated for storage.' }, - toolDurationMs: { type: 'number' }, - }, - required: ['toolCallId', 'toolName', 'toolArgs'], - }, - - DebugMetrics: { - type: 'object', - description: 'Aggregate metrics for a single AI response.', - properties: { - iterations: { type: 'integer', description: 'Number of orchestration loop iterations.' }, - totalToolCalls: { type: 'integer' }, - totalDurationMs: { type: 'number' }, - inputTokens: { type: 'integer' }, - outputTokens: { type: 'integer' }, - inputCostPerMillion: { type: 'number' }, - outputCostPerMillion: { type: 'number' }, - }, - required: ['iterations', 'totalToolCalls', 'totalDurationMs'], + ], }, - GetAIMessagesSuccessResponse: { + GetRepoAgentSessionConfigSuccessResponse: { allOf: [ { $ref: '#/components/schemas/SuccessApiResponse' }, { @@ -3133,280 +3898,68 @@ export const openApiSpecificationForV2Api: OAS3Options = { data: { type: 'object', properties: { - messages: { - type: 'array', - items: { $ref: '#/components/schemas/ConversationMessage' }, - }, - lastActivity: { type: 'integer', nullable: true }, + repoFullName: { type: 'string' }, + config: { $ref: '#/components/schemas/AgentSessionControlPlaneConfig' }, }, - required: ['messages'], + required: ['repoFullName', 'config'], }, }, - required: ['data'], }, ], }, - DeleteAISessionSuccessResponse: { + GetAdminAgentToolInventorySuccessResponse: { allOf: [ { $ref: '#/components/schemas/SuccessApiResponse' }, { type: 'object', properties: { data: { - type: 'object', - properties: { - success: { type: 'boolean' }, - messagesCleared: { type: 'integer' }, - }, - required: ['success', 'messagesCleared'], + type: 'array', + items: { $ref: '#/components/schemas/AgentSessionToolInventoryEntry' }, }, }, - required: ['data'], }, ], }, - // =================================================================== - // AI Chat SSE Event Schemas - // =================================================================== - - SSEChunkEvent: { - type: 'object', - description: 'Streamed text content fragment from the AI response.', - properties: { - type: { type: 'string', enum: ['chunk'] }, - content: { type: 'string', description: 'A fragment of the AI response text.' }, - }, - required: ['type', 'content'], - }, - - SSEToolCallEvent: { - type: 'object', - description: - 'Emitted when the AI invokes a tool. The toolCallId can be correlated with a later SSEProcessingEvent.', - properties: { - type: { type: 'string', enum: ['tool_call'] }, - message: { type: 'string', description: 'Human-readable description of the tool being called.' }, - toolCallId: { type: 'string' }, - }, - required: ['type', 'message'], - }, - - SSEProcessingEvent: { - type: 'object', - description: 'Emitted when a tool call completes. Messages starting with a checkmark indicate success.', - properties: { - type: { type: 'string', enum: ['processing'] }, - message: { type: 'string' }, - details: { + GetAdminAgentCapabilitiesSuccessResponse: { + allOf: [ + { $ref: '#/components/schemas/SuccessApiResponse' }, + { type: 'object', properties: { - toolDurationMs: { type: 'number' }, - totalDurationMs: { type: 'number' }, + data: { $ref: '#/components/schemas/AgentCapabilityGovernanceResponse' }, }, }, - resultPreview: { type: 'string', description: 'Truncated preview of the tool result.' }, - toolCallId: { type: 'string', description: 'Correlates with the original SSEToolCallEvent.' }, - }, - required: ['type', 'message'], - }, - - SSEThinkingEvent: { - type: 'object', - description: 'Emitted when the AI is reasoning before producing output.', - properties: { - type: { type: 'string', enum: ['thinking'] }, - message: { type: 'string' }, - }, - required: ['type', 'message'], - }, - - SSEActivityErrorEvent: { - type: 'object', - description: 'Emitted for non-fatal processing errors during investigation.', - properties: { - type: { type: 'string', enum: ['error'] }, - message: { type: 'string' }, - }, - required: ['type', 'message'], - }, - - SSEEvidenceFileEvent: { - type: 'object', - description: 'A source file referenced as evidence during investigation.', - properties: { - type: { type: 'string', enum: ['evidence_file'] }, - toolCallId: { type: 'string' }, - filePath: { type: 'string' }, - repository: { type: 'string' }, - branch: { type: 'string' }, - lineStart: { type: 'integer' }, - lineEnd: { type: 'integer' }, - language: { type: 'string' }, - }, - required: ['type', 'toolCallId', 'filePath', 'repository'], + ], }, - SSEEvidenceCommitEvent: { - type: 'object', - description: 'A git commit referenced as evidence during investigation.', - properties: { - type: { type: 'string', enum: ['evidence_commit'] }, - toolCallId: { type: 'string' }, - commitUrl: { type: 'string' }, - commitMessage: { type: 'string' }, - filePaths: { type: 'array', items: { type: 'string' } }, - }, - required: ['type', 'toolCallId', 'commitUrl', 'commitMessage', 'filePaths'], - }, + // =================================================================== + // AI Runtime Config Schemas + // =================================================================== - SSEEvidenceResourceEvent: { + AIConfigStatus: { type: 'object', - description: 'A Kubernetes resource referenced as evidence during investigation.', properties: { - type: { type: 'string', enum: ['evidence_resource'] }, - toolCallId: { type: 'string' }, - resourceType: { type: 'string', description: 'Kubernetes resource kind (e.g. Pod, Deployment).' }, - resourceName: { type: 'string' }, - namespace: { type: 'string' }, - status: { type: 'string' }, + enabled: { type: 'boolean' }, + provider: { type: 'string' }, + configured: { type: 'boolean' }, }, - required: ['type', 'toolCallId', 'resourceType', 'resourceName', 'namespace'], + required: ['enabled'], }, - SSEDebugContextEvent: { - type: 'object', - description: 'Debug info about the system prompt and model selection for this response.', - properties: { - type: { type: 'string', enum: ['debug_context'] }, - systemPrompt: { type: 'string' }, - maskingStats: { + GetAIConfigSuccessResponse: { + allOf: [ + { $ref: '#/components/schemas/SuccessApiResponse' }, + { type: 'object', - nullable: true, properties: { - totalTokensBefore: { type: 'integer' }, - totalTokensAfter: { type: 'integer' }, - maskedParts: { type: 'integer' }, - savedTokens: { type: 'integer' }, + data: { $ref: '#/components/schemas/AIConfigStatus' }, }, + required: ['data'], }, - provider: { type: 'string' }, - modelId: { type: 'string' }, - }, - required: ['type', 'systemPrompt', 'provider', 'modelId'], - }, - - SSEDebugToolCallEvent: { - type: 'object', - description: 'Raw tool invocation data for debugging.', - properties: { - type: { type: 'string', enum: ['debug_tool_call'] }, - toolCallId: { type: 'string' }, - toolName: { type: 'string' }, - toolArgs: { type: 'object' }, - }, - required: ['type', 'toolCallId', 'toolName', 'toolArgs'], - }, - - SSEDebugToolResultEvent: { - type: 'object', - description: 'Raw tool result data for debugging.', - properties: { - type: { type: 'string', enum: ['debug_tool_result'] }, - toolCallId: { type: 'string' }, - toolName: { type: 'string' }, - toolResult: { description: 'The raw tool result value.' }, - toolDurationMs: { type: 'number' }, - }, - required: ['type', 'toolCallId', 'toolName', 'toolResult'], - }, - - SSEDebugMetricsEvent: { - type: 'object', - description: 'Aggregate metrics emitted once per response.', - properties: { - type: { type: 'string', enum: ['debug_metrics'] }, - iterations: { type: 'integer' }, - totalToolCalls: { type: 'integer' }, - totalDurationMs: { type: 'number' }, - inputTokens: { type: 'integer' }, - outputTokens: { type: 'integer' }, - inputCostPerMillion: { type: 'number' }, - outputCostPerMillion: { type: 'number' }, - }, - required: ['type', 'iterations', 'totalToolCalls', 'totalDurationMs'], - }, - - SSECompleteEvent: { - type: 'object', - description: 'Signals the end of a plain-text AI response. This is the final event in the stream.', - properties: { - type: { type: 'string', enum: ['complete'] }, - totalInvestigationTimeMs: { type: 'number' }, - assistantTimestamp: { - type: 'number', - description: 'Server-side timestamp (epoch ms) used for persisting the final assistant message.', - }, - }, - required: ['type', 'totalInvestigationTimeMs'], - }, - - SSECompleteJsonEvent: { - type: 'object', - description: - 'Signals the end of a structured JSON AI response (e.g. investigation_complete). ' + - 'The content field contains the full JSON string. Sent before SSECompleteEvent.', - properties: { - type: { type: 'string', enum: ['complete_json'] }, - content: { type: 'string', description: 'The full JSON response as a string.' }, - preamble: { - type: 'string', - description: - 'Optional plain-text summary emitted before structured JSON when mixed text+JSON model output is split.', - }, - totalInvestigationTimeMs: { type: 'number' }, - }, - required: ['type', 'content', 'totalInvestigationTimeMs'], - }, - - SSEErrorEvent: { - type: 'object', - description: - 'Streamed error event. Errors during SSE streaming are sent as events (not HTTP errors) ' + - 'because the HTTP 200 status has already been committed.', - properties: { - error: { type: 'boolean', enum: [true] }, - userMessage: { type: 'string', description: 'Human-readable error description.' }, - category: { - type: 'string', - enum: ['rate-limited', 'transient', 'deterministic', 'ambiguous'], - description: - 'Error classification. rate-limited: provider rate limit hit (retryable). ' + - 'transient: temporary provider outage (retryable). ' + - 'deterministic: auth error, bad request, or config issue (not retryable). ' + - 'ambiguous: unknown error state (retryable).', - }, - suggestedAction: { - type: 'string', - enum: ['retry', 'switch-model', 'check-config'], - nullable: true, - description: 'Recommended client action.', - }, - retryAfter: { - type: 'number', - nullable: true, - description: 'Seconds to wait before retrying (only for rate-limited errors).', - }, - modelName: { type: 'string' }, - code: { - type: 'string', - description: - 'Machine-readable error code. Known codes: AI_AGENT_DISABLED, CONTEXT_ERROR, ' + - 'LLM_INIT_ERROR, LLM_API_ERROR, CIRCUIT_BREAKER_OPEN.', - }, - }, - required: ['error', 'userMessage', 'category', 'suggestedAction', 'retryAfter', 'modelName'], + ], }, // =================================================================== diff --git a/src/shared/types/aiChat.ts b/src/shared/types/aiChat.ts deleted file mode 100644 index 9ef54ec3..00000000 --- a/src/shared/types/aiChat.ts +++ /dev/null @@ -1,164 +0,0 @@ -/** - * Copyright 2025 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. - */ - -export interface ChunkEvent { - type: 'chunk'; - content: string; -} - -export interface CompleteEvent { - type: 'complete'; - totalInvestigationTimeMs: number; - assistantTimestamp?: number; -} - -export interface CompleteJsonEvent { - type: 'complete_json'; - content: string; - totalInvestigationTimeMs: number; - preamble?: string; -} - -export interface ToolCallEvent { - type: 'tool_call'; - message: string; - toolCallId?: string; -} - -export interface ProcessingEvent { - type: 'processing'; - message: string; - details?: { - toolDurationMs?: number; - totalDurationMs?: number; - }; - resultPreview?: string; - toolCallId?: string; -} - -export interface ThinkingEvent { - type: 'thinking'; - message: string; -} - -export interface ActivityErrorEvent { - type: 'error'; - message: string; -} - -export interface EvidenceFileEvent { - type: 'evidence_file'; - toolCallId: string; - filePath: string; - repository: string; - branch?: string; - lineStart?: number; - lineEnd?: number; - language?: string; -} - -export interface EvidenceCommitEvent { - type: 'evidence_commit'; - toolCallId: string; - commitUrl: string; - commitMessage: string; - filePaths: string[]; -} - -export interface EvidenceResourceEvent { - type: 'evidence_resource'; - toolCallId: string; - resourceType: string; - resourceName: string; - namespace: string; - status?: string; -} - -export type AIChatEvidenceEvent = EvidenceFileEvent | EvidenceCommitEvent | EvidenceResourceEvent; - -export interface SSEErrorEvent { - error: true; - userMessage: string; - category: 'rate-limited' | 'transient' | 'deterministic' | 'ambiguous'; - suggestedAction: 'retry' | 'switch-model' | 'check-config' | null; - retryAfter: number | null; - modelName: string; - code?: string; -} - -export interface DebugContextEvent { - type: 'debug_context'; - systemPrompt: string; - maskingStats: { - totalTokensBefore: number; - totalTokensAfter: number; - maskedParts: number; - savedTokens: number; - } | null; - provider: string; - modelId: string; -} - -export interface DebugToolCallEvent { - type: 'debug_tool_call'; - toolCallId: string; - toolName: string; - toolArgs: Record; -} - -export interface DebugToolResultEvent { - type: 'debug_tool_result'; - toolCallId: string; - toolName: string; - toolResult: unknown; - toolDurationMs?: number; -} - -export interface DebugMetricsEvent { - type: 'debug_metrics'; - iterations: number; - totalToolCalls: number; - totalDurationMs: number; - inputTokens?: number; - outputTokens?: number; - inputCostPerMillion?: number; - outputCostPerMillion?: number; -} - -export type AIChatDebugEvent = DebugContextEvent | DebugToolCallEvent | DebugToolResultEvent | DebugMetricsEvent; - -export type AIChatActivityEvent = ToolCallEvent | ProcessingEvent | ThinkingEvent | ActivityErrorEvent; - -export type AIChatSSEEvent = - | ChunkEvent - | CompleteEvent - | CompleteJsonEvent - | AIChatActivityEvent - | AIChatEvidenceEvent - | AIChatDebugEvent; - -export function isEvidenceEvent(event: AIChatSSEEvent): event is AIChatEvidenceEvent { - return event.type === 'evidence_file' || event.type === 'evidence_commit' || event.type === 'evidence_resource'; -} - -export function isDebugEvent(event: AIChatSSEEvent): event is AIChatDebugEvent { - return ( - event.type === 'debug_context' || - event.type === 'debug_tool_call' || - event.type === 'debug_tool_result' || - event.type === 'debug_metrics' - ); -}