feat: use session-reported Claude CLI context windows
This commit is contained in:
@@ -0,0 +1,71 @@
|
||||
import { randomUUID } from "node:crypto";
|
||||
|
||||
const INITIALIZE_TIMEOUT_MS = 5_000;
|
||||
|
||||
/** A best-effort, local-only CLI control query. Wire contract matches the
|
||||
* official SDK0.3.259 source; this does not import or execute the SDK.
|
||||
* Summary deliberately avoids the default full mode's token-count requests.
|
||||
* CLI maxTokens is an effective/auto-compaction budget, not a nominal model max. */
|
||||
export class ContextWindowProbe {
|
||||
private readonly initializeId = randomUUID();
|
||||
private summaryId: string | undefined;
|
||||
private initialized = false;
|
||||
private timer: NodeJS.Timeout | undefined;
|
||||
private ready = false;
|
||||
private closed = false;
|
||||
|
||||
constructor(
|
||||
private readonly write: (line: string) => void,
|
||||
private readonly sendPrompt: () => void,
|
||||
private readonly onWindow: (tokens: number) => void,
|
||||
) {}
|
||||
|
||||
start(): void {
|
||||
this.timer = setTimeout(() => this.proceed(false), INITIALIZE_TIMEOUT_MS);
|
||||
this.timer.unref();
|
||||
this.request(this.initializeId, { subtype: "initialize" });
|
||||
}
|
||||
|
||||
/** Consume control replies separately from model/tool output, including
|
||||
* late/foreign replies while a native question tool is already waiting. */
|
||||
handle(value: unknown): boolean {
|
||||
if (!value || typeof value !== "object" || !("type" in value) || value.type !== "control_response") return false;
|
||||
if (this.closed || !("response" in value) || !value.response || typeof value.response !== "object") return true;
|
||||
const response = value.response as Record<string, unknown>;
|
||||
if (response.request_id === this.initializeId) this.proceed(response.subtype === "success");
|
||||
else if (this.summaryId && response.request_id === this.summaryId) {
|
||||
this.summaryId = undefined;
|
||||
if (response.subtype !== "success") return true;
|
||||
const body = response.response as Record<string, unknown> | undefined;
|
||||
const tokens = body?.maxTokens;
|
||||
if (typeof tokens === "number" && Number.isSafeInteger(tokens) && tokens > 0) this.onWindow(tokens);
|
||||
}
|
||||
return true;
|
||||
}
|
||||
|
||||
/** Retry when the CLI has begun a model response: pre-first-turn metadata
|
||||
* is not guaranteed by the protocol. Refresh also follows later responses
|
||||
* within the same held process, without opening a second process. */
|
||||
refresh(): void {
|
||||
if (!this.initialized || this.closed) return;
|
||||
this.summaryId = randomUUID();
|
||||
this.request(this.summaryId, { subtype: "get_context_usage", detail: "summary" });
|
||||
}
|
||||
|
||||
close(): void { this.closed = true; clearTimeout(this.timer); }
|
||||
|
||||
private proceed(query: boolean): void {
|
||||
if (this.ready || this.closed) return;
|
||||
this.ready = true;
|
||||
clearTimeout(this.timer);
|
||||
this.initialized = query;
|
||||
if (query) this.refresh();
|
||||
// Never wait for metadata to run the actual user request. No support,
|
||||
// no response, or malformed metadata simply leaves the window unknown.
|
||||
this.sendPrompt();
|
||||
}
|
||||
|
||||
private request(id: string, request: Record<string, string>): void {
|
||||
this.write(JSON.stringify({ type: "control_request", request_id: id, request }));
|
||||
}
|
||||
}
|
||||
@@ -24,9 +24,11 @@ import type { AnthropicEffort } from "@earendil-works/pi-ai/api/anthropic-messag
|
||||
import { parseStreamingJson, parseToolCallArguments } from "@earendil-works/pi-ai/utils/json-parse";
|
||||
|
||||
import { externalCliEnv } from "../../core/auth-env.ts";
|
||||
import type { LhMessage } from "../../core/native-messages.ts";
|
||||
import { ClaudeLaunchError, createClaudeLauncher } from "./launch.ts";
|
||||
import { parseCliLine, type CliBlock, type CliEvent, type CliUsage } from "./cli-events.ts";
|
||||
import { ClaudeCodeError } from "./constants.ts";
|
||||
import { ContextWindowProbe } from "./context-window.ts";
|
||||
import { addLiveProcess, getLiveProcess, removeLiveProcess, type LiveProcess } from "./live-processes.ts";
|
||||
import { configurationForReuse, conversationForReuse } from "./reuse-context.ts";
|
||||
import { BRIDGE_SERVER_NAME, createToolBridge, toolNameFromMcp, type PendingToolCall, type ToolBridge } from "./tool-bridge.ts";
|
||||
@@ -137,7 +139,7 @@ function getAssistantErrorMessage(error: string, text: string): string {
|
||||
}
|
||||
|
||||
class ClaudeTurn {
|
||||
private readonly output: AssistantMessage;
|
||||
private readonly output: AssistantMessage & Pick<LhMessage, "cliContextWindow">;
|
||||
private readonly contentIndexes = new Map<number, number>();
|
||||
private readonly partialJson = new Map<number, string>();
|
||||
private readonly toolBlocks = new Map<number, ToolCall>();
|
||||
@@ -233,6 +235,10 @@ class ClaudeTurn {
|
||||
|
||||
getMessage(): AssistantMessage { return this.output; }
|
||||
|
||||
setContextWindow(tokens: number | undefined): void {
|
||||
if (!this.finished && tokens !== undefined) this.output.cliContextWindow = tokens;
|
||||
}
|
||||
|
||||
admit(calls: ToolCall[]): void {
|
||||
this.start();
|
||||
for (const call of calls) {
|
||||
@@ -398,6 +404,8 @@ class ClaudeProcess implements LiveProcess {
|
||||
private admitted: PendingToolCall[] = [];
|
||||
private readonly cancelledIds = new Set<string>();
|
||||
private pendingError: string | undefined;
|
||||
private contextWindow: number | undefined;
|
||||
private contextProbe: ContextWindowProbe | undefined;
|
||||
|
||||
constructor(
|
||||
private readonly child: ChildProcessWithoutNullStreams,
|
||||
@@ -420,12 +428,17 @@ class ClaudeProcess implements LiveProcess {
|
||||
this.child.stderr.setEncoding("utf8");
|
||||
this.child.stderr.on("data", (chunk: string) => { this.stderr = (this.stderr + chunk).slice(-65_536); });
|
||||
this.child.stdin.on("error", (error) => console.error("claude CLI stdin closed early:", error));
|
||||
this.child.stdin.end(`${stdinLine}\n`);
|
||||
this.contextProbe = new ContextWindowProbe(
|
||||
(line) => this.child.stdin.write(`${line}\n`),
|
||||
() => { if (!this.stopping) this.child.stdin.write(`${stdinLine}\n`); },
|
||||
(tokens) => { this.contextWindow = tokens; this.turn.setContextWindow(tokens); },
|
||||
);
|
||||
createInterface({ input: this.child.stdout }).on("line", (line) => {
|
||||
try { this.handleLine(line); }
|
||||
catch { this.fail("claude CLI emitted malformed stream output"); }
|
||||
});
|
||||
this.setSignal(signal);
|
||||
if (!this.stopping) this.contextProbe.start();
|
||||
}
|
||||
|
||||
canContinue(owner: symbol, sessionId: string | undefined, configuration: unknown, messages: Message[], results: ToolResultMessage[]): boolean {
|
||||
@@ -445,6 +458,7 @@ class ClaudeProcess implements LiveProcess {
|
||||
this.admitted = [];
|
||||
this.conversation = structuredClone(conversationForReuse(messages));
|
||||
this.turn = turn;
|
||||
this.turn.setContextWindow(this.contextWindow);
|
||||
this.sourceComplete = this.sourceCalls.length > 0;
|
||||
this.setSignal(signal);
|
||||
if (this.stopping) return;
|
||||
@@ -459,6 +473,7 @@ class ClaudeProcess implements LiveProcess {
|
||||
removeLiveProcess(this);
|
||||
if (this.stopping || this.closed) return;
|
||||
this.stopping = true;
|
||||
this.contextProbe?.close();
|
||||
this.child.kill();
|
||||
this.killTimer = setTimeout(() => this.child.kill("SIGKILL"), 1000);
|
||||
this.killTimer.unref();
|
||||
@@ -526,7 +541,8 @@ class ClaudeProcess implements LiveProcess {
|
||||
}
|
||||
|
||||
private handleLine(line: string): void {
|
||||
if (this.stopping) return;
|
||||
if (this.stopping || !line.trim()) return;
|
||||
if (this.contextProbe?.handle(JSON.parse(line))) return;
|
||||
const event = parseCliLine(line);
|
||||
if (!event) return;
|
||||
if (event.type === "toolResults") {
|
||||
@@ -545,6 +561,7 @@ class ClaudeProcess implements LiveProcess {
|
||||
this.fail("claude CLI moved on without admitting its harness tool calls");
|
||||
return;
|
||||
}
|
||||
if (event.type === "messageStart") this.contextProbe?.refresh();
|
||||
this.turn.handle(event);
|
||||
if (this.turn.finished) { this.kill(); return; }
|
||||
if (event.type === "messageStop" && this.turn.waitingForAdmission) {
|
||||
@@ -562,6 +579,7 @@ class ClaudeProcess implements LiveProcess {
|
||||
private handleClose(message: string): void {
|
||||
if (this.closed) return;
|
||||
this.closed = true;
|
||||
this.contextProbe?.close();
|
||||
clearTimeout(this.killTimer);
|
||||
removeLiveProcess(this);
|
||||
if (this.pendingError) addLiveProcess(this);
|
||||
|
||||
Reference in New Issue
Block a user