Add support of partial results in MCP server. Make WM clients identifiable, add client/initialize request in WM protocol. Store runner logs in separate directory. Fix LSP start with tcp. Fix per-project requests in LSP server (use project path instead of project name).

Aksem · Aksem · commit 1102c9b3416f · 2026-03-18T21:15:32.000+01:00
diff --git a/finecode_extension_runner/src/finecode_extension_runner/cli.py b/finecode_extension_runner/src/finecode_extension_runner/cli.py
@@ -60,6 +60,7 @@ def start(
         / ".venvs"
         / env_name
         / "logs"
+        / "runner"
         / "runner.log")
     
     logs.setup_logging(log_level="INFO" if trace is False else "TRACE", log_file_path=log_file_path)
diff --git a/finecode_extension_runner/src/finecode_extension_runner/lsp_server.py b/finecode_extension_runner/src/finecode_extension_runner/lsp_server.py
@@ -192,6 +192,38 @@ async def tcp_server(h: str, p: int):
         except asyncio.CancelledError:
             logger.debug("Server was cancelled")
 
+    async def start_tcp_async(self, host: str, port: int) -> None:
+        """Starts TCP server from within an existing event loop."""
+        logger.info("Starting TCP server on %s:%s", host, port)
+
+        self._stop_event = stop_event = threading.Event()
+
+        async def lsp_connection(
+            reader: asyncio.StreamReader, writer: asyncio.StreamWriter
+        ):
+            logger.debug("Connected to client")
+            self.protocol.set_writer(writer)  # type: ignore
+            await run_async(
+                stop_event=stop_event,
+                reader=reader,
+                protocol=self.protocol,
+                logger=logger,
+                error_handler=self.report_server_error,
+            )
+            logger.debug("Main loop finished")
+            self.shutdown()
+
+        self._server = await asyncio.start_server(lsp_connection, host, port)
+
+        addrs = ", ".join(str(sock.getsockname()) for sock in self._server.sockets)
+        logger.info(f"Serving on {addrs}")
+
+        try:
+            async with self._server:
+                await self._server.serve_forever()
+        finally:
+            await self._finecode_exit_stack.aclose()
+
 
 
 def file_editor_file_change_to_lsp_text_edit(file_change: ifileeditor.FileChange) -> types.TextEdit:
diff --git a/src/finecode/lsp_server/endpoints/diagnostics.py b/src/finecode/lsp_server/endpoints/diagnostics.py
@@ -12,8 +12,8 @@
 from finecode_extension_api.actions import lint as lint_action
 
 
-async def _find_project_name_for_file(file_path: Path) -> str | None:
-    """Return the project name containing *file_path*.
+async def _find_project_dir_for_file(file_path: Path) -> str | None:
+    """Return the absolute directory path of the project containing *file_path*.
 
     This helper delegates the lookup to the WM server via
     ``workspace/findProjectForFile``; the server applies the same logic that
