NES: experiment-driven global token budget for the inline-edits prompt (#323948)

* nes: fold currentFile budgeting into the global budget cascade

Under the opt-in global budget experiment, the current file now draws its
clip budget from the shared pool (its `shares.currentFile` slice of
`totalTokens`) and donates whatever it does not use to the cascade as the
initial surplus. When the global budget is disabled (prod default) behavior
is byte-identical, and the new defaults (totalTokens 8000, shares in eighths)
reproduce today's per-part `maxTokens` caps exactly.

- xtabPromptOptions: add `GlobalBudgetSharePart`, `currentFileBudget()` and
  `validate()`; grow `shares` to include `currentFile`; bump
  `DEFAULT_TOTAL_TOKENS` to 8000 and rebalance `DEFAULT_SHARES`.
- promptCrafting: seed the cascade's initial surplus from the current file's
  leftover budget; delegate validation to the namespace helper.
- xtabProvider: size and clip the current file from the pool and compute the
  surplus, gated on `globalBudget` being defined.
- Tests: volume-neutral invariant, `currentFileBudget`/`validate` units, and
  provider-level regression + new-behavior coverage.
- Docs: rewrite globalBudgetCascade.md for the seeded-surplus design.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* nes: collapse global-budget settings into a single exp-driven JSON string

Replace the two experiment-driven global-budget settings
(globalBudget.enabled boolean + globalBudget.totalTokens number) with a
single JSON-encoded string setting, modelled after modelConfigurationString,
that defines totalTokens, order, and shares together.

- configurationService: remove InlineEditsXtabGlobalBudgetEnabled and
  InlineEditsXtabGlobalBudgetTotalTokens; add InlineEditsXtabGlobalBudget
  (string | undefined, ExperimentBased, default undefined).
- xtabPromptOptions: add GlobalBudgetOptions.VALIDATOR (all top-level fields
  optional; shares must be complete when present) and a pure
  GlobalBudgetOptions.fromConfigString(): Result<GlobalBudgetOptions, string>
  that JSON-parses, structurally validates, merges over the defaults, then
  runs the semantic validate(). Never throws.
- xtabProvider: inject ITelemetryService; resolve the budget via
  getGlobalBudget(), returning undefined (disabled, identical to prod) when
  the string is unset/empty/invalid and emitting incorrectNesGlobalBudgetConfig
  telemetry on parse/validation failure.
- Tests: migrate the provider global-budget tests to the JSON string and add
  fromConfigString unit tests.
- Docs: rewrite the globalBudgetCascade.md Wiring section and migration note.

No regression: an unset string keeps the byte-identical legacy path. '{}'
enables the budget with the volume-neutral defaults; {"totalTokens":N} only
overrides the pool size.

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>

* nes: clip the current file last under a global budget to reuse cascade leftover

Under a global budget the current file is now clipped LAST, sized to its base
share plus whatever the budget cascade left unused (`currentFileBudget +
finalSurplus`), so it trims less and reuses leftover budget. The cascade runs
first (seeded with 0, the current file never donates) and is threaded into
`getUserPrompt` via `precomputedCascade` so it renders exactly once.

This replaces the earlier two-mode design: the opt-in `reuseLeftoverForCurrentFile`
flag and the seeded-surplus "donate-forward" path are dropped, along with all
`currentFileBudgetSurplus`/`initialSurplus` threading. Clip-last is now the single
behavior under a global budget. The validator hardening (rejecting non-finite or
negative `totalTokens`/shares) is kept.

No production regression: when `globalBudget` is undefined the current file is
still clipped to its own `currentFile.maxTokens` with the legacy await ordering,
byte-identical to before. The global budget remains experiment-gated and off in
prod.

Tests and globalBudgetCascade.md updated for the clip-last-only behavior.

Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>

* nes: address review — restore nLinesOfCurrentFileInPrompt telemetry timing

- xtabProvider: record setNLinesOfCurrentFileInPrompt inside both budget
  branches. In the legacy/prod branch the clip happens before the
  context-gathering awaits, so the telemetry call is moved back there to
  match prod timing exactly (it is still emitted when a request is
  cancelled mid-gathering). The clip-last branch records it after the
  cascade, as the clip is inherently last there.
- xtabPromptOptions: fix two stale JSDoc blocks that still described the
  removed donate-forward behavior; describe clip-last instead.
- globalBudgetCascade.md: correct the conservation bound to T·(Σ shares)
  and document the share-sum tolerance (over-allocation ≤ ~1e-3·T).
- promptCrafting.spec: add a cascade test asserting finalSurplus shrinks
  as the cascade consumes budget (current file reuses less leftover).

Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>

* nes: fix stale DEFAULT_ORDER doc comment to say clip-last

Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>

* nes: rename clipCurrentFileToBudget param to overriddenMaxTokens

Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>

* nes: extract current-file clip + context gathering into a method

Pulls the ~55-line global-budget/legacy if-else out of doGetNextEditWithSelection
into a private gatherContextAndClipCurrentFile method returning
Result<{ clippedTaggedCurrentDoc, areaAroundCodeToEdit, precomputedCascade,
langCtx, neighborSnippets }, NoNextEditReason>. Behavior, clip-first/clip-last
ordering, and nLinesOfCurrentFileInPrompt telemetry timing are unchanged.

Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>

* nes: recalibrate global-budget defaults to merged currentFile cap (1500)

main (#322382) lowered currentFile.maxTokens 2000->1500, which broke the
global-budget 'volume-neutral defaults' invariant (floor(total*share) must
equal each part's legacy cap). Recalibrate DEFAULT_SHARES to cap/7500 and set
DEFAULT_TOTAL_TOKENS to 7500 (the sum of the merged per-part caps: 1500 + 2000
+ 2000 + 1000 + 1000), so shares still sum to exactly 1 and every part's base
allocation reproduces its cap. finalSurplus stays 6000 (7500 - 1500).

Update dependent specs (literal 8000 -> DEFAULT_TOTAL_TOKENS, stale 2000 cap
comments -> 1500) and the design doc (defaults table, worked examples, effective
caps, ordering caveat, migration note) for T=7500 and currentFile base 1500.

Co-authored-by: Copilot App <223556219+Copilot@users.noreply.github.com>

---------

Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
This commit is contained in:
Ulugbek Abdullaev
2026-07-02 21:48:34 +05:00
committed by GitHub
parent 881d17aced
commit 3073c2bb48
8 changed files with 1001 additions and 250 deletions
@@ -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<GlobalBudgetPart, number>`. 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<GlobalBudgetSharePart, number>` — 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.
@@ -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<DocumentId>;
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<string>();
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\`\`\``;
}
@@ -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<LanguageContextResponse | undefined>,
gatherNeighborSnippets: () => Promise<readonly INeighborFileSnippet[] | undefined>,
clipCurrentFileToBudget: (overriddenMaxTokens: number | undefined) => Result<{ clippedTaggedCurrentDoc: ClippedDocument; areaAroundCodeToEdit: string }, 'outOfBudget'>,
): Promise<Result<{
clippedTaggedCurrentDoc: ClippedDocument;
areaAroundCodeToEdit: string;
precomputedCascade: CascadeResult | undefined;
langCtx: LanguageContextResponse | undefined;
neighborSnippets: readonly INeighborFileSnippet[] | undefined;
}, NoNextEditReason>> {
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);
@@ -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<typeof runGlobalBudgetCascade> }): PromptPieces {
const { activeDoc, currentDocument, currentDocLines } = makeActiveDoc();
const promptOptions: PromptOptions = {
...DEFAULT_OPTIONS,
@@ -854,14 +857,28 @@ describe('getUserPrompt — globalBudget cascade', () => {
[],
currentDocLines,
'<area>some code</area>',
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);
});
});
@@ -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<string> {
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();
@@ -885,8 +885,7 @@ export namespace ConfigKey {
export const InlineEditsXtabLanguageContextMaxTokens = defineTeamInternalSetting<number>('chat.advanced.inlineEdits.xtabProvider.languageContext.maxTokens', ConfigType.ExperimentBased, xtabPromptOptions.DEFAULT_OPTIONS.languageContext.maxTokens);
export const InlineEditsXtabIncludeNeighborFiles = defineTeamInternalSetting<boolean>('chat.advanced.inlineEdits.xtabProvider.neighborFiles.enabled', ConfigType.ExperimentBased, xtabPromptOptions.DEFAULT_OPTIONS.neighborFiles.enabled);
export const InlineEditsXtabNeighborFilesMaxTokens = defineTeamInternalSetting<number>('chat.advanced.inlineEdits.xtabProvider.neighborFiles.maxTokens', ConfigType.ExperimentBased, xtabPromptOptions.DEFAULT_OPTIONS.neighborFiles.maxTokens);
export const InlineEditsXtabGlobalBudgetEnabled = defineTeamInternalSetting<boolean>('chat.advanced.inlineEdits.xtabProvider.globalBudget.enabled', ConfigType.ExperimentBased, false);
export const InlineEditsXtabGlobalBudgetTotalTokens = defineTeamInternalSetting<number>('chat.advanced.inlineEdits.xtabProvider.globalBudget.totalTokens', ConfigType.ExperimentBased, xtabPromptOptions.GlobalBudgetOptions.DEFAULT_TOTAL_TOKENS);
export const InlineEditsXtabGlobalBudget = defineTeamInternalSetting<string | undefined>('chat.advanced.inlineEdits.xtabProvider.globalBudget', ConfigType.ExperimentBased, undefined);
export const InlineEditsXtabMaxMergeConflictLines = defineTeamInternalSetting<number | undefined>('chat.advanced.inlineEdits.xtabProvider.maxMergeConflictLines', ConfigType.ExperimentBased, undefined);
export const InlineEditsXtabOnlyMergeConflictLines = defineTeamInternalSetting<boolean>('chat.advanced.inlineEdits.xtabProvider.onlyMergeConflictLines', ConfigType.ExperimentBased, false);
export const InlineEditsXtabDuplicateAdditionsMode = defineTeamInternalSetting<DuplicateAdditionsMode>('chat.advanced.inlineEdits.xtabProvider.diffPatch.duplicateAdditionsMode', ConfigType.ExperimentBased, DuplicateAdditionsMode.Off, DuplicateAdditionsMode.VALIDATOR);
@@ -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<Record<GlobalBudgetPart, number>>;
/** Share of `totalTokens` allocated to each part. Must sum to ~1 across `order` plus `currentFile`. */
readonly shares: Readonly<Record<GlobalBudgetSharePart, number>>;
};
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<Record<GlobalBudgetPart, number>> = {
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<Record<GlobalBudgetSharePart, number>> = {
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<string>();
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<GlobalBudgetPart[]>('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<GlobalBudgetOptions, string> {
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 };
@@ -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> = {}): ModelConfiguration {
return {
@@ -102,3 +102,155 @@ describe('MODEL_CONFIGURATION_VALIDATOR', () => {
expect(result.error).toBeDefined();
});
});
describe('GlobalBudgetOptions', () => {
function gb(overrides: Partial<GlobalBudgetOptions> = {}): 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);
});
});
});