VS Code companion drops initial IDE context during MCP connect

[acoliver]

Summary

A freshly built and installed LLxprt Code Companion can accept an authenticated MCP connection from LLxprt, but the CLI never receives the editor context that already exists when it connects. The client reports connected while the shared IDE context remains null, so the first model turn has no active file, selection, or open-file context.

This was reproduced with a fully automated real VS Code extension-host/integrated-terminal test, not mocks.

Environment

• OS: Windows 11 Pro 64-bit, build 26200
• VS Code: 1.126.0 x64
• Bun: 1.3.14
• Repository commit: a25965779a4cff6241ae8ade6d5d3c7036440dfa
• Companion: locally built from that commit, vybestack.llxprt-code-vscode-ide-companion@0.10.0
• CLI: repository scripts/start.ts
• Profile: gpt56high (codex / gpt-5.6-sol); profile authentication independently verified with IDE mode disabled

Reproduction

1. Build the current companion:

   npm run build:vscode
   

2. Install the generated VSIX:

   code --install-extension packages/vscode-ide-companion/llxprt-code-vscode-ide-companion-0.10.0.vsix --force
   

3. Open the repository in VS Code.
4. Open README.md and select the text LLxprt Code.
5. Create a new integrated-terminal process after the companion has activated.
6. Confirm the terminal environment contains all expected companion values:
   • TERM_PROGRAM=vscode
   • numeric LLXPRT_CODE_IDE_SERVER_PORT
   • LLXPRT_CODE_IDE_WORKSPACE_PATH equals the open repository
   • LLXPRT_CODE_IDE_AUTH_TOKEN is present
7. From that terminal, connect directly and inspect context:

   import {
     IdeClient,
     ideContext,
   } from '@vybestack/llxprt-code-ide-integration';
   
   const client = await IdeClient.getInstance();
   await client.connect();
   await Bun.sleep(1000);
   console.log({
     connection: client.getConnectionStatus(),
     detectedIde: client.getDetectedIdeDisplayName(),
     context: ideContext.getIdeContext(),
   });

Actual direct diagnostic result:

{
  "connection": { "status": "connected" },
  "detectedIde": "VS Code",
  "context": null
}

The same behavior is visible end-to-end:

bun scripts/start.ts --profile-load gpt56high --ide-mode enable --output-format json "Using only the IDE context supplied by the connected VS Code companion, reply with exactly IDE_OK|README.md|LLxprt Code if the active file is README.md and the selected text is LLxprt Code. Otherwise reply with exactly IDE_FAIL followed by a brief reason. Do not use tools to read files; rely only on IDE context."

Actual model response:

IDE_FAIL|No connected VS Code companion context is available.

A baseline run with --ide-mode disable and the same profile successfully returned PROFILE_OK, ruling out profile/provider/authentication failure.

Expected behavior

Immediately after IdeClient.connect() resolves, the initial active editor/open files/selection should be available through ideContext.getIdeContext(). The first model turn should receive the selected README.md context without requiring a subsequent editor event.

Diagnosis

The likely cause is an initial-notification race between the MCP client and companion server:

• packages/ide-integration/src/ide/ide-client.ts calls await this.client.connect(transport) before this.registerClientHandlers().
• packages/vscode-ide-companion/src/ide-server.ts sends ide/contextUpdate as a one-time initial notification during the MCP session handshake/request flow, then records that the session has received its initial notification.
• If the server's initial notification arrives while Client.connect() is still completing, no IdeContextNotificationSchema handler exists yet. The update is dropped, while the transport still becomes connected.
• With no later editor change, the companion does not broadcast another context update, leaving ideContext undefined for the first model turn.

The current unit tests do not appear to exercise a real initial-context notification arriving during connection setup.

Suggested fix / test

• Register notification/error/close handlers immediately after constructing Client and before awaiting client.connect(transport), or introduce an explicit post-connect request/ack for the current editor state.
• Add a behavioral test that starts with an already-active file and selection, connects a real client/server pair, and asserts the initial context is available as soon as connect() completes without a later editor event.

Additional notes

• The extension itself activates successfully.
• Its MCP server port, workspace path, and auth token propagate into the integrated terminal correctly.
• The direct client is authenticated and reports connected; this is specifically loss of initial editor context, not server discovery or authorization failure.
• No token values or credentials are included in this report.

