diff --git a/extensions/copilot/src/extension/xtab/common/globalBudgetCascade.md b/extensions/copilot/src/extension/xtab/common/globalBudgetCascade.md index 8dbacc21ec3..0394f2e899e 100644 --- a/extensions/copilot/src/extension/xtab/common/globalBudgetCascade.md +++ b/extensions/copilot/src/extension/xtab/common/globalBudgetCascade.md @@ -19,79 +19,117 @@ The cascade is **opt-in**. When `PromptOptions.globalBudget` is `undefined`, ## Scope -### Parts that participate +### Parts rendered by the cascade -`GlobalBudgetPart` includes only: +`GlobalBudgetPart` (the parts listed in `order` and emitted by the cascade loop) +includes only: - `recentlyViewedDocuments` - `languageContext` - `neighborFiles` - `diffHistory` -### Parts intentionally excluded +### Parts that draw a share but are not rendered by the cascade -| Part | Why | -|---|---| -| `currentFile` | Essential context for every prediction; allowing donation to/from it would either bloat the prompt or starve the most important section. Keeps its own cap (`currentFile.maxTokens`) and is clipped independently around the cursor by `createTaggedCurrentFileContentUsingPagedClipping`. | -| `lintOptions` | Optional, formatted separately, and small. Keeps its own per-part shape. | +`currentFile` participates in the **allocation** (`shares`) but not in the +**render order** (`order`). It is clipped **last** — after the cascade has run — +around the cursor / edit window, and sized to its share of the pool **plus** +whatever budget the cascade left unused. Concretely, the cascade runs first +seeded with `0` (so the current file donates nothing), and the current file is +then clipped to `currentFileBudget + cascadeFinalSurplus`. Because budget flows +in a single direction (cascade → current file, never back), the current file +"reuses" the cascade's leftover and trims less. See +[Clip the current file last](#clip-the-current-file-last). + +The set of parts that get a `shares` entry is +`GlobalBudgetSharePart = GlobalBudgetPart | 'currentFile'`. + +| Part | How it relates to the pool | +| --- | --- | +| `currentFile` | Sized from the pool but clipped **outside and after** the cascade: clipped to `floor(totalTokens * shares.currentFile) + cascadeFinalSurplus` around the cursor / edit window by `createTaggedCurrentFileContentUsingPagedClipping` (the clip cap `currentFile.maxTokens` is overridden with that pool budget in `xtabProvider`). It is **not** in `order`, so the cascade loop never renders it, and it absorbs the cascade's leftover rather than donating into it. When `globalBudget` is `undefined` it falls back to its own `currentFile.maxTokens` cap. | +| `lintOptions` | Optional, formatted separately, and small. Excluded entirely — no `shares` entry. Keeps its own per-part shape. | ## Inputs | Input | Description | -|---|---| -| `globalBudget.totalTokens` | The single pool size. Default `6000`. Configurable via experiment `chat.advanced.inlineEdits.xtabProvider.globalBudget.totalTokens`. | -| `globalBudget.order` | Ordered list of parts. Earlier parts get budget first; their surplus flows to later parts. | -| `globalBudget.shares` | `Record`. Each part's fraction of `totalTokens`. Must sum to `1 ± 1e-3` across `order`. | +| --- | --- | +| `globalBudget.totalTokens` | The single pool size. Default `7500`. Set via the `totalTokens` field of the experiment JSON string `chat.advanced.inlineEdits.xtabProvider.globalBudget`. | +| `globalBudget.order` | Ordered list of **rendered** parts. Earlier parts get budget first; their surplus flows to later parts. `currentFile` is not listed here. | +| `globalBudget.shares` | `Record` — one fraction of `totalTokens` per rendered part **and** for `currentFile`. Must sum to `1 ± 1e-3` across `order` plus `currentFile`. | ### Defaults `GlobalBudgetOptions.DEFAULT_ORDER`: -``` +```javascript ['languageContext', 'recentlyViewedDocuments', 'neighborFiles', 'diffHistory'] ``` `GlobalBudgetOptions.DEFAULT_SHARES` (volume-neutral with today's per-part caps): -| Part | Share | -|---|---| -| `recentlyViewedDocuments` | 2/6 | -| `languageContext` | 2/6 | -| `neighborFiles` | 1/6 | -| `diffHistory` | 1/6 | +| Part | Share | Base budget at `totalTokens = 7500` | +| --- | --- | --- | +| `currentFile` | 1500/7500 | 1500 | +| `recentlyViewedDocuments` | 2000/7500 | 2000 | +| `languageContext` | 2000/7500 | 2000 | +| `neighborFiles` | 1000/7500 | 1000 | +| `diffHistory` | 1000/7500 | 1000 | -`GlobalBudgetOptions.DEFAULT_TOTAL_TOKENS` = `6000`. +`GlobalBudgetOptions.DEFAULT_TOTAL_TOKENS` = `7500`. + +These shares reproduce today's per-part caps exactly: `currentFile.maxTokens` 1500, +`recentlyViewedDocuments` 2000, `languageContext` 2000, `neighborFiles` 1000, +`diffHistory` 1000. The pool total (`7500`) is the sum of those caps, so enabling +the default global budget neither grows nor shrinks any rendered part's base +allocation. The default order places `languageContext` first because it is often disabled or empty, donating its share to the always-on `recentlyViewedDocuments` next in -line. +line. Whatever the cascade does not use carries to `finalSurplus` and is handed +to the current file's clip. + +`GlobalBudgetOptions.currentFileBudget(gb)` returns \`floor(totalTokens \* +shares.currentFile)\` — the single source of truth for the current file's **base** +share. The current file's actual clip cap is this base plus the cascade's +`finalSurplus`. ## Algorithm -``` -surplus ← 0 +```javascript +surplus ← 0 // the current file never donates, so the cascade always seeds 0 for part in order: budget ← max(0, floor(surplus + totalTokens * shares[part])) consumed ← runSubBuilder(part, maxTokens: budget) // sub-builder ≤ budget surplus ← max(0, budget - consumed) // → next part only +// end-of-loop surplus is returned as finalSurplus and added to the current file's clip cap ``` Notes: +- The cascade is **always seeded with `0`**. The current file is clipped after the + cascade and only ever *receives* leftover, so it has nothing to donate forward. - The cascade calls the existing sub-builder for each part, overriding only its `maxTokens` (other options are inherited from `opts`): - - `recentlyViewedDocuments` → `buildCodeSnippetsUsingPagedClipping` - - `languageContext` → `appendLanguageContextSnippets` - - `neighborFiles` → `appendNeighborFileSnippets` - - `diffHistory` → `getEditDiffHistory` + - `recentlyViewedDocuments` → `buildCodeSnippetsUsingPagedClipping` + - `languageContext` → `appendLanguageContextSnippets` + - `neighborFiles` → `appendNeighborFileSnippets` + - `diffHistory` → `getEditDiffHistory` - Behavior of each sub-builder is unchanged. - Each sub-builder returns `tokensConsumed` using the **same internal accounting** it uses to make budget decisions (paged-clipping line cost for recently-viewed, raw-snippet cost for appenders, per-entry diff cost for diff history). Using that reported value to compute `surplus` keeps the cascade aligned with how each part actually charges against its budget. -- `surplus` is **forward-only**. Unused tokens at the *last* part are lost — the - cascade does not back-flow to earlier parts. +- `surplus` is **forward-only** between cascade parts. Unused tokens at the *last* + part are **not** lost: they form `finalSurplus`, which the current file's clip + reuses. +- The cascade returns its end-of-loop `surplus` as `CascadeResult.finalSurplus`. + `xtabProvider` grows the current file's clip budget by exactly this amount. See + [Clip the current file last](#clip-the-current-file-last). +- Defensive invariant: each sub-builder must report `0 ≤ tokensConsumed ≤ budget`. + The cascade `softAsserts` this per part (it does **not** silently clamp an + overspend down, which would hide the bug) so the conservation argument the + current-file clip relies on holds. ### Document tracking @@ -99,7 +137,7 @@ The cascade seeds `docsInPrompt` with the active document, and the `recentlyViewedDocuments` step adds the documents it includes. `neighborFiles` reads this set to avoid duplicating files already present, and then `appendNeighborFileSnippets` adds each included neighbor document to the set as well. This dependency is -why `validateGlobalBudget` rejects orderings where `neighborFiles` precedes +why `GlobalBudgetOptions.validate` rejects orderings where `neighborFiles` precedes `recentlyViewedDocuments`. The accumulated `docsInPrompt` is then passed to `getEditDiffHistory`, so any @@ -110,9 +148,10 @@ whose document is in `docsInPrompt`). ### Output The cascade's output mirrors the legacy `getRecentCodeSnippets` shape so the -rest of `getUserPrompt` is identical: +rest of `getUserPrompt` is identical, plus the `finalSurplus` used by the +current-file clip: -``` +```javascript { codeSnippets, // recentlyViewed + langCtx + neighbor joined by "\n\n" documents, // docsInPrompt @@ -120,14 +159,14 @@ rest of `getUserPrompt` is identical: editDiffHistory, nDiffsInPrompt, diffTokensInPrompt, + finalSurplus, // end-of-loop surplus, reused by the current-file clip } ``` ## Guarantees and limits -With `totalTokens = T` and shares `s_i` (assumed non-negative; negative shares -would be clamped to 0 by `max(0, floor(…))` in the budget computation, so the -floor guarantee below would not hold for them): +With `totalTokens = T` and shares `s_i` (validation guarantees they are finite and +non-negative — see [Validation](#validation)): - **Per-part floor**: part at index `i` always receives at least `floor(T * s_i)` tokens, regardless of what earlier parts do. @@ -135,111 +174,163 @@ floor guarantee below would not hold for them): floored sum `floor(…floor(floor(T * s_0) + T * s_1) + … + T * s_i)` — its own share plus everything donated by earlier parts, with `floor` applied at every step. Note this is generally smaller than `floor(T * (s_0 + … + s_i))`. +- **Current-file floor**: the current file always receives at least + `floor(T * shares.currentFile)` (its base share), and `finalSurplus ≥ 0` on top. - **Pool ceiling**: total cascade-managed tokens ≤ the ceiling of the last part - (always ≤ `T`, and typically strictly less due to per-step `floor` rounding). - `currentFile` and lint live outside this budget and add to the final prompt - size. -- **No back-flow**: surplus at the last part is wasted. + (always ≤ `T - floor(T * shares.currentFile)`). The current file then draws its + own base slice plus the cascade's `finalSurplus`, so the budgeted parts together + consume ≤ `T`; lint and scaffolding live outside the budget. +- **No back-flow between cascade parts**: surplus at the last cascade part is not + redistributed to earlier cascade parts — it flows to the current file as + `finalSurplus`. - **No intra-part fairness**: a single large item inside one part can consume that part's entire allocation; the cascade only addresses cross-part donation. ## Validation -`validateGlobalBudget` runs at the start of every cascade invocation and throws -on misconfiguration. The config is runtime-tunable (experiments), so failing -loudly is preferable to silent under/over-allocation. +`GlobalBudgetOptions.validate` runs at the start of every cascade invocation (and +again in `xtabProvider` before the current-file clip) and throws on +misconfiguration. The config is runtime-tunable (experiments), so failing loudly +is preferable to silent under/over-allocation. | Rule | Error | -|---|---| +| --- | --- | +| `totalTokens` is finite and `>= 0` | `globalBudget.totalTokens must be a finite, non-negative number, got X` | | `order` has no duplicate parts | `globalBudget.order contains duplicate part 'X'` | | Every part in `order` has a numeric `shares[part]` | `globalBudget.shares is missing entry for 'X'` | +| `shares.currentFile` is a number | `globalBudget.shares is missing entry for 'currentFile'` | +| Every share (order parts **and** `currentFile`) is finite and `>= 0` | `globalBudget.shares['X'] must be a finite, non-negative number, got Y` | | If both present, `recentlyViewedDocuments` precedes `neighborFiles` | `globalBudget.order must place 'recentlyViewedDocuments' before 'neighborFiles'` | -| Sum of `shares[part]` for parts in `order` ≈ 1 (epsilon `1e-3`) | `globalBudget.shares across order must sum to ~1, got ${sharesSum}` | +| Sum of `shares[part]` for parts in `order` plus `shares.currentFile` ≈ 1 (epsilon `1e-3`) | `globalBudget.shares across order must sum to ~1, got ${sharesSum}` | + +> **Why the non-negativity rule matters:** a negative share can still pass the +> "sum ≈ 1" check (e.g. one part `-0.25`, another `1.0`). At allocation time the +> negative part clamps to a `0` budget, but it still counted toward the sum, so the +> other parts over-allocate past the pool — which would let the current file's +> `finalSurplus` exceed the true leftover. Rejecting negative/non-finite shares up +> front keeps `Σ consumed ≤ totalTokens` provable. ## Wiring -The cascade is enabled via `ConfigKey.TeamInternal.InlineEditsXtabGlobalBudgetEnabled` -in `xtabProvider.ts` (~L1444): +The cascade is configured by a **single experiment-driven JSON string**, +`ConfigKey.TeamInternal.InlineEditsXtabGlobalBudget`, modelled after +`modelConfigurationString`. `xtabProvider.ts` (`getGlobalBudget()`) reads it and +parses it with `GlobalBudgetOptions.fromConfigString`: ```ts -globalBudget: globalBudgetEnabled - ? { - totalTokens: configService.getExperimentBasedConfig(InlineEditsXtabGlobalBudgetTotalTokens, expService), - order: GlobalBudgetOptions.DEFAULT_ORDER, - shares: GlobalBudgetOptions.DEFAULT_SHARES, +private getGlobalBudget(): GlobalBudgetOptions | undefined { + const configString = configService.getExperimentBasedConfig(InlineEditsXtabGlobalBudget, expService); + if (!configString) { + return undefined; // unset/empty → disabled, identical to prod } - : undefined, + const result = GlobalBudgetOptions.fromConfigString(configString); + if (result.isError()) { + telemetryService.sendMSFTTelemetryEvent('incorrectNesGlobalBudgetConfig', { errorMessage: result.err, configValue: configString }); + return undefined; // bad config → disabled, never crashes + } + return result.val; +} ``` +The JSON value defines the budget knobs together — `totalTokens`, `order`, and +`shares` — and every field is optional. Omitted fields fall back to +`DEFAULT_TOTAL_TOKENS` / `DEFAULT_ORDER` / `DEFAULT_SHARES`, so: + +- `undefined` / unset / `""` → global budget **disabled** (prod default, byte-identical legacy path). +- `{}` → **enabled** with the volume-neutral defaults. +- `{"totalTokens":6000}` → enabled, only the pool size overridden. +- `{"totalTokens":12000,"order":[…],"shares":{…}}` → fully custom. + +`fromConfigString` structurally validates the JSON (via `GlobalBudgetOptions.VALIDATOR`), +merges it over the defaults, then runs the semantic `GlobalBudgetOptions.validate` +(see [Validation](#validation)); any parse, structural, or semantic failure +returns a `Result.error` and disables the budget. When `shares` is provided it +must list **every** part (the rendered parts plus `currentFile`) — partial +`shares` objects are rejected so the pool stays fully allocated. + +When the budget is enabled, `xtabProvider` gathers the cascade inputs, runs the +cascade, and then clips the current file with cap +`currentFileBudget(globalBudget) + cascade.finalSurplus` (instead of the +standalone `currentFile.maxTokens`). The already-run cascade is threaded into +`getUserPrompt` as `precomputedCascade` so it renders exactly once. See +[Clip the current file last](#clip-the-current-file-last). + Experiment-controlled settings: | Setting | Default | Purpose | -|---|---|---| -| `chat.advanced.inlineEdits.xtabProvider.globalBudget.enabled` | `false` | Master switch | -| `chat.advanced.inlineEdits.xtabProvider.globalBudget.totalTokens` | `6000` | Pool size | +| --- | --- | --- | +| `chat.advanced.inlineEdits.xtabProvider.globalBudget` | `undefined` | JSON string defining `totalTokens`, `order`, and `shares`; unset/empty disables the budget | + +> **Migration note:** the old `globalBudget.enabled` (boolean) and +> `globalBudget.totalTokens` (number) settings have been **replaced** by this +> single JSON string. Any live experiment treatment that pinned those keys must +> migrate: `enabled:true` + `totalTokens:N` becomes the JSON `{"totalTokens":N}`, +> and a bare `enabled:true` becomes `{}`. Because `totalTokens` also funds +> `currentFile`, treatments that pinned an old total (e.g. `6000` or `8000`) +> should move to `{}` (or `7500`) to stay volume-neutral. ## Worked examples -All examples use `DEFAULT_ORDER`, `DEFAULT_SHARES`, and `totalTokens = 5000`. +All examples use `DEFAULT_ORDER`, `DEFAULT_SHARES`, and `totalTokens = 7500`. -Base allocations: `floor(5000 * share)` per part → +Base allocations: `floor(7500 * share)` per part → | Part | Base allocation | -|---|---| -| `languageContext` | 1666 | -| `recentlyViewedDocuments` | 1666 | -| `neighborFiles` | 833 | -| `diffHistory` | 833 | +| --- | --- | +| `languageContext` | 2000 | +| `recentlyViewedDocuments` | 2000 | +| `neighborFiles` | 1000 | +| `diffHistory` | 1000 | +| `currentFile` (clipped last) | 1500 | -Sum = 4998 (2 tokens lost to per-step `floor`). +Sum = 7500. The cascade iterates only the four rendered parts, seeded with `0`. +Its end-of-loop surplus (`finalSurplus`) is added to the current file's base +allocation, shown as the final row's `budget`. -### Example A — `languageContext` disabled; recently-viewed wants a lot +### Example A — cascade parts modest; current file absorbs the leftover | Part | surplus in | budget | consumed | surplus out | -|---|---|---|---|---| -| `languageContext` | 0 | 1666 | 0 | 1666 | -| `recentlyViewedDocuments` | 1666 | 3332 | 3332 | 0 | -| `neighborFiles` | 0 | 833 | 500 | 333 | -| `diffHistory` | 333 | 1166 | 1166 | 0 | +| --- | --- | --- | --- | --- | +| `languageContext` | 0 | 2000 | 0 | 2000 | +| `recentlyViewedDocuments` | 2000 | 4000 | 1500 | 2500 | +| `neighborFiles` | 2500 | 3500 | 500 | 3000 | +| `diffHistory` | 3000 | 4000 | 500 | 3500 | +| `currentFile` (clipped last) | 3500 (`finalSurplus`) | 1500 + 3500 = 5000 | 5000 | 0 | -Tokens placed: 0 + 3332 + 500 + 1166 = **4998**. The disabled language-context -share doubled what recently-viewed got. +Tokens placed: 0 + 1500 + 500 + 500 + 5000 = **7500**. The cascade consumed only +2500, so its 3500 leftover flowed into the current file, which grew from its 1500 +base to 5000 and trimmed less. -### Example B — modest language context; large single recently-viewed file; no neighbors +### Example B — empty cascade; current file reuses the whole pool | Part | surplus in | budget | consumed | surplus out | -|---|---|---|---|---| -| `languageContext` | 0 | 1666 | 800 | 866 | -| `recentlyViewedDocuments` | 866 | 2532 | 2532 | 0 | -| `neighborFiles` | 0 | 833 | 0 | 833 | -| `diffHistory` | 833 | 1666 | 1500 | 166 (wasted) | +| --- | --- | --- | --- | --- | +| `languageContext` | 0 | 2000 | 0 | 2000 | +| `recentlyViewedDocuments` | 2000 | 4000 | 0 | 4000 | +| `neighborFiles` | 4000 | 5000 | 0 | 5000 | +| `diffHistory` | 5000 | 6000 | 0 | 6000 | +| `currentFile` (clipped last) | 6000 (`finalSurplus`) | 1500 + 6000 = 7500 | ≤ 7500 | — | -Tokens placed: 800 + 2532 + 0 + 1500 = **4832**. The trailing 166 from the last -part is lost (no back-flow). +With no language context, empty history, and neighbors disabled, every cascade +part consumes `0`, so the entire non-currentFile pool (`2000 + 2000 + 1000 + 1000 += 6000`) carries to `finalSurplus`. The current file's clip cap becomes `1500 + +6000 = 7500 = T` — it effectively reuses the whole pool. This is the common case +for a file edited in isolation. -### Example C — everything modest +### Example C — cascade fills its pool; current file gets only its base | Part | surplus in | budget | consumed | surplus out | -|---|---|---|---|---| -| `languageContext` | 0 | 1666 | 400 | 1266 | -| `recentlyViewedDocuments` | 1266 | 2932 | 1500 | 1432 | -| `neighborFiles` | 1432 | 2265 | 900 | 1365 | -| `diffHistory` | 1365 | 2198 | 600 | 1598 (wasted) | +| --- | --- | --- | --- | --- | +| `languageContext` | 0 | 2000 | 2000 | 0 | +| `recentlyViewedDocuments` | 0 | 2000 | 2000 | 0 | +| `neighborFiles` | 0 | 1000 | 1000 | 0 | +| `diffHistory` | 0 | 1000 | 1000 | 0 | +| `currentFile` (clipped last) | 0 (`finalSurplus`) | 1500 + 0 = 1500 | 1500 | 0 | -Tokens placed: 400 + 1500 + 900 + 600 = **3400**. Each later part received a -generous inflated cap but had no material to fill it. - -### Example D — large `languageContext`, no donation - -| Part | surplus in | budget | consumed | surplus out | -|---|---|---|---|---| -| `languageContext` | 0 | 1666 | 1666 | 0 | -| `recentlyViewedDocuments` | 0 | 1666 | 1500 | 166 | -| `neighborFiles` | 166 | 999 | 900 | 99 | -| `diffHistory` | 99 | 932 | 800 | 132 (wasted) | - -Tokens placed: 1666 + 1500 + 900 + 800 = **4866**. Each part stayed at or below -its floor; nothing was starved. +Tokens placed: 2000 + 2000 + 1000 + 1000 + 1500 = **7500**. Every cascade part +filled its own share exactly, so `finalSurplus = 0` and the current file falls +back to its 1500 base — the per-part floor. The current file never shrinks below +this base. ## Disabled parts @@ -248,51 +339,53 @@ Parts can be disabled by configuration (`languageContext.enabled = false`, context response, empty neighbor snippets, no edit history). Those parts remain in `order` — the wired-in code always uses `DEFAULT_ORDER`/`DEFAULT_SHARES` — so their slot still runs, just with `consumed = 0`, and their full share donates -forward. +forward. Whatever reaches the end of the cascade becomes `finalSurplus` and is +reused by the current file's clip rather than wasted. ### Effective caps when both `languageContext` and `neighborFiles` are off With `DEFAULT_ORDER` and pool `T`, applying the per-step floor at every donation -step (let `C_rv = floor(floor(T·2/6) + T·2/6)` be the effective cap on -recently-viewed): +step (the cascade is seeded `0`, so let \`C\_rv = floor(floor(T·langCtxShare) + T·rvShare)\` be +the effective cap on recently-viewed): | Part | Effective cap | -|---|---| +| --- | --- | | `languageContext` | 0 (consumed) | -| `recentlyViewedDocuments` | `C_rv` (own 2/6 + langCtx's 2/6, with per-step `floor`) | +| `recentlyViewedDocuments` | `C_rv` (langCtx's share + own share, with per-step `floor`) | | `neighborFiles` | 0 (consumed) | -| `diffHistory` | `floor(floor((C_rv − consumed_rv) + T·1/6) + T·1/6)` (own 1/6 + neighbors' 1/6 + recently-viewed's leftover, with per-step `floor`) | +| `diffHistory` | `floor(floor((C_rv − consumed_rv) + T·neighborShare) + T·diffShare)` (own share + neighbors' share + recently-viewed's leftover, with per-step `floor`) | -At `T = 5000` (so `C_rv = floor(1666 + 1666.666…) = 3332` and the total -cascade pool ceiling is `floor(floor(3332 + 833.333…) + 833.333…) = 4998`, -not `5000`, because of per-step flooring): +At `T = 7500`, \`C\_rv = floor(2000 + 2000) = 4000\` and the total cascade pool +ceiling is \`floor(floor(4000 + 1000) + 1000) = 6000\`. Whatever `diffHistory` +leaves becomes `finalSurplus` for the current file: -| `recentlyViewedDocuments` consumed | `diffHistory` cap | -|---|---| -| 3332 (fills cap) | 1666 | -| 1500 | 3498 | -| 0 | 4998 (whole cascade pool) | +| `recentlyViewedDocuments` consumed | `diffHistory` cap | `finalSurplus` → current file | +| --- | --- | --- | +| 4000 (fills cap) | 2000 | 0 (current file at base 1500) | +| 1500 | 4500 | up to 4500 | +| 0 | 6000 (whole cascade pool) | up to 6000 (current file up to 7500) | -Worked example with both enabled parts hungry at `T = 5000`: +Worked example with both enabled parts hungry at `T = 7500`: | Part | surplus in | budget | consumed | surplus out | -|---|---|---|---|---| -| `languageContext` | 0 | 1666 | 0 | 1666 | -| `recentlyViewedDocuments` | 1666 | 3332 | 3332 | 0 | -| `neighborFiles` | 0 | 833 | 0 | 833 | -| `diffHistory` | 833 | 1666 | 1666 | 0 | +| --- | --- | --- | --- | --- | +| `languageContext` | 0 | 2000 | 0 | 2000 | +| `recentlyViewedDocuments` | 2000 | 4000 | 4000 | 0 | +| `neighborFiles` | 0 | 1000 | 0 | 1000 | +| `diffHistory` | 1000 | 2000 | 2000 | 0 | +| `currentFile` (clipped last) | 0 (`finalSurplus`) | 1500 | 1500 | 0 | -Total placed: **4998**. Recently-viewed absorbs langCtx's donation; diff history -absorbs neighbors' donation. Nothing is permanently wasted unless the *last* -part (`diffHistory`) also under-fills its budget — diff's surplus has no next -part to flow to. +Cascade placed **6000** (recently-viewed absorbed langCtx's donation; diff history +absorbed neighbors' donation), leaving `finalSurplus = 0`, so the current file +stays at its 1500 base for a total of 7500. ### Ordering caveat -`recentlyViewedDocuments` has **first claim** on the donation pool from -`languageContext` because it sits earlier in `order`. If recently-viewed is -hungry, diff history only sees `neighbors' 1/6 + own 1/6 = 2/6` (1666 at -`T = 5000`) — it cannot reach `languageContext`'s share directly. +`recentlyViewedDocuments` has **first claim** on the cascade donation pool +(`languageContext`'s share) because it sits earlier in `order`. If recently-viewed +is hungry, diff history only sees \`neighbors' + own share = 2000 at T = 7500\` +— it cannot reach `languageContext`'s share directly. Whatever +diff history (the last cascade part) leaves still flows to the current file. To make diff history share the donation, you would need to either reorder so `diffHistory` precedes `recentlyViewedDocuments`, or rebalance `shares`. The @@ -305,14 +398,98 @@ rebalance `shares` to sum to 1 — the donation would then be baked in staticall rather than emerging from the cascade. The wired-in code does not do this; it always passes `DEFAULT_ORDER` and `DEFAULT_SHARES`. -## Composition with `currentFile` +## Clip the current file last -`currentFile` is sized **outside** the cascade by -`createTaggedCurrentFileContentUsingPagedClipping`, which clips around the -cursor / edit window up to `currentFile.maxTokens`. The clipped string is -passed into `getUserPrompt` via `taggedCurrentDocLines` and concatenated -between the cascade-managed `recent_files` and `edit_history` blocks. The -cascade never sees nor influences current-file sizing. +The current file is the natural sink for the cascade's leftover because it is the +one part clipped *around* a point of interest (the cursor), so it can always +absorb more context. Rather than build the prompt, measure what is unused, then +rebuild the current file bigger (a two-pass approach that risks double-counting +the leftover), the implementation simply **clips the current file last**: run the +cascade first, then size the current file with all the budget the cascade did not +use. -Final prompt size ≈ `currentFile.maxTokens` + ≤ `totalTokens` (cascade) + -lint + tags/scaffolding + postscript. +### Mechanism + +In `xtabProvider`, under a global budget: + +1. Gather the cascade inputs (language context, neighbor snippets) — these do + **not** depend on the current-file clip, so they can be produced first. +2. Run `runGlobalBudgetCascade(...)` (seeded with `0`; the current file donates + nothing, so the cascade only ever *gives*). +3. Clip the current file **last** with cap + `currentFileBudget + cascade.finalSurplus`. +4. Assemble the prompt, passing the already-computed cascade through to + `getUserPrompt` as `precomputedCascade`. `getUserPrompt` honors + `precomputedCascade` only when `globalBudget` is set, so the cascade runs + exactly **once** and the rendered snippets match the sizing. + +When `globalBudget` is `undefined` (prod default) none of this applies: the +current file is clipped to its own `currentFile.maxTokens`, the inputs are +gathered after, and the cascade is not run at all — byte-identical to the legacy +path. + +### Conservation (proved) + +Because the current file does not donate, budget flows in a single direction +(cascade → current file), so there is no double-counting. With the cascade seeded +at `0`, the end-of-loop surplus telescopes to + +``` +finalSurplus ≤ Σ(totalTokens · shareᵢ) for i in order + − Σ(consumedᵢ) = (T − T·share_cf) − C_cascade +``` + +so the current file's clip cap is + +``` +cfBudget = floor(T · share_cf) + finalSurplus ≤ T · (Σ all shares) − C_cascade +⇒ C_cf + C_cascade ≤ T · (Σ all shares) (total bounded by the pool) ✅ +``` + +and since `finalSurplus ≥ 0`, `cfBudget ≥ floor(T · share_cf)` — the current file +**never shrinks below** its base share. Non-negative, validated shares (see +[Validation](#validation)) are what make the first inequality hold. + +> **Caveat — share-sum tolerance.** `validate` accepts `|Σ shares − 1| ≤ 1e-3`, so +> the bound above is `T · (Σ shares)`, not exactly `T`. A config whose shares sum +> slightly above 1 can over-allocate by at most `~1e-3 · T` (≈ 7.5 tokens at the +> default `T = 7500`). Shares summing to exactly 1 (the defaults do) give the clean +> `≤ T` bound. Under-allocation (sum < 1) simply wastes a little budget. + +> **Caveat — internal accounting, not the full rendered prompt.** "≤ `T`" is over +> the budgeted parts' *internal* token accounting (paged-clipping line cost, +> raw-snippet cost, diff-entry cost). The fully rendered prompt also carries tag +> wrappers, related-info scaffolding, lint, and the postscript, which live outside +> the pool. So the guarantee is "the budgeted parts together consume ≤ `T`", not +> "the whole prompt is ≤ `T` characters". Tests assert current-file-region +> **growth** and internal accounting, not an absolute full-prompt bound. + +### Worked example + +`T = 6000`, `languageContext` disabled, cascade consumes `rv 1500 + neighbor 500 ++ diff 500 = 2500` (`C_cascade = 2500`), `share_cf = 1500/7500 = 1/5` ⇒ base +`currentFileBudget = floor(6000 · 1/5) = 1200`: + +- The cascade is seeded `0`; its `finalSurplus = (1600 + 1600 + 800 + 800) − 2500 + = 2300`. +- The current file is clipped last to `1200 + 2300 = 3500` (= `6000 − 2500`). + +The current file grows `1200 → 3500`, absorbing exactly the 2300 the cascade left +unused — matching the intuition "budget 6k, first build consumed 4k ⇒ give the +current file the remaining 2k". + +### Trade-offs and caveats + +- **Current file does not donate.** Budget only ever flows cascade → current file. + When the current file is small, its base share is **not** handed to the cascade + parts (they get only their own shares). A deployment optimizing for rich + *neighbor/history* context (rather than current-file context) would need to + rebalance `shares` to give those parts more. +- **Reordered awaits.** Because the cascade runs before the current-file clip, + language-context and neighbor-snippet gathering happen before the current file is + clipped. Any `PromptTooLarge('currentFile')` early-return / cancellation reason + therefore fires *after* those awaits under a global budget. This is an acceptable + consequence of the opt-in feature; the prod path keeps the legacy ordering. +- **Next-cursor predictor unaffected.** `xtabNextCursorPredictor` keeps its own + dedicated current-file cap and never carries a global budget into its prompt, so + it is byte-identical regardless of this feature. diff --git a/extensions/copilot/src/extension/xtab/common/promptCrafting.ts b/extensions/copilot/src/extension/xtab/common/promptCrafting.ts index 869fbe8f3a5..aaa17174a0f 100644 --- a/extensions/copilot/src/extension/xtab/common/promptCrafting.ts +++ b/extensions/copilot/src/extension/xtab/common/promptCrafting.ts @@ -12,7 +12,7 @@ import { IXtabHistoryEntry } from '../../../platform/inlineEdits/common/workspac import { ContextKind, TraitContext } from '../../../platform/languageServer/common/languageContextService'; import { Result } from '../../../util/common/result'; import { range } from '../../../util/vs/base/common/arrays'; -import { assertNever } from '../../../util/vs/base/common/assert'; +import { assertNever, softAssert } from '../../../util/vs/base/common/assert'; import { StringEdit, StringReplacement } from '../../../util/vs/editor/common/core/edits/stringEdit'; import { OffsetRange } from '../../../util/vs/editor/common/core/ranges/offsetRange'; import { getEditDiffHistory } from './diffHistoryForPrompt'; @@ -38,6 +38,15 @@ export class PromptPieces { public readonly computeTokens: (s: string) => number, public readonly opts: PromptOptions, public readonly neighborSnippets?: readonly INeighborFileSnippet[], + /** + * A cascade result computed by the caller (the provider, which runs the + * cascade first so it can clip `currentFile` last to `currentFileBudget + + * finalSurplus`). When provided, {@link getUserPrompt} renders these + * snippets instead of running the cascade itself, guaranteeing the surplus + * used to size the current file matches the snippets that end up in the + * prompt. Only honored when `opts.globalBudget` is set. + */ + public readonly precomputedCascade?: CascadeResult, ) { } } @@ -51,7 +60,7 @@ export interface UserPromptResult { export function getUserPrompt(promptPieces: PromptPieces): UserPromptResult { - const { activeDoc, xtabHistory, taggedCurrentDocLines, areaAroundCodeToEdit, langCtx, aggressivenessLevel, lintErrors, computeTokens, opts, neighborSnippets } = promptPieces; + const { activeDoc, xtabHistory, taggedCurrentDocLines, areaAroundCodeToEdit, langCtx, aggressivenessLevel, lintErrors, computeTokens, opts, neighborSnippets, precomputedCascade } = promptPieces; const currentFileContent = taggedCurrentDocLines.join('\n'); let recentlyViewedCodeSnippets: string; @@ -62,7 +71,10 @@ export function getUserPrompt(promptPieces: PromptPieces): UserPromptResult { let diffTokensInPrompt: number; if (opts.globalBudget !== undefined) { - const cascade = runGlobalBudgetCascade(activeDoc, xtabHistory, langCtx, computeTokens, opts, neighborSnippets, opts.globalBudget); + // Reuse a cascade the caller already ran (the provider runs it first so it can + // clip the current file last from `finalSurplus`), or run it now for callers + // that set a global budget without precomputing (e.g. tests). + const cascade = precomputedCascade ?? runGlobalBudgetCascade(activeDoc, xtabHistory, langCtx, computeTokens, opts, neighborSnippets, opts.globalBudget); recentlyViewedCodeSnippets = cascade.codeSnippets; docsInPrompt = cascade.documents; neighborSnippetsResult = cascade.neighborSnippetsResult; @@ -159,13 +171,19 @@ ${PromptTags.EDIT_HISTORY.end}`; return { prompt: trimmedPrompt, nDiffsInPrompt, diffTokensInPrompt, neighborSnippetsResult }; } -interface CascadeResult { +export interface CascadeResult { readonly codeSnippets: string; readonly documents: Set; readonly neighborSnippetsResult: AppendNeighborFileSnippetsResult | undefined; readonly editDiffHistory: string; readonly nDiffsInPrompt: number; readonly diffTokensInPrompt: number; + /** + * Budget left unused after the last part in `order` ran. The provider adds it + * to the current file's clip budget (`currentFileBudget + finalSurplus`) so the + * current file, which is clipped last, reuses whatever the cascade left unused. + */ + readonly finalSurplus: number; } /** @@ -174,17 +192,21 @@ interface CascadeResult { * `surplus + totalTokens * shares[part]` and run that sub-builder with the override. * Unspent budget cascades to the next part. * - * Sub-builders are invoked using existing helpers so behavior of each individual - * part is unchanged. `currentFile` and `lintOptions` are intentionally excluded - * and continue using their own per-part caps. + * The cascade starts with surplus `0` and renders only the parts in `order`. + * `currentFile` is not rendered here (it is clipped separately by the caller) and + * `lintOptions` is excluded entirely; both keep their own per-part caps for + * clipping. Whatever budget is unused after the last part is returned as + * {@link CascadeResult.finalSurplus}; the provider adds it to the current file's + * clip budget so the current file reuses the leftover (it is clipped last). * - * Each sub-builder reports `tokensConsumed` using the same internal accounting - * it uses to make budget decisions (paged-clipping line cost, raw-snippet cost - * for appenders, diff-entry cost for history). The cascade uses that reported - * value to compute `surplus`, which keeps the cascade aligned with how each - * part actually charges against its budget. + * Sub-builders are invoked using existing helpers so behavior of each individual + * part is unchanged. Each sub-builder reports `tokensConsumed` using the same + * internal accounting it uses to make budget decisions (paged-clipping line cost, + * raw-snippet cost for appenders, diff-entry cost for history). The cascade uses + * that reported value to compute `surplus`, which keeps the cascade aligned with + * how each part actually charges against its budget. */ -function runGlobalBudgetCascade( +export function runGlobalBudgetCascade( activeDoc: StatelessNextEditDocument, xtabHistory: readonly IXtabHistoryEntry[], langCtx: LanguageContextResponse | undefined, @@ -193,7 +215,7 @@ function runGlobalBudgetCascade( neighborSnippets: readonly INeighborFileSnippet[] | undefined, globalBudget: GlobalBudgetOptions, ): CascadeResult { - validateGlobalBudget(globalBudget); + GlobalBudgetOptions.validate(globalBudget); const recentlyViewedSnippets: string[] = []; const langCtxSnippets: string[] = []; @@ -250,6 +272,11 @@ function runGlobalBudgetCascade( default: assertNever(part); } + // The conservation guarantee (current file + cascade <= totalTokens) relies + // on every sub-builder reporting `0 <= tokensConsumed <= budget`. Surface a + // violation (never silently clamp `tokensConsumed` down — that would hide + // real overspend and let the prompt exceed the pool). + softAssert(tokensConsumed >= 0 && tokensConsumed <= budget, `globalBudget part '${part}' reported tokensConsumed=${tokensConsumed} outside [0, ${budget}]`); surplus = Math.max(0, budget - tokensConsumed); } @@ -262,44 +289,10 @@ function runGlobalBudgetCascade( editDiffHistory, nDiffsInPrompt, diffTokensInPrompt, + finalSurplus: surplus, }; } -/** - * Validate {@link GlobalBudgetOptions} since it is runtime-configurable - * (e.g. via experiments). Catches misconfigurations that would otherwise - * cause silent, hard-to-debug behavior: - * - duplicate parts in `order` (would render the same part twice) - * - missing share for any part in `order` - * - shares not summing to ~1 across `order` (would over/under-allocate) - * - `neighborFiles` ordered before `recentlyViewedDocuments` (the former - * consults `docsInPrompt` populated by the latter) - */ -function validateGlobalBudget(globalBudget: GlobalBudgetOptions): void { - const seen = new Set(); - for (const part of globalBudget.order) { - if (seen.has(part)) { - throw new Error(`globalBudget.order contains duplicate part '${part}'`); - } - seen.add(part); - if (typeof globalBudget.shares[part] !== 'number') { - throw new Error(`globalBudget.shares is missing entry for '${part}'`); - } - } - - const recentIdx = globalBudget.order.indexOf('recentlyViewedDocuments'); - const neighborIdx = globalBudget.order.indexOf('neighborFiles'); - if (recentIdx !== -1 && neighborIdx !== -1 && neighborIdx < recentIdx) { - throw new Error(`globalBudget.order must place 'recentlyViewedDocuments' before 'neighborFiles'`); - } - - const sharesSum = globalBudget.order.reduce((sum, part) => sum + globalBudget.shares[part], 0); - const epsilon = 1e-3; - if (Math.abs(sharesSum - 1) > epsilon) { - throw new Error(`globalBudget.shares across order must sum to ~1, got ${sharesSum}`); - } -} - function wrapInBackticks(content: string) { return `\`\`\`\n${content}\n\`\`\``; } diff --git a/extensions/copilot/src/extension/xtab/node/xtabProvider.ts b/extensions/copilot/src/extension/xtab/node/xtabProvider.ts index 959e2a31a93..fa773eb2c09 100644 --- a/extensions/copilot/src/extension/xtab/node/xtabProvider.ts +++ b/extensions/copilot/src/extension/xtab/node/xtabProvider.ts @@ -33,6 +33,7 @@ import { OptionalChatRequestParams, Prediction } from '../../../platform/network import { IChatEndpoint } from '../../../platform/networking/common/networking'; import { ISimulationTestContext } from '../../../platform/simulationTestContext/common/simulationTestContext'; import { IExperimentationService } from '../../../platform/telemetry/common/nullExperimentationService'; +import { ITelemetryService } from '../../../platform/telemetry/common/telemetry'; import { IWorkspaceService } from '../../../platform/workspace/common/workspaceService'; import { raceFilter } from '../../../util/common/async'; import { AsyncIterUtils, AsyncIterUtilsExt } from '../../../util/common/asyncIterableUtils'; @@ -62,9 +63,9 @@ import { IgnoreImportChangesAspect } from '../../inlineEdits/node/importFilterin import { FetchStreamError } from '../common/fetchStreamError'; import { determineIsInlineSuggestionPosition } from '../common/inlineSuggestion'; import { LintErrors } from '../common/lintErrors'; -import { ClippedDocument, constructTaggedFile, getUserPrompt, N_LINES_ABOVE, N_LINES_AS_CONTEXT, N_LINES_BELOW, PromptPieces } from '../common/promptCrafting'; +import { ClippedDocument, constructTaggedFile, getUserPrompt, N_LINES_ABOVE, N_LINES_AS_CONTEXT, N_LINES_BELOW, PromptPieces, runGlobalBudgetCascade, CascadeResult } from '../common/promptCrafting'; import { countTokensForLines, toUniquePath } from '../common/promptCraftingUtils'; -import { ISimilarFilesContextService } from '../common/similarFilesContextService'; +import { INeighborFileSnippet, ISimilarFilesContextService } from '../common/similarFilesContextService'; import { nes41Miniv3SystemPrompt, simplifiedPrompt, systemPromptTemplate, unifiedModelSystemPrompt, xtab275SystemPrompt } from '../common/systemMessages'; import { PromptTags } from '../common/tags'; import { TerminalMonitor } from '../common/terminalOutput'; @@ -177,6 +178,7 @@ export class XtabProvider implements IStatelessNextEditProvider { @ILanguageDiagnosticsService private readonly langDiagService: ILanguageDiagnosticsService, @IIgnoreService private readonly ignoreService: IIgnoreService, @ISimilarFilesContextService private readonly similarFilesContextService: ISimilarFilesContextService, + @ITelemetryService private readonly _telemetryService: ITelemetryService, ) { this.userInteractionMonitor = this.instaService.createInstance(UserInteractionMonitor); this.terminalMonitor = this.instaService.createInstance(TerminalMonitor); @@ -250,6 +252,87 @@ export class XtabProvider implements IStatelessNextEditProvider { ); } + /** + * Gathers language context and neighbor snippets and clips the current file, returning the + * pieces {@link getUserPrompt} needs, or a {@link NoNextEditReason} when the request is + * cancelled mid-gathering or the current file cannot fit its budget. + * + * Under a global budget the current file is clipped LAST: the cascade runs first so the + * current file can reuse whatever budget it leaves unused (via `finalSurplus`), and the + * already-run cascade is returned as `precomputedCascade` so {@link getUserPrompt} renders + * it exactly once. With no global budget (prod default) the current file is clipped to its + * own per-part cap first, byte-identical to the legacy path. + */ + private async gatherContextAndClipCurrentFile( + globalBudget: xtabPromptOptions.GlobalBudgetOptions | undefined, + activeDocument: StatelessNextEditDocument, + request: StatelessNextEditRequest, + promptOptions: xtabPromptOptions.PromptOptions, + telemetry: StatelessNextEditTelemetryBuilder, + cancellationToken: CancellationToken, + gatherLanguageContext: () => Promise, + gatherNeighborSnippets: () => Promise, + clipCurrentFileToBudget: (overriddenMaxTokens: number | undefined) => Result<{ clippedTaggedCurrentDoc: ClippedDocument; areaAroundCodeToEdit: string }, 'outOfBudget'>, + ): Promise> { + if (globalBudget !== undefined) { + // Clip the current file LAST. Gather the cascade inputs (which do not depend + // on the current-file clip) and run the cascade first, then size the current + // file to `currentFileBudget + finalSurplus` so it reuses whatever budget the + // cascade left unused. The already-run cascade is threaded into + // `getUserPrompt` as `precomputedCascade` so it renders exactly once. + const langCtx = await gatherLanguageContext(); + if (cancellationToken.isCancellationRequested) { + return Result.error(new NoNextEditReason.GotCancelled('afterLanguageContextAwait')); + } + + const neighborSnippets = await gatherNeighborSnippets(); + if (cancellationToken.isCancellationRequested) { + return Result.error(new NoNextEditReason.GotCancelled('afterNeighborSnippetsAwait')); + } + + const cascade = runGlobalBudgetCascade(activeDocument, request.xtabEditHistory, langCtx, XtabProvider.computeTokens, promptOptions, neighborSnippets, globalBudget); + const currentFileBudget = xtabPromptOptions.GlobalBudgetOptions.currentFileBudget(globalBudget); + + const taggedCurrentFileContentResult = clipCurrentFileToBudget(currentFileBudget + cascade.finalSurplus); + if (taggedCurrentFileContentResult.isError()) { + return Result.error(new NoNextEditReason.PromptTooLarge('currentFile')); + } + const { clippedTaggedCurrentDoc, areaAroundCodeToEdit } = taggedCurrentFileContentResult.val; + telemetry.setNLinesOfCurrentFileInPrompt(clippedTaggedCurrentDoc.lines.length); + return Result.ok({ clippedTaggedCurrentDoc, areaAroundCodeToEdit, precomputedCascade: cascade, langCtx, neighborSnippets }); + } else { + // No global budget (prod default): clip the current file to its own per-part + // cap, then gather context. Byte-identical to the legacy path. + const taggedCurrentFileContentResult = clipCurrentFileToBudget(undefined); + if (taggedCurrentFileContentResult.isError()) { + return Result.error(new NoNextEditReason.PromptTooLarge('currentFile')); + } + const { clippedTaggedCurrentDoc, areaAroundCodeToEdit } = taggedCurrentFileContentResult.val; + // Record the clipped line count BEFORE the context-gathering awaits so it is + // still emitted when the request is cancelled mid-gathering (matches the + // legacy prod timing — the cancellation path below also reports telemetry). + telemetry.setNLinesOfCurrentFileInPrompt(clippedTaggedCurrentDoc.lines.length); + + const langCtx = await gatherLanguageContext(); + if (cancellationToken.isCancellationRequested) { + return Result.error(new NoNextEditReason.GotCancelled('afterLanguageContextAwait')); + } + + const neighborSnippets = await gatherNeighborSnippets(); + if (cancellationToken.isCancellationRequested) { + return Result.error(new NoNextEditReason.GotCancelled('afterNeighborSnippetsAwait')); + } + + return Result.ok({ clippedTaggedCurrentDoc, areaAroundCodeToEdit, precomputedCascade: undefined, langCtx, neighborSnippets }); + } + } + private async *doGetNextEditWithSelection( request: StatelessNextEditRequest, selection: Range | null, @@ -310,27 +393,37 @@ export class XtabProvider implements IStatelessNextEditProvider { const doesIncludeCursorTag = editWindowLines.some(line => line.includes(PromptTags.CURSOR)); const shouldRemoveCursorTagFromResponse = !doesIncludeCursorTag; // we'd like to remove the tag only if the original edit-window didn't include the tag - const taggedCurrentFileContentResult = constructTaggedFile( - currentDocument, - editWindowLinesRange, - areaAroundEditWindowLinesRange, - promptOptions, - XtabProvider.computeTokens, - { - includeLineNumbers: { - areaAroundCodeToEdit: xtabPromptOptions.IncludeLineNumbersOption.None, - currentFileContent: promptOptions.currentFile.includeLineNumbers, - } - } - ); - - if (taggedCurrentFileContentResult.isError()) { - return new NoNextEditReason.PromptTooLarge('currentFile'); + // Under a global budget the current file is clipped LAST so it reuses whatever + // budget the cascade parts (recently-viewed, language context, neighbors, diff + // history) leave unused: the cascade runs first, then the current file is + // clipped to `currentFileBudget + cascadeFinalSurplus`, so it trims less. With + // no global budget (prod default) the current file keeps its own per-part cap. + const globalBudget = promptOptions.globalBudget; + if (globalBudget !== undefined) { + xtabPromptOptions.GlobalBudgetOptions.validate(globalBudget); } - const { clippedTaggedCurrentDoc, areaAroundCodeToEdit } = taggedCurrentFileContentResult.val; - - telemetry.setNLinesOfCurrentFileInPrompt(clippedTaggedCurrentDoc.lines.length); + // Clips the current file to `overriddenMaxTokens` (or its per-part + // `currentFile.maxTokens` cap when `overriddenMaxTokens` is undefined), + // returning the tagged lines plus the area around the code to edit. + const clipCurrentFileToBudget = (overriddenMaxTokens: number | undefined) => { + const cfPromptOptions = overriddenMaxTokens !== undefined + ? { ...promptOptions, currentFile: { ...promptOptions.currentFile, maxTokens: overriddenMaxTokens } } + : promptOptions; + return constructTaggedFile( + currentDocument, + editWindowLinesRange, + areaAroundEditWindowLinesRange, + cfPromptOptions, + XtabProvider.computeTokens, + { + includeLineNumbers: { + areaAroundCodeToEdit: xtabPromptOptions.IncludeLineNumbersOption.None, + currentFileContent: promptOptions.currentFile.includeLineNumbers, + } + } + ); + }; const { aggressivenessLevel, userHappinessScore } = this.userInteractionMonitor.getAggressivenessLevel(); @@ -344,7 +437,7 @@ export class XtabProvider implements IStatelessNextEditProvider { telemetry.setXtabUserHappinessScore(userHappinessScore); } - const langCtx = await this.getAndProcessLanguageContext( + const gatherLanguageContext = () => this.getAndProcessLanguageContext( request, delaySession, activeDocument, @@ -354,23 +447,31 @@ export class XtabProvider implements IStatelessNextEditProvider { cancellationToken, ); - if (cancellationToken.isCancellationRequested) { - return new NoNextEditReason.GotCancelled('afterLanguageContextAwait'); - } - - const neighborSnippets = promptOptions.neighborFiles.enabled - ? await raceCancellation( + const gatherNeighborSnippets = () => promptOptions.neighborFiles.enabled + ? raceCancellation( raceTimeout( this.similarFilesContextService.getSnippetsForPrompt(activeDocument.id.uri, activeDocument.languageId, activeDocument.documentAfterEdits.value, currentDocument.cursorOffset), delaySession.getDebounceTime() ), cancellationToken, ) - : undefined; + : Promise.resolve(undefined); - if (cancellationToken.isCancellationRequested) { - return new NoNextEditReason.GotCancelled('afterNeighborSnippetsAwait'); + const contextResult = await this.gatherContextAndClipCurrentFile( + globalBudget, + activeDocument, + request, + promptOptions, + telemetry, + cancellationToken, + gatherLanguageContext, + gatherNeighborSnippets, + clipCurrentFileToBudget, + ); + if (contextResult.isError()) { + return contextResult.err; } + const { clippedTaggedCurrentDoc, areaAroundCodeToEdit, precomputedCascade, langCtx, neighborSnippets } = contextResult.val; const lintErrors = new LintErrors(activeDocument.id, currentDocument, this.langDiagService, request.xtabEditHistory); @@ -388,6 +489,7 @@ export class XtabProvider implements IStatelessNextEditProvider { XtabProvider.computeTokens, promptOptions, neighborSnippets, + precomputedCascade, ); const { prompt: userPrompt, nDiffsInPrompt, diffTokensInPrompt, neighborSnippetsResult } = getUserPrompt(promptPieces); @@ -1456,13 +1558,7 @@ export class XtabProvider implements IStatelessNextEditProvider { }, lintOptions: undefined, includePostScript: true, - globalBudget: this.configService.getExperimentBasedConfig(ConfigKey.TeamInternal.InlineEditsXtabGlobalBudgetEnabled, this.expService) - ? { - totalTokens: this.configService.getExperimentBasedConfig(ConfigKey.TeamInternal.InlineEditsXtabGlobalBudgetTotalTokens, this.expService), - order: xtabPromptOptions.GlobalBudgetOptions.DEFAULT_ORDER, - shares: xtabPromptOptions.GlobalBudgetOptions.DEFAULT_SHARES, - } - : undefined, + globalBudget: this.getGlobalBudget(), }; const selectedModelConfig = this.modelService.selectedModelConfiguration(); @@ -1473,6 +1569,36 @@ export class XtabProvider implements IStatelessNextEditProvider { }; } + /** + * Resolve the opt-in global budget from its single experiment-driven JSON + * config string (mirrors `modelConfigurationString`). Returns `undefined` + * — disabling the global budget, identical to prod — when the string is + * unset, empty, or fails to parse/validate. Parse failures are reported via + * telemetry so misconfigured experiments are observable. + */ + private getGlobalBudget(): xtabPromptOptions.GlobalBudgetOptions | undefined { + const configString = this.configService.getExperimentBasedConfig(ConfigKey.TeamInternal.InlineEditsXtabGlobalBudget, this.expService); + if (!configString) { + return undefined; + } + + const result = xtabPromptOptions.GlobalBudgetOptions.fromConfigString(configString); + if (result.isError()) { + /* __GDPR__ + "incorrectNesGlobalBudgetConfig" : { + "owner": "ulugbekna", + "comment": "Capture if the experiment-driven NES global budget config string is invalid or malformed, so the global budget was disabled.", + "errorMessage": { "classification": "SystemMetaData", "purpose": "FeatureInsight", "comment": "Error message from parsing or validation." }, + "configValue": { "classification": "SystemMetaData", "purpose": "FeatureInsight", "comment": "The invalid config string so the bad experiment value can be identified." } + } + */ + this._telemetryService.sendMSFTTelemetryEvent('incorrectNesGlobalBudgetConfig', { errorMessage: result.err, configValue: configString }); + return undefined; + } + + return result.val; + } + private getEndpointWithLogging(configuredModelName: string | undefined, logContext: InlineEditRequestLogContext, telemetry: StatelessNextEditTelemetryBuilder): ChatEndpoint { const endpoint = this.getEndpoint(configuredModelName); logContext.setEndpointInfo(typeof endpoint.urlOrRequestMetadata === 'string' ? endpoint.urlOrRequestMetadata : JSON.stringify(endpoint.urlOrRequestMetadata.type), endpoint.model); diff --git a/extensions/copilot/src/extension/xtab/test/common/promptCrafting.spec.ts b/extensions/copilot/src/extension/xtab/test/common/promptCrafting.spec.ts index ccae96a0a98..9deffc6eaed 100644 --- a/extensions/copilot/src/extension/xtab/test/common/promptCrafting.spec.ts +++ b/extensions/copilot/src/extension/xtab/test/common/promptCrafting.spec.ts @@ -8,6 +8,8 @@ import { DocumentId } from '../../../../platform/inlineEdits/common/dataTypes/do import { Edits } from '../../../../platform/inlineEdits/common/dataTypes/edit'; import { LanguageId } from '../../../../platform/inlineEdits/common/dataTypes/languageId'; import { AggressivenessLevel, CurrentFileOptions, DEFAULT_OPTIONS, GlobalBudgetOptions, IncludeLineNumbersOption, PromptingStrategy, PromptOptions } from '../../../../platform/inlineEdits/common/dataTypes/xtabPromptOptions'; +import { LanguageContextResponse } from '../../../../platform/inlineEdits/common/dataTypes/languageContext'; +import { ContextKind } from '../../../../platform/languageServer/common/languageContextService'; import { StatelessNextEditDocument } from '../../../../platform/inlineEdits/common/statelessNextEditProvider'; import { TestLanguageDiagnosticsService } from '../../../../platform/languages/common/testLanguageDiagnosticsService'; import { Result } from '../../../../util/common/result'; @@ -16,8 +18,9 @@ import { StringEdit } from '../../../../util/vs/editor/common/core/edits/stringE import { Position } from '../../../../util/vs/editor/common/core/position'; import { OffsetRange } from '../../../../util/vs/editor/common/core/ranges/offsetRange'; import { StringText } from '../../../../util/vs/editor/common/core/text/abstractText'; +import { Uri } from '../../../../vscodeTypes'; import { LintErrors } from '../../common/lintErrors'; -import { constructTaggedFile, createTaggedCurrentFileContentUsingPagedClipping, expandRangeToPageRange, getUserPrompt, PromptPieces } from '../../common/promptCrafting'; +import { constructTaggedFile, createTaggedCurrentFileContentUsingPagedClipping, expandRangeToPageRange, getUserPrompt, PromptPieces, runGlobalBudgetCascade } from '../../common/promptCrafting'; import { PromptTags } from '../../common/tags'; import { CurrentDocument } from '../../common/xtabCurrentDocument'; @@ -839,7 +842,7 @@ describe('getUserPrompt — globalBudget cascade', () => { return { activeDoc, currentDocument, currentDocLines }; } - function makePieces(globalBudget: PromptOptions['globalBudget']): PromptPieces { + function makePieces(globalBudget: PromptOptions['globalBudget'], extra?: { langCtx?: LanguageContextResponse; precomputedCascade?: ReturnType }): PromptPieces { const { activeDoc, currentDocument, currentDocLines } = makeActiveDoc(); const promptOptions: PromptOptions = { ...DEFAULT_OPTIONS, @@ -854,14 +857,28 @@ describe('getUserPrompt — globalBudget cascade', () => { [], currentDocLines, 'some code', - undefined, + extra?.langCtx, AggressivenessLevel.Medium, new LintErrors(activeDoc.id, currentDocument, new TestLanguageDiagnosticsService()), s => Math.ceil(s.length / 4), promptOptions, + undefined, + extra?.precomputedCascade, ); } + function makeLangCtxWithSnippet(value: string): LanguageContextResponse { + return { + start: 0, + end: 0, + items: [{ + context: { kind: ContextKind.Snippet, priority: 1, uri: Uri.parse('file:///test/ctx.ts'), value }, + timeStamp: 0, + onTimeout: false, + }], + }; + } + test('produces the same prompt as legacy path when budgets are large', () => { const legacy = getUserPrompt(makePieces(undefined)); const cascaded = getUserPrompt(makePieces({ @@ -898,6 +915,7 @@ describe('getUserPrompt — globalBudget cascade', () => { totalTokens: 1000, order: GlobalBudgetOptions.DEFAULT_ORDER, shares: { + currentFile: 0.5, recentlyViewedDocuments: 0.5, languageContext: 0.5, neighborFiles: 0.5, @@ -913,11 +931,79 @@ describe('getUserPrompt — globalBudget cascade', () => { order: GlobalBudgetOptions.DEFAULT_ORDER, // missing 'diffHistory' shares: { + currentFile: 0.2, languageContext: 0.4, - recentlyViewedDocuments: 0.4, + recentlyViewedDocuments: 0.2, neighborFiles: 0.2, - } as Record<'languageContext' | 'recentlyViewedDocuments' | 'neighborFiles' | 'diffHistory', number>, + } as Record<'currentFile' | 'languageContext' | 'recentlyViewedDocuments' | 'neighborFiles' | 'diffHistory', number>, }); expect(() => getUserPrompt(pieces)).toThrow(/shares is missing entry for 'diffHistory'/); }); + + test('throws when shares is missing an entry for currentFile', () => { + const pieces = makePieces({ + totalTokens: 1000, + order: GlobalBudgetOptions.DEFAULT_ORDER, + // missing 'currentFile' + shares: { + languageContext: 0.25, + recentlyViewedDocuments: 0.25, + neighborFiles: 0.25, + diffHistory: 0.25, + } as Record<'currentFile' | 'languageContext' | 'recentlyViewedDocuments' | 'neighborFiles' | 'diffHistory', number>, + }); + expect(() => getUserPrompt(pieces)).toThrow(/shares is missing entry for 'currentFile'/); + }); + + function runCascade(globalBudget: GlobalBudgetOptions, extra?: { langCtx?: LanguageContextResponse }) { + const { activeDoc } = makeActiveDoc(); + const opts: PromptOptions = { ...DEFAULT_OPTIONS, globalBudget }; + return runGlobalBudgetCascade(activeDoc, [], extra?.langCtx, s => Math.ceil(s.length / 4), opts, undefined, globalBudget); + } + + test('finalSurplus carries the full unused pool when no cascade part consumes budget', () => { + // No langCtx, empty history and neighbors disabled ⇒ every cascade part + // consumes 0, so the entire non-currentFile pool carries to finalSurplus. + // DEFAULT_TOTAL_TOKENS (7500) − currentFileBudget (1500) = 6000. This is + // exactly the budget the provider adds to the current file's clip + // (currentFileBudget 1500 + finalSurplus 6000 = 7500 = T). + const cascade = runCascade({ + totalTokens: GlobalBudgetOptions.DEFAULT_TOTAL_TOKENS, + order: GlobalBudgetOptions.DEFAULT_ORDER, + shares: GlobalBudgetOptions.DEFAULT_SHARES, + }); + expect(cascade.finalSurplus).toBe(6000); + }); + + test('finalSurplus shrinks by what the cascade consumes, so the current file reuses less leftover', () => { + const globalBudget: GlobalBudgetOptions = { + totalTokens: GlobalBudgetOptions.DEFAULT_TOTAL_TOKENS, + order: GlobalBudgetOptions.DEFAULT_ORDER, + shares: GlobalBudgetOptions.DEFAULT_SHARES, + }; + // Empty cascade ⇒ the whole non-currentFile pool (6000) carries to finalSurplus. + const empty = runCascade(globalBudget); + // A rendered language-context snippet consumes budget from the pool, so less + // leftover carries to finalSurplus. The provider clips the current file last to + // currentFileBudget + finalSurplus, so a smaller finalSurplus ⇒ the current file + // reuses less leftover (it is trimmed more) once other parts have content. + const consuming = runCascade(globalBudget, { langCtx: makeLangCtxWithSnippet('const ctxMarker = 1;\n'.repeat(40)) }); + expect(consuming.finalSurplus).toBeLessThan(empty.finalSurplus); + }); + + test('precomputedCascade produces an identical prompt to computing the cascade internally', () => { + const snippet = 'const sharedCtxMarker = 42;'; + const globalBudget: PromptOptions['globalBudget'] = { + totalTokens: 100000, + order: GlobalBudgetOptions.DEFAULT_ORDER, + shares: GlobalBudgetOptions.DEFAULT_SHARES, + }; + const internal = getUserPrompt(makePieces(globalBudget, { langCtx: makeLangCtxWithSnippet(snippet) })); + + const cascade = runCascade(globalBudget, { langCtx: makeLangCtxWithSnippet(snippet) }); + const precomputed = getUserPrompt(makePieces(globalBudget, { langCtx: makeLangCtxWithSnippet(snippet), precomputedCascade: cascade })); + + expect(precomputed.prompt).toBe(internal.prompt); + expect(precomputed.nDiffsInPrompt).toBe(internal.nDiffsInPrompt); + }); }); diff --git a/extensions/copilot/src/extension/xtab/test/node/xtabProvider.spec.ts b/extensions/copilot/src/extension/xtab/test/node/xtabProvider.spec.ts index 47dd5c519fe..82acfbff683 100644 --- a/extensions/copilot/src/extension/xtab/test/node/xtabProvider.spec.ts +++ b/extensions/copilot/src/extension/xtab/test/node/xtabProvider.spec.ts @@ -42,6 +42,7 @@ import { DelaySession } from '../../../inlineEdits/common/delay'; import { createExtensionUnitTestingServices } from '../../../test/node/services'; import { N_LINES_AS_CONTEXT } from '../../common/promptCrafting'; import { nes41Miniv3SystemPrompt, simplifiedPrompt, systemPromptTemplate, unifiedModelSystemPrompt, xtab275SystemPrompt } from '../../common/systemMessages'; +import { PromptTags } from '../../common/tags'; import { CurrentDocument } from '../../common/xtabCurrentDocument'; import { computeAreaAroundEditWindowLinesRange, @@ -1138,6 +1139,69 @@ describe('XtabProvider integration', () => { // Group 4: Filter Pipeline // ======================================================================== + describe('global budget', () => { + + /** Drives the provider once and returns the captured user-message text. */ + async function captureUserPrompt(provider: XtabProvider, request: StatelessNextEditRequest): Promise { + streamingFetcher.setStreamingLines(['x']); + const capturesBefore = streamingFetcher.capturedOptions.length; + const gen = provider.provideNextEdit(request, createMockLogger(), createLogContext(), CancellationToken.None); + await AsyncIterUtils.drainUntilReturn(gen); + // Guard against silently comparing a stale capture: the run must have fetched. + expect(streamingFetcher.capturedOptions.length).toBeGreaterThan(capturesBefore); + const messages = streamingFetcher.capturedOptions.at(-1)?.messages; + const userMessage = messages?.find(m => m.role === Raw.ChatRole.User); + expect(userMessage).toBeDefined(); + return getMessageText(userMessage!); + } + + /** Number of lines in the `<|current_file_content|>` region of the prompt. */ + function currentFileRegionLineCount(prompt: string): number { + const start = prompt.indexOf(PromptTags.CURRENT_FILE.start); + const end = prompt.indexOf(PromptTags.CURRENT_FILE.end); + expect(start).toBeGreaterThanOrEqual(0); + expect(end).toBeGreaterThan(start); + return prompt.slice(start, end).split('\n').length; + } + + const bigFile = Array.from({ length: 400 }, (_, i) => `const value${i} = ${i};`); + + // Under a global budget the current file is clipped LAST, so it absorbs whatever + // budget the cascade parts leave unused. With the (here empty) cascade the + // current file therefore reuses essentially the whole pool and keeps strictly + // MORE of the file than the legacy path, which caps it at its own + // currentFile.maxTokens (1500) and trims the tail. + it('absorbs leftover cascade budget so it keeps more of the current file than the legacy cap', async () => { + const legacy = await captureUserPrompt(createProvider(), createRequestWithEdit(bigFile, { insertionOffset: 3, insertedText: 'a' })); + + await configService.setConfig(ConfigKey.TeamInternal.InlineEditsXtabGlobalBudget, '{}'); + const enabledAtDefault = await captureUserPrompt(createProvider(), createRequestWithEdit(bigFile, { insertionOffset: 3, insertedText: 'a' })); + + // Legacy caps the current file at 2000 tokens → the tail is trimmed. + expect(legacy).not.toContain('const value399 = 399;'); + // Clip-last lets the current file reuse the whole pool → the entire file fits. + expect(enabledAtDefault).toContain('const value399 = 399;'); + expect(currentFileRegionLineCount(legacy)).toBeLessThan(currentFileRegionLineCount(enabledAtDefault)); + }); + + // New behavior: because the current file is sized to its share PLUS the cascade + // leftover (≈ the whole pool when the cascade is empty), a larger total budget + // keeps more of the file. A small pool still trims the tail; a generous pool + // fits the entire file. + it('keeps more of the current file as the total budget grows', async () => { + await configService.setConfig(ConfigKey.TeamInternal.InlineEditsXtabGlobalBudget, JSON.stringify({ totalTokens: 2000 })); + const smallBudget = await captureUserPrompt(createProvider(), createRequestWithEdit(bigFile, { insertionOffset: 3, insertedText: 'a' })); + + await configService.setConfig(ConfigKey.TeamInternal.InlineEditsXtabGlobalBudget, JSON.stringify({ totalTokens: 8000 })); + const wideBudget = await captureUserPrompt(createProvider(), createRequestWithEdit(bigFile, { insertionOffset: 3, insertedText: 'a' })); + + // Small pool trims the tail; wide pool fits the whole file. + expect(smallBudget).not.toContain('const value399 = 399;'); + expect(wideBudget).toContain('const value399 = 399;'); + expect(currentFileRegionLineCount(smallBudget)).toBeLessThan(currentFileRegionLineCount(wideBudget)); + }); + }); + describe('filter pipeline', () => { it('filters out import-only changes', async () => { const provider = createProvider(); diff --git a/extensions/copilot/src/platform/configuration/common/configurationService.ts b/extensions/copilot/src/platform/configuration/common/configurationService.ts index 82e103edfeb..de1853ef334 100644 --- a/extensions/copilot/src/platform/configuration/common/configurationService.ts +++ b/extensions/copilot/src/platform/configuration/common/configurationService.ts @@ -885,8 +885,7 @@ export namespace ConfigKey { export const InlineEditsXtabLanguageContextMaxTokens = defineTeamInternalSetting('chat.advanced.inlineEdits.xtabProvider.languageContext.maxTokens', ConfigType.ExperimentBased, xtabPromptOptions.DEFAULT_OPTIONS.languageContext.maxTokens); export const InlineEditsXtabIncludeNeighborFiles = defineTeamInternalSetting('chat.advanced.inlineEdits.xtabProvider.neighborFiles.enabled', ConfigType.ExperimentBased, xtabPromptOptions.DEFAULT_OPTIONS.neighborFiles.enabled); export const InlineEditsXtabNeighborFilesMaxTokens = defineTeamInternalSetting('chat.advanced.inlineEdits.xtabProvider.neighborFiles.maxTokens', ConfigType.ExperimentBased, xtabPromptOptions.DEFAULT_OPTIONS.neighborFiles.maxTokens); - export const InlineEditsXtabGlobalBudgetEnabled = defineTeamInternalSetting('chat.advanced.inlineEdits.xtabProvider.globalBudget.enabled', ConfigType.ExperimentBased, false); - export const InlineEditsXtabGlobalBudgetTotalTokens = defineTeamInternalSetting('chat.advanced.inlineEdits.xtabProvider.globalBudget.totalTokens', ConfigType.ExperimentBased, xtabPromptOptions.GlobalBudgetOptions.DEFAULT_TOTAL_TOKENS); + export const InlineEditsXtabGlobalBudget = defineTeamInternalSetting('chat.advanced.inlineEdits.xtabProvider.globalBudget', ConfigType.ExperimentBased, undefined); export const InlineEditsXtabMaxMergeConflictLines = defineTeamInternalSetting('chat.advanced.inlineEdits.xtabProvider.maxMergeConflictLines', ConfigType.ExperimentBased, undefined); export const InlineEditsXtabOnlyMergeConflictLines = defineTeamInternalSetting('chat.advanced.inlineEdits.xtabProvider.onlyMergeConflictLines', ConfigType.ExperimentBased, false); export const InlineEditsXtabDuplicateAdditionsMode = defineTeamInternalSetting('chat.advanced.inlineEdits.xtabProvider.diffPatch.duplicateAdditionsMode', ConfigType.ExperimentBased, DuplicateAdditionsMode.Off, DuplicateAdditionsMode.VALIDATOR); diff --git a/extensions/copilot/src/platform/inlineEdits/common/dataTypes/xtabPromptOptions.ts b/extensions/copilot/src/platform/inlineEdits/common/dataTypes/xtabPromptOptions.ts index aeeb53e8c50..8e6bf8f62c5 100644 --- a/extensions/copilot/src/platform/inlineEdits/common/dataTypes/xtabPromptOptions.ts +++ b/extensions/copilot/src/platform/inlineEdits/common/dataTypes/xtabPromptOptions.ts @@ -3,8 +3,9 @@ * Licensed under the MIT License. See License.txt in the project root for license information. *--------------------------------------------------------------------------------------------*/ +import { Result } from '../../../../util/common/result'; import { assertNever } from '../../../../util/vs/base/common/assert'; -import { IValidator, vBoolean, vEnum, vNumber, vObj, vRequired, vString, vUndefined, vUnion } from '../../../configuration/common/validator'; +import { IValidator, vArray, vBoolean, vEnum, vNumber, vObj, vRequired, vString, vUndefined, vUnion } from '../../../configuration/common/validator'; import { ImportChanges } from './importFilteringOptions'; export enum IncludeLineNumbersOption { @@ -82,8 +83,14 @@ export type DiffHistoryOptions = { }; /** - * Parts that participate in the global-budget cascade. `currentFile` and lint - * output are intentionally excluded and continue to use their own per-part caps. + * Parts that are rendered by the global-budget cascade and listed in `order`. + * Lint output is intentionally excluded and keeps its own per-part shape. + * + * `currentFile` is NOT in this set: it is clipped outside the cascade (around the + * cursor) but still draws an allocation from the pool via {@link GlobalBudgetSharePart}. + * It is clipped LAST (after the cascade) to its share plus whatever budget the + * cascade left unused, so it only ever receives leftover and never donates into + * the cascade. */ export type GlobalBudgetPart = | 'recentlyViewedDocuments' @@ -91,30 +98,43 @@ export type GlobalBudgetPart = | 'neighborFiles' | 'diffHistory'; +/** + * Parts that receive a `shares` allocation from the pool: every rendered + * {@link GlobalBudgetPart} plus `currentFile`, which is sized from the pool but + * clipped externally (see {@link GlobalBudgetOptions.currentFileBudget}). + */ +export type GlobalBudgetSharePart = GlobalBudgetPart | 'currentFile'; + /** * Opt-in global-budget allocation modelled after the cascade in * `CascadingPromptFactory` (completions-core): every participating part gets a - * percentage share of a single `totalTokens` pool, parts are rendered in - * `order`, and any unused tokens in one part cascade as surplus to the next. + * percentage share of a single `totalTokens` pool, the rendered parts are emitted + * in `order`, and any unused tokens in one part cascade as surplus to the next. + * + * `currentFile` participates through `shares` only: it is clipped LAST, after the + * cascade runs, to its share of the pool plus whatever budget the cascade left + * unused (its `finalSurplus`), so it reuses the leftover and trims less. * * When `undefined` (the default), each part uses its own `maxTokens` cap as * before and no cross-part budget reuse happens. */ export type GlobalBudgetOptions = { readonly totalTokens: number; - /** Cascade order. Earlier parts get budget first; their surplus flows to later parts. */ + /** Cascade render order. Earlier parts get budget first; their surplus flows to later parts. `currentFile` is not listed here. */ readonly order: readonly GlobalBudgetPart[]; - /** Share of `totalTokens` allocated to each part. Must sum to 1 across `order`. */ - readonly shares: Readonly>; + /** Share of `totalTokens` allocated to each part. Must sum to ~1 across `order` plus `currentFile`. */ + readonly shares: Readonly>; }; export namespace GlobalBudgetOptions { /** - * Default cascade: language context donates first (often disabled), then - * recently-viewed documents (always-on, accepts most of the surplus), then - * neighbor files (must run after recently-viewed because it consults - * `docsInPrompt` to avoid duplicating recently-viewed documents), then - * diff history. + * Default render order: language context first (often disabled, donates its + * share onward), then recently-viewed documents (always-on, accepts most of + * the surplus), then neighbor files (must run after recently-viewed because it + * consults `docsInPrompt` to avoid duplicating recently-viewed documents), + * then diff history. `currentFile` is sized from the pool but clipped last + * (after the cascade, so it can reuse the cascade's leftover), so it is not + * part of this order. */ export const DEFAULT_ORDER: readonly GlobalBudgetPart[] = [ 'languageContext', @@ -123,15 +143,149 @@ export namespace GlobalBudgetOptions { 'diffHistory', ]; - /** Shares matching today's per-part `maxTokens` ratios (volume-neutral baseline). */ - export const DEFAULT_SHARES: Readonly> = { - recentlyViewedDocuments: 2 / 6, - languageContext: 2 / 6, - neighborFiles: 1 / 6, - diffHistory: 1 / 6, + /** + * Shares matching today's per-part `maxTokens` ratios (volume-neutral + * baseline): each part's share is its legacy cap divided by the pool total + * ({@link DEFAULT_TOTAL_TOKENS}), so `floor(totalTokens * share)` reproduces + * that cap exactly. Sum to 1 across all parts. + */ + export const DEFAULT_SHARES: Readonly> = { + currentFile: 1500 / 7500, + recentlyViewedDocuments: 2000 / 7500, + languageContext: 2000 / 7500, + neighborFiles: 1000 / 7500, + diffHistory: 1000 / 7500, }; - export const DEFAULT_TOTAL_TOKENS = 6000; + /** + * The pool size: the sum of today's per-part `maxTokens` caps + * (`currentFile` 1500 + `recentlyViewedDocuments` 2000 + `languageContext` + * 2000 + `neighborFiles` 1000 + `diffHistory` 1000), so the default global + * budget neither grows nor shrinks the total available budget. + */ + export const DEFAULT_TOTAL_TOKENS = 7500; + + /** + * The token budget allotted to `currentFile` from the pool: `floor(totalTokens + * * shares.currentFile)`. Used to override the per-part `currentFile.maxTokens` + * cap when clipping the current file under a global budget. + */ + export function currentFileBudget(globalBudget: GlobalBudgetOptions): number { + return Math.max(0, Math.floor(globalBudget.totalTokens * (globalBudget.shares.currentFile ?? 0))); + } + + /** + * Validate {@link GlobalBudgetOptions} since it is runtime-configurable + * (e.g. via experiments). Catches misconfigurations that would otherwise + * cause silent, hard-to-debug behavior: + * - `totalTokens` not a finite, non-negative number + * - duplicate parts in `order` (would render the same part twice) + * - missing share for any part in `order` or for `currentFile` + * - any share not a finite, non-negative number (negative shares would break + * budget conservation: a negative allocation clamps to 0 yet still counts + * toward the share sum, letting the remaining parts over-allocate past the pool) + * - shares not summing to ~1 across `order` plus `currentFile` (would over/under-allocate) + * - `neighborFiles` ordered before `recentlyViewedDocuments` (the former + * consults `docsInPrompt` populated by the latter) + */ + export function validate(globalBudget: GlobalBudgetOptions): void { + if (!Number.isFinite(globalBudget.totalTokens) || globalBudget.totalTokens < 0) { + throw new Error(`globalBudget.totalTokens must be a finite, non-negative number, got ${globalBudget.totalTokens}`); + } + + const seen = new Set(); + for (const part of globalBudget.order) { + if (seen.has(part)) { + throw new Error(`globalBudget.order contains duplicate part '${part}'`); + } + seen.add(part); + if (typeof globalBudget.shares[part] !== 'number') { + throw new Error(`globalBudget.shares is missing entry for '${part}'`); + } + if (!Number.isFinite(globalBudget.shares[part]) || globalBudget.shares[part] < 0) { + throw new Error(`globalBudget.shares['${part}'] must be a finite, non-negative number, got ${globalBudget.shares[part]}`); + } + } + + if (typeof globalBudget.shares.currentFile !== 'number') { + throw new Error(`globalBudget.shares is missing entry for 'currentFile'`); + } + if (!Number.isFinite(globalBudget.shares.currentFile) || globalBudget.shares.currentFile < 0) { + throw new Error(`globalBudget.shares['currentFile'] must be a finite, non-negative number, got ${globalBudget.shares.currentFile}`); + } + + const recentIdx = globalBudget.order.indexOf('recentlyViewedDocuments'); + const neighborIdx = globalBudget.order.indexOf('neighborFiles'); + if (recentIdx !== -1 && neighborIdx !== -1 && neighborIdx < recentIdx) { + throw new Error(`globalBudget.order must place 'recentlyViewedDocuments' before 'neighborFiles'`); + } + + const sharesSum = globalBudget.order.reduce((sum, part) => sum + globalBudget.shares[part], 0) + globalBudget.shares.currentFile; + const epsilon = 1e-3; + if (Math.abs(sharesSum - 1) > epsilon) { + throw new Error(`globalBudget.shares across order must sum to ~1, got ${sharesSum}`); + } + } + + /** + * Structural validator for the experiment-driven JSON config string. Every + * top-level field is optional; omitted fields fall back to the defaults in + * {@link fromConfigString}. When `shares` is provided it must list every + * part (the rendered cascade parts plus `currentFile`) so the pool stays + * fully allocated; partial `shares` objects are rejected. + */ + export const VALIDATOR = vObj({ + 'totalTokens': vUnion(vNumber(), vUndefined()), + 'order': vUnion(vArray(vEnum('languageContext', 'recentlyViewedDocuments', 'neighborFiles', 'diffHistory')), vUndefined()), + 'shares': vUnion(vObj({ + 'currentFile': vRequired(vNumber()), + 'recentlyViewedDocuments': vRequired(vNumber()), + 'languageContext': vRequired(vNumber()), + 'neighborFiles': vRequired(vNumber()), + 'diffHistory': vRequired(vNumber()), + }), vUndefined()), + }); + + /** + * Parse the single experiment-driven JSON config string (modelled after + * `modelConfigurationString`) into a fully-populated {@link GlobalBudgetOptions}. + * + * Any field absent from the JSON falls back to its `DEFAULT_*` value, so + * `'{}'` enables the global budget with the volume-neutral defaults and + * `'{"totalTokens":6000}'` only overrides the pool size. The merged result is + * then run through {@link validate} so semantic misconfigurations (bad share + * sums, duplicate/misordered parts) are reported. + * + * Returns a {@link Result.error} (never throws) so callers can fall back to a + * disabled global budget and report the error via telemetry. + */ + export function fromConfigString(configString: string): Result { + let parsed: unknown; + try { + parsed = JSON.parse(configString); + } catch (e) { + return Result.error(`Failed to parse globalBudget config string: ${e instanceof Error ? e.message : String(e)}`); + } + + const { content, error } = VALIDATOR.validate(parsed); + if (error) { + return Result.error(`globalBudget config validation failed: ${error.message}`); + } + + const options: GlobalBudgetOptions = { + totalTokens: content.totalTokens ?? DEFAULT_TOTAL_TOKENS, + order: content.order ?? DEFAULT_ORDER, + shares: content.shares ?? DEFAULT_SHARES, + }; + + try { + validate(options); + } catch (e) { + return Result.error(e instanceof Error ? e.message : String(e)); + } + + return Result.ok(options); + } } export type PagedClipping = { pageSize: number }; diff --git a/extensions/copilot/src/platform/inlineEdits/test/common/xtabPromptOptions.spec.ts b/extensions/copilot/src/platform/inlineEdits/test/common/xtabPromptOptions.spec.ts index c2f8bb27877..31c46134d8b 100644 --- a/extensions/copilot/src/platform/inlineEdits/test/common/xtabPromptOptions.spec.ts +++ b/extensions/copilot/src/platform/inlineEdits/test/common/xtabPromptOptions.spec.ts @@ -5,7 +5,7 @@ import { describe, expect, it } from 'vitest'; import { ImportChanges } from '../../common/dataTypes/importFilteringOptions'; -import { applyStrategyConfig, IncludeLineNumbersOption, MODEL_CONFIGURATION_VALIDATOR, ModelConfiguration, PromptingStrategy } from '../../common/dataTypes/xtabPromptOptions'; +import { applyStrategyConfig, DEFAULT_OPTIONS, GlobalBudgetOptions, IncludeLineNumbersOption, MODEL_CONFIGURATION_VALIDATOR, ModelConfiguration, PromptingStrategy } from '../../common/dataTypes/xtabPromptOptions'; function baseConfig(overrides: Partial = {}): ModelConfiguration { return { @@ -102,3 +102,155 @@ describe('MODEL_CONFIGURATION_VALIDATOR', () => { expect(result.error).toBeDefined(); }); }); + +describe('GlobalBudgetOptions', () => { + + function gb(overrides: Partial = {}): GlobalBudgetOptions { + return { + totalTokens: GlobalBudgetOptions.DEFAULT_TOTAL_TOKENS, + order: GlobalBudgetOptions.DEFAULT_ORDER, + shares: GlobalBudgetOptions.DEFAULT_SHARES, + ...overrides, + }; + } + + describe('volume-neutral defaults', () => { + // Guards the core no-regression promise: enabling the global budget with the + // default total + shares must reproduce today's per-part `maxTokens` caps + // exactly. If anyone changes DEFAULT_SHARES or DEFAULT_TOTAL_TOKENS in a way + // that shifts a part's budget, this fails loudly instead of silently + // shrinking/growing prompts in the experiment arm. + it('reproduce the legacy per-part caps', () => { + const total = GlobalBudgetOptions.DEFAULT_TOTAL_TOKENS; + const shares = GlobalBudgetOptions.DEFAULT_SHARES; + const computed = { + currentFile: Math.floor(total * shares.currentFile), + recentlyViewedDocuments: Math.floor(total * shares.recentlyViewedDocuments), + languageContext: Math.floor(total * shares.languageContext), + neighborFiles: Math.floor(total * shares.neighborFiles), + diffHistory: Math.floor(total * shares.diffHistory), + }; + expect(computed).toEqual({ + currentFile: DEFAULT_OPTIONS.currentFile.maxTokens, + recentlyViewedDocuments: DEFAULT_OPTIONS.recentlyViewedDocuments.maxTokens, + languageContext: DEFAULT_OPTIONS.languageContext.maxTokens, + neighborFiles: DEFAULT_OPTIONS.neighborFiles.maxTokens, + diffHistory: DEFAULT_OPTIONS.diffHistory.maxTokens, + }); + }); + + it('shares sum to exactly 1', () => { + const shares = GlobalBudgetOptions.DEFAULT_SHARES; + const sum = shares.currentFile + shares.recentlyViewedDocuments + shares.languageContext + shares.neighborFiles + shares.diffHistory; + expect(sum).toBe(1); + }); + }); + + describe('currentFileBudget', () => { + it('floors totalTokens * shares.currentFile', () => { + expect(GlobalBudgetOptions.currentFileBudget(gb({ totalTokens: 8000, shares: { ...GlobalBudgetOptions.DEFAULT_SHARES, currentFile: 2 / 8 } }))).toBe(2000); + expect(GlobalBudgetOptions.currentFileBudget(gb({ totalTokens: 999, shares: { ...GlobalBudgetOptions.DEFAULT_SHARES, currentFile: 1 / 3 } }))).toBe(333); + }); + + it('clamps at 0 for a zero share', () => { + expect(GlobalBudgetOptions.currentFileBudget(gb({ totalTokens: 8000, shares: { ...GlobalBudgetOptions.DEFAULT_SHARES, currentFile: 0 } }))).toBe(0); + }); + }); + + describe('validate', () => { + it('accepts the defaults', () => { + expect(() => GlobalBudgetOptions.validate(gb())).not.toThrow(); + }); + + it('throws on a duplicate part in order', () => { + expect(() => GlobalBudgetOptions.validate(gb({ + order: ['languageContext', 'languageContext', 'recentlyViewedDocuments', 'neighborFiles', 'diffHistory'], + }))).toThrow(/duplicate part 'languageContext'/); + }); + + it('throws when shares omit currentFile', () => { + expect(() => GlobalBudgetOptions.validate(gb({ + shares: { languageContext: 0.25, recentlyViewedDocuments: 0.25, neighborFiles: 0.25, diffHistory: 0.25 } as GlobalBudgetOptions['shares'], + }))).toThrow(/missing entry for 'currentFile'/); + }); + + it('throws when shares do not sum to ~1 across order plus currentFile', () => { + expect(() => GlobalBudgetOptions.validate(gb({ + shares: { ...GlobalBudgetOptions.DEFAULT_SHARES, currentFile: 0.9 }, + }))).toThrow(/shares across order must sum to ~1/); + }); + + it('throws when neighborFiles is ordered before recentlyViewedDocuments', () => { + expect(() => GlobalBudgetOptions.validate(gb({ + order: ['languageContext', 'neighborFiles', 'recentlyViewedDocuments', 'diffHistory'], + }))).toThrow(/must place 'recentlyViewedDocuments' before 'neighborFiles'/); + }); + + it('throws on a negative totalTokens', () => { + expect(() => GlobalBudgetOptions.validate(gb({ totalTokens: -1 }))).toThrow(/totalTokens must be a finite, non-negative number/); + }); + + it('throws on a negative share (which would break budget conservation)', () => { + // Sums to 1 so the legacy sum check passes, but a negative share clamps to + // a 0 allocation yet still counts toward the sum, letting other parts + // over-allocate past the pool. Must be rejected. + expect(() => GlobalBudgetOptions.validate(gb({ + shares: { currentFile: 0.25, languageContext: -0.25, recentlyViewedDocuments: 1.0, neighborFiles: 0, diffHistory: 0 }, + }))).toThrow(/must be a finite, non-negative number/); + }); + + it('throws on a non-finite share', () => { + expect(() => GlobalBudgetOptions.validate(gb({ + shares: { ...GlobalBudgetOptions.DEFAULT_SHARES, currentFile: Number.NaN }, + }))).toThrow(/must be a finite, non-negative number/); + }); + }); + + describe('fromConfigString', () => { + it('fills every omitted field with the defaults', () => { + const result = GlobalBudgetOptions.fromConfigString('{}'); + expect(result.isOk() && result.val).toEqual({ + totalTokens: GlobalBudgetOptions.DEFAULT_TOTAL_TOKENS, + order: GlobalBudgetOptions.DEFAULT_ORDER, + shares: GlobalBudgetOptions.DEFAULT_SHARES, + }); + }); + + it('overrides only the fields present in the JSON', () => { + const result = GlobalBudgetOptions.fromConfigString(JSON.stringify({ totalTokens: 6000 })); + expect(result.isOk() && result.val).toEqual({ + totalTokens: 6000, + order: GlobalBudgetOptions.DEFAULT_ORDER, + shares: GlobalBudgetOptions.DEFAULT_SHARES, + }); + }); + + it('parses a fully-specified config and ignores unknown keys', () => { + const order: GlobalBudgetOptions['order'] = ['languageContext', 'recentlyViewedDocuments', 'neighborFiles', 'diffHistory']; + const shares: GlobalBudgetOptions['shares'] = { currentFile: 0.4, languageContext: 0.2, recentlyViewedDocuments: 0.2, neighborFiles: 0.1, diffHistory: 0.1 }; + const result = GlobalBudgetOptions.fromConfigString(JSON.stringify({ totalTokens: 12000, order, shares, unknown: 'ignored' })); + expect(result.isOk() && result.val).toEqual({ totalTokens: 12000, order, shares }); + }); + + it('errors on malformed JSON', () => { + expect(GlobalBudgetOptions.fromConfigString('{ not json').isError()).toBe(true); + }); + + it('errors when a field has the wrong type', () => { + expect(GlobalBudgetOptions.fromConfigString(JSON.stringify({ totalTokens: '6000' })).isError()).toBe(true); + }); + + it('errors on an unknown part in order', () => { + expect(GlobalBudgetOptions.fromConfigString(JSON.stringify({ order: ['languageContext', 'bogus'] })).isError()).toBe(true); + }); + + it('errors when shares are partial (every part is required)', () => { + expect(GlobalBudgetOptions.fromConfigString(JSON.stringify({ shares: { currentFile: 1 } })).isError()).toBe(true); + }); + + it('errors when the merged config is semantically invalid', () => { + const shares = { currentFile: 0.9, languageContext: 0.2, recentlyViewedDocuments: 0.2, neighborFiles: 0.1, diffHistory: 0.1 }; + expect(GlobalBudgetOptions.fromConfigString(JSON.stringify({ shares })).isError()).toBe(true); + }); + }); +});