25 lines
1.2 KiB
TypeScript
25 lines
1.2 KiB
TypeScript
import "server-only";
|
|
import { unlinkSync, writeFileSync } from "node:fs";
|
|
import { tmpdir } from "node:os";
|
|
import { join } from "node:path";
|
|
|
|
const state = globalThis as typeof globalThis & { __aiChatbotActiveWork?: number; __aiChatbotDurableWork?: number };
|
|
const path = join(process.platform === "win32" ? tmpdir() : "/tmp", `ai-chatbot-active-work-${process.pid}.json`);
|
|
|
|
function update(delta: number, durable: boolean) {
|
|
state.__aiChatbotActiveWork = (state.__aiChatbotActiveWork || 0) + delta;
|
|
if (durable) state.__aiChatbotDurableWork = (state.__aiChatbotDurableWork || 0) + delta;
|
|
try {
|
|
writeFileSync(path, JSON.stringify({ pid: process.pid, active: state.__aiChatbotActiveWork, durable: state.__aiChatbotDurableWork || 0 }), { mode: 0o600 });
|
|
} catch (error) {
|
|
try { unlinkSync(path); } catch {}
|
|
console.warn("[Active Work] Failed to record running work:", error);
|
|
}
|
|
}
|
|
|
|
/** Includes generation that continues after the client connection closes. */
|
|
export async function trackRuntimeWork<T>(execute: () => T | Promise<T>, options: { durable?: boolean } = {}): Promise<T> {
|
|
update(1, options.durable === true);
|
|
try { return await execute(); } finally { update(-1, options.durable === true); }
|
|
}
|