feat: add User-Agent tracking for analytics

package.json

@@ -68,7 +68,7 @@
     }
   },
   "dependencies": {
-    "@augmentcode/auggie-sdk": "^0.1.14",
+    "@augmentcode/auggie-sdk": "^0.1.15",
     "cheerio": "^1.1.2",
     "commander": "^12.0.0",
     "ignore": "^5.3.0",

src/bin/cmd-agent.ts

@@ -7,6 +7,7 @@ import * as readline from "readline";
 import { CLIAgent, type Provider } from "../clients/cli-agent.js";
 import { MultiIndexRunner } from "../clients/multi-index-runner.js";
 import { CompositeStoreReader, parseIndexSpecs } from "../stores/index.js";
+import { buildClientUserAgent } from "../core/utils.js";
 
 const PROVIDER_DEFAULTS: Record<Provider, string> = {
   openai: "gpt-5-mini",
@@ -50,9 +51,13 @@ export const agentCommand = new Command("agent")
       const store = await CompositeStoreReader.fromSpecs(specs);
 
       // Create multi-index runner
+      // Build User-Agent for analytics tracking
+      const clientUserAgent = buildClientUserAgent("cli");
+
       const runner = await MultiIndexRunner.create({
         store,
         searchOnly: options.searchOnly,
+        clientUserAgent,
       });
 
       console.log("\x1b[1;36mContext Connectors Minimal Agent\x1b[0m");
@@ -73,6 +78,7 @@ export const agentCommand = new Command("agent")
         model,
         maxSteps: options.maxSteps,
         verbose: options.verbose,
+        clientUserAgent,
       });
       await agent.initialize();
 

src/bin/cmd-index.ts

@@ -7,6 +7,7 @@ import { Indexer } from "../core/indexer.js";
 import { Source } from "../sources/types.js";
 import { FilesystemStore } from "../stores/filesystem.js";
 import { getS3Config } from "../stores/s3-config.js";
