feat(mcp): Add command-line options for MCP.

Author: Zhe Yu

docs/cli.md

@@ -636,6 +636,10 @@ is available in the shell as `vectorcode-mcp-server`, and make sure you're using
 a [standalone chromadb server](#chromadb) configured in the [JSON](#configuring-vectorcode) 
 via the `host` and `port` options. 
 
+The MCP server entry point (`vectorcode-mcp-server`) provides some CLI options
+that you can use to customise the default behaviour of the server. To view the
+supported options, run `vectorcode-mcp-server -h` in your shell.
+
 ### Writing Prompts
 
 If you want to integrate VectorCode in your LLM application, you may want to

src/vectorcode/mcp_main.py

@@ -1,7 +1,9 @@
+import argparse
 import asyncio
 import logging
 import os
 import sys
+from dataclasses import dataclass
 from pathlib import Path
 from typing import Optional
 
@@ -31,7 +33,33 @@
 from vectorcode.subcommands.query import get_query_result_files
 
 logger = logging.getLogger(name=__name__)
-mcp = FastMCP("VectorCode", instructions="\n".join(prompt_strings))
+
+
+@dataclass
+class MCPConfig:
+    n_results: int = 10
+    ls_on_start: bool = False
+
+
+mcp_config = MCPConfig()
+
+
+def get_arg_parser():
+    parser = argparse.ArgumentParser(prog="vectorcode-mcp-server")
+    parser.add_argument(
+        "--number",
+        "-n",
+        type=int,
+        default=10,
+        help="Default number of files to retrieve.",
+    )
+    parser.add_argument(
+        "--ls-on-start",
+        action="store_true",
+        default=False,
+        help="Whether to include the output of `vectorcode ls` in the tool description.",
+    )
+    return parser
 
 
 default_config: Optional[Config] = None
@@ -114,6 +142,7 @@ async def query_tool(
 
 async def mcp_server():
     global default_config, default_client, default_collection
+
     local_config_dir = await find_project_config_dir(".")
 
     if local_config_dir is not None:
@@ -128,9 +157,23 @@ async def mcp_server():
         try:
             default_collection = await get_collection(default_client, default_config)
             logger.info("Collection initialised for %s.", project_root)
-        except InvalidCollectionException:
+        except InvalidCollectionException:  # pragma: nocover
             default_collection = None
 
+    default_instructions = "\n".join(prompt_strings)
+    if default_client is None:
+        if mcp_config.ls_on_start:  # pragma: nocover
+            logger.warning(
+                "Failed to initialise a chromadb client. Ignoring --ls-on-start flag."
+            )
+    else:
+        if mcp_config.ls_on_start:
+            logger.info("Adding available collections to the server instructions.")
+            default_instructions += "\nYou have access to the following collections:\n"
+            for name in await list_collections():
+                default_instructions += f"<collection>{name}</collection>"
+
+    mcp = FastMCP("VectorCode", instructions=default_instructions)
     mcp.add_tool(
         fn=list_collections,
         name="ls",
@@ -140,24 +183,36 @@ async def mcp_server():
     mcp.add_tool(
         fn=query_tool,
         name="query",
-        description="""
+        description=f"""
 Use VectorCode to perform vector similarity search on repositories and return a list of relevant file paths and contents. 
 Make sure `project_root` is one of the values from the `ls` tool. 
+Unless the user requested otherwise, start your retrievals by {mcp_config.n_results} files.
 The result contains the relative paths for the files and their corresponding contents.
 """,
     )
 
     return mcp
 
 
+def parse_cli_args(args: Optional[list[str]] = None) -> MCPConfig:
+    parser = get_arg_parser()
+    parsed_args = parser.parse_args(args or sys.argv)
+    return MCPConfig(n_results=parsed_args.number, ls_on_start=parsed_args.ls_on_start)
+
+
 async def run_server():  # pragma: nocover
     mcp = await mcp_server()
     await mcp.run_stdio_async()
     return 0
 
 
 def main():  # pragma: nocover
+    global mcp_config
     config_logging("vectorcode-mcp-server", stdio=False)
+    mcp_config = parse_cli_args()
+    assert mcp_config.n_results > 0 and mcp_config.n_results % 1 == 0, (
+        "--number must be used with a positive integer!"
+    )
     return asyncio.run(run_server())
 
 

tests/test_mcp.py

@@ -1,10 +1,17 @@
+from argparse import ArgumentParser
 from unittest.mock import AsyncMock, MagicMock, patch
 
 import pytest
 from mcp import McpError
 
 from vectorcode.cli_utils import Config
-from vectorcode.mcp_main import list_collections, mcp_server, query_tool
+from vectorcode.mcp_main import (
+    get_arg_parser,
+    list_collections,
+    mcp_server,
+    parse_cli_args,
+    query_tool,
+)
 
 
 @pytest.mark.asyncio
@@ -44,11 +51,11 @@ async def test_list_collections_no_metadata():
         mock_collection2 = AsyncMock()
         mock_collection2.metadata = None
 
-        async def async_generator():
+        async def async_generator(cli):
             yield mock_collection1
             yield mock_collection2
 
-        mock_get_collections.return_value = async_generator()
+        mock_get_collections.side_effect = async_generator
 
         result = await list_collections()
         assert result == ["path1"]
@@ -182,3 +189,50 @@ async def test_mcp_server():
         await mcp_server()
 
         assert mock_add_tool.call_count == 2
+
+
+@pytest.mark.asyncio
+async def test_mcp_server_ls_on_start():
+    with (
+        patch(
+            "vectorcode.mcp_main.find_project_config_dir"
+        ) as mock_find_project_config_dir,
+        patch("vectorcode.mcp_main.load_config_file") as mock_load_config_file,
+        patch("vectorcode.mcp_main.get_client") as mock_get_client,
+        patch("vectorcode.mcp_main.get_collection") as mock_get_collection,
+        patch(
+            "vectorcode.mcp_main.get_collections", spec=AsyncMock
+        ) as mock_get_collections,
+        patch("mcp.server.fastmcp.FastMCP.add_tool") as mock_add_tool,
+    ):
+        from vectorcode.mcp_main import mcp_config
+
+        mcp_config.ls_on_start = True
+        mock_find_project_config_dir.return_value = "/path/to/config"
+        mock_load_config_file.return_value = Config(project_root="/path/to/project")
+        mock_client = AsyncMock()
+        mock_get_client.return_value = mock_client
+        mock_collection = AsyncMock()
+        mock_collection.metadata = {"path": "/path/to/project"}
+        mock_get_collection.return_value = mock_collection
+
+        async def new_get_collections(clients):
+            yield mock_collection
+
+        mock_get_collections.side_effect = new_get_collections
+
+        await mcp_server()
+
+        assert mock_add_tool.call_count == 2
+        mock_get_collections.assert_called()
+
+
+def test_arg_parser():
+    assert isinstance(get_arg_parser(), ArgumentParser)
+
+
+def test_args_parsing():
+    args = ["--number", "15", "--ls-on-start"]
+    parsed = parse_cli_args(args)
+    assert parsed.n_results == 15
+    assert parsed.ls_on_start