[UPDATE PRIMITIVE] Fix bugs 1–5 and implement design improvements 1–3, 5 from v2.25.1-next.2 evaluation

extensions/vscode/README.md

@@ -45,6 +45,7 @@ All settings are under the `codeql-mcp` namespace in VS Code settings:
 | `codeql-mcp.serverCommand`                 | `"node"`   | Command to launch the server. Override to `"npx"` or a custom path. |
 | `codeql-mcp.serverArgs`                    | `[]`       | Custom args. When empty, the bundled entry point is used.           |
 | `codeql-mcp.watchCodeqlExtension`          | `true`     | Watch for databases and results from the CodeQL extension.          |
+| `codeql-mcp.enableAnnotationTools`         | `true`     | Enable annotation, audit, and cache tools.                          |
 | `codeql-mcp.additionalEnv`                 | `{}`       | Extra environment variables passed to the server process.           |
 | `codeql-mcp.additionalDatabaseDirs`        | `[]`       | Additional directories to search for CodeQL databases.              |
 | `codeql-mcp.additionalMrvaRunResultsDirs`  | `[]`       | Additional directories containing MRVA run results.                 |

extensions/vscode/package.json

@@ -90,6 +90,11 @@
           "default": true,
           "markdownDescription": "Copy CodeQL databases from the `GitHub.vscode-codeql` extension storage into a managed directory, removing query-server lock files so the MCP server CLI can operate without contention. Disable to use databases in-place (may fail when the CodeQL query server is running)."
         },
