@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.
- package/README.md +56 -1
- package/dist/auggie/sdk-acp-client.d.ts +7 -6
- package/dist/auggie/sdk-acp-client.js +299 -332
- package/dist/auggie/sdk-mcp-server.d.ts +5 -3
- package/dist/auggie/sdk-mcp-server.js +102 -112
- package/dist/context/direct-context.d.ts +82 -22
- package/dist/context/direct-context.js +675 -562
- package/dist/context/filesystem-context.d.ts +5 -3
- package/dist/context/filesystem-context.js +187 -209
- package/dist/context/internal/__mocks__/api-client.d.ts +17 -11
- package/dist/context/internal/__mocks__/api-client.js +104 -91
- package/dist/context/internal/api-client.d.ts +14 -11
- package/dist/context/internal/api-client.js +234 -239
- package/dist/context/internal/blob-name-calculator.d.ts +5 -4
- package/dist/context/internal/blob-name-calculator.js +41 -38
- package/dist/context/internal/chat-utils.d.ts +6 -3
- package/dist/context/internal/chat-utils.js +5 -18
- package/dist/context/internal/credentials.d.ts +5 -4
- package/dist/context/internal/credentials.js +24 -38
- package/dist/context/internal/retry-utils.d.ts +5 -4
- package/dist/context/internal/retry-utils.js +60 -114
- package/dist/context/internal/search-utils.d.ts +3 -2
- package/dist/context/internal/search-utils.js +8 -9
- package/dist/context/internal/session-reader.d.ts +4 -3
- package/dist/context/internal/session-reader.js +14 -22
- package/dist/context/types.d.ts +132 -13
- package/dist/context/types.js +0 -5
- package/dist/index.d.ts +8 -7
- package/dist/index.js +14 -9
- package/dist/version.d.ts +3 -2
- package/dist/version.js +24 -38
- package/package.json +3 -2
- package/dist/auggie/sdk-acp-client.d.ts.map +0 -1
- package/dist/auggie/sdk-acp-client.js.map +0 -1
- package/dist/auggie/sdk-mcp-server.d.ts.map +0 -1
- package/dist/auggie/sdk-mcp-server.js.map +0 -1
- package/dist/context/direct-context.d.ts.map +0 -1
- package/dist/context/direct-context.js.map +0 -1
- package/dist/context/filesystem-context.d.ts.map +0 -1
- package/dist/context/filesystem-context.js.map +0 -1
- package/dist/context/internal/__mocks__/api-client.d.ts.map +0 -1
- package/dist/context/internal/__mocks__/api-client.js.map +0 -1
- package/dist/context/internal/api-client.d.ts.map +0 -1
- package/dist/context/internal/api-client.js.map +0 -1
- package/dist/context/internal/blob-name-calculator.d.ts.map +0 -1
- package/dist/context/internal/blob-name-calculator.js.map +0 -1
- package/dist/context/internal/chat-utils.d.ts.map +0 -1
- package/dist/context/internal/chat-utils.js.map +0 -1
- package/dist/context/internal/credentials.d.ts.map +0 -1
- package/dist/context/internal/credentials.js.map +0 -1
- package/dist/context/internal/retry-utils.d.ts.map +0 -1
- package/dist/context/internal/retry-utils.js.map +0 -1
- package/dist/context/internal/search-utils.d.ts.map +0 -1
- package/dist/context/internal/search-utils.js.map +0 -1
- package/dist/context/internal/session-reader.d.ts.map +0 -1
- package/dist/context/internal/session-reader.js.map +0 -1
- package/dist/context/types.d.ts.map +0 -1
- package/dist/context/types.js.map +0 -1
- package/dist/index.d.ts.map +0 -1
- package/dist/index.js.map +0 -1
- package/dist/version.d.ts.map +0 -1
- 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
|
-
|
|
10
|
-
|
|
11
|
-
|
|
12
|
-
*
|
|
13
|
-
|
|
14
|
-
|
|
15
|
-
|
|
16
|
-
|
|
17
|
-
|
|
18
|
-
|
|
19
|
-
|
|
20
|
-
|
|
21
|
-
|
|
22
|
-
|
|
23
|
-
|
|
24
|
-
|
|
25
|
-
|
|
26
|
-
|
|
27
|
-
|
|
28
|
-
|
|
29
|
-
|
|
30
|
-
|
|
31
|
-
|
|
32
|
-
|
|
33
|
-
|
|
34
|
-
|
|
35
|
-
|
|
36
|
-
|
|
37
|
-
|
|
38
|
-
|
|
39
|
-
|
|
40
|
-
|
|
41
|
-
|
|
42
|
-
|
|
43
|
-
|
|
44
|
-
|
|
45
|
-
|
|
46
|
-
|
|
47
|
-
|
|
48
|
-
|
|
49
|
-
|
|
50
|
-
|
|
51
|
-
|
|
52
|
-
|
|
53
|
-
|
|
54
|
-
|
|
55
|
-
|
|
56
|
-
|
|
57
|
-
|
|
58
|
-
|
|
59
|
-
|
|
60
|
-
|
|
61
|
-
|
|
62
|
-
|
|
63
|
-
|
|
64
|
-
|
|
65
|
-
|
|
66
|
-
|
|
67
|
-
|
|
68
|
-
|
|
69
|
-
|
|
70
|
-
|
|
71
|
-
|
|
72
|
-
|
|
73
|
-
|
|
74
|
-
|
|
75
|
-
|
|
76
|
-
|
|
77
|
-
|
|
78
|
-
|
|
79
|
-
|
|
80
|
-
|
|
81
|
-
|
|
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
|
-
|
|
85
|
-
|
|
86
|
-
|
|
87
|
-
|
|
88
|
-
|
|
89
|
-
|
|
90
|
-
|
|
91
|
-
|
|
92
|
-
|
|
93
|
-
|
|
94
|
-
|
|
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
|
-
|
|
98
|
-
|
|
99
|
-
|
|
100
|
-
|
|
101
|
-
|
|
102
|
-
|
|
103
|
-
|
|
104
|
-
|
|
105
|
-
|
|
106
|
-
|
|
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
|
-
|
|
110
|
-
|
|
111
|
-
|
|
112
|
-
|
|
113
|
-
|
|
114
|
-
|
|
115
|
-
|
|
116
|
-
|
|
117
|
-
|
|
118
|
-
|
|
119
|
-
|
|
120
|
-
|
|
121
|
-
|
|
122
|
-
|
|
123
|
-
|
|
124
|
-
|
|
125
|
-
|
|
126
|
-
|
|
127
|
-
|
|
128
|
-
|
|
129
|
-
|
|
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
|
-
|
|
157
|
-
|
|
158
|
-
|
|
159
|
-
|
|
160
|
-
|
|
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
|
-
|
|
165
|
-
|
|
166
|
-
|
|
167
|
-
|
|
168
|
-
|
|
169
|
-
|
|
170
|
-
|
|
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
|
-
|
|
173
|
-
|
|
174
|
-
|
|
175
|
-
|
|
176
|
-
|
|
177
|
-
|
|
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
|
-
|
|
190
|
-
*/
|
|
191
|
-
getIndexedPaths() {
|
|
192
|
-
return Array.from(this.blobMap.keys());
|
|
250
|
+
if (currentBatch.length > 0) {
|
|
251
|
+
batches.push(currentBatch);
|
|
193
252
|
}
|
|
194
|
-
|
|
195
|
-
|
|
196
|
-
|
|
197
|
-
|
|
198
|
-
|
|
199
|
-
|
|
200
|
-
|
|
201
|
-
|
|
202
|
-
|
|
203
|
-
|
|
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
|
-
|
|
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
|
-
|
|
221
|
-
|
|
222
|
-
|
|
223
|
-
|
|
224
|
-
|
|
225
|
-
|
|
226
|
-
|
|
227
|
-
|
|
228
|
-
|
|
229
|
-
|
|
230
|
-
|
|
231
|
-
|
|
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
|
-
|
|
241
|
-
|
|
242
|
-
|
|
243
|
-
|
|
244
|
-
|
|
245
|
-
|
|
246
|
-
|
|
247
|
-
|
|
248
|
-
|
|
249
|
-
|
|
250
|
-
|
|
251
|
-
|
|
252
|
-
|
|
253
|
-
|
|
254
|
-
|
|
255
|
-
|
|
256
|
-
|
|
257
|
-
|
|
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
|
-
|
|
263
|
-
|
|
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
|
-
|
|
271
|
-
|
|
272
|
-
|
|
273
|
-
|
|
274
|
-
|
|
275
|
-
|
|
276
|
-
|
|
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
|
-
|
|
282
|
-
|
|
283
|
-
|
|
284
|
-
|
|
285
|
-
|
|
286
|
-
|
|
287
|
-
|
|
288
|
-
|
|
289
|
-
|
|
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
|
-
|
|
314
|
-
|
|
315
|
-
|
|
316
|
-
|
|
317
|
-
|
|
318
|
-
|
|
319
|
-
|
|
320
|
-
|
|
321
|
-
|
|
322
|
-
|
|
323
|
-
|
|
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
|
-
|
|
328
|
-
|
|
329
|
-
|
|
330
|
-
|
|
331
|
-
|
|
332
|
-
|
|
333
|
-
|
|
334
|
-
|
|
335
|
-
|
|
336
|
-
|
|
337
|
-
|
|
338
|
-
|
|
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
|
-
|
|
423
|
+
this.blobMap.delete(path);
|
|
424
|
+
this.clientBlobMap.delete(path);
|
|
425
|
+
}
|
|
342
426
|
}
|
|
343
|
-
|
|
344
|
-
|
|
345
|
-
|
|
346
|
-
|
|
347
|
-
|
|
348
|
-
|
|
349
|
-
|
|
350
|
-
|
|
351
|
-
|
|
352
|
-
|
|
353
|
-
|
|
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
|
-
|
|
356
|
-
|
|
357
|
-
|
|
358
|
-
|
|
359
|
-
|
|
360
|
-
|
|
361
|
-
|
|
362
|
-
|
|
363
|
-
|
|
364
|
-
|
|
365
|
-
|
|
366
|
-
|
|
367
|
-
|
|
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
|
-
|
|
371
|
-
|
|
372
|
-
|
|
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
|
-
|
|
377
|
-
|
|
378
|
-
|
|
379
|
-
|
|
380
|
-
|
|
381
|
-
|
|
382
|
-
|
|
383
|
-
|
|
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
|
-
|
|
405
|
-
|
|
406
|
-
|
|
407
|
-
|
|
408
|
-
|
|
409
|
-
|
|
410
|
-
|
|
411
|
-
|
|
412
|
-
|
|
413
|
-
|
|
414
|
-
|
|
415
|
-
|
|
416
|
-
|
|
417
|
-
|
|
418
|
-
|
|
419
|
-
|
|
420
|
-
|
|
421
|
-
|
|
422
|
-
|
|
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
|
-
|
|
426
|
-
|
|
427
|
-
|
|
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
|
-
|
|
463
|
-
|
|
464
|
-
|
|
465
|
-
|
|
466
|
-
|
|
467
|
-
|
|
468
|
-
|
|
469
|
-
|
|
470
|
-
|
|
471
|
-
|
|
472
|
-
|
|
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
|
-
|
|
476
|
-
|
|
477
|
-
|
|
478
|
-
|
|
479
|
-
|
|
480
|
-
|
|
481
|
-
|
|
482
|
-
|
|
483
|
-
|
|
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
|
-
|
|
508
|
-
|
|
509
|
-
|
|
510
|
-
|
|
511
|
-
|
|
512
|
-
|
|
513
|
-
|
|
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
|
-
|
|
536
|
-
|
|
537
|
-
|
|
538
|
-
|
|
539
|
-
|
|
540
|
-
|
|
541
|
-
|
|
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
|
-
|
|
544
|
-
|
|
545
|
-
|
|
546
|
-
|
|
547
|
-
|
|
548
|
-
|
|
549
|
-
|
|
550
|
-
|
|
551
|
-
|
|
552
|
-
|
|
553
|
-
}
|
|
554
|
-
|
|
555
|
-
|
|
556
|
-
|
|
557
|
-
|
|
558
|
-
|
|
559
|
-
|
|
560
|
-
|
|
561
|
-
|
|
562
|
-
|
|
563
|
-
|
|
564
|
-
|
|
565
|
-
|
|
566
|
-
|
|
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
|
-
|
|
590
|
-
DirectContext
|
|
591
|
-
|
|
592
|
-
DirectContext.MAX_FIND_MISSING_SIZE = 1000;
|
|
593
|
-
DirectContext.CHECKPOINT_THRESHOLD = 1000;
|
|
594
|
-
//# sourceMappingURL=direct-context.js.map
|
|
705
|
+
export {
|
|
706
|
+
DirectContext
|
|
707
|
+
};
|