[coderabbitai]

🔗 Related PRs

#704 - Fix VSCode companion activation by bundling tiktoken wasm (#692) [merged]
#802 - Fix keytar handling for MCP OAuth (#770) [merged]
#2071 - Fix dumpcontext invocation propagation (Fixes #2066) [merged]
#2451 - Make MCP server discovery non-blocking at startup (Fixes #2325) [merged]
#2583 - Require explicit provider selection at startup (Fixes #2481) [merged]

---

📝 Issue Planner

_{Check the box below or use the @coderabbitai plan command to generate an implementation plan and prompts that you can use with your favorite coding assistant.}

✅ Successfully generated implementation plan.

💡 Use @coderabbitai <feedback> to replan

---

🧪 Issue enrichment is currently in open beta.

You can configure auto-planning by selecting labels in the issue_enrichment configuration.

To disable automatic issue enrichment, add the following to your .coderabbit.yaml:

issue_enrichment:
  auto_enrich:
    enabled: false

💬 Have feedback or questions? Drop into our discord!

[coderabbitai]

Coding Plan

Summary

• Fix the root cause by reordering establishConnection() in ide-client.ts so registerClientHandlers() runs before client.connect(transport), ensuring the notification handler exists before the companion's one-time initial ide/contextUpdate push arrives.
• Avoid introducing a new pull-style protocol surface, keeping the change minimal and consistent with standard MCP SDK handler-before-connect usage.
• Add a real cross-package behavioral test in packages/vscode-ide-companion that pre-seeds active editor state, connects a genuine IdeClient to a genuine IDEServer, and asserts initial context is present right after connect() resolves with no later editor event.

Design Choices

Design Choice 1: Should the fix rely on reordering handler registration, or introduce an explicit post-connect pull request for current context?

Options Considered:

1. Reorder establishConnection() so handlers are attached to the constructed Client before awaiting connect(transport). No protocol/schema changes; closes the race for the server's initial push.
2. Add a new ide/getContext request/response (or a getIdeContext tool) on the server and have the client explicitly pull context after connect. Requires new schemas in both packages, wiring openFilesManager into createMcpServer, and client-side request logic.

Chosen Option: 1

Rationale: The ticket offers both options ("Register handlers before connect... or introduce an explicit post-connect request/ack"). The exploration confirmed registerClientHandlers() only needs the constructed Client object and does not depend on the transport or an established connection, so reordering is safe and minimal. It is the smallest, most idiomatic change, directly addresses the diagnosed race, requires no new protocol surface across two independently-maintained schema copies, and matches standard SDK usage where handlers are registered before connecting.

💡 User Tips

Regenerate the plan with different choices with @coderabbitai <feedback>.

Implementation Steps

Phase 1: Eliminate the handler-registration-order race

The goal of this phase is to guarantee that all server→client notification and lifecycle handlers are wired onto the MCP Client before the transport connection is established, so the companion's one-time initial ide/contextUpdate push can never arrive before a handler exists. This is a localized change confined to the single connection code path.

Register client handlers before connecting the transport

Reorder the connection sequence so notification/error/close handlers are attached immediately after the Client is constructed and before the transport handshake begins.

