diff --git a/api/server/controllers/agents/v1.js b/api/server/controllers/agents/v1.js index efc36c9940..3f6b4c5f53 100644 --- a/api/server/controllers/agents/v1.js +++ b/api/server/controllers/agents/v1.js @@ -1166,6 +1166,12 @@ const getListAgentsHandler = async (req, res) => { requiredPermission = PermissionBits.VIEW; } const canReturnSkillConfig = hasEditBit(requiredPermission); + /** + * Derived from the same bit as `canReturnSkillConfig` but answering a different question: + * skill-config exposure versus edit-permission reporting. An EDIT-scoped request matches + * only editable agents, so it needs no second lookup to know which ones those are. + */ + const needsEditableLookup = !hasEditBit(requiredPermission); // Base filter const filter = {}; @@ -1188,33 +1194,99 @@ const getListAgentsHandler = async (req, res) => { filter.$or = [{ name: regex }, { description: regex }]; } - // Get agent IDs the user has VIEW access to via ACL - const accessibleIds = await findAccessibleResources({ - userId, - role: req.user.role, - resourceType: ResourceType.AGENT, - requiredPermissions: requiredPermission, - }); + const cache = getLogStores(CacheKeys.S3_EXPIRY_INTERVAL); + const refreshKey = `${userId}:agents_avatar_refresh`; - const publiclyAccessibleIds = await findPubliclyAccessibleResources({ - resourceType: ResourceType.AGENT, - requiredPermissions: PermissionBits.VIEW, - }); + /** + * These reads share no inputs, so they resolve together rather than chaining round + * trips ahead of the list query. The viewer skill scope and the editable set are only + * consumed when the page is non-empty; dispatching them here trades a wasted lookup on + * the (cheap) zero-agent path for one less serial hop on every populated page. + * + * `editableIds` lets a VIEW-scoped response mark which agents the caller may also edit, + * so consumers wanting just the editable subset can filter one shared VIEW fetch rather + * than issuing a second full paginated walk under an EDIT-scoped cache key. Requests + * that already ask for EDIT get it for free: everything they match is editable. + * + * `idOnTheSource` is forwarded so `getUserPrincipals` resolves identity without reading + * the user document; the auth strategies already normalize it to a value or null. Each + * omission would cost this handler another `User.findById`, once per lookup. + */ + const { idOnTheSource } = req.user; + const [ + accessibleIds, + publiclyAccessibleIds, + cachedRefreshEntry, + accessibleSkillIds, + editableIds, + ] = await Promise.all([ + findAccessibleResources({ + userId, + role: req.user.role, + idOnTheSource, + resourceType: ResourceType.AGENT, + requiredPermissions: requiredPermission, + }), + findPubliclyAccessibleResources({ + resourceType: ResourceType.AGENT, + requiredPermissions: PermissionBits.VIEW, + }), + cache.get(refreshKey), + canReturnSkillConfig + ? null + : findAccessibleResources({ + userId, + role: req.user.role, + idOnTheSource, + resourceType: ResourceType.SKILL, + requiredPermissions: PermissionBits.VIEW, + }), + needsEditableLookup + ? findAccessibleResources({ + userId, + role: req.user.role, + idOnTheSource, + resourceType: ResourceType.AGENT, + requiredPermissions: PermissionBits.EDIT, + }) + : null, + ]); + + const isValidCachedRefresh = + cachedRefreshEntry != null && + typeof cachedRefreshEntry === 'object' && + cachedRefreshEntry.urlCache != null; /** * Refresh all S3 avatars for this user's accessible agent set (not only the current page) - * This addresses page-size limits preventing refresh of agents beyond the first page + * This addresses page-size limits preventing refresh of agents beyond the first page. + * + * Scoped to agents that actually carry an S3 avatar so the `MAX_AVATAR_REFRESH_AGENTS` + * budget is spent on agents that can do work. Unfiltered, that budget is the most + * recently updated accessible agents regardless of avatar, and because a refresh writes + * through `updateAgent` and advances `updatedAt`, the window is self-reinforcing: an + * S3-avatar agent ranked past the budget never enters it and its presigned URL is never + * regenerated. The predicate is not indexed (`avatar` is `Mixed`), so this trades docs + * examined for that coverage. + * + * Must settle BEFORE the list query below, and is deliberately not parallelized with + * it. `updateAgent` writes through `findOneAndUpdate` on a `timestamps: true` schema, + * so refreshing an avatar advances `updatedAt`, the very field + * `getListAgentsByAccess` sorts and cursors on. A refresh landing after the first + * page's snapshot would move that agent ahead of the returned cursor, dropping it + * from every later page and silently truncating the caller's flattened list. + * Serializing costs nothing on the common path: a cache hit returns below without + * issuing any query, so only the once-per-30-minutes miss pays for the ordering. */ - const cache = getLogStores(CacheKeys.S3_EXPIRY_INTERVAL); - const refreshKey = `${userId}:agents_avatar_refresh`; - let cachedRefresh = await cache.get(refreshKey); - const isValidCachedRefresh = - cachedRefresh != null && typeof cachedRefresh === 'object' && cachedRefresh.urlCache != null; - if (!isValidCachedRefresh) { + const resolveAvatarRefresh = async () => { + if (isValidCachedRefresh) { + logger.debug('[/Agents] S3 avatar refresh already checked, skipping'); + return cachedRefreshEntry; + } try { const fullList = await db.getListAgentsByAccess({ accessibleIds, - otherParams: {}, + otherParams: { 'avatar.source': FileSources.s3 }, limit: MAX_AVATAR_REFRESH_AGENTS, after: null, }); @@ -1224,14 +1296,16 @@ const getListAgentsHandler = async (req, res) => { refreshS3Url, updateAgent: db.updateAgent, }); - cachedRefresh = { urlCache }; - await cache.set(refreshKey, cachedRefresh, Time.THIRTY_MINUTES); + const refreshEntry = { urlCache }; + await cache.set(refreshKey, refreshEntry, Time.THIRTY_MINUTES); + return refreshEntry; } catch (err) { logger.error('[/Agents] Error refreshing avatars for full list: %o', err); + return null; } - } else { - logger.debug('[/Agents] S3 avatar refresh already checked, skipping'); - } + }; + + const cachedRefresh = await resolveAvatarRefresh(); // Use the new ACL-aware function const data = await db.getListAgentsByAccess({ @@ -1247,20 +1321,13 @@ const getListAgentsHandler = async (req, res) => { return res.json(data); } - let accessibleSkillSet = null; - if (!canReturnSkillConfig) { - const accessibleSkillIds = await findAccessibleResources({ - userId, - role: req.user.role, - resourceType: ResourceType.SKILL, - requiredPermissions: PermissionBits.VIEW, - }); - accessibleSkillSet = new Set( - mergeDeploymentSkillIds(accessibleSkillIds).map((oid) => oid.toString()), - ); - } + const accessibleSkillSet = canReturnSkillConfig + ? null + : new Set(mergeDeploymentSkillIds(accessibleSkillIds).map((oid) => oid.toString())); const publicSet = new Set(publiclyAccessibleIds.map((oid) => oid.toString())); + /** Null for EDIT-scoped requests, where every matched agent is editable by definition. */ + const editableSet = editableIds ? new Set(editableIds.map((oid) => oid.toString())) : null; const agentsWithContacts = await attachOwnerContacts(agents); const urlCache = cachedRefresh?.urlCache; @@ -1272,6 +1339,7 @@ const getListAgentsHandler = async (req, res) => { if (agent?._id && publicSet.has(agent._id.toString())) { agent.isPublic = true; } + agent.isEditable = editableSet == null || editableSet.has(agent?._id?.toString()); if ( urlCache && agent?.id && @@ -1280,9 +1348,8 @@ const getListAgentsHandler = async (req, res) => { ) { agent.avatar = { ...agent.avatar, filepath: urlCache[agent.id] }; } - } catch (e) { - // Silently ignore mapping errors - void e; + } catch (err) { + logger.warn('[/Agents] Error mapping agent %s for list response: %o', agent?.id, err); } return agent; }); diff --git a/api/server/controllers/agents/v1.spec.js b/api/server/controllers/agents/v1.spec.js index 02f58d9929..67000ea4b3 100644 --- a/api/server/controllers/agents/v1.spec.js +++ b/api/server/controllers/agents/v1.spec.js @@ -1676,6 +1676,65 @@ describe('Agent Controllers - Mass Assignment Protection', () => { expect(response.data[0].owner_contact).toBeUndefined(); }); + test('should mark isEditable per agent on a VIEW-scoped list', async () => { + mockReq.user.id = userA.toString(); + mockReq.query = { requiredPermission: String(PermissionBits.VIEW) }; + /** VIEW reaches all three; the EDIT lookup only reaches agentA1. */ + findAccessibleResources.mockImplementation(({ resourceType, requiredPermissions }) => { + if (resourceType === 'agent' && requiredPermissions === PermissionBits.EDIT) { + return Promise.resolve([agentA1._id]); + } + if (resourceType === 'agent') { + return Promise.resolve([agentA1._id, agentA2._id, agentA3._id]); + } + return Promise.resolve([]); + }); + findPubliclyAccessibleResources.mockResolvedValue([]); + + await getListAgentsHandler(mockReq, mockRes); + + const byId = Object.fromEntries( + mockRes.json.mock.calls[0][0].data.map((a) => [a.id, a.isEditable]), + ); + expect(byId[agentA1.id]).toBe(true); + expect(byId[agentA2.id]).toBe(false); + expect(byId[agentA3.id]).toBe(false); + }); + + test('should forward idOnTheSource to every ACL lookup', async () => { + /** Without it `getUserPrincipals` reads the user document once per lookup, so the + * handler pays an extra `User.findById` for each permission it resolves. */ + mockReq.user.id = userA.toString(); + mockReq.user.idOnTheSource = 'external-oid-1'; + findAccessibleResources.mockResolvedValue([agentA1._id]); + findPubliclyAccessibleResources.mockResolvedValue([]); + + await getListAgentsHandler(mockReq, mockRes); + + expect(findAccessibleResources.mock.calls.length).toBeGreaterThan(1); + for (const [args] of findAccessibleResources.mock.calls) { + expect(args.idOnTheSource).toBe('external-oid-1'); + } + }); + + test('should mark every agent editable when the request is already EDIT-scoped', async () => { + mockReq.user.id = userA.toString(); + mockReq.query = { requiredPermission: String(PermissionBits.EDIT) }; + findAccessibleResources.mockResolvedValue([agentA1._id, agentA2._id]); + findPubliclyAccessibleResources.mockResolvedValue([]); + + await getListAgentsHandler(mockReq, mockRes); + + const response = mockRes.json.mock.calls[0][0]; + expect(response.data.every((a) => a.isEditable === true)).toBe(true); + /** No extra EDIT lookup: an EDIT-scoped match is editable by definition. */ + const editCalls = findAccessibleResources.mock.calls.filter( + ([args]) => + args.resourceType === 'agent' && args.requiredPermissions === PermissionBits.EDIT, + ); + expect(editCalls).toHaveLength(1); + }); + test('should return only expected safe list fields for VIEW callers', async () => { const hiddenSkillId = new mongoose.Types.ObjectId(); await Agent.findByIdAndUpdate(agentA1._id, { @@ -1721,6 +1780,7 @@ describe('Agent Controllers - Mass Assignment Protection', () => { 'conversation_starters', 'description', 'id', + 'isEditable', 'is_promoted', 'name', 'support_contact', @@ -2235,6 +2295,110 @@ describe('Agent Controllers - Mass Assignment Protection', () => { expect(mockRes.json).toHaveBeenCalled(); }); + test('should finish avatar writes before snapshotting the paginated list query', async () => { + /** `updateAgent` bumps `updatedAt`, which is the field `getListAgentsByAccess` + * sorts and cursors on. If the list query snapshots before a refresh write + * lands, that agent jumps ahead of the returned cursor and vanishes from every + * later page. Assert the ordering rather than the symptom, which only shows up + * on multi-page S3 accounts under a specific interleaving. */ + const db = require('~/models'); + const order = []; + /** Yield a macrotask so a parallelized refresh would lose the race, the way a real + * S3 presign round trip does. */ + refreshS3Url.mockImplementation(async () => { + await new Promise((resolve) => setTimeout(resolve, 5)); + order.push('avatar-write'); + return 'new-s3-path.jpg'; + }); + const realList = db.getListAgentsByAccess; + const listSpy = jest.spyOn(db, 'getListAgentsByAccess').mockImplementation(async (params) => { + if (params.includeSkillConfig) { + order.push('list-query'); + return { object: 'list', data: [], has_more: false, after: null }; + } + return realList(params); + }); + mockCache.get.mockResolvedValue(false); + findAccessibleResources.mockResolvedValue([agentWithS3Avatar._id]); + findPubliclyAccessibleResources.mockResolvedValue([]); + + const mockReq = { user: { id: userA.toString(), role: 'USER' }, query: {} }; + const mockRes = { status: jest.fn().mockReturnThis(), json: jest.fn().mockReturnThis() }; + + try { + await getListAgentsHandler(mockReq, mockRes); + expect(order).toContain('avatar-write'); + expect(order.indexOf('avatar-write')).toBeLessThan(order.indexOf('list-query')); + } finally { + listSpy.mockRestore(); + refreshS3Url.mockReset(); + } + }); + + test('should serve the refreshed filepath in the same response on cache miss', async () => { + const agentId = agentWithS3Avatar.id; + mockCache.get.mockResolvedValue(false); + findAccessibleResources.mockResolvedValue([agentWithS3Avatar._id]); + findPubliclyAccessibleResources.mockResolvedValue([]); + refreshS3Url.mockResolvedValue('new-s3-path.jpg'); + + const mockReq = { + user: { id: userA.toString(), role: 'USER' }, + query: {}, + }; + const mockRes = { + status: jest.fn().mockReturnThis(), + json: jest.fn().mockReturnThis(), + }; + + await getListAgentsHandler(mockReq, mockRes); + + const responseData = mockRes.json.mock.calls[0][0]; + const agent = responseData.data.find((a) => a.id === agentId); + /** The refresh runs alongside the list query, so the refreshed path must reach the + * response through `urlCache` rather than through what the list query read. */ + expect(agent.avatar.filepath).toBe('new-s3-path.jpg'); + }); + + test('should scope the refresh query to S3 avatars without filtering the list query', async () => { + const db = require('~/models'); + const listSpy = jest.spyOn(db, 'getListAgentsByAccess'); + mockCache.get.mockResolvedValue(false); + findAccessibleResources.mockResolvedValue([agentWithLocalAvatar._id]); + findPubliclyAccessibleResources.mockResolvedValue([]); + + const mockReq = { + user: { id: userA.toString(), role: 'USER' }, + query: {}, + }; + const mockRes = { + status: jest.fn().mockReturnThis(), + json: jest.fn().mockReturnThis(), + }; + + try { + await getListAgentsHandler(mockReq, mockRes); + + /** The refresh pass must query only S3-avatar agents — `refreshListAvatars` + * skips non-S3 entries anyway, so without this assertion the filter could + * regress to `{}` (reloading the whole accessible set) unnoticed. */ + expect(listSpy).toHaveBeenCalledWith( + expect.objectContaining({ otherParams: { 'avatar.source': FileSources.s3 } }), + ); + /** The user-facing list query keeps the request filter, not the refresh scope. */ + expect(listSpy).toHaveBeenCalledWith( + expect.objectContaining({ includeSkillConfig: true, otherParams: {} }), + ); + + expect(refreshS3Url).not.toHaveBeenCalled(); + const responseData = mockRes.json.mock.calls[0][0]; + const agent = responseData.data.find((a) => a.id === agentWithLocalAvatar.id); + expect(agent.avatar.filepath).toBe('local-path.jpg'); + } finally { + listSpy.mockRestore(); + } + }); + test('should refresh avatars for all accessible agents (VIEW permission)', async () => { mockCache.get.mockResolvedValue(false); // User A has access to both their own agent and userB's agent diff --git a/api/server/services/PermissionService.js b/api/server/services/PermissionService.js index efc1e5013d..6900a90702 100644 --- a/api/server/services/PermissionService.js +++ b/api/server/services/PermissionService.js @@ -238,11 +238,19 @@ const getResourcePermissionsMap = async ({ userId, role, resourceType, resourceI * @param {Object} params - Parameters for finding accessible resources * @param {string|mongoose.Types.ObjectId} params.userId - The ID of the user * @param {string} [params.role] - Optional user role (if not provided, will query from DB) + * @param {string|null} [params.idOnTheSource] - Optional external member id. `null` means "known to + * be absent" (local user); only `undefined` makes `getUserPrincipals` read the user document. * @param {string} params.resourceType - Type of resource (e.g., 'agent') * @param {number} params.requiredPermissions - The minimum permission bits required (e.g., 1 for VIEW, 3 for VIEW+EDIT) * @returns {Promise} Array of resource IDs */ -const findAccessibleResources = async ({ userId, role, resourceType, requiredPermissions }) => { +const findAccessibleResources = async ({ + userId, + role, + idOnTheSource, + resourceType, + requiredPermissions, +}) => { try { if (typeof requiredPermissions !== 'number' || requiredPermissions < 1) { throw new Error('requiredPermissions must be a positive number'); @@ -251,7 +259,7 @@ const findAccessibleResources = async ({ userId, role, resourceType, requiredPer validateResourceType(resourceType); // Get all principals for the user (user + groups + public) - const principalsList = await db.getUserPrincipals({ userId, role }); + const principalsList = await db.getUserPrincipals({ userId, role, idOnTheSource }); if (principalsList.length === 0) { return []; diff --git a/api/server/services/PermissionService.spec.js b/api/server/services/PermissionService.spec.js index f1220faf63..53b281aefb 100644 --- a/api/server/services/PermissionService.spec.js +++ b/api/server/services/PermissionService.spec.js @@ -632,6 +632,24 @@ describe('PermissionService', () => { }); }); + test('should forward idOnTheSource so principal resolution can skip the user lookup', async () => { + getUserPrincipals.mockResolvedValue([ + { principalType: PrincipalType.USER, principalId: userId }, + ]); + + await findAccessibleResources({ + userId, + role: 'USER', + idOnTheSource: null, + resourceType: ResourceType.AGENT, + requiredPermissions: 1, // VIEW + }); + + expect(getUserPrincipals).toHaveBeenCalledWith( + expect.objectContaining({ idOnTheSource: null }), + ); + }); + test('should find resources user can view', async () => { // Mock getUserPrincipals to return user principal getUserPrincipals.mockResolvedValue([ diff --git a/client/src/components/Chat/Menus/Endpoints/ModelSelectorContext.tsx b/client/src/components/Chat/Menus/Endpoints/ModelSelectorContext.tsx index b09bef94be..9e450fc75e 100644 --- a/client/src/components/Chat/Menus/Endpoints/ModelSelectorContext.tsx +++ b/client/src/components/Chat/Menus/Endpoints/ModelSelectorContext.tsx @@ -1,6 +1,11 @@ import React, { createContext, useContext, useState, useMemo, useCallback } from 'react'; import debounce from 'lodash/debounce'; -import { EModelEndpoint, isAgentsEndpoint, isAssistantsEndpoint } from 'librechat-data-provider'; +import { + EModelEndpoint, + PermissionBits, + isAgentsEndpoint, + isAssistantsEndpoint, +} from 'librechat-data-provider'; import type * as t from 'librechat-data-provider'; import type { Endpoint, SelectedValues } from '~/common'; import { @@ -82,11 +87,27 @@ export function ModelSelectorProvider({ children, startupConfig }: ModelSelector }, [startupConfig, agentsMap]); const permissionLevel = useAgentDefaultPermissionLevel(); - const { data: agents = null } = useListAgentsQuery( - { requiredPermission: permissionLevel }, - { - select: (data) => data?.data, + /** + * Always query the VIEW scope so this shares one cache entry (and one paginated walk) + * with `useAgentsMap` and `useMentions`. Asking for EDIT here spawned a second full + * fetch under its own key, holding a duplicate copy of the whole agent list. The + * marketplace's "my agents" framing is preserved by filtering on `isEditable`, which + * the list endpoint resolves from the same ACL read it already performs. + */ + const wantsEditableOnly = permissionLevel === PermissionBits.EDIT; + const selectAgents = useCallback( + (data: t.AgentListResponse) => { + const list = data?.data; + if (!wantsEditableOnly) { + return list; + } + return list?.filter((agent) => agent.isEditable !== false); }, + [wantsEditableOnly], + ); + const { data: agents = null } = useListAgentsQuery( + { requiredPermission: PermissionBits.VIEW }, + { select: selectAgents }, ); const { mappedEndpoints, endpointRequiresUserKey } = useEndpoints({ diff --git a/client/src/components/Chat/Menus/Endpoints/components/EndpointItem.tsx b/client/src/components/Chat/Menus/Endpoints/components/EndpointItem.tsx index 3f446cef87..6692541d7b 100644 --- a/client/src/components/Chat/Menus/Endpoints/components/EndpointItem.tsx +++ b/client/src/components/Chat/Menus/Endpoints/components/EndpointItem.tsx @@ -1,4 +1,4 @@ -import { useMemo } from 'react'; +import { useCallback, useMemo } from 'react'; import { VisuallyHidden } from '@ariakit/react'; import { Spinner, TooltipAnchor } from '@librechat/client'; import { CheckCircle2, MousePointerClick, SettingsIcon } from 'lucide-react'; @@ -6,12 +6,13 @@ import { EModelEndpoint, isAgentsEndpoint, isAssistantsEndpoint } from 'librecha import type { TModelSpec } from 'librechat-data-provider'; import type { Endpoint } from '~/common'; import { CustomMenu as Menu, CustomMenuItem as MenuItem, CustomMenuSeparator } from '../CustomMenu'; +import { renderEndpointModels, VIRTUALIZE_THRESHOLD } from './EndpointModelItem'; import MarketplaceItem, { marketplaceSearchMatches } from './Marketplace'; import { filterModels, shouldRenderEndpointOption } from '../utils'; import { useModelSelectorContext } from '../ModelSelectorContext'; -import { renderEndpointModels } from './EndpointModelItem'; +import VirtualizedModelList from './VirtualizedModelList'; +import { useFavorites, useLocalize } from '~/hooks'; import { ModelSpecItem } from './ModelSpecItem'; -import { useLocalize } from '~/hooks'; import { cn } from '~/utils'; interface EndpointItemProps { @@ -133,21 +134,123 @@ function EndpointMenuContent({ endpoint.showMarketplace === true && marketplaceSearchMatches(searchValue, localize); const hasSelectableRows = endpointSpecs.length > 0 || renderedModels.length > 0; + /** + * Once the model list is windowed, the DOM no longer holds every option, so a screen + * reader would infer position and total from the mounted slice alone. Declare them + * explicitly across the whole listbox — mixing declared and inferred values within one + * set is worse than either — and leave them off entirely when nothing is virtualized. + */ + const precedingOptionCount = (showMarketplace ? 1 : 0) + endpointSpecs.length; + const isVirtualized = renderedModels.length > VIRTUALIZE_THRESHOLD; + const listboxSetSize = isVirtualized ? precedingOptionCount + renderedModels.length : undefined; + return ( <> - {showMarketplace && } + {showMarketplace && ( + + )} {showMarketplace && hasSelectableRows && } - {endpointSpecs.map((spec: TModelSpec) => ( - + {endpointSpecs.map((spec: TModelSpec, specIndex: number) => ( + ))} - {filteredModels - ? renderEndpointModels(endpoint, endpoint.models || [], filteredModels, endpointIndex) - : endpoint.models && - renderEndpointModels(endpoint, endpoint.models, undefined, endpointIndex)} + ); } +/** + * Owns the model rows for one endpoint. `useFavorites` is called once here rather + * than inside each row: it opens a jotai subscription, a React Query subscription + * and a mutation per call site, which at agent-list scale was thousands of live + * subscriptions for one dropdown. + */ +function EndpointModels({ + endpoint, + renderedModels, + endpointIndex, + searchValue, + precedingOptionCount, +}: { + endpoint: Endpoint; + renderedModels: string[]; + endpointIndex: number; + searchValue: string; + precedingOptionCount: number; +}) { + const { isFavoriteModel, toggleFavoriteModel, isFavoriteAgent, toggleFavoriteAgent } = + useFavorites(); + const isAgent = isAgentsEndpoint(endpoint.value); + + const isFavorite = useCallback( + (modelId: string) => + isAgent ? isFavoriteAgent(modelId) : isFavoriteModel(modelId, endpoint.value), + [isAgent, isFavoriteAgent, isFavoriteModel, endpoint.value], + ); + const onToggleFavorite = useCallback( + (modelId: string) => { + if (isAgent) { + toggleFavoriteAgent(modelId); + } else { + toggleFavoriteModel({ model: modelId, endpoint: endpoint.value }); + } + }, + [isAgent, toggleFavoriteAgent, toggleFavoriteModel, endpoint.value], + ); + + const models = useMemo(() => endpoint.models ?? [], [endpoint.models]); + const globalByName = useMemo( + () => new Map(models.map((model) => [model.name, model.isGlobal ?? false])), + [models], + ); + + if (!renderedModels.length) { + return null; + } + + if (renderedModels.length > VIRTUALIZE_THRESHOLD) { + return ( + /** + * Keyed on the filter so a new result set starts at the top. `Grid` keeps its + * scroll offset across prop changes and, when the row count shrinks, clamps it to + * `totalRowsHeight - height` — the END of the shorter list. Without this a user who + * scrolled deep and then searched would land on the tail matches, with only those + * rows mounted and therefore reachable by keyboard. + */ + + ); + } + + return renderEndpointModels(endpoint, models, renderedModels, endpointIndex, { + isFavorite, + onToggleFavorite, + }); +} + export function EndpointItem({ endpoint, endpointIndex }: EndpointItemProps) { const localize = useLocalize(); const { diff --git a/client/src/components/Chat/Menus/Endpoints/components/EndpointModelItem.tsx b/client/src/components/Chat/Menus/Endpoints/components/EndpointModelItem.tsx index aa4cb8efaa..92ba3aa421 100644 --- a/client/src/components/Chat/Menus/Endpoints/components/EndpointModelItem.tsx +++ b/client/src/components/Chat/Menus/Endpoints/components/EndpointModelItem.tsx @@ -3,17 +3,38 @@ import { VisuallyHidden } from '@ariakit/react'; import { CheckCircle2, EarthIcon, Pin, PinOff } from 'lucide-react'; import { isAgentsEndpoint, isAssistantsEndpoint } from 'librechat-data-provider'; import type { Endpoint } from '~/common'; -import { useFavorites, useLocalize, useIsActiveItem } from '~/hooks'; import { useModelSelectorContext } from '../ModelSelectorContext'; import { CustomMenuItem as MenuItem } from '../CustomMenu'; +import useActiveItem from '../useActiveItem'; +import { useLocalize } from '~/hooks'; import { cn } from '~/utils'; interface EndpointModelItemProps { modelId: string | null; endpoint: Endpoint; + /** Resolved by the parent from the same array it maps, so the row does not rescan it. */ + isGlobal?: boolean; + isFavorite: boolean; + onToggleFavorite: (modelId: string) => void; + /** + * Only set when the list is virtualized. The mounted rows are then a small window over + * a much larger set, so the position a screen reader would infer from the DOM is wrong; + * these carry the real position and total. Left undefined otherwise, where the DOM holds + * every option and the implicit values are already correct. + */ + posInSet?: number; + setSize?: number; } -export function EndpointModelItem({ modelId, endpoint }: EndpointModelItemProps) { +function EndpointModelItemComponent({ + modelId, + endpoint, + isGlobal = false, + isFavorite, + onToggleFavorite, + posInSet, + setSize, +}: EndpointModelItemProps) { const localize = useLocalize(); const { handleSelectModel, selectedValues } = useModelSelectorContext(); const { @@ -23,21 +44,15 @@ export function EndpointModelItem({ modelId, endpoint }: EndpointModelItemProps) } = selectedValues; const isSelected = !selectedSpec && selectedEndpoint === endpoint.value && selectedModel === modelId; - const { isFavoriteModel, toggleFavoriteModel, isFavoriteAgent, toggleFavoriteAgent } = - useFavorites(); - const { ref: itemRef, isActive } = useIsActiveItem(); + const { ref: itemRef, isActive } = useActiveItem(); - let isGlobal = false; let modelName = modelId; const avatarUrl = endpoint?.modelIcons?.[modelId ?? ''] || null; // Use custom names if available if (endpoint && modelId && isAgentsEndpoint(endpoint.value) && endpoint.agentNames?.[modelId]) { modelName = endpoint.agentNames[modelId]; - - const modelInfo = endpoint?.models?.find((m) => m.name === modelId); - isGlobal = modelInfo?.isGlobal ?? false; } else if ( endpoint && modelId && @@ -47,26 +62,12 @@ export function EndpointModelItem({ modelId, endpoint }: EndpointModelItemProps) modelName = endpoint.assistantNames[modelId]; } - const isAgent = isAgentsEndpoint(endpoint.value); - const isFavorite = isAgent - ? isFavoriteAgent(modelId ?? '') - : isFavoriteModel(modelId ?? '', endpoint.value); - - const handleFavoriteToggle = () => { + const handleFavoriteClick = (e: React.MouseEvent) => { + e.stopPropagation(); if (!modelId) { return; } - - if (isAgent) { - toggleFavoriteAgent(modelId); - } else { - toggleFavoriteModel({ model: modelId, endpoint: endpoint.value }); - } - }; - - const handleFavoriteClick = (e: React.MouseEvent) => { - e.stopPropagation(); - handleFavoriteToggle(); + onToggleFavorite(modelId); }; const renderAvatar = () => { @@ -101,6 +102,8 @@ export function EndpointModelItem({ modelId, endpoint }: EndpointModelItemProps) ref={itemRef} onClick={() => handleSelectModel(endpoint, modelId ?? '')} aria-selected={isSelected || undefined} + aria-posinset={posInSet} + aria-setsize={setSize} className="group flex w-full cursor-pointer items-center justify-between rounded-lg px-2 text-sm" >
@@ -141,23 +144,45 @@ export function EndpointModelItem({ modelId, endpoint }: EndpointModelItemProps) ); } +export const EndpointModelItem = React.memo(EndpointModelItemComponent); + +/** + * Above this many rows the list is windowed. Below it, rendering everything keeps + * Ariakit's composite registry complete, so arrow-key navigation and typeahead + * reach every row — which is the behaviour virtualization has to work to preserve. + */ +export const VIRTUALIZE_THRESHOLD = 100; + export function renderEndpointModels( endpoint: Endpoint | null, models: Array<{ name: string; isGlobal?: boolean }>, filteredModels?: string[], endpointIndex?: number, + favorites?: { + isFavorite: (modelId: string) => boolean; + onToggleFavorite: (modelId: string) => void; + }, ) { + if (!endpoint) { + return null; + } const modelsToRender = filteredModels || models.map((model) => model.name); const indexSuffix = endpointIndex != null ? `-${endpointIndex}` : ''; + const isFavorite = favorites?.isFavorite ?? (() => false); + const onToggleFavorite = favorites?.onToggleFavorite ?? (() => {}); - return modelsToRender.map( - (modelId, modelIndex) => - endpoint && ( - - ), - ); + /** `models` carries `isGlobal`; without this map each row rescanned the whole + * array to recover it, which is quadratic in the number of agents. */ + const globalByName = new Map(models.map((model) => [model.name, model.isGlobal ?? false])); + + return modelsToRender.map((modelId, modelIndex) => ( + + )); } diff --git a/client/src/components/Chat/Menus/Endpoints/components/Marketplace.tsx b/client/src/components/Chat/Menus/Endpoints/components/Marketplace.tsx index 577e537866..60a02687ab 100644 --- a/client/src/components/Chat/Menus/Endpoints/components/Marketplace.tsx +++ b/client/src/components/Chat/Menus/Endpoints/components/Marketplace.tsx @@ -22,9 +22,14 @@ export function marketplaceSearchMatches(searchValue: string, localize: Localize export default function MarketplaceItem({ className, label, + posInSet, + setSize, }: { className?: string; label: string; + /** Set when the sibling model list is virtualized; see `VirtualizedModelList`. */ + posInSet?: number; + setSize?: number; }) { const navigate = useNavigate(); @@ -32,6 +37,8 @@ export default function MarketplaceItem({ navigate('/agents')} aria-label={label} + aria-posinset={posInSet} + aria-setsize={setSize} data-testid="model-selector-marketplace-item" className={cn( 'flex w-full cursor-pointer items-center justify-between rounded-lg px-2 text-sm', diff --git a/client/src/components/Chat/Menus/Endpoints/components/ModelSpecItem.tsx b/client/src/components/Chat/Menus/Endpoints/components/ModelSpecItem.tsx index 09b29502b7..e85b4b62de 100644 --- a/client/src/components/Chat/Menus/Endpoints/components/ModelSpecItem.tsx +++ b/client/src/components/Chat/Menus/Endpoints/components/ModelSpecItem.tsx @@ -12,9 +12,12 @@ import { cn } from '~/utils'; interface ModelSpecItemProps { spec: TModelSpec; isSelected: boolean; + /** Set when the sibling model list is virtualized; see `VirtualizedModelList`. */ + posInSet?: number; + setSize?: number; } -export function ModelSpecItem({ spec, isSelected }: ModelSpecItemProps) { +export function ModelSpecItem({ spec, isSelected, posInSet, setSize }: ModelSpecItemProps) { const localize = useLocalize(); const { handleSelectSpec, endpointsConfig } = useModelSelectorContext(); const { isFavoriteSpec, toggleFavoriteSpec } = useFavorites(); @@ -34,6 +37,8 @@ export function ModelSpecItem({ spec, isSelected }: ModelSpecItemProps) { ref={itemRef} onClick={() => handleSelectSpec(spec)} aria-selected={isSelected || undefined} + aria-posinset={posInSet} + aria-setsize={setSize} className="group flex w-full cursor-pointer items-center justify-between rounded-lg px-2 text-sm" >
; + isFavorite: (modelId: string) => boolean; + onToggleFavorite: (modelId: string) => void; + endpointIndex?: number; + /** Count of options rendered ahead of this list in the same listbox (marketplace entry, + * model specs), so `aria-posinset` is relative to the whole listbox and not just this list. */ + precedingOptionCount: number; +} + +/** + * Windowed model list for endpoints with very large model sets (agents, mainly). + * + * Only the visible slice is mounted, so the per-row costs — Ariakit composite + * registration, the active-item subscription, and ~12 DOM nodes each — stay + * bounded no matter how many agents the user can see. + * + * Ariakit's composite only knows about mounted rows, so arrow-keying to the edge + * of the window would otherwise find no next item and let focus escape the nested + * menu, closing it. `handleBoundaryNavigation` catches that case: it scrolls the + * next index into the window, waits for the row to mount, then moves the composite + * onto it. Navigation inside the window is left entirely to Ariakit, whose own + * `scrollIntoView` drives the list's scroll position. + */ +export default function VirtualizedModelList({ + endpoint, + modelIds, + globalByName, + isFavorite, + onToggleFavorite, + endpointIndex, + precedingOptionCount, +}: VirtualizedModelListProps) { + const listRef = useRef(null); + const containerRef = useRef(null); + const combobox = Ariakit.useComboboxContext(); + const indexSuffix = endpointIndex != null ? `-${endpointIndex}` : ''; + const rowCount = modelIds.length; + + const rowAt = useCallback( + (index: number) => + containerRef.current?.querySelector( + `[data-row-index="${index}"] [role="option"], [data-row-index="${index}"] [role="menuitem"]`, + ) ?? null, + [], + ); + + useEffect(() => { + if (!combobox) { + return; + } + const deltaFor = (key: string) => { + if (key === 'ArrowDown') { + return 1; + } + return key === 'ArrowUp' ? -1 : 0; + }; + const handleBoundaryNavigation = (event: KeyboardEvent) => { + const delta = deltaFor(event.key); + if (delta === 0) { + return; + } + const activeId = combobox.getState().activeId; + const activeRow = activeId ? document.getElementById(activeId) : null; + const wrapper = activeRow?.closest('[data-row-index]'); + if (!wrapper || !containerRef.current?.contains(wrapper)) { + return; + } + const next = Number(wrapper.dataset.rowIndex) + delta; + /** Let Ariakit own both the in-window case and the ends of the list. */ + if (next < 0 || next >= rowCount || rowAt(next)) { + return; + } + event.preventDefault(); + event.stopPropagation(); + listRef.current?.scrollToRow(next); + requestAnimationFrame(() => { + requestAnimationFrame(() => { + const id = rowAt(next)?.id; + if (id) { + combobox.move(id); + } + }); + }); + }; + document.addEventListener('keydown', handleBoundaryNavigation, true); + return () => document.removeEventListener('keydown', handleBoundaryNavigation, true); + }, [combobox, rowCount, rowAt]); + + const height = useMemo( + () => Math.min(MAX_LIST_HEIGHT, Math.max(ROW_HEIGHT, rowCount * ROW_HEIGHT)), + [rowCount], + ); + + const rowRenderer = useCallback( + ({ index, key, style }: ListRowProps) => { + const modelId = modelIds[index]; + return ( +
+ +
+ ); + }, + [ + endpoint, + globalByName, + isFavorite, + modelIds, + onToggleFavorite, + precedingOptionCount, + rowCount, + ], + ); + + return ( +
+ +
+ ); +} diff --git a/client/src/components/Chat/Menus/Endpoints/components/__tests__/EndpointModelItem.test.tsx b/client/src/components/Chat/Menus/Endpoints/components/__tests__/EndpointModelItem.test.tsx index 786d937006..822bb5c3c8 100644 --- a/client/src/components/Chat/Menus/Endpoints/components/__tests__/EndpointModelItem.test.tsx +++ b/client/src/components/Chat/Menus/Endpoints/components/__tests__/EndpointModelItem.test.tsx @@ -26,15 +26,24 @@ jest.mock('~/components/Chat/Menus/Endpoints/CustomMenu', () => { jest.mock('~/hooks', () => ({ useLocalize: () => (key: string) => key, - useFavorites: () => ({ - isFavoriteModel: () => false, - toggleFavoriteModel: jest.fn(), - isFavoriteAgent: () => false, - toggleFavoriteAgent: jest.fn(), - }), - useIsActiveItem: () => ({ ref: { current: null }, isActive: false }), })); +jest.mock('~/components/Chat/Menus/Endpoints/useActiveItem', () => ({ + __esModule: true, + default: () => ({ ref: { current: null }, isActive: false }), +})); + +const renderItem = (props: Partial> = {}) => + render( + , + ); + const baseEndpoint: Endpoint = { value: 'anthropic', label: 'Anthropic', @@ -50,7 +59,7 @@ describe('EndpointModelItem', () => { it('renders checkmark when model and endpoint match with no active spec', () => { mockSelectedValues = { endpoint: 'anthropic', model: 'claude-opus-4-6', modelSpec: '' }; - render(); + renderItem(); const menuItem = screen.getByRole('menuitem'); expect(menuItem).toHaveAttribute('aria-selected', 'true'); @@ -62,7 +71,7 @@ describe('EndpointModelItem', () => { model: 'claude-opus-4-6', modelSpec: 'my-anthropic-spec', }; - render(); + renderItem(); const menuItem = screen.getByRole('menuitem'); expect(menuItem).not.toHaveAttribute('aria-selected'); @@ -70,7 +79,7 @@ describe('EndpointModelItem', () => { it('does NOT render checkmark when model matches but endpoint differs', () => { mockSelectedValues = { endpoint: 'openai', model: 'claude-opus-4-6', modelSpec: '' }; - render(); + renderItem(); const menuItem = screen.getByRole('menuitem'); expect(menuItem).not.toHaveAttribute('aria-selected'); @@ -78,7 +87,7 @@ describe('EndpointModelItem', () => { it('does NOT render checkmark when endpoint matches but model differs', () => { mockSelectedValues = { endpoint: 'anthropic', model: 'claude-sonnet-4-5', modelSpec: '' }; - render(); + renderItem(); const menuItem = screen.getByRole('menuitem'); expect(menuItem).not.toHaveAttribute('aria-selected'); diff --git a/client/src/components/Chat/Menus/Endpoints/useActiveItem.ts b/client/src/components/Chat/Menus/Endpoints/useActiveItem.ts new file mode 100644 index 0000000000..b4c7b2ba7a --- /dev/null +++ b/client/src/components/Chat/Menus/Endpoints/useActiveItem.ts @@ -0,0 +1,33 @@ +import { useRef } from 'react'; +import * as Ariakit from '@ariakit/react'; +import type { RefObject } from 'react'; + +/** + * Reports whether this row is the composite's active item, read from the Ariakit + * store rather than from a per-row `MutationObserver`. + * + * `useIsActiveItem` allocates one observer per mounted row, which is fine for a + * handful of rows and ruinous for a large model list — and under virtualization + * rows mount and unmount on every scroll frame, so the observers churn as well. + * The store already tracks `activeId`; the selector returns a boolean so a row + * only re-renders when its own active state flips, not on every arrow key. + */ +export default function useActiveItem(): { + ref: RefObject; + isActive: boolean; +} { + const ref = useRef(null); + const combobox = Ariakit.useComboboxContext(); + const menu = Ariakit.useMenuContext(); + /** Endpoint submenus render as a combobox list; plain menus fall back to the menu store. */ + const store = combobox ?? menu; + + const isActive = + Ariakit.useStoreState( + store, + (state) => + state?.activeId != null && ref.current != null && state.activeId === ref.current.id, + ) ?? false; + + return { ref, isActive }; +} diff --git a/client/src/data-provider/Agents/__tests__/mutations.test.ts b/client/src/data-provider/Agents/__tests__/mutations.test.ts index ac4bb45a1d..d16963cc90 100644 --- a/client/src/data-provider/Agents/__tests__/mutations.test.ts +++ b/client/src/data-provider/Agents/__tests__/mutations.test.ts @@ -1,10 +1,14 @@ import { createElement } from 'react'; -import { dataService, QueryKeys } from 'librechat-data-provider'; import { act, renderHook, waitFor } from '@testing-library/react'; import { QueryClient, QueryClientProvider } from '@tanstack/react-query'; -import type { Agent, GraphEdge } from 'librechat-data-provider'; +import { dataService, PermissionBits, QueryKeys } from 'librechat-data-provider'; +import type { Agent, AgentListResponse, GraphEdge } from 'librechat-data-provider'; import type { ReactNode } from 'react'; -import { useDeleteAgentMutation } from '../mutations'; +import { + useDeleteAgentMutation, + useDuplicateAgentMutation, + useUpdateAgentMutation, +} from '../mutations'; jest.mock('librechat-data-provider', () => { const actual = jest.requireActual('librechat-data-provider'); @@ -13,11 +17,13 @@ jest.mock('librechat-data-provider', () => { dataService: { ...actual.dataService, deleteAgent: jest.fn(), + updateAgent: jest.fn(), + duplicateAgent: jest.fn(), }, }; }); -const createAgent = (id: string, edges: GraphEdge[] = []): Agent => ({ +const createAgent = (id: string, edges: GraphEdge[] = [], isEditable?: boolean): Agent => ({ id, name: id, description: null, @@ -35,6 +41,7 @@ const createAgent = (id: string, edges: GraphEdge[] = []): Agent => ({ presence_penalty: null, }, edges, + ...(isEditable !== undefined ? { isEditable } : {}), }); const createWrapper = (queryClient: QueryClient) => @@ -111,3 +118,88 @@ describe('useDeleteAgentMutation', () => { expect(queryClient.getQueryData([QueryKeys.agent, targetId, 'expanded'])).toBeUndefined(); }); }); + +describe('useUpdateAgentMutation', () => { + it('preserves the list-cache isEditable flag after a successful update', async () => { + /** MANAGE_AGENTS can PATCH agents the ACL marks non-editable. Mutation success must + * not promote those VIEW rows into the editable-only "My Agents" subset. */ + const queryClient = new QueryClient({ + defaultOptions: { + queries: { + retry: false, + }, + }, + }); + const agentId = 'agent_view_only'; + const listKey = [QueryKeys.agents, { requiredPermission: PermissionBits.VIEW }]; + const cachedList: AgentListResponse = { + object: 'list', + data: [createAgent(agentId, [], false)], + first_id: agentId, + last_id: agentId, + has_more: false, + }; + queryClient.setQueryData(listKey, cachedList); + + const updatedAgent = createAgent(agentId); + updatedAgent.name = 'Renamed'; + jest.mocked(dataService.updateAgent).mockResolvedValue(updatedAgent); + + const { result } = renderHook(() => useUpdateAgentMutation(), { + wrapper: createWrapper(queryClient), + }); + + await act(async () => { + await result.current.mutateAsync({ agent_id: agentId, data: { name: 'Renamed' } }); + }); + + const listRes = queryClient.getQueryData(listKey); + expect(listRes?.data).toHaveLength(1); + expect(listRes?.data[0]).toMatchObject({ + id: agentId, + name: 'Renamed', + isEditable: false, + }); + }); +}); + +describe('useDuplicateAgentMutation', () => { + it('marks the duplicated agent editable in the list cache', async () => { + /** Duplicating grants ownership, so the new row is editable. Without this the row + * carries no `isEditable` and only survives the selector filter by failing open. */ + const queryClient = new QueryClient({ + defaultOptions: { + queries: { + retry: false, + }, + }, + }); + const sourceId = 'agent_source'; + const duplicateId = 'agent_duplicate'; + const listKey = [QueryKeys.agents, { requiredPermission: PermissionBits.VIEW }]; + queryClient.setQueryData(listKey, { + object: 'list', + data: [createAgent(sourceId, [], false)], + first_id: sourceId, + last_id: sourceId, + has_more: false, + } satisfies AgentListResponse); + + jest + .mocked(dataService.duplicateAgent) + .mockResolvedValue({ agent: createAgent(duplicateId), actions: [] }); + + const { result } = renderHook(() => useDuplicateAgentMutation(), { + wrapper: createWrapper(queryClient), + }); + + await act(async () => { + await result.current.mutateAsync({ agent_id: sourceId }); + }); + + const listRes = queryClient.getQueryData(listKey); + expect(listRes?.data[0]).toMatchObject({ id: duplicateId, isEditable: true }); + /** The untouched source row keeps its own ACL flag. */ + expect(listRes?.data[1]).toMatchObject({ id: sourceId, isEditable: false }); + }); +}); diff --git a/client/src/data-provider/Agents/__tests__/queries.test.ts b/client/src/data-provider/Agents/__tests__/queries.test.ts new file mode 100644 index 0000000000..f4ee1dd02b --- /dev/null +++ b/client/src/data-provider/Agents/__tests__/queries.test.ts @@ -0,0 +1,91 @@ +import { createElement } from 'react'; +import { renderHook, waitFor } from '@testing-library/react'; +import { QueryClient, QueryClientProvider } from '@tanstack/react-query'; +import { dataService, QueryKeys, EModelEndpoint, PermissionBits } from 'librechat-data-provider'; +import type { AgentListResponse } from 'librechat-data-provider'; +import type { ReactNode } from 'react'; +import { defaultAgentParams, useListAgentsQuery } from '../queries'; + +jest.mock('librechat-data-provider', () => { + const actual = jest.requireActual('librechat-data-provider'); + return { + ...actual, + dataService: { + ...actual.dataService, + listAgents: jest.fn(), + }, + }; +}); + +const listAgents = dataService.listAgents as jest.MockedFunction; + +const page = (ids: string[], after: string | null): AgentListResponse => + ({ + object: 'list', + data: ids.map((id) => ({ id, name: id })), + has_more: after != null, + after, + first_id: ids[0] ?? '', + last_id: ids[ids.length - 1] ?? '', + }) as unknown as AgentListResponse; + +const createWrapper = (queryClient: QueryClient) => + function Wrapper({ children }: { children: ReactNode }) { + return createElement(QueryClientProvider, { client: queryClient }, children); + }; + +const renderListAgents = (params: Parameters[0]) => { + const queryClient = new QueryClient({ + defaultOptions: { queries: { retry: false } }, + }); + /** The hook is gated on the agents endpoint being configured. */ + queryClient.setQueryData([QueryKeys.endpoints], { [EModelEndpoint.agents]: {} }); + return renderHook(() => useListAgentsQuery(params), { + wrapper: createWrapper(queryClient), + }); +}; + +describe('useListAgentsQuery', () => { + beforeEach(() => { + jest.clearAllMocks(); + }); + + it('requests the server maximum page size so a typical agent set resolves in one round trip', async () => { + listAgents.mockResolvedValue(page(['a', 'b'], null)); + + const { result } = renderListAgents({ requiredPermission: PermissionBits.VIEW }); + + await waitFor(() => expect(result.current.isSuccess).toBe(true)); + expect(listAgents).toHaveBeenCalledTimes(1); + expect(listAgents).toHaveBeenCalledWith( + expect.objectContaining({ limit: 1000, requiredPermission: PermissionBits.VIEW }), + ); + }); + + it('keeps the walk page size when a caller supplies a smaller limit', async () => { + listAgents.mockResolvedValue(page(['a'], null)); + + const { result } = renderListAgents({ limit: 10, requiredPermission: PermissionBits.VIEW }); + + await waitFor(() => expect(result.current.isSuccess).toBe(true)); + expect(listAgents).toHaveBeenCalledWith(expect.objectContaining({ limit: 1000 })); + }); + + it('does not carry a page size in the default params', async () => { + expect(defaultAgentParams.limit).toBeUndefined(); + }); + + it('still walks every page and flattens the result when the server returns a cursor', async () => { + listAgents + .mockResolvedValueOnce(page(['a', 'b'], 'cursor-1')) + .mockResolvedValueOnce(page(['c'], null)); + + const { result } = renderListAgents({ requiredPermission: PermissionBits.VIEW }); + + await waitFor(() => expect(result.current.isSuccess).toBe(true)); + expect(listAgents).toHaveBeenCalledTimes(2); + expect(listAgents).toHaveBeenLastCalledWith(expect.objectContaining({ cursor: 'cursor-1' })); + expect(result.current.data?.data.map((agent) => agent.id)).toEqual(['a', 'b', 'c']); + expect(result.current.data?.has_more).toBe(false); + }); +}); diff --git a/client/src/data-provider/Agents/mutations.ts b/client/src/data-provider/Agents/mutations.ts index 133d57a59a..9ef291f18c 100644 --- a/client/src/data-provider/Agents/mutations.ts +++ b/client/src/data-provider/Agents/mutations.ts @@ -30,6 +30,18 @@ const hasEdgeWithAgent = (data: unknown, agentId: string): boolean => { ); }; +/** + * Mutation responses omit list-only `isEditable`. When merging into a cached list + * row, keep the ACL flag the list endpoint set rather than inferring it from + * write success (`MANAGE_AGENTS` can PATCH agents the ACL marks non-editable). + */ +const mergeAgentListRow = (previous: t.Agent, next: t.Agent): t.Agent => { + if (previous.isEditable === undefined) { + return next; + } + return { ...next, isEditable: previous.isEditable }; +}; + /** * Create a new agent */ @@ -47,7 +59,12 @@ export const useCreateAgentMutation = ( if (!listRes) { return options?.onSuccess?.(newAgent, variables, context); } - const currentAgents = [newAgent, ...JSON.parse(JSON.stringify(listRes.data))]; + /** The create succeeded, so the caller can edit it. Mutation responses carry no + * `isEditable`; without this the cached row loses the field the list sets. */ + const currentAgents = [ + { ...newAgent, isEditable: true }, + ...JSON.parse(JSON.stringify(listRes.data)), + ]; queryClient.setQueryData([QueryKeys.agents, key], { ...listRes, @@ -94,7 +111,7 @@ export const useUpdateAgentMutation = ( ...listRes, data: listRes.data.map((agent) => { if (agent.id === variables.agent_id) { - return updatedAgent; + return mergeAgentListRow(agent, updatedAgent); } return agent; }), @@ -186,7 +203,9 @@ export const useDuplicateAgentMutation = ( keys.forEach((key) => { const listRes = queryClient.getQueryData([QueryKeys.agents, key]); if (listRes) { - const currentAgents = [agent, ...listRes.data]; + /** Duplicating grants the caller ownership, so the new row is editable. + * The response omits list-only `isEditable`; see `mergeAgentListRow`. */ + const currentAgents = [{ ...agent, isEditable: true }, ...listRes.data]; queryClient.setQueryData([QueryKeys.agents, key], { ...listRes, data: currentAgents, @@ -235,7 +254,7 @@ export const useUploadAgentAvatarMutation = ( ...listRes, data: listRes.data.map((agent) => { if (agent.id === variables.agent_id) { - return updatedAgent; + return mergeAgentListRow(agent, updatedAgent); } return agent; }), @@ -286,7 +305,7 @@ export const useUpdateAgentAction = ( ...listRes, data: listRes.data.map((agent) => { if (agent.id === variables.agent_id) { - return updatedAgent; + return mergeAgentListRow(agent, updatedAgent); } return agent; }), @@ -422,7 +441,7 @@ export const useRevertAgentVersionMutation = ( ...listRes, data: listRes.data.map((agent) => { if (agent.id === variables.agent_id) { - return revertedAgent; + return mergeAgentListRow(agent, revertedAgent); } return agent; }), diff --git a/client/src/data-provider/Agents/queries.ts b/client/src/data-provider/Agents/queries.ts index 4dc03b7bd2..d51328aa25 100644 --- a/client/src/data-provider/Agents/queries.ts +++ b/client/src/data-provider/Agents/queries.ts @@ -12,10 +12,18 @@ import { isEphemeralAgent } from '~/common'; * AGENTS */ export const defaultAgentParams: t.AgentListParams = { - limit: 10, requiredPermission: PermissionBits.EDIT, }; +/** + * Page size for the internal pagination walk. Callers consume the flattened result, so + * every page costs a serial round trip with no benefit: request the server's maximum + * (`getListAgentsByAccess` caps at 1000) so realistic agent sets resolve in one request. + * Kept out of the query key, and applied last so a caller-supplied `limit` cannot shrink + * it: this is a transport detail, and a caller limit never bounds what the walk returns. + */ +const WALK_PAGE_SIZE = 1000; + /** Walk the cursor pagination and return all pages flattened into one `AgentListResponse`. */ async function fetchAllAgentPages(params: t.AgentListParams): Promise { const pages: t.AgentListResponse[] = []; @@ -24,6 +32,7 @@ async function fetchAllAgentPages(params: t.AgentListParams): Promise { + /* An empty anchor can never match a parent chain, so walking the + * conversation only buys an unbounded read whose result is discarded. + * `req.body.parentMessageId` reaches this layer unnormalized. */ + const { agent, req, res, loadTools, db } = setupExecuteCodeAgent(); + + const getMessages = jest.fn().mockResolvedValue([]); + const getConvoFiles = jest.fn().mockResolvedValue([]); + + await initializeAgent( + { + req, + res, + agent, + loadTools, + endpointOption: { endpoint: EModelEndpoint.agents }, + conversationId: 'conv-1', + parentMessageId: '', + allowedProviders: new Set([Providers.OPENAI]), + isInitialAgent: true, + codeEnvAvailable: true, + }, + { ...db, getMessages, getConvoFiles }, + ); + + expect(getMessages).not.toHaveBeenCalled(); + expect(mockGetThreadData).not.toHaveBeenCalled(); + /* The conversation read is unconditional and must survive the guard. */ + expect(getConvoFiles).toHaveBeenCalledTimes(1); + }); + + it('dispatches the convo-file read and the thread walk concurrently', async () => { + /* Both reads gate the model call, so serializing them costs + * time-to-first-token on every turn. Holding BOTH unresolved is what + * makes this fail under either ordering: whichever runs first blocks, + * and the second is never dispatched. + * + * DELETE this test, do not repair it, if the thread walk ever gains a + * data dependency on the convo file ids — serializing becomes correct. */ + const { agent, req, res, loadTools, db } = setupExecuteCodeAgent(); + + let releaseConvoFiles!: (fileIds: string[]) => void; + let releaseMessages!: (messages: Array<{ messageId: string }>) => void; + const getConvoFiles = jest + .fn() + .mockReturnValue(new Promise((resolve) => (releaseConvoFiles = resolve))); + const getMessages = jest + .fn() + .mockReturnValue( + new Promise>((resolve) => (releaseMessages = resolve)), + ); + + const initialized = initializeAgent( + { + req, + res, + agent, + loadTools, + endpointOption: { endpoint: EModelEndpoint.agents }, + conversationId: 'conv-1', + parentMessageId: 'msgN', + allowedProviders: new Set([Providers.OPENAI]), + isInitialAgent: true, + codeEnvAvailable: true, + }, + { ...db, getConvoFiles, getMessages }, + ); + + /* Drain pending microtasks so the mocked chain runs up to the first + * genuinely-pending await. */ + await new Promise((resolve) => setImmediate(resolve)); + + expect(getConvoFiles).toHaveBeenCalledTimes(1); + expect(getMessages).toHaveBeenCalledTimes(1); + + releaseConvoFiles([]); + releaseMessages([]); + await initialized; + }); }); describe('initializeAgent — run-scoped MCP tool definitions', () => { diff --git a/packages/api/src/agents/initialize.ts b/packages/api/src/agents/initialize.ts index b9b5eed8c8..03b21a0e17 100644 --- a/packages/api/src/agents/initialize.ts +++ b/packages/api/src/agents/initialize.ts @@ -707,7 +707,6 @@ export async function initializeAgent( * on handoff agents would fail to find previously attached files. */ if (conversationId != null && resendFiles) { - const fileIds = (await db.getConvoFiles(conversationId)) ?? []; const toolResourceSet = new Set(); for (const tool of agent.tools ?? []) { if (EToolResources[tool as keyof typeof EToolResources]) { @@ -715,74 +714,76 @@ export async function initializeAgent( } } - const toolFiles = requestFileOwnerScope - ? ((await db.getToolFilesByIds( - fileIds, - toolResourceSet, - requestFileOwnerScope, - )) as IMongoFile[]) - : []; + const getThreadMessages = db.getMessages; + /** Falsy anchors cannot match a parent chain, so they get no walk. */ + const threadAnchor = + parentMessageId && parentMessageId !== Constants.NO_PARENT ? parentMessageId : null; + const needsThreadWalk = + toolResourceSet.has(EToolResources.execute_code) && + threadAnchor != null && + getThreadMessages != null; + + /** + * The conversation's file refs and the thread walk share no inputs, so they resolve + * together. Both gate the model call, and this runs on every turn — each serialized + * round trip here is time-to-first-token the user waits through. + * + * Thread walk selects only the fields traversal needs. Both `files` (user uploads) + * and `attachments` (code-execution outputs from `processCodeOutput`) carry the + * `file_id` refs the next turn must prime — selecting only `files` silently drops + * every code-output ref. + */ + const [convoFileIds, threadMessages] = await Promise.all([ + db.getConvoFiles(conversationId), + needsThreadWalk && getThreadMessages + ? getThreadMessages({ conversationId }, 'messageId parentMessageId files attachments') + : null, + ]); + const fileIds = convoFileIds ?? []; + + /** Walk the parent chain and collect file_ids referenced by + * any message in the thread (`messages.files[].file_id` + + * `messages.attachments[].file_id`). Used as the primary + * anchor for both `getCodeGeneratedFiles` and + * `getUserCodeFiles` — message ids no longer needed at + * this layer. */ + const threadFileIds = + threadMessages && threadMessages.length > 0 + ? getThreadData(threadMessages, threadAnchor).fileIds + : undefined; /** * Retrieve execute_code files filtered to the current thread. * This includes both code-generated files and user-uploaded execute_code files. + * + * Code-generated and user-uploaded execute_code files share the same primary anchor: + * file_ids referenced by messages in the current thread. The two queries differ only + * by `context` (`execute_code` for generated outputs, others for uploads). Anchoring + * both on `threadFileIds` reaches files regardless of which sibling first generated + * them — see `getCodeGeneratedFiles` for the branched-conversation rationale. */ - let codeGeneratedFiles: IMongoFile[] = []; - let userCodeFiles: IMongoFile[] = []; - - if (toolResourceSet.has(EToolResources.execute_code)) { - let threadFileIds: string[] | undefined; - - if (parentMessageId && parentMessageId !== Constants.NO_PARENT && db.getMessages) { - /** Only select fields needed for thread traversal. Both - * `files` (user uploads) and `attachments` (code-execution - * outputs from `processCodeOutput`) carry the `file_id` - * refs the next turn must prime — selecting only `files` - * silently drops every code-output ref. */ - const messages = await db.getMessages( - { conversationId }, - 'messageId parentMessageId files attachments', - ); - if (messages && messages.length > 0) { - /** Walk the parent chain and collect file_ids referenced by - * any message in the thread (`messages.files[].file_id` + - * `messages.attachments[].file_id`). Used as the primary - * anchor for both `getCodeGeneratedFiles` and - * `getUserCodeFiles` — message ids no longer needed at - * this layer. */ - threadFileIds = getThreadData(messages, parentMessageId).fileIds; - } - } - - /** Code-generated and user-uploaded execute_code files share the - * same primary anchor: file_ids referenced by messages in the - * current thread. The two queries differ only by `context` - * (`execute_code` for generated outputs, others for uploads). - * Anchoring both on `threadFileIds` reaches files regardless of - * which sibling first generated them — see `getCodeGeneratedFiles` - * for the branched-conversation rationale. */ - if (db.getCodeGeneratedFiles) { - codeGeneratedFiles = requestFileOwnerScope - ? ((await db.getCodeGeneratedFiles( - conversationId, - threadFileIds, - requestFileOwnerScope, - )) as IMongoFile[]) - : []; - } - - if ( - db.getUserCodeFiles && - requestFileOwnerScope && - threadFileIds && - threadFileIds.length > 0 - ) { - userCodeFiles = (await db.getUserCodeFiles( - threadFileIds, - requestFileOwnerScope, - )) as IMongoFile[]; - } - } + const wantsCodeFiles = toolResourceSet.has(EToolResources.execute_code); + const [toolFiles, codeGeneratedFiles, userCodeFiles] = await Promise.all([ + requestFileOwnerScope + ? (db.getToolFilesByIds(fileIds, toolResourceSet, requestFileOwnerScope) as Promise< + IMongoFile[] + >) + : ([] as IMongoFile[]), + wantsCodeFiles && db.getCodeGeneratedFiles && requestFileOwnerScope + ? (db.getCodeGeneratedFiles( + conversationId, + threadFileIds, + requestFileOwnerScope, + ) as Promise) + : ([] as IMongoFile[]), + wantsCodeFiles && + db.getUserCodeFiles && + requestFileOwnerScope && + threadFileIds && + threadFileIds.length > 0 + ? (db.getUserCodeFiles(threadFileIds, requestFileOwnerScope) as Promise) + : ([] as IMongoFile[]), + ]); const allToolFiles = toolFiles.concat(codeGeneratedFiles, userCodeFiles); if (requestFiles.length || allToolFiles.length) { diff --git a/packages/data-provider/src/types/assistants.ts b/packages/data-provider/src/types/assistants.ts index 4180eb8ad3..a634358a8d 100644 --- a/packages/data-provider/src/types/assistants.ts +++ b/packages/data-provider/src/types/assistants.ts @@ -304,6 +304,19 @@ export type Agent = { artifacts?: ArtifactModes; recursion_limit?: number; isPublic?: boolean; + /** + * Whether the requesting user holds EDIT on this agent, so a single VIEW-scoped fetch can + * serve consumers that only need the editable subset instead of issuing a second full + * paginated walk under an EDIT-scoped cache key. + * + * Set by the list endpoint only; single-agent responses omit it. Treat absence as unknown + * and fail open (`isEditable !== false`), never as `false`, since a client on an older + * server would otherwise see an empty list rather than too many rows. + * + * Reflects the caller's ACL grant. The `MANAGE_AGENTS` capability bypasses ACL on write, + * so a capability holder can edit agents this flag reports as not editable. + */ + isEditable?: boolean; version?: number; category?: string; support_contact?: SupportContact;