feat(ai): improve chart runtime UX

apps/cli/src/ai/agent-stream-runner.test.ts

@@ -0,0 +1,42 @@
+import { describe, it } from 'node:test'
+import assert from 'node:assert/strict'
+import { getChartCapabilityAllowedBuiltinTools } from '@openchatlab/core'
+import { getAvailableToolDefs } from './agent-stream-runner'
+
+describe('CLI chart capability tool filtering', () => {
+  it('does not expose uncategorized raw SQL in chart-only turns', () => {
+    const allowedToolSet = new Set(getChartCapabilityAllowedBuiltinTools())
+
+    const toolNames = getAvailableToolDefs(true, allowedToolSet).map((tool) => tool.name)
+
+    assert.deepEqual(toolNames.sort(), ['get_schema', 'render_chart'])
+    assert.ok(!toolNames.includes('execute_sql'))
+  })
+
+  it('keeps only chart core tools plus explicitly allowed analysis tools', () => {
+    const allowedToolSet = new Set(getChartCapabilityAllowedBuiltinTools(['keyword_frequency', 'execute_sql']))
+
+    const toolNames = getAvailableToolDefs(true, allowedToolSet).map((tool) => tool.name)
+
+    assert.deepEqual(toolNames.sort(), ['get_schema', 'keyword_frequency', 'render_chart'])
+    assert.ok(!toolNames.includes('execute_sql'))
+  })
+
+  it('can expose render_chart for auto skill turns with restrictive assistant tools', () => {
+    const allowedToolSet = new Set(getChartCapabilityAllowedBuiltinTools(['keyword_frequency']))
+
+    const toolNames = getAvailableToolDefs(false, allowedToolSet).map((tool) => tool.name)
+
+    assert.ok(toolNames.includes('keyword_frequency'))
+    assert.ok(toolNames.includes('render_chart'))
+  })
+
+  it('does not expose raw SQL when the assistant did not allow it', () => {
+    const allowedToolSet = new Set(['keyword_frequency'])
+
+    const toolNames = getAvailableToolDefs(false, allowedToolSet).map((tool) => tool.name)
+
+    assert.ok(toolNames.includes('keyword_frequency'))
+    assert.ok(!toolNames.includes('execute_sql'))
+  })
+})

apps/cli/src/ai/agent-stream-runner.ts

@@ -4,12 +4,19 @@
  */
 
 import type { DatabaseManager, AIConversationManager, AgentStreamChunk } from '@openchatlab/node-runtime'
-import { SkillManager, createActivateSkillTool } from '@openchatlab/node-runtime'
 import {
+  SkillManager,
+  buildSkillMenuWithBuiltinChart,
+  createActivateSkillTool,
+  getSkillConfigWithBuiltinChart,
+} from '@openchatlab/node-runtime'
+import {
+  CHART_CAPABILITY_CORE_TOOLS,
   CHART_CAPABILITY_SKILL_ID,
   getChartCapabilityAllowedBuiltinTools,
   getChartCapabilitySkill,
   getChatOverview,
+  shouldUseChartCapabilityForMessage,
 } from '@openchatlab/core'
 import type { DataSnapshot } from '@openchatlab/node-runtime'
 import type { AgentStreamRequest } from '@openchatlab/http-routes'
@@ -27,6 +34,23 @@ function getAiDir(dbManager: DatabaseManager): string {
   return pathProvider.getAiDataDir()
 }
 