• Update packages/ide-integration/src/ide/ide-client.ts in establishConnection() to call this.registerClientHandlers() immediately after this.client = new Client({...}) and before await this.client.connect(transport).
• Leave registerClientHandlers() itself unchanged; its early-return guard already only requires this.client to exist and it never touches the transport, so it is safe to invoke pre-connect.
• Confirm this is the only call site of both client.connect( and registerClientHandlers(), so no reconnect/retry path needs mirrored changes.
• Keep setState(Connected) after connect() resolves so connection status semantics are unchanged.

🤖 Prompt for AI agents

Implement the fix for the initial `ide/contextUpdate` race by reordering handler
registration relative to the transport connection in the MCP client.

- In `packages/ide-integration/src/ide/ide-client.ts`, inside
`establishConnection()`, move the call to `this.registerClientHandlers()` so it
happens immediately after `this.client = new Client({...})` is constructed, and
before the `await this.client.connect(transport)` call.
- Do not modify the internals of `registerClientHandlers()`; its early-return
guard only checks that `this.client` exists and does not depend on the
transport, so it is safe to call before connecting.
- Verify there is only one call site for both `client.connect(` and
`registerClientHandlers()` in this file, so there is no other reconnect/retry
code path that needs the same reordering.
- Preserve the existing behavior of calling `setState(Connected)` after
`connect()` resolves, so connection status semantics remain unchanged.

Phase 2: Prove initial context delivery with a real client/server behavioral test

The goal of this phase is to add regression coverage that reproduces the ticket scenario end-to-end — an already-open, already-selected file present before connection — using a genuine IdeClient connected to a genuine IDEServer, and asserting the context is available as soon as connect() resolves without any subsequent editor event. This test lives in the companion package because dependency direction allows it to import the client, and because IDEServer.start() already spins up a real HTTP/MCP server.

Task 1: Establish real client/server test scaffolding in the companion package

Create the infrastructure to start a real server and point a real client at it, reusing the existing companion test's mocking approach for the vscode module and filesystem.

• Add a new behavioral test file under packages/vscode-ide-companion/src/ (e.g. ide-client-integration.test.ts), mirroring the existing ide-server.test.ts setup that mocks vscode and node:fs/promises/node:os but leaves the MCP SDK and HTTP layer real.
• Pre-seed vscode.window.activeTextEditor (file-scheme README.md, a selection whose getText() returns LLxprt Code, and a cursor position) before calling ideServer.start(mockContext), so OpenFilesManager captures the active file/selection at construction time.
• Extract the server's real ephemeral port and auth token from the mocked environmentVariableCollection.replace calls (as existing tests already read written values), and set the LLXPRT_CODE_IDE_SERVER_PORT, LLXPRT_CODE_IDE_WORKSPACE_PATH, and LLXPRT_CODE_IDE_AUTH_TOKEN environment variables consumed by IdeClient.connect()'s discovery logic.
• Import IdeClient and ideContext from @vybestack/llxprt-code-ide-integration, reset the client singleton between tests, and ensure teardown stops the server and disconnects the client.

Task 2: Assert initial context is available immediately after connect

Add the assertion that captures the ticket's core failure and verifies the fix.

• After await client.connect() resolves, assert client.getConnectionStatus() reports connected AND ideContext.getIdeContext() reflects the pre-seeded README.md file with selectedText LLxprt Code — without triggering any additional editor change event.
• Ensure the assertion does not rely on arbitrary long sleeps beyond what is needed for the SSE push to arrive; prefer awaiting the context store's subscription/update or a short bounded poll so the test is deterministic.
• Optionally guard against regression by confirming the test fails when handler registration is placed after connect() (documented in the test intent, not necessarily as a separate case).

🤖 Prompt for AI agents

Add a real client/server behavioral test in the companion package that
reproduces the ticket's race condition scenario and verifies the fix, without
adding any new protocol surface.

- Create a new test file under `packages/vscode-ide-companion/src/` (e.g.
`ide-client-integration.test.ts`), mirroring the mocking approach in the
existing `ide-server.test.ts`: mock the `vscode` module and
`node:fs/promises`/`node:os`, but keep the MCP SDK and HTTP layer real so both a
genuine `IdeClient` and a genuine `IDEServer` are exercised.
- Before calling `ideServer.start(mockContext)`, pre-seed
`vscode.window.activeTextEditor` with a file-scheme `README.md` document, a
selection whose `getText()` returns `LLxprt Code`, and a cursor position, so
that `OpenFilesManager` captures this active file/selection state at
construction time.
- Extract the server's real ephemeral port and auth token from the mocked
`environmentVariableCollection.replace` calls (following the pattern of existing
tests that read these written values), then set the
`LLXPRT_CODE_IDE_SERVER_PORT`, `LLXPRT_CODE_IDE_WORKSPACE_PATH`, and
`LLXPRT_CODE_IDE_AUTH_TOKEN` environment variables that `IdeClient.connect()`'s
discovery logic consumes.
- Import `IdeClient` and `ideContext` from
`@vybestack/llxprt-code-ide-integration`; reset the `IdeClient` singleton
between tests, and ensure teardown stops the server and disconnects the client.
- After `await client.connect()` resolves, assert both that
`client.getConnectionStatus()` reports connected and that
`ideContext.getIdeContext()` reflects the pre-seeded `README.md` file with
`selectedText` equal to `LLxprt Code`, without triggering any additional editor
change event after connect.
- Make the assertion deterministic by awaiting the context store's
subscription/update (or a short bounded poll) rather than relying on arbitrary
long sleeps.
- Optionally note in the test's intent/documentation that this test should fail
if handler registration were placed after `connect()`, as a way to guard against
regression of the root cause, though this does not need to be a separate test
case.

Research

The IDE integration spans two workspace packages.

packages/ide-integration provides the client-side IdeClient singleton (src/ide/ide-client.ts) and the shared in-memory ideContext pub/sub store plus zod wire schemas (src/ide/ideContext.ts), exported publicly via src/ide/index.ts.

packages/vscode-ide-companion provides the server-side IDEServer (src/ide-server.ts), which runs a real Streamable HTTP MCP server, tracks editor state via OpenFilesManager (src/open-files-manager.ts), and pushes ide/contextUpdate notifications.

The companion already depends on @vybestack/llxprt-code-ide-integration, and both packages pin compatible MCP SDK versions, making a genuine cross-package client/server test feasible from within the companion package.

---

🚀 Next Steps

🤖 All AI agent prompts combined

Task: 1

Implement the fix for the initial `ide/contextUpdate` race by reordering handler
registration relative to the transport connection in the MCP client.

- In `packages/ide-integration/src/ide/ide-client.ts`, inside
`establishConnection()`, move the call to `this.registerClientHandlers()` so it
happens immediately after `this.client = new Client({...})` is constructed, and
before the `await this.client.connect(transport)` call.
- Do not modify the internals of `registerClientHandlers()`; its early-return
guard only checks that `this.client` exists and does not depend on the
transport, so it is safe to call before connecting.
- Verify there is only one call site for both `client.connect(` and
`registerClientHandlers()` in this file, so there is no other reconnect/retry
code path that needs the same reordering.
- Preserve the existing behavior of calling `setState(Connected)` after
`connect()` resolves, so connection status semantics remain unchanged.
===============================================================================

Task: 2

Add a real client/server behavioral test in the companion package that
reproduces the ticket's race condition scenario and verifies the fix, without
adding any new protocol surface.

- Create a new test file under `packages/vscode-ide-companion/src/` (e.g.
`ide-client-integration.test.ts`), mirroring the mocking approach in the
existing `ide-server.test.ts`: mock the `vscode` module and
`node:fs/promises`/`node:os`, but keep the MCP SDK and HTTP layer real so both a
genuine `IdeClient` and a genuine `IDEServer` are exercised.
- Before calling `ideServer.start(mockContext)`, pre-seed
`vscode.window.activeTextEditor` with a file-scheme `README.md` document, a
selection whose `getText()` returns `LLxprt Code`, and a cursor position, so
that `OpenFilesManager` captures this active file/selection state at
construction time.
- Extract the server's real ephemeral port and auth token from the mocked
`environmentVariableCollection.replace` calls (following the pattern of existing
tests that read these written values), then set the
`LLXPRT_CODE_IDE_SERVER_PORT`, `LLXPRT_CODE_IDE_WORKSPACE_PATH`, and
`LLXPRT_CODE_IDE_AUTH_TOKEN` environment variables that `IdeClient.connect()`'s
discovery logic consumes.
- Import `IdeClient` and `ideContext` from
`@vybestack/llxprt-code-ide-integration`; reset the `IdeClient` singleton
between tests, and ensure teardown stops the server and disconnects the client.
- After `await client.connect()` resolves, assert both that
`client.getConnectionStatus()` reports connected and that
`ideContext.getIdeContext()` reflects the pre-seeded `README.md` file with
`selectedText` equal to `LLxprt Code`, without triggering any additional editor
change event after connect.
- Make the assertion deterministic by awaiting the context store's
subscription/update (or a short bounded poll) rather than relying on arbitrary
long sleeps.
- Optionally note in the test's intent/documentation that this test should fail
if handler registration were placed after `connect()`, as a way to guard against
regression of the root cause, though this does not need to be a separate test
case.

💡 Iterate on the plan with: `@coderabbitai` <feedback>

Example Feedback
- `@coderabbitai` You can skip phase 3. Add a simple unit test case for phase 2.
- `@coderabbitai` For design choice 1 go ahead with option 3 and replan.

---

💬 Have feedback or questions? Drop into our discord!