+        "codeql-mcp.enableAnnotationTools": {
+          "type": "boolean",
+          "default": true,
+          "markdownDescription": "Enable annotation, audit, and query results caching tools. When enabled, the MCP server registers `annotation_*`, `audit_*`, and `query_results_cache_*` tools. Disable to reduce the tool surface if these capabilities are not needed."
+        },
         "codeql-mcp.serverArgs": {
           "type": "array",
           "items": {

extensions/vscode/src/bridge/environment-builder.ts

@@ -147,7 +147,17 @@ export class EnvironmentBuilder extends DisposableObject {
     queryDirs.push(...userQueryDirs);
     env.CODEQL_QUERY_RUN_RESULTS_DIRS = queryDirs.join(delimiter);
 
-    // User-configured additional environment variables
+    // Annotation, audit, and cache tools — enabled by default (Design 5).
+    // The setting controls ENABLE_ANNOTATION_TOOLS and defaults
+    // MONITORING_STORAGE_LOCATION to the scratch directory so tools work
+    // out-of-the-box without manual env var configuration.
+    const enableAnnotations = config.get<boolean>('enableAnnotationTools', true);
+    env.ENABLE_ANNOTATION_TOOLS = enableAnnotations ? 'true' : 'false';
+    if (enableAnnotations && env.CODEQL_MCP_SCRATCH_DIR) {
+      env.MONITORING_STORAGE_LOCATION = env.CODEQL_MCP_SCRATCH_DIR;
+    }
+
+    // User-configured additional environment variables (overrides above defaults)
     const additionalEnv = config.get<Record<string, string>>('additionalEnv', {});
     for (const [key, value] of Object.entries(additionalEnv)) {
       env[key] = value;

extensions/vscode/test/bridge/environment-builder.test.ts

@@ -272,4 +272,70 @@ describe('EnvironmentBuilder', () => {
   it('should be disposable', () => {
     expect(() => builder.dispose()).not.toThrow();
   });
+
+  it('should set ENABLE_ANNOTATION_TOOLS=true by default', async () => {
+    const env = await builder.build();
+    expect(env.ENABLE_ANNOTATION_TOOLS).toBe('true');
+  });
+
+  it('should set ENABLE_ANNOTATION_TOOLS=false when setting is disabled', async () => {
+    const vscode = await import('vscode');
+    const originalGetConfig = vscode.workspace.getConfiguration;
+    vscode.workspace.getConfiguration = () => ({
+      get: (_key: string, defaultVal?: any) => {
+        if (_key === 'enableAnnotationTools') return false;
+        if (_key === 'additionalDatabaseDirs') return [];
+        if (_key === 'additionalQueryRunResultsDirs') return [];
+        if (_key === 'additionalMrvaRunResultsDirs') return [];
+        return defaultVal;
+      },
+      has: () => false,
+      inspect: () => undefined as any,
+      update: () => Promise.resolve(),
+    }) as any;
+
+    builder.invalidate();
+    const env = await builder.build();
+    expect(env.ENABLE_ANNOTATION_TOOLS).toBe('false');
+
+    vscode.workspace.getConfiguration = originalGetConfig;
+  });
+
+  it('should set MONITORING_STORAGE_LOCATION to scratch dir when annotations enabled with workspace', async () => {
+    const vscode = await import('vscode');
+    const origFolders = vscode.workspace.workspaceFolders;
+    (vscode.workspace.workspaceFolders as any) = [
+      { uri: { fsPath: '/mock/workspace' }, name: 'ws', index: 0 },
+    ];
+
+    builder.invalidate();
+    const env = await builder.build();
+    expect(env.MONITORING_STORAGE_LOCATION).toBe('/mock/workspace/.codeql/ql-mcp');
+
+    (vscode.workspace.workspaceFolders as any) = origFolders;
+  });
+
+  it('should allow additionalEnv to override ENABLE_ANNOTATION_TOOLS', async () => {
+    const vscode = await import('vscode');
+    const originalGetConfig = vscode.workspace.getConfiguration;
+    vscode.workspace.getConfiguration = () => ({
+      get: (_key: string, defaultVal?: any) => {
+        if (_key === 'additionalEnv') return { ENABLE_ANNOTATION_TOOLS: 'false' };
+        if (_key === 'additionalDatabaseDirs') return [];
+        if (_key === 'additionalQueryRunResultsDirs') return [];
+        if (_key === 'additionalMrvaRunResultsDirs') return [];
+        return defaultVal;
+      },
+      has: () => false,
+      inspect: () => undefined as any,
+      update: () => Promise.resolve(),
+    }) as any;
+
+    builder.invalidate();
+    const env = await builder.build();
+    // additionalEnv comes after the default, so it should override
+    expect(env.ENABLE_ANNOTATION_TOOLS).toBe('false');
+
+    vscode.workspace.getConfiguration = originalGetConfig;
+  });
 });

server/dist/codeql-development-mcp-server.js

@@ -183584,8 +183584,9 @@ var SqliteStore = class _SqliteStore {
       params.$entity_key_prefix = filter.entityKeyPrefix + "%";
     }
     if (filter?.search) {
-      conditions.push("id IN (SELECT rowid FROM annotations_fts WHERE annotations_fts MATCH $search)");
+      conditions.push("(id IN (SELECT rowid FROM annotations_fts WHERE annotations_fts MATCH $search) OR category = $search_cat)");
       params.$search = filter.search;
+      params.$search_cat = filter.search;
     }
     let sql = "SELECT * FROM annotations";
     if (conditions.length > 0) {
@@ -184496,6 +184497,14 @@ Interpreted output saved to: ${outputFilePath}`;
               queryPath
             });
             const store = sessionDataManager.getStore();
+            let resultCount = null;
+            if (outputFormat.includes("sarif")) {
+              try {
+                const sarif = JSON.parse(resultContent);
+                resultCount = sarif?.runs?.[0]?.results?.length ?? null;
+              } catch {
+              }
+            }
             store.putCacheEntry({
               bqrsPath,
               cacheKey: cacheKey2,
@@ -184507,7 +184516,8 @@ Interpreted output saved to: ${outputFilePath}`;
               outputFormat,
               queryName: queryName || basename4(queryPath, ".ql"),
               queryPath,
-              resultContent
+              resultContent,
+              resultCount
             });
             enhancedOutput += `
 Results cached with key: ${cacheKey2}`;
@@ -184574,6 +184584,54 @@ Warning: Query processing error - ${error2}`
     };
   }
 }
+function cacheDatabaseAnalyzeResults(params, logger2) {
+  try {
+    const config2 = sessionDataManager.getConfig();
+    if (!config2.enableAnnotationTools) return;
+    const outputPath = params.output;
+    const format = params.format;
+    const dbPath = params.database;
+    const queries = params.queries;
+    if (!outputPath || !format || !dbPath) return;
+    if (!format.includes("sarif")) return;
+    let resultContent;
+    try {
+      resultContent = readFileSync6(outputPath, "utf8");
+    } catch {
+      return;
+    }
+    const codeqlVersion = getActualCodeqlVersion();
+    const queryName = queries ? basename4(queries, ".qls") : "database-analyze";
+    let resultCount = null;
+    try {
+      const sarif = JSON.parse(resultContent);
+      resultCount = sarif?.runs?.[0]?.results?.length ?? null;
+    } catch {
+    }
+    const cacheKey2 = computeQueryCacheKey({
+      codeqlVersion,
+      databasePath: dbPath,
+      outputFormat: format,
+      queryPath: queries || "database-analyze"
+    });
+    const store = sessionDataManager.getStore();
+    store.putCacheEntry({
+      cacheKey: cacheKey2,
+      codeqlVersion,
+      databasePath: dbPath,
+      interpretedPath: outputPath,
+      language: "unknown",
+      outputFormat: format,
+      queryName,
+      queryPath: queries || "database-analyze",
+      resultContent,
+      resultCount
+    });
+    logger2.info(`Cached database-analyze results with key: ${cacheKey2} (${resultCount ?? 0} results)`);
+  } catch (err) {
+    logger2.error("Failed to cache database-analyze results:", err);
+  }
+}
 
 // src/lib/cli-tool-registry.ts
 init_package_paths();
@@ -187168,6 +187226,21 @@ var safeDump = renamed("safeDump", "dump");
 
 // src/lib/cli-tool-registry.ts
 init_temp_dir();
+var databaseLocks = /* @__PURE__ */ new Map();
+function acquireDatabaseLock(dbPath) {
+  const key = dbPath;
+  const previous = databaseLocks.get(key) ?? Promise.resolve();
+  let release;
+  const gate = new Promise((resolve15) => {
+    release = resolve15;
+  });
+  databaseLocks.set(key, gate);
+  return {
+    ready: previous.then(() => {
+    }),
+    release
+  };
+}
 var defaultCLIResultProcessor = (result, _params) => {
   if (!result.success) {
     return `Command failed (exit code ${result.exitCode || "unknown"}):
@@ -187407,6 +187480,13 @@ function registerCLITool(server, definition) {
             }
             break;
           }
+          case "codeql_bqrs_interpret":
+            if (options.database) {
+              const dbPath = resolveDatabasePath(options.database);
+              options["source-archive"] = join10(dbPath, "src");
+              delete options.database;
+            }
+            break;
           case "codeql_query_compile":
           case "codeql_resolve_metadata":
             if (query) {
@@ -187501,7 +187581,16 @@ function registerCLITool(server, definition) {
           if (name === "codeql_test_run") {
             options["keep-databases"] = true;
           }
-          result = await executeCodeQLCommand(subcommand, options, [...positionalArgs, ...userAdditionalArgs], cwd);
+          let dbLock;
+          if (name === "codeql_database_analyze" && positionalArgs.length > 0) {
+            dbLock = acquireDatabaseLock(positionalArgs[0]);
+            await dbLock.ready;
+          }
+          try {
+            result = await executeCodeQLCommand(subcommand, options, [...positionalArgs, ...userAdditionalArgs], cwd);
+          } finally {
+            dbLock?.release();
+          }
         } else if (command === "qlt") {
           result = await executeQLTCommand(subcommand, options, [...positionalArgs, ...userAdditionalArgs]);
         } else {
@@ -187531,6 +187620,9 @@ function registerCLITool(server, definition) {
             }
           }
         }
+        if (name === "codeql_database_analyze" && result.success && options.output) {
+          cacheDatabaseAnalyzeResults({ ...params, output: options.output, format: options.format }, logger);
+        }
         const processedResult = resultProcessor(result, params);
         return {
           content: [{
@@ -187654,7 +187746,7 @@ var codeqlBqrsInfoTool = {
   command: "codeql",
   subcommand: "bqrs info",
   inputSchema: {
-    files: external_exports.array(external_exports.string()).describe("BQRS file(s) to examine"),
+    file: external_exports.string().describe("BQRS file to examine"),
     format: external_exports.enum(["text", "json"]).optional().describe("Output format: text (default) or json. Use json for machine-readable output and pagination offset computation."),
     "paginate-rows": external_exports.number().optional().describe("Compute byte offsets for pagination at intervals of this many rows. Use with --format=json. Offsets can be passed to codeql_bqrs_decode --start-at."),
     "paginate-result-set": external_exports.string().optional().describe("Compute pagination offsets only for this result set name"),
@@ -187679,7 +187771,8 @@ var codeqlBqrsInterpretTool = {
     file: external_exports.string().describe("The BQRS file to interpret"),
     format: external_exports.enum(["csv", "sarif-latest", "sarifv2.1.0", "graphtext", "dgml", "dot"]).describe("Output format: csv (comma-separated), sarif-latest/sarifv2.1.0 (SARIF), graphtext/dgml/dot (graph formats, only for @kind graph queries)"),
     output: createCodeQLSchemas.output(),
-    t: external_exports.array(external_exports.string()).describe('Query metadata key=value pairs. At least "kind" and "id" must be specified (e.g., ["kind=graph", "id=js/print-ast"])'),
+    database: external_exports.string().optional().describe("Path to the CodeQL database, used to resolve source archive context for SARIF interpretation (provides file contents and snippets)"),
+    t: external_exports.array(external_exports.string()).describe('Query metadata key=value pairs in KEY=VALUE format. At least "kind" and "id" must be specified. Example: ["kind=problem", "id=js/sql-injection"]. Common keys: kind (problem|path-problem|graph|metric|diagnostic), id (query identifier like js/xss)'),
     "max-paths": external_exports.number().optional().describe("Maximum number of paths to produce for each alert with paths (default: 4)"),
     "sarif-add-file-contents": external_exports.boolean().optional().describe("[SARIF only] Include full file contents for all files referenced in results"),
     "sarif-add-snippets": external_exports.boolean().optional().describe("[SARIF only] Include code snippets for each location with context"),
@@ -193667,16 +193760,35 @@ function registerAuditListFindingsTool(server) {
 function registerAuditAddNotesTool(server) {
   server.tool(
     "audit_add_notes",
-    "Append notes to an existing audit finding. The notes are appended to the annotation content.",
+    "Append notes to an existing audit finding. Identify the finding by findingId (preferred) or by owner+repo+sourceLocation+line.",
     {
-      owner: external_exports.string().describe("Repository owner."),
-      repo: external_exports.string().describe("Repository name."),
-      sourceLocation: external_exports.string().describe("File path of the finding."),
-      line: external_exports.number().int().min(1).describe("Line number of the finding (integer >= 1)."),
+      findingId: external_exports.number().int().positive().optional().describe("Annotation ID of the finding (returned by audit_store_findings and audit_list_findings). Preferred lookup method."),
+      owner: external_exports.string().optional().describe("Repository owner (required when findingId is not provided)."),
+      repo: external_exports.string().optional().describe("Repository name (required when findingId is not provided)."),
+      sourceLocation: external_exports.string().optional().describe("File path of the finding (required when findingId is not provided)."),
+      line: external_exports.number().int().min(1).optional().describe("Line number of the finding (required when findingId is not provided)."),
       notes: external_exports.string().describe("Notes to append.")
     },
-    async ({ owner, repo, sourceLocation, line, notes }) => {
+    async ({ findingId, owner, repo, sourceLocation, line, notes }) => {
       const store = sessionDataManager.getStore();
+      if (findingId) {
+        const annotation2 = store.getAnnotation(findingId);
+        if (!annotation2 || annotation2.category !== AUDIT_CATEGORY) {
+          return {
+            content: [{ type: "text", text: `No audit finding found with id ${findingId}.` }]
+          };
+        }
+        const updatedContent2 = (annotation2.content || "") + "\n" + notes;
+        store.updateAnnotation(annotation2.id, { content: updatedContent2.trim() });
+        return {
+          content: [{ type: "text", text: `Updated notes for finding id ${findingId}.` }]
+        };
+      }
+      if (!owner || !repo || !sourceLocation || !line) {
+        return {
+          content: [{ type: "text", text: "Either findingId or all of owner+repo+sourceLocation+line must be provided." }]
+        };
+      }
       const entityKey = `${repoKey(owner, repo)}:${sourceLocation}:L${line}`;
       const existing = store.listAnnotations({
         category: AUDIT_CATEGORY,
@@ -193879,14 +193991,32 @@ function registerQueryResultsCacheCompareTool(server) {
         if (!byDatabase.has(key)) byDatabase.set(key, []);
         byDatabase.get(key).push(entry);
       }
-      const comparison = Array.from(byDatabase.entries()).map(([db, dbEntries]) => ({
-        database: db,
-        languages: [...new Set(dbEntries.map((e) => e.language))],
-        formats: [...new Set(dbEntries.map((e) => e.outputFormat))],
-        totalResultCount: dbEntries.reduce((sum, e) => sum + (e.resultCount ?? 0), 0),
-        cachedRuns: dbEntries.length,
-        latestCachedAt: dbEntries[0].createdAt
-      }));
+      const comparison = Array.from(byDatabase.entries()).map(([db, dbEntries]) => {
+        let totalResultCount = 0;
+        for (const e of dbEntries) {
+          if (e.resultCount != null) {
+            totalResultCount += e.resultCount;
+          } else {
+            try {
+              const content = store.getCacheContent(e.cacheKey);
+              if (content) {
+                const sarif = JSON.parse(content);
+                const count = sarif?.runs?.[0]?.results?.length ?? 0;
+                totalResultCount += count;
+              }
+            } catch {
+            }
+          }
+        }
+        return {
+          database: db,
+          languages: [...new Set(dbEntries.map((e) => e.language))],
+          formats: [...new Set(dbEntries.map((e) => e.outputFormat))],
+          totalResultCount,
+          cachedRuns: dbEntries.length,
+          latestCachedAt: dbEntries[0].createdAt
+        };
+      });
       return {
         content: [{
           type: "text",