[fix: chart/table data] resolve data unavailable when chart/table changes on live stories

apps/backend/src/mcp/embed/sandbox-html.ts

@@ -1,6 +1,6 @@
 import type { QueryDataMap } from '../../utils/story-download';
 import { generateStoryHtml } from '../../utils/story-html';
-import { resolveStoryQueryDataForSandbox } from '../../utils/story-query-data';
+import { backfillMissingQueryDataForSandbox } from '../../utils/story-query-data';
 import { storyEmbedUrls } from '../urls';
 import { MAX_SANDBOX_HTML_CHARS } from './embed-payload';
 import { SANDBOX_EMBED_ROOT_STYLES, SANDBOX_ICON_DOWNLOAD, SANDBOX_ICON_EXTERNAL_LINK } from './header';
@@ -28,7 +28,7 @@ export async function buildStorySandboxHtml(params: {
 	chatId?: string | null;
 	userId?: string;
 }): Promise<string | null> {
-	const queryData = await resolveStoryQueryDataForSandbox(params.code, {
+	const queryData = await backfillMissingQueryDataForSandbox(params.code, {
 		storyId: params.storyId,
 		chatId: params.chatId,
 		projectId: params.projectId,

apps/backend/src/mcp/tools/context-layer.ts

@@ -12,7 +12,7 @@ import { upsertMcpQueryData } from '../../queries/mcp-query-data.queries';
 import * as storyQueries from '../../queries/story.queries';
 import * as storyFolderQueries from '../../queries/story-folder.queries';
 import { pinQueryDataToChat, pinStoryMessageToChat } from '../../utils/chat-message-story';
-import { resolveStoryQueryData, type StoryQueryDataMap } from '../../utils/story-query-data';
+import { backfillMissingQueryData, type StoryQueryDataMap } from '../../utils/story-query-data';
 import { STORY_OUTPUT_SCHEMA, type StoryMcpToolPayload } from '../embed/embed-tool-result';
 import { STORY_APP_URI, uiToolMeta } from '../embed/ui-resources';
 import type { McpContext } from '../logging';
@@ -273,11 +273,10 @@ async function cacheStoryQueryData(
 		...((existingCache?.queryData as StoryQueryDataMap | null) ?? {}),
 		...(queryData ?? {}),
 	};
-	const resolvedQueryData = await resolveStoryQueryData(
+	const resolvedQueryData = await backfillMissingQueryData(
 		code,
 		Object.keys(seededQueryData).length > 0 ? seededQueryData : null,
-		ctx.projectId,
-		ctx.userId,
+		{ projectId: ctx.projectId, userId: ctx.userId },
 	);
 	if (!resolvedQueryData) {
 		return;

apps/backend/src/queries/shared-story.queries.ts

@@ -1,3 +1,4 @@
+import { extractQueryIds } from '@nao/shared/story-segments';
 import { and, count, desc, eq, isNull, max, or, type SQL, sql } from 'drizzle-orm';
 
 import s, { type DBSharedStory } from '../db/abstractSchema';
@@ -130,13 +131,7 @@ export async function getQueryDataFromCode(
 	chatId: string,
 	code: string,
 ): Promise<Record<string, { data: unknown[]; columns: string[] }> | null> {
-	const chartRegex = /<(?:chart|table)\s+[^>]*query_id="([^"]*)"[^>]*\/?>/g;
-	const queryIds = new Set<string>();
-	let match;
-	while ((match = chartRegex.exec(code)) !== null) {
-		queryIds.add(match[1]);
-	}
-
+	const queryIds = extractQueryIds(code);
 	if (queryIds.size === 0) {
 		return null;
 	}

apps/backend/src/queries/story.queries.ts

@@ -1,3 +1,4 @@
+import { extractQueryIds } from '@nao/shared/story-segments';
 import { type StorySharingInfo } from '@nao/shared/types';
 import { and, asc, desc, eq, inArray, isNull, max, or, type SQL, sql } from 'drizzle-orm';
 
@@ -557,13 +558,7 @@ export async function getSqlQueriesFromCode(
 	chatId: string,
 	code: string,
 ): Promise<Record<string, { sqlQuery: string; databaseId?: string }>> {
-	const chartRegex = /<(?:chart|table)\s+[^>]*query_id="([^"]*)"[^>]*\/?>/g;
-	const queryIds = new Set<string>();
-	let match;
-	while ((match = chartRegex.exec(code)) !== null) {
-		queryIds.add(match[1]);
-	}
-
+	const queryIds = extractQueryIds(code);
 	if (queryIds.size === 0) {
 		return {};
 	}

apps/backend/src/services/automation-tools.ts

@@ -1,4 +1,5 @@
 import type { DateFormatSettings } from '@nao/shared/date';
+import { extractQueryIds } from '@nao/shared/story-segments';
 import type { displayChart } from '@nao/shared/tools';
 import { z } from 'zod/v4';
 
@@ -316,8 +317,8 @@ async function buildStoryPdfAttachments(
 }
 
 async function getStoryQueryData(context: ToolContext, code: string): Promise<QueryDataMap | null> {
-	const queryIds = extractStoryQueryIds(code);
-	if (queryIds.length === 0) {
+	const queryIds = extractQueryIds(code);
+	if (queryIds.size === 0) {
 		return null;
 	}
 
@@ -332,16 +333,6 @@ async function getStoryQueryData(context: ToolContext, code: string): Promise<Qu
 	return Object.keys(queryData).length > 0 ? queryData : null;
 }
 
-function extractStoryQueryIds(code: string): string[] {
-	const queryIds = new Set<string>();
-	const chartRegex = /<(?:chart|table)\s+[^>]*query_id="([^"]*)"[^>]*\/?>/g;
-	let match: RegExpExecArray | null;
-	while ((match = chartRegex.exec(code)) !== null) {
-		queryIds.add(match[1]);
-	}
-	return [...queryIds];
-}
-
 function appendInlineChartImages(html: string, attachments: GeneratedArtifactAttachment[]): string {
 	const charts = attachments.filter((attachment) => attachment.kind === 'chart' && attachment.cid);
 	if (charts.length === 0) {

apps/backend/src/services/live-story.ts

@@ -5,6 +5,7 @@ import { z } from 'zod';
 
 import { llmTelemetry } from '../agents/telemetry';
 import { LiveStoryRefreshPrompt } from '../components/ai/live-story-refresh-prompt';
+import type { DBStoryDataCache } from '../db/abstractSchema';
 import { env } from '../env';
 import { renderToMarkdown } from '../lib/markdown';
 import * as chatQueries from '../queries/chat.queries';
@@ -13,6 +14,7 @@ import * as llmConfigQueries from '../queries/project-llm-config.queries';
 import { getQueryDataFromCode } from '../queries/shared-story.queries';
 import * as storyQueries from '../queries/story.queries';
 import { getDefaultModelId, resolveProviderModel } from '../utils/llm';
+import { backfillMissingQueryData, findMissingQueryIds } from '../utils/story-query-data';
 import { MAX_OUTPUT_TOKENS } from './agent';
 const MAX_RENDERED_ROWS = 60;
 
@@ -105,7 +107,7 @@ export async function getStoryQueryData(
 	const cache = await storyQueries.getStoryDataCacheByChatAndSlug(chatId, slug);
 
 	if (cache && !isCacheExpired(cache.cachedAt, cacheSchedule)) {
-		return { queryData: cache.queryData, cachedAt: cache.cachedAt };
+		return resolveFromCache(chatId, code, cache);
 	}
 
 	try {
@@ -116,12 +118,19 @@ export async function getStoryQueryData(
 		};
 	} catch {
 		if (cache) {
-			return { queryData: cache.queryData, cachedAt: cache.cachedAt };
+			return resolveFromCache(chatId, code, cache);
 		}
 		return { queryData: await getQueryDataFromCode(chatId, code), cachedAt: null };
 	}
 }
 
+async function resolveFromCache(chatId: string, code: string, cache: DBStoryDataCache): Promise<StoryQueryDataResult> {
+	const missing = findMissingQueryIds(code, cache.queryData);
+	const queryData =
+		missing.length > 0 ? await backfillMissingQueryData(code, cache.queryData, { chatId }) : cache.queryData;
+	return { queryData, cachedAt: cache.cachedAt };
+}
+
 async function executeRawSql(
 	sqlQuery: string,
 	projectFolder: string,

apps/backend/src/trpc/story.routes.ts

@@ -17,6 +17,7 @@ import { executeLiveQuery, getStoryQueryData, refreshStoryData } from '../servic
 import { nextCronTick } from '../services/scheduler.service';
 import { logAnalyticsEvent } from '../utils/analytics-event';
 import { buildDownloadResponse } from '../utils/story-download';
+import { backfillMissingQueryData } from '../utils/story-query-data';
 import { extractStorySummary } from '../utils/story-summary';
 import { canSendProcedure, ownedResourceProcedure, projectProtectedProcedure, protectedProcedure } from './trpc';
 
@@ -121,7 +122,11 @@ export const storyRoutes = {
 			});
 		}
 
-		return { ...story, queryData: cache?.queryData ?? null };
+		const queryData = story.chatId
+			? await backfillMissingQueryData(story.code, cache?.queryData ?? null, { chatId: story.chatId })
+			: (cache?.queryData ?? null);
+
+		return { ...story, queryData };
 	}),
 
 	getLatest: chatOwnerProcedure

apps/backend/src/utils/embed-story.ts

@@ -4,7 +4,7 @@ import * as projectQueries from '../queries/project.queries';
 import * as storyQueries from '../queries/story.queries';
 import { assertProjectMcpEnabled, verifyEmbedToken } from './embed-token';
 import { HandlerError } from './error';
-import { resolveStoryQueryDataForSandbox, type StoryQueryDataMap } from './story-query-data';
+import { backfillMissingQueryDataForSandbox, type StoryQueryDataMap } from './story-query-data';
 
 export type EmbedStoryContent = {
 	storyId: string;
@@ -46,7 +46,7 @@ export async function loadEmbedStoryContent(storyId: string, token: string): Pro
 	}
 
 	const [queryData, displaySettings] = await Promise.all([
-		resolveStoryQueryDataForSandbox(version.code, {
+		backfillMissingQueryDataForSandbox(version.code, {
 			storyId,
 			chatId: version.chatId,
 			projectId,

apps/backend/src/utils/story-query-data.ts

@@ -1,10 +1,14 @@
+import { extractQueryIds } from '@nao/shared/story-segments';
+
 import { getMcpQueryData } from '../queries/mcp-query-data.queries';
 import { getQueryDataFromCode } from '../queries/shared-story.queries';
 import * as storyQueries from '../queries/story.queries';
 
 export type StoryQueryDataMap = Record<string, { data: unknown[]; columns: string[] }>;
 
-export async function resolveStoryQueryDataForSandbox(
+export type StoryQueryDataSource = { chatId: string } | { projectId: string; userId?: string };
+
+export async function backfillMissingQueryDataForSandbox(
 	code: string,
 	opts: { storyId?: string; chatId?: string | null; projectId: string; userId?: string },
 ): Promise<StoryQueryDataMap | null> {
@@ -23,44 +27,54 @@ export async function resolveStoryQueryDataForSandbox(
 		}
 	}
 	const seeded = Object.keys(seed).length > 0 ? seed : null;
-	return resolveStoryQueryData(code, seeded, opts.projectId, opts.userId);
+	return backfillMissingQueryData(code, seeded, { projectId: opts.projectId, userId: opts.userId });
 }
 
-export async function resolveStoryQueryData(
+export function findMissingQueryIds(code: string, cachedQueryData: StoryQueryDataMap | null): string[] {
+	return [...extractQueryIds(code)].filter((id) => !cachedQueryData?.[id]);
+}
+
+export async function backfillMissingQueryData(
 	code: string,
 	cachedQueryData: StoryQueryDataMap | null,
-	projectId: string,
-	userId?: string,
+	source: StoryQueryDataSource,
 ): Promise<StoryQueryDataMap | null> {
-	const referencedIds = extractQueryIdsFromStoryCode(code);
-	if (referencedIds.size === 0) {
+	const missing = findMissingQueryIds(code, cachedQueryData);
+	if (missing.length === 0) {
 		return cachedQueryData;
 	}
 
-	const merged: StoryQueryDataMap = { ...(cachedQueryData ?? {}) };
-	const missing = [...referencedIds].filter((id) => !merged[id]);
-	if (missing.length === 0) {
-		return merged;
+	const filled =
+		'chatId' in source
+			? await fetchFromChat(source.chatId, code, missing)
+			: await fetchFromMcp(missing, source.projectId, source.userId);
+
+	const merged = { ...(cachedQueryData ?? {}), ...filled };
+	return Object.keys(merged).length > 0 ? merged : null;
+}
+
+async function fetchFromChat(chatId: string, code: string, ids: string[]): Promise<StoryQueryDataMap> {
+	const fromChat = await getQueryDataFromCode(chatId, code).catch(() => null);
+
+	const filled: StoryQueryDataMap = {};
+	for (const id of ids) {
+		if (fromChat?.[id]) {
+			filled[id] = fromChat[id];
+		}
 	}
+	return filled;
+}
 
+async function fetchFromMcp(ids: string[], projectId: string, userId?: string): Promise<StoryQueryDataMap> {
 	const fetchOptions = userId ? { userId } : undefined;
-	const fetched = await Promise.all(missing.map((id) => getMcpQueryData(id, projectId, fetchOptions)));
-	missing.forEach((id, idx) => {
-		const row = fetched[idx];
+	const rows = await Promise.all(ids.map((id) => getMcpQueryData(id, projectId, fetchOptions)));
+
+	const filled: StoryQueryDataMap = {};
+	ids.forEach((id, idx) => {
+		const row = rows[idx];
 		if (row) {
-			merged[id] = { columns: row.columns, data: row.data };
+			filled[id] = { columns: row.columns, data: row.data };
 		}
 	});
-
-	return Object.keys(merged).length > 0 ? merged : null;
-}
-
-function extractQueryIdsFromStoryCode(code: string): Set<string> {
-	const ids = new Set<string>();
-	const regex = /<(?:chart|table)\s+[^>]*?\bquery_id\s*=\s*"([^"]+)"/g;
-	let match: RegExpExecArray | null;
-	while ((match = regex.exec(code)) !== null) {
-		ids.add(match[1]);
-	}
-	return ids;
+	return filled;
 }

apps/frontend/src/lib/story-share.utils.ts

@@ -1,3 +1,5 @@
+import { extractQueryIds } from '@nao/shared/story-segments';
+
 import type { UIMessage } from '@nao/backend/chat';
 
 /**
@@ -9,13 +11,7 @@ export function getQueryDataFromCodeFromMessages(
 	messages: UIMessage[],
 	code: string,
 ): Record<string, unknown[]> | null {
-	const chartRegex = /<(?:chart|table)\s+[^>]*query_id="([^"]*)"[^>]*\/?>/g;
-	const queryIds = new Set<string>();
-	let match;
-	while ((match = chartRegex.exec(code)) !== null) {
-		queryIds.add(match[1]);
-	}
-
+	const queryIds = extractQueryIds(code);
 	if (queryIds.size === 0) {
 		return null;
 	}

apps/shared/src/story-segments.ts

@@ -195,6 +195,16 @@ function extractSeriesFromRawAttrs(attrString: string): ParsedChartBlock['series
 	return null;
 }
 
+export function extractQueryIds(code: string): Set<string> {
+	const ids = new Set<string>();
+	const regex = /<(?:chart|table)\s+[^>]*?\bquery_id\s*=\s*"([^"]+)"/g;
+	let match: RegExpExecArray | null;
+	while ((match = regex.exec(code)) !== null) {
+		ids.add(match[1]);
+	}
+	return ids;
+}
+
 export function splitCodeIntoSegments(code: string): Segment[] {
 	const segments: Segment[] = [];
 	const blockRegex = new RegExp(