fix(tools): preserve datasource id tool inputs

This commit is contained in:
matevip 2026-08-13 03:58:05 -04:00
parent 7bf18e6af8
commit 7a34f3a502
4 changed files with 74 additions and 5 deletions

View File

@ -61,7 +61,7 @@ public class DatasourceTool {
""")
public String query_datasource(
@ToolParam(description = "动作list_datasources / list_tables / describe_table") String action,
@ToolParam(description = "数据源 IDlist_tables 和 describe_table 时必填)", required = false) Long datasourceId,
@ToolParam(description = "数据源 IDlist_tables 和 describe_table 时必填)。必须作为字符串传入,避免大整数精度丢失", required = false) String datasourceId,
@ToolParam(description = "表名describe_table 时必填)", required = false) String tableName) {
try {
@ -208,4 +208,24 @@ public class DatasourceTool {
private String error(String message) {
return JSONUtil.toJsonStr(new JSONObject().set("error", message));
}
private Long parseDatasourceId(String datasourceId, String action) {
String trimmed = datasourceId != null ? datasourceId.trim() : "";
if (trimmed.isEmpty()) {
throw new IllegalArgumentException(action + " 需要 datasourceId 参数");
}
try {
return Long.parseLong(trimmed);
} catch (NumberFormatException e) {
throw new IllegalArgumentException("datasourceId 必须是数字字符串");
}
}
private String listTables(String datasourceId) throws SQLException {
return listTables(parseDatasourceId(datasourceId, "list_tables"));
}
private String describeTable(String datasourceId, String tableName) throws SQLException {
return describeTable(parseDatasourceId(datasourceId, "describe_table"), tableName);
}
}

View File

@ -47,16 +47,17 @@ public class SqlQueryTool {
如果数据适合可视化会自动附带一个 echarts 图表配置代码块前端会自动渲染为交互式图表
""")
public String execute_sql(
@ToolParam(description = "目标数据源 ID") Long datasourceId,
@ToolParam(description = "目标数据源 ID。必须作为字符串传入,避免大整数精度丢失") String datasourceId,
@ToolParam(description = "要执行的 SQL 查询(仅允许 SELECT") String sql) {
try {
Long parsedDatasourceId = parseDatasourceId(datasourceId);
// 1. 验证并规范化 SQL
String safeSql = sqlValidationService.validateAndNormalize(sql);
log.info("执行 SQL 查询 [数据源 {}]: {}", datasourceId, safeSql);
log.info("执行 SQL 查询 [数据源 {}]: {}", parsedDatasourceId, safeSql);
// 2. 获取数据源连接
DatasourceEntity entity = datasourceService.getDecrypted(datasourceId);
DatasourceEntity entity = datasourceService.getDecrypted(parsedDatasourceId);
// 3. 执行查询
long startTime = System.currentTimeMillis();
@ -79,6 +80,18 @@ public class SqlQueryTool {
}
}
private Long parseDatasourceId(String datasourceId) {
String trimmed = datasourceId != null ? datasourceId.trim() : "";
if (trimmed.isEmpty()) {
throw new IllegalArgumentException("datasourceId is required");
}
try {
return Long.parseLong(trimmed);
} catch (NumberFormatException e) {
throw new IllegalArgumentException("datasourceId must be a numeric string");
}
}
private String formatResult(ResultSet rs, String sql, long elapsedMs) throws SQLException {
ResultSetMetaData meta = rs.getMetaData();
int colCount = meta.getColumnCount();

View File

@ -1,14 +1,18 @@
package vip.mate.tool.builtin;
import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.fasterxml.jackson.databind.json.JsonMapper;
import com.fasterxml.jackson.databind.ser.std.ToStringSerializer;
import com.fasterxml.jackson.databind.module.SimpleModule;
import org.junit.jupiter.api.DisplayName;
import org.junit.jupiter.api.Test;
import org.springframework.ai.support.ToolCallbacks;
import org.springframework.ai.tool.ToolCallback;
import vip.mate.datasource.model.DatasourceEntity;
import vip.mate.datasource.service.DatasourceConnectionManager;
import vip.mate.datasource.service.DatasourceService;
import vip.mate.datasource.service.SqlValidationService;
import java.util.List;
@ -56,4 +60,36 @@ class DatasourceToolIdPrecisionTest {
assertFalse(out.contains(": " + bigId) || out.contains(":" + bigId),
"id must NOT appear as a bare JSON number (precision-lossy across double/JS Number)");
}
@Test
@DisplayName("datasource tools publish datasourceId as a string parameter so LLM tool calls preserve precision")
void datasourceIdSchemasAreString() throws Exception {
DatasourceTool datasourceTool = new DatasourceTool(
mock(DatasourceService.class),
mock(DatasourceConnectionManager.class),
idSafeMapper());
SqlQueryTool sqlQueryTool = new SqlQueryTool(
mock(DatasourceService.class),
mock(DatasourceConnectionManager.class),
mock(SqlValidationService.class));
assertDatasourceIdIsString(datasourceTool, "query_datasource");
assertDatasourceIdIsString(sqlQueryTool, "execute_sql");
}
private static void assertDatasourceIdIsString(Object tool, String name) throws Exception {
JsonNode root = idSafeMapper().readTree(callback(tool, name).getToolDefinition().inputSchema());
assertTrue("string".equals(root.at("/properties/datasourceId/type").asText()),
name + " datasourceId must be a string schema");
}
private static ToolCallback callback(Object tool, String name) {
for (ToolCallback callback : ToolCallbacks.from(tool)) {
if (name.equals(callback.getToolDefinition().name())) {
return callback;
}
}
throw new AssertionError("Missing tool callback: " + name);
}
}

View File

@ -52,7 +52,7 @@ class DatasourceToolPostgresqlViewDiscoveryTest {
DatasourceTool tool = new DatasourceTool(service, connectionManager, JsonMapper.builder().build());
// When the agent asks MateClaw to discover available relations.
String output = tool.query_datasource("list_tables", 1L, null);
String output = tool.query_datasource("list_tables", "1", null);
// Then the metadata query must include views and return the discovered view.
ArgumentCaptor<String> sql = ArgumentCaptor.forClass(String.class);