Skip to content

Commit 3afed56

Browse files
koesie10koesie10
committed
Convert remaining extension host code to handle multiple models
This converts all remaining extension host code to handle multiple models per method. The only place where we're using the legacy format is in the webview and in the boundary between the webview and the extension host.
1 parent c5175e0 commit 3afed56

12 files changed

Lines changed: 182 additions & 87 deletions

File tree

‎extensions/ql-vscode/src/model-editor/auto-model.ts‎

Lines changed: 5 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -14,13 +14,13 @@ import { groupMethods, sortGroupNames, sortMethods } from "./shared/sorting";
1414
* the order in the UI.
1515
* @param mode Whether it is application or framework mode.
1616
* @param methods all methods.
17-
* @param modeledMethods the currently modeled methods.
17+
* @param modeledMethodsBySignature the currently modeled methods.
1818
* @returns list of modeled methods that are candidates for modeling.
1919
*/
2020
export function getCandidates(
2121
mode: Mode,
2222
methods: Method[],
23-
modeledMethods: Record<string, ModeledMethod>,
23+
modeledMethodsBySignature: Record<string, ModeledMethod[]>,
2424
): MethodSignature[] {
2525
// Sort the same way as the UI so we send the first ones listed in the UI first
2626
const grouped = groupMethods(methods, mode);
@@ -32,12 +32,11 @@ export function getCandidates(
3232
const candidates: MethodSignature[] = [];
3333

3434
for (const method of sortedMethods) {
35-
const modeledMethod: ModeledMethod = modeledMethods[method.signature] ?? {
36-
type: "none",
37-
};
35+
const modeledMethods: ModeledMethod[] =
36+
modeledMethodsBySignature[method.signature] ?? [];
3837

3938
// Anything that is modeled is not a candidate
40-
if (modeledMethod.type !== "none") {
39+
if (modeledMethods.some((m) => m.type !== "none")) {
4140
continue;
4241
}
4342

‎extensions/ql-vscode/src/model-editor/auto-modeler.ts‎

Lines changed: 19 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,6 @@ import { QueryRunner } from "../query-server";
1616
import { DatabaseItem } from "../databases/local-databases";
1717
import { Mode } from "./shared/mode";
1818
import { CancellationTokenSource } from "vscode";
19-
import { convertToLegacyModeledMethods } from "./modeled-methods-legacy";
2019

2120
// Limit the number of candidates we send to the model in each request
2221
// to avoid long requests.
@@ -43,7 +42,7 @@ export class AutoModeler {
4342
inProgressMethods: string[],
4443
) => Promise<void>,
4544
private readonly addModeledMethods: (
46-
modeledMethods: Record<string, ModeledMethod>,
45+
modeledMethods: Record<string, ModeledMethod[]>,
4746
) => Promise<void>,
4847
) {
4948
this.jobs = new Map<string, CancellationTokenSource>();
@@ -60,7 +59,7 @@ export class AutoModeler {
6059
public async startModeling(
6160
packageName: string,
6261
methods: Method[],
63-
modeledMethods: Record<string, ModeledMethod>,
62+
modeledMethods: Record<string, ModeledMethod[]>,
6463
mode: Mode,
6564
): Promise<void> {
6665
if (this.jobs.has(packageName)) {
@@ -107,7 +106,7 @@ export class AutoModeler {
107106
private async modelPackage(
108107
packageName: string,
109108
methods: Method[],
110-
modeledMethods: Record<string, ModeledMethod>,
109+
modeledMethods: Record<string, ModeledMethod[]>,
111110
mode: Mode,
112111
cancellationTokenSource: CancellationTokenSource,
113112
): Promise<void> {
@@ -193,31 +192,31 @@ export class AutoModeler {
193192
filename: "auto-model.yml",
194193
});
195194

196-
const rawLoadedMethods = loadDataExtensionYaml(models);
197-
if (!rawLoadedMethods) {
195+
const loadedMethods = loadDataExtensionYaml(models);
196+
if (!loadedMethods) {
198197
return;
199198
}
200199

201-
const loadedMethods = convertToLegacyModeledMethods(rawLoadedMethods);
202-
203200
// Any candidate that was part of the response is a negative result
204201
// meaning that the canidate is not a sink for the kinds that the LLM is checking for.
205202
// For now we model this as a sink neutral method, however this is subject
206203
// to discussion.
207204
for (const candidate of candidateMethods) {
208205
if (!(candidate.signature in loadedMethods)) {
209-
loadedMethods[candidate.signature] = {
210-
type: "neutral",
211-
kind: "sink",
212-
input: "",
213-
output: "",
214-
provenance: "ai-generated",
215-
signature: candidate.signature,
216-
packageName: candidate.packageName,
217-
typeName: candidate.typeName,
218-
methodName: candidate.methodName,
219-
methodParameters: candidate.methodParameters,
220-
};
206+
loadedMethods[candidate.signature] = [
207+
{
208+
type: "neutral",
209+
kind: "sink",
210+
input: "",
211+
output: "",
212+
provenance: "ai-generated",
213+
signature: candidate.signature,
214+
packageName: candidate.packageName,
215+
typeName: candidate.typeName,
216+
methodName: candidate.methodName,
217+
methodParameters: candidate.methodParameters,
218+
},
219+
];
221220
}
222221
}
223222

‎extensions/ql-vscode/src/model-editor/method-modeling/method-modeling-view-provider.ts‎

Lines changed: 17 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,10 @@ import { DbModelingState, ModelingStore } from "../modeling-store";
1212
import { AbstractWebviewViewProvider } from "../../common/vscode/abstract-webview-view-provider";
1313
import { assertNever } from "../../common/helpers-pure";
1414
import { ModelEditorViewTracker } from "../model-editor-view-tracker";
15+
import {
16+
convertFromLegacyModeledMethod,
17+
convertToLegacyModeledMethod,
18+
} from "../modeled-methods-legacy";
1519

1620
export class MethodModelingViewProvider extends AbstractWebviewViewProvider<
1721
ToMethodModelingMessage,
@@ -51,7 +55,9 @@ export class MethodModelingViewProvider extends AbstractWebviewViewProvider<
5155
void this.postMessage({
5256
t: "setSelectedMethod",
5357
method: selectedMethod.method,
54-
modeledMethod: selectedMethod.modeledMethod,
58+
modeledMethod: convertToLegacyModeledMethod(
59+
selectedMethod.modeledMethods,
60+
),
5561
isModified: selectedMethod.isModified,
5662
});
5763
}
@@ -82,9 +88,10 @@ export class MethodModelingViewProvider extends AbstractWebviewViewProvider<
8288
case "setModeledMethod": {
8389
const activeState = this.ensureActiveState();
8490

85-
this.modelingStore.updateModeledMethod(
91+
this.modelingStore.updateModeledMethods(
8692
activeState.databaseItem,
87-
msg.method,
93+
msg.method.signature,
94+
convertFromLegacyModeledMethod(msg.method),
8895
);
8996
break;
9097
}
@@ -123,11 +130,11 @@ export class MethodModelingViewProvider extends AbstractWebviewViewProvider<
123130
this.push(
124131
this.modelingStore.onModeledMethodsChanged(async (e) => {
125132
if (this.webviewView && e.isActiveDb) {
126-
const modeledMethod = e.modeledMethods[this.method?.signature ?? ""];
127-
if (modeledMethod) {
128-
await this.webviewView.webview.postMessage({
133+
const modeledMethods = e.modeledMethods[this.method?.signature ?? ""];
134+
if (modeledMethods) {
135+
await this.postMessage({
129136
t: "setModeledMethod",
130-
method: modeledMethod,
137+
method: convertToLegacyModeledMethod(modeledMethods),
131138
});
132139
}
133140
}
@@ -138,7 +145,7 @@ export class MethodModelingViewProvider extends AbstractWebviewViewProvider<
138145
this.modelingStore.onModifiedMethodsChanged(async (e) => {
139146
if (this.webviewView && e.isActiveDb && this.method) {
140147
const isModified = e.modifiedMethods.has(this.method.signature);
141-
await this.webviewView.webview.postMessage({
148+
await this.postMessage({
142149
t: "setMethodModified",
143150
isModified,
144151
});
@@ -150,10 +157,10 @@ export class MethodModelingViewProvider extends AbstractWebviewViewProvider<
150157
this.modelingStore.onSelectedMethodChanged(async (e) => {
151158
if (this.webviewView) {
152159
this.method = e.method;
153-
await this.webviewView.webview.postMessage({
160+
await this.postMessage({
154161
t: "setSelectedMethod",
155162
method: e.method,
156-
modeledMethod: e.modeledMethod,
163+
modeledMethod: convertToLegacyModeledMethod(e.modeledMethods),
157164
isModified: e.isModified,
158165
});
159166
}

‎extensions/ql-vscode/src/model-editor/methods-usage/methods-usage-data-provider.ts‎

Lines changed: 8 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,7 @@ import { DatabaseItem } from "../../databases/local-databases";
1414
import { relative } from "path";
1515
import { CodeQLCliServer } from "../../codeql-cli/cli";
1616
import { INITIAL_HIDE_MODELED_METHODS_VALUE } from "../shared/hide-modeled-methods";
17-
import { getModelingStatus } from "../shared/modeling-status";
17+
import { getModelingStatusForModeledMethods } from "../shared/modeling-status";
1818
import { assertNever } from "../../common/helpers-pure";
1919
import { ModeledMethod } from "../modeled-method";
2020

@@ -26,7 +26,7 @@ export class MethodsUsageDataProvider
2626
private databaseItem: DatabaseItem | undefined = undefined;
2727
private sourceLocationPrefix: string | undefined = undefined;
2828
private hideModeledMethods: boolean = INITIAL_HIDE_MODELED_METHODS_VALUE;
29-
private modeledMethods: Record<string, ModeledMethod> = {};
29+
private modeledMethods: Record<string, ModeledMethod[]> = {};
3030
private modifiedMethodSignatures: Set<string> = new Set();
3131

3232
private readonly onDidChangeTreeDataEmitter = this.push(
@@ -52,7 +52,7 @@ export class MethodsUsageDataProvider
5252
methods: Method[],
5353
databaseItem: DatabaseItem,
5454
hideModeledMethods: boolean,
55-
modeledMethods: Record<string, ModeledMethod>,
55+
modeledMethods: Record<string, ModeledMethod[]>,
5656
modifiedMethodSignatures: Set<string>,
5757
): Promise<void> {
5858
if (
@@ -99,10 +99,13 @@ export class MethodsUsageDataProvider
9999
}
100100

101101
private getModelingStatusIcon(method: Method): ThemeIcon {
102-
const modeledMethod = this.modeledMethods[method.signature];
102+
const modeledMethods = this.modeledMethods[method.signature];
103103
const modifiedMethod = this.modifiedMethodSignatures.has(method.signature);
104104

105-
const status = getModelingStatus(modeledMethod, modifiedMethod);
105+
const status = getModelingStatusForModeledMethods(
106+
modeledMethods,
107+
modifiedMethod,
108+
);
106109
switch (status) {
107110
case "unmodeled":
108111
return new ThemeIcon("error", new ThemeColor("errorForeground"));

‎extensions/ql-vscode/src/model-editor/methods-usage/methods-usage-panel.ts‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -34,7 +34,7 @@ export class MethodsUsagePanel extends DisposableObject {
3434
methods: Method[],
3535
databaseItem: DatabaseItem,
3636
hideModeledMethods: boolean,
37-
modeledMethods: Record<string, ModeledMethod>,
37+
modeledMethods: Record<string, ModeledMethod[]>,
3838
modifiedMethodSignatures: Set<string>,
3939
): Promise<void> {
4040
await this.dataProvider.setState(

‎extensions/ql-vscode/src/model-editor/model-editor-view.ts‎

Lines changed: 25 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -44,6 +44,7 @@ import { telemetryListener } from "../common/vscode/telemetry";
4444
import { ModelingStore } from "./modeling-store";
4545
import { ModelEditorViewTracker } from "./model-editor-view-tracker";
4646
import {
47+
convertFromLegacyModeledMethod,
4748
convertFromLegacyModeledMethods,
4849
convertToLegacyModeledMethods,
4950
} from "./modeled-methods-legacy";
@@ -282,7 +283,7 @@ export class ModelEditorView extends AbstractWebview<
282283
await this.generateModeledMethodsFromLlm(
283284
msg.packageName,
284285
msg.methods,
285-
msg.modeledMethods,
286+
convertFromLegacyModeledMethods(msg.modeledMethods),
286287
);
287288
void telemetryListener?.sendUIInteraction(
288289
"model-editor-generate-methods-from-llm",
@@ -326,7 +327,10 @@ export class ModelEditorView extends AbstractWebview<
326327
);
327328
break;
328329
case "setModeledMethod": {
329-
this.setModeledMethod(msg.method);
330+
this.setModeledMethods(
331+
msg.method.signature,
332+
convertFromLegacyModeledMethod(msg.method),
333+
);
330334
break;
331335
}
332336
default:
@@ -385,10 +389,7 @@ export class ModelEditorView extends AbstractWebview<
385389
this.cliServer,
386390
this.app.logger,
387391
);
388-
this.modelingStore.setModeledMethods(
389-
this.databaseItem,
390-
convertToLegacyModeledMethods(modeledMethods),
391-
);
392+
this.modelingStore.setModeledMethods(this.databaseItem, modeledMethods);
392393
} catch (e: unknown) {
393394
void showAndLogErrorMessage(
394395
this.app.logger,
@@ -460,10 +461,16 @@ export class ModelEditorView extends AbstractWebview<
460461
queryStorageDir: this.queryStorageDir,
461462
databaseItem: addedDatabase ?? this.databaseItem,
462463
onResults: async (modeledMethods) => {
463-
const modeledMethodsByName: Record<string, ModeledMethod> = {};
464+
const modeledMethodsByName: Record<string, ModeledMethod[]> = {};
464465

465466
for (const modeledMethod of modeledMethods) {
466-
modeledMethodsByName[modeledMethod.signature] = modeledMethod;
467+
if (!(modeledMethod.signature in modeledMethodsByName)) {
468+
modeledMethodsByName[modeledMethod.signature] = [];
469+
}
470+
471+
modeledMethodsByName[modeledMethod.signature].push(
472+
modeledMethod,
473+
);
467474
}
468475

469476
this.addModeledMethods(modeledMethodsByName);
@@ -488,7 +495,7 @@ export class ModelEditorView extends AbstractWebview<
488495
private async generateModeledMethodsFromLlm(
489496
packageName: string,
490497
methods: Method[],
491-
modeledMethods: Record<string, ModeledMethod>,
498+
modeledMethods: Record<string, ModeledMethod[]>,
492499
): Promise<void> {
493500
await this.autoModeler.startModeling(
494501
packageName,
@@ -625,7 +632,7 @@ export class ModelEditorView extends AbstractWebview<
625632
if (event.dbUri === this.databaseItem.databaseUri.toString()) {
626633
await this.postMessage({
627634
t: "setModeledMethods",
628-
methods: event.modeledMethods,
635+
methods: convertToLegacyModeledMethods(event.modeledMethods),
629636
});
630637
}
631638
}),
@@ -643,7 +650,7 @@ export class ModelEditorView extends AbstractWebview<
643650
);
644651
}
645652

646-
private addModeledMethods(modeledMethods: Record<string, ModeledMethod>) {
653+
private addModeledMethods(modeledMethods: Record<string, ModeledMethod[]>) {
647654
this.modelingStore.addModeledMethods(this.databaseItem, modeledMethods);
648655

649656
this.modelingStore.addModifiedMethods(
@@ -652,13 +659,17 @@ export class ModelEditorView extends AbstractWebview<
652659
);
653660
}
654661

655-
private setModeledMethod(method: ModeledMethod) {
662+
private setModeledMethods(signature: string, methods: ModeledMethod[]) {
656663
const state = this.modelingStore.getStateForActiveDb();
657664
if (!state) {
658665
throw new Error("Attempting to set modeled method without active db");
659666
}
660667

661-
this.modelingStore.updateModeledMethod(state.databaseItem, method);
662-
this.modelingStore.addModifiedMethod(state.databaseItem, method.signature);
668+
this.modelingStore.updateModeledMethods(
669+
state.databaseItem,
670+
signature,
671+
methods,
672+
);
673+
this.modelingStore.addModifiedMethod(state.databaseItem, signature);
663674
}
664675
}

0 commit comments

Comments
 (0)