Fix bugs 1-5 and implement design improvements 2, 3, 5

Bug 1: Improve bqrs_interpret -t parameter docs, add database param
Bug 2: Fix query_results_cache_compare to derive resultCount from SARIF
Bug 3: Include category field in annotation_search FTS matching
Bug 4: Add findingId as primary lookup for audit_add_notes
Bug 5: Change bqrs_info files param to single file string
Design 2: Serialize concurrent database_analyze on same database
Design 3: Already implemented (cacheKey/queryName in cache_lookup)
Design 5: Auto-enable annotation tools in VSIX with new setting

Agent-Logs-Url: https://github.com/advanced-security/codeql-development-mcp-server/sessions/0d52c27d-21be-49ea-bf90-c8fb6f50a3da

Co-authored-by: data-douser <70299490+data-douser@users.noreply.github.com>

Author: Copilot

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}`;
@@ -187168,6 +187178,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 +187432,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 +187533,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 {
@@ -187654,7 +187695,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 +187720,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 +193709,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 +193940,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",

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

@@ -18,6 +18,32 @@ import { createProjectTempDir } from '../utils/temp-dir';
 
 export type { CLIExecutionResult } from './cli-executor';
 
+/**
+ * Per-database mutex map — serializes concurrent CLI operations that
+ * acquire an exclusive lock on a CodeQL database (e.g. database analyze).
+ */
+const databaseLocks = new Map<string, Promise<unknown>>();
+
+/**
+ * Acquire a serialized slot for the given database path.
+ * Returns a release function that must be called when the operation completes.
+ */
+function acquireDatabaseLock(dbPath: string): { ready: Promise<void>; release: () => void } {
+  const key = dbPath;
+  const previous = databaseLocks.get(key) ?? Promise.resolve();
+
+  let release!: () => void;
+  const gate = new Promise<void>(resolve => {
+    release = resolve;
+  });
+  databaseLocks.set(key, gate);
+
+  return {
+    ready: previous.then(() => {}),
+    release,
+  };
+}
+
 export interface CLIToolDefinition {
   name: string;
   description: string;
@@ -352,6 +378,15 @@ export function registerCLITool(server: McpServer, definition: CLIToolDefinition
             break;
           }
 
+          case 'codeql_bqrs_interpret':
+            // Map 'database' to '--source-archive' for codeql bqrs interpret
+            if (options.database) {
+              const dbPath = resolveDatabasePath(options.database as string);
+              options['source-archive'] = join(dbPath, 'src');
+              delete options.database;
+            }
+            break;
+
           case 'codeql_query_compile':
           case 'codeql_resolve_metadata':
             // Handle query parameter as positional argument for query compilation and metadata tools
@@ -519,7 +554,19 @@ export function registerCLITool(server: McpServer, definition: CLIToolDefinition
             options['keep-databases'] = true;
           }
 
-          result = await executeCodeQLCommand(subcommand, options, [...positionalArgs, ...userAdditionalArgs], cwd);
+          // Serialize concurrent database_analyze calls to the same database
+          // to prevent "cache directory is already locked" errors from the CLI.
+          let dbLock: { ready: Promise<void>; release: () => void } | undefined;
+          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 {

server/src/lib/cli-tool-registry.ts

@@ -232,6 +232,16 @@ export async function processQueryRunResults(
             });
 
             const store = sessionDataManager.getStore();
+
+            // Compute result count from content for SARIF formats
+            let resultCount: number | null = null;
+            if (outputFormat.includes('sarif')) {
+              try {
+                const sarif = JSON.parse(resultContent);
+                resultCount = (sarif?.runs?.[0]?.results as unknown[] | undefined)?.length ?? null;
+              } catch { /* non-SARIF content — leave count null */ }
+            }
+
             store.putCacheEntry({
               bqrsPath,
               cacheKey,
@@ -244,6 +254,7 @@ export async function processQueryRunResults(
               queryName: (queryName as string) || basename(queryPath, '.ql'),
               queryPath,
               resultContent,
+              resultCount,
             });
             enhancedOutput += `\nResults cached with key: ${cacheKey}`;
             logger.info(`Cached query results with key: ${cacheKey}`);