fix: keep tool_use/tool_result pairs atomic when trimming context

Author: jeremydw

.changeset/fix-context-trim-tool-pairing.md

@@ -0,0 +1,12 @@
+---
+'vscode-agent-platform-connector': patch
+---
+
+Fix a second cause of the Anthropic tool-pairing 400 ("tool_use ids were found
+without tool_result blocks" / "unexpected tool_use_id in tool_result blocks").
+`fitRequestToContext`'s oldest-message trimming (used to keep long
+conversations within the model's context window) could drop an assistant
+`tool_use` turn while keeping its paired `tool_result` reply, or vice versa,
+producing an invalid sequence upstream. Trimming now treats a `tool_use`
+assistant turn and its immediately-following `tool_result` turn as one atomic
+unit that is always dropped or kept together.

src/vertex.ts

@@ -1068,30 +1068,57 @@ export function fitRequestToContext(
   let total = estimateRequestTokens(req);
   if (total <= budget) return req;
 
-  const messages = [...req.messages];
+  const messages = req.messages;
   // Indices we must never drop: all system messages and the final message.
   const lastIdx = messages.length - 1;
   const protectedIdx = new Set<number>([lastIdx]);
   messages.forEach((m, i) => {
     if (m.role === 'system') protectedIdx.add(i);
   });
 
-  // Walk from oldest to newest, removing droppable messages until we fit.
+  // Group messages into atomic drop units: an assistant `tool_use` turn and
+  // its immediately-following `tool_result` reply must be dropped (or kept)
+  // together. Trimming just one half leaves the other as an orphan and
+  // upstream rejects the request with a 400 ("tool_use ids were found
+  // without tool_result blocks" / "unexpected tool_use_id in tool_result").
+  const units: Array<{
+    indices: number[];
+    tokens: number;
+    protected: boolean;
+  }> = [];
+  for (let i = 0; i < messages.length; ) {
+    const m = messages[i];
+    const pairsWithNext =
+      m.role === 'assistant' &&
+      !!m.toolCalls?.length &&
+      !!messages[i + 1]?.toolResults?.length;
+    const indices = pairsWithNext ? [i, i + 1] : [i];
+    const tokens = indices.reduce(
+      (sum, idx) => sum + estimateMessageTokens(messages[idx]),
+      0,
+    );
+    units.push({
+      indices,
+      tokens,
+      protected: indices.some((idx) => protectedIdx.has(idx)),
+    });
+    i += indices.length;
+  }
+
+  // Walk from oldest to newest, removing droppable units until we fit.
   let removed = 0;
-  for (let i = 0; i < messages.length && total > budget; i++) {
-    if (protectedIdx.has(i)) continue;
-    if (!messages[i]) continue;
-    total -= estimateMessageTokens(messages[i]);
-    // Tombstone; filtered out below to keep indices stable during the loop.
-    (messages as Array<NormMessage | null>)[i] = null;
-    removed++;
+  const dropped = new Set<number>();
+  for (const unit of units) {
+    if (total <= budget) break;
+    if (unit.protected) continue;
+    total -= unit.tokens;
+    for (const idx of unit.indices) dropped.add(idx);
+    removed += unit.indices.length;
   }
 
   if (removed === 0) return req;
 
-  const trimmed = (messages as Array<NormMessage | null>).filter(
-    (m): m is NormMessage => m !== null,
-  );
+  const trimmed = messages.filter((_, i) => !dropped.has(i));
   log(
     `context safeguard: trimmed ${removed} oldest message(s); ` +
       `~${total} est tokens <= ${budget} budget (limit ${limit})`,

test/context-fit.test.ts

@@ -64,4 +64,40 @@ describe('fitRequestToContext', () => {
     });
     expect(large).toBeGreaterThan(small);
   });
+
+  it('drops a tool_use/tool_result pair together, never splitting them', () => {
+    // Each old turn (assistant tool_use + user tool_result) is large enough
+    // that dropping just one of the two wouldn't be enough to hit budget,
+    // forcing the trimmer to consider the pair as a unit.
+    const big = 'x'.repeat(4000);
+    const req: NormRequest = {
+      messages: [
+        {role: 'system', text: 'system prompt'},
+        {
+          role: 'assistant',
+          toolCalls: [{id: 'call_1', name: 'do_thing', input: {a: big}}],
+        },
+        {role: 'user', toolResults: [{callId: 'call_1', content: big}]},
+        userMessage('latest ' + big),
+      ],
+    };
+
+    const fitted = fitRequestToContext(model, req);
+    expect(fitted).not.toBe(req);
+
+    const hasOrphanToolUse = fitted.messages.some(
+      (m, i) =>
+        m.role === 'assistant' &&
+        m.toolCalls?.length &&
+        !fitted.messages[i + 1]?.toolResults?.length,
+    );
+    const hasOrphanToolResult = fitted.messages.some((m, i) => {
+      if (!m.toolResults?.length) return false;
+      const prev = fitted.messages[i - 1];
+      const ids = new Set((prev?.toolCalls ?? []).map((t) => t.id));
+      return m.toolResults.some((r) => !ids.has(r.callId));
+    });
+    expect(hasOrphanToolUse).toBe(false);
+    expect(hasOrphanToolResult).toBe(false);
+  });
 });