Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
10 changes: 10 additions & 0 deletions extensions/mssql/l10n/bundle.l10n.json
Original file line number Diff line number Diff line change
Expand Up @@ -2654,6 +2654,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"]
Expand All @@ -2671,6 +2679,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.",
Expand Down
424 changes: 210 additions & 214 deletions extensions/mssql/package.json

Large diffs are not rendered by default.

15 changes: 15 additions & 0 deletions extensions/mssql/src/constants/locConstants.ts
Original file line number Diff line number Diff line change
Expand Up @@ -2031,6 +2031,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({
Expand Down Expand Up @@ -2061,6 +2070,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.",
);
Expand Down
24 changes: 23 additions & 1 deletion extensions/mssql/src/controllers/mainController.ts
Original file line number Diff line number Diff line change
Expand Up @@ -855,7 +855,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),
),
),
);
}

Expand All @@ -874,6 +881,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
*/
Expand Down
2 changes: 1 addition & 1 deletion extensions/mssql/src/copilot/prompts.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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.`;
Expand Down
62 changes: 59 additions & 3 deletions extensions/mssql/src/copilot/tools/dabTool.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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";
Expand All @@ -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";
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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);
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -46,7 +46,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>("");
Expand All @@ -56,17 +57,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]);
Expand All @@ -75,13 +71,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();
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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 };
}

Expand Down Expand Up @@ -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();
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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;
Expand Down Expand Up @@ -535,6 +537,8 @@ const SchemaDesignerStateProvider: React.FC<SchemaDesignerProviderProps> = ({ ch
value={{
...getCoreRPCs(extensionRpc),
extensionRpc,
isInitializedRef,
waitForInitialization,
schemaNames,
datatypes,
findTableText,
Expand Down
2 changes: 1 addition & 1 deletion extensions/mssql/src/sharedInterfaces/dab.ts
Original file line number Diff line number Diff line change
Expand Up @@ -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 };
Expand Down
Loading
Loading