Step 254: response budget system (truncation, continuation, pagination)

Optional budget parameter on getAST and getDiagnostics. Responses exceeding
the budget are truncated with continuation tokens for paginated follow-up.
Binary search finds max array elements fitting the budget. Diagnostics sorted
by severity (errors first) before truncation.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
This commit is contained in:
Bill
2026-02-11 15:54:01 +00:00
parent d6941da793
commit ba9da3fd55
6 changed files with 522 additions and 3 deletions

View File

@@ -1468,4 +1468,15 @@ target_link_libraries(step253_test PRIVATE
tree_sitter_java tree_sitter_rust tree_sitter_go
tree_sitter_org)
# Step 254: Response budget system
add_executable(step254_test tests/step254_test.cpp)
target_include_directories(step254_test PRIVATE src)
target_link_libraries(step254_test PRIVATE
nlohmann_json::nlohmann_json
unofficial::tree-sitter::tree-sitter
tree_sitter_python tree_sitter_cpp tree_sitter_elisp
tree_sitter_javascript tree_sitter_typescript
tree_sitter_java tree_sitter_rust tree_sitter_go
tree_sitter_org)
# Step 12: Dear ImGui shell scaffolding created (main.cpp exists but not built due to dependencies)

View File

@@ -81,6 +81,11 @@ inline json handleHeadlessAgentRequest(HeadlessEditorState& state,
}
result["version"] = state.active()->versionTracker.version;
result["tokenEstimate"] = tokenEstimate(result);
int budget = params.value("budget", 0);
if (budget > 0) {
auto br = applyBudget(result, budget);
return headlessRpcResult(id, br.result);
}
return headlessRpcResult(id, result);
}
@@ -601,11 +606,19 @@ inline json handleHeadlessAgentRequest(HeadlessEditorState& state,
// Record snapshot for delta tracking
auto allDiagsUnfiltered = collectAllDiagnostics(state.activeAST());
state.active()->diagTracker.recordSnapshot(allDiagsUnfiltered);
return headlessRpcResult(id, {
{"diagnostics", diagnosticsToJson(diags)},
json diagJson = diagnosticsToJson(diags);
sortDiagnosticsByPriority(diagJson);
json result = {
{"diagnostics", diagJson},
{"count", (int)diags.size()},
{"version", state.active()->diagTracker.version}
});
};
int budget = params.value("budget", 0);
if (budget > 0) {
auto br = applyBudget(result, budget);
return headlessRpcResult(id, br.result);
}
return headlessRpcResult(id, result);
}
// --- getDiagnosticsDelta ---

View File

@@ -25,6 +25,7 @@
#include "FileOperations.h"
#include "CompactAST.h"
#include "StructuredDiagnostics.h"
#include "ResponseBudget.h"
#include "Orchestrator.h"
#include "IncrementalOptimizer.h"
#include "TextASTSync.h"

168
editor/src/ResponseBudget.h Normal file
View File

