Add support for auto-model provenance

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

@@ -110,9 +110,10 @@ export function parsePredictedClassifications(
       // For now, model any method for which none of its arguments are modeled as sinks as neutral
       modeledMethods[signature] = {
         type: "neutral",
-        kind: "",
+        kind: "summary",
         input: "",
         output: "",
+        provenance: "ai-generated",
       };
       continue;
     }
@@ -129,6 +130,7 @@ export function parsePredictedClassifications(
       kind: sink.classification?.kind ?? "",
       input: sink.input ?? "",
       output: sink.output ?? "",
+      provenance: "ai-generated",
     };
   }
 

extensions/ql-vscode/src/data-extensions-editor/modeled-method.ts

@@ -10,6 +10,10 @@ export type Provenance =
   | "df-generated"
   // Generated by the dataflow model and manually edited
   | "df-manual"
+  // Generated by the auto-model
+  | "ai-generated"
+  // Generated by the auto-model and manually edited
+  | "ai-manual"
   // Entered by the user in the editor manually
   | "manual";
 

extensions/ql-vscode/src/view/data-extensions-editor/MethodRow.tsx

@@ -13,6 +13,7 @@ import { ExternalApiUsage } from "../../data-extensions-editor/external-api-usag
 import {
   ModeledMethod,
   ModeledMethodType,
+  Provenance,
 } from "../../data-extensions-editor/modeled-method";
 import { KindInput } from "./KindInput";
 import { extensiblePredicateDefinitions } from "../../data-extensions-editor/predicates";
@@ -63,6 +64,13 @@ export const MethodRow = ({
     (e: InputEvent) => {
       const target = e.target as HTMLSelectElement;
 
+      let newProvenance: Provenance = "manual";
+      if (modeledMethod?.provenance === "df-generated") {
+        newProvenance = "df-manual";
+      } else if (modeledMethod?.provenance === "ai-generated") {
+        newProvenance = "ai-manual";
+      }
+
       onChange(externalApiUsage, {
         // If there are no arguments, we will default to "Argument[this]"
         input: argumentsList.length === 0 ? "Argument[this]" : "Argument[0]",
@@ -70,8 +78,7 @@ export const MethodRow = ({
         kind: "value",
         ...modeledMethod,
         type: target.value as ModeledMethodType,
-        provenance:
-          modeledMethod?.provenance === "df-generated" ? "df-manual" : "manual",
+        provenance: newProvenance,
       });
     },
     [onChange, externalApiUsage, modeledMethod, argumentsList],

extensions/ql-vscode/test/unit-tests/data-extensions-editor/auto-model.test.ts

@@ -191,12 +191,14 @@ describe("createAutoModelRequest", () => {
       kind: "",
       input: "",
       output: "",
+      provenance: "manual",
     },
     "org.sql2o.Sql2o#Sql2o(String)": {
       type: "sink",
       kind: "jndi-injection",
       input: "Argument[0]",
       output: "",
+      provenance: "manual",
     },
   };
 
@@ -407,18 +409,21 @@ describe("parsePredictedClassifications", () => {
         kind: "sql injection sink",
         input: "Argument[0]",
         output: "",
+        provenance: "ai-generated",
       },
       "org.sql2o.Sql2o#executeScalar(Class)": {
         type: "neutral",
-        kind: "",
+        kind: "summary",
         input: "",
         output: "",
+        provenance: "ai-generated",
       },
       "org.sql2o.Sql2o#Sql2o(String,String,String)": {
         type: "sink",
         kind: "sql injection sink",
         input: "Argument[1]",
         output: "",
+        provenance: "ai-generated",
       },
     });
   });