feat: refactor server tools and add health checks

.gitignore

@@ -30,3 +30,4 @@ gateway/gateway
 .DS_Store
 *.tmp
 *.log
+config/config.yaml

gateway/internal/mcp/proxy_test.go

@@ -14,17 +14,22 @@ import (
 	"github.com/endemics/limnos/gateway/internal/queue"
 )
 
-// emptyPool returns a zero-value WorkerPool. Next() always returns false
-// because the workers slice is nil (len == 0).
-func emptyPool() *queue.WorkerPool {
-	return &queue.WorkerPool{}
-}
-
 // silentLogger discards all log output during tests.
 func silentLogger() *slog.Logger {
 	return slog.New(slog.NewTextHandler(io.Discard, &slog.HandlerOptions{Level: slog.LevelError + 1}))
 }
 
+func dummyAuth() *auth.APIKeyAuth {
+	return auth.NewAPIKeyAuth(auth.APIKeyAuthConfig{})
+}
+
+func emptyPool() *queue.WorkerPool {
+	pool, _ := queue.NewWorkerPool(queue.WorkerPoolConfig{
+		Size: 0,
+	})
+	return pool
+}
+
 // ── No healthy workers ─────────────────────────────────────────────────────────
 
 func TestProxy_NoWorkers_Returns503(t *testing.T) {

server/catalog/iceberg.py

@@ -190,5 +190,8 @@ def _parse_s3_path(s3_path: str) -> Tuple[str, str]:
     without_scheme = s3_path.removeprefix("s3://")
     parts = without_scheme.split("/", 1)
     bucket = parts[0]
-    prefix = parts[1].rstrip("/") + "/" if len(parts) > 1 else ""
+    if len(parts) > 1 and parts[1].rstrip("/"):
+        prefix = parts[1].rstrip("/") + "/"
+    else:
+        prefix = ""
     return bucket, prefix

server/main.py

@@ -3,8 +3,8 @@
 
 Usage:
     python main.py                          # stdio (Claude Desktop)
-    python main.py --transport http         # HTTP on port 8000
-    python main.py --transport http --port 9000
+    python main.py --transport sse          # SSE on port 8000
+    python main.py --transport sse --port 9000
 """
 
 import argparse
@@ -23,9 +23,7 @@
     list_datasets,
     describe_table,
     sample_data,
-    estimate_query,
     query,
-    refresh_schema,
 )
 from config import load_config
 
@@ -85,8 +83,8 @@ async def app_lifespan(server: FastMCP):
 Workflow for answering data questions:
 1. Call datalake_list_datasets to discover available tables.
 2. Call datalake_describe_table on the relevant table(s) to understand schema and partitions.
-3. Call datalake_estimate_query to preview cost before executing.
-4. Call datalake_query with your natural language question or SQL.
+3. Call datalake_query with explain_only=true to preview cost before executing.
+4. Call datalake_query with your natural language question or SQL to execute.
 
 Always filter on partition columns when possible — this dramatically reduces cost and latency.
 If a query is estimated to be expensive, explain the cost to the user and ask for confirmation.
@@ -98,9 +96,7 @@ async def app_lifespan(server: FastMCP):
 list_datasets.register(mcp)
 describe_table.register(mcp)
 sample_data.register(mcp)
-estimate_query.register(mcp)
 query.register(mcp)
-refresh_schema.register(mcp)
 
 
 # ─── CLI ──────────────────────────────────────────────────────────────────────
@@ -110,22 +106,34 @@ def main():
     parser = argparse.ArgumentParser(description="Limnos MCP Server")
     parser.add_argument(
         "--transport",
-        choices=["stdio", "http"],
+        choices=["stdio", "sse"],
         default="stdio",
         help="Transport to use (default: stdio for Claude Desktop)",
     )
     parser.add_argument(
         "--port",
         type=int,
         default=8000,
-        help="Port for HTTP transport (default: 8000)",
+        help="Port for SSE transport (default: 8000)",
     )
     args = parser.parse_args()
 
     if args.transport == "stdio":
         mcp.run(transport="stdio")
     else:
-        mcp.run(transport="streamable_http", port=args.port)
+        # Get the underlying FastAPI/Starlette app
+        app = mcp.sse_app()
+
+        # Add health check for the Go gateway
+        from starlette.responses import Response
+
+        @app.route("/health")
+        async def health(request):
+            return Response(status_code=200)
+
+        import uvicorn
+
+        uvicorn.run(app, host="0.0.0.0", port=args.port)
 
 
 if __name__ == "__main__":

server/tests/test_iceberg_reader.py

@@ -0,0 +1,160 @@
+"""Tests for catalog.iceberg — Iceberg metadata reader."""
+
+from __future__ import annotations
+
+import json
+from io import BytesIO
+from unittest.mock import MagicMock
+
+import pytest
+
+from catalog.iceberg import (
+    read_iceberg_metadata,
+    _parse_s3_path,
+    _read_version_hint,
+    _parse_schema,
+    _parse_partition_spec,
+)
+
+
+def test_parse_s3_path():
+    assert _parse_s3_path("s3://my-bucket/path/to/table") == (
+        "my-bucket",
+        "path/to/table/",
+    )
+    assert _parse_s3_path("s3://my-bucket/") == ("my-bucket", "")
+    assert _parse_s3_path("s3://my-bucket") == ("my-bucket", "")
+
+
+def test_read_version_hint_from_text_file():
+    s3 = MagicMock()
+    s3.get_object.return_value = {"Body": BytesIO(b"5\n")}
+
+    version = _read_version_hint(s3, "bucket", "prefix/")
+    assert version == 5
+    s3.get_object.assert_called_with(Bucket="bucket", Key="prefix/version-hint.text")
+
+
+def test_read_version_hint_from_listing_fallback():
+    s3 = MagicMock()
+    # version-hint.text missing
+    s3.get_object.side_effect = Exception("404")
+    # v1, v2, v3 metadata files exist
+    s3.list_objects_v2.return_value = {
+        "Contents": [
+            {"Key": "prefix/v1.metadata.json"},
+            {"Key": "prefix/v2.metadata.json"},
+            {"Key": "prefix/v3.metadata.json"},
+            {"Key": "prefix/not-a-meta-file.txt"},
+        ]
+    }
+
+    version = _read_version_hint(s3, "bucket", "prefix/")
+    assert version == 3
+
+
+def test_read_version_hint_not_found_raises():
+    s3 = MagicMock()
+    s3.get_object.side_effect = Exception("404")
+    s3.list_objects_v2.return_value = {"Contents": []}
+
+    with pytest.raises(FileNotFoundError, match="No Iceberg metadata found"):
+        _read_version_hint(s3, "bucket", "prefix/")
+
+
+def test_parse_schema_basic():
+    schema_json = {
+        "fields": [
+            {"id": 1, "name": "id", "type": "int", "required": True},
+            {"id": 2, "name": "data", "type": "string", "required": False},
+            {"id": 3, "name": "ts", "type": "timestamp", "required": True},
+        ]
+    }
+    columns = _parse_schema(schema_json)
+
+    assert len(columns) == 3
+    assert columns[0].name == "id"
+    assert columns[0].dtype == "INTEGER"
+    assert columns[0].required is True
+    assert columns[1].name == "data"
+    assert columns[1].dtype == "VARCHAR"
+    assert columns[2].dtype == "TIMESTAMP"
+
+
+def test_parse_partition_spec():
+    spec_json = {
+        "fields": [
+            {
+                "source-id": 1,
+                "field-id": 1000,
+                "name": "id_bucket",
+                "transform": "bucket[16]",
+            },
+            {"source-id": 3, "field-id": 1001, "name": "ts_day", "transform": "day"},
+        ]
+    }
+    parts = _parse_partition_spec(spec_json)
+
+    assert len(parts) == 2
+    assert parts[0].name == "id_bucket"
+    assert parts[0].transform == "bucket[16]"
+    assert parts[1].name == "ts_day"
+    assert parts[1].transform == "day"
+
+
+def test_read_iceberg_metadata_full_flow():
+    s3 = MagicMock()
+
+    # 1. Mock version hint
+    s3.get_object.side_effect = [
+        {"Body": BytesIO(b"1")},  # version-hint.text
+        {
+            "Body": BytesIO(
+                json.dumps(
+                    {
+                        "table-uuid": "test-uuid",
+                        "format-version": 2,
+                        "location": "s3://bucket/table/",
+                        "current-schema-id": 0,
+                        "schemas": [
+                            {
+                                "schema-id": 0,
+                                "fields": [
+                                    {"id": 1, "name": "id", "type": "int"},
+                                    {"id": 2, "name": "val", "type": "float"},
+                                ],
+                            }
+                        ],
+                        "default-spec-id": 0,
+                        "partition-specs": [{"spec-id": 0, "fields": []}],
+                        "current-snapshot-id": 123,
+                        "snapshots": [
+                            {
+                                "snapshot-id": 123,
+                                "timestamp-ms": 1640995200000,
+                                "manifest-list": "s3://bucket/table/metadata/snap-123.avro",
+                                "summary": {
+                                    "total-data-files": "5",
+                                    "total-records": "1000",
+                                    "total-files-size": "1048576",
+                                },
+                            }
+                        ],
+                    }
+                ).encode()
+            )
+        },  # v1.metadata.json
+    ]
+
+    meta = read_iceberg_metadata("s3://bucket/table/", s3_client=s3)
+
+    assert meta.table_uuid == "test-uuid"
+    assert len(meta.schema_columns) == 2
+    assert meta.total_files == 5
+    assert meta.total_rows == 1000
+    assert meta.total_bytes == 1048576
+    assert meta.current_snapshot.snapshot_id == 123
+    assert (
+        meta.current_snapshot.manifest_list
+        == "s3://bucket/table/metadata/snap-123.avro"
+    )

server/tools/__init__.py

@@ -0,0 +1,15 @@
+"""
+Limnos MCP tools package.
+"""
+
+from __future__ import annotations
+
+# Tool submodules (explicit re-export for linting)
+from . import list_datasets as list_datasets
+from . import describe_table as describe_table
+from . import query as query
+from . import sample_data as sample_data
+
+# Re-export formatting helpers (explicit re-export for linting)
+from .formatting import format_table as format_table
+from .formatting import format_query_result as format_query_result

server/tools/describe_table.py

@@ -215,23 +215,6 @@ async def _scan_metadata(
             bytes_per_row_estimate=bpr,
         )
 
-        # Auto-provision Glue external table so Athena fallback works.
-        # TXT is excluded — Athena has no useful capability over unstructured text.
-        if table_cfg.format != "txt" and config is not None:
-            try:
-                await asyncio.to_thread(
-                    GlueProvisioner(config).sync_table,
-                    table_cfg,
-                    columns,
-                    partition_cols,
-                )
-            except Exception:
-                logger.warning(
-                    "glue_provision_failed",
-                    table=table_cfg.name,
-                    exc_info=True,
-                )
-
     else:
         # Parquet — read schema via DuckDB, discover partitions via S3
         raw_schema = await asyncio.to_thread(
@@ -278,6 +261,23 @@ async def _scan_metadata(
             description=table_cfg.description,
         )
 
+    # Auto-provision Glue external table so Athena fallback works.
+    # Exclude Iceberg (has own catalog) and TXT (no Athena support).
+    if table_cfg.format not in ("iceberg", "txt") and config is not None:
+        try:
+            await asyncio.to_thread(
+                GlueProvisioner(config).sync_table,
+                table_cfg,
+                columns,
+                partition_cols,
+            )
+        except Exception:
+            logger.warning(
+                "glue_provision_failed",
+                table=table_cfg.name,
+                exc_info=True,
+            )
+
     cache.upsert(meta)
     return meta
 

server/tools/formatting.py

@@ -0,0 +1,45 @@
+"""
+Formatting helpers for the Limnos MCP server.
+"""
+
+from __future__ import annotations
+
+from typing import Any, Dict, List
+
+
+def format_table(rows: List[Dict[str, Any]], columns: List[str]) -> str:
+    """Format a list of dictionaries as a Markdown table."""
+    if not rows:
+        return "_No results found._"
+
+    # Header
+    header = "| " + " | ".join(columns) + " |"
+    separator = "| " + " | ".join(["---"] * len(columns)) + " |"
+
+    # Rows
+    md_rows = []
+    for row in rows:
+        md_row = "| " + " | ".join(str(row.get(col, "")) for col in columns) + " |"
+        md_rows.append(md_row)
+
+    return "\n".join([header, separator] + md_rows)
+
+
+def format_query_result(result: Any, summary: str) -> str:
+    """Format a QueryResult object as a Markdown response with metadata."""
+    table_md = format_table(result.rows, result.columns)
+
+    response = (
+        f"### Query Results\n\n"
+        f"**{summary}**\n\n"
+        f"| Metric | Value |\n"
+        f"|--------|-------|\n"
+        f"| Engine | {result.engine} |\n"
+        f"| Rows | {result.row_count}{' (truncated)' if result.truncated else ''} |\n"
+        f"| Duration | {result.duration_ms}ms |\n"
+        f"| Bytes scanned | {result.bytes_scanned if result.bytes_scanned >= 0 else 'unknown'} |\n\n"
+        f"{table_md}\n\n"
+        f"**SQL executed:**\n```sql\n{result.sql_executed}\n```"
+    )
+
+    return response