Add mssql_dab show flow and improve tool descripiton (#21511)

* Add DAB show flow and align tool contract

* Stabilize DAB tool open and init behavior

* Align DAB no-active-designer messaging

* Improve DAB tool coverage

* Fix DAB tool test coverage assertions

Author: caohai

extensions/mssql/l10n/bundle.l10n.json

@@ -2673,6 +2673,14 @@
             "{2} is the operation name 'apply_edits' and must not be translated"
         ]
     },
+    "Data API Builder opened. Continue with {0} operations ({1}/{2})./{0} is the command identifier 'mssql_dab' and must not be translated{1} is the operation name 'get_state' and must not be translated{2} is the operation name 'apply_changes' and must not be translated": {
+        "message": "Data API Builder opened. Continue with {0} operations ({1}/{2}).",
+        "comment": [
+            "{0} is the command identifier 'mssql_dab' and must not be translated",
+            "{1} is the operation name 'get_state' and must not be translated",
+            "{2} is the operation name 'apply_changes' and must not be translated"
+        ]
+    },
     "Execute '{0}' operation on the schema designer?/{0} is the operation name": {
         "message": "Execute '{0}' operation on the schema designer?",
         "comment": ["{0} is the operation name"]
@@ -2690,6 +2698,8 @@
         "message": "Executing '{0}' operation on Data API Builder",
         "comment": ["{0} is the operation name"]
     },
+    "No active schema designer found. Please open Data API Builder first using mssql_dab with operation 'show' or from the UI.": "No active schema designer found. Please open Data API Builder first using mssql_dab with operation 'show' or from the UI.",
+    "Missing connectionId. Please provide a connectionId to open Data API Builder.": "Missing connectionId. Please provide a connectionId to open Data API Builder.",
     "No active schema designer found. Please open one first using mssql_schema_designer with operation 'show' or from the UI.": "No active schema designer found. Please open one first using mssql_schema_designer with operation 'show' or from the UI.",
     "Schema designer state changed. Fetch the latest schema and retry the operation.": "Schema designer state changed. Fetch the latest schema and retry the operation.",
     "Missing connectionId. Please provide a connectionId to open the schema designer.": "Missing connectionId. Please provide a connectionId to open the schema designer.",

extensions/mssql/package.json

@@ -2055,6 +2055,15 @@ export class MssqlChatAgent {
             "{2} is the operation name 'apply_edits' and must not be translated",
         ],
     });
