osiman/SuperBizAgent-java
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat(harness): add readonly mysql tool

Browse Source
This commit is contained in:
zhuyongxin
2026-07-21 21:24:42 +08:00
parent 3e602781d6
commit 85029d96a7
31 changed files with 1896 additions and 53 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/test/java/com/superbiz/agent/harness/tool/adapter/MysqlToolAdapterTest.java
+120
View File
@@ -0,0 +1,120 @@
package com.superbiz.agent.harness.tool.adapter;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.superbiz.agent.harness.contract.EvidenceStatus;
import com.superbiz.agent.harness.core.HarnessCoreFixtures;
import com.superbiz.agent.harness.core.MutableClock;
import com.superbiz.agent.harness.core.RunContext;
import com.superbiz.agent.harness.tool.boundary.ToolBoundary;
import com.superbiz.agent.harness.tool.boundary.ToolBoundaryErrorCode;
import com.superbiz.agent.harness.tool.boundary.ToolBoundaryResult;
import com.superbiz.agent.harness.tool.boundary.ToolCallRequestEnvelope;
import com.superbiz.agent.harness.tool.mysql.MysqlDataSourceDefinition;
import com.superbiz.agent.harness.tool.mysql.MysqlRawResult;
import com.superbiz.agent.harness.tool.mysql.MysqlResultProjector;
import com.superbiz.agent.harness.tool.mysql.MysqlSqlValidator;
import com.superbiz.agent.harness.tool.mysql.MysqlToolLimits;
import com.superbiz.agent.harness.tool.store.CanonicalInvocationLimits;
import com.superbiz.agent.harness.tool.store.CanonicalInvocationStore;
import com.superbiz.agent.harness.tool.store.CanonicalToolInvocation;
import com.superbiz.agent.harness.tool.store.DuplicateInvocationException;
import com.superbiz.agent.harness.tool.store.ToolCallKeyFactory;
import org.junit.jupiter.api.Test;
import java.time.Duration;
import java.time.Instant;
import java.util.HashMap;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;
import java.util.Optional;
import java.util.Set;
import java.util.concurrent.atomic.AtomicInteger;
import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertFalse;
import static org.junit.jupiter.api.Assertions.assertTrue;
class MysqlToolAdapterTest {
private final ObjectMapper objectMapper = new ObjectMapper();
@Test
void reusesBoundaryAndRejectsInvalidSqlBeforeExecution() {
MutableClock clock = new MutableClock(Instant.parse("2026-07-21T10:00:00Z"));
FakeStore store = new FakeStore();
ToolBoundary boundary = new ToolBoundary(HarnessCoreFixtures.core(clock),
new ToolCallKeyFactory("superbiz:harness:tool-call"), store, objectMapper, clock);
MysqlDataSourceDefinition definition = new MysqlDataSourceDefinition(
"order_readonly", "order_db",
Map.of("order_db", Map.of("biz_order", Set.of("order_id", "password"))),
MysqlToolLimits.defaults());
AtomicInteger executions = new AtomicInteger();
MysqlToolAdapter adapter = new MysqlToolAdapter(boundary, objectMapper,
new MysqlSqlValidator(Map.of("order_readonly", definition)),
(plan, context) -> {
executions.incrementAndGet();
Map<String, Object> row = new LinkedHashMap<>();
row.put("order_id", "order-1");
row.put("password", "raw-secret");
return new MysqlRawResult(List.of("order_id", "password"), List.of(row), false);
}, new MysqlResultProjector(objectMapper));
RunContext context = HarnessCoreFixtures.core(clock).startRun("session", "run-mysql");
ToolBoundaryResult valid = adapter.execute(context, envelope("run-mysql", "framework-mysql-1",
"{\"data_source\":\"order_readonly\",\"sql\":\"SELECT order_id, password FROM order_db.biz_order WHERE order_id = ?\",\"params\":[\"order-1\"]}"));
ToolBoundaryResult invalid = adapter.execute(context, envelope("run-mysql", "framework-mysql-2",
"{\"data_source\":\"order_readonly\",\"sql\":\"SELECT * FROM order_db.biz_order\",\"params\":[]}"));
assertEquals(EvidenceStatus.EVIDENCE_FOUND, valid.evidenceStatus());
assertEquals("framework-mysql-1", valid.toolCallId());
assertFalse(valid.agentResult().contains("raw-secret"));
assertTrue(store.records.values().iterator().next().rawResponse().contains("raw-secret"));
assertEquals(ToolBoundaryErrorCode.INVALID_REQUEST.name(), invalid.errorCode());
assertEquals(1, executions.get());
}
private ToolCallRequestEnvelope envelope(String runId, String callId, String requestJson) {
return new ToolCallRequestEnvelope(runId, callId, "query_mysql", requestJson, true, true);
}
private static final class FakeStore implements CanonicalInvocationStore {
private final CanonicalInvocationLimits limits =
new CanonicalInvocationLimits(Duration.ofHours(1), 100_000, 64_000);
private final Map<String, CanonicalToolInvocation> records = new HashMap<>();
@Override
public CanonicalInvocationLimits limits() {
return limits;
}
@Override
public void begin(String key, CanonicalToolInvocation invocation) {
if (records.putIfAbsent(key, invocation) != null) {
throw new DuplicateInvocationException();
}
}
@Override
public Optional<CanonicalToolInvocation> find(String key) {
return Optional.ofNullable(records.get(key));
}
@Override
public CanonicalToolInvocation markReady(String key, String rawResponse, String agentResult,
EvidenceStatus evidenceStatus, Instant completedAt) {
CanonicalToolInvocation updated = records.get(key)
.markReady(rawResponse, agentResult, evidenceStatus, completedAt);
records.put(key, updated);
return updated;
}
@Override
public CanonicalToolInvocation markError(String key, String rawResponse,
String errorCode, Instant completedAt) {
CanonicalToolInvocation updated = records.get(key).markError(rawResponse, errorCode, completedAt);
records.put(key, updated);
return updated;
}
}
}
src/test/java/com/superbiz/agent/harness/tool/mysql/JdbcMysqlReadOnlyExecutorTest.java
+112
View File
@@ -0,0 +1,112 @@
package com.superbiz.agent.harness.tool.mysql;
import com.superbiz.agent.harness.core.HarnessCoreFixtures;
import com.superbiz.agent.harness.core.MutableClock;
import com.superbiz.agent.harness.core.RunCancellationReason;
import com.superbiz.agent.harness.core.RunContext;
import com.superbiz.agent.harness.tool.contract.MysqlToolRequest;
import org.h2.jdbcx.JdbcDataSource;
import org.junit.jupiter.api.Test;
import java.sql.SQLException;
import java.sql.Connection;
import java.sql.PreparedStatement;
import java.sql.ResultSet;
import java.sql.ResultSetMetaData;
import java.time.Instant;
import java.util.List;
import java.util.Map;
import java.util.Set;
import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertThrows;
import static org.junit.jupiter.api.Assertions.assertTrue;
import static org.mockito.ArgumentMatchers.anyString;
import static org.mockito.ArgumentMatchers.eq;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.when;
class JdbcMysqlReadOnlyExecutorTest {
@Test
void executesBoundPreparedSelectAndTruncatesRows() throws Exception {
JdbcDataSource dataSource = dataSource();
MysqlDataSourceDefinition definition = definition(new MysqlToolLimits(1, 20, 1024, 2));
MysqlSqlValidator validator = new MysqlSqlValidator(Map.of("order_readonly", definition));
MysqlQueryPlan plan = validator.validate(new MysqlToolRequest(
"order_readonly", "SELECT order_id, status FROM biz_order WHERE status = ? ORDER BY order_id",
List.of("FAILED")));
MutableClock clock = new MutableClock(Instant.parse("2026-07-21T10:00:00Z"));
RunContext context = HarnessCoreFixtures.core(clock).startRun("session", "run-mysql");
MysqlRawResult result = new JdbcMysqlReadOnlyExecutor(Map.of("order_readonly", dataSource), clock)
.execute(plan, context);
assertEquals(List.of("ORDER_ID", "STATUS"), result.columns());
assertEquals(1, result.rows().size());
assertTrue(result.truncated());
}
@Test
void refusesExecutionAfterRunCancellation() throws Exception {
JdbcDataSource dataSource = dataSource();
MysqlDataSourceDefinition definition = definition(MysqlToolLimits.defaults());
MysqlSqlValidator validator = new MysqlSqlValidator(Map.of("order_readonly", definition));
MysqlQueryPlan plan = validator.validate(new MysqlToolRequest(
"order_readonly", "SELECT order_id FROM biz_order", List.of()));
MutableClock clock = new MutableClock(Instant.parse("2026-07-21T10:00:00Z"));
RunContext context = HarnessCoreFixtures.core(clock).startRun("session", "run-cancelled");
context.cancellation().cancel(RunCancellationReason.USER_REQUESTED);
assertThrows(SQLException.class,
() -> new JdbcMysqlReadOnlyExecutor(Map.of("order_readonly", dataSource), clock)
.execute(plan, context));
}
@Test
void appliesReadOnlyTimeoutMaxRowsAndParameterBinding() throws Exception {
javax.sql.DataSource dataSource = mock(javax.sql.DataSource.class);
Connection connection = mock(Connection.class);
PreparedStatement statement = mock(PreparedStatement.class);
ResultSet resultSet = mock(ResultSet.class);
ResultSetMetaData metadata = mock(ResultSetMetaData.class);
when(dataSource.getConnection()).thenReturn(connection);
when(connection.prepareStatement(anyString(), eq(ResultSet.TYPE_FORWARD_ONLY), eq(ResultSet.CONCUR_READ_ONLY)))
.thenReturn(statement);
when(statement.executeQuery()).thenReturn(resultSet);
when(resultSet.getMetaData()).thenReturn(metadata);
when(metadata.getColumnCount()).thenReturn(1);
when(metadata.getColumnLabel(1)).thenReturn("order_id");
when(resultSet.next()).thenReturn(false);
MysqlDataSourceDefinition definition = definition(new MysqlToolLimits(7, 20, 1024, 3));
MysqlQueryPlan plan = new MysqlSqlValidator(Map.of("order_readonly", definition))
.validate(new MysqlToolRequest("order_readonly",
"SELECT order_id FROM biz_order WHERE status = ?", List.of("FAILED")));
MutableClock clock = new MutableClock(Instant.parse("2026-07-21T10:00:00Z"));
RunContext context = HarnessCoreFixtures.core(clock).startRun("session", "run-controls");
new JdbcMysqlReadOnlyExecutor(Map.of("order_readonly", dataSource), clock).execute(plan, context);
verify(connection).setReadOnly(true);
verify(statement).setQueryTimeout(3);
verify(statement).setMaxRows(8);
verify(statement).setObject(1, "FAILED");
}
private JdbcDataSource dataSource() throws Exception {
JdbcDataSource dataSource = new JdbcDataSource();
dataSource.setURL("jdbc:h2:mem:mysqltool;MODE=MySQL;DB_CLOSE_DELAY=-1");
try (var connection = dataSource.getConnection(); var statement = connection.createStatement()) {
statement.execute("DROP TABLE IF EXISTS biz_order");
statement.execute("CREATE TABLE biz_order(order_id VARCHAR(32), status VARCHAR(20))");
statement.execute("INSERT INTO biz_order VALUES ('order-1','FAILED'),('order-2','FAILED'),('order-3','OK')");
}
return dataSource;
}
private MysqlDataSourceDefinition definition(MysqlToolLimits limits) {
return new MysqlDataSourceDefinition("order_readonly", "PUBLIC",
Map.of("PUBLIC", Map.of("biz_order", Set.of("order_id", "status"))), limits);
}
}
src/test/java/com/superbiz/agent/harness/tool/mysql/MysqlResultProjectorTest.java
+67
View File
@@ -0,0 +1,67 @@
package com.superbiz.agent.harness.tool.mysql;
import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.superbiz.agent.harness.contract.EvidenceStatus;
import com.superbiz.agent.harness.tool.boundary.ProjectedToolResult;
import com.superbiz.agent.harness.tool.contract.MysqlToolRequest;
import org.junit.jupiter.api.Test;
import java.util.List;
import static org.junit.jupiter.api.Assertions.assertEquals;
import static org.junit.jupiter.api.Assertions.assertFalse;
import static org.junit.jupiter.api.Assertions.assertTrue;
class MysqlResultProjectorTest {
private final ObjectMapper objectMapper = new ObjectMapper();
@Test
void projectsRowsRedactsSensitiveColumnsAndBoundsValues() throws Exception {
MysqlResultProjector projector = new MysqlResultProjector(objectMapper,
new MysqlToolLimits(1, 5, 1024, 5));
String raw = "{\"columns\":[\"order_id\",\"password\",\"status\"],\"rows\":["
+ "{\"order_id\":\"order-123456\",\"password\":\"secret\",\"status\":\"FAILED\"},"
+ "{\"order_id\":\"order-2\",\"password\":\"secret2\",\"status\":\"OK\"}],\"truncated\":false}";
ProjectedToolResult projected = projector.project(
new MysqlToolRequest("order_readonly", "SELECT order_id FROM biz_order", List.of()),
"framework-mysql-1", raw);
JsonNode json = objectMapper.readTree(projected.agentResult());
assertEquals(EvidenceStatus.EVIDENCE_FOUND, projected.evidenceStatus());
assertEquals(1, json.path("returned_count").asInt());
assertTrue(json.path("truncated").asBoolean());
assertEquals("[REDACTED]", json.path("rows").get(0).path("password").asText());
assertEquals("order", json.path("rows").get(0).path("order_id").asText());
assertFalse(projected.agentResult().contains("secret"));
}
@Test
void returnsNoEvidenceForEmptyRows() throws Exception {
MysqlResultProjector projector = new MysqlResultProjector(objectMapper);
ProjectedToolResult projected = projector.project(
new MysqlToolRequest("order_readonly", "SELECT order_id FROM biz_order", List.of()),
"framework-mysql-2", "{\"columns\":[\"order_id\"],\"rows\":[],\"truncated\":false}");
assertEquals(EvidenceStatus.NO_EVIDENCE, projected.evidenceStatus());
}
@Test
void enforcesTotalUtf8BudgetWithValidJson() throws Exception {
MysqlResultProjector projector = new MysqlResultProjector(objectMapper,
new MysqlToolLimits(10, 100, 230, 5));
String raw = "{\"columns\":[\"order_id\"],\"rows\":["
+ "{\"order_id\":\"" + "a".repeat(80) + "\"},"
+ "{\"order_id\":\"" + "b".repeat(80) + "\"}],\"truncated\":false}";
ProjectedToolResult projected = projector.project(
new MysqlToolRequest("order_readonly", "SELECT order_id FROM biz_order", List.of()),
"framework-mysql-3", raw);
JsonNode json = objectMapper.readTree(projected.agentResult());
assertTrue(json.path("truncated").asBoolean());
assertTrue(projected.agentResult().getBytes(java.nio.charset.StandardCharsets.UTF_8).length <= 230);
}
}
src/test/java/com/superbiz/agent/harness/tool/mysql/MysqlSqlValidatorTest.java
+68
View File
@@ -0,0 +1,68 @@
package com.superbiz.agent.harness.tool.mysql;
import com.superbiz.agent.harness.tool.contract.MysqlToolRequest;
import org.junit.jupiter.api.Test;
import java.util.List;
import java.util.Map;
import java.util.Set;
import static org.junit.jupiter.api.Assertions.assertDoesNotThrow;
import static org.junit.jupiter.api.Assertions.assertThrows;
class MysqlSqlValidatorTest {
private final MysqlSqlValidator validator = new MysqlSqlValidator(Map.of(
"order_readonly", new MysqlDataSourceDefinition(
"order_readonly", "order_db",
Map.of("order_db", Map.of(
"biz_order", Set.of("order_id", "status", "updated_at"),
"payment_record", Set.of("order_id", "status"))),
MysqlToolLimits.defaults())));
@Test
void acceptsExplicitSelectAndCountStar() {
assertDoesNotThrow(() -> validator.validate(new MysqlToolRequest(
"order_readonly",
"SELECT o.order_id, o.status FROM order_db.biz_order o WHERE o.order_id = ?",
List.of("order-1"))));
assertDoesNotThrow(() -> validator.validate(new MysqlToolRequest(
"order_readonly",
"SELECT COUNT(*) AS total FROM order_db.biz_order",
List.of())));
assertDoesNotThrow(() -> validator.validate(new MysqlToolRequest(
"order_readonly",
"SELECT o.order_id, p.status FROM order_db.biz_order o INNER JOIN order_db.payment_record p ON o.order_id = p.order_id WHERE o.status = ?",
List.of("FAILED"))));
}
@Test
void rejectsWritesCompoundQueriesWildcardsAndAllowlistBypass() {
List<String> unsafe = List.of(
"UPDATE order_db.biz_order SET status = ?",
"WITH x AS (SELECT order_id FROM order_db.biz_order) SELECT order_id FROM x",
"SELECT * FROM order_db.biz_order",
"SELECT order_id FROM order_db.biz_order WHERE order_id IN (SELECT order_id FROM order_db.payment_record)",
"SELECT order_id FROM order_db.biz_order UNION SELECT order_id FROM order_db.payment_record",
"SELECT order_id FROM order_db.biz_order WHERE secret = ?",
"SELECT order_id FROM order_db.unknown_table",
"SELECT SLEEP(?) FROM order_db.biz_order",
"SELECT order_id FROM order_db.biz_order WHERE status = 'FAILED'",
"SELECT CASE WHEN status = ? THEN order_id END FROM order_db.biz_order",
"SELECT order_id FROM order_db.biz_order FOR UPDATE",
"SELECT order_id FROM order_db.biz_order; SELECT status FROM order_db.biz_order");
for (String sql : unsafe) {
assertThrows(MysqlSecurityException.class,
() -> validator.validate(new MysqlToolRequest("order_readonly", sql, List.of("x"))), sql);
}
}
@Test
void rejectsPlaceholderMismatchAndAmbiguousColumns() {
assertThrows(MysqlSecurityException.class, () -> validator.validate(new MysqlToolRequest(
"order_readonly", "SELECT order_id FROM order_db.biz_order WHERE order_id = ?", List.of())));
assertThrows(MysqlSecurityException.class, () -> validator.validate(new MysqlToolRequest(
"order_readonly", "SELECT order_id FROM order_db.biz_order o INNER JOIN order_db.payment_record p ON o.order_id = p.order_id", List.of())));
}
}