+const RAW_SQL_TOOL_NAMES = new Set(['execute_sql'])
+
+export function getAvailableToolDefs(isChartCapability: boolean, allowedToolSet: Set<string> | null) {
+  if (isChartCapability) {
+    const chartCoreTools = new Set<string>(CHART_CAPABILITY_CORE_TOOLS)
+    return AGENT_TOOL_REGISTRY.filter(
+      (tool) =>
+        chartCoreTools.has(tool.name) ||
+        (tool.category === 'analysis' && allowedToolSet?.has(tool.name) && !RAW_SQL_TOOL_NAMES.has(tool.name))
+    )
+  }
+
+  return allowedToolSet
+    ? AGENT_TOOL_REGISTRY.filter((tool) => tool.category !== 'analysis' || allowedToolSet.has(tool.name))
+    : AGENT_TOOL_REGISTRY
+}
+
 export function createCliRunAgentStream(
   dbManager: DatabaseManager,
   convManager: AIConversationManager
@@ -66,15 +90,18 @@ export function createCliRunAgentStream(
     const maxToolResultTokens = Math.floor(contextWindow * (maxToolResultPercent / 100))
 
     const db = (dbManager as any).open?.(sessionId)
-    const isChartCapability = skillId === CHART_CAPABILITY_SKILL_ID
+    const isChartCapability =
+      skillId === CHART_CAPABILITY_SKILL_ID || (!skillId && shouldUseChartCapabilityForMessage(userMessage))
+    const autoSkillAllowedTools =
+      !skillId && enableAutoSkill && assistantAllowedTools && assistantAllowedTools.length > 0
+        ? getChartCapabilityAllowedBuiltinTools(assistantAllowedTools)
+        : assistantAllowedTools
     const allowedToolSet = isChartCapability
       ? new Set(getChartCapabilityAllowedBuiltinTools(assistantAllowedTools))
-      : assistantAllowedTools && assistantAllowedTools.length > 0
-        ? new Set(assistantAllowedTools)
+      : autoSkillAllowedTools && autoSkillAllowedTools.length > 0
+        ? new Set(autoSkillAllowedTools)
         : null
-    const availableToolDefs = allowedToolSet
-      ? AGENT_TOOL_REGISTRY.filter((tool) => tool.category !== 'analysis' || allowedToolSet.has(tool.name))
-      : AGENT_TOOL_REGISTRY
+    const availableToolDefs = getAvailableToolDefs(isChartCapability, allowedToolSet)
 
     const agentTools = db
       ? adaptToolsForAgent(availableToolDefs, () => ({ db, sessionId, locale }), { maxToolResultTokens })
@@ -93,7 +120,7 @@ export function createCliRunAgentStream(
       const def = skillMgr.getSkillConfig(skillId)
       if (def) resolvedSkillDef = { name: def.name, prompt: def.prompt }
     } else if (enableAutoSkill) {
-      const menu = skillMgr.getSkillMenu(chatType ?? 'group', toolNames)
+      const menu = buildSkillMenuWithBuiltinChart(skillMgr.getSkillMenu(chatType ?? 'group', toolNames), locale)
       if (menu) resolvedSkillMenu = menu
     }
 
@@ -101,8 +128,10 @@ export function createCliRunAgentStream(
       const activateSkillTool = createActivateSkillTool({
         chatType: chatType ?? 'group',
         allowedTools: toolNames,
+        coreToolNames: new Set<string>(CHART_CAPABILITY_CORE_TOOLS),
         locale,
-        getSkillConfig: (id) => skillMgr.getSkillConfig(id),
+        getSkillConfig: (id) =>
+          getSkillConfigWithBuiltinChart(id, locale, (skillConfigId) => skillMgr.getSkillConfig(skillConfigId)),
       })
       agentTools.push(activateSkillTool as any)
     }

apps/cli/src/ai/tool-adapter.test.ts

@@ -3,6 +3,7 @@ import assert from 'node:assert/strict'
 import { adaptToolsForAgent } from './tool-adapter'
 import type { ToolDefinition } from '@openchatlab/tools'
 import type { ChartPayload, DatabaseAdapter } from '@openchatlab/core'
+import { CHART_SCHEMA_REQUIRED_MESSAGE } from '@openchatlab/node-runtime'
 
 const chart: ChartPayload = {
   version: 1,
@@ -30,7 +31,40 @@ const chart: ChartPayload = {
 }
 
 describe('adaptToolsForAgent', () => {
-  it('preserves chart payloads in tool result details', async () => {
+  it('requires get_schema before render_chart', async () => {
+    let chartCalls = 0
+    const tool: ToolDefinition = {
+      name: 'render_chart',
+      description: 'Render a chart',
+      inputSchema: { type: 'object', properties: {} },
+      async handler() {
+        chartCalls += 1
+        return { content: 'Generated chart.' }
+      },
+    }
+
+    const [agentTool] = adaptToolsForAgent([tool], () => ({
+      db: {} as DatabaseAdapter,
+      sessionId: 'session-1',
+      locale: 'en-US',
+    }))
+
+    const result = await agentTool.execute('call-1', {})
+
+    assert.equal(chartCalls, 0)
+    assert.deepEqual(result.content, [{ type: 'text', text: CHART_SCHEMA_REQUIRED_MESSAGE }])
+    assert.equal(result.details, null)
+  })
+
+  it('preserves chart payloads in tool result details after schema lookup', async () => {
+    const getSchemaTool: ToolDefinition = {
+      name: 'get_schema',
+      description: 'Get schema',
+      inputSchema: { type: 'object', properties: {} },
+      async handler() {
+        return { content: 'message(id, ts)' }
+      },
+    }
     const tool: ToolDefinition = {
       name: 'render_chart',
       description: 'Render a chart',
@@ -44,13 +78,14 @@ describe('adaptToolsForAgent', () => {
       },
     }
 
-    const [agentTool] = adaptToolsForAgent([tool], () => ({
+    const [getSchema, agentTool] = adaptToolsForAgent([getSchemaTool, tool], () => ({
       db: {} as DatabaseAdapter,
       sessionId: 'session-1',
       locale: 'en-US',
     }))
 
     assert.ok(agentTool)
+    await getSchema.execute('call-0', {})
     const result = await agentTool.execute('call-1', {})
 
     assert.deepEqual(result.content, [{ type: 'text', text: 'Generated chart.' }])

apps/cli/src/ai/tool-adapter.ts

@@ -15,6 +15,8 @@ import {
   type AgentToolResult,
   type PreprocessableMessage,
   type TruncationStrategy,
+  createChartSchemaGateState,
+  wrapWithChartSchemaGate,
 } from '@openchatlab/node-runtime'
 import { getServerAiLogger } from './logger'
 
@@ -57,62 +59,68 @@ export function adaptToolsForAgent(
   options?: AdaptToolsOptions
 ): AgentTool<any, any>[] {
   const tokenBudget = options?.maxToolResultTokens ?? DEFAULT_MAX_TOOL_RESULT_TOKENS
+  const chartSchemaGateState = createChartSchemaGateState()
 
-  return tools.map((tool) => ({
-    name: tool.name,
-    label: tool.name,
-    description: tool.description,
-    parameters: convertJsonSchemaToParameters(tool.inputSchema) as any,
-    async execute(_toolCallId: string, params: unknown): Promise<AgentToolResult<unknown>> {
-      const toolParams = (params && typeof params === 'object' ? params : {}) as Record<string, unknown>
-      const ctx = getContext()
-      const execCtx: ToolExecutionContext = {
-        db: ctx.db,
-        dataProvider: new CoreDataProvider(ctx.db),
-        sessionId: ctx.sessionId,
-        locale: ctx.locale,
-        segmentText: (texts, locale, options) => batchSegmentWithFrequency(texts, locale as any, options as any),
-      }
-      try {
-        const result = await tool.handler(toolParams, execCtx)
-        const chartDetails =
-          result.chart || result.charts
-            ? {
-                ...(result.chart ? { chart: result.chart } : {}),
-                ...(result.charts ? { charts: result.charts } : {}),
-              }
-            : {}
-
-        if (result.rawMessages && result.rawMessages.length > 0) {
-          const pipelineResult = applyPreprocessingPipeline({
-            rawMessages: result.rawMessages as PreprocessableMessage[],
+  return tools.map((tool) =>
+    wrapWithChartSchemaGate(
+      {
+        name: tool.name,
+        label: tool.name,
+        description: tool.description,
+        parameters: convertJsonSchemaToParameters(tool.inputSchema) as any,
+        async execute(_toolCallId: string, params: unknown): Promise<AgentToolResult<unknown>> {
+          const toolParams = (params && typeof params === 'object' ? params : {}) as Record<string, unknown>
+          const ctx = getContext()
+          const execCtx: ToolExecutionContext = {
+            db: ctx.db,
+            dataProvider: new CoreDataProvider(ctx.db),
+            sessionId: ctx.sessionId,
             locale: ctx.locale,
-            maxToolResultTokens: tokenBudget,
-            truncationStrategy: TOOL_TRUNCATION_STRATEGY[tool.name] ?? 'keep_last',
-            extraDetails: (result.data ?? {}) as Record<string, unknown>,
-            logger: getServerAiLogger() ?? undefined,
-          })
-          return {
-            content: [{ type: 'text', text: pipelineResult.text }],
-            details: Object.keys(chartDetails).length > 0 ? chartDetails : null,
+            segmentText: (texts, locale, options) => batchSegmentWithFrequency(texts, locale as any, options as any),
           }
-        }
+          try {
+            const result = await tool.handler(toolParams, execCtx)
+            const chartDetails =
+              result.chart || result.charts
+                ? {
+                    ...(result.chart ? { chart: result.chart } : {}),
+                    ...(result.charts ? { charts: result.charts } : {}),
+                  }
+                : {}
 
-        const baseDetails =
-          typeof result.data === 'object' && result.data !== null
-            ? (result.data as Record<string, unknown>)
-            : result.data === undefined
-              ? null
-              : { value: result.data }
+            if (result.rawMessages && result.rawMessages.length > 0) {
+              const pipelineResult = applyPreprocessingPipeline({
+                rawMessages: result.rawMessages as PreprocessableMessage[],
+                locale: ctx.locale,
+                maxToolResultTokens: tokenBudget,
+                truncationStrategy: TOOL_TRUNCATION_STRATEGY[tool.name] ?? 'keep_last',
+                extraDetails: (result.data ?? {}) as Record<string, unknown>,
+                logger: getServerAiLogger() ?? undefined,
+              })
+              return {
+                content: [{ type: 'text', text: pipelineResult.text }],
+                details: Object.keys(chartDetails).length > 0 ? chartDetails : null,
+              }
+            }
 
-        return {
-          content: [{ type: 'text', text: result.content }],
-          details: Object.keys(chartDetails).length > 0 ? { ...(baseDetails ?? {}), ...chartDetails } : baseDetails,
-        }
-      } catch (error) {
-        const msg = error instanceof Error ? error.message : String(error)
-        return { content: [{ type: 'text', text: `Error: ${msg}` }], details: null }
-      }
-    },
-  }))
+            const baseDetails =
+              typeof result.data === 'object' && result.data !== null
+                ? (result.data as Record<string, unknown>)
+                : result.data === undefined
+                  ? null
+                  : { value: result.data }
+
+            return {
+              content: [{ type: 'text', text: result.content }],
+              details: Object.keys(chartDetails).length > 0 ? { ...(baseDetails ?? {}), ...chartDetails } : baseDetails,
+            }
+          } catch (error) {
+            const msg = error instanceof Error ? error.message : String(error)
+            return { content: [{ type: 'text', text: `Error: ${msg}` }], details: null }
+          }
+        },
+      },
+      chartSchemaGateState
+    )
+  )
 }