+import { buildClientUserAgent } from "../core/utils.js";
 
 // Shared store options
 interface StoreOptions {
@@ -49,7 +50,9 @@ async function runIndex(
   sourceType: string
 ) {
   console.log(`Indexing ${sourceType} source...`);
-  const indexer = new Indexer();
+  const indexer = new Indexer({
+    clientUserAgent: buildClientUserAgent("cli"),
+  });
   const result = await indexer.index(source, store, indexKey);
 
   console.log(`\nIndexing complete!`);

src/bin/cmd-search.ts

@@ -9,6 +9,7 @@ import { getSourceIdentifier } from "../core/types.js";
 import { getS3Config } from "../stores/s3-config.js";
 import { parseIndexSpec } from "../stores/index-spec.js";
 import type { IndexStoreReader } from "../stores/types.js";
+import { buildClientUserAgent } from "../core/utils.js";
 
 export const searchCommand = new Command("search")
   .description("Search indexed content and answer questions (use --raw for raw results)")
@@ -83,6 +84,7 @@ export const searchCommand = new Command("search")
       const client = new SearchClient({
         store,
         indexName: indexKey,
+        clientUserAgent: buildClientUserAgent("cli"),
       });
 
       await client.initialize();

src/clients/cli-agent.ts

@@ -79,6 +79,12 @@ export interface CLIAgentSingleConfig {
   stream?: boolean;
   /** Custom system prompt. Uses a sensible default if not provided. */
   systemPrompt?: string;
+  /**
+   * Custom User-Agent string for analytics tracking.
+   * When provided, this is sent to the Augment API for usage analytics.
+   * Only used when provider is 'augment'.
+   */
+  clientUserAgent?: string;
 }
 
 /**
@@ -110,6 +116,12 @@ export interface CLIAgentMultiConfig {
   stream?: boolean;
   /** Custom system prompt. Uses a sensible default if not provided. */
   systemPrompt?: string;
+  /**
+   * Custom User-Agent string for analytics tracking.
+   * When provided, this is sent to the Augment API for usage analytics.
+   * Only used when provider is 'augment'.
+   */
+  clientUserAgent?: string;
 }
 
 /** Configuration for the CLI agent */
@@ -136,7 +148,8 @@ Be concise but thorough. Reference specific files and line numbers when helpful.
  */
 async function loadModel(
   provider: Provider,
-  modelName: string
+  modelName: string,
+  clientUserAgent?: string
 ): Promise<LanguageModel> {
   switch (provider) {
     case "openai": {
@@ -179,6 +192,7 @@ async function loadModel(
       return new AugmentLanguageModel(modelName, {
         apiKey: credentials.apiKey,
         apiUrl: credentials.apiUrl,
+        clientUserAgent,
       }) as unknown as LanguageModel;
     }
     default:
@@ -223,6 +237,7 @@ export class CLIAgent {
   private readonly verbose: boolean;
   private readonly stream: boolean;
   private readonly systemPrompt: string;
+  private readonly clientUserAgent?: string;
   private readonly tools: ToolSet;
   private messages: CoreMessage[] = [];
 
@@ -242,6 +257,7 @@ export class CLIAgent {
     this.verbose = config.verbose ?? false;
     this.stream = config.stream ?? true;
     this.systemPrompt = config.systemPrompt ?? DEFAULT_SYSTEM_PROMPT;
+    this.clientUserAgent = config.clientUserAgent;
     this.tools = this.runner ? this.createMultiIndexTools() : this.createSingleClientTools();
   }
 
@@ -393,7 +409,7 @@ export class CLIAgent {
    * @throws Error if the provider package is not installed
    */
   async initialize(): Promise<void> {
-    this.model = await loadModel(this.provider, this.modelName);
+    this.model = await loadModel(this.provider, this.modelName, this.clientUserAgent);
   }
 
   /**

src/clients/mcp-server.ts

@@ -39,13 +39,13 @@ import {
 } from "@modelcontextprotocol/sdk/types.js";
 import type { IndexStoreReader } from "../stores/types.js";
 import { MultiIndexRunner } from "./multi-index-runner.js";
+import { buildClientUserAgent, type MCPClientInfo } from "../core/utils.js";
 import {
   SEARCH_DESCRIPTION,
   LIST_FILES_DESCRIPTION,
   READ_FILE_DESCRIPTION,
   withIndexList,
 } from "./tool-descriptions.js";
-
 /**
  * Configuration for the MCP server.
  */
@@ -72,7 +72,6 @@ export interface MCPServerConfig {
    */
   version?: string;
 }
-
 /**
  * Create an MCP server instance.
  *
@@ -96,18 +95,19 @@ export async function createMCPServer(
   config: MCPServerConfig
 ): Promise<Server> {
   // Create shared runner for multi-index operations
+  // Build User-Agent for analytics tracking
+  const clientUserAgent = buildClientUserAgent("mcp");
+
   const runner = await MultiIndexRunner.create({
     store: config.store,
     indexNames: config.indexNames,
     searchOnly: config.searchOnly,
+    clientUserAgent,
   });
-
   const { indexNames, indexes } = runner;
   const searchOnly = !runner.hasFileOperations();
-
   // Format index list for tool descriptions
   const indexListStr = runner.getIndexListString();
-
   // Create MCP server
   const server = new Server(
     {
@@ -120,7 +120,19 @@ export async function createMCPServer(
       },
     }
   );
-
+  // Use the SDK's oninitialized callback to capture MCP client info
+  // This preserves the SDK's protocol version negotiation
+  server.oninitialized = () => {
+    const clientInfo = server.getClientVersion();
+    if (clientInfo) {
+      const mcpClientInfo: MCPClientInfo = {
+        name: clientInfo.name,
+        version: clientInfo.version,
+      };
+      const updatedUserAgent = buildClientUserAgent("mcp", mcpClientInfo);
+      runner.updateClientUserAgent(updatedUserAgent);
+    }
+  };
   // Define tool type for type safety
   type Tool = {
     name: string;
@@ -134,12 +146,10 @@ export async function createMCPServer(
       required?: string[];
     };
   };
-
   // Tool descriptions with available indexes (from shared module)
   const searchDescription = withIndexList(SEARCH_DESCRIPTION, indexListStr);
   const listFilesDescription = withIndexList(LIST_FILES_DESCRIPTION, indexListStr);
   const readFileDescription = withIndexList(READ_FILE_DESCRIPTION, indexListStr);
-
   // List available tools
   server.setRequestHandler(ListToolsRequestSchema, async () => {
     const tools: Tool[] = [
@@ -167,7 +177,6 @@ export async function createMCPServer(
         },
       },
     ];
-
     // Only advertise file tools if not in search-only mode
     if (!searchOnly) {
       tools.push(
@@ -247,18 +256,14 @@ export async function createMCPServer(
         }
       );
     }
-
     return { tools };
   });
-
   // Handle tool calls
   server.setRequestHandler(CallToolRequestSchema, async (request) => {
     const { name, arguments: args } = request.params;
-
     try {
       const indexName = args?.index_name as string;
       const client = await runner.getClient(indexName);
-
       switch (name) {
         case "search": {
           const result = await client.search(args?.query as string, {
@@ -270,7 +275,6 @@ export async function createMCPServer(
             ],
           };
         }
-
         case "list_files": {
           if (searchOnly) {
             return {
@@ -291,7 +295,6 @@ export async function createMCPServer(
             content: [{ type: "text", text }],
           };
         }
-
         case "read_file": {
           if (searchOnly) {
             return {
@@ -321,7 +324,6 @@ export async function createMCPServer(
             content: [{ type: "text", text: result.contents ?? "" }],
           };
         }
-
         default:
           return {
             content: [{ type: "text", text: `Unknown tool: ${name}` }],
@@ -335,10 +337,8 @@ export async function createMCPServer(
       };
     }
   });
-
   return server;
 }
-
 /**
  * Run an MCP server with stdio transport.
  *
@@ -369,4 +369,3 @@ export async function runMCPServer(config: MCPServerConfig): Promise<void> {
   const transport = new StdioServerTransport();
   await server.connect(transport);
 }
-

src/clients/multi-index-runner.ts

@@ -36,6 +36,11 @@ export interface MultiIndexRunnerConfig {
    * When true, only search is available.
    */
   searchOnly?: boolean;
+  /**
+   * Custom User-Agent string for analytics tracking.
+   * When provided, this is passed to SearchClient instances for API requests.
+   */
+  clientUserAgent?: string;
 }
 
 /** Create a Source from index state metadata */
@@ -65,6 +70,7 @@ async function createSourceFromState(state: IndexStateSearchOnly): Promise<Sourc
 export class MultiIndexRunner {
   private readonly store: IndexStoreReader;
   private readonly searchOnly: boolean;
+  private clientUserAgent?: string;
   private readonly clientCache = new Map<string, SearchClient>();
 
   /** Available index names */
@@ -77,12 +83,14 @@ export class MultiIndexRunner {
     store: IndexStoreReader,
     indexNames: string[],
     indexes: IndexInfo[],
-    searchOnly: boolean
+    searchOnly: boolean,
+    clientUserAgent?: string
   ) {
     this.store = store;
     this.indexNames = indexNames;
     this.indexes = indexes;
     this.searchOnly = searchOnly;
+    this.clientUserAgent = clientUserAgent;
   }
 
   /**
@@ -132,7 +140,18 @@ export class MultiIndexRunner {
       throw new Error("No valid indexes available (all indexes failed to load)");
     }
 
-    return new MultiIndexRunner(store, validIndexNames, indexes, searchOnly);
+    return new MultiIndexRunner(store, validIndexNames, indexes, searchOnly, config.clientUserAgent);
+  }
+
+
+  /**
+   * Update the User-Agent string.
+   * 
+   * Call this after receiving MCP client info to include the client name/version.
+   * Note: Only affects future client creations, not existing cached clients.
+   */
+  updateClientUserAgent(newUserAgent: string): void {
+    this.clientUserAgent = newUserAgent;
   }
 
   /**
@@ -158,6 +177,7 @@ export class MultiIndexRunner {
         store: this.store,
         source,
         indexName,
+        clientUserAgent: this.clientUserAgent,
       });
       await client.initialize();
       this.clientCache.set(indexName, client);

src/clients/search-client.ts

@@ -66,6 +66,11 @@ export interface SearchClientConfig {
    * @default process.env.AUGMENT_API_URL
    */
   apiUrl?: string;
+  /**
+   * Custom User-Agent string for analytics tracking.
+   * When provided, this is sent to the Augment API for usage analytics.
+   */
+  clientUserAgent?: string;
 }
 
 /**
@@ -106,6 +111,7 @@ export class SearchClient {
   private indexName: string;
   private apiKey: string;
   private apiUrl: string;
+  private clientUserAgent?: string;
 
   private context: DirectContext | null = null;
   private state: IndexStateSearchOnly | null = null;
@@ -123,6 +129,7 @@ export class SearchClient {
     this.indexName = config.indexName;
     this.apiKey = config.apiKey ?? process.env.AUGMENT_API_TOKEN ?? "";
     this.apiUrl = config.apiUrl ?? process.env.AUGMENT_API_URL ?? "";
+    this.clientUserAgent = config.clientUserAgent;
   }
 
   /**
@@ -162,6 +169,7 @@ export class SearchClient {
     this.context = await DirectContext.import(this.state.contextState, {
       apiKey: this.apiKey,
       apiUrl: this.apiUrl,
+      clientUserAgent: this.clientUserAgent,
     });
   }
 

src/core/index.ts

@@ -20,7 +20,8 @@ export {
   shouldFilterFile,
 } from "./file-filter.js";
 
-export { sanitizeKey, isoTimestamp } from "./utils.js";
+export { sanitizeKey, isoTimestamp, buildClientUserAgent } from "./utils.js";
+export type { ClientProduct, MCPClientInfo } from "./utils.js";
 
 export { Indexer } from "./indexer.js";
 export type { IndexerConfig } from "./indexer.js";