@augmentcode/auggie-sdk 0.1.10 → 0.1.12

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (62) hide show
  1. package/README.md +56 -1
  2. package/dist/auggie/sdk-acp-client.d.ts +7 -6
  3. package/dist/auggie/sdk-acp-client.js +299 -332
  4. package/dist/auggie/sdk-mcp-server.d.ts +5 -3
  5. package/dist/auggie/sdk-mcp-server.js +102 -112
  6. package/dist/context/direct-context.d.ts +82 -22
  7. package/dist/context/direct-context.js +675 -562
  8. package/dist/context/filesystem-context.d.ts +5 -3
  9. package/dist/context/filesystem-context.js +187 -209
  10. package/dist/context/internal/__mocks__/api-client.d.ts +17 -11
  11. package/dist/context/internal/__mocks__/api-client.js +104 -91
  12. package/dist/context/internal/api-client.d.ts +14 -11
  13. package/dist/context/internal/api-client.js +234 -239
  14. package/dist/context/internal/blob-name-calculator.d.ts +5 -4
  15. package/dist/context/internal/blob-name-calculator.js +41 -38
  16. package/dist/context/internal/chat-utils.d.ts +6 -3
  17. package/dist/context/internal/chat-utils.js +5 -18
  18. package/dist/context/internal/credentials.d.ts +5 -4
  19. package/dist/context/internal/credentials.js +24 -38
  20. package/dist/context/internal/retry-utils.d.ts +5 -4
  21. package/dist/context/internal/retry-utils.js +60 -114
  22. package/dist/context/internal/search-utils.d.ts +3 -2
  23. package/dist/context/internal/search-utils.js +8 -9
  24. package/dist/context/internal/session-reader.d.ts +4 -3
  25. package/dist/context/internal/session-reader.js +14 -22
  26. package/dist/context/types.d.ts +132 -13
  27. package/dist/context/types.js +0 -5
  28. package/dist/index.d.ts +8 -7
  29. package/dist/index.js +14 -9
  30. package/dist/version.d.ts +3 -2
  31. package/dist/version.js +24 -38
  32. package/package.json +3 -2
  33. package/dist/auggie/sdk-acp-client.d.ts.map +0 -1
  34. package/dist/auggie/sdk-acp-client.js.map +0 -1
  35. package/dist/auggie/sdk-mcp-server.d.ts.map +0 -1
  36. package/dist/auggie/sdk-mcp-server.js.map +0 -1
  37. package/dist/context/direct-context.d.ts.map +0 -1
  38. package/dist/context/direct-context.js.map +0 -1
  39. package/dist/context/filesystem-context.d.ts.map +0 -1
  40. package/dist/context/filesystem-context.js.map +0 -1
  41. package/dist/context/internal/__mocks__/api-client.d.ts.map +0 -1
  42. package/dist/context/internal/__mocks__/api-client.js.map +0 -1
  43. package/dist/context/internal/api-client.d.ts.map +0 -1
  44. package/dist/context/internal/api-client.js.map +0 -1
  45. package/dist/context/internal/blob-name-calculator.d.ts.map +0 -1
  46. package/dist/context/internal/blob-name-calculator.js.map +0 -1
  47. package/dist/context/internal/chat-utils.d.ts.map +0 -1
  48. package/dist/context/internal/chat-utils.js.map +0 -1
  49. package/dist/context/internal/credentials.d.ts.map +0 -1
  50. package/dist/context/internal/credentials.js.map +0 -1
  51. package/dist/context/internal/retry-utils.d.ts.map +0 -1
  52. package/dist/context/internal/retry-utils.js.map +0 -1
  53. package/dist/context/internal/search-utils.d.ts.map +0 -1
  54. package/dist/context/internal/search-utils.js.map +0 -1
  55. package/dist/context/internal/session-reader.d.ts.map +0 -1
  56. package/dist/context/internal/session-reader.js.map +0 -1
  57. package/dist/context/types.d.ts.map +0 -1
  58. package/dist/context/types.js.map +0 -1
  59. package/dist/index.d.ts.map +0 -1
  60. package/dist/index.js.map +0 -1
  61. package/dist/version.d.ts.map +0 -1
  62. package/dist/version.js.map +0 -1
@@ -1,594 +1,707 @@
1
1
  import { readFile, writeFile } from "node:fs/promises";
2
2
  import { Mutex } from "async-mutex";
