diff --git a/.changeset/get-tree-react-agent.md b/.changeset/get-tree-react-agent.md new file mode 100644 index 00000000..c347754c --- /dev/null +++ b/.changeset/get-tree-react-agent.md @@ -0,0 +1,6 @@ +--- +'@rozenite/middleware': patch +'rozenite': patch +--- + +Add `getTree` to the React agent so the current tree can be fetched and paged from the live app. diff --git a/packages/agent-sdk/src/constants.ts b/packages/agent-sdk/src/constants.ts index b1be5ba4..63c2b6fe 100644 --- a/packages/agent-sdk/src/constants.ts +++ b/packages/agent-sdk/src/constants.ts @@ -41,6 +41,7 @@ export const STATIC_DOMAIN_TOOL_PREFIXES: Record = {}; export const STATIC_DOMAIN_TOOL_NAMES: Record = { console: ['clearMessages', 'getMessages'], react: [ + 'getTree', 'searchNodes', 'getNode', 'getChildren', diff --git a/packages/cli/skills/rozenite-agent/domains/react.md b/packages/cli/skills/rozenite-agent/domains/react.md index 842d7da4..4407fe42 100644 --- a/packages/cli/skills/rozenite-agent/domains/react.md +++ b/packages/cli/skills/rozenite-agent/domains/react.md @@ -3,6 +3,7 @@ Search and traverse the React component tree, read props, state, and hooks for a ## Tools - `searchNodes` -> `{"query":""}` | `{"query":"","cursor":""}` | `{"query":"","limit":20}` +- `getTree` -> `{}` | `{"depth":2}` | `{"root":123}` | `{"cursor":""}` - `getNode` -> `{"nodeId":123}` - `getChildren` -> `{"nodeId":123}` | `{"nodeId":123,"cursor":""}` | `{"nodeId":123,"limit":20}` - `getProps` -> `{"nodeId":123}` | `{"nodeId":123,"cursor":""}` | `{"nodeId":123,"limit":20}` @@ -16,7 +17,7 @@ Search and traverse the React component tree, read props, state, and hooks for a ## Flow Search and inspect: -`searchNodes` -> `getNode` / `getChildren` -> `getProps` / `getState` / `getHooks`. +`getTree` / `searchNodes` -> `getNode` / `getChildren` -> `getProps` / `getState` / `getHooks`. Profile: `startProfiling` -> reproduce interaction -> `stopProfiling` -> `getRenderData`. diff --git a/packages/middleware/src/agent/local-domains.ts b/packages/middleware/src/agent/local-domains.ts index afc60e59..78c9650c 100644 --- a/packages/middleware/src/agent/local-domains.ts +++ b/packages/middleware/src/agent/local-domains.ts @@ -648,6 +648,33 @@ export const createReactDomainService = (deps: { }) => void; }): LocalAgentToolService => { const tools: AgentTool[] = [ + { + name: 'getTree', + description: + 'Get the current React component tree with optional depth limiting and cursor-based pagination.', + inputSchema: { + type: 'object', + properties: { + root: { + type: 'integer', + description: 'Optional root node ID to scope the tree to a subtree.', + }, + depth: { + type: 'integer', + description: + 'Optional max descendant depth relative to the selected root. 0 returns only roots.', + }, + limit: { + type: 'integer', + description: 'Page size. Default 20, max 100.', + }, + cursor: { + type: 'string', + description: 'Opaque cursor returned by the previous page.', + }, + }, + }, + }, { name: 'getNode', description: 'Get a single React node summary by ID.', @@ -884,6 +911,8 @@ export const createReactDomainService = (deps: { getTools: () => tools, callTool: async (toolName, args) => { switch (toolName) { + case 'getTree': + return store.getTree(sessionDeviceId, args); case 'getNode': return store.getNode(sessionDeviceId, args); case 'getChildren': diff --git a/packages/middleware/src/agent/runtime/react/__tests__/store.test.ts b/packages/middleware/src/agent/runtime/react/__tests__/store.test.ts new file mode 100644 index 00000000..565c86a2 --- /dev/null +++ b/packages/middleware/src/agent/runtime/react/__tests__/store.test.ts @@ -0,0 +1,147 @@ +import { describe, expect, it } from 'vitest'; +import { createReactTreeStore } from '../store.js'; + +const DEVICE_ID = 'device-1'; + +const createStoreWithTree = () => { + const store = createReactTreeStore(); + store.registerDevice(DEVICE_ID); + store.syncTree(DEVICE_ID, { + roots: [1], + nodes: [ + { + nodeId: 1, + displayName: 'Root', + elementType: 'root', + childIds: [2, 3], + }, + { + nodeId: 2, + displayName: 'App', + elementType: 'function', + parentId: 1, + childIds: [4], + }, + { + nodeId: 3, + displayName: 'Sidebar', + elementType: 'function', + parentId: 1, + childIds: [], + }, + { + nodeId: 4, + displayName: 'Button', + elementType: 'host', + parentId: 2, + childIds: [], + }, + ], + }); + + return store; +}; + +describe('React tree store', () => { + it('returns the current React tree in traversal order', () => { + const store = createStoreWithTree(); + + expect(store.getTree(DEVICE_ID, {})).toMatchObject({ + roots: [1], + totalCount: 4, + items: [ + { + nodeId: 1, + displayName: 'Root', + elementType: 'root', + childIds: [2, 3], + childCount: 2, + depth: 0, + }, + { + nodeId: 2, + displayName: 'App', + elementType: 'function', + parentId: 1, + childIds: [4], + childCount: 1, + depth: 1, + }, + { + nodeId: 4, + displayName: 'Button', + elementType: 'host', + parentId: 2, + childIds: [], + childCount: 0, + depth: 2, + }, + { + nodeId: 3, + displayName: 'Sidebar', + elementType: 'function', + parentId: 1, + childIds: [], + childCount: 0, + depth: 1, + }, + ], + page: { + limit: 20, + hasMore: false, + }, + }); + }); + + it('limits tree traversal by depth', () => { + const store = createStoreWithTree(); + + const result = store.getTree(DEVICE_ID, { depth: 1 }); + + expect(result.totalCount).toBe(3); + expect(result.items.map((item) => item.nodeId)).toEqual([1, 2, 3]); + expect(result.items.map((item) => item.depth)).toEqual([0, 1, 1]); + }); + + it('returns a subtree scoped to a root node', () => { + const store = createStoreWithTree(); + + const result = store.getTree(DEVICE_ID, { root: 2 }); + + expect(result.roots).toEqual([2]); + expect(result.totalCount).toBe(2); + expect(result.items.map((item) => item.nodeId)).toEqual([2, 4]); + expect(result.items.map((item) => item.depth)).toEqual([0, 1]); + }); + + it('paginates tree results with stable cursors', () => { + const store = createStoreWithTree(); + + const firstPage = store.getTree(DEVICE_ID, { limit: 2 }); + + expect(firstPage.items.map((item) => item.nodeId)).toEqual([1, 2]); + expect(firstPage.page.hasMore).toBe(true); + expect(firstPage.page.nextCursor).toEqual(expect.any(String)); + + const secondPage = store.getTree(DEVICE_ID, { + limit: 2, + cursor: firstPage.page.nextCursor, + }); + + expect(secondPage.items.map((item) => item.nodeId)).toEqual([4, 3]); + expect(secondPage.page.hasMore).toBe(false); + }); + + it('rejects cursors from a different tree request', () => { + const store = createStoreWithTree(); + const firstPage = store.getTree(DEVICE_ID, { limit: 1 }); + + expect(() => + store.getTree(DEVICE_ID, { + depth: 1, + cursor: firstPage.page.nextCursor, + }), + ).toThrow('Cursor does not match this request context'); + }); +}); + diff --git a/packages/middleware/src/agent/runtime/react/store.ts b/packages/middleware/src/agent/runtime/react/store.ts index 15fee1ec..c376825c 100644 --- a/packages/middleware/src/agent/runtime/react/store.ts +++ b/packages/middleware/src/agent/runtime/react/store.ts @@ -4,6 +4,7 @@ import type { ReactGetChildrenResult, ReactGetInspectableResult, ReactGetRenderDataResult, + ReactGetTreeResult, ReactInspectedNodeRecord, ReactProfilingCursorPayload, ReactProfilingStatusResult, @@ -14,6 +15,7 @@ import type { ReactSearchNodesResult, ReactStartProfilingResult, ReactStopProfilingResult, + ReactTreeNode, ReactTreeNodeInput, ReactTreeSyncPayload, } from './types.js'; @@ -24,6 +26,7 @@ import { const DEFAULT_SEARCH_LIMIT = 20; const MAX_SEARCH_LIMIT = 100; +const GET_TREE_TOOL_NAME = 'getTree'; const SEARCH_TOOL_NAME = 'searchNodes'; const GET_CHILDREN_TOOL_NAME = 'getChildren'; const GET_PROPS_TOOL_NAME = 'getProps'; @@ -131,6 +134,19 @@ const normalizeLimit = (value: unknown): number => { return Math.min(parsed, MAX_SEARCH_LIMIT); }; +const normalizeDepth = (value: unknown): number | undefined => { + if (value === undefined) { + return undefined; + } + + const parsed = Number(value); + if (!Number.isFinite(parsed) || !Number.isInteger(parsed) || parsed < 0) { + throw new Error('"depth" must be a non-negative integer'); + } + + return parsed; +}; + const normalizeMatch = (value: unknown): 'name' | 'name-or-key' => { if (value === 'name-or-key') { return 'name-or-key'; @@ -200,6 +216,20 @@ const ensureNodeSummary = (node: ReactNodeRecord): ReactNodeSummary => { }; }; +const ensureTreeNode = ( + node: ReactNodeRecord, + options: { + depth: number; + childIds: number[]; + }, +): ReactTreeNode => { + return { + ...ensureNodeSummary(node), + childIds: options.childIds, + depth: options.depth, + }; +}; + const getNodeId = (value: unknown): number => { if (!Number.isInteger(value)) { throw new Error('"nodeId" must be an integer'); @@ -222,6 +252,18 @@ const ensureNodeExists = ( return node; }; +const getOptionalRootId = (value: unknown): number | undefined => { + if (value === undefined) { + return undefined; + } + + if (!Number.isInteger(value)) { + throw new Error('"root" must be an integer node ID'); + } + + return Number(value); +}; + const createSerializableSnapshot = ( value: unknown, depth = 3, @@ -942,6 +984,106 @@ export const createReactTreeStore = (options?: { }; }; + const getTree = ( + deviceId: string, + rawRequest: unknown, + ): ReactGetTreeResult => { + const request = getRecord(rawRequest) || {}; + const state = getOrCreateState(deviceId); + const rootId = getOptionalRootId(request.root); + const depth = normalizeDepth(request.depth); + const limit = normalizeLimit(request.limit); + + const traversalRoots = + rootId === undefined + ? state.rootIds.filter((id) => state.nodesById.has(id)) + : [ensureNodeExists(state, rootId).nodeId]; + + const filtersHash = hashFilters({ + rootId, + depth, + }); + + let offset = 0; + if ( + typeof request.cursor === 'string' && + request.cursor.trim().length > 0 + ) { + const decoded = decodeCursor(request.cursor); + if ( + decoded.deviceId !== deviceId || + decoded.tool !== GET_TREE_TOOL_NAME || + decoded.filtersHash !== filtersHash + ) { + throw new Error( + 'Cursor does not match this request context. Restart pagination without cursor.', + ); + } + offset = decoded.offset; + } + + const visited = new Set(); + const allItems: ReactTreeNode[] = []; + + const walk = (nodeId: number, currentDepth: number): void => { + if (visited.has(nodeId)) { + return; + } + if (depth !== undefined && currentDepth > depth) { + return; + } + + const node = state.nodesById.get(nodeId); + if (!node) { + return; + } + + visited.add(nodeId); + const childIds = node.childIds.filter((childId) => + state.nodesById.has(childId), + ); + allItems.push( + ensureTreeNode(node, { + depth: currentDepth, + childIds, + }), + ); + + for (const childId of childIds) { + walk(childId, currentDepth + 1); + } + }; + + for (const traversalRoot of traversalRoots) { + walk(traversalRoot, 0); + } + + const safeOffset = Math.max(0, Math.min(offset, allItems.length)); + const end = Math.min(safeOffset + limit, allItems.length); + const items = allItems.slice(safeOffset, end); + const hasMore = end < allItems.length; + const nextCursor = hasMore + ? encodeCursor({ + v: 1, + tool: GET_TREE_TOOL_NAME, + deviceId, + offset: end, + filtersHash, + }) + : undefined; + + return { + roots: traversalRoots, + items, + totalCount: allItems.length, + page: { + limit, + hasMore, + ...(nextCursor ? { nextCursor } : {}), + }, + }; + }; + const getNode = (deviceId: string, rawRequest: unknown): ReactNodeSummary => { const request = getRecord(rawRequest) || {}; const state = getOrCreateState(deviceId); @@ -1324,6 +1466,7 @@ export const createReactTreeStore = (options?: { isProfilingStarted, stopProfiling, getRenderData, + getTree, searchNodes, getNode, getChildren, diff --git a/packages/middleware/src/agent/runtime/react/types.ts b/packages/middleware/src/agent/runtime/react/types.ts index 109288ed..f7ba3e18 100644 --- a/packages/middleware/src/agent/runtime/react/types.ts +++ b/packages/middleware/src/agent/runtime/react/types.ts @@ -47,6 +47,22 @@ export interface ReactSearchNodesResult { }; } +export interface ReactTreeNode extends ReactNodeSummary { + childIds: number[]; + depth: number; +} + +export interface ReactGetTreeResult { + roots: number[]; + items: ReactTreeNode[]; + totalCount: number; + page: { + limit: number; + hasMore: boolean; + nextCursor?: string; + }; +} + export interface ReactGetChildrenResult { items: ReactNodeSummary[]; page: {