diff --git a/README.md b/README.md index 4b4d0cb..2515f98 100644 --- a/README.md +++ b/README.md @@ -129,7 +129,18 @@ uv run ruff check . uv run ruff format . # Build package -uv build +uv build# Clone the repo +git clone https://github.com/kapillamba4/code-memory.git +cd code-memory + +# Install dependencies +uv sync + +# Run the MCP server (stdio transport) +uv run mcp run server.py + +# Run the MCP server (SSE transport for shared usage) +uv run python server.py --transport sse # Build standalone binary (requires pyinstaller) pip install pyinstaller @@ -228,8 +239,77 @@ For Windows: } ``` +## Shared SSE Server (Reduce Memory Usage) + +By default, each MCP host project launches its own `code-memory` process, which loads the embedding model (~1–2 GB) once per project. To avoid this, you can run a **single shared instance** over SSE (Server-Sent Events) and point all your MCP hosts at it. + +### Start the shared server + +```bash +# Using uvx (recommended) +uvx code-memory --transport sse + +# Custom port and host +uvx code-memory --transport sse --port 8765 --host 127.0.0.1 + +# Using standalone binary +./code-memory-linux-x86_64 --transport sse +``` + +The server listens on `http://127.0.0.1:8765/sse` by default. + +### Configure MCP hosts to use the shared server + +Instead of launching a new process, point your MCP host at the running SSE endpoint. + +#### Claude Desktop + +```json +{ + "mcpServers": { + "code-memory": { + "url": "http://127.0.0.1:8765/sse" + } + } +} +``` + +#### VS Code (Copilot / Continue) + +```json +{ + "servers": { + "code-memory": { + "url": "http://127.0.0.1:8765/sse" + } + } +} +``` + +#### Claude Code (CLI) — `.mcp.json` + +```json +{ + "mcpServers": { + "code-memory": { + "url": "http://127.0.0.1:8765/sse" + } + } +} +``` + +> **Tip:** Configure `uvx code-memory --transport sse` to start via a single-instance service manager (e.g. systemd user service, launchd agent, or another one-time login/startup mechanism) so the shared server starts automatically. + ## Configuration +### CLI Options + +| Option | Description | Default | +|--------|-------------|---------| +| `--transport` | Transport protocol: `stdio` or `sse` | `stdio` | +| `--port` | Port for SSE transport (only when `--transport sse` is used) | `8765` | +| `--host` | Host/bind address for SSE transport (only when `--transport sse` is used) | `127.0.0.1` | + ### Environment Variables | Variable | Description | Default | diff --git a/server.py b/server.py index 0eac9e8..0d1ac74 100644 --- a/server.py +++ b/server.py @@ -772,10 +772,44 @@ def search_history( # ── Entrypoint ──────────────────────────────────────────────────────────── +def build_arg_parser() -> "argparse.ArgumentParser": + """Build and return the CLI argument parser for code-memory.""" + import argparse + + parser = argparse.ArgumentParser( + prog="code-memory", + description="code-memory MCP server", + ) + parser.add_argument( + "--transport", + choices=["stdio", "sse"], + default="stdio", + help="Transport protocol to use (default: stdio). Use 'sse' to run a shared HTTP server.", + ) + parser.add_argument( + "--port", + type=int, + default=8765, + help="Port for SSE transport (default: 8765). Only used when --transport=sse.", + ) + parser.add_argument( + "--host", + default="127.0.0.1", + help="Host for SSE transport (default: 127.0.0.1). Only used when --transport=sse.", + ) + return parser + + def main(): """Entry point for the MCP server when installed as a package.""" + args = build_arg_parser().parse_args() + + if args.transport == "sse": + mcp.settings.host = args.host + mcp.settings.port = args.port + # Warmup is now done lazily on first index_codebase call - mcp.run() + mcp.run(transport=args.transport) if __name__ == "__main__": diff --git a/tests/test_main.py b/tests/test_main.py new file mode 100644 index 0000000..25eb961 --- /dev/null +++ b/tests/test_main.py @@ -0,0 +1,107 @@ +"""Tests for the main() entrypoint CLI argument parsing.""" + +from __future__ import annotations + +import argparse +from unittest.mock import MagicMock, patch + +import pytest + +from server import build_arg_parser + + +class TestMainArgParsing: + """Tests for CLI argument parsing in main().""" + + def _parse_args(self, argv: list[str]) -> argparse.Namespace: + """Parse args using the real CLI parser from server.py.""" + return build_arg_parser().parse_args(argv) + + def test_default_transport_is_stdio(self): + """Default transport should be stdio.""" + args = self._parse_args([]) + assert args.transport == "stdio" + + def test_default_port_is_8765(self): + """Default SSE port should be 8765.""" + args = self._parse_args([]) + assert args.port == 8765 + + def test_default_host_is_localhost(self): + """Default SSE host should be 127.0.0.1.""" + args = self._parse_args([]) + assert args.host == "127.0.0.1" + + def test_sse_transport_flag(self): + """--transport sse should set transport to sse.""" + args = self._parse_args(["--transport", "sse"]) + assert args.transport == "sse" + + def test_custom_port(self): + """--port should accept a custom port number.""" + args = self._parse_args(["--transport", "sse", "--port", "9000"]) + assert args.port == 9000 + + def test_custom_host(self): + """--host should accept a custom host.""" + args = self._parse_args(["--transport", "sse", "--host", "0.0.0.0"]) + assert args.host == "0.0.0.0" + + def test_invalid_transport_raises_error(self): + """Invalid transport should raise SystemExit.""" + with pytest.raises(SystemExit): + self._parse_args(["--transport", "invalid"]) + + +class TestMainRunsBehavior: + """Tests that main() calls mcp.run() with the correct arguments.""" + + @pytest.fixture(autouse=True) + def require_server(self): + """Skip tests when server module dependencies are unavailable.""" + pytest.importorskip("sqlite_vec", reason="sqlite_vec not installed") + + @pytest.fixture + def server_main(self): + """Import and return the server.main function.""" + import importlib + + import server as server_mod + importlib.reload(server_mod) + return server_mod.main + + def test_stdio_calls_run_with_stdio(self, server_main): + """main() with no args should call mcp.run(transport='stdio').""" + import server as server_mod + with patch("sys.argv", ["code-memory"]): + with patch.object(server_mod, "mcp") as mock_mcp: + mock_mcp.settings = MagicMock() + server_main() + mock_mcp.run.assert_called_once_with(transport="stdio") + + def test_sse_calls_run_with_sse(self, server_main): + """main() with --transport sse should call mcp.run(transport='sse').""" + import server as server_mod + with patch("sys.argv", ["code-memory", "--transport", "sse"]): + with patch.object(server_mod, "mcp") as mock_mcp: + mock_mcp.settings = MagicMock() + server_main() + mock_mcp.run.assert_called_once_with(transport="sse") + + def test_sse_sets_port_on_settings(self, server_main): + """main() with --transport sse --port 9000 should set mcp.settings.port.""" + import server as server_mod + with patch("sys.argv", ["code-memory", "--transport", "sse", "--port", "9000"]): + with patch.object(server_mod, "mcp") as mock_mcp: + mock_mcp.settings = MagicMock() + server_main() + assert mock_mcp.settings.port == 9000 + + def test_sse_sets_host_on_settings(self, server_main): + """main() with --transport sse --host 0.0.0.0 should set mcp.settings.host.""" + import server as server_mod + with patch("sys.argv", ["code-memory", "--transport", "sse", "--host", "0.0.0.0"]): + with patch.object(server_mod, "mcp") as mock_mcp: + mock_mcp.settings = MagicMock() + server_main() + assert mock_mcp.settings.host == "0.0.0.0" diff --git a/uv.lock b/uv.lock index e521bb7..362b515 100644 --- a/uv.lock +++ b/uv.lock @@ -109,7 +109,7 @@ wheels = [ [[package]] name = "code-memory" -version = "1.0.20" +version = "1.0.28" source = { editable = "." } dependencies = [ { name = "einops" }, @@ -1638,6 +1638,10 @@ dependencies = [ ] wheels = [ { url = "https://files.pythonhosted.org/packages/ec/23/2c9fe0c9c27f7f6cb865abcea8a4568f29f00acaeadfc6a37f6801f84cb4/torch-2.10.0-2-cp313-none-macosx_11_0_arm64.whl", hash = "sha256:e521c9f030a3774ed770a9c011751fb47c4d12029a3d6522116e48431f2ff89e", size = 79498254, upload-time = "2026-02-10T21:44:44.095Z" }, + { url = "https://files.pythonhosted.org/packages/ab/c6/4dfe238342ffdcec5aef1c96c457548762d33c40b45a1ab7033bb26d2ff2/torch-2.10.0-3-cp313-cp313-manylinux_2_28_x86_64.whl", hash = "sha256:80b1b5bfe38eb0e9f5ff09f206dcac0a87aadd084230d4a36eea5ec5232c115b", size = 915627275, upload-time = "2026-03-11T14:16:11.325Z" }, + { url = "https://files.pythonhosted.org/packages/d8/f0/72bf18847f58f877a6a8acf60614b14935e2f156d942483af1ffc081aea0/torch-2.10.0-3-cp313-cp313t-manylinux_2_28_x86_64.whl", hash = "sha256:46b3574d93a2a8134b3f5475cfb98e2eb46771794c57015f6ad1fb795ec25e49", size = 915523474, upload-time = "2026-03-11T14:17:44.422Z" }, + { url = "https://files.pythonhosted.org/packages/f4/39/590742415c3030551944edc2ddc273ea1fdfe8ffb2780992e824f1ebee98/torch-2.10.0-3-cp314-cp314-manylinux_2_28_x86_64.whl", hash = "sha256:b1d5e2aba4eb7f8e87fbe04f86442887f9167a35f092afe4c237dfcaaef6e328", size = 915632474, upload-time = "2026-03-11T14:15:13.666Z" }, + { url = "https://files.pythonhosted.org/packages/b6/8e/34949484f764dde5b222b7fe3fede43e4a6f0da9d7f8c370bb617d629ee2/torch-2.10.0-3-cp314-cp314t-manylinux_2_28_x86_64.whl", hash = "sha256:0228d20b06701c05a8f978357f657817a4a63984b0c90745def81c18aedfa591", size = 915523882, upload-time = "2026-03-11T14:14:46.311Z" }, { url = "https://files.pythonhosted.org/packages/c9/6f/f2e91e34e3fcba2e3fc8d8f74e7d6c22e74e480bbd1db7bc8900fdf3e95c/torch-2.10.0-cp313-cp313-manylinux_2_28_aarch64.whl", hash = "sha256:5c4d217b14741e40776dd7074d9006fd28b8a97ef5654db959d8635b2fe5f29b", size = 146004247, upload-time = "2026-01-21T16:24:29.335Z" }, { url = "https://files.pythonhosted.org/packages/98/fb/5160261aeb5e1ee12ee95fe599d0541f7c976c3701d607d8fc29e623229f/torch-2.10.0-cp313-cp313-manylinux_2_28_x86_64.whl", hash = "sha256:6b71486353fce0f9714ca0c9ef1c850a2ae766b409808acd58e9678a3edb7738", size = 915716445, upload-time = "2026-01-21T16:22:45.353Z" }, { url = "https://files.pythonhosted.org/packages/6a/16/502fb1b41e6d868e8deb5b0e3ae926bbb36dab8ceb0d1b769b266ad7b0c3/torch-2.10.0-cp313-cp313-win_amd64.whl", hash = "sha256:c2ee399c644dc92ef7bc0d4f7e74b5360c37cdbe7c5ba11318dda49ffac2bc57", size = 113757050, upload-time = "2026-01-21T16:24:19.204Z" },