3
- import { ContextAPIClient } from "./internal/api-client";
4
- import { BlobNameCalculator } from "./internal/blob-name-calculator";
5
- import { chatWithRetry } from "./internal/chat-utils";
6
- import { resolveCredentials } from "./internal/credentials";
7
- import { formatSearchPrompt } from "./internal/search-utils";
8
- /**
9
- * Direct Context - API-based indexing with import/export state
10
- *
11
- * ⚠️ **EXPERIMENTAL API** - This API is experimental and may change without notice.
12
- *
13
- * @experimental
14
- *
15
- * This class provides explicit file indexing via API calls with the ability
16
- * to export and restore state to avoid re-indexing between sessions.
17
- *
18
- * ## Creating a Context
19
- *
20
- * **Start fresh:**
21
- * ```typescript
22
- * const context = await DirectContext.create();
23
- * ```
24
- *
25
- * **Restore from saved state:**
26
- * ```typescript
27
- * const context = await DirectContext.importFromFile("state.json");
28
- * // or
29
- * const context = await DirectContext.import(stateObject);
30
- * ```
31
- *
32
- * ## Indexing Flow
33
- *
34
- * Files are uploaded to the backend and then indexed asynchronously. You have
35
- * explicit control over when to wait for indexing:
36
- *
37
- * **Option 1: Wait after each upload**
38
- * ```typescript
39
- * await context.addToIndex(files); // waits for indexing by default
40
- * await context.search("query"); // files are already indexed
41
- * ```
42
- *
43
- * **Option 2: Batch uploads then wait**
44
- * ```typescript
45
- * await context.addToIndex(files1, { waitForIndexing: false });
46
- * await context.addToIndex(files2, { waitForIndexing: false });
47
- * await context.waitForIndexing(); // wait for all files to be indexed
48
- * await context.search("query"); // files are now indexed
49
- * ```
50
- *
51
- * Note: `search()` and `searchAndAsk()` do not wait for indexing automatically.
52
- *
53
- * ## State Persistence
54
- *
55
- * **Export state:**
56
- * ```typescript
57
- * const state = context.export();
58
- * // or save to file
59
- * await context.exportToFile("state.json");
60
- * ```
61
- *
62
- * **Import state:**
63
- * ```typescript
64
- * const context = await DirectContext.importFromFile("state.json");
65
- * ```
66
- */
67
- export class DirectContext {
68
- /**
69
- * Create and initialize a new DirectContext instance
70
- *
71
- * Authentication priority:
72
- * 1. options.apiKey / options.apiUrl
73
- * 2. AUGMENT_API_TOKEN / AUGMENT_API_URL environment variables
74
- * 3. ~/.augment/session.json (created by `auggie login`)
75
- *
76
- * @param options Configuration options
77
- * @returns Promise that resolves to a DirectContext instance
78
- */
79
- static async create(options = {}) {
80
- const { apiKey, apiUrl } = await resolveCredentials(options);
81
- return new DirectContext(apiKey, apiUrl, options.debug ?? false);
3
+ import { ContextAPIClient } from "./internal/api-client.js";
4
+ import { BlobNameCalculator } from "./internal/blob-name-calculator.js";
5
+ import { chatWithRetry } from "./internal/chat-utils.js";
6
+ import { resolveCredentials } from "./internal/credentials.js";
7
+ import { formatSearchPrompt } from "./internal/search-utils.js";
8
+ class DirectContext {
9
+ static MAX_FILE_SIZE_BYTES = 1048576;
10
+ // 1MB
11
+ static MAX_BATCH_UPLOAD_SIZE = 1e3;
12
+ static MAX_BATCH_CONTENT_BYTES = 2 * 1024 * 1024;
13
+ // 2MB
14
+ static MAX_FIND_MISSING_SIZE = 1e3;
15
+ static CHECKPOINT_THRESHOLD = 1e3;
16
+ apiClient;
17
+ blobCalculator;
18
+ debug;
19
+ /**
20
+ * State management:
21
+ *
22
+ * blobMap: Tracks all files that have been uploaded (path -> serverBlobId).
23
+ * This includes both indexed and non-indexed blobs.
24
+ * Also, includes both checkpointed and non-checkpointed blobs.
25
+ * The server blob ID is the authoritative identifier.
26
+ *
27
+ * clientBlobMap: Tracks client-computed blob IDs (path -> clientBlobId).
28
+ * Used for detecting unchanged files on the client side.
29
+ * May differ from serverBlobId due to encoding differences.
30
+ *
31
+ * checkpointId: The current checkpoint ID from the backend.
32
+ * Checkpoints are named snapshots of the blob set for performance.
33
+ * Updated after each successful checkpoint operation.
34
+ *
35
+ * pendingAdded: Blob IDs that have been uploaded but not yet checkpointed.
36
+ * Accumulated until CHECKPOINT_THRESHOLD is reached.
37
+ *
38
+ * pendingDeleted: Blob IDs that have been removed but not yet checkpointed.
39
+ * Accumulated until CHECKPOINT_THRESHOLD is reached.
40
+ *
41
+ * Note: Checkpointing and indexing are independent operations.
42
+ * A blob can be checkpointed but not indexed, indexed but not checkpointed,
43
+ * both, or neither.
44
+ */
45
+ blobMap = /* @__PURE__ */ new Map();
46
+ clientBlobMap = /* @__PURE__ */ new Map();
47
+ checkpointId;
48
+ pendingAdded = /* @__PURE__ */ new Set();
49
+ pendingDeleted = /* @__PURE__ */ new Set();
50
+ // Track if this instance was imported from search-only state
51
+ isSearchOnly = false;
52
+ // Mutex for serializing state-modifying operations (upload, checkpoint, etc.)
53
+ mutex = new Mutex();
54
+ // Mutex for serializing polling operations to prevent unbounded concurrent requests
55
+ // to the backend's findMissing endpoint. Without this, concurrent addToIndex() calls
56
+ // could result in many simultaneous polling loops hitting the backend.
57
+ pollingMutex = new Mutex();
58
+ /**
59
+ * Create and initialize a new DirectContext instance
60
+ *
61
+ * Authentication priority:
62
+ * 1. options.apiKey / options.apiUrl
63
+ * 2. AUGMENT_API_TOKEN / AUGMENT_API_URL environment variables
64
+ * 3. ~/.augment/session.json (created by `auggie login`)
65
+ *
66
+ * @param options Configuration options
67
+ * @returns Promise that resolves to a DirectContext instance
68
+ */
69
+ static async create(options = {}) {
70
+ const { apiKey, apiUrl } = await resolveCredentials(options);
71
+ return new DirectContext(apiKey, apiUrl, options.debug ?? false);
72
+ }
73
+ /**
74
+ * Import a DirectContext instance from a saved state object
75
+ *
76
+ * @param state The state object to restore from
77
+ * @param options Configuration options
78
+ * @returns Promise that resolves to a DirectContext instance with restored state
79
+ */
80
+ static async import(state, options = {}) {
81
+ const { apiKey, apiUrl } = await resolveCredentials(options);
82
+ const instance = new DirectContext(apiKey, apiUrl, options.debug ?? false);
83
+ await instance.doImport(state);
84
+ return instance;
85
+ }
86
+ /**
87
+ * Import a DirectContext instance from a saved state file (Node.js only)
88
+ *
89
+ * @param filePath Path to the state file
90
+ * @param options Configuration options
91
+ * @returns Promise that resolves to a DirectContext instance with restored state
92
+ */
93
+ static async importFromFile(filePath, options = {}) {
94
+ const content = await readFile(filePath, "utf-8");
95
+ const state = JSON.parse(content);
96
+ return DirectContext.import(state, options);
97
+ }
98
+ /**
99
+ * Private constructor - use DirectContext.create() instead
100
+ * @param apiKey API key for authentication
101
+ * @param apiUrl API URL for the tenant
102
+ * @param debug Enable debug logging
103
+ */
104
+ constructor(apiKey, apiUrl, debug) {
105
+ this.debug = debug;
106
+ this.apiClient = new ContextAPIClient({
107
+ apiKey,
108
+ apiUrl,
109
+ debug
110
+ });
111
+ this.blobCalculator = new BlobNameCalculator(
112
+ DirectContext.MAX_FILE_SIZE_BYTES
113
+ );
114
+ }
115
+ /**
116
+ * Log a debug message if debug mode is enabled
117
+ */
118
+ log(message) {
119
+ if (this.debug) {
120
+ console.log(`[DirectContext] ${message}`);
82
121
  }
83
- /**
84
- * Import a DirectContext instance from a saved state object
85
- *
86
- * @param state The state object to restore from
87
- * @param options Configuration options
88
- * @returns Promise that resolves to a DirectContext instance with restored state
89
- */
90
- static async import(state, options = {}) {
91
- const { apiKey, apiUrl } = await resolveCredentials(options);
92
- const instance = new DirectContext(apiKey, apiUrl, options.debug ?? false);
93
- await instance.doImport(state);
94
- return instance;
122
+ }
123
+ /**
124
+ * Create a progress context for an operation
125
+ */
126
+ createProgressContext(onProgress, total, startedAt) {
127
+ return { onProgress, total, startedAt, uploaded: 0, indexed: 0 };
128
+ }
129
+ /**
130
+ * Create a checkpoint if threshold is reached
131
+ */
132
+ async maybeCheckpoint() {
133
+ const pendingChanges = this.pendingAdded.size + this.pendingDeleted.size;
134
+ if (pendingChanges < DirectContext.CHECKPOINT_THRESHOLD) {
135
+ this.log(
136
+ `Skipping checkpoint: ${pendingChanges} pending changes (threshold: ${DirectContext.CHECKPOINT_THRESHOLD})`
137
+ );
138
+ return;
95
139
  }
96
- /**
97
- * Import a DirectContext instance from a saved state file (Node.js only)
98
- *
99
- * @param filePath Path to the state file
100
- * @param options Configuration options
101
- * @returns Promise that resolves to a DirectContext instance with restored state
102
- */
103
- static async importFromFile(filePath, options = {}) {
104
- const content = await readFile(filePath, "utf-8");
105
- const state = JSON.parse(content);
106
- return DirectContext.import(state, options);
140
+ const addedBlobs = Array.from(this.pendingAdded);
141
+ const deletedBlobs = Array.from(this.pendingDeleted);
142
+ this.log(
143
+ `Creating checkpoint: ${addedBlobs.length} added, ${deletedBlobs.length} deleted blobs`
144
+ );
145
+ const blobs = {
146
+ checkpointId: this.checkpointId,
147
+ addedBlobs: addedBlobs.sort(),
148
+ deletedBlobs: deletedBlobs.sort()
149
+ };
150
+ const result = await this.apiClient.checkpointBlobs(blobs);
151
+ this.checkpointId = result.newCheckpointId;
152
+ this.log(`Checkpoint created: ${this.checkpointId}`);
153
+ this.pendingAdded.clear();
154
+ this.pendingDeleted.clear();
155
+ }
156
+ /**
157
+ * Get the list of indexed file paths
158
+ *
159
+ * @throws Error if the context was imported from search-only state
160
+ */
161
+ getIndexedPaths() {
162
+ if (this.isSearchOnly) {
163
+ throw new Error(
164
+ "Cannot call getIndexedPaths() on a context imported from search-only state. This operation requires full state with blob information. Import from a full state export instead."
165
+ );
107
166
  }
108
- /**
109
- * Private constructor - use DirectContext.create() instead
110
- * @param apiKey API key for authentication
111
- * @param apiUrl API URL for the tenant
112
- * @param debug Enable debug logging
113
- */
114
- constructor(apiKey, apiUrl, debug) {
115
- /**
116
- * State management:
117
- *
118
- * blobMap: Tracks all files that have been uploaded (path -> blobName).
119
- * This includes both indexed and non-indexed blobs.
120
- * Also, includes both checkpointed and non-checkpointed blobs.
121
- *
122
- * checkpointId: The current checkpoint ID from the backend.
123
- * Checkpoints are named snapshots of the blob set for performance.
124
- * Updated after each successful checkpoint operation.
125
- *
126
- * pendingAdded: Blob names that have been uploaded but not yet checkpointed.
127
- * Accumulated until CHECKPOINT_THRESHOLD is reached.
128
- *
129
- * pendingDeleted: Blob names that have been removed but not yet checkpointed.
130
- * Accumulated until CHECKPOINT_THRESHOLD is reached.
131
- *
132
- * Note: Checkpointing and indexing are independent operations.
133
- * A blob can be checkpointed but not indexed, indexed but not checkpointed,
134
- * both, or neither.
135
- */
136
- this.blobMap = new Map();
137
- this.pendingAdded = new Set();
138
- this.pendingDeleted = new Set();
139
- // Mutex for serializing state-modifying operations (upload, checkpoint, etc.)
140
- this.mutex = new Mutex();
141
- // Mutex for serializing polling operations to prevent unbounded concurrent requests
142
- // to the backend's findMissing endpoint. Without this, concurrent addToIndex() calls
143
- // could result in many simultaneous polling loops hitting the backend.
144
- this.pollingMutex = new Mutex();
145
- this.debug = debug;
146
- // Initialize API client
147
- this.apiClient = new ContextAPIClient({
148
- apiKey,
149
- apiUrl,
150
- debug,
151
- });
152
- // Initialize blob calculator
153
- this.blobCalculator = new BlobNameCalculator(DirectContext.MAX_FILE_SIZE_BYTES);
167
+ return Array.from(this.blobMap.keys());
168
+ }
169
+ /**
170
+ * Add files to the index by uploading them to the backend.
171
+ *
172
+ * By default, this method waits for the uploaded files to be fully indexed
173
+ * on the backend before returning. Set `waitForIndexing: false` to return
174
+ * immediately after upload completes (indexing will continue asynchronously).
175
+ *
176
+ * @param files - Array of files to add to the index
177
+ * @param options - Optional configuration
178
+ * @param options.waitForIndexing - If true (default), waits for the newly added files to be indexed before returning
179
+ * @param options.timeout - Timeout in milliseconds for indexing operation (default: undefined = no timeout)
180
+ * @param options.onProgress - Optional callback to receive progress updates
181
+ * @returns Result indicating which files were newly uploaded vs already uploaded
182
+ * @throws Error if the context was imported from search-only state
183
+ */
184
+ async addToIndex(files, options) {
185
+ if (this.isSearchOnly) {
186
+ throw new Error(
187
+ "Cannot call addToIndex() on a context imported from search-only state. This operation requires full state with blob information for deduplication. Import from a full state export instead, or create a new context with DirectContext.create()."
188
+ );
154
189
  }
155
- /**
156
- * Log a debug message if debug mode is enabled
157
- */
158
- log(message) {
159
- if (this.debug) {
160
- console.log(`[DirectContext] ${message}`);
161
- }
190
+ const waitForIndexing = options?.waitForIndexing ?? true;
191
+ const timeout = options?.timeout;
192
+ const ctx = this.createProgressContext(
193
+ options?.onProgress,
194
+ files.length,
195
+ /* @__PURE__ */ new Date()
196
+ );
197
+ const result = await this.mutex.runExclusive(
198
+ () => this.doAddToIndex(files, ctx)
199
+ );
200
+ if (waitForIndexing && result.newlyUploaded.length > 0) {
201
+ const newlyUploadedBlobNames = result.newlyUploaded.map((path) => this.blobMap.get(path)).filter((blobName) => blobName !== void 0);
202
+ await this.waitForSpecificBlobs(newlyUploadedBlobNames, timeout, ctx);
162
203
  }
163
- /**
164
- * Create a checkpoint if threshold is reached
165
- */
166
- async maybeCheckpoint() {
167
- const pendingChanges = this.pendingAdded.size + this.pendingDeleted.size;
168
- if (pendingChanges < DirectContext.CHECKPOINT_THRESHOLD) {
169
- this.log(`Skipping checkpoint: ${pendingChanges} pending changes (threshold: ${DirectContext.CHECKPOINT_THRESHOLD})`);
170
- return;
204
+ return result;
205
+ }
206
+ /**
207
+ * Find blobs that are missing or not yet indexed on the backend (detailed result)
208
+ * Batches requests in chunks of 1000
209
+ *
210
+ * @param blobNames - Array of blob names to check
211
+ * @returns Object with unknownBlobNames and nonindexedBlobNames arrays
212
+ */
213
+ async findMissingBlobsDetailed(blobNames) {
214
+ const allUnknown = [];
215
+ const allNonIndexed = [];
216
+ for (let i = 0; i < blobNames.length; i += DirectContext.MAX_FIND_MISSING_SIZE) {
217
+ const batch = blobNames.slice(i, i + DirectContext.MAX_FIND_MISSING_SIZE);
218
+ const result = await this.apiClient.findMissing(batch);
219
+ allUnknown.push(...result.unknownBlobNames);
220
+ allNonIndexed.push(...result.nonindexedBlobNames);
221
+ }
222
+ return {
223
+ unknownBlobNames: allUnknown,
224
+ nonindexedBlobNames: allNonIndexed
225
+ };
226
+ }
227
+ /**
228
+ * Upload files in batches respecting backend limits
229
+ * Returns a map of client-computed blob IDs to backend-returned blob IDs
230
+ */
231
+ async batchUploadFiles(filesToUpload, ctx) {
232
+ const batches = [];
233
+ let currentBatch = [];
234
+ let currentBatchSize = 0;
235
+ for (const file of filesToUpload) {
236
+ const fileSize = Buffer.byteLength(file.text, "utf-8");
237
+ const wouldExceedCount = currentBatch.length >= DirectContext.MAX_BATCH_UPLOAD_SIZE;
238
+ const wouldExceedSize = currentBatchSize + fileSize > DirectContext.MAX_BATCH_CONTENT_BYTES;
239
+ if (wouldExceedCount || wouldExceedSize) {
240
+ if (currentBatch.length > 0) {
241
+ batches.push(currentBatch);
171
242
  }
172
- // Get blob names for checkpoint (both sets already contain blob names)
173
- const addedBlobs = Array.from(this.pendingAdded);
174
- const deletedBlobs = Array.from(this.pendingDeleted);
175
- this.log(`Creating checkpoint: ${addedBlobs.length} added, ${deletedBlobs.length} deleted blobs`);
176
- const blobs = {
177
- checkpointId: this.checkpointId,
178
- addedBlobs: addedBlobs.sort(),
179
- deletedBlobs: deletedBlobs.sort(),
180
- };
181
- const result = await this.apiClient.checkpointBlobs(blobs);
182
- this.checkpointId = result.newCheckpointId;
183
- this.log(`Checkpoint created: ${this.checkpointId}`);
184
- // Clear pending changes after successful checkpoint
185
- this.pendingAdded.clear();
186
- this.pendingDeleted.clear();
243
+ currentBatch = [file];
244
+ currentBatchSize = fileSize;
245
+ } else {
246
+ currentBatch.push(file);
247
+ currentBatchSize += fileSize;
248
+ }
187
249
  }
188
- /**
189
- * Get the list of indexed file paths
190
- */
191
- getIndexedPaths() {
192
- return Array.from(this.blobMap.keys());
250
+ if (currentBatch.length > 0) {
251
+ batches.push(currentBatch);
193
252
  }
194
- /**
195
- * Add files to the index by uploading them to the backend.
196
- *
197
- * By default, this method waits for the uploaded files to be fully indexed
198
- * on the backend before returning. Set `waitForIndexing: false` to return
199
- * immediately after upload completes (indexing will continue asynchronously).
200
- *
201
- * @param files - Array of files to add to the index
202
- * @param options - Optional configuration
203
- * @param options.waitForIndexing - If true (default), waits for the newly added files to be indexed before returning
204
- * @returns Result indicating which files were newly uploaded vs already uploaded
205
- */
206
- async addToIndex(files, options) {
207
- const waitForIndexing = options?.waitForIndexing ?? true;
208
- const result = await this.mutex.runExclusive(() => this.doAddToIndex(files));
209
- if (waitForIndexing && result.newlyUploaded.length > 0) {
210
- // Wait for the newly uploaded files to be indexed
211
- // These paths are guaranteed to be in blobMap since they were just added in doAddToIndex
212
- const newlyUploadedBlobNames = result.newlyUploaded
213
- .map((path) => this.blobMap.get(path))
214
- .filter((blobName) => blobName !== undefined);
215
- await this.waitForSpecificBlobs(newlyUploadedBlobNames);
253
+ const blobIdMap = /* @__PURE__ */ new Map();
254
+ let filesProcessed = 0;
255
+ let bytesUploaded = 0;
256
+ for (const batch of batches) {
257
+ const result = await this.apiClient.batchUpload(batch);
258
+ for (const [i, batchItem] of batch.entries()) {
259
+ const clientBlobId = batchItem.blobName;
260
+ const backendBlobId = result.blobNames[i];
261
+ if (backendBlobId) {
262
+ blobIdMap.set(clientBlobId, backendBlobId);
216
263
  }
217
- return result;
264
+ }
265
+ filesProcessed += batch.length;
266
+ ctx.uploaded = filesProcessed;
267
+ bytesUploaded += batch.reduce(
268
+ (sum, file) => sum + Buffer.byteLength(file.text, "utf-8"),
269
+ 0
270
+ );
271
+ if (ctx.onProgress) {
272
+ ctx.onProgress({
273
+ stage: "uploading",
274
+ uploaded: ctx.uploaded,
275
+ indexed: ctx.indexed,
276
+ total: ctx.total,
277
+ currentFile: batch.at(-1)?.pathName,
278
+ bytesUploaded,
279
+ startedAt: ctx.startedAt
280
+ });
281
+ }
218
282
  }
219
- /**
220
- * Check which blobs are missing or not yet indexed on the server
221
- *
222
- * @param blobNames - Array of blob names to check
223
- * @param includeNonIndexed - If true, includes blobs that are uploaded but not yet indexed.
224
- * If false, only returns blobs that need to be uploaded.
225
- * @returns Array of blob names that are either missing or not indexed (depending on flag)
226
- */
227
- async findMissingBlobs(blobNames, includeNonIndexed = false) {
228
- const allMissing = [];
229
- for (let i = 0; i < blobNames.length; i += DirectContext.MAX_FIND_MISSING_SIZE) {
230
- const batch = blobNames.slice(i, i + DirectContext.MAX_FIND_MISSING_SIZE);
231
- const result = await this.apiClient.findMissing(batch);
232
- allMissing.push(...result.unknownBlobNames);
233
- if (includeNonIndexed) {
234
- allMissing.push(...result.nonindexedBlobNames);
235
- }
236
- }
237
- return allMissing;
283
+ return blobIdMap;
284
+ }
285
+ /**
286
+ * Validate that all files are within size limits
287
+ */
288
+ validateFileSizes(files) {
289
+ for (const file of files) {
290
+ const sizeBytes = Buffer.byteLength(file.contents, "utf-8");
291
+ if (sizeBytes > DirectContext.MAX_FILE_SIZE_BYTES) {
292
+ throw new Error(
293
+ `File ${file.path} is too large (${sizeBytes} bytes). Maximum size is ${DirectContext.MAX_FILE_SIZE_BYTES} bytes (1MB).`
294
+ );
295
+ }
238
296
  }
239
- /**
240
- * Upload files in batches respecting backend limits
241
- */
242
- async batchUploadFiles(filesToUpload) {
243
- const batches = [];
244
- let currentBatch = [];
245
- let currentBatchSize = 0;
246
- for (const file of filesToUpload) {
247
- const fileSize = Buffer.byteLength(file.text, "utf-8");
248
- const wouldExceedCount = currentBatch.length >= DirectContext.MAX_BATCH_UPLOAD_SIZE;
249
- const wouldExceedSize = currentBatchSize + fileSize > DirectContext.MAX_BATCH_CONTENT_BYTES;
250
- if (wouldExceedCount || wouldExceedSize) {
251
- if (currentBatch.length > 0) {
252
- batches.push(currentBatch);
253
- }
254
- currentBatch = [file];
255
- currentBatchSize = fileSize;
256
- }
257
- else {
258
- currentBatch.push(file);
259
- currentBatchSize += fileSize;
260
- }
297
+ }
298
+ /**
299
+ * Calculate blob names and prepare upload data for files
300
+ * Returns files that need uploading, new blob entries, and already uploaded files
301
+ */
302
+ prepareBlobsForUpload(files) {
303
+ const filesToUpload = [];
304
+ const newBlobEntries = [];
305
+ const alreadyUploaded = [];
306
+ for (const file of files) {
307
+ const clientBlobId = this.blobCalculator.calculate(
308
+ file.path,
309
+ file.contents
310
+ );
311
+ if (clientBlobId) {
312
+ const existingClientBlobId = this.clientBlobMap.get(file.path);
313
+ if (existingClientBlobId && existingClientBlobId === clientBlobId) {
314
+ alreadyUploaded.push(file.path);
315
+ continue;
261
316
  }
262
- if (currentBatch.length > 0) {
263
- batches.push(currentBatch);
264
- }
265
- for (const batch of batches) {
266
- await this.apiClient.batchUpload(batch);
317
+ const existingServerBlobId = this.blobMap.get(file.path);
318
+ if (existingServerBlobId) {
319
+ this.handleBlobReplacement(existingServerBlobId);
267
320
  }
321
+ newBlobEntries.push([file.path, clientBlobId]);
322
+ filesToUpload.push({
323
+ blobName: clientBlobId,
324
+ // API expects 'blobName' field
325
+ pathName: file.path,
326
+ text: file.contents,
327
+ metadata: []
328
+ });
329
+ }
268
330
  }
269
- /**
270
- * Validate that all files are within size limits
271
- */
272
- validateFileSizes(files) {
273
- for (const file of files) {
274
- const sizeBytes = Buffer.byteLength(file.contents, "utf-8");
275
- if (sizeBytes > DirectContext.MAX_FILE_SIZE_BYTES) {
276
- throw new Error(`File ${file.path} is too large (${sizeBytes} bytes). Maximum size is ${DirectContext.MAX_FILE_SIZE_BYTES} bytes (1MB).`);
277
- }
278
- }
331
+ return { filesToUpload, newBlobEntries, alreadyUploaded };
332
+ }
333
+ /**
334
+ * Handle replacing an existing blob (track deletion of old blob)
335
+ */
336
+ handleBlobReplacement(oldBlobName) {
337
+ if (this.pendingAdded.has(oldBlobName)) {
338
+ this.pendingAdded.delete(oldBlobName);
339
+ } else {
340
+ this.pendingDeleted.add(oldBlobName);
279
341
  }
280
- /**
281
- * Calculate blob names and prepare upload data for files
282
- * Returns files that need uploading, new blob entries, and already uploaded files
283
- */
284
- prepareBlobsForUpload(files) {
285
- const filesToUpload = [];
286
- const newBlobEntries = [];
287
- const alreadyUploaded = [];
288
- for (const file of files) {
289
- const blobName = this.blobCalculator.calculate(file.path, file.contents);
290
- if (blobName) {
291
- const existingBlobName = this.blobMap.get(file.path);
292
- if (existingBlobName && existingBlobName === blobName) {
293
- // File content hasn't changed
294
- alreadyUploaded.push(file.path);
295
- continue;
296
- }
297
- // Handle file replacement: if a file with this path already exists
298
- if (existingBlobName) {
299
- this.handleBlobReplacement(existingBlobName);
300
- }
301
- newBlobEntries.push([file.path, blobName]);
302
- filesToUpload.push({
303
- blobName,
304
- pathName: file.path,
305
- text: file.contents,
306
- metadata: [],
307
- });
308
- }
309
- }
310
- return { filesToUpload, newBlobEntries, alreadyUploaded };
342
+ }
343
+ /**
344
+ * Internal implementation of addToIndex
345
+ */
346
+ async doAddToIndex(files, ctx) {
347
+ this.log(`Adding ${files.length} files to index`);
348
+ this.validateFileSizes(files);
349
+ const { filesToUpload, newBlobEntries, alreadyUploaded } = this.prepareBlobsForUpload(files);
350
+ if (newBlobEntries.length === 0) {
351
+ return { newlyUploaded: [], alreadyUploaded };
311
352
  }
312
- /**
313
- * Handle replacing an existing blob (track deletion of old blob)
314
- */
315
- handleBlobReplacement(oldBlobName) {
316
- // Check if the old blob is in pendingAdded
317
- if (this.pendingAdded.has(oldBlobName)) {
318
- // Old blob was added in this session, just remove it from pendingAdded
319
- this.pendingAdded.delete(oldBlobName);
320
- }
321
- else {
322
- // Old blob was from a previous checkpoint, add to deletedBlobs
323
- this.pendingDeleted.add(oldBlobName);
324
- }
353
+ const blobNamesToCheck = newBlobEntries.map(
354
+ ([_, clientBlobId]) => clientBlobId
355
+ );
356
+ this.log(`Checking ${blobNamesToCheck.length} blobs with server`);
357
+ const result = await this.findMissingBlobsDetailed(blobNamesToCheck);
358
+ const missingBlobSet = /* @__PURE__ */ new Set([
359
+ ...result.unknownBlobNames,
360
+ ...result.nonindexedBlobNames
361
+ ]);
362
+ this.log(`Server is missing ${missingBlobSet.size} blobs`);
363
+ const filesToActuallyUpload = filesToUpload.filter(
364
+ (file) => missingBlobSet.has(file.blobName)
365
+ );
366
+ const blobIdMap = /* @__PURE__ */ new Map();
367
+ if (filesToActuallyUpload.length > 0) {
368
+ this.log(`Uploading ${filesToActuallyUpload.length} files to backend`);
369
+ const uploadedBlobIdMap = await this.batchUploadFiles(
370
+ filesToActuallyUpload,
371
+ ctx
372
+ );
373
+ this.log("Upload complete");
374
+ for (const [clientId, serverId] of uploadedBlobIdMap) {
375
+ blobIdMap.set(clientId, serverId);
376
+ }
325
377
  }
326
- /**
327
- * Categorize files into newly uploaded vs already on server based on server response
328
- */
329
- categorizeFiles(newBlobEntries, // [path, blobName]
330
- missingBlobSet) {
331
- const newlyUploaded = [];
332
- const alreadyOnServer = [];
333
- for (const [path, blobName] of newBlobEntries) {
334
- if (missingBlobSet.has(blobName)) {
335
- newlyUploaded.push(path);
336
- }
337
- else {
338
- alreadyOnServer.push(path);
339
- }
378
+ const newlyUploaded = [];
379
+ const alreadyOnServer = [];
380
+ for (const [path, clientBlobId] of newBlobEntries) {
381
+ const serverBlobId = blobIdMap.get(clientBlobId) ?? clientBlobId;
382
+ this.pendingAdded.add(serverBlobId);
383
+ this.pendingDeleted.delete(serverBlobId);
384
+ this.blobMap.set(path, serverBlobId);
385
+ this.clientBlobMap.set(path, clientBlobId);
386
+ if (missingBlobSet.has(clientBlobId)) {
387
+ newlyUploaded.push(path);
388
+ } else {
389
+ alreadyOnServer.push(path);
390
+ }
391
+ }
392
+ await this.maybeCheckpoint();
393
+ return {
394
+ newlyUploaded,
395
+ alreadyUploaded: [...alreadyUploaded, ...alreadyOnServer]
396
+ };
397
+ }
398
+ /**
399
+ * Remove paths from the index
400
+ *
401
+ * @throws Error if the context was imported from search-only state
402
+ */
403
+ async removeFromIndex(paths) {
404
+ if (this.isSearchOnly) {
405
+ throw new Error(
406
+ "Cannot call removeFromIndex() on a context imported from search-only state. This operation requires full state with blob information. Import from a full state export instead."
407
+ );
408
+ }
409
+ return await this.mutex.runExclusive(() => this.doRemoveFromIndex(paths));
410
+ }
411
+ /**
412
+ * Internal implementation of removeFromIndex
413
+ */
414
+ async doRemoveFromIndex(paths) {
415
+ for (const path of paths) {
416
+ const serverBlobId = this.blobMap.get(path);
417
+ if (serverBlobId) {
418
+ if (this.pendingAdded.has(serverBlobId)) {
419
+ this.pendingAdded.delete(serverBlobId);
420
+ } else {
421
+ this.pendingDeleted.add(serverBlobId);
340
422
  }
341
- return { newlyUploaded, alreadyOnServer };
423
+ this.blobMap.delete(path);
424
+ this.clientBlobMap.delete(path);
425
+ }
342
426
  }
343
- /**
344
- * Internal implementation of addToIndex
345
- */
346
- async doAddToIndex(files) {
347
- this.log(`Adding ${files.length} files to index`);
348
- // Validate file sizes
349
- this.validateFileSizes(files);
350
- // Calculate blob names and prepare uploads
351
- const { filesToUpload, newBlobEntries, alreadyUploaded } = this.prepareBlobsForUpload(files);
352
- if (newBlobEntries.length === 0) {
353
- return { newlyUploaded: [], alreadyUploaded };
427
+ await this.maybeCheckpoint();
428
+ }
429
+ /**
430
+ * Clear the entire index
431
+ */
432
+ async clearIndex() {
433
+ return await this.mutex.runExclusive(() => this.doClearIndex());
434
+ }
435
+ /**
436
+ * Internal implementation of clearIndex
437
+ */
438
+ doClearIndex() {
439
+ this.log(`Clearing index (${this.blobMap.size} files)`);
440
+ this.checkpointId = void 0;
441
+ this.blobMap.clear();
442
+ this.clientBlobMap.clear();
443
+ this.pendingAdded.clear();
444
+ this.pendingDeleted.clear();
445
+ this.log("Index cleared");
446
+ return Promise.resolve();
447
+ }
448
+ /**
449
+ * Wait for specific blobs to be indexed on the backend
450
+ *
451
+ * This method is serialized via pollingMutex to prevent unbounded concurrent
452
+ * polling requests to the backend.
453
+ *
454
+ * @param blobNames - Array of blob names to wait for
455
+ * @param timeoutMs - Timeout in milliseconds (default: undefined = no timeout)
456
+ * @param ctx - Progress context for reporting
457
+ */
458
+ waitForSpecificBlobs(blobNames, timeoutMs, ctx) {
459
+ return this.pollingMutex.runExclusive(async () => {
460
+ if (blobNames.length === 0) {
461
+ this.log("No blobs to wait for");
462
+ return;
463
+ }
464
+ const timeoutMsg = timeoutMs ? `timeout: ${timeoutMs / 1e3}s` : "no timeout";
465
+ this.log(
466
+ `Waiting for ${blobNames.length} blobs to be indexed on backend (${timeoutMsg})`
467
+ );
468
+ const initialPollIntervalMs = 3e3;
469
+ const backoffThresholdMs = 6e4;
470
+ const backoffPollIntervalMs = 6e4;
471
+ const startTime = Date.now();
472
+ while (true) {
473
+ const result = await this.findMissingBlobsDetailed(blobNames);
474
+ const unknownBlobs = result.unknownBlobNames.filter(
475
+ (id) => this.pendingAdded.has(id)
476
+ );
477
+ if (unknownBlobs.length > 0) {
478
+ this.log(
479
+ `WARNING: Backend doesn't recognize ${unknownBlobs.length} blob IDs. This may indicate a blob ID mismatch.`
480
+ );
481
+ this.log(`Unknown blob IDs: ${unknownBlobs.join(", ")}`);
482
+ this.log(
483
+ "This is unexpected but not necessarily an error. Continuing to poll..."
484
+ );
354
485
  }
355
- // Check which blobs the server already has
356
- const blobNamesToCheck = newBlobEntries.map(([_, blobName]) => blobName);
357
- this.log(`Checking ${blobNamesToCheck.length} blobs with server`);
358
- const missingBlobNames = await this.findMissingBlobs(blobNamesToCheck);
359
- const missingBlobSet = new Set(missingBlobNames);
360
- this.log(`Server is missing ${missingBlobNames.length} blobs`);
361
- // Categorize files
362
- const { newlyUploaded, alreadyOnServer } = this.categorizeFiles(newBlobEntries, missingBlobSet);
363
- // Upload only missing files
364
- const filesToActuallyUpload = filesToUpload.filter((file) => missingBlobSet.has(file.blobName));
365
- if (filesToActuallyUpload.length > 0) {
366
- this.log(`Uploading ${filesToActuallyUpload.length} files to backend`);
367
- await this.batchUploadFiles(filesToActuallyUpload);
368
- this.log("Upload complete");
486
+ const stillPending = [
487
+ ...result.unknownBlobNames,
488
+ ...result.nonindexedBlobNames
489
+ ];
490
+ ctx.indexed = blobNames.length - stillPending.length;
491
+ if (ctx.onProgress) {
492
+ ctx.onProgress({
493
+ stage: "indexing",
494
+ uploaded: ctx.uploaded,
495
+ indexed: ctx.indexed,
496
+ total: ctx.total,
497
+ startedAt: ctx.startedAt
498
+ });
369
499
  }
370
- // Update blob tracking state
371
- for (const [path, blobName] of newBlobEntries) {
372
- this.pendingAdded.add(blobName);
373
- this.pendingDeleted.delete(blobName);
374
- this.blobMap.set(path, blobName);
500
+ if (stillPending.length === 0) {
501
+ this.log("All blobs indexed successfully");
502
+ return;
375
503
  }
376
- await this.maybeCheckpoint();
377
- return {
378
- newlyUploaded,
379
- alreadyUploaded: [...alreadyUploaded, ...alreadyOnServer],
380
- };
381
- }
382
- /**
383
- * Remove paths from the index
384
- */
385
- async removeFromIndex(paths) {
386
- return await this.mutex.runExclusive(() => this.doRemoveFromIndex(paths));
387
- }
388
- /**
389
- * Internal implementation of removeFromIndex
390
- */
391
- async doRemoveFromIndex(paths) {
392
- for (const path of paths) {
393
- const blobName = this.blobMap.get(path);
394
- if (blobName) {
395
- if (this.pendingAdded.has(blobName)) {
396
- this.pendingAdded.delete(blobName);
397
- }
398
- else {
399
- this.pendingDeleted.add(blobName);
400
- }
401
- this.blobMap.delete(path);
402
- }
504
+ const elapsedMs = Date.now() - startTime;
505
+ this.log(
506
+ `Still waiting for ${stillPending.length} blobs to be indexed (elapsed: ${Math.round(elapsedMs / 1e3)}s)`
507
+ );
508
+ if (timeoutMs !== void 0 && elapsedMs >= timeoutMs) {
509
+ throw new Error(
510
+ `Indexing timeout: Backend did not finish indexing within ${timeoutMs / 1e3} seconds. You can increase the timeout by passing a 'timeout' option to addToIndex() or waitForIndexing().`
511
+ );
403
512
  }
404
- await this.maybeCheckpoint();
405
- }
406
- /**
407
- * Clear the entire index
408
- */
409
- async clearIndex() {
410
- return await this.mutex.runExclusive(() => this.doClearIndex());
411
- }
412
- /**
413
- * Internal implementation of clearIndex
414
- */
415
- doClearIndex() {
416
- this.log(`Clearing index (${this.blobMap.size} files)`);
417
- this.checkpointId = undefined;
418
- this.blobMap.clear();
419
- this.pendingAdded.clear();
420
- this.pendingDeleted.clear();
421
- this.log("Index cleared");
422
- return Promise.resolve();
513
+ const pollIntervalMs = elapsedMs < backoffThresholdMs ? initialPollIntervalMs : backoffPollIntervalMs;
514
+ await new Promise((resolve) => setTimeout(resolve, pollIntervalMs));
515
+ }
516
+ });
517
+ }
518
+ /**
519
+ * Wait for all indexed files to be fully indexed on the backend.
520
+ *
521
+ * This method polls the backend until all files that have been added to the index
522
+ * are confirmed to be indexed and searchable.
523
+ *
524
+ * @param timeout - Timeout in milliseconds (default: undefined = no timeout)
525
+ * @param onProgress - Optional callback to receive progress updates
526
+ * @returns Promise that resolves when all files are indexed
527
+ * @throws Error if indexing times out (only when timeout is specified)
528
+ */
529
+ async waitForIndexing(timeout, onProgress) {
530
+ const blobNames = Array.from(this.blobMap.values());
531
+ const ctx = this.createProgressContext(onProgress, blobNames.length, /* @__PURE__ */ new Date());
532
+ await this.waitForSpecificBlobs(blobNames, timeout, ctx);
533
+ }
534
+ /**
535
+ * Search the codebase using natural language and return formatted results.
536
+ *
537
+ * The results are returned as a formatted string designed for use in LLM prompts.
538
+ * The format includes file paths, line numbers, and code content in a structured,
539
+ * readable format that can be passed directly to LLM APIs like `generate()`.
540
+ *
541
+ * Note: This method does not wait for indexing. Ensure files are indexed before
542
+ * searching by either:
543
+ * - Using `addToIndex()` with `waitForIndexing: true` (default)
544
+ * - Calling `waitForIndexing()` explicitly before searching
545
+ *
546
+ * @param query - The search query describing what code you're looking for
547
+ * @param options - Optional search options
548
+ * @param options.maxOutputLength - Maximum character length of the formatted output (default: 20000, max: 80000)
549
+ * @returns A formatted string containing the search results, ready for LLM consumption
550
+ */
551
+ async search(query, options) {
552
+ this.log(`Searching for: "${query}"`);
553
+ if (!this.isSearchOnly && this.blobMap.size === 0) {
554
+ throw new Error(
555
+ "Index not initialized. Add files to index first using addToIndex()."
556
+ );
423
557
  }
424
- /**
425
- * Wait for specific blobs to be indexed on the backend
426
- *
427
- * This method is serialized via pollingMutex to prevent unbounded concurrent
428
- * polling requests to the backend.
429
- */
430
- waitForSpecificBlobs(blobNames) {
431
- return this.pollingMutex.runExclusive(async () => {
432
- if (blobNames.length === 0) {
433
- this.log("No blobs to wait for");
434
- return;
435
- }
436
- this.log(`Waiting for ${blobNames.length} blobs to be indexed on backend`);
437
- const initialPollIntervalMs = 3000;
438
- const backoffThresholdMs = 60000;
439
- const backoffPollIntervalMs = 60000;
440
- const maxWaitTimeMs = 600000;
441
- const startTime = Date.now();
442
- while (true) {
443
- // Check for blobs that are not yet indexed (either not uploaded or uploaded but not indexed)
444
- const stillPending = await this.findMissingBlobs(blobNames, true);
445
- if (stillPending.length === 0) {
446
- this.log("All blobs indexed successfully");
447
- return;
448
- }
449
- const elapsedMs = Date.now() - startTime;
450
- this.log(`Still waiting for ${stillPending.length} blobs to be indexed (elapsed: ${Math.round(elapsedMs / 1000)}s)`);
451
- if (elapsedMs >= maxWaitTimeMs) {
452
- throw new Error(`Indexing timeout: Backend did not finish indexing within ${maxWaitTimeMs / 1000} seconds`);
453
- }
454
- const pollIntervalMs = elapsedMs < backoffThresholdMs
455
- ? initialPollIntervalMs
456
- : backoffPollIntervalMs;
457
- await new Promise((resolve) => setTimeout(resolve, pollIntervalMs));
458
- }
459
- });
558
+ if (this.isSearchOnly && !this.checkpointId && this.pendingAdded.size === 0) {
559
+ throw new Error(
560
+ "Index not initialized. This search-only context has no checkpoint or pending changes."
561
+ );
460
562
  }
461
- /**
462
- * Wait for all indexed files to be fully indexed on the backend.
463
- *
464
- * This method polls the backend until all files that have been added to the index
465
- * are confirmed to be indexed and searchable.
466
- *
467
- * @returns Promise that resolves when all files are indexed
468
- * @throws Error if indexing times out (default: 10 minutes)
469
- */
470
- async waitForIndexing() {
471
- const blobNames = Array.from(this.blobMap.values());
472
- await this.waitForSpecificBlobs(blobNames);
563
+ const blobs = {
564
+ checkpointId: this.checkpointId,
565
+ addedBlobs: Array.from(this.pendingAdded).sort(),
566
+ deletedBlobs: Array.from(this.pendingDeleted).sort()
567
+ };
568
+ this.log(
569
+ `Executing search with checkpoint ${this.checkpointId || "(none)"}, ${this.blobMap.size} indexed files`
570
+ );
571
+ const result = await this.apiClient.agentCodebaseRetrieval(
572
+ query,
573
+ blobs,
574
+ options?.maxOutputLength
575
+ );
576
+ this.log("Search completed successfully");
577
+ return result.formattedRetrieval;
578
+ }
579
+ /**
580
+ * Search the indexed codebase and ask an LLM a question about the results.
581
+ *
582
+ * This is a convenience method that combines search() with an LLM call to answer
583
+ * questions about your codebase.
584
+ *
585
+ * Note: This method does not wait for indexing. Ensure files are indexed before
586
+ * searching by either:
587
+ * - Using `addToIndex()` with `waitForIndexing: true` (default)
588
+ * - Calling `waitForIndexing()` explicitly before searching
589
+ *
590
+ * @param searchQuery - The semantic search query to find relevant code (also used as the prompt if no separate prompt is provided)
591
+ * @param prompt - Optional prompt to ask the LLM about the search results. If not provided, searchQuery is used as the prompt.
592
+ * @returns The LLM's answer to your question
593
+ *
594
+ * @example
595
+ * ```typescript
596
+ * const answer = await context.searchAndAsk(
597
+ * "How does the authentication flow work?"
598
+ * );
599
+ * console.log(answer);
600
+ * ```
601
+ */
602
+ async searchAndAsk(searchQuery, prompt) {
603
+ const results = await this.search(searchQuery);
604
+ const llmPrompt = formatSearchPrompt(prompt ?? searchQuery, results);
605
+ return await chatWithRetry(this.apiClient, llmPrompt, this.debug);
606
+ }
607
+ /**
608
+ * Export the current state to a JSON object
609
+ *
610
+ * @param options Export options
611
+ * @param options.mode Export mode - 'full' (default) includes all blob information,
612
+ * 'search-only' excludes blobs array for minimal storage
613
+ * @returns The exported state object
614
+ *
615
+ * @example
616
+ * ```typescript
617
+ * // Full export (default) - supports all operations
618
+ * const fullState = context.export();
619
+ * const fullState = context.export({ mode: 'full' });
620
+ *
621
+ * // Search-only export - much smaller, only supports search operations
622
+ * const searchState = context.export({ mode: 'search-only' });
623
+ * ```
624
+ */
625
+ export(options) {
626
+ const mode = options?.mode ?? "full";
627
+ if (mode === "full" && this.isSearchOnly) {
628
+ throw new Error(
629
+ "Cannot export as 'full' from a context imported from search-only state. The blob information required for full state is not available. Use export({ mode: 'search-only' }) instead."
630
+ );
473
631
  }
474
- /**
475
- * Search the codebase using natural language and return formatted results.
476
- *
477
- * The results are returned as a formatted string designed for use in LLM prompts.
478
- * The format includes file paths, line numbers, and code content in a structured,
479
- * readable format that can be passed directly to LLM APIs like `generate()`.
480
- *
481
- * Note: This method does not wait for indexing. Ensure files are indexed before
482
- * searching by either:
483
- * - Using `addToIndex()` with `waitForIndexing: true` (default)
484
- * - Calling `waitForIndexing()` explicitly before searching
485
- *
486
- * @param query - The search query describing what code you're looking for
487
- * @param options - Optional search options
488
- * @param options.maxOutputLength - Maximum character length of the formatted output (default: 20000, max: 80000)
489
- * @returns A formatted string containing the search results, ready for LLM consumption
490
- */
491
- async search(query, options) {
492
- this.log(`Searching for: "${query}"`);
493
- if (this.blobMap.size === 0) {
494
- throw new Error("Index not initialized. Add files to index first using addToIndex().");
495
- }
496
- const blobs = {
497
- checkpointId: this.checkpointId,
498
- addedBlobs: Array.from(this.pendingAdded).sort(),
499
- deletedBlobs: Array.from(this.pendingDeleted).sort(),
500
- };
501
- this.log(`Executing search with checkpoint ${this.checkpointId || "(none)"}, ${this.blobMap.size} indexed files`);
502
- const result = await this.apiClient.agentCodebaseRetrieval(query, blobs, options?.maxOutputLength);
503
- this.log("Search completed successfully");
504
- return result.formattedRetrieval;
632
+ const addedBlobs = Array.from(this.pendingAdded);
633
+ const deletedBlobs = Array.from(this.pendingDeleted);
634
+ if (mode === "search-only") {
635
+ const state2 = {
636
+ mode: "search-only",
637
+ checkpointId: this.checkpointId,
638
+ addedBlobs,
639
+ deletedBlobs
640
+ };
641
+ return state2;
505
642
  }
506
- /**
507
- * Search the indexed codebase and ask an LLM a question about the results.
508
- *
509
- * This is a convenience method that combines search() with an LLM call to answer
510
- * questions about your codebase.
511
- *
512
- * Note: This method does not wait for indexing. Ensure files are indexed before
513
- * searching by either:
514
- * - Using `addToIndex()` with `waitForIndexing: true` (default)
515
- * - Calling `waitForIndexing()` explicitly before searching
516
- *
517
- * @param searchQuery - The semantic search query to find relevant code (also used as the prompt if no separate prompt is provided)
518
- * @param prompt - Optional prompt to ask the LLM about the search results. If not provided, searchQuery is used as the prompt.
519
- * @returns The LLM's answer to your question
520
- *
521
- * @example
522
- * ```typescript
523
- * const answer = await context.searchAndAsk(
524
- * "How does the authentication flow work?"
525
- * );
526
- * console.log(answer);
527
- * ```
528
- */
529
- async searchAndAsk(searchQuery, prompt) {
530
- const results = await this.search(searchQuery);
531
- const llmPrompt = formatSearchPrompt(prompt ?? searchQuery, results);
532
- return await chatWithRetry(this.apiClient, llmPrompt, this.debug);
643
+ const blobs = [];
644
+ for (const [path, serverBlobId] of this.blobMap.entries()) {
645
+ const clientBlobId = this.clientBlobMap.get(path);
646
+ if (clientBlobId && clientBlobId !== serverBlobId) {
647
+ blobs.push([serverBlobId, path, clientBlobId]);
648
+ } else {
649
+ blobs.push([serverBlobId, path]);
650
+ }
533
651
  }
534
- /**
535
- * Export the current state to a JSON object
536
- */
537
- export() {
538
- // Convert blobMap to array of [blobName, path] tuples
539
- const blobs = [];
540
- for (const [path, blobName] of this.blobMap.entries()) {
541
- blobs.push([blobName, path]);
652
+ const state = {
653
+ mode: "full",
654
+ checkpointId: this.checkpointId,
655
+ addedBlobs,
656
+ deletedBlobs,
657
+ blobs
658
+ };
659
+ return state;
660
+ }
661
+ /**
662
+ * Internal method to import state from a JSON object
663
+ */
664
+ async doImport(state) {
665
+ return await this.mutex.runExclusive(() => {
666
+ const mode = "mode" in state ? state.mode : "full";
667
+ this.isSearchOnly = mode === "search-only";
668
+ this.checkpointId = state.checkpointId;
669
+ this.blobMap.clear();
670
+ this.clientBlobMap.clear();
671
+ this.pendingAdded.clear();
672
+ this.pendingDeleted.clear();
673
+ if ("blobs" in state && state.blobs) {
674
+ for (const entry of state.blobs) {
675
+ const [serverBlobId, path, clientBlobId] = entry;
676
+ this.blobMap.set(path, serverBlobId);
677
+ this.clientBlobMap.set(path, clientBlobId ?? serverBlobId);
542
678
  }
543
- // Export pending added blobs (blobs added since last checkpoint)
544
- const addedBlobs = Array.from(this.pendingAdded);
545
- // Export pending deleted blobs (blobs deleted since last checkpoint)
546
- const deletedBlobs = Array.from(this.pendingDeleted);
547
- return {
548
- checkpointId: this.checkpointId,
549
- addedBlobs,
550
- deletedBlobs,
551
- blobs,
552
- };
553
- }
554
- /**
555
- * Internal method to import state from a JSON object
556
- */
557
- async doImport(state) {
558
- return await this.mutex.runExclusive(() => {
559
- this.checkpointId = state.checkpointId;
560
- this.blobMap.clear();
561
- this.pendingAdded.clear();
562
- this.pendingDeleted.clear();
563
- // Restore blobMap from blobs array [blobName, path]
564
- if (state.blobs) {
565
- for (const [blobName, path] of state.blobs) {
566
- this.blobMap.set(path, blobName);
567
- }
568
- }
569
- // Restore pending added blobs - directly copy the array to set
570
- if (state.addedBlobs && state.addedBlobs.length > 0) {
571
- this.pendingAdded = new Set(state.addedBlobs);
572
- }
573
- // Restore pending deleted blobs - directly copy the array to set
574
- if (state.deletedBlobs && state.deletedBlobs.length > 0) {
575
- this.pendingDeleted = new Set(state.deletedBlobs);
576
- }
577
- this.log(`State imported: checkpoint ${this.checkpointId}, ${this.blobMap.size} files, ${this.pendingAdded.size} pending added, ${this.pendingDeleted.size} pending deleted`);
578
- });
579
- }
580
- /**
581
- * Export state to a file (Node.js only)
582
- */
583
- async exportToFile(filePath) {
584
- const state = this.export();
585
- await writeFile(filePath, JSON.stringify(state, null, 2), "utf-8");
586
- this.log(`State saved to ${filePath}`);
587
- }
679
+ }
680
+ if (state.addedBlobs && state.addedBlobs.length > 0) {
681
+ this.pendingAdded = new Set(state.addedBlobs);
682
+ }
683
+ if (state.deletedBlobs && state.deletedBlobs.length > 0) {
684
+ this.pendingDeleted = new Set(state.deletedBlobs);
685
+ }
686
+ this.log(
687
+ `State imported (mode: ${mode}): checkpoint ${this.checkpointId}, ${this.blobMap.size} files, ${this.pendingAdded.size} pending added, ${this.pendingDeleted.size} pending deleted`
688
+ );
689
+ });
690
+ }
691
+ /**
692
+ * Export state to a file (Node.js only)
693
+ *
694
+ * @param filePath Path to save the state file
695
+ * @param options Export options
696
+ * @param options.mode Export mode - 'full' (default) includes all blob information,
697
+ * 'search-only' excludes blobs array for minimal storage
698
+ */
699
+ async exportToFile(filePath, options) {
700
+ const state = this.export(options);
701
+ await writeFile(filePath, JSON.stringify(state, null, 2), "utf-8");
702
+ this.log(`State saved to ${filePath}`);
703
+ }
588
704
  }
589
- DirectContext.MAX_FILE_SIZE_BYTES = 1048576; // 1MB
590
- DirectContext.MAX_BATCH_UPLOAD_SIZE = 1000;
591
- DirectContext.MAX_BATCH_CONTENT_BYTES = 2 * 1024 * 1024; // 2MB
592
- DirectContext.MAX_FIND_MISSING_SIZE = 1000;
593
- DirectContext.CHECKPOINT_THRESHOLD = 1000;
594
- //# sourceMappingURL=direct-context.js.map
705
+ export {
706
+ DirectContext
707
+ };