@@ -0,0 +1,168 @@
#pragma once
// Step 254: Response Budget System
//
// Lets agents request only what they need. Query methods gain an optional
// `budget` parameter (max response chars). If exceeded, the response is
// truncated with `truncated: true` and a `continuation` token for follow-up.
#include <string>
#include <nlohmann/json.hpp>
using json = nlohmann::json;
// -----------------------------------------------------------------------
// Budget application: truncate a JSON result if it exceeds the budget
// -----------------------------------------------------------------------
struct BudgetResult {
json result; // the (possibly truncated) response
bool truncated = false;
std::string continuation; // opaque token for follow-up
};
// Apply budget to a result JSON. If the serialized size exceeds
// maxChars, tries to trim array fields to fit.
inline BudgetResult applyBudget(const json& fullResult, int maxChars) {
BudgetResult br;
if (maxChars <= 0) {
// No budget → return as-is
br.result = fullResult;
return br;
}
std::string dump = fullResult.dump();
if ((int)dump.size() <= maxChars) {
// Fits within budget
br.result = fullResult;
return br;
}
// Need to truncate. Strategy: find the largest array in the result
// and progressively shrink it until we fit.
br.result = fullResult;
br.truncated = true;
// Find the array field to truncate
std::string arrayField;
int maxArraySize = 0;
for (auto& [key, val] : br.result.items()) {
if (val.is_array() && (int)val.size() > maxArraySize) {
arrayField = key;
maxArraySize = (int)val.size();
}
}
if (arrayField.empty() || maxArraySize == 0) {
// No array to truncate — mark as truncated but return as-is
br.result["truncated"] = true;
br.continuation = "none";
return br;
}
// Binary search for the right number of elements
int lo = 0, hi = maxArraySize;
int bestFit = 0;
while (lo <= hi) {
int mid = (lo + hi) / 2;
json trial = br.result;
json sliced = json::array();
for (int i = 0; i < mid && i < maxArraySize; ++i)
sliced.push_back(trial[arrayField][i]);
trial[arrayField] = sliced;
trial["truncated"] = true;
trial["totalCount"] = maxArraySize;
trial["returnedCount"] = mid;
if ((int)trial.dump().size() <= maxChars) {
bestFit = mid;
lo = mid + 1;
} else {
hi = mid - 1;
}
}
// Build the truncated result
json sliced = json::array();
for (int i = 0; i < bestFit && i < maxArraySize; ++i)
sliced.push_back(br.result[arrayField][i]);
br.result[arrayField] = sliced;
br.result["truncated"] = true;
br.result["totalCount"] = maxArraySize;
br.result["returnedCount"] = bestFit;
// Continuation token: "field:offset:total"
br.continuation = arrayField + ":" +
std::to_string(bestFit) + ":" +
std::to_string(maxArraySize);
br.result["continuation"] = br.continuation;
return br;
}
// -----------------------------------------------------------------------
// Continue from a previous truncated response
// -----------------------------------------------------------------------
inline json applyContinuation(const json& fullResult,
const std::string& continuation,
int maxChars) {
// Parse continuation token: "field:offset:total"
auto pos1 = continuation.find(':');
auto pos2 = continuation.find(':', pos1 + 1);
if (pos1 == std::string::npos || pos2 == std::string::npos)
return {{"error", "Invalid continuation token"}};
std::string field = continuation.substr(0, pos1);
int offset = std::stoi(continuation.substr(pos1 + 1, pos2 - pos1 - 1));
int total = std::stoi(continuation.substr(pos2 + 1));
if (!fullResult.contains(field) || !fullResult[field].is_array())
return {{"error", "Continuation field not found"}};
int arraySize = (int)fullResult[field].size();
json result = fullResult;
// Slice from offset
json sliced = json::array();
int count = 0;
for (int i = offset; i < arraySize; ++i) {
sliced.push_back(fullResult[field][i]);
++count;
// Check if we're approaching the budget
json trial = result;
trial[field] = sliced;
trial["returnedCount"] = count;
trial["offset"] = offset;
if (maxChars > 0 && (int)trial.dump().size() > maxChars) {
sliced.erase(sliced.size() - 1); // remove last
--count;
break;
}
}
result[field] = sliced;
result["returnedCount"] = count;
result["offset"] = offset;
result["totalCount"] = total;
if (offset + count < arraySize) {
result["truncated"] = true;
result["continuation"] = field + ":" +
std::to_string(offset + count) + ":" +
std::to_string(total);
} else {
result["truncated"] = false;
}
return result;
}
// -----------------------------------------------------------------------
// Priority-sort diagnostics (highest severity first) for budget
// -----------------------------------------------------------------------
inline void sortDiagnosticsByPriority(json& diagArray) {
if (!diagArray.is_array()) return;
std::sort(diagArray.begin(), diagArray.end(),
[](const json& a, const json& b) {
return a.value("severityLevel", 4) <
b.value("severityLevel", 4);
});
}

View File

@@ -0,0 +1,292 @@
// Step 254 TDD Test: Response Budget System
//
// Tests budget parameter on getAST and getDiagnostics: truncation,
// continuation tokens, pagination, and backward compatibility.
#include "HeadlessEditorState.h"
#include "MCPBridge.h"
#include <iostream>
#include <string>
static void expect(bool cond, const std::string& name,
int& passed, int& failed) {
if (cond) {
std::cout << "Test " << (passed + failed + 1)
<< " PASS: " << name << "\n";
++passed;
} else {
std::cout << "Test " << (passed + failed + 1)
<< " FAIL: " << name << "\n";
++failed;
}
}
static json rpc(HeadlessEditorState& state, const std::string& session,
const std::string& method,
json params = json::object()) {
json request = {{"jsonrpc", "2.0"}, {"id", 1},
{"method", method}, {"params", params}};
return state.processAgentRequest(request, session);
}
int main() {
int passed = 0;
int failed = 0;
// Build a large module for testing
std::string bigSrc;
for (int i = 0; i < 50; ++i) {
bigSrc += "def func_" + std::to_string(i) + "(a, b):\n";
bigSrc += " return a + b\n\n";
}
// ---------------------------------------------------------------
// Test 1: getAST without budget → full response (backward compat)
// ---------------------------------------------------------------
{
HeadlessEditorState state;
state.defaultLanguage = "python";
state.openBuffer("big.py", bigSrc, "python");
state.setAgentRole("s1", AgentRole::Refactor);
json resp = rpc(state, "s1", "getAST");
bool hasAst = resp.contains("result") &&
resp["result"].contains("ast");
bool noTruncated = !resp["result"].contains("truncated") ||
resp["result"].value("truncated", false) == false;
expect(hasAst && noTruncated,
"getAST without budget returns full response",
passed, failed);
}
// ---------------------------------------------------------------
// Test 2: getAST with small budget → truncated response
// ---------------------------------------------------------------
{
HeadlessEditorState state;
state.defaultLanguage = "python";
state.openBuffer("big.py", bigSrc, "python");
state.setAgentRole("s1", AgentRole::Refactor);
json resp = rpc(state, "s1", "getAST",
{{"compact", true}, {"budget", 500}});
bool isTruncated = resp.contains("result") &&
resp["result"].value("truncated", false) == true;
bool hasContinuation = resp.contains("result") &&
resp["result"].contains("continuation");
int size = (int)resp["result"].dump().size();
expect(isTruncated && hasContinuation && size <= 600,
"getAST budget=500 truncated (size=" +
std::to_string(size) + ")",
passed, failed);
}
// ---------------------------------------------------------------
// Test 3: Truncated response includes totalCount and returnedCount
// ---------------------------------------------------------------
{
HeadlessEditorState state;
state.defaultLanguage = "python";
state.openBuffer("big.py", bigSrc, "python");
state.setAgentRole("s1", AgentRole::Refactor);
json resp = rpc(state, "s1", "getAST",
{{"compact", true}, {"budget", 500}});
bool hasTotalCount = resp["result"].contains("totalCount");
bool hasReturnedCount = resp["result"].contains("returnedCount");
int total = resp["result"].value("totalCount", 0);
int returned = resp["result"].value("returnedCount", 0);
expect(hasTotalCount && hasReturnedCount &&
returned < total && returned > 0,
"Truncated has totalCount=" + std::to_string(total) +
" returnedCount=" + std::to_string(returned),
passed, failed);
}
// ---------------------------------------------------------------
// Test 4: Continuation token can be used for follow-up
// ---------------------------------------------------------------
{
HeadlessEditorState state;
state.defaultLanguage = "python";
state.openBuffer("big.py", bigSrc, "python");
state.setAgentRole("s1", AgentRole::Refactor);
// Get first page
json resp1 = rpc(state, "s1", "getAST",
{{"compact", true}, {"budget", 500}});
std::string cont = resp1["result"].value("continuation", "");
int returned1 = resp1["result"].value("returnedCount", 0);
// Use continuation token to get the full result and apply
json fullResult;
{
json compactNodes = toJsonCompactSummary(state.activeAST());
fullResult = {{"nodes", compactNodes},
{"nodeCount", (int)compactNodes.size()}};
}
json page2 = applyContinuation(fullResult, cont, 2000);
bool hasOffset = page2.contains("offset");
int offset = page2.value("offset", 0);
expect(!cont.empty() && hasOffset && offset == returned1,
"Continuation token resumes from offset=" +
std::to_string(offset),
passed, failed);
}
// ---------------------------------------------------------------
// Test 5: Budget on getAST full mode also truncates
// ---------------------------------------------------------------
{
HeadlessEditorState state;
state.defaultLanguage = "python";
state.openBuffer("big.py", bigSrc, "python");
state.setAgentRole("s1", AgentRole::Refactor);
json fullResp = rpc(state, "s1", "getAST");
int fullSize = (int)fullResp["result"].dump().size();
json budgetResp = rpc(state, "s1", "getAST",
{{"budget", 1000}});
bool truncated = budgetResp["result"].value("truncated", false);
// Full response is huge for 50 functions, budget should set truncated
expect(fullSize > 1000 && truncated,
"Full-mode getAST with budget sets truncated flag (full=" +
std::to_string(fullSize) + ")",
passed, failed);
}
// ---------------------------------------------------------------
// Test 6: getDiagnostics with budget truncates
// ---------------------------------------------------------------
{
HeadlessEditorState state;
state.defaultLanguage = "python";
state.openBuffer("diag.py",
"def cleanup(ptr):\n x = ptr\n", "python");
state.setAgentRole("s1", AgentRole::Refactor);
// Add multiple annotations to create multiple diagnostics
Module* ast = state.activeAST();
for (auto* c : ast->allChildren()) {
if (c->conceptType == "Function") {
c->addChild("annotations", new DeallocateAnnotation());
break;
}
}
// Without budget → normal response
json fullResp = rpc(state, "s1", "getDiagnostics");
int fullCount = fullResp["result"].value("count", 0);
// With very small budget
json budgetResp = rpc(state, "s1", "getDiagnostics",
{{"budget", 100}});
bool hasDiags = budgetResp.contains("result") &&
budgetResp["result"].contains("diagnostics");
expect(fullCount > 0 && hasDiags,
"getDiagnostics with budget returns diagnostics",
passed, failed);
}
// ---------------------------------------------------------------
// Test 7: applyBudget with no array returns truncated flag only
// ---------------------------------------------------------------
{
json simple = {{"message", "hello"}, {"value", 42}};
auto br = applyBudget(simple, 10);
expect(br.truncated,
"applyBudget on non-array result sets truncated=true",
passed, failed);
}
// ---------------------------------------------------------------
// Test 8: applyBudget with large budget returns full result
// ---------------------------------------------------------------
{
json result = {{"items", json::array({1, 2, 3})}};
auto br = applyBudget(result, 10000);
expect(!br.truncated,
"applyBudget with large budget returns full result",
passed, failed);
}
// ---------------------------------------------------------------
// Test 9: applyContinuation with invalid token returns error
// ---------------------------------------------------------------
{
json full = {{"items", json::array({1, 2, 3})}};
json result = applyContinuation(full, "invalid", 1000);
bool hasError = result.contains("error");
expect(hasError,
"applyContinuation with invalid token returns error",
passed, failed);
}
// ---------------------------------------------------------------
// Test 10: sortDiagnosticsByPriority puts errors first
// ---------------------------------------------------------------
{
json diags = json::array({
{{"code", "W001"}, {"severity", "warning"},
{"severityLevel", 2}},
{{"code", "E001"}, {"severity", "error"},
{"severityLevel", 1}},
{{"code", "I001"}, {"severity", "info"},
{"severityLevel", 3}}
});
sortDiagnosticsByPriority(diags);
bool sorted = diags[0].value("severity", "") == "error" &&
diags[1].value("severity", "") == "warning" &&
diags[2].value("severity", "") == "info";
expect(sorted,
"sortDiagnosticsByPriority: error → warning → info",
passed, failed);
}
// ---------------------------------------------------------------
// Test 11: Budget parameter zero treated as unlimited
// ---------------------------------------------------------------
{
HeadlessEditorState state;
state.defaultLanguage = "python";
state.openBuffer("zero.py", bigSrc, "python");
state.setAgentRole("s1", AgentRole::Refactor);
json resp = rpc(state, "s1", "getAST",
{{"compact", true}, {"budget", 0}});
bool noTruncated = !resp["result"].contains("truncated") ||
resp["result"].value("truncated", false) == false;
expect(noTruncated,
"budget=0 treated as unlimited (no truncation)",
passed, failed);
}
// ---------------------------------------------------------------
// Test 12: Multiple pages cover all data
// ---------------------------------------------------------------
{
// Test that continuation covers remaining items
json items = json::array();
for (int i = 0; i < 100; ++i)
items.push_back({{"id", i}, {"value", "item_" + std::to_string(i)}});
json full = {{"items", items}};
auto br = applyBudget(full, 200);
int firstReturned = 0;
if (br.truncated && br.result.contains("returnedCount"))
firstReturned = br.result.value("returnedCount", 0);
json page2 = applyContinuation(full, br.continuation, 10000);
int secondReturned = page2.value("returnedCount", 0);
bool coversAll = (firstReturned + secondReturned) == 100;
expect(br.truncated && coversAll,
"Two pages cover all 100 items (first=" +
std::to_string(firstReturned) + " second=" +
std::to_string(secondReturned) + ")",
passed, failed);
}
std::cout << "\n=== Step 254 Results: " << passed << " passed, "
<< failed << " failed ===\n";
return failed == 0 ? 0 : 1;
}

View File

@@ -278,3 +278,37 @@ complete flow: introduce error → diagnose → fix → verify clearing.
- Phase 9b complete: all 4 steps pass (44/44 tests across steps 250253)
- Full diagnostic pipeline validated: diagnose → fix → verify → delta
- 21 MCP tools, structured error codes (E01xx/E02xx/E03xx), delta streaming
---
## Phase 9c: Token-Efficient Agent Protocol
### Step 254: Response Budget System
**Status:** PASS (12/12 tests)
Optional `budget` parameter on query methods. If the response exceeds the
budget, it's truncated with `truncated: true` and a `continuation` token
for paginated follow-up. Arrays are shrunk via binary search; responses
sorted by priority (errors first).
**Files created:**
- `editor/src/ResponseBudget.h` — applyBudget (binary search truncation of
arrays), applyContinuation (resume from offset), sortDiagnosticsByPriority
- `editor/tests/step254_test.cpp` — 12 test cases: no-budget backward compat,
compact AST truncation, totalCount/returnedCount, continuation pagination,
full-mode truncated flag, diagnostics with budget, non-array truncation,
large budget passthrough, invalid continuation error, priority sorting,
budget=0 unlimited, multi-page coverage
**Files modified:**
- `editor/src/HeadlessEditorState.h` — include ResponseBudget.h
- `editor/src/HeadlessAgentRPCHandler.h` — getAST and getDiagnostics gain
optional budget parameter with truncation
- `editor/CMakeLists.txt` — step254_test target
**Key design decisions:**
- Budget in chars (not tokens) — simple, predictable, cheap to compute
- Binary search finds max array elements that fit within budget
- Continuation token format: "field:offset:total" (opaque to agent)
- Diagnostics sorted by severity before truncation (errors first)
- budget=0 or omitted → unlimited (backward compatible)