Skip to content

Commit 58bcd78

Browse files
committed
merge: raw fallback refuses on an independent byte proxy (a guard that cannot measure must not serve) + narrow LKG tool-boundary correctness
2 parents 10a2ae0 + dcaa1fb commit 58bcd78

4 files changed

Lines changed: 213 additions & 16 deletions

File tree

‎packages/plugin/src/hooks/magic-context/lkg-replay.ts‎

Lines changed: 30 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -363,9 +363,33 @@ function partIsAnthropicThinking(part: unknown): boolean {
363363
return type === "thinking" || type === "reasoning" || type === "redacted_thinking";
364364
}
365365

366+
function partEndsAnthropicAssistantRun(part: unknown): boolean {
367+
if (!part || typeof part !== "object") return false;
368+
const value = part as Record<string, unknown>;
369+
if (
370+
value.type !== "tool" ||
371+
value.providerExecuted === true ||
372+
!value.state ||
373+
typeof value.state !== "object"
374+
) {
375+
return false;
376+
}
377+
const status = (value.state as Record<string, unknown>).status;
378+
return status === "completed" || status === "error";
379+
}
380+
381+
function partIsOpenCodeStepMetadata(part: unknown): boolean {
382+
if (!part || typeof part !== "object") return false;
383+
const type = (part as Record<string, unknown>).type;
384+
return type === "step-start" || type === "step-finish";
385+
}
386+
366387
/**
367-
* The Anthropic adapter merges adjacent assistant messages before sending them.
368-
* A merged run may contain only one leading thinking block; moving or retaining a
388+
* The Anthropic adapter merges adjacent assistant content before sending it. A
389+
* completed non-provider-executed tool result materializes as user content and
390+
* starts a new assistant run, while OpenCode's step markers do not materialize
391+
* on the provider wire.
392+
* Each resulting assistant run may contain only one leading thinking block; a
369393
* later signed block would invalidate its provider signature, so recovery declines
370394
* the entire replay instead of attempting a rewrite.
371395
*/
@@ -383,7 +407,10 @@ export function validateAnthropicReasoningRuns(messages: MessageLike[]): boolean
383407
if (partIsAnthropicThinking(part)) {
384408
thinkingBlocks += 1;
385409
if (thinkingBlocks > 1 || sawOtherContent) return false;
386-
} else {
410+
} else if (partEndsAnthropicAssistantRun(part)) {
411+
thinkingBlocks = 0;
412+
sawOtherContent = false;
413+
} else if (!partIsOpenCodeStepMetadata(part)) {
387414
sawOtherContent = true;
388415
}
389416
}

‎packages/plugin/src/hooks/magic-context/lkg-transform-replay.test.ts‎

Lines changed: 64 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,7 @@ import {
66
captureLkgSlot,
77
projectLkgEntry,
88
replayLkg,
9+
validateAnthropicReasoningRuns,
910
validateLkgEntry,
1011
validateLkgSeam,
1112
} from "./lkg-replay";
@@ -298,6 +299,69 @@ describe("LKG transform replay", () => {
298299
expect(getSlot("anthropic-invalid")).toBeUndefined();
299300
});
300301

