fix metric sneaking in

This commit is contained in:
dal 2025-09-24 13:58:23 -06:00
parent ee7acd68f4
commit d29f077bdf
No known key found for this signature in database
GPG Key ID: 16F4B0E1E9F61122
3 changed files with 156 additions and 28 deletions

View File

@ -1,5 +1,6 @@
import { db } from '@buster/database/connection';
import { chats, messages } from '@buster/database/schema';
import type { AssetType } from '@buster/database/schema-types';
import { createTestChat, createTestMessage } from '@buster/test-utils';
import { eq } from 'drizzle-orm';
import { afterEach, beforeEach, describe, expect, it, vi } from 'vitest';
@ -50,7 +51,7 @@ describe('markMessageComplete integration test', () => {
it('should update chat with selected file information', async () => {
const selectedFile = {
fileId: '123e4567-e89b-12d3-a456-426614174000',
fileType: 'dashboard_file',
fileType: 'dashboard_file' as AssetType,
versionNumber: 1,
};
@ -109,7 +110,7 @@ describe('markMessageComplete integration test', () => {
it('should not update chat when chatId is missing', async () => {
const selectedFile = {
fileId: '123e4567-e89b-12d3-a456-426614174000',
fileType: 'metric_file',
fileType: 'metric_file' as AssetType,
versionNumber: 2,
};

View File

@ -508,6 +508,81 @@ describe('done-tool-file-selection - report filtering functionality', () => {
});
describe('BUG-1885: Metrics used in reports appearing in response messages', () => {
it('should filter out metric when content has escaped quotes (actual production bug)', () => {
// This test reproduces the EXACT bug scenario from production where escaped quotes
// in the report content were causing the metric ID regex to fail
const messages: ModelMessage[] = [
// Create a metric first
{
role: 'tool',
content: [
{
type: 'tool-result',
toolName: 'createMetrics',
toolCallId: 'toolu_01R7jSBNXiNd1mdN162Kk5To',
output: {
type: 'json',
value: {
files: [
{
id: '229f7b5d-c660-42a9-b4f2-46a0bf1f8726',
name: 'Total Customers',
version_number: 1,
},
],
},
},
},
],
},
// Create report with ESCAPED quotes in the content (as seen in production)
{
role: 'assistant',
content: [
{
type: 'tool-call',
toolName: 'createReports',
toolCallId: 'toolu_019dwAk5Ls2XHMpCfZVJg9z3',
input: {
name: 'Total Customers',
// This is the exact format from production with escaped quotes
content:
'<metric metricId=\\"229f7b5d-c660-42a9-b4f2-46a0bf1f8726\\"/>\\n\\nAdventure Works has **19,820 total customers** in their database.',
},
},
],
},
{
role: 'tool',
content: [
{
type: 'tool-result',
toolName: 'createReports',
toolCallId: 'toolu_019dwAk5Ls2XHMpCfZVJg9z3',
output: {
type: 'json',
value: {
file: {
id: 'a41ae0e9-2215-4ba8-8045-7b5e68e6f4b8',
name: 'Total Customers',
version_number: 1,
},
},
},
},
],
},
];
const files = extractFilesFromToolCalls(messages);
// The metric should be filtered out because it's referenced in the report
// The report should also be filtered out (all reports are filtered)
expect(files).toHaveLength(0);
expect(files.find((f) => f.id === '229f7b5d-c660-42a9-b4f2-46a0bf1f8726')).toBeUndefined();
expect(files.find((f) => f.id === 'a41ae0e9-2215-4ba8-8045-7b5e68e6f4b8')).toBeUndefined();
});
it('should filter out metric when using new single-file report structure (user reported bug)', () => {
// This test reproduces the exact bug scenario reported by the user
const messages: ModelMessage[] = [

View File

@ -18,10 +18,12 @@ import {
} from '../../../visualization-tools/metrics/modify-metrics-tool/modify-metrics-tool';
import {
CREATE_REPORTS_TOOL_NAME,
type CreateReportsInput,
type CreateReportsOutput,
} from '../../../visualization-tools/reports/create-reports-tool/create-reports-tool';
import {
MODIFY_REPORTS_TOOL_NAME,
type ModifyReportsInput,
type ModifyReportsOutput,
} from '../../../visualization-tools/reports/modify-reports-tool/modify-reports-tool';
@ -184,6 +186,56 @@ export function extractAllFilesForChatUpdate(messages: ModelMessage[]): Extracte
function extractReferencedMetricIds(messages: ModelMessage[]): Set<string> {
const referencedMetricIds = new Set<string>();
// First, collect all metric IDs that exist from tool results
const allMetricIds = new Set<string>();
for (const message of messages) {
if (message.role === 'tool') {
const toolContent = message.content;
if (Array.isArray(toolContent)) {
for (const content of toolContent) {
if (
content &&
typeof content === 'object' &&
'type' in content &&
content.type === 'tool-result'
) {
const toolName = (content as unknown as Record<string, unknown>).toolName;
const output = (content as unknown as Record<string, unknown>).output;
if (
(toolName === CREATE_METRICS_TOOL_NAME || toolName === MODIFY_METRICS_TOOL_NAME) &&
output
) {
const outputObj = output as Record<string, unknown>;
if (outputObj.type === 'json' && outputObj.value) {
try {
const parsedOutput =
typeof outputObj.value === 'string'
? JSON.parse(outputObj.value)
: outputObj.value;
const metricsOutput = parsedOutput as CreateMetricsOutput | ModifyMetricsOutput;
if (metricsOutput.files && Array.isArray(metricsOutput.files)) {
for (const file of metricsOutput.files) {
if (file.id) {
allMetricIds.add(file.id);
}
}
}
} catch (_error) {
// Ignore parsing errors
}
}
}
}
}
}
}
}
// Then check if any metric IDs appear in report content
for (const message of messages) {
if (message.role === 'assistant' && Array.isArray(message.content)) {
for (const content of message.content) {
@ -197,21 +249,19 @@ function extractReferencedMetricIds(messages: ModelMessage[]): Set<string> {
const toolName = content.toolName;
const input = (content as { input?: unknown }).input;
// Extract from CREATE_REPORTS
// Extract from CREATE_REPORTS with type safety
if (toolName === CREATE_REPORTS_TOOL_NAME && input) {
const reportInput = input as {
content?: string;
// Type the input as CreateReportsInput with fallback for legacy
const reportInput = input as Partial<CreateReportsInput> & {
files?: Array<{ yml_content?: string; content?: string }>;
};
// Handle new structure
if (reportInput.content) {
// Use a more flexible pattern that matches both UUID and simple IDs
const metricIdPattern = /<metric\s+metricId\s*=\s*["']([a-zA-Z0-9-]+)["']\s*\/>/gi;
const matches = reportInput.content.matchAll(metricIdPattern);
for (const match of matches) {
if (match[1]) {
referencedMetricIds.add(match[1]);
// Simply check if any metric ID appears in the content
for (const metricId of allMetricIds) {
if (reportInput.content.includes(metricId)) {
referencedMetricIds.add(metricId);
}
}
}
@ -221,12 +271,10 @@ function extractReferencedMetricIds(messages: ModelMessage[]): Set<string> {
for (const file of reportInput.files) {
const content = file.yml_content || file.content;
if (content) {
const metricIdPattern =
/<metric\s+metricId\s*=\s*["']([a-zA-Z0-9-]+)["']\s*\/>/gi;
const matches = content.matchAll(metricIdPattern);
for (const match of matches) {
if (match[1]) {
referencedMetricIds.add(match[1]);
// Simply check if any metric ID appears in the content
for (const metricId of allMetricIds) {
if (content.includes(metricId)) {
referencedMetricIds.add(metricId);
}
}
}
@ -234,24 +282,28 @@ function extractReferencedMetricIds(messages: ModelMessage[]): Set<string> {
}
}
// Extract from MODIFY_REPORTS
// Extract from MODIFY_REPORTS with type safety
if (toolName === MODIFY_REPORTS_TOOL_NAME && input) {
const modifyInput = input as {
edits?: Array<{ code?: string; new_content?: string; content?: string }>;
// Type the input as ModifyReportsInput with support for variant field names
const modifyInput = input as Partial<ModifyReportsInput> & {
edits?: Array<{
operation?: 'replace' | 'append';
code?: string;
code_to_replace?: string;
new_content?: string;
content?: string;
}>;
};
if (modifyInput.edits && Array.isArray(modifyInput.edits)) {
for (const edit of modifyInput.edits) {
// The field is 'code' based on the schema
// Check all possible field names for content
const content = edit.code || edit.new_content || edit.content;
if (content) {
// Use a more flexible pattern that matches both UUID and simple IDs
const metricIdPattern =
/<metric\s+metricId\s*=\s*["']([a-zA-Z0-9-]+)["']\s*\/>/gi;
const matches = content.matchAll(metricIdPattern);
for (const match of matches) {
if (match[1]) {
referencedMetricIds.add(match[1]);
// Simply check if any metric ID appears in the content
for (const metricId of allMetricIds) {
if (content.includes(metricId)) {
referencedMetricIds.add(metricId);
}
}
}