repositories / pi-ext
pi-ext
bugabingas pi extensions
owned by admin
extensions/klaus/src/coordinator.ts
Rawimport { join } from "node:path";
import {
type Api,
type AssistantMessage,
type AssistantMessageEventStream,
calculateCost,
createAssistantMessageEventStream,
getCurrentSystemPrompt,
getCurrentTools,
type Model,
type SimpleStreamOptions,
type TranscriptContext,
} from "@earendil-works/pi-ai";
import type { KlausQueryHandle } from "./agent-sdk.js";
import {
FileSessionStore,
KlausCacheCorruptionError,
type KlausSessionStore,
loadCheckpoints,
MemorySessionStore,
StagingSessionStore,
saveCheckpoint,
} from "./cache.js";
import { dbg, diagnosticKind } from "./debug.js";
import {
contextFingerprint,
createAssistant,
isKlausModelId,
type KlausContentEvent,
type KlausContext,
type KlausQueryCallbacks,
type KlausRequest,
klausSelector,
modelFacingContent,
requiredConstrainedTool,
SDK_VERSION,
toKlausTools,
} from "./protocol.js";
import { replayImages, replayPrompt } from "./replay.js";
import { rewriteSubscriptionPrompt } from "./subscription-prompt.js";
import type { PiToolResult } from "./tool-bridge.js";
interface ActiveQuery {
id: symbol;
sessionId?: string;
handle?: KlausQueryHandle;
sink?: PiSink;
toolCallIds: Set<string>;
tools: KlausRequest["tools"];
boundaryFingerprint?: string;
boundaryMessageCount?: number;
closed: boolean;
closePromise?: Promise<void>;
timer?: ReturnType<typeof setTimeout>;
timeoutMs?: number;
idleLease?: ReturnType<typeof setTimeout>;
/** Set while this child is idle between turns and may take the next one. */
reuseLease?: ReturnType<typeof setTimeout>;
sdkSessionId?: string;
position?: string;
messageCount?: number;
/** Lineage of the turns this child already holds, computed without tools so
* a tool-list change can be answered with a swap instead of a replay. */
lineageFingerprint?: string;
abort?: () => void;
signal?: AbortSignal;
secondary: boolean;
committable: boolean;
modelId: string;
cwd: string;
thinking: SimpleStreamOptions["reasoning"];
thinkingBudget?: number;
maxTokens?: number;
systemPrompt: string;
headers: SimpleStreamOptions["headers"];
env: SimpleStreamOptions["env"];
metadata?: Record<string, string>;
emittedContent: boolean;
context: KlausContext;
store: KlausSessionStore;
staging: StagingSessionStore;
}
interface CompletedQuery {
sdkSessionId: string;
fingerprint: string;
position?: string;
piLeafId?: string;
messageCount: number;
protocol: 1;
store: KlausSessionStore;
}
interface SessionScope {
id: string;
persisted: boolean;
dir: string;
}
/** How long a finished child stays alive waiting for the next turn. Tests may
* shorten it, because proving expiry must not cost two minutes per run. */
function reuseLeaseMs(): number {
const override = process.env.VITEST
? Number(process.env.KLAUS_REUSE_LEASE_MS)
: Number.NaN;
return Number.isSafeInteger(override) && override > 0 ? override : 2 * 60_000;
}
type ResolveOAuth = (modelId: string) => Promise<string>;
type CurrentCwd = () => string;
type CurrentSession = () => SessionScope | undefined;
type Warn = (message: string) => void;
class PiSink {
readonly stream = createAssistantMessageEventStream();
readonly output: AssistantMessage;
private indices = new Map<number, number>();
private toolJson = new Map<number, string>();
private finished = false;
constructor(private readonly model: Model<Api>) {
dbg?.("coordinator.sink.create");
this.output = createAssistant(model);
this.stream.push({ type: "start", partial: this.output });
}
content(event: KlausContentEvent): void {
dbg?.("coordinator.sink.content", {
index: event.index,
length:
"delta" in event
? event.delta.length
: "text" in event
? event.text.length
: "thinking" in event
? event.thinking.length
: undefined,
});
if (this.finished) return;
switch (event.type) {
case "text-start": {
const index = this.output.content.length;
this.indices.set(event.index, index);
this.output.content.push({ type: "text", text: "" });
this.stream.push({
type: "text_start",
contentIndex: index,
partial: this.output,
});
break;
}
case "text-delta": {
const index = this.indices.get(event.index);
const block =
index === undefined ? undefined : this.output.content[index];
if (index === undefined || block?.type !== "text") return;
block.text += event.delta;
this.stream.push({
type: "text_delta",
contentIndex: index,
delta: event.delta,
partial: this.output,
});
break;
}
case "text-end": {
const index = this.indices.get(event.index);
if (index === undefined) return;
this.stream.push({
type: "text_end",
contentIndex: index,
content: event.text,
partial: this.output,
});
break;
}
case "thinking-start": {
const index = this.output.content.length;
this.indices.set(event.index, index);
this.output.content.push({ type: "thinking", thinking: "" });
this.stream.push({
type: "thinking_start",
contentIndex: index,
partial: this.output,
});
break;
}
case "thinking-delta": {
const index = this.indices.get(event.index);
const block =
index === undefined ? undefined : this.output.content[index];
if (index === undefined || block?.type !== "thinking") return;
block.thinking += event.delta;
this.stream.push({
type: "thinking_delta",
contentIndex: index,
delta: event.delta,
partial: this.output,
});
break;
}
case "thinking-end": {
const index = this.indices.get(event.index);
const block =
index === undefined ? undefined : this.output.content[index];
if (index === undefined || block?.type !== "thinking") return;
block.thinkingSignature = event.signature;
block.redacted = event.redacted;
this.stream.push({
type: "thinking_end",
contentIndex: index,
content: event.thinking,
partial: this.output,
});
break;
}
case "tool-start": {
const index = this.output.content.length;
this.indices.set(event.index, index);
this.toolJson.set(event.index, "");
this.output.content.push({
type: "toolCall",
id: event.id,
name: event.name,
arguments: {},
});
this.stream.push({
type: "toolcall_start",
contentIndex: index,
partial: this.output,
});
break;
}
case "tool-delta": {
const index = this.indices.get(event.index);
if (index === undefined) return;
this.toolJson.set(
event.index,
`${this.toolJson.get(event.index) ?? ""}${event.delta}`,
);
this.stream.push({
type: "toolcall_delta",
contentIndex: index,
delta: event.delta,
partial: this.output,
});
break;
}
case "tool-end": {
const index = this.indices.get(event.index);
const block =
index === undefined ? undefined : this.output.content[index];
if (index === undefined || block?.type !== "toolCall") return;
block.arguments = event.arguments;
this.stream.push({
type: "toolcall_end",
contentIndex: index,
toolCall: block,
partial: this.output,
});
break;
}
}
}
toolUse(
usage: {
input: number;
output: number;
cacheRead: number;
cacheWrite: number;
},
responseId?: string,
responseModel?: string,
): void {
dbg?.("coordinator.sink.toolUse", { finished: this.finished });
if (this.finished) return;
this.finished = true;
this.applyUsage(usage, responseId, responseModel);
this.output.stopReason = "toolUse";
this.stream.push({ type: "done", reason: "toolUse", message: this.output });
this.stream.end();
}
done(
usage: {
input: number;
output: number;
cacheRead: number;
cacheWrite: number;
},
responseId?: string,
responseModel?: string,
stopReason: "stop" | "length" = "stop",
): void {
dbg?.("coordinator.sink.done", { finished: this.finished });
if (this.finished) return;
this.finished = true;
this.applyUsage(usage, responseId, responseModel);
this.output.stopReason = stopReason;
this.stream.push({
type: "done",
reason: stopReason,
message: this.output,
});
this.stream.end();
}
private applyUsage(
usage: {
input: number;
output: number;
cacheRead: number;
cacheWrite: number;
},
responseId?: string,
responseModel?: string,
): void {
Object.assign(this.output.usage, usage);
this.output.usage.totalTokens =
usage.input + usage.output + usage.cacheRead + usage.cacheWrite;
const fallbackCost =
this.model.compat &&
"allowedFallbackModels" in this.model.compat &&
responseModel !== this.model.id
? this.model.compat.allowedFallbackModels?.find(
(fallback) =>
fallback.provider === "anthropic" &&
fallback.model === responseModel,
)?.cost
: undefined;
calculateCost(
fallbackCost && responseModel
? { ...this.model, id: responseModel, cost: fallbackCost }
: this.model,
this.output.usage,
);
this.output.responseId = responseId;
this.output.responseModel = responseModel;
}
error(error: unknown, aborted = false): void {
dbg?.("coordinator.sink.error", {
finished: this.finished,
aborted,
kind: diagnosticKind(error),
});
if (this.finished) return;
this.finished = true;
this.output.stopReason = aborted ? "aborted" : "error";
this.output.errorMessage =
error instanceof Error ? error.message : String(error);
this.stream.push({
type: "error",
reason: this.output.stopReason,
error: this.output,
});
this.stream.end();
}
}
export class QueryCoordinator {
private readonly bySession = new Map<string, Set<ActiveQuery>>();
private readonly anonymous = new Set<ActiveQuery>();
private readonly completed = new Map<string, CompletedQuery>();
private readonly memoryStores = new Map<string, MemorySessionStore>();
private readonly diskStores = new Map<string, FileSessionStore>();
private readonly pendingCheckpoints = new Map<string, CompletedQuery>();
private cacheReady: Promise<void> = Promise.resolve();
constructor(
private readonly resolveOAuth: ResolveOAuth,
private readonly currentCwd: CurrentCwd,
private readonly currentSession: CurrentSession = () => undefined,
private readonly warn: Warn = () => undefined,
) {
dbg?.("coordinator.create");
}
stream(
model: Model<Api>,
context: TranscriptContext,
options: SimpleStreamOptions = {},
): AssistantMessageEventStream {
const currentContext: KlausContext = {
messages: context.messages.filter(
(message): message is KlausContext["messages"][number] =>
message.role !== "system",
),
systemPrompt: getCurrentSystemPrompt(context.messages),
tools: getCurrentTools(context.messages),
};
dbg?.("coordinator.stream", {
messageCount: currentContext.messages.length,
toolCount: currentContext.tools?.length ?? 0,
});
const sink = new PiSink(model);
void this.invoke(sink, model, currentContext, options);
return sink.stream;
}
preparePrimaryCache(scope: SessionScope): Promise<void> {
dbg?.("coordinator.preparePrimaryCache", { persisted: scope.persisted });
if (!scope.persisted) {
this.cacheReady = Promise.resolve();
return this.cacheReady;
}
const store = new FileSessionStore(
join(scope.dir, ".klaus", SDK_VERSION, "sessions"),
);
this.diskStores.set(scope.id, store);
this.cacheReady = store.prepare();
return this.cacheReady;
}
async persistCheckpoint(sessionId: string, piLeafId: string): Promise<void> {
dbg?.("coordinator.persistCheckpoint.start");
const checkpoint = this.pendingCheckpoints.get(sessionId);
const root = this.checkpointRoot(sessionId);
if (!checkpoint || !root) {
dbg?.("coordinator.persistCheckpoint.skip", {
hasCheckpoint: Boolean(checkpoint),
hasRoot: Boolean(root),
});
return;
}
checkpoint.piLeafId = piLeafId;
dbg?.("coordinator.persistCheckpoint.position", {
messageCount: checkpoint.messageCount,
});
await saveCheckpoint(root, sessionId, checkpoint);
this.pendingCheckpoints.delete(sessionId);
dbg?.("coordinator.persistCheckpoint.end");
}
async closeAll(reason: string): Promise<void> {
dbg?.("coordinator.closeAll.start");
const states = [
...this.anonymous,
...[...this.bySession.values()].flatMap((states) => [...states]),
];
await Promise.all(states.map((state) => this.closeState(state, reason)));
this.completed.clear();
this.memoryStores.clear();
this.diskStores.clear();
this.pendingCheckpoints.clear();
dbg?.("coordinator.closeAll.end", { stateCount: states.length });
}
private async invoke(
sink: PiSink,
model: Model<Api>,
context: KlausContext,
options: SimpleStreamOptions,
): Promise<void> {
dbg?.("coordinator.invoke.start", {
messageCount: context.messages.length,
});
try {
if (!isKlausModelId(model.id))
throw new Error(`Unsupported Klaus model ${model.id}.`);
const required = requiredConstrainedTool(context.tools);
if (required)
throw new Error(
`Tool ${required} requires constrained sampling, which Klaus cannot provide.`,
);
const batch = recentToolResults(context);
dbg?.("coordinator.invoke.batch", { resultCount: batch.results.length });
const continuation = this.findContinuation(
options.sessionId,
batch.results,
);
const tools = toKlausTools(context.tools);
dbg?.("coordinator.invoke.continuation", {
found: Boolean(continuation),
hasHandle: Boolean(continuation?.handle),
});
if (continuation?.handle) {
const resultCount = batch.results.length;
const boundaryMessageCount = context.messages.length - resultCount;
const boundaryFingerprint = contextFingerprint(
{
systemPrompt: context.systemPrompt,
messages: context.messages.slice(0, boundaryMessageCount),
tools: context.tools,
},
`${model.id}\0${this.currentCwd()}`,
);
const incompatible =
batch.userAfter ||
continuation.boundaryMessageCount !== boundaryMessageCount ||
continuation.boundaryFingerprint !== boundaryFingerprint ||
JSON.stringify(continuation.tools) !== JSON.stringify(tools) ||
continuation.modelId !== model.id ||
continuation.cwd !== this.currentCwd() ||
continuation.systemPrompt !== (context.systemPrompt ?? "") ||
continuation.thinking !== options.reasoning ||
continuation.thinkingBudget !== thinkingBudget(options) ||
continuation.maxTokens !== options.maxTokens ||
!sameStringRecord(continuation.headers, options.headers) ||
!sameStringRecord(continuation.env, options.env);
dbg?.("coordinator.invoke.compatibility", { incompatible });
if (incompatible) {
await this.closeState(
continuation,
"Pi request changed across a tool boundary; replaying canonically.",
);
} else {
continuation.sink = sink;
continuation.context = context;
if (continuation.idleLease) clearTimeout(continuation.idleLease);
continuation.idleLease = undefined;
this.replaceAbort(continuation, options.signal);
continuation.timeoutMs = options.timeoutMs;
this.armTimeout(continuation, continuation.timeoutMs);
if (options.signal?.aborted) {
await this.closeState(
continuation,
"Pi aborted the Klaus query.",
true,
);
return;
}
if (options.onResponse && continuation.metadata) {
await options.onResponse(
{ status: 200, headers: continuation.metadata },
model,
);
}
if (options.signal?.aborted) {
await this.closeState(
continuation,
"Pi aborted the Klaus query.",
true,
);
return;
}
for (const result of batch.results) {
await continuation.handle.bridge.waitForPending(
[result.id],
options.signal,
);
if (options.signal?.aborted) {
await this.closeState(
continuation,
"Pi aborted the Klaus query.",
true,
);
return;
}
if (!continuation.handle.bridge.settle(result)) {
throw new Error(`Unknown Klaus tool result ${result.id}.`);
}
}
return;
}
}
await this.start(
sink,
model,
context,
options,
tools,
false,
batch.results.length === 0
? this.findReusable(options.sessionId)
: undefined,
);
} catch (error) {
dbg?.("coordinator.invoke.error", { kind: diagnosticKind(error) });
sink.error(error, options.signal?.aborted);
}
}
/** `reusable` is an idle child from this session's previous turn. It is used
* only when this turn resolves to exactly the delta the child already
* expects, and it is closed on any other outcome. */
private async start(
sink: PiSink,
model: Model<Api>,
context: KlausContext,
options: SimpleStreamOptions,
tools: KlausRequest["tools"],
forceCold = false,
reusable?: ActiveQuery,
): Promise<void> {
dbg?.("coordinator.start.begin", { messageCount: context.messages.length });
await this.cacheReady;
dbg?.("coordinator.start.cacheReady");
const store = this.storeFor(options.sessionId);
const cwd = this.currentCwd();
const identity = `${model.id}\0${cwd}`;
const candidates: CompletedQuery[] = [];
const memory =
!forceCold && options.sessionId
? this.completed.get(options.sessionId)
: undefined;
if (memory?.position) candidates.push(memory);
dbg?.("coordinator.start.memoryCandidate", { found: Boolean(memory) });
const checkpointRoot = this.checkpointRoot(options.sessionId);
if (!forceCold && options.sessionId && checkpointRoot) {
try {
for (const checkpoint of await loadCheckpoints(
checkpointRoot,
options.sessionId,
)) {
if (checkpoint.position) candidates.push({ ...checkpoint, store });
}
} catch (error) {
dbg?.("coordinator.start.checkpointLoadError", {
kind: diagnosticKind(error),
});
candidates.length = memory?.position ? 1 : 0;
}
}
dbg?.("coordinator.start.candidates", { count: candidates.length });
const prefixFingerprints = new Map<number, string>();
let completed: CompletedQuery | undefined;
for (const checkpoint of candidates) {
if (checkpoint.messageCount >= context.messages.length) continue;
let fingerprint = prefixFingerprints.get(checkpoint.messageCount);
if (!fingerprint) {
fingerprint = contextFingerprint(
{
systemPrompt: context.systemPrompt,
messages: context.messages.slice(0, checkpoint.messageCount),
tools: context.tools,
},
identity,
);
prefixFingerprints.set(checkpoint.messageCount, fingerprint);
}
if (
checkpoint.fingerprint === fingerprint &&
(!completed || checkpoint.messageCount > completed.messageCount)
) {
completed = checkpoint;
}
}
dbg?.("coordinator.start.completed", {
found: Boolean(completed),
messageCount: completed?.messageCount,
});
const offered =
reusable && !forceCold && reusable.handle?.isIdle() === true
? this.reuseFit(reusable, context, model, options, tools, identity)
: undefined;
dbg?.("coordinator.start.reuseFit");
const reuse = offered ? reusable : undefined;
if (reusable && !reuse) {
await this.closeState(
reusable,
"Klaus cannot reuse the idle child for this turn.",
);
}
const reuseBoundary = reuse?.messageCount;
const deltaContext: KlausContext = {
systemPrompt: context.systemPrompt,
messages:
reuseBoundary !== undefined
? context.messages.slice(reuseBoundary)
: completed
? context.messages.slice(completed.messageCount)
: context.messages,
tools: context.tools,
};
const linearResume =
(reuse !== undefined || completed !== undefined) &&
deltaContext.messages.length === 1 &&
isTextUser(deltaContext.messages[0]);
const canResume = reuse === undefined && completed !== undefined;
dbg?.("coordinator.start.resumePlan", { linearResume });
const staging = new StagingSessionStore(completed?.store ?? store);
const request: KlausRequest = {
modelId: model.id,
selector: klausSelector(model.id),
systemPrompt: context.systemPrompt ?? "",
prompt: linearResume
? finalUserPrompt(deltaContext)
: replayPrompt(deltaContext),
images: replayImages(deltaContext),
tools,
thinking: options.reasoning,
thinkingBudget: thinkingBudget(options),
maxTokens: options.maxTokens,
cwd,
sessionId: options.sessionId,
headers: options.headers ?? {},
env: options.env ?? {},
timeoutMs: options.timeoutMs,
sessionStore: staging,
resume: canResume ? completed?.sdkSessionId : undefined,
resumeAt: canResume ? completed?.position : undefined,
forkSession: canResume,
};
dbg?.("coordinator.start.request", {
imageCount: request.images?.length ?? 0,
});
const payloadRequest: KlausRequest = { ...request };
delete payloadRequest.sessionStore;
delete payloadRequest.resume;
delete payloadRequest.resumeAt;
delete payloadRequest.forkSession;
dbg?.("coordinator.start.payloadHook", {
enabled: options.onPayload !== undefined,
});
const hookPayload = options.onPayload
? await options.onPayload(payloadRequest, model)
: undefined;
const transformedPayload =
hookPayload === undefined ? payloadRequest : hookPayload;
const transformed = isRequest(transformedPayload)
? {
...transformedPayload,
sessionStore: request.sessionStore,
resume: request.resume,
resumeAt: request.resumeAt,
forkSession: request.forkSession,
}
: transformedPayload;
dbg?.("coordinator.start.payloadTransformed", {
transformed: hookPayload !== undefined,
});
if (!isRequest(transformed))
throw new Error("Klaus onPayload returned an invalid request.");
if (
transformed.resume &&
(transformed.modelId !== request.modelId ||
transformed.selector !== request.selector ||
transformed.systemPrompt !== request.systemPrompt ||
transformed.cwd !== request.cwd ||
JSON.stringify(transformed.tools) !== JSON.stringify(request.tools))
) {
await this.start(sink, model, context, options, tools, true);
return;
}
const scope = this.currentSession();
dbg?.("coordinator.start.scope", { persisted: scope?.persisted });
if (
reuse &&
(transformed.prompt !== request.prompt ||
transformed.systemPrompt !== request.systemPrompt ||
transformed.modelId !== request.modelId ||
transformed.selector !== request.selector ||
transformed.cwd !== request.cwd ||
JSON.stringify(transformed.tools) !== JSON.stringify(request.tools))
) {
dbg?.("coordinator.start.reuseRejectedByHook");
await this.closeState(
reuse,
"Klaus onPayload changed a reused turn; replaying canonically.",
);
await this.start(sink, model, context, options, tools, true);
return;
}
const state: ActiveQuery = reuse ?? {
id: Symbol("klaus-query"),
sessionId: options.sessionId,
sink,
toolCallIds: new Set(),
tools,
closed: false,
timeoutMs: transformed.timeoutMs,
secondary: Boolean(options.sessionId && options.sessionId !== scope?.id),
committable: true,
modelId: model.id,
cwd,
thinking: options.reasoning,
thinkingBudget: thinkingBudget(options),
maxTokens: options.maxTokens,
systemPrompt: context.systemPrompt ?? "",
headers: options.headers,
env: options.env,
emittedContent: false,
context,
store: completed?.store ?? store,
staging,
};
if (reuse) {
if (state.reuseLease) clearTimeout(state.reuseLease);
state.reuseLease = undefined;
state.sink = sink;
state.context = context;
state.tools = tools;
state.toolCallIds = new Set();
state.emittedContent = false;
state.timeoutMs = transformed.timeoutMs;
state.boundaryFingerprint = undefined;
state.boundaryMessageCount = undefined;
} else {
this.addState(state);
}
dbg?.("coordinator.start.stateAdded");
this.armTimeout(state, transformed.timeoutMs);
this.replaceAbort(state, options.signal);
if (options.signal?.aborted) {
await this.closeState(state, "Pi aborted the Klaus query.", true);
return;
}
const callbacks: KlausQueryCallbacks = {
onReady: async (metadata) => {
dbg?.("coordinator.callback.ready");
if (state.closed) return;
state.metadata = metadata;
await options.onResponse?.({ status: 200, headers: metadata }, model);
},
onActivity: () => {
dbg?.("coordinator.callback.activity", { closed: state.closed });
this.touchState(state, state.timeoutMs);
},
onNotice: (message) => {
dbg?.("coordinator.callback.notice");
this.warn(message);
},
onContent: (event) => {
dbg?.("coordinator.callback.content", { index: event.index });
if (state.closed) return;
state.emittedContent = true;
if (event.type === "tool-start") state.toolCallIds.add(event.id);
state.sink?.content(event);
},
onToolBoundary: (result) => {
dbg?.("coordinator.callback.toolBoundary", { closed: state.closed });
if (state.closed) return;
if (!state.sessionId) {
state.sink?.error(
new Error("Klaus cannot continue a sessionless tool call."),
);
void this.closeState(
state,
"Sessionless tool continuation is unsupported.",
);
return;
}
if (state.timer) clearTimeout(state.timer);
state.timer = undefined;
state.sink?.toolUse(
result.usage,
result.responseId,
result.responseModel,
);
if (state.sink) {
state.boundaryMessageCount = state.context.messages.length + 1;
state.boundaryFingerprint = contextFingerprint(
{
systemPrompt: state.context.systemPrompt,
messages: [...state.context.messages, state.sink.output],
tools: state.context.tools,
},
`${state.modelId}\0${state.cwd}`,
);
}
state.sink = undefined;
if (state.secondary) {
state.idleLease = setTimeout(
() =>
void this.closeState(
state,
"Klaus secondary query idle lease expired.",
),
15 * 60_000,
);
}
},
onResult: (result) => {
dbg?.("coordinator.callback.result", { closed: state.closed });
if (state.closed) return;
const resultSink = state.sink;
void (async () => {
if (
state.committable &&
state.sessionId &&
result.sdkSessionId &&
resultSink
) {
try {
await state.staging.commit();
if (!state.closed && state.committable) {
const completed: CompletedQuery = {
sdkSessionId: result.sdkSessionId,
position: result.position,
messageCount: state.context.messages.length + 1,
protocol: 1,
fingerprint: contextFingerprint(
{
systemPrompt: state.context.systemPrompt,
messages: [...state.context.messages, resultSink.output],
tools: state.context.tools,
},
`${state.modelId}\0${state.cwd}`,
),
store: state.store,
};
dbg?.("coordinator.checkpoint.published", {
messageCount: completed.messageCount,
});
this.completed.set(state.sessionId, completed);
state.sdkSessionId = completed.sdkSessionId;
state.position = completed.position;
state.messageCount = completed.messageCount;
state.lineageFingerprint = contextFingerprint(
{
systemPrompt: state.context.systemPrompt,
messages: [...state.context.messages, resultSink.output],
},
`${state.modelId}\0${state.cwd}`,
);
if (this.checkpointRoot(state.sessionId)) {
this.pendingCheckpoints.set(state.sessionId, completed);
}
}
} catch (error) {
this.warn(
`Klaus could not cache the completed session: ${error instanceof Error ? error.message : String(error)}`,
);
}
}
if (!state.closed) {
resultSink?.done(
result.usage,
result.responseId,
result.responseModel,
result.stopReason,
);
}
if (this.keepForReuse(state)) return;
await this.closeState(state, "Klaus query completed.");
})().catch((error) => {
resultSink?.error(error);
void this.closeState(state, "Klaus checkpoint failed.");
});
},
onError: (error) => {
dbg?.("coordinator.callback.error", {
kind: diagnosticKind(error),
closed: state.closed,
});
if (state.closed) return;
if (
transformed.resume &&
!state.emittedContent &&
state.sink &&
isReplayableResumeError(error)
) {
const retrySink = state.sink;
state.sink = undefined;
if (state.sessionId) {
this.completed.delete(state.sessionId);
this.pendingCheckpoints.delete(state.sessionId);
}
void (async () => {
await this.closeState(
state,
"Klaus resume failed; replaying canonically.",
);
await this.start(retrySink, model, context, options, tools, true);
})().catch((retryError) => retrySink.error(retryError));
return;
}
state.sink?.error(error);
void this.closeState(state, error.message);
},
};
if (reuse) {
try {
if (offered?.fit === "swap") {
dbg?.("coordinator.reuse.swapTools");
await state.handle?.setTools(tools);
}
if (state.metadata) {
await options.onResponse?.(
{ status: 200, headers: state.metadata },
model,
);
}
if (state.closed) return;
dbg?.("coordinator.reuse.send");
state.handle?.send(
{ prompt: transformed.prompt, images: transformed.images },
callbacks,
);
} catch (error) {
dbg?.("coordinator.reuse.error", { kind: diagnosticKind(error) });
state.sink = undefined;
await this.closeState(state, "Klaus could not reuse the idle child.");
await this.start(sink, model, context, options, tools, forceCold);
}
return;
}
try {
dbg?.("coordinator.oauth.resolve.start");
const oauth = await this.resolveOAuth(model.id);
dbg?.("coordinator.oauth.resolve.end");
if (state.closed) return;
dbg?.("coordinator.sdk.start");
const outgoing = rewriteSubscriptionPrompt(transformed, this.warn);
// Claude Agent SDK costs about 200 ms to parse and is needed only once
// Klaus handles a request, not while every Pi process starts.
const { startSdkQuery } = await import("./agent-sdk.js");
const handle = await startSdkQuery(outgoing, oauth, callbacks);
dbg?.("coordinator.sdk.started");
if (state.closed) {
await handle.close("Klaus query closed during startup.");
return;
}
state.handle = handle;
dbg?.("coordinator.handle.attached");
} catch (error) {
dbg?.("coordinator.start.error", { kind: diagnosticKind(error) });
state.sink?.error(error, options.signal?.aborted);
await this.closeState(state, "Klaus query startup failed.");
}
}
/** Decides whether an idle child can take this turn, and whether it needs a
* tool swap first. Lineage is compared without tools so a Pi tool-list change
* alone does not force a canonical replay. */
private reuseFit(
state: ActiveQuery,
context: KlausContext,
model: Model<Api>,
options: SimpleStreamOptions,
tools: KlausRequest["tools"],
identity: string,
): { fit: "reuse" | "swap" } | undefined {
const boundary = state.messageCount;
if (boundary === undefined || state.closed) return undefined;
const delta = context.messages.slice(boundary);
if (delta.length !== 1 || !isTextUser(delta[0])) return undefined;
if (
state.modelId !== model.id ||
state.cwd !== this.currentCwd() ||
state.systemPrompt !== (context.systemPrompt ?? "") ||
state.thinking !== options.reasoning ||
state.thinkingBudget !== thinkingBudget(options) ||
state.maxTokens !== options.maxTokens ||
!sameStringRecord(state.headers, options.headers) ||
!sameStringRecord(state.env, options.env)
) {
return undefined;
}
const lineage = contextFingerprint(
{
systemPrompt: context.systemPrompt,
messages: context.messages.slice(0, boundary),
},
identity,
);
if (state.lineageFingerprint !== lineage) return undefined;
return {
fit:
JSON.stringify(state.tools) === JSON.stringify(tools)
? "reuse"
: "swap",
};
}
/** A finished child keeps running for a short window so the next turn in the
* same session skips the child spawn and initialize round trip. */
private keepForReuse(state: ActiveQuery): boolean {
const reusable =
!state.closed &&
state.committable &&
Boolean(state.sessionId) &&
Boolean(state.sdkSessionId) &&
Boolean(state.position) &&
state.handle !== undefined;
dbg?.("coordinator.keepForReuse");
if (!reusable) return false;
state.sink = undefined;
if (state.timer) clearTimeout(state.timer);
state.timer = undefined;
if (state.reuseLease) clearTimeout(state.reuseLease);
state.reuseLease = setTimeout(
() => void this.closeState(state, "Klaus idle child lease expired."),
reuseLeaseMs(),
);
return true;
}
private findReusable(sessionId: string | undefined): ActiveQuery | undefined {
if (!sessionId) return undefined;
const found = [...(this.bySession.get(sessionId) ?? [])].find(
(state) => !state.closed && state.reuseLease !== undefined,
);
dbg?.("coordinator.findReusable", { found: Boolean(found) });
return found;
}
private findContinuation(
sessionId: string | undefined,
results: PiToolResult[],
): ActiveQuery | undefined {
dbg?.("coordinator.findContinuation");
if (!sessionId || results.length === 0) return undefined;
return [...(this.bySession.get(sessionId) ?? [])].find(
(state) =>
!state.closed &&
/** An idle child between turns holds no parked call, so its tool ids
* from the finished turn must never capture a continuation. */
state.reuseLease === undefined &&
results.every((result) => state.toolCallIds.has(result.id)),
);
}
private addState(state: ActiveQuery): void {
dbg?.("coordinator.addState");
if (!state.sessionId) {
this.anonymous.add(state);
return;
}
const states = this.bySession.get(state.sessionId) ?? new Set();
if ([...states].some((existing) => !existing.closed)) {
state.committable = false;
for (const existing of states) existing.committable = false;
}
states.add(state);
this.bySession.set(state.sessionId, states);
}
private replaceAbort(
state: ActiveQuery,
signal: AbortSignal | undefined,
): void {
dbg?.("coordinator.replaceAbort", {
hadSignal: Boolean(state.signal),
aborted: signal?.aborted,
});
if (state.signal && state.abort) {
state.signal.removeEventListener("abort", state.abort);
}
state.signal = signal;
state.abort = signal
? () => void this.closeState(state, "Pi aborted the Klaus query.", true)
: undefined;
if (signal && state.abort) {
signal.addEventListener("abort", state.abort, { once: true });
}
}
private touchState(state: ActiveQuery, timeoutMs: number | undefined): void {
dbg?.("coordinator.touchState", { timeoutMs });
this.armTimeout(state, timeoutMs);
if (!state.idleLease) return;
clearTimeout(state.idleLease);
state.idleLease = setTimeout(
() =>
void this.closeState(
state,
"Klaus secondary query idle lease expired.",
),
15 * 60_000,
);
}
private checkpointRoot(sessionId: string | undefined): string | undefined {
const scope = this.currentSession();
dbg?.("coordinator.checkpointRoot", { persisted: scope?.persisted });
return sessionId && scope?.id === sessionId && scope.persisted
? join(scope.dir, ".klaus", SDK_VERSION, "checkpoints")
: undefined;
}
private storeFor(sessionId: string | undefined): KlausSessionStore {
dbg?.("coordinator.storeFor");
if (!sessionId) return new MemorySessionStore();
const scope = this.currentSession();
if (scope?.id === sessionId && scope.persisted) {
const existing = this.diskStores.get(sessionId);
if (existing) return existing;
const store = new FileSessionStore(
join(scope.dir, ".klaus", SDK_VERSION, "sessions"),
);
this.diskStores.set(sessionId, store);
return store;
}
const existing = this.memoryStores.get(sessionId);
if (existing) return existing;
const store = new MemorySessionStore();
this.memoryStores.set(sessionId, store);
return store;
}
private armTimeout(state: ActiveQuery, timeoutMs: number | undefined): void {
dbg?.("coordinator.armTimeout", { timeoutMs, closed: state.closed });
if (state.timer) clearTimeout(state.timer);
if (!timeoutMs || state.closed || !state.sink) return;
state.timer = setTimeout(() => {
state.sink?.error(
new Error("Klaus timed out waiting for Claude activity."),
);
void this.closeState(state, "Klaus activity timeout.");
}, timeoutMs);
}
private async closeState(
state: ActiveQuery,
reason: string,
aborted = false,
): Promise<void> {
dbg?.("coordinator.closeState.start", {
aborted,
hasHandle: Boolean(state.handle),
});
if (state.closePromise) return state.closePromise;
state.closed = true;
state.closePromise = (async () => {
if (state.timer) clearTimeout(state.timer);
if (state.idleLease) clearTimeout(state.idleLease);
if (state.reuseLease) clearTimeout(state.reuseLease);
if (state.signal && state.abort) {
state.signal.removeEventListener("abort", state.abort);
}
state.sink?.error(new Error(reason), aborted);
if (state.handle) await state.handle.close(reason);
if (state.sessionId) {
const states = this.bySession.get(state.sessionId);
states?.delete(state);
if (states?.size === 0) this.bySession.delete(state.sessionId);
} else {
this.anonymous.delete(state);
}
dbg?.("coordinator.closeState.end");
})();
return state.closePromise;
}
}
function sameStringRecord(
left: Record<string, string | null> | undefined,
right: Record<string, string | null> | undefined,
): boolean {
const entries = (value: Record<string, string | null> | undefined) =>
Object.entries(value ?? {}).sort(([leftKey], [rightKey]) =>
leftKey < rightKey ? -1 : leftKey > rightKey ? 1 : 0,
);
return JSON.stringify(entries(left)) === JSON.stringify(entries(right));
}
export function isReplayableResumeError(error: Error): boolean {
return (
error instanceof KlausCacheCorruptionError ||
/^Klaus cache corruption:/i.test(error.message) ||
/^Resume rejected by --resume-drops-turn:/i.test(error.message) ||
/^No message found with message\.uuid of:/i.test(error.message)
);
}
function thinkingBudget(options: SimpleStreamOptions): number | undefined {
dbg?.("coordinator.thinkingBudget");
switch (options.reasoning) {
case "minimal":
return options.thinkingBudgets?.minimal;
case "low":
return options.thinkingBudgets?.low;
case "medium":
return options.thinkingBudgets?.medium;
case "high":
case "xhigh":
case "max":
return options.thinkingBudgets?.high;
default:
return undefined;
}
}
function isTextUser(
message: KlausContext["messages"][number] | undefined,
): boolean {
return (
message?.role === "user" &&
(typeof message.content === "string" ||
message.content.every((item) => item.type === "text"))
);
}
function finalUserPrompt(context: KlausContext): string {
dbg?.("coordinator.finalUserPrompt", {
messageCount: context.messages.length,
});
const message = context.messages.at(-1);
if (message?.role !== "user") return replayPrompt(context);
return typeof message.content === "string"
? message.content
: message.content
.filter((item) => item.type === "text")
.map((item) => item.text)
.join("\n");
}
function recentToolResults(context: KlausContext): {
results: PiToolResult[];
userAfter: boolean;
} {
dbg?.("coordinator.recentToolResults.start", {
messageCount: context.messages.length,
});
const results: PiToolResult[] = [];
let userAfter = false;
for (let index = context.messages.length - 1; index >= 0; index -= 1) {
const message = context.messages[index];
if (message.role === "assistant") break;
if (message.role === "user") {
userAfter = true;
continue;
}
results.unshift({
id: message.toolCallId,
content: modelFacingContent(message.content),
isError: message.isError,
});
}
dbg?.("coordinator.recentToolResults.end");
return { results, userAfter };
}
function isRequest(value: unknown): value is KlausRequest {
dbg?.("coordinator.isRequest");
if (typeof value !== "object" || value === null || Array.isArray(value))
return false;
const request = value as Partial<KlausRequest>;
return (
typeof request.modelId === "string" &&
isKlausModelId(request.modelId) &&
typeof request.selector === "string" &&
typeof request.systemPrompt === "string" &&
typeof request.prompt === "string" &&
(request.images === undefined ||
(Array.isArray(request.images) &&
request.images.every(
(image) =>
typeof image === "object" &&
image !== null &&
typeof image.data === "string" &&
typeof image.mimeType === "string",
))) &&
(request.thinking === undefined ||
["minimal", "low", "medium", "high", "xhigh", "max"].includes(
request.thinking,
)) &&
(request.thinkingBudget === undefined ||
(Number.isInteger(request.thinkingBudget) &&
request.thinkingBudget > 0)) &&
(request.maxTokens === undefined ||
(Number.isSafeInteger(request.maxTokens) && request.maxTokens > 0)) &&
(request.timeoutMs === undefined ||
(Number.isFinite(request.timeoutMs) && request.timeoutMs >= 0)) &&
(request.sessionId === undefined ||
typeof request.sessionId === "string") &&
Array.isArray(request.tools) &&
request.tools.every(
(tool) =>
typeof tool === "object" &&
tool !== null &&
typeof tool.name === "string" &&
typeof tool.description === "string" &&
typeof tool.inputSchema === "object" &&
tool.inputSchema !== null &&
!Array.isArray(tool.inputSchema),
) &&
typeof request.cwd === "string" &&
typeof request.headers === "object" &&
request.headers !== null &&
Object.values(request.headers).every(
(value) => typeof value === "string" || value === null,
) &&
(request.env === undefined ||
(typeof request.env === "object" &&
request.env !== null &&
Object.values(request.env).every((value) => typeof value === "string")))
);
}