From 7a9a00343044b4e50e8a3c32a3b18f5725e43518 Mon Sep 17 00:00:00 2001 From: Dmitry Lyalin Date: Tue, 10 Feb 2026 20:47:11 -0500 Subject: [PATCH] Harden visualization input normalization and table mapping --- .../messages/VisualizationDisplay.test.tsx | 32 ++++++++ .../messages/VisualizationDisplay.tsx | 40 +++++++++- .../src/tools/render-visualization.test.ts | 53 ++++++++++++ .../core/src/tools/render-visualization.ts | 80 ++++++++++++++++++- 4 files changed, 202 insertions(+), 3 deletions(-) diff --git a/packages/cli/src/ui/components/messages/VisualizationDisplay.test.tsx b/packages/cli/src/ui/components/messages/VisualizationDisplay.test.tsx index 092f543f32..1fa8582da9 100644 --- a/packages/cli/src/ui/components/messages/VisualizationDisplay.test.tsx +++ b/packages/cli/src/ui/components/messages/VisualizationDisplay.test.tsx @@ -128,6 +128,38 @@ describe('VisualizationResultDisplay', () => { expect(frame).toContain('Showing truncated data (8 original items)'); }); + it('renders table object rows with humanized columns without blank cells', () => { + const visualization: VisualizationDisplay = { + type: 'visualization', + kind: 'table', + title: 'Server Health Check', + data: { + columns: ['Server Name', 'CPU %', 'Memory GB', 'Status'], + rows: [ + { + server_name: 'api-1', + cpu_percent: 62, + memoryGb: 8, + status: 'healthy', + }, + ], + }, + meta: { + truncated: false, + originalItemCount: 1, + }, + }; + + const { lastFrame } = render( + , + ); + const frame = lastFrame(); + + expect(frame).toContain('Server Name'); + expect(frame).toContain('api-1'); + expect(frame).toContain('healthy'); + }); + it('renders diagram visualization with UML-like nodes', () => { const visualization: VisualizationDisplay = { type: 'visualization', diff --git a/packages/cli/src/ui/components/messages/VisualizationDisplay.tsx b/packages/cli/src/ui/components/messages/VisualizationDisplay.tsx index 53656c7328..14ecf8d252 100644 --- a/packages/cli/src/ui/components/messages/VisualizationDisplay.tsx +++ b/packages/cli/src/ui/components/messages/VisualizationDisplay.tsx @@ -137,6 +137,42 @@ function normalizeCell(value: unknown): PrimitiveCell { return JSON.stringify(value); } +function normalizeLookupKey(value: string): string { + return value.toLowerCase().replace(/[^a-z0-9]/g, ''); +} + +function getRecordValueByColumn( + record: Record, + column: string, +): unknown { + if (Object.prototype.hasOwnProperty.call(record, column)) { + return record[column]; + } + + const normalizedColumn = normalizeLookupKey(column); + if (normalizedColumn.length === 0) { + return undefined; + } + + const entries = Object.entries(record).map(([key, value]) => ({ + key, + value, + normalized: normalizeLookupKey(key), + })); + + const exact = entries.find((entry) => entry.normalized === normalizedColumn); + if (exact) { + return exact.value; + } + + const partial = entries.find( + (entry) => + entry.normalized.includes(normalizedColumn) || + normalizedColumn.includes(entry.normalized), + ); + return partial?.value; +} + function makeBar(value: number, max: number, width: number): string { const safeMax = max <= 0 ? 1 : max; const ratio = Math.max(0, Math.min(1, value / safeMax)); @@ -234,7 +270,9 @@ function asTable(data: Record): { const record = row; const keys = columns.length > 0 ? columns : Object.keys(record); - return keys.map((key) => normalizeCell(record[key])); + return keys.map((key) => + normalizeCell(getRecordValueByColumn(record, key)), + ); }); const inferredColumns = diff --git a/packages/core/src/tools/render-visualization.test.ts b/packages/core/src/tools/render-visualization.test.ts index d494e41efe..7bb2005b8e 100644 --- a/packages/core/src/tools/render-visualization.test.ts +++ b/packages/core/src/tools/render-visualization.test.ts @@ -206,6 +206,34 @@ describe('RenderVisualizationTool', () => { }); }); + it('maps table object rows when column labels are humanized', async () => { + const result = await tool.buildAndExecute( + { + visualizationKind: 'table', + data: { + columns: ['Server Name', 'CPU %', 'Memory GB', 'Status'], + rows: [ + { + server_name: 'api-1', + cpu_percent: 62, + memoryGb: 8, + status: 'healthy', + }, + ], + }, + }, + signal, + ); + + expect(result.returnDisplay).toMatchObject({ + type: 'visualization', + kind: 'table', + data: { + rows: [['api-1', 62, 8, 'healthy']], + }, + }); + }); + it('renders diagram visualization', async () => { const params: RenderVisualizationToolParams = { visualizationKind: 'diagram', @@ -257,6 +285,31 @@ describe('RenderVisualizationTool', () => { }); }); + it('accepts top-level diagram direction alias', async () => { + const result = await tool.buildAndExecute( + { + visualizationKind: 'diagram', + direction: 'TB', + data: { + nodes: [ + { id: 'start', label: 'Start' }, + { id: 'end', label: 'End' }, + ], + edges: [{ from: 'start', to: 'end' }], + }, + }, + signal, + ); + + expect(result.returnDisplay).toMatchObject({ + type: 'visualization', + kind: 'diagram', + data: { + direction: 'TB', + }, + }); + }); + it('converts legacy dashboard payloads into table', async () => { const params: RenderVisualizationToolParams = { visualizationKind: 'table', diff --git a/packages/core/src/tools/render-visualization.ts b/packages/core/src/tools/render-visualization.ts index badd8620cd..468d48db82 100644 --- a/packages/core/src/tools/render-visualization.ts +++ b/packages/core/src/tools/render-visualization.ts @@ -38,6 +38,8 @@ export interface RenderVisualizationToolParams { xLabel?: string; yLabel?: string; unit?: string; + direction?: 'LR' | 'TB' | string; + diagramKind?: 'architecture' | 'flowchart' | string; data: unknown; sourceContext?: string; sort?: SortMode; @@ -129,6 +131,44 @@ function normalizeCell(value: unknown): PrimitiveCell { return JSON.stringify(value); } +function normalizeLookupKey(value: string): string { + return value.toLowerCase().replace(/[^a-z0-9]/g, ''); +} + +function getRecordValueByColumn( + record: Record, + column: string, +): unknown { + if (Object.prototype.hasOwnProperty.call(record, column)) { + return record[column]; + } + + const normalizedColumn = normalizeLookupKey(column); + if (normalizedColumn.length === 0) { + return undefined; + } + + const normalizedEntries = Object.entries(record).map(([key, value]) => ({ + key, + value, + normalized: normalizeLookupKey(key), + })); + + const exactNormalized = normalizedEntries.find( + (entry) => entry.normalized === normalizedColumn, + ); + if (exactNormalized) { + return exactNormalized.value; + } + + const includesNormalized = normalizedEntries.find( + (entry) => + entry.normalized.includes(normalizedColumn) || + normalizedColumn.includes(entry.normalized), + ); + return includesNormalized?.value; +} + function applySort( points: VisualizationPoint[], sort: SortMode, @@ -805,7 +845,9 @@ function normalizeTable( if (columns.length === 0) { columns = Object.keys(rawRow); } - return columns.map((column) => normalizeCell(rawRow[column])); + return columns.map((column) => + normalizeCell(getRecordValueByColumn(rawRow, column)), + ); } throw new Error(`table row ${rowIndex} must be an array or object.`); @@ -1090,9 +1132,33 @@ class RenderVisualizationToolInvocation extends BaseToolInvocation< Math.min(this.params.maxItems ?? DEFAULT_MAX_ITEMS, MAX_ALLOWED_ITEMS), ); + const toolData = (() => { + const data = ensureRecord(this.params.data, '`data` must be an object.'); + if (this.params.visualizationKind !== 'diagram') { + return data; + } + + const diagramData: Record = { ...data }; + if ( + diagramData['direction'] === undefined && + typeof this.params.direction === 'string' && + this.params.direction.trim().length > 0 + ) { + diagramData['direction'] = this.params.direction; + } + if ( + diagramData['diagramKind'] === undefined && + typeof this.params.diagramKind === 'string' && + this.params.diagramKind.trim().length > 0 + ) { + diagramData['diagramKind'] = this.params.diagramKind; + } + return diagramData; + })(); + const normalized = normalizeByKind( this.params.visualizationKind, - this.params.data, + toolData, sort, maxItems, ); @@ -1199,6 +1265,16 @@ export class RenderVisualizationTool extends BaseDeclarativeTool< type: 'string', description: 'Optional unit label for values.', }, + direction: { + type: 'string', + description: + 'Optional top-level direction alias for diagram calls; copied into data.direction when provided.', + }, + diagramKind: { + type: 'string', + description: + 'Optional top-level diagram kind alias for diagram calls; copied into data.diagramKind when provided.', + }, data: { type: 'object', description: