fix: tools use org service W-19828433 (#287)

* chore: stricter e2e test

* chore: org resume returns username

* chore: stricter access

* chore: manual allowlist org check

Author: cristiand391

packages/mcp-provider-dx-core/src/index.ts

@@ -38,12 +38,12 @@ export class DxCoreMcpProvider extends McpProvider {
     return Promise.resolve([
       new AssignPermissionSetMcpTool(services),
       new CreateOrgSnapshotMcpTool(services),
-      new CreateScratchOrgMcpTool(),
-      new DeleteOrgMcpTool(),
+      new CreateScratchOrgMcpTool(services),
+      new DeleteOrgMcpTool(services),
       new DeployMetadataMcpTool(services),
       new GetUsernameMcpTool(services),
       new ListAllOrgsMcpTool(services),
-      new OrgOpenMcpTool(),
+      new OrgOpenMcpTool(services),
       new QueryOrgMcpTool(services),
       new ResumeMcpTool(services),
       new RetrieveMetadataMcpTool(services),

packages/mcp-provider-dx-core/src/tools/create_org_snapshot.ts

@@ -88,7 +88,9 @@ create a snapshot of my MyScratch in myDevHub`,
     try {
       process.chdir(input.directory);
 
-      const sourceOrgId = (await Org.create({ aliasOrUsername: input.sourceOrg })).getOrgId();
+      const connection = await this.services.getOrgService().getConnection(input.sourceOrg);
+
+      const sourceOrgId = (await Org.create({ connection })).getOrgId();
       const devHubConnection = await this.services.getOrgService().getConnection(input.devHub);
       const createResponse = await devHubConnection.sobject('OrgSnapshot').create({
         SourceOrg: sourceOrgId,

packages/mcp-provider-dx-core/src/tools/create_scratch_org.ts

@@ -20,7 +20,7 @@ import { z } from 'zod';
 import { Org, scratchOrgCreate, ScratchOrgCreateOptions } from '@salesforce/core';
 import { Duration } from '@salesforce/kit';
 import { CallToolResult } from '@modelcontextprotocol/sdk/types.js';
-import { McpTool, McpToolConfig, ReleaseState, Toolset } from '@salesforce/mcp-provider-api';
+import { McpTool, McpToolConfig, ReleaseState, Services, Toolset } from '@salesforce/mcp-provider-api';
 import { ensureString } from '@salesforce/ts-types/lib/narrowing/ensure.js';
 import { textResponse } from '../shared/utils.js';
 import { directoryParam, usernameOrAliasParam } from '../shared/params.js';
@@ -95,6 +95,10 @@ type InputArgsShape = typeof createScratchOrgParams.shape;
 type OutputArgsShape = z.ZodRawShape;
 
 export class CreateScratchOrgMcpTool extends McpTool<InputArgsShape, OutputArgsShape> {
+  public constructor(private readonly services: Services) {
+    super();
+  }
+
   public getReleaseState(): ReleaseState {
     return ReleaseState.NON_GA;
   }
@@ -127,6 +131,27 @@ create a scratch org aliased as MyNewOrg and set as default and don't wait for i
   public async exec(input: InputArgs): Promise<CallToolResult> {
     try {
       process.chdir(input.directory);
+
+      // NOTE: 
+      // this should be:
+      // ```ts
+      // const connection = await this.services.getOrgService().getConnection(input.devHub);
+      // const hubOrProd = await Org.create({ connection });
+      // ```
+      //
+      // but there's a bug where if you create scratch synchronously, sfdx-core throws while polling;
+      // ```
+      // [NamedOrgNotFoundError]: No authorization information found for <devhub-username>.
+      // ```
+      //
+      // it doesn't happen when creating asynchronously.
+      // will be fixed in W-19828802
+      const allowedOrgs = await this.services.getOrgService().getAllowedOrgs()
+      if (!allowedOrgs.find(o => o.aliases?.includes(input.devHub) || o.username === input.devHub)) {
+        throw new Error(
+          'No org found with the provided devhub username/alias. Ask the user to specify one or check their MCP Server startup config.'
+        )}
+
       const hubOrProd = await Org.create({ aliasOrUsername: input.devHub });
 
       const requestParams: ScratchOrgCreateOptions = {

packages/mcp-provider-dx-core/src/tools/delete_org.ts

@@ -17,7 +17,7 @@
 import { z } from 'zod';
 import { AuthRemover, Org } from '@salesforce/core';
 import { CallToolResult } from '@modelcontextprotocol/sdk/types.js';
-import { McpTool, McpToolConfig, ReleaseState, Toolset } from '@salesforce/mcp-provider-api';
+import { McpTool, McpToolConfig, ReleaseState, Services, Toolset } from '@salesforce/mcp-provider-api';
 import { textResponse } from '../shared/utils.js';
 import { directoryParam, usernameOrAliasParam } from '../shared/params.js';
 
@@ -42,6 +42,10 @@ type InputArgsShape = typeof deleteOrgParams.shape;
 type OutputArgsShape = z.ZodRawShape;
 
 export class DeleteOrgMcpTool extends McpTool<InputArgsShape, OutputArgsShape> {
+  public constructor(private readonly services: Services) {
+    super();
+  }
+
   public getReleaseState(): ReleaseState {
     return ReleaseState.NON_GA;
   }
@@ -75,7 +79,9 @@ Can you delete [email protected]`,
   public async exec(input: InputArgs): Promise<CallToolResult> {
     try {
       process.chdir(input.directory);
-      const org = await Org.create({ aliasOrUsername: input.usernameOrAlias });
+      const connection = await this.services.getOrgService().getConnection(input.usernameOrAlias);
+      const org = await Org.create({ connection });
+
       await org.delete();
       return textResponse(`Successfully deleted ${input.usernameOrAlias}`);
     } catch (e) {

packages/mcp-provider-dx-core/src/tools/open_org.ts

@@ -18,7 +18,7 @@ import { z } from 'zod';
 import { Org } from '@salesforce/core';
 import { MetadataResolver } from '@salesforce/source-deploy-retrieve';
 import open from 'open';
-import { McpTool, McpToolConfig, ReleaseState, Toolset } from '@salesforce/mcp-provider-api';
+import { McpTool, McpToolConfig, ReleaseState, Services, Toolset } from '@salesforce/mcp-provider-api';
 import { CallToolResult } from '@modelcontextprotocol/sdk/types.js';
 import { textResponse } from '../shared/utils.js';
 import { directoryParam, usernameOrAliasParam } from '../shared/params.js';
@@ -37,6 +37,10 @@ type InputArgsShape = typeof orgOpenParamsSchema.shape;
 type OutputArgsShape = z.ZodRawShape;
 
 export class OrgOpenMcpTool extends McpTool<InputArgsShape, OutputArgsShape> {
+  public constructor(private readonly services: Services) {
+    super();
+  }
+
   public getReleaseState(): ReleaseState {
     return ReleaseState.NON_GA;
   }
@@ -67,9 +71,11 @@ You can specify a metadata file you want to open.`,
   public async exec(input: InputArgs): Promise<CallToolResult> {
     process.chdir(input.directory);
 
+    const connection = await this.services.getOrgService().getConnection(input.usernameOrAlias);
+
     const org = await Org.create({
-      aliasOrUsername: input.usernameOrAlias,
-    });
+      connection
+    })
 
     if (input.filePath) {
       const metadataResolver = new MetadataResolver();

packages/mcp-provider-dx-core/src/tools/resume_tool_operation.ts

@@ -21,6 +21,7 @@ import { Duration } from '@salesforce/kit';
 import { MetadataApiDeploy } from '@salesforce/source-deploy-retrieve';
 import { McpTool, McpToolConfig, ReleaseState, Services, Toolset } from '@salesforce/mcp-provider-api';
 import { CallToolResult } from '@modelcontextprotocol/sdk/types.js';
+import { ensureString } from '@salesforce/ts-types';
 import { textResponse } from '../shared/utils.js';
 import { directoryParam, usernameOrAliasParam } from '../shared/params.js';
 import { type ToolTextResponse } from '../shared/types.js';
@@ -46,7 +47,7 @@ const resumableIdPrefixes = new Map<string, string>([
  * Returns:
  * - textResponse: Username/alias and org configuration
  */
-const resumeParamsSchema = z.object({
+export const resumeParamsSchema = z.object({
   jobId: z.string().describe('The job id of the long running operation to resume (required)'),
   wait: z
     .number()
@@ -178,7 +179,7 @@ async function resumeOrgSnapshot(connection: Connection, jobId: string, wait: nu
 async function resumeScratchOrg(jobId: string, wait: number): Promise<ToolTextResponse> {
   try {
     const result = await scratchOrgResume(jobId, Duration.minutes(wait));
-    return textResponse(`Scratch org created: ${JSON.stringify(result)}`, false);
+    return textResponse(`Successfully created scratch org, username: ${ensureString(result.username)}`);
   } catch (error) {
     return textResponse(
       `Resumed scratch org creation failed: ${error instanceof Error ? error.message : 'Unknown error'}`,

packages/mcp-provider-dx-core/test/e2e/create_scratch_org.test.ts

@@ -14,22 +14,23 @@
  * limitations under the License.
  */
 
-import { expect } from 'chai';
+import { assert, expect } from 'chai';
 import { McpTestClient, DxMcpTransport } from '@salesforce/mcp-test-client';
 import { TestSession } from '@salesforce/cli-plugins-testkit';
 import { z } from 'zod';
 import { matchesAccessToken } from '@salesforce/core';
+import { ensureString } from '@salesforce/ts-types';
 import { createScratchOrgParams } from '../../src/tools/create_scratch_org.js';
+import { resumeParamsSchema } from '../../src/tools/resume_tool_operation.js';
 
 describe('create_scratch_org', () => {
   const client = new McpTestClient({
     timeout: 120_000,
   });
 
-  const devHubUsername = process.env.TESTKIT_HUB_USERNAME as string;
+  let devHubUsername: string;
 
   let testSession: TestSession;
-  let resolvedDevHubUsername: string;
 
   const createScratchOrgSchema = {
     name: z.literal('create_scratch_org'),
@@ -44,15 +45,10 @@ describe('create_scratch_org', () => {
       devhubAuthStrategy: 'AUTO',
     });
 
-    const hubUsername = testSession.hubOrg?.username;
-    resolvedDevHubUsername = hubUsername ?? devHubUsername;
-    
-    if (!resolvedDevHubUsername) {
-      throw new Error('No DevHub username available from TestSession or TESTKIT_HUB_USERNAME environment variable');
-    }
+    devHubUsername = ensureString(testSession.hubOrg?.username);
 
     const transport = DxMcpTransport({
-      args: ['--orgs', 'ALLOW_ALL_ORGS', '--no-telemetry', '--toolsets', 'all', '--allow-non-ga-tools'],
+      args: ['--orgs', `DEFAULT_TARGET_ORG,${devHubUsername}`, '--no-telemetry', '--toolsets', 'all', '--allow-non-ga-tools'],
     });
 
     await client.connect(transport);
@@ -72,7 +68,7 @@ describe('create_scratch_org', () => {
       name: 'create_scratch_org',
       params: {
         directory: testSession.project.dir,
-        devHub: resolvedDevHubUsername,
+        devHub: devHubUsername,
       },
     });
 
@@ -81,7 +77,7 @@ describe('create_scratch_org', () => {
     expect(result.content[0].type).to.equal('text');
 
     const responseText = result.content[0].text as string;
-    expect(matchesAccessToken(responseText)).to.be.false;
+    assertNoSensitiveInfo(responseText)
     expect(responseText).to.include('Successfully created scratch org');
   });
 
@@ -90,25 +86,56 @@ describe('create_scratch_org', () => {
       name: 'create_scratch_org',
       params: {
         directory: testSession.project.dir,
-        devHub: resolvedDevHubUsername,
+        devHub: devHubUsername,
         async: true,
         alias: 'test-async-org'
       },
     });
     expect(asyncResult.isError).to.be.false;
     expect(asyncResult.content.length).to.equal(1);
-    expect(asyncResult.content[0].type).to.equal('text');
 
+    if (asyncResult.content[0].type !== 'text') assert.fail();
+
     const asyncResponseText = asyncResult.content[0].text;
     expect(asyncResponseText).to.include('Successfully enqueued scratch org with job Id:');
+
+    // now validate it was created by resuming the operation
+
+    const jobIdMatch = asyncResponseText.match(/job Id: ([A-Za-z0-9]+)/);
+    expect(jobIdMatch).to.not.be.null;
+    
+    const jobId: string =  jobIdMatch![1]
+
+    const asyncResumeResult = await client.callTool({
+      name: z.literal('resume_tool_operation'),
+      params: resumeParamsSchema
+    }, {
+      name: 'resume_tool_operation',
+      params: {
+        directory: testSession.project.dir,
+        jobId,
+        usernameOrAlias: ensureString(testSession.hubOrg.username)
+      },
+    });
+
+    expect(asyncResumeResult.isError).to.be.false;
+    expect(asyncResumeResult.content.length).to.equal(1);
+    
+    if (asyncResumeResult.content[0].type !== 'text') assert.fail();
+
+    const asyncResumeResponseText = asyncResumeResult.content[0].text;
+    
+    // tool output shouldn't access tokens/auth info other than the username
+    assertNoSensitiveInfo(asyncResumeResponseText);
+    expect(asyncResumeResponseText).to.include('Successfully created scratch org');
   });
 
   it('should create scratch org with optional parameters', async () => {
     const result = await client.callTool(createScratchOrgSchema, {
       name: 'create_scratch_org',
       params: {
         directory: testSession.project.dir,
-        devHub: resolvedDevHubUsername,
+        devHub: devHubUsername,
         alias: 'test-custom-org',
         duration: 3,
         edition: 'developer',
@@ -142,4 +169,16 @@ describe('create_scratch_org', () => {
     const responseText = result.content[0].text;
     expect(responseText).to.include('Failed to create org:');
   });
-});
+});
+
+/**
+ * Helper function to assert that response text doesn't contain sensitive authentication information
+ */
+function assertNoSensitiveInfo(responseText: string): void {
+  expect(matchesAccessToken(responseText)).to.be.false;
+  expect(responseText).to.not.match(/authcode/i);
+  expect(responseText).to.not.match(/token/i);
+  expect(responseText).to.not.match(/privatekey/i);
+  expect(responseText).to.not.match(/clientid/i);
+  expect(responseText).to.not.match(/connectedappconsumerkey/i);
+}

packages/mcp-provider-dx-core/test/e2e/retrieve_metadata.test.ts

@@ -25,7 +25,7 @@ import { retrieveMetadataParams } from '../../src/tools/retrieve_metadata.js';
 
 describe('retrieve_metadata', () => {
   const client = new McpTestClient({
-    timeout: 60000,
+    timeout: 60_000,
   });
 
   let testSession: TestSession;