@posthog/agent
Version:
TypeScript agent framework wrapping Claude Agent SDK with Git-based task execution for PostHog
196 lines (193 loc) • 6.43 kB
JavaScript
import { Logger } from './utils/logger.js';
class SessionStore {
posthogAPI;
pendingEntries = new Map();
flushTimeouts = new Map();
configs = new Map();
logger;
constructor(posthogAPI, logger) {
this.posthogAPI = posthogAPI;
this.logger =
logger ?? new Logger({ debug: false, prefix: "[SessionStore]" });
// Flush all pending on process exit
const flushAllAndExit = async () => {
const flushPromises = [];
for (const sessionId of this.configs.keys()) {
flushPromises.push(this.flush(sessionId));
}
await Promise.all(flushPromises);
process.exit(0);
};
process.on("beforeExit", () => {
flushAllAndExit().catch((e) => this.logger.error("Flush failed:", e));
});
process.on("SIGINT", () => {
flushAllAndExit().catch((e) => this.logger.error("Flush failed:", e));
});
process.on("SIGTERM", () => {
flushAllAndExit().catch((e) => this.logger.error("Flush failed:", e));
});
}
/** Register a session for persistence */
register(sessionId, config) {
this.configs.set(sessionId, config);
}
/** Unregister and flush pending */
async unregister(sessionId) {
await this.flush(sessionId);
this.configs.delete(sessionId);
}
/** Check if a session is registered for persistence */
isRegistered(sessionId) {
return this.configs.has(sessionId);
}
/**
* Append a raw JSON-RPC line for persistence.
* Parses and wraps as StoredNotification for the API.
*/
appendRawLine(sessionId, line) {
const config = this.configs.get(sessionId);
if (!config) {
return;
}
try {
const message = JSON.parse(line);
const entry = {
type: "notification",
timestamp: new Date().toISOString(),
notification: message,
};
const pending = this.pendingEntries.get(sessionId) ?? [];
pending.push(entry);
this.pendingEntries.set(sessionId, pending);
this.scheduleFlush(sessionId);
}
catch {
this.logger.warn("Failed to parse raw line for persistence", {
sessionId,
lineLength: line.length,
});
}
}
/** Load raw JSON-RPC messages from S3 */
async load(logUrl) {
const response = await fetch(logUrl);
if (!response.ok) {
return [];
}
const content = await response.text();
if (!content.trim())
return [];
return content
.trim()
.split("\n")
.map((line) => {
try {
return JSON.parse(line);
}
catch {
return null;
}
})
.filter((entry) => entry !== null);
}
/** Force flush pending entries */
async flush(sessionId) {
const config = this.configs.get(sessionId);
const pending = this.pendingEntries.get(sessionId);
if (!config || !pending?.length)
return;
this.pendingEntries.delete(sessionId);
const timeout = this.flushTimeouts.get(sessionId);
if (timeout) {
clearTimeout(timeout);
this.flushTimeouts.delete(sessionId);
}
if (!this.posthogAPI) {
this.logger.debug("No PostHog API configured, skipping flush");
return;
}
try {
await this.posthogAPI.appendTaskRunLog(config.taskId, config.runId, pending);
}
catch (error) {
this.logger.error("Failed to persist session logs:", error);
}
}
scheduleFlush(sessionId) {
const existing = this.flushTimeouts.get(sessionId);
if (existing)
clearTimeout(existing);
const timeout = setTimeout(() => this.flush(sessionId), 500);
this.flushTimeouts.set(sessionId, timeout);
}
/** Get the persistence config for a session */
getConfig(sessionId) {
return this.configs.get(sessionId);
}
/**
* Start a session for persistence.
* Loads the task run and updates status to "in_progress".
*/
async start(sessionId, taskId, runId) {
if (!this.posthogAPI) {
this.logger.debug("No PostHog API configured, registering session without persistence");
this.register(sessionId, {
taskId,
runId,
logUrl: "",
});
return undefined;
}
const taskRun = await this.posthogAPI.getTaskRun(taskId, runId);
this.register(sessionId, {
taskId,
runId,
logUrl: taskRun.log_url,
});
await this.updateTaskRun(sessionId, { status: "in_progress" });
return taskRun;
}
/**
* Mark a session as completed.
*/
async complete(sessionId) {
await this.flush(sessionId);
await this.updateTaskRun(sessionId, { status: "completed" });
}
/**
* Mark a session as failed.
*/
async fail(sessionId, error) {
await this.flush(sessionId);
const message = typeof error === "string" ? error : error.message;
await this.updateTaskRun(sessionId, {
status: "failed",
error_message: message,
});
this.logger.error("Session failed", { sessionId, error: message });
}
/**
* Update the task run associated with a session.
*/
async updateTaskRun(sessionId, update) {
const config = this.configs.get(sessionId);
if (!config) {
this.logger.error(`Cannot update task run: session ${sessionId} not registered`);
return undefined;
}
if (!this.posthogAPI) {
this.logger.debug("No PostHog API configured, skipping task run update");
return undefined;
}
try {
return await this.posthogAPI.updateTaskRun(config.taskId, config.runId, update);
}
catch (error) {
this.logger.error("Failed to update task run:", error);
return undefined;
}
}
}
export { SessionStore };
//# sourceMappingURL=session-store.js.map