@@ -70,15 +70,15 @@ async def document_diagnostic_with_full_result(
         logger.error("Diagnostics requested but WM client not connected")
         return None
 
-    project_name = await _find_project_name_for_file(file_path)
-    if project_name is None:
+    project_dir = await _find_project_dir_for_file(file_path)
+    if project_dir is None:
         logger.error(f"Cannot determine project for diagnostics: {file_path}")
         return None
 
     try:
         response = await global_state.wm_client.run_action(
             action="lint",
-            project=project_name,
+            project=project_dir,
             params={
                 "target": "files",
                 "file_paths": [str(file_path)],
@@ -142,8 +142,8 @@ async def document_diagnostic_with_partial_results(
         logger.error("Diagnostics requested but WM client not connected")
         return None
 
-    project_name = await _find_project_name_for_file(file_path)
-    if project_name is None:
+    project_dir = await _find_project_dir_for_file(file_path)
+    if project_dir is None:
         logger.error(f"Cannot determine project for diagnostics: {file_path}")
         return None
 
@@ -155,7 +155,7 @@ async def document_diagnostic_with_partial_results(
             "actions/runWithPartialResults",
             {
                 "action": "lint",
-                "project": project_name,
+                "project": project_dir,
                 "params": {"file_paths": [str(file_path)]},
                 "partialResultToken": partial_result_token,
                 "options": {"resultFormats": ["json"], "trigger": "system", "devEnv": "ide"},
diff --git a/src/finecode/lsp_server/endpoints/formatting.py b/src/finecode/lsp_server/endpoints/formatting.py
@@ -22,15 +22,15 @@ async def format_document(ls: LanguageServer, params: types.DocumentFormattingPa
         logger.error("Formatting requested but WM client not connected")
         return None
 
-    project_name = await global_state.wm_client.find_project_for_file(str(file_path))
-    if project_name is None:
+    project_dir = await global_state.wm_client.find_project_for_file(str(file_path))
+    if project_dir is None:
         logger.error(f"Cannot determine project for formatting: {file_path}")
         return []
 
     try:
         response = await global_state.wm_client.run_action(
             action="format",
-            project=project_name,
+            project=project_dir,
             params={"file_paths": [str(file_path)], "save": False, "target": "files"},
             options={"trigger": "user", "devEnv": "ide"},
         )
diff --git a/src/finecode/lsp_server/endpoints/inlay_hints.py b/src/finecode/lsp_server/endpoints/inlay_hints.py
@@ -53,15 +53,15 @@ async def document_inlay_hint(
         logger.error("Inlay hints requested but WM client not connected")
         return None
 
-    project_name = await global_state.wm_client.find_project_for_file(str(file_path))
-    if project_name is None:
+    project_dir = await global_state.wm_client.find_project_for_file(str(file_path))
+    if project_dir is None:
         # Not all files belong to a project with this action — not an error.
         return []
 
     try:
         response = await global_state.wm_client.run_action(
             action="text_document_inlay_hint",
-            project=project_name,
+            project=project_dir,
             params=inlay_hint_params_to_dict(params),
             options={"trigger": "system", "devEnv": "ide"},
         )
diff --git a/src/finecode/lsp_server/lsp_server.py b/src/finecode/lsp_server/lsp_server.py
@@ -220,7 +220,7 @@ async def _on_initialized(ls: LanguageServer, params: types.InitializedParams):
 
     try:
         global_state.wm_client = ApiClient()
-        await global_state.wm_client.connect("127.0.0.1", port)
+        await global_state.wm_client.connect("127.0.0.1", port, client_id="lsp")
     except (ConnectionRefusedError, OSError) as exc:
         logger.error(f"Could not connect to FineCode WM server: {exc}")
         global_state.wm_client = None
@@ -234,18 +234,14 @@ async def _on_initialized(ls: LanguageServer, params: types.InitializedParams):
             )
         )
 
-    try:
-        info = await global_state.wm_client.get_info()
-        log_path = info.get("logFilePath")
-        if log_path:
-            ls.window_log_message(
-                types.LogMessageParams(
-                    type=types.MessageType.Info,
-                    message=f"FineCode WM Server log: {log_path}",
-                )
+    log_path = global_state.wm_client.server_info.get("logFilePath")
+    if log_path:
+        ls.window_log_message(
+            types.LogMessageParams(
+                type=types.MessageType.Info,
+                message=f"FineCode WM Server log: {log_path}",
             )
-    except Exception:
-        pass
+        )
 
     # Register notification handlers for server→client push messages.
     async def on_tree_changed(params: dict) -> None:
diff --git a/src/finecode/lsp_server/main.py b/src/finecode/lsp_server/main.py
@@ -15,4 +15,7 @@ async def start(
     global_state.lsp_log_file_path = logger_utils.init_logger(log_name="lsp_server", log_level=log_level)
     global_state.wm_log_level = log_level
     server = create_lsp_server()
-    await server.start_io_async()
+    if comm_type == communication_utils.CommunicationType.TCP:
+        await server.start_tcp_async(host, port)
+    else:
+        await server.start_io_async()
diff --git a/src/finecode/mcp_server.py b/src/finecode/mcp_server.py
@@ -11,19 +11,79 @@
 import json
 import pathlib
 import sys
+import uuid
 
+from finecode.wm_client import ApiClient
+from finecode.wm_server import wm_lifecycle
 from loguru import logger
 from mcp.server import Server
 from mcp.server.stdio import stdio_server
 from mcp.types import TextContent, Tool
 
-from finecode.wm_client import ApiClient
-from finecode.wm_server import wm_lifecycle
-
-
 _wm_client = ApiClient()
 server = Server("FineCode")
 
+_partial_result_queues: dict[str, asyncio.Queue] = {}
+
+
+def _setup_partial_result_forwarding() -> None:
+    """Register the WM partial-result notification handler.
+
+    Must be called once after ``_wm_client.connect()``.  Each ``actions/partialResult``
+    notification is routed by token to the matching per-call asyncio.Queue.
+    """
+
+    async def _on_partial_result(params: dict) -> None:
+        token = params.get("token")
+        value = params.get("value")
+        if token and value is not None:
+            queue = _partial_result_queues.get(token)
+            if queue is not None:
+                queue.put_nowait(value)
+
+    _wm_client.on_notification("actions/partialResult", _on_partial_result)
+
+
+async def _run_with_progress(
+    action: str,
+    project: str,
+    params: dict,
+    options: dict,
+    session,
+) -> dict:
+    """Run a WM action with streaming partial results forwarded as MCP log messages.
+
+    ``project`` may be ``""`` to run across all projects that expose the action.
+    Each ``actions/partialResult`` notification is forwarded to the MCP client as a
+    ``notifications/message`` log message while the call blocks waiting for the final result.
+    """
+    token = str(uuid.uuid4())
+    queue: asyncio.Queue = asyncio.Queue()
+    _partial_result_queues[token] = queue
+
+    async def _forward() -> None:
+        try:
+            while True:
+                value = await queue.get()
+                await session.send_log_message(
+                    level="info", data=value, logger="finecode"
+                )
+        except asyncio.CancelledError:
+            pass
+
+    result_task = asyncio.create_task(
+        _wm_client.run_action_with_partial_results(
+            action, project, token, params, options
+        )
+    )
+    forward_task = asyncio.create_task(_forward())
+    try:
+        return await result_task
+    finally:
+        forward_task.cancel()
+        await asyncio.gather(forward_task, return_exceptions=True)
+        _partial_result_queues.pop(token, None)
+
 
 @server.list_tools()
 async def list_tools() -> list[Tool]:
@@ -167,22 +227,22 @@ async def call_tool(name: str, arguments: dict) -> list[TextContent]:
             params={
                 "source_file_path": str(project_path / "pyproject.toml"),
                 "project_raw_config": raw_config,
-                "target_file_path": str(project_path / "finecode_config_dump" / "pyproject.toml"),
+                "target_file_path": str(
+                    project_path / "finecode_config_dump" / "pyproject.toml"
+                ),
             },
             options={"resultFormats": ["json"], "trigger": "user", "devEnv": "ai"},
         )
         return [TextContent(type="text", text=json.dumps(result))]
 
+    from mcp.server.lowlevel.server import request_ctx
+
+    session = request_ctx.get().session
     project = arguments.pop("project", None)
     options = {"resultFormats": ["json"], "trigger": "user", "devEnv": "ai"}
-    if project is not None:
-        result = await _wm_client.run_action(
-            name, project, params=arguments or None, options=options
-        )
-    else:
-        result = await _wm_client.run_batch(
-            [name], params=arguments or None, options=options
-        )
+    result = await _run_with_progress(
+        name, project or "", arguments or {}, options, session
+    )
     return [TextContent(type="text", text=json.dumps(result))]
 
 
@@ -209,10 +269,13 @@ def start(workdir: pathlib.Path, port_file: pathlib.Path | None = None) -> None:
 
     async def _run() -> None:
         try:
-            await _wm_client.connect("127.0.0.1", port)
+            await _wm_client.connect("127.0.0.1", port, client_id="mcp")
         except (ConnectionRefusedError, OSError) as exc:
-            logger.error(f"Could not connect to FineCode WM server on port {port}: {exc}")
+            logger.error(
+                f"Could not connect to FineCode WM server on port {port}: {exc}"
+            )
             sys.exit(1)
+        _setup_partial_result_forwarding()
         logger.debug(f"Add dir to API Client: {workdir}")
         await _wm_client.add_dir(workdir)
         logger.debug("Added dir")
diff --git a/src/finecode/wm_client.py b/src/finecode/wm_client.py
@@ -79,22 +79,26 @@ def __init__(self) -> None:
             str, collections.abc.Callable[..., collections.abc.Coroutine]
         ] = {}
         self._reader_task: asyncio.Task | None = None
+        self.server_info: dict = {}
 
     # -- Connection lifecycle -----------------------------------------------
 
-    async def connect(self, host: str, port: int) -> None:
+    async def connect(self, host: str, port: int, client_id: str | None = None) -> None:
         self._reader, self._writer = await asyncio.open_connection(host, port)
         self._reader_task = asyncio.create_task(self._read_loop())
         logger.info(f"Connected to FineCode API at {host}:{port}")
         try:
-            info = await self.get_info()
-            log_path = info.get("logFilePath")
+            params: dict = {}
+            if client_id is not None:
+                params["clientId"] = client_id
+            self.server_info = await self.request("client/initialize", params) or {}
+            log_path = self.server_info.get("logFilePath")
             if log_path:
                 logger.info(f"WM Server log file: {log_path}")
             else:
                 logger.info("WM Server returned no log file path")
         except Exception as exception:
-            logger.info(f"Failed to get WM Server log file path: {exception}")
+            logger.info(f"Failed to initialize with WM Server: {exception}")
 
     async def close(self) -> None:
         if self._reader_task is not None:
@@ -140,7 +144,7 @@ async def list_projects(self) -> list[dict]:
         return await self.request("workspace/listProjects")
 
     async def find_project_for_file(self, file_path: str) -> str | None:
-        """Return the project name containing a given file.
+        """Return the absolute directory path of the project containing a given file.
 
         An empty string or null result indicates that the file does not belong to
         any project.  This mirrors the server's
@@ -277,6 +281,30 @@ async def run_action(
             body["params"] = params
         return await self.request("actions/run", body)
 
+    async def run_action_with_partial_results(
+        self,
+        action: str,
+        project: str,
+        partial_result_token: str,
+        params: dict | None = None,
+        options: dict | None = None,
+    ) -> dict:
+        """Run an action with streaming partial results via notifications.
+
+        Pass ``project=""`` to run across all projects that expose the action.
+        Partial results are delivered as ``actions/partialResult`` notifications
+        before this coroutine returns the aggregated final result.
+        """
+        body: dict = {
+            "action": action,
+            "project": project,
+            "partialResultToken": partial_result_token,
+            "options": options or {},
+        }
+        if params:
+            body["params"] = params
+        return await self.request("actions/runWithPartialResults", body)
+
     async def add_dir(
         self,
         dir_path: pathlib.Path,
diff --git a/src/finecode/wm_server/wm_server.py b/src/finecode/wm_server/wm_server.py
