diff --git a/src/dbjavagenix/database/atomic_codegen_tools.py b/src/dbjavagenix/database/atomic_codegen_tools.py index 2b55c2b..327a32d 100644 --- a/src/dbjavagenix/database/atomic_codegen_tools.py +++ b/src/dbjavagenix/database/atomic_codegen_tools.py @@ -31,6 +31,7 @@ from ..core.exceptions import DatabaseConnectionError, MCPServiceError from ..database.connection_manager import connection_manager +from ..utils.json_serialization import dumps as _json_dumps logger = logging.getLogger(__name__) @@ -273,13 +274,13 @@ async def handle_codegen_build_context(arguments: Dict[str, Any]) -> List[TextCo ], }, } - return [TextContent(type="text", text=json.dumps(result, ensure_ascii=False, indent=2))] + return [TextContent(type="text", text=_json_dumps(result, indent=2))] except (DatabaseConnectionError, MCPServiceError) as e: return [ TextContent( type="text", - text=json.dumps({"error": str(e), "stage": "build_context"}, ensure_ascii=False), + text=_json_dumps({"error": str(e), "stage": "build_context"}), ) ] except Exception as e: # noqa: BLE001 @@ -287,9 +288,7 @@ async def handle_codegen_build_context(arguments: Dict[str, Any]) -> List[TextCo return [ TextContent( type="text", - text=json.dumps( - {"error": f"unexpected: {e}", "stage": "build_context"}, ensure_ascii=False - ), + text=_json_dumps({"error": f"unexpected: {e}", "stage": "build_context"}), ) ] @@ -314,19 +313,18 @@ async def handle_codegen_render_dto(arguments: Dict[str, Any]) -> List[TextConte """Render the Java 21 record DTO exposed by the sb35-java21 template family.""" context = _extract_context(arguments) if not isinstance(context, dict): - return [TextContent(type="text", text=json.dumps({"error": "context missing or invalid"}))] + return [TextContent(type="text", text=_json_dumps({"error": "context missing or invalid"}))] if context.get("templateCategory") != "sb35-java21": return [ TextContent( type="text", - text=json.dumps( + text=_json_dumps( { "files": [], "language": "java", "note": "template_category does not provide a record DTO", }, - ensure_ascii=False, indent=2, ), ) @@ -345,7 +343,7 @@ async def handle_codegen_render_mapper(arguments: Dict[str, Any]) -> List[TextCo """ context = _extract_context(arguments) if not isinstance(context, dict): - return [TextContent(type="text", text=json.dumps({"error": "context missing or invalid"}))] + return [TextContent(type="text", text=_json_dumps({"error": "context missing or invalid"}))] category = context.get("templateCategory", "") templates: List[str] = [] @@ -363,13 +361,12 @@ async def handle_codegen_render_mapper(arguments: Dict[str, Any]) -> List[TextCo return [ TextContent( type="text", - text=json.dumps( + text=_json_dumps( { "files": [], "language": "java", "note": f"template_category={category} 不需要 mapper 层(BaseMapper/JpaRepository 内置)", }, - ensure_ascii=False, indent=2, ), ) @@ -406,7 +403,7 @@ async def _render_single_layer( return [ TextContent( type="text", - text=json.dumps({"error": "context missing or not a dict"}, ensure_ascii=False), + text=_json_dumps({"error": "context missing or not a dict"}), ) ] @@ -419,13 +416,12 @@ async def _render_single_layer( return [ TextContent( type="text", - text=json.dumps( + text=_json_dumps( { "error": "unsupported template category", "template_category": category, "supported_categories": supported_categories, }, - ensure_ascii=False, indent=2, ), ) @@ -492,7 +488,7 @@ async def _render_single_layer( content = TextContent( type="text", - text=json.dumps({"files": files, "language": "java"}, ensure_ascii=False, indent=2), + text=_json_dumps({"files": files, "language": "java"}, indent=2), ) return [attach_meta(content, diff_meta)] diff --git a/src/dbjavagenix/database/mcp_tools.py b/src/dbjavagenix/database/mcp_tools.py index c74be1f..edfa186 100644 --- a/src/dbjavagenix/database/mcp_tools.py +++ b/src/dbjavagenix/database/mcp_tools.py @@ -12,7 +12,8 @@ from functools import lru_cache from pathlib import Path, PureWindowsPath from typing import Dict, Any, List, Optional -from uuid import UUID +from ..utils.json_serialization import default as _query_result_json_default +from ..utils.json_serialization import dumps as _json_dumps from mcp.types import Tool, TextContent, ImageContent, EmbeddedResource @@ -59,29 +60,6 @@ _DOLLAR_QUOTE_PATTERN = re.compile(r"\$(?:[A-Za-z_][A-Za-z0-9_]*)?\$") -def _query_result_json_default(value: object) -> object: - """Encode driver result values without losing precision or binary identity.""" - if isinstance(value, Decimal): - return str(value) - if isinstance(value, (datetime, date, time)): - return value.isoformat() - if isinstance(value, UUID): - return str(value) - if isinstance(value, (bytes, bytearray, memoryview)): - return { - "encoding": "base64", - "data": b64encode(bytes(value)).decode("ascii"), - } - if isinstance(value, timedelta): - return str(value) - raise TypeError(f"Object of type {type(value).__name__} is not JSON serializable") - - -def _json_dumps(value: object) -> str: - """Serialize MCP payloads with the same driver-value rules as query results.""" - return json.dumps(value, ensure_ascii=False, default=_query_result_json_default) - - def _resolve_codegen_output_path(base_dir: Path, relative_path: object) -> Path: """Resolve a generated filename while keeping it inside its output directory.""" if not isinstance(relative_path, str) or not relative_path.strip(): @@ -750,7 +728,7 @@ async def handle_db_connect_test(arguments: Dict[str, Any]) -> List[TextContent] f"- Host: {config.host}:{config.port}\n" f"- Type: {config.type.value}\n\n" f"Use this connection_id for subsequent database operations.\n\n" - f"Raw Response: {json.dumps(response, ensure_ascii=False)}" + f"Raw Response: {_json_dumps(response, ensure_ascii=False)}" )] except DatabaseConnectionError as e: @@ -762,7 +740,7 @@ async def handle_db_connect_test(arguments: Dict[str, Any]) -> List[TextContent] } return [TextContent( type="text", - text=f"Database connection failed: {safe_error}\n\nRaw Response: {json.dumps(error_response, ensure_ascii=False)}" + text=f"Database connection failed: {safe_error}\n\nRaw Response: {_json_dumps(error_response, ensure_ascii=False)}" )] except Exception as e: @@ -775,7 +753,7 @@ async def handle_db_connect_test(arguments: Dict[str, Any]) -> List[TextContent] } return [TextContent( type="text", - text=f"Unexpected error: {safe_error}\n\nRaw Response: {json.dumps(error_response, ensure_ascii=False)}" + text=f"Unexpected error: {safe_error}\n\nRaw Response: {_json_dumps(error_response, ensure_ascii=False)}" )] @@ -817,7 +795,7 @@ async def handle_db_query_databases(arguments: Dict[str, Any]) -> List[TextConte } return [TextContent( type="text", - text=f"SQLite databases: [{config.database}]\n\nRaw Response: {json.dumps(response, ensure_ascii=False)}" + text=f"SQLite databases: [{config.database}]\n\nRaw Response: {_json_dumps(response, ensure_ascii=False)}" )] else: raise MCPServiceError(f"Listing databases not implemented for {config.type}") @@ -842,7 +820,7 @@ async def handle_db_query_databases(arguments: Dict[str, Any]) -> List[TextConte type="text", text=f"Found {len(databases)} databases:\n" + "\n".join(f"- {db}" for db in databases) + - f"\n\nRaw Response: {json.dumps(response, ensure_ascii=False)}" + f"\n\nRaw Response: {_json_dumps(response, ensure_ascii=False)}" )] except (DatabaseConnectionError, DatabaseQueryError) as e: @@ -853,7 +831,7 @@ async def handle_db_query_databases(arguments: Dict[str, Any]) -> List[TextConte } return [TextContent( type="text", - text=f"Failed to list databases: {str(e)}\n\nRaw Response: {json.dumps(error_response, ensure_ascii=False)}" + text=f"Failed to list databases: {str(e)}\n\nRaw Response: {_json_dumps(error_response, ensure_ascii=False)}" )] except Exception as e: @@ -865,7 +843,7 @@ async def handle_db_query_databases(arguments: Dict[str, Any]) -> List[TextConte } return [TextContent( type="text", - text=f"Unexpected error: {str(e)}\n\nRaw Response: {json.dumps(error_response, ensure_ascii=False)}" + text=f"Unexpected error: {str(e)}\n\nRaw Response: {_json_dumps(error_response, ensure_ascii=False)}" )] @@ -946,7 +924,7 @@ async def handle_db_query_tables(arguments: Dict[str, Any]) -> List[TextContent] text=f"Found {len(tables)} tables in database '{database}'" + (f" schema '{schema}'" if schema else "") + ":\n" + "\n".join(f"- {table}" for table in tables) + - f"\n\nRaw Response: {json.dumps(response, ensure_ascii=False)}" + f"\n\nRaw Response: {_json_dumps(response, ensure_ascii=False)}" )] except (DatabaseConnectionError, DatabaseQueryError) as e: @@ -957,7 +935,7 @@ async def handle_db_query_tables(arguments: Dict[str, Any]) -> List[TextContent] } return [TextContent( type="text", - text=f"Failed to list tables: {str(e)}\n\nRaw Response: {json.dumps(error_response, ensure_ascii=False)}" + text=f"Failed to list tables: {str(e)}\n\nRaw Response: {_json_dumps(error_response, ensure_ascii=False)}" )] except Exception as e: @@ -969,7 +947,7 @@ async def handle_db_query_tables(arguments: Dict[str, Any]) -> List[TextContent] } return [TextContent( type="text", - text=f"Unexpected error: {str(e)}\n\nRaw Response: {json.dumps(error_response, ensure_ascii=False)}" + text=f"Unexpected error: {str(e)}\n\nRaw Response: {_json_dumps(error_response, ensure_ascii=False)}" )] @@ -1044,7 +1022,7 @@ async def handle_db_query_table_exists(arguments: Dict[str, Any]) -> List[TextCo type="text", text=( f"Table '{table}' {status} in database '{database}'\n\n" - f"Raw Response: {json.dumps(response, ensure_ascii=False)}" + f"Raw Response: {_json_dumps(response, ensure_ascii=False)}" ) )] @@ -1056,7 +1034,7 @@ async def handle_db_query_table_exists(arguments: Dict[str, Any]) -> List[TextCo } return [TextContent( type="text", - text=f"Failed to check table existence: {str(e)}\n\nRaw Response: {json.dumps(error_response, ensure_ascii=False)}" + text=f"Failed to check table existence: {str(e)}\n\nRaw Response: {_json_dumps(error_response, ensure_ascii=False)}" )] except Exception as e: @@ -1068,7 +1046,7 @@ async def handle_db_query_table_exists(arguments: Dict[str, Any]) -> List[TextCo } return [TextContent( type="text", - text=f"Unexpected error: {str(e)}\n\nRaw Response: {json.dumps(error_response, ensure_ascii=False)}" + text=f"Unexpected error: {str(e)}\n\nRaw Response: {_json_dumps(error_response, ensure_ascii=False)}" )] @@ -1134,7 +1112,7 @@ async def handle_db_query_execute(arguments: Dict[str, Any]) -> List[TextContent } return [TextContent( type="text", - text=f"Failed to execute query: {str(e)}\n\nRaw Response: {json.dumps(error_response, ensure_ascii=False)}" + text=f"Failed to execute query: {str(e)}\n\nRaw Response: {_json_dumps(error_response, ensure_ascii=False)}" )] except Exception as e: @@ -1146,7 +1124,7 @@ async def handle_db_query_execute(arguments: Dict[str, Any]) -> List[TextContent } return [TextContent( type="text", - text=f"Unexpected error: {str(e)}\n\nRaw Response: {json.dumps(error_response, ensure_ascii=False)}" + text=f"Unexpected error: {str(e)}\n\nRaw Response: {_json_dumps(error_response, ensure_ascii=False)}" )] @@ -1323,7 +1301,7 @@ async def handle_db_table_describe(arguments: Dict[str, Any]) -> List[TextConten } return [TextContent( type="text", - text=f"Failed to describe table: {str(e)}\n\nRaw Response: {json.dumps(error_response, ensure_ascii=False)}" + text=f"Failed to describe table: {str(e)}\n\nRaw Response: {_json_dumps(error_response, ensure_ascii=False)}" )] except Exception as e: @@ -1335,7 +1313,7 @@ async def handle_db_table_describe(arguments: Dict[str, Any]) -> List[TextConten } return [TextContent( type="text", - text=f"Unexpected error: {str(e)}\n\nRaw Response: {json.dumps(error_response, ensure_ascii=False)}" + text=f"Unexpected error: {str(e)}\n\nRaw Response: {_json_dumps(error_response, ensure_ascii=False)}" )] @@ -1467,7 +1445,7 @@ async def handle_db_table_columns(arguments: Dict[str, Any]) -> List[TextContent } return [TextContent( type="text", - text=f"Failed to get column information: {str(e)}\n\nRaw Response: {json.dumps(error_response, ensure_ascii=False)}" + text=f"Failed to get column information: {str(e)}\n\nRaw Response: {_json_dumps(error_response, ensure_ascii=False)}" )] except Exception as e: @@ -1479,7 +1457,7 @@ async def handle_db_table_columns(arguments: Dict[str, Any]) -> List[TextContent } return [TextContent( type="text", - text=f"Unexpected error: {str(e)}\n\nRaw Response: {json.dumps(error_response, ensure_ascii=False)}" + text=f"Unexpected error: {str(e)}\n\nRaw Response: {_json_dumps(error_response, ensure_ascii=False)}" )] @@ -1577,7 +1555,7 @@ async def handle_db_table_primary_keys(arguments: Dict[str, Any]) -> List[TextCo } return [TextContent( type="text", - text=f"Failed to get primary keys: {str(e)}\n\nRaw Response: {json.dumps(error_response, ensure_ascii=False)}" + text=f"Failed to get primary keys: {str(e)}\n\nRaw Response: {_json_dumps(error_response, ensure_ascii=False)}" )] except Exception as e: @@ -1589,7 +1567,7 @@ async def handle_db_table_primary_keys(arguments: Dict[str, Any]) -> List[TextCo } return [TextContent( type="text", - text=f"Unexpected error: {str(e)}\n\nRaw Response: {json.dumps(error_response, ensure_ascii=False)}" + text=f"Unexpected error: {str(e)}\n\nRaw Response: {_json_dumps(error_response, ensure_ascii=False)}" )] @@ -1727,7 +1705,7 @@ async def handle_db_table_foreign_keys(arguments: Dict[str, Any]) -> List[TextCo } return [TextContent( type="text", - text=f"Failed to get foreign keys: {str(e)}\n\nRaw Response: {json.dumps(error_response, ensure_ascii=False)}" + text=f"Failed to get foreign keys: {str(e)}\n\nRaw Response: {_json_dumps(error_response, ensure_ascii=False)}" )] except Exception as e: @@ -1739,7 +1717,7 @@ async def handle_db_table_foreign_keys(arguments: Dict[str, Any]) -> List[TextCo } return [TextContent( type="text", - text=f"Unexpected error: {str(e)}\n\nRaw Response: {json.dumps(error_response, ensure_ascii=False)}" + text=f"Unexpected error: {str(e)}\n\nRaw Response: {_json_dumps(error_response, ensure_ascii=False)}" )] @@ -1891,7 +1869,7 @@ async def handle_db_table_indexes(arguments: Dict[str, Any]) -> List[TextContent } return [TextContent( type="text", - text=f"Failed to get table indexes: {str(e)}\n\nRaw Response: {json.dumps(error_response, ensure_ascii=False)}" + text=f"Failed to get table indexes: {str(e)}\n\nRaw Response: {_json_dumps(error_response, ensure_ascii=False)}" )] except Exception as e: @@ -1903,7 +1881,7 @@ async def handle_db_table_indexes(arguments: Dict[str, Any]) -> List[TextContent } return [TextContent( type="text", - text=f"Unexpected error: {str(e)}\n\nRaw Response: {json.dumps(error_response, ensure_ascii=False)}" + text=f"Unexpected error: {str(e)}\n\nRaw Response: {_json_dumps(error_response, ensure_ascii=False)}" )] @@ -2160,7 +2138,7 @@ async def handle_db_codegen_analyze(arguments: Dict[str, Any]) -> List[TextConte } return [TextContent( type="text", - text=f"Failed to analyze table for code generation: {str(e)}\n\nRaw Response: {json.dumps(error_response, ensure_ascii=False)}" + text=f"Failed to analyze table for code generation: {str(e)}\n\nRaw Response: {_json_dumps(error_response, ensure_ascii=False)}" )] except Exception as e: @@ -2172,7 +2150,7 @@ async def handle_db_codegen_analyze(arguments: Dict[str, Any]) -> List[TextConte } return [TextContent( type="text", - text=f"Unexpected error: {str(e)}\n\nRaw Response: {json.dumps(error_response, ensure_ascii=False)}" + text=f"Unexpected error: {str(e)}\n\nRaw Response: {_json_dumps(error_response, ensure_ascii=False)}" )] @@ -2593,7 +2571,7 @@ def with_suffix(kind: str) -> str: } return [TextContent( type="text", - text=f"Failed to generate code: {str(e)}\n\nRaw Response: {json.dumps(error_response, ensure_ascii=False)}" + text=f"Failed to generate code: {str(e)}\n\nRaw Response: {_json_dumps(error_response, ensure_ascii=False)}" )] except Exception as e: @@ -2605,7 +2583,7 @@ def with_suffix(kind: str) -> str: } return [TextContent( type="text", - text=f"Unexpected error: {str(e)}\n\nRaw Response: {json.dumps(error_response, ensure_ascii=False)}" + text=f"Unexpected error: {str(e)}\n\nRaw Response: {_json_dumps(error_response, ensure_ascii=False)}" )] @@ -2989,7 +2967,7 @@ async def handle_springboot_validate_project(arguments: Dict[str, Any]) -> List[ } return [TextContent( type="text", - text=f"Project validation failed: {str(e)}\n\nRaw Response: {json.dumps(error_response, ensure_ascii=False)}" + text=f"Project validation failed: {str(e)}\n\nRaw Response: {_json_dumps(error_response, ensure_ascii=False)}" )] @@ -3129,7 +3107,7 @@ async def handle_springboot_analyze_dependencies(arguments: Dict[str, Any]) -> L } return [TextContent( type="text", - text=f"Dependency analysis failed: {str(e)}\n\nRaw Response: {json.dumps(error_response, ensure_ascii=False)}" + text=f"Dependency analysis failed: {str(e)}\n\nRaw Response: {_json_dumps(error_response, ensure_ascii=False)}" )] @@ -3404,7 +3382,7 @@ def _get(d: Dict[str, Any], path: str, default=None): return [TextContent( type="text", text='\n'.join(text_lines) - + f"\n\nRaw Response: {json.dumps(response, ensure_ascii=False)}" + + f"\n\nRaw Response: {_json_dumps(response, ensure_ascii=False)}" )] except Exception as e: @@ -3415,6 +3393,6 @@ def _get(d: Dict[str, Any], path: str, default=None): } return [TextContent( type="text", - text=f"Failed to read Spring Boot config: {e}\n\nRaw Response: {json.dumps(err, ensure_ascii=False)}" + text=f"Failed to read Spring Boot config: {e}\n\nRaw Response: {_json_dumps(err, ensure_ascii=False)}" )] diff --git a/src/dbjavagenix/utils/json_serialization.py b/src/dbjavagenix/utils/json_serialization.py new file mode 100644 index 0000000..8c5c487 --- /dev/null +++ b/src/dbjavagenix/utils/json_serialization.py @@ -0,0 +1,33 @@ +"""JSON encoding for values returned by database drivers.""" + +from __future__ import annotations + +from base64 import b64encode +from datetime import date, datetime, time, timedelta +from decimal import Decimal +import json +from typing import Any +from uuid import UUID + + +def default(value: object) -> object: + """Encode common driver values without losing precision or binary identity.""" + if isinstance(value, Decimal): + return str(value) + if isinstance(value, (datetime, date, time)): + return value.isoformat() + if isinstance(value, UUID): + return str(value) + if isinstance(value, (bytes, bytearray, memoryview)): + return { + "encoding": "base64", + "data": b64encode(bytes(value)).decode("ascii"), + } + if isinstance(value, timedelta): + return str(value) + raise TypeError(f"Object of type {type(value).__name__} is not JSON serializable") + + +def dumps(value: Any, *, ensure_ascii: bool = False, indent: int | str | None = None) -> str: + """Serialize a payload using the shared MCP driver-value contract.""" + return json.dumps(value, ensure_ascii=ensure_ascii, indent=indent, default=default) diff --git a/tests/unit/test_json_serialization.py b/tests/unit/test_json_serialization.py new file mode 100644 index 0000000..c2bb358 --- /dev/null +++ b/tests/unit/test_json_serialization.py @@ -0,0 +1,100 @@ +"""Tests for the shared MCP JSON driver-value contract.""" + +from datetime import date, datetime, time, timedelta, timezone +from decimal import Decimal +import json +from types import SimpleNamespace +from uuid import UUID + +import pytest + +from dbjavagenix.core.models import DatabaseType +from dbjavagenix.database import atomic_codegen_tools +from dbjavagenix.utils.json_serialization import default, dumps + + +def test_dumps_preserves_supported_driver_values(): + payload = { + "amount": Decimal("12.30"), + "created_at": datetime(2026, 9, 8, 10, 0, tzinfo=timezone.utc), + "event_date": date(2026, 9, 8), + "event_time": time(10, 0, 1, 234_000), + "event_id": UUID("12345678-1234-5678-1234-567812345678"), + "payload": b"\x00\xff", + "mutable_payload": bytearray(b"ok"), + "view_payload": memoryview(b"view"), + "duration": timedelta(seconds=2), + } + + encoded = json.loads(dumps(payload)) + + assert encoded == { + "amount": "12.30", + "created_at": "2026-09-08T10:00:00+00:00", + "event_date": "2026-09-08", + "event_time": "10:00:01.234000", + "event_id": "12345678-1234-5678-1234-567812345678", + "payload": {"encoding": "base64", "data": "AP8="}, + "mutable_payload": {"encoding": "base64", "data": "b2s="}, + "view_payload": {"encoding": "base64", "data": "dmlldw=="}, + "duration": "0:00:02", + } + + +def test_default_rejects_unknown_objects(): + class UnknownDriverValue: + pass + + with pytest.raises(TypeError, match="UnknownDriverValue"): + default(UnknownDriverValue()) + + +@pytest.mark.asyncio +async def test_atomic_context_serializes_nested_driver_values(monkeypatch): + class Analyzer: + def __init__(self, _manager): + pass + + async def analyze_table_for_codegen(self, *_args, **_kwargs): + return { + "template_context": { + "packageSuffix": "", + "columns": [ + { + "defaultValue": Decimal("12.30"), + "observedAt": datetime(2026, 9, 8, tzinfo=timezone.utc), + "eventId": UUID("12345678-1234-5678-1234-567812345678"), + "payload": memoryview(b"view"), + } + ], + } + } + + monkeypatch.setattr( + "dbjavagenix.database.codegen_tools.CodegenAnalyzer", + Analyzer, + ) + monkeypatch.setattr( + atomic_codegen_tools.connection_manager, + "get_connection_info", + lambda _connection_id: SimpleNamespace( + type=DatabaseType.SQLITE, + database="app", + ), + ) + monkeypatch.setattr(atomic_codegen_tools, "_collect_all_table_names", lambda *_args: []) + + response = await atomic_codegen_tools.handle_codegen_build_context( + { + "connection_id": "sqlite-1", + "table_name": "events", + "template_category": "Default", + } + ) + + payload = json.loads(response[0].text) + column = payload["context"]["columns"][0] + assert column["defaultValue"] == "12.30" + assert column["observedAt"] == "2026-09-08T00:00:00+00:00" + assert column["eventId"] == "12345678-1234-5678-1234-567812345678" + assert column["payload"] == {"encoding": "base64", "data": "dmlldw=="}