302+
test("serves a new Anthropic thinking run after a completed tool result", () => {
303+
resetLkgSlotsForTest();
304+
const prefix = assistant("a-prefix", 1, [
305+
{ type: "step-start" },
306+
{ type: "thinking", thinking: "first signed trace", signature: "sig-a" },
307+
{ type: "text", text: "calling a tool" },
308+
{
309+
type: "tool",
310+
callID: "call-1",
311+
tool: "read",
312+
providerExecuted: false,
313+
state: { status: "completed", input: {}, output: "result" },
314+
},
315+
{ type: "step-finish" },
316+
]);
317+
captureSlot("anthropic-tool-boundary", {
318+
jsonPrefix: JSON.stringify([prefix]),
319+
inputIdSeq: ["u-anchor"],
320+
inputContentDigests: [
321+
lkgContentDigest(
322+
user("u-anchor", 2, { providerID: "anthropic", modelID: "claude-test" }),
323+
)!,
324+
],
325+
lastInputMessageId: "u-anchor",
326+
modelKey: "anthropic/claude-test",
327+
providerKey: "anthropic",
328+
capturedAt: 1,
329+
});
330+
const current = [
331+
user("u-anchor", 2, { providerID: "anthropic", modelID: "claude-test" }),
332+
assistant("a-tail", 3, [
333+
{ type: "thinking", thinking: "second signed trace", signature: "sig-b" },
334+
{ type: "text", text: "second response" },
335+
]),
336+
];
337+
338+
const replay = replayLkg({
339+
sessionId: "anthropic-tool-boundary",
340+
messages: current,
341+
modelKey: "anthropic/claude-test",
342+
providerKey: "anthropic",
343+
});
344+
expect(replay).toEqual({ ok: true, messages: [prefix, current[1]] });
345+
});
346+
347+
test("declines a new thinking run after a provider-executed tool", () => {
348+
const first = assistant("a-prefix", 1, [
349+
{ type: "thinking", thinking: "first signed trace", signature: "sig-a" },
350+
{
351+
type: "tool",
352+
callID: "call-1",
353+
tool: "provider-tool",
354+
providerExecuted: true,
355+
state: { status: "completed", input: {}, output: "result" },
356+
},
357+
]);
358+
const second = assistant("a-tail", 2, [
359+
{ type: "thinking", thinking: "second signed trace", signature: "sig-b" },
360+
]);
361+
362+
expect(validateAnthropicReasoningRuns([first, second])).toBe(false);
363+
});
364+
301365
test("serves an Anthropic replay with one leading thinking block in an assistant run", () => {
302366
resetLkgSlotsForTest();
303367
captureSlot("anthropic-valid", {

‎packages/plugin/src/hooks/magic-context/rust-mode-transform.test.ts‎

Lines changed: 83 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1864,6 +1864,89 @@ describe("Rust mode authority adapter", () => {
18641864
expect(output.messages).toEqual([]);
18651865
});
18661866

1867+
it("refuses a large raw fallback when the token estimator is unavailable", async () => {
1868+
const sessionId = `rust-raw-estimator-unavailable-${Date.now()}`;
1869+
sessions.push(sessionId);
1870+
const db = makeDb();
1871+
installRawProvider(sessionId);
1872+
recordDetectedContextLimit(db, sessionId, 1_000, "test-provider/test-model");
1873+
const moduleClient: RustModeModuleClient = {
1874+
call: async ({ method }) => {
1875+
if (method === "transform") throw new Error("client closed");
1876+
return { ok: true };
1877+
},
1878+
};
1879+
const deps = makeDeps(db, moduleClient);
1880+
let estimatorCalls = 0;
1881+
const transform = createRustModeTransform(deps, {
1882+
moduleClient,
1883+
rawFallbackEstimatorForTests: () => {
1884+
estimatorCalls += 1;
1885+
throw new Error("tokenizer unavailable");
1886+
},
1887+
});
1888+
const input = [
1889+
{
1890+
info: {
1891+
id: "m1",
1892+
role: "user",
1893+
sessionID: sessionId,
1894+
model: { providerID: "test-provider", modelID: "test-model" },
1895+
},
1896+
parts: [{ type: "text", text: "x".repeat(5_000) }],
1897+
},
1898+
] as MessageLike[];
1899+
const output = { messages: [] as unknown[] };
1900+
1901+
await expect(
1902+
transform.run(sessionId, input, output, makeMeta(db, sessionId)),
1903+
).rejects.toBeInstanceOf(RawFallbackContextLimitError);
1904+
expect(estimatorCalls).toBe(1);
1905+
expect(output.messages).toEqual([]);
1906+
});
1907+
1908+
it("refuses a byte-large raw fallback when the token estimate is materially low", async () => {
1909+
const sessionId = `rust-raw-estimator-low-${Date.now()}`;
1910+
sessions.push(sessionId);
1911+
const db = makeDb();
1912+
installRawProvider(sessionId);
1913+
recordDetectedContextLimit(db, sessionId, 1_000, "test-provider/test-model");
1914+
const moduleClient: RustModeModuleClient = {
1915+
call: async ({ method }) => {
1916+
if (method === "transform") throw new Error("client closed");
1917+
return { ok: true };
1918+
},
1919+
};
1920+
const deps = makeDeps(db, moduleClient);
1921+
const transform = createRustModeTransform(deps, {
1922+
moduleClient,
1923+
rawFallbackEstimatorForTests: () => ({
1924+
tokens: 900,
1925+
trusted: true,
1926+
messageTokens: { conversation: 900, toolCall: 0 },
1927+
systemTokens: 0,
1928+
toolDefinitionTokens: 0,
1929+
}),
1930+
});
1931+
const input = [
1932+
{
1933+
info: {
1934+
id: "m1",
1935+
role: "user",
1936+
sessionID: sessionId,
1937+
model: { providerID: "test-provider", modelID: "test-model" },
1938+
},
1939+
parts: [{ type: "text", text: "x".repeat(5_000) }],
1940+
},
1941+
] as MessageLike[];
1942+
const output = { messages: [] as unknown[] };
1943+
1944+
await expect(
1945+
transform.run(sessionId, input, output, makeMeta(db, sessionId)),
1946+
).rejects.toBeInstanceOf(RawFallbackContextLimitError);
1947+
expect(output.messages).toEqual([]);
1948+
});
1949+
18671950
it("preserves raw fail-open when the estimate fits the known context limit", async () => {
18681951
const sessionId = `rust-raw-fits-${Date.now()}`;
18691952
sessions.push(sessionId);

‎packages/plugin/src/hooks/magic-context/rust-mode-transform.ts‎

Lines changed: 36 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -102,6 +102,17 @@ export const RUST_PARK_RETRY_INTERVAL = 5;
102102
export const RUST_EMERGENCY_WALL_PCT = 95;
103103
export const RUST_PARK_PROBE_PRESSURE_BYPASS_PCT = 90;
104104
const RUST_SEND_TIMEOUT_MS = 15_000;
105+
const RAW_FALLBACK_BYTES_PER_CONTEXT_TOKEN = 4;
106+
107+
function rawFallbackSerializedBytes(messages: readonly MessageLike[]): number | null {
108+
try {
109+
const serialized = JSON.stringify(messages);
110+
return typeof serialized === "string" ? Buffer.byteLength(serialized) : null;
111+
} catch {
112+
// Serialization is itself required before these messages can reach a provider.
113+
return null;
114+
}
115+
}
105116

106117
export interface RustModeModuleClient extends ModuleStateSyncClient {
107118
authorityStatus?(args: {
@@ -212,6 +223,8 @@ export interface RustModeTransformOptions {
212223
allowAuthorityProtocolBypassForTests?: boolean;
213224
/** Override only for deterministic capture scheduling in tests. */
214225
scheduleLkgCapture?: (capture: () => void) => void;
226+
/** Override only to exercise raw-fallback estimator failures in tests. */
227+
rawFallbackEstimatorForTests?: typeof estimateFinalWireInputTokens;
215228
}
216229

217230
function isRecord(value: unknown): value is Record<string, unknown> {
@@ -1110,6 +1123,8 @@ export function createRustModeTransform(
11101123
const wireCaches = new Map<string, RustWireCache>();
11111124
const scheduleLkgCapture =
11121125
options.scheduleLkgCapture ?? ((capture: () => void) => setImmediate(capture));
1126+
const rawFallbackEstimator =
1127+
options.rawFallbackEstimatorForTests ?? estimateFinalWireInputTokens;
11131128
const timeoutMs = Math.max(1, options.moduleTimeoutMs ?? RUST_SEND_TIMEOUT_MS);
11141129

11151130
const logStage = (
@@ -1421,26 +1436,34 @@ export function createRustModeTransform(
14211436
? overflowState.detectedContextLimit
14221437
: undefined);
14231438
if (contextLimit !== undefined) {
1439+
let estimate: ReturnType<typeof estimateFinalWireInputTokens> | undefined;
14241440
try {
1425-
const estimate = estimateFinalWireInputTokens({
1441+
estimate = rawFallbackEstimator({
14261442
messages,
14271443
systemPromptTokens: sessionMeta.systemPromptTokens,
14281444
providerID: model?.providerID,
14291445
modelID: model?.modelID,
14301446
agentName: deps.getNotificationParams?.(sessionId)?.agent,
14311447
});
1432-
if (estimate.tokens > contextLimit) {
1433-
sessionLog(
1434-
sessionId,
1435-
`raw_fallback_over_context_limit estimated=${estimate.tokens} limit=${contextLimit}`,
1436-
);
1437-
throw new RawFallbackContextLimitError(estimate.tokens, contextLimit, {
1438-
cause,
1439-
});
1440-
}
1441-
} catch (error) {
1442-
if (error instanceof RawFallbackContextLimitError) throw error;
1443-
// An unavailable estimate leaves the ordinary raw fail-open contract intact.
1448+
} catch {
1449+
// The byte proxy below remains available when tokenization does not.
1450+
}
1451+
const rawBytes = rawFallbackSerializedBytes(messages);
1452+
// The local tokenizer is telemetry-grade and can materially undercount a new
1453+
// provider tokenizer. Four serialized bytes per context token is an independent,
1454+
// conservative risk budget for a raw full-history fallback.
1455+
const proxyTokens =
1456+
rawBytes === null
1457+
? contextLimit + 1
1458+
: Math.ceil(rawBytes / RAW_FALLBACK_BYTES_PER_CONTEXT_TOKEN);
1459+
const refusalTokens = Math.max(estimate?.tokens ?? 0, proxyTokens);
1460+
if (refusalTokens > contextLimit) {
1461+
sessionLog(
1462+
sessionId,
1463+
`raw_fallback_over_context_limit estimated=${estimate?.tokens ?? "unavailable"} ` +
1464+
`proxy_bytes=${rawBytes ?? "unavailable"} proxy_tokens=${proxyTokens} limit=${contextLimit}`,
1465+
);
1466+
throw new RawFallbackContextLimitError(refusalTokens, contextLimit, { cause });
14441467
}
14451468
}
14461469
replaceMessagesInPlace(output, messages);

0 commit comments

Comments
 (0)