99using System . Threading . Tasks ;
1010using Microsoft . ServiceHub . Framework ;
1111using Microsoft . VisualStudio . Copilot ;
12+ using Microsoft . VisualStudio . Copilot . Internal . Mcp ;
1213using Microsoft . VisualStudio . Shell ;
1314using Microsoft . VisualStudio . Shell . ServiceBroker ;
1415using NuGet . VisualStudio ;
@@ -22,7 +23,7 @@ internal class CopilotToolInvocationService : ICopilotToolInvocationService
2223 private const string AuthStatusDetermined = "c936efcc-6baa-4ad3-9c2b-7ba750acf18f" ;
2324 private static readonly Guid CopilotReadyUIContext = new ( AuthStatusDetermined ) ;
2425
25- [ Import ( typeof ( SVsFullAccessServiceBroker ) ) ]
26+ [ Import ( typeof ( SVsFullAccessServiceBroker ) , AllowDefault = true ) ]
2627 public IServiceBroker ? ServiceBroker { get ; set ; }
2728
2829 public async Task < CopilotToolSessionResult > TryCreateToolSessionAsync (
@@ -46,7 +47,25 @@ public async Task<CopilotToolSessionResult> TryCreateToolSessionAsync(
4647 return CopilotToolSessionResult . Failure ( CopilotToolSessionError . ServiceBrokerNotAvailable ) ;
4748 }
4849
49- // 3. Acquire Copilot service, ownership transfers to CopilotToolSession on success
50+ // 3. Verify the required MCP server is registered and active
51+ IMcpServerInfoService ? mcpServerInfoService = await ServiceBroker . GetProxyAsync < IMcpServerInfoService > ( McpServiceIdentities . ServerInfoService . Descriptor , cancellationToken ) ;
52+ using ( mcpServerInfoService as IDisposable )
53+ {
54+ if ( mcpServerInfoService is null )
55+ {
56+ return CopilotToolSessionResult . Failure ( CopilotToolSessionError . McpServerInfoServiceNotAvailable ) ;
57+ }
58+
59+ // The NuGet MCP server may be registered under different names depending on how it
60+ // was installed (in-VS vs. Anthropic/GitHub MCP registry). It is considered available
61+ // if any of the acceptable names reports an Active or Suspended state.
62+ if ( ! await IsServerAvailableAsync ( mcpServerInfoService , acceptableMcpServerNames , cancellationToken ) )
63+ {
64+ return CopilotToolSessionResult . Failure ( CopilotToolSessionError . McpServerNotActive ) ;
65+ }
66+ }
67+
68+ // 4. Acquire Copilot service, ownership transfers to CopilotToolSession on success
5069#pragma warning disable ISB001 // Dispose objects before losing scope - ownership is transferred to CopilotToolSession on success
5170 ICopilotService ? copilotService = await ServiceBroker . GetProxyAsync < ICopilotService > ( CopilotDescriptors . CopilotService , cancellationToken ) ;
5271#pragma warning restore ISB001
@@ -59,7 +78,7 @@ public async Task<CopilotToolSessionResult> TryCreateToolSessionAsync(
5978 return CopilotToolSessionResult . Failure ( CopilotToolSessionError . CopilotServiceNotAvailable ) ;
6079 }
6180
62- // 4 . Acquire MCP tool function provider and get available functions
81+ // 5 . Acquire MCP tool function provider and get available functions
6382 ICopilotFunctionProvider ? cfp = await ServiceBroker . GetProxyAsync < ICopilotFunctionProvider > ( CopilotDescriptors . McpToolService , cancellationToken ) ;
6483 using ( cfp as IDisposable )
6584 {
@@ -68,16 +87,16 @@ public async Task<CopilotToolSessionResult> TryCreateToolSessionAsync(
6887 return CopilotToolSessionResult . Failure ( CopilotToolSessionError . McpToolServiceNotAvailable ) ;
6988 }
7089
71- // 5 . Verify the required tool is available. We match on ServerNameOfFunction + Group
90+ // 6 . Verify the required tool is available. We match on ServerNameOfFunction + Group
7291 // (the same logical NuGet MCP tool can be exposed under different Group values
73- // depending on how it was installed — in-VS vs. Anthropic/GitHub MCP registry).
92+ // depending on how it was installed - in-VS vs. Anthropic/GitHub MCP registry).
7493 IReadOnlyList < CopilotFunctionDescriptor > ? functions = await cfp . GetFunctionsAsync ( correlationId , cancellationToken ) ;
7594 if ( ! IsToolAvailable ( functions , mcpToolName , acceptableMcpServerNames ) )
7695 {
7796 return CopilotToolSessionResult . Failure ( CopilotToolSessionError . ToolNotAvailable ) ;
7897 }
7998
80- // 6 . Start Copilot thread
99+ // 7 . Start Copilot thread
81100 CopilotThreadOptions options = new ( clientId ) ;
82101 CopilotThread thread = await copilotService . StartThreadAsync ( options , cancellationToken ) ;
83102
@@ -109,5 +128,22 @@ internal static bool IsToolAvailable(
109128 && f . Group is not null
110129 && acceptableMcpServerNames . Contains ( f . Group , StringComparer . OrdinalIgnoreCase ) ) ?? false ;
111130 }
131+
132+ internal static async Task < bool > IsServerAvailableAsync (
133+ IMcpServerInfoService mcpServerInfoService ,
134+ IReadOnlyCollection < string > acceptableMcpServerNames ,
135+ CancellationToken cancellationToken )
136+ {
137+ foreach ( string serverName in acceptableMcpServerNames )
138+ {
139+ McpServerState ? state = await mcpServerInfoService . GetServerStateAsync ( serverName , cancellationToken ) ;
140+ if ( state is McpServerState . Active or McpServerState . Suspended )
141+ {
142+ return true ;
143+ }
144+ }
145+
146+ return false ;
147+ }
112148 }
113149}
0 commit comments