+    public static dabToolShowSuccessMessage = l10n.t({
+        message: "Data API Builder opened. Continue with {0} operations ({1}/{2}).",
+        args: ["mssql_dab", "get_state", "apply_changes"],
+        comment: [
+            "{0} is the command identifier 'mssql_dab' and must not be translated",
+            "{1} is the operation name 'get_state' and must not be translated",
+            "{2} is the operation name 'apply_changes' and must not be translated",
+        ],
+    });
     public static schemaDesignerToolConfirmationTitle = l10n.t("Schema Designer");
     public static schemaDesignerToolConfirmationMessage = (operation: string) => {
         return l10n.t({
@@ -2085,6 +2094,12 @@ export class MssqlChatAgent {
             comment: ["{0} is the operation name"],
         });
     };
+    public static dabToolNoActiveDesigner = l10n.t(
+        "No active schema designer found. Please open Data API Builder first using mssql_dab with operation 'show' or from the UI.",
+    );
+    public static dabToolMissingConnectionId = l10n.t(
+        "Missing connectionId. Please provide a connectionId to open Data API Builder.",
+    );
     public static schemaDesignerNoActiveDesigner = l10n.t(
         "No active schema designer found. Please open one first using mssql_schema_designer with operation 'show' or from the UI.",
     );

extensions/mssql/src/constants/locConstants.ts

@@ -856,7 +856,14 @@ export default class MainController implements vscode.Disposable {
 
         // Register mssql_dab tool
         this._context.subscriptions.push(
-            vscode.lm.registerTool(Constants.copilotDabToolName, new DabTool()),
+            vscode.lm.registerTool(
+                Constants.copilotDabToolName,
+                new DabTool(
+                    this.connectionManager,
+                    async (connectionUri: string, database: string) =>
+                        this.openDabDesigner(connectionUri, database),
+                ),
+            ),
         );
     }
 
@@ -875,6 +882,21 @@ export default class MainController implements vscode.Disposable {
         return designer;
     }
 
+    private async openDabDesigner(connectionUri: string, database: string) {
+        const designer = await SchemaDesignerWebviewManager.getInstance().getSchemaDesigner(
+            this._context,
+            this._vscodeWrapper,
+            this,
+            this.schemaDesignerService,
+            database,
+            undefined,
+            connectionUri,
+        );
+        designer.showView(SchemaDesigner.SchemaDesignerActiveView.Dab);
+        designer.revealToForeground();
+        return designer;
+    }
+
     /**
      * Returns a flag indicating if the extension is initialized
      */

extensions/mssql/src/controllers/mainController.ts

@@ -10,7 +10,7 @@ export const schemaDesignerAgentPrompt = `Use mssql_schema_designer for the acti
 Run get_overview, briefly summarize the schema, then ask what I want to explore or change.
 Do not apply edits until I confirm.`;
 
-export const dabAgentPrompt = `Use mssql_dab for the active Data API Builder configuration.
+export const dabAgentPrompt = `Use mssql_dab for the Data API Builder configuration in the active schema designer.
 Run get_state, briefly summarize the current configuration, then ask what I want to review or change.
 Prefer mssql_dab for configuration tasks, and only use mssql_schema_designer if schema context is needed.
 Do not apply changes until I confirm.`;

extensions/mssql/src/copilot/prompts.ts

@@ -7,6 +7,7 @@ import * as vscode from "vscode";
 import { ToolBase } from "./toolBase";
 import * as Constants from "../../constants/constants";
 import { MssqlChatAgent as loc } from "../../constants/locConstants";
+import ConnectionManager from "../../controllers/connectionManager";
 import { SchemaDesignerWebviewManager } from "../../schemaDesigner/schemaDesignerWebviewManager";
 import { SchemaDesignerWebviewController } from "../../schemaDesigner/schemaDesignerWebviewController";
 import { SchemaDesigner } from "../../sharedInterfaces/schemaDesigner";
@@ -15,9 +16,10 @@ import { TelemetryActions, TelemetryViews } from "../../sharedInterfaces/telemet
 import { Dab } from "../../sharedInterfaces/dab";
 import { matchesStrictTargetHint, ToolTargetHint } from "./toolsUtils";
 
-type DabToolOperation = "get_state" | "apply_changes";
+type DabToolOperation = "show" | "get_state" | "apply_changes";
 
 export type DabToolParams =
+    | { operation: "show"; connectionId: string }
     | { operation: "get_state" }
     | {
           operation: "apply_changes";
@@ -84,6 +86,16 @@ interface DabToolError {
 export class DabTool extends ToolBase<DabToolParams> {
     public readonly toolName = Constants.copilotDabToolName;
 
+    constructor(
+        private _connectionManager: ConnectionManager,
+        private _showDab: (
+            connectionUri: string,
+            database: string,
+        ) => Promise<SchemaDesignerWebviewController>,
+    ) {
+        super();
+    }
+
     async call(
         options: vscode.LanguageModelToolInvocationOptions<DabToolParams>,
         _token: vscode.CancellationToken,
@@ -171,13 +183,57 @@ export class DabTool extends ToolBase<DabToolParams> {
         });
 
         try {
+            if (operation === "show") {
+                const { connectionId } = options.input;
+                if (!connectionId) {
+                    const err: DabToolError = {
+                        success: false,
+                        reason: "invalid_request",
+                        message: loc.dabToolMissingConnectionId,
+                    };
+                    sendToolTelemetry({ operation, success: false, reason: err.reason });
+                    return json(err);
+                }
+
+                const connInfo = this._connectionManager.getConnectionInfo(connectionId);
+                const connCreds = connInfo?.credentials;
+                if (!connCreds) {
+                    const err: DabToolError = {
+                        success: false,
+                        reason: "invalid_request",
+                        message: loc.noConnectionError(connectionId),
+                    };
+                    sendToolTelemetry({ operation, success: false, reason: err.reason });
+                    return json(err);
+                }
+
+                const designer = await this._showDab(connectionId, connCreds.database);
+                sendToolTelemetry({
+                    operation,
+                    success: true,
+                });
+
+                return json(
+                    withTarget(
+                        {
+                            success: true,
+                            message: loc.dabToolShowSuccessMessage,
+                            recommendedTool: this.toolName,
+                            recommendedNextCall: {
+                                operation: "get_state",
+                            },
+                        },
+                        designer,
+                    ),
+                );
+            }
+
             const activeDesigner = schemaDesignerManager.getActiveDesigner();
             if (!activeDesigner) {
                 const err: DabToolError = {
                     success: false,
                     reason: "no_active_designer",
-                    message:
-                        "No active schema designer found. Please open a schema designer first.",
+                    message: loc.dabToolNoActiveDesigner,
                 };
                 sendToolTelemetry({ operation, success: false, reason: err.reason });
                 return json(err);

extensions/mssql/src/copilot/tools/dabTool.ts

@@ -47,7 +47,8 @@ interface DabProviderProps {
 
 export const DabProvider: React.FC<DabProviderProps> = ({ children }) => {
     const schemaDesignerContext = useContext(SchemaDesignerContext);
-    const { extensionRpc, extractSchema, isInitialized } = schemaDesignerContext;
+    const { extensionRpc, extractSchema, isInitialized, isInitializedRef, waitForInitialization } =
+        schemaDesignerContext;
 
     const [dabConfig, setDabConfig] = useState<Dab.DabConfig | null>(null);
     const [dabTextFilter, setDabTextFilter] = useState<string>("");
@@ -57,17 +58,12 @@ export const DabProvider: React.FC<DabProviderProps> = ({ children }) => {
     );
 
     const dabConfigRef = useRef<Dab.DabConfig | null>(dabConfig);
-    const isInitializedRef = useRef<boolean>(isInitialized);
     const extractSchemaRef = useRef<() => ReturnType<typeof extractSchema>>(extractSchema);
 
     useEffect(() => {
         dabConfigRef.current = dabConfig;
     }, [dabConfig]);
 
-    useEffect(() => {
-        isInitializedRef.current = isInitialized;
-    }, [isInitialized]);
-
     useEffect(() => {
         extractSchemaRef.current = extractSchema;
     }, [extractSchema]);
@@ -76,13 +72,14 @@ export const DabProvider: React.FC<DabProviderProps> = ({ children }) => {
         registerSchemaDesignerDabToolHandlers({
             extensionRpc,
             isInitializedRef,
+            waitForInitialization,
             getCurrentDabConfig: () => dabConfigRef.current,
             getCurrentSchemaTables: () => extractSchemaRef.current().tables,
             commitDabConfig: (config) => {
                 setDabConfig(config);
             },
         });
-    }, [extensionRpc]);
+    }, [extensionRpc, waitForInitialization]);
 
     const initializeDabConfig = useCallback(() => {
         const schema = extractSchema();

extensions/mssql/src/reactviews/pages/SchemaDesigner/dab/dabContext.tsx

@@ -1266,6 +1266,28 @@ function resolveEntityRef(
     };
 }
 
+function createDabValidationError(message: string): {
+    success: false;
+    reason: DabApplyFailureReason;
+    message: string;
+} {
+    return {
+        success: false,
+        reason: "validation_error",
+        message,
+    };
+}
+
+function createEntityWithEnabledActions(
+    entity: Dab.DabEntityConfig,
+    enabledActions: Dab.EntityAction[],
+): Dab.DabEntityConfig {
+    return {
+        ...entity,
+        enabledActions: [...enabledActions],
+    };
+}
+
 function applyDabToolChange(
     config: Dab.DabConfig,
     change: Dab.DabToolChange,
@@ -1319,33 +1341,21 @@ function applyDabToolChange(
                 return resolvedEntity;
             }
 
-            if (!Array.isArray(change.actions) || change.actions.length === 0) {
-                return {
-                    success: false,
-                    reason: "validation_error",
-                    message: "actions must be a non-empty array.",
-                };
+            if (!Array.isArray(change.enabledActions) || change.enabledActions.length === 0) {
+                return createDabValidationError("enabledActions must be a non-empty array.");
             }
-            const uniqueActions = new Set(change.actions);
-            if (uniqueActions.size !== change.actions.length) {
-                return {
-                    success: false,
-                    reason: "validation_error",
-                    message: "actions must be unique.",
-                };
+            const uniqueActions = new Set(change.enabledActions);
+            if (uniqueActions.size !== change.enabledActions.length) {
+                return createDabValidationError("enabledActions must be unique.");
             }
-            if (change.actions.some((action) => !allowedActions.has(action))) {
-                return {
-                    success: false,
-                    reason: "validation_error",
-                    message: "actions contains unsupported values.",
-                };
+            if (change.enabledActions.some((action) => !allowedActions.has(action))) {
+                return createDabValidationError("enabledActions contains unsupported values.");
             }
 
-            config.entities[resolvedEntity.index] = {
-                ...resolvedEntity.entity,
-                enabledActions: [...change.actions],
-            };
+            config.entities[resolvedEntity.index] = createEntityWithEnabledActions(
+                resolvedEntity.entity,
+                change.enabledActions,
+            );
             return { success: true };
         }
 
@@ -1509,21 +1519,26 @@ function applyDabToolChange(
 export function registerSchemaDesignerDabToolHandlers(params: {
     extensionRpc: WebviewRpc<SchemaDesigner.SchemaDesignerReducers>;
     isInitializedRef: { current: boolean };
+    waitForInitialization: () => Promise<boolean>;
     getCurrentDabConfig: () => Dab.DabConfig | null;
     getCurrentSchemaTables: () => SchemaDesigner.Table[];
     commitDabConfig: (config: Dab.DabConfig) => void;
 }) {
     const {
         extensionRpc,
         isInitializedRef,
+        waitForInitialization,
         getCurrentDabConfig,
         getCurrentSchemaTables,
         commitDabConfig,
     } = params;
 
     const handleGetState = async (): Promise<Dab.GetDabToolStateResponse> => {
         if (!isInitializedRef.current) {
-            throw new Error(locConstants.schemaDesigner.schemaDesignerNotInitialized);
+            const initialized = await waitForInitialization();
+            if (!initialized || !isInitializedRef.current) {
+                throw new Error(locConstants.schemaDesigner.schemaDesignerNotInitialized);
+            }
         }
 
         const baseSnapshot = getCurrentDabConfig();

extensions/mssql/src/reactviews/pages/SchemaDesigner/schemaDesignerRpcHandlers.ts

@@ -30,6 +30,8 @@ import { stateStack } from "./schemaDesignerUndoState";
 
 export interface SchemaDesignerContextProps extends CoreRPCs {
     extensionRpc: WebviewRpc<SchemaDesigner.SchemaDesignerReducers>;
+    isInitializedRef: { current: boolean };
+    waitForInitialization: () => Promise<boolean>;
     schemaNames: string[];
     datatypes: string[];
     findTableText: string;
@@ -535,6 +537,8 @@ const SchemaDesignerStateProvider: React.FC<SchemaDesignerProviderProps> = ({ ch
             value={{
                 ...getCoreRPCs(extensionRpc),
                 extensionRpc,
+                isInitializedRef,
+                waitForInitialization,
                 schemaNames,
                 datatypes,
                 findTableText,

extensions/mssql/src/reactviews/pages/SchemaDesigner/schemaDesignerStateProvider.tsx

@@ -229,7 +229,7 @@ export namespace Dab {
     export type DabToolChange =
         | { type: "set_api_types"; apiTypes: ApiType[] }
         | { type: "set_entity_enabled"; entity: DabEntityRef; isEnabled: boolean }
-        | { type: "set_entity_actions"; entity: DabEntityRef; actions: EntityAction[] }
+        | { type: "set_entity_actions"; entity: DabEntityRef; enabledActions: EntityAction[] }
         | { type: "patch_entity_settings"; entity: DabEntityRef; set: DabEntitySettingsPatch }
         | { type: "set_only_enabled_entities"; entities: DabEntityRef[] }
         | { type: "set_all_entities_enabled"; isEnabled: boolean };