RPC: handle tau auto-compaction retries

main
Peter Steinberger 2025-12-10 11:40:32 +00:00
parent 81385cf820
commit 8f37f15a33
1 changed files with 68 additions and 7 deletions

View File

@ -22,6 +22,10 @@ class TauRpcClient {
private stderr = ""; private stderr = "";
private buffer: string[] = []; private buffer: string[] = [];
private idleTimer: NodeJS.Timeout | null = null; private idleTimer: NodeJS.Timeout | null = null;
private resolveTimer: NodeJS.Timeout | null = null;
private compactionRunning = false;
private pendingRetryCount = 0;
private seenAgentEnd = false;
private pending: private pending:
| { | {
resolve: (r: TauRpcResult) => void; resolve: (r: TauRpcResult) => void;
@ -37,6 +41,14 @@ class TauRpcClient {
private readonly cwd: string | undefined, private readonly cwd: string | undefined,
) {} ) {}
private resetRunState() {
this.buffer = [];
this.compactionRunning = false;
this.pendingRetryCount = 0;
this.seenAgentEnd = false;
this.clearResolveTimer();
}
private ensureChild() { private ensureChild() {
if (this.child) return; if (this.child) return;
this.child = spawn(this.argv[0], this.argv.slice(1), { this.child = spawn(this.argv[0], this.argv.slice(1), {
@ -49,6 +61,7 @@ class TauRpcClient {
this.stderr += d.toString(); this.stderr += d.toString();
}); });
this.child.on("exit", (code, signal) => { this.child.on("exit", (code, signal) => {
this.clearResolveTimer();
if (this.idleTimer) clearTimeout(this.idleTimer); if (this.idleTimer) clearTimeout(this.idleTimer);
if (this.pending) { if (this.pending) {
const pending = this.pending; const pending = this.pending;
@ -63,6 +76,7 @@ class TauRpcClient {
signal, signal,
}); });
} }
this.resetRunState();
this.dispose(); this.dispose();
}); });
} }
@ -84,6 +98,7 @@ class TauRpcClient {
success?: boolean; success?: boolean;
error?: string; error?: string;
message?: unknown; message?: unknown;
willRetry?: boolean;
}; };
if (evt.type === "response" && evt.command === "prompt") { if (evt.type === "response" && evt.command === "prompt") {
@ -91,6 +106,8 @@ class TauRpcClient {
const pending = this.pending; const pending = this.pending;
this.pending = undefined; this.pending = undefined;
this.buffer = []; this.buffer = [];
this.clearResolveTimer();
this.resetRunState();
if (pending) { if (pending) {
clearTimeout(pending.timer); clearTimeout(pending.timer);
pending.reject( pending.reject(
@ -102,14 +119,25 @@ class TauRpcClient {
} }
} }
if (evt.type === "auto_compaction_start") {
this.compactionRunning = true;
this.clearResolveTimer();
return;
}
if (evt.type === "auto_compaction_end") {
this.compactionRunning = false;
if (evt.willRetry) this.pendingRetryCount += 1;
this.scheduleMaybeResolve();
return;
}
if (evt?.type === "agent_end") { if (evt?.type === "agent_end") {
// Tau signals the end of the prompt/response cycle; resolve with all buffered output. this.seenAgentEnd = true;
const pending = this.pending; if (this.pendingRetryCount > 0) {
this.pending = undefined; this.pendingRetryCount -= 1;
const out = this.buffer.join("\n"); }
this.buffer = []; this.scheduleMaybeResolve();
clearTimeout(pending.timer);
pending.resolve({ stdout: out, stderr: this.stderr, code: 0 });
return; return;
} }
} catch { } catch {
@ -117,6 +145,37 @@ class TauRpcClient {
} }
} }
private scheduleMaybeResolve() {
if (!this.pending) return;
this.clearResolveTimer();
// Allow a short window for auto-compaction events to arrive after agent_end.
this.resolveTimer = setTimeout(() => {
this.resolveTimer = null;
this.maybeResolve();
}, 150);
}
private maybeResolve() {
if (!this.pending) return;
if (!this.seenAgentEnd) return;
if (this.compactionRunning) return;
if (this.pendingRetryCount > 0) return;
const pending = this.pending;
this.pending = undefined;
const out = this.buffer.join("\n");
this.buffer = [];
clearTimeout(pending.timer);
pending.resolve({ stdout: out, stderr: this.stderr, code: 0 });
}
private clearResolveTimer() {
if (this.resolveTimer) {
clearTimeout(this.resolveTimer);
this.resolveTimer = null;
}
}
private resetTimeout() { private resetTimeout() {
if (!this.pending) return; if (!this.pending) return;
const capMs = this.pending.capMs; const capMs = this.pending.capMs;
@ -142,6 +201,7 @@ class TauRpcClient {
} }
const child = this.child; const child = this.child;
if (!child) throw new Error("tau rpc child not initialized"); if (!child) throw new Error("tau rpc child not initialized");
this.resetRunState();
await new Promise<void>((resolve, reject) => { await new Promise<void>((resolve, reject) => {
const ok = child.stdin.write( const ok = child.stdin.write(
`${JSON.stringify({ `${JSON.stringify({
@ -168,6 +228,7 @@ class TauRpcClient {
} }
dispose() { dispose() {
this.clearResolveTimer();
this.rl?.close(); this.rl?.close();
this.rl = null; this.rl = null;
if (this.child && !this.child.killed) { if (this.child && !this.child.killed) {