Menu
popagent
publicLatest change 7cecf6a89b6f39dae8c5678f5b369014199aeb3c - Add self-hosted Mastra observability by AkurAI Build
import { Cron } from "croner";
import { agentRuntimeSettings, type AgentRuntimeSettingsStore } from "./agent-runtime-settings";
import { DEFAULT_WORKSPACE_ID, type AgentRuntimeSettings, type AgentSchedule, type AgentTask, type AgentTaskProgress } from "./api-types";
import { storage } from "./storage";
import { appLogger } from "./observability";
import { retryableInit } from "./retryable-init";
const HOST_TIMEZONE = Intl.DateTimeFormat().resolvedOptions().timeZone;
const TASK_COLUMNS = `id, session_id AS "sessionId", schedule_id AS "scheduleId",
workspace_id AS "workspaceId", prompt, model, status, output, error,
steps_completed AS "stepsCompleted", progress, recovery_count AS "recoveryCount",
created_at::text AS "createdAt", started_at::text AS "startedAt",
completed_at::text AS "completedAt"`;
const SCHEDULE_COLUMNS = `id, session_id AS "sessionId", workspace_id AS "workspaceId",
name, prompt, model, cron, timezone, enabled, next_run_at::text AS "nextRunAt",
last_run_at::text AS "lastRunAt", created_at::text AS "createdAt",
updated_at::text AS "updatedAt"`;
export class AgentTaskStore {
readonly storage = storage;
private readonly initializeOnce = retryableInit(() => this.initialize());
init(): Promise<void> {
return this.initializeOnce();
}
private async initialize() {
await this.storage.init();
await this.storage.db.none(`
CREATE TABLE IF NOT EXISTS popagent_agent_schedules (
id TEXT PRIMARY KEY,
session_id TEXT,
workspace_id TEXT NOT NULL DEFAULT '${DEFAULT_WORKSPACE_ID}',
name TEXT NOT NULL,
prompt TEXT NOT NULL,
model TEXT NOT NULL,
cron TEXT NOT NULL,
timezone TEXT NOT NULL DEFAULT 'UTC',
enabled BOOLEAN NOT NULL DEFAULT TRUE,
next_run_at TIMESTAMPTZ,
last_run_at TIMESTAMPTZ,
created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
updated_at TIMESTAMPTZ NOT NULL DEFAULT NOW()
);
CREATE TABLE IF NOT EXISTS popagent_agent_tasks (
id TEXT PRIMARY KEY,
session_id TEXT,
workspace_id TEXT NOT NULL DEFAULT '${DEFAULT_WORKSPACE_ID}',
schedule_id TEXT REFERENCES popagent_agent_schedules(id) ON DELETE SET NULL,
prompt TEXT NOT NULL,
model TEXT NOT NULL,
status TEXT NOT NULL,
output TEXT,
error TEXT,
steps_completed INTEGER NOT NULL DEFAULT 0,
progress TEXT,
recovery_count INTEGER NOT NULL DEFAULT 0,
created_at TIMESTAMPTZ NOT NULL DEFAULT NOW(),
started_at TIMESTAMPTZ,
completed_at TIMESTAMPTZ
);
ALTER TABLE popagent_agent_schedules
ADD COLUMN IF NOT EXISTS workspace_id TEXT NOT NULL DEFAULT '${DEFAULT_WORKSPACE_ID}';
ALTER TABLE popagent_agent_tasks
ADD COLUMN IF NOT EXISTS workspace_id TEXT NOT NULL DEFAULT '${DEFAULT_WORKSPACE_ID}';
CREATE INDEX IF NOT EXISTS popagent_agent_tasks_workspace_created_idx
ON popagent_agent_tasks (workspace_id, created_at DESC);
ALTER TABLE popagent_agent_tasks
ADD COLUMN IF NOT EXISTS steps_completed INTEGER NOT NULL DEFAULT 0;
ALTER TABLE popagent_agent_tasks ADD COLUMN IF NOT EXISTS progress TEXT;
ALTER TABLE popagent_agent_tasks
ADD COLUMN IF NOT EXISTS recovery_count INTEGER NOT NULL DEFAULT 0;
ALTER TABLE popagent_agent_tasks DROP COLUMN IF EXISTS agent_id;
ALTER TABLE popagent_agent_schedules DROP COLUMN IF EXISTS agent_id;
`);
const schedules = await this.storage.db.any<{ id: string; cron: string; enabled: boolean }>(
"SELECT id, cron, enabled FROM popagent_agent_schedules",
);
for (const schedule of schedules) {
const next = schedule.enabled ? new Cron(schedule.cron, { timezone: HOST_TIMEZONE }).nextRun() : null;
await this.storage.db.none(
"UPDATE popagent_agent_schedules SET timezone=$2,next_run_at=$3 WHERE id=$1",
[schedule.id, HOST_TIMEZONE, next],
);
}
}
async listTasks(workspaceId?: string): Promise<AgentTask[]> {
await this.init();
return this.storage.db.any<AgentTask>(
`SELECT ${TASK_COLUMNS} FROM popagent_agent_tasks
WHERE ($1::text IS NULL OR workspace_id=$1)
ORDER BY created_at DESC LIMIT 100`,
[workspaceId ?? null],
);
}
async listQueuedTasks(): Promise<AgentTask[]> {
await this.init();
return this.storage.db.any<AgentTask>(
`SELECT ${TASK_COLUMNS} FROM popagent_agent_tasks WHERE status='queued' ORDER BY created_at`,
);
}
async recoverInterruptedTasks(staleAfterMs: number): Promise<number> {
if (!Number.isFinite(staleAfterMs) || staleAfterMs <= 0) {
throw new RangeError("Task stale threshold must be positive");
}
await this.init();
return this.storage.db.tx(async (db) => {
await db.none(
`UPDATE popagent_agent_tasks
SET status='failed', error='stale after restart', progress='Failed', completed_at=NOW()
WHERE status='running'
AND started_at < NOW() - ($1::double precision * INTERVAL '1 millisecond')`,
[staleAfterMs],
);
const recovered = await db.any<{ id: string }>(
`UPDATE popagent_agent_tasks
SET status='queued',started_at=NULL,steps_completed=0,
progress='Queued after restart',recovery_count=recovery_count+1
WHERE status='running' RETURNING id`,
);
await db.none(
"UPDATE popagent_agent_tasks SET status='cancelled',progress='Cancelled',completed_at=NOW() WHERE status='cancelling'",
);
return recovered.length;
});
}
async createTask(input: {
sessionId?: string;
workspaceId?: string;
scheduleId?: string;
prompt: string;
model: string;
}): Promise<AgentTask> {
await this.init();
return this.storage.db.one<AgentTask>(
`INSERT INTO popagent_agent_tasks
(id, session_id, schedule_id, workspace_id, prompt, model, status)
VALUES ($1,$2,$3,$4,$5,$6,'queued') RETURNING ${TASK_COLUMNS}`,
[
crypto.randomUUID(),
input.sessionId ?? null,
input.scheduleId ?? null,
input.workspaceId ?? DEFAULT_WORKSPACE_ID,
input.prompt,
input.model,
],
);
}
async claim(id: string): Promise<boolean> {
await this.init();
return Boolean(await this.storage.db.oneOrNone(
"UPDATE popagent_agent_tasks SET status='running', started_at=NOW(), error=NULL, steps_completed=0, progress='Starting' WHERE id=$1 AND status='queued' RETURNING id",
[id],
));
}
async updateProgress(id: string, update: AgentTaskProgress) {
await this.init();
await this.storage.db.none(
`UPDATE popagent_agent_tasks
SET steps_completed=GREATEST(steps_completed,$2),progress=$3
WHERE id=$1 AND status='running'`,
[id, update.stepsCompleted, update.progress],
);
}
async complete(id: string, output: string): Promise<"completed" | "cancelled" | undefined> {
await this.init();
const result = await this.storage.db.oneOrNone<{ status: "completed" | "cancelled" }>(
`UPDATE popagent_agent_tasks SET status=CASE WHEN status='cancelling' THEN 'cancelled' ELSE 'completed' END,
output=$2, progress=CASE WHEN status='cancelling' THEN 'Cancelled' ELSE 'Completed' END, completed_at=NOW()
WHERE id=$1 AND status IN ('running','cancelling') RETURNING status`, [id, output]);
return result?.status;
}
async fail(id: string, error: string) {
await this.init();
await this.storage.db.none(
"UPDATE popagent_agent_tasks SET status='failed', error=$2, progress='Failed', completed_at=NOW() WHERE id=$1 AND status='running'",
[id, error],
);
}
async cancel(id: string): Promise<"cancelling" | "cancelled" | undefined> {
await this.init();
const result = await this.storage.db.oneOrNone<{ status: "cancelling" | "cancelled" }>(
`UPDATE popagent_agent_tasks
SET status=CASE status WHEN 'queued' THEN 'cancelled' ELSE 'cancelling' END,
progress=CASE status WHEN 'queued' THEN 'Cancelled' ELSE 'Cancelling' END,
completed_at=CASE status WHEN 'queued' THEN NOW() ELSE NULL END
WHERE id=$1 AND status IN ('queued','running')
RETURNING status`,
[id],
);
return result?.status;
}
async finishCancellation(id: string): Promise<void> {
await this.init();
await this.storage.db.none(
"UPDATE popagent_agent_tasks SET status='cancelled',progress='Cancelled',completed_at=NOW() WHERE id=$1 AND status='cancelling'",
[id],
);
}
async listSchedules(workspaceId?: string): Promise<AgentSchedule[]> {
await this.init();
return this.storage.db.any<AgentSchedule>(
`SELECT ${SCHEDULE_COLUMNS} FROM popagent_agent_schedules
WHERE ($1::text IS NULL OR workspace_id=$1)
ORDER BY name`,
[workspaceId ?? null],
);
}
async createSchedule(input: {
sessionId?: string;
workspaceId?: string;
name: string;
prompt: string;
model: string;
cron: string;
enabled: boolean;
}): Promise<AgentSchedule> {
await this.init();
const next = input.enabled ? new Cron(input.cron, { timezone: HOST_TIMEZONE }).nextRun() : null;
return this.storage.db.one<AgentSchedule>(
`INSERT INTO popagent_agent_schedules
(id,session_id,workspace_id,name,prompt,model,cron,timezone,enabled,next_run_at)
VALUES ($1,$2,$3,$4,$5,$6,$7,$8,$9,$10) RETURNING ${SCHEDULE_COLUMNS}`,
[
crypto.randomUUID(),
input.sessionId ?? null,
input.workspaceId ?? DEFAULT_WORKSPACE_ID,
input.name,
input.prompt,
input.model,
input.cron,
HOST_TIMEZONE,
input.enabled,
next,
],
);
}
async updateSchedule(
id: string,
input: { workspaceId?: string; name: string; prompt: string; model: string; cron: string; enabled: boolean },
): Promise<AgentSchedule | undefined> {
await this.init();
const next = input.enabled ? new Cron(input.cron, { timezone: HOST_TIMEZONE }).nextRun() : null;
return (await this.storage.db.oneOrNone<AgentSchedule>(
`UPDATE popagent_agent_schedules
SET workspace_id=COALESCE($2, workspace_id),name=$3,prompt=$4,model=$5,cron=$6,timezone=$7,
enabled=$8,next_run_at=$9,updated_at=NOW()
WHERE id=$1 RETURNING ${SCHEDULE_COLUMNS}`,
[id, input.workspaceId ?? null, input.name, input.prompt, input.model, input.cron, HOST_TIMEZONE, input.enabled, next],
)) ?? undefined;
}
async deleteSchedule(id: string): Promise<boolean> {
await this.init();
return Boolean(await this.storage.db.oneOrNone(
"DELETE FROM popagent_agent_schedules WHERE id=$1 RETURNING id",
[id],
));
}
async materializeDueTasks(): Promise<AgentTask[]> {
await this.init();
return this.storage.db.tx(async (db) => {
const schedules = await db.any<AgentSchedule>(
`SELECT ${SCHEDULE_COLUMNS} FROM popagent_agent_schedules
WHERE enabled AND next_run_at <= NOW()
ORDER BY next_run_at, id
FOR UPDATE SKIP LOCKED`,
);
const tasks: AgentTask[] = [];
for (const schedule of schedules) {
const next = new Cron(schedule.cron, { timezone: HOST_TIMEZONE }).nextRun();
const task = await db.one<AgentTask>(
`INSERT INTO popagent_agent_tasks
(id, session_id, schedule_id, workspace_id, prompt, model, status)
VALUES ($1,$2,$3,$4,$5,$6,'queued') RETURNING ${TASK_COLUMNS}`,
[crypto.randomUUID(), schedule.sessionId, schedule.id, schedule.workspaceId, schedule.prompt, schedule.model],
);
await db.none(
"UPDATE popagent_agent_schedules SET last_run_at=NOW(),next_run_at=$2,timezone=$3,updated_at=NOW() WHERE id=$1",
[schedule.id, next, HOST_TIMEZONE],
);
tasks.push(task);
}
return tasks;
});
}
}
type TaskRepository = Pick<
AgentTaskStore,
"init" | "recoverInterruptedTasks" | "listQueuedTasks" | "claim" | "updateProgress"
| "complete" | "fail" | "cancel" | "finishCancellation" | "materializeDueTasks"
>;
type AgentTaskExecutor = (
task: AgentTask,
signal: AbortSignal,
turnId: string,
reportProgress: (progress: AgentTaskProgress) => Promise<void>,
) => Promise<string>;
export type AgentTaskLifecycle = (event: {
task: AgentTask;
turnId: string;
status: "running" | "completed" | "failed" | "cancelled" | "interrupted";
error?: string;
}) => Promise<void>;
type TaskRuntimeSettingsRepository = Pick<AgentRuntimeSettingsStore, "get">;
export class AgentTaskRuntime {
private timer?: Timer;
private started = false;
private polling = false;
private active = 0;
private pending: AgentTask[] = [];
private readonly scheduled = new Set<string>();
private readonly controllers = new Map<string, AbortController>();
private readonly jobs = new Set<Promise<void>>();
private configuration?: AgentRuntimeSettings;
constructor(
readonly store: TaskRepository,
private readonly executeTask: AgentTaskExecutor,
private readonly lifecycle?: AgentTaskLifecycle,
private readonly settings: TaskRuntimeSettingsRepository = agentRuntimeSettings,
) {}
configure(settings: AgentRuntimeSettings) {
if (!Number.isInteger(settings.taskConcurrency) || settings.taskConcurrency < 1) {
throw new RangeError("Task concurrency must be a positive integer");
}
if (!Number.isFinite(settings.taskPollIntervalMs) || settings.taskPollIntervalMs <= 0) {
throw new RangeError("Task poll interval must be positive");
}
if (!Number.isFinite(settings.taskTimeoutMs) || settings.taskTimeoutMs <= 0) {
throw new RangeError("Task timeout must be positive");
}
if (!Number.isFinite(settings.taskStaleAfterMs) || settings.taskStaleAfterMs <= 0) {
throw new RangeError("Task stale threshold must be positive");
}
this.configuration = settings;
this.restartPolling();
this.drain();
}
private configured(): AgentRuntimeSettings {
if (!this.configuration) throw new Error("Task runtime settings are not loaded");
return this.configuration;
}
private restartPolling() {
clearInterval(this.timer);
this.timer = this.started
? setInterval(() => this.poll(), this.configured().taskPollIntervalMs)
: undefined;
}
async start() {
if (this.started) return;
this.configure(await this.settings.get());
this.started = true;
try {
await this.store.init();
const recovered = await this.store.recoverInterruptedTasks(this.configured().taskStaleAfterMs);
if (recovered > 0) appLogger().info("task.recovered", { count: recovered });
for (const task of await this.store.listQueuedTasks()) this.enqueue(task);
this.drain();
this.restartPolling();
} catch (error) {
this.started = false;
this.pending = [];
this.scheduled.clear();
this.restartPolling();
throw error;
}
}
async stop() {
clearInterval(this.timer);
this.timer = undefined;
this.started = false;
this.pending = [];
for (const controller of this.controllers.values()) {
controller.abort(new Error("Background task interrupted"));
}
await Promise.allSettled([...this.jobs]);
if (this.configuration) {
await this.store.recoverInterruptedTasks(this.configuration.taskStaleAfterMs);
}
this.scheduled.clear();
}
enqueue(task: AgentTask) {
if (this.scheduled.has(task.id)) return;
this.scheduled.add(task.id);
this.pending.push(task);
this.drain();
}
async cancel(id: string) {
const cancellation = await this.store.cancel(id);
if (!cancellation) return false;
const controller = this.controllers.get(id);
if (controller) {
controller.abort(new Error("Background task cancelled"));
} else {
this.pending = this.pending.filter((task) => task.id !== id);
this.scheduled.delete(id);
}
return true;
}
private drain() {
while (this.started && this.active < this.configured().taskConcurrency) {
const task = this.pending.shift();
if (!task) return;
this.active++;
const job = this.execute(task)
.catch((error) => {
appLogger().error("task.failed", { error: error instanceof Error ? error.message : String(error) });
})
.finally(() => {
this.active--;
this.scheduled.delete(task.id);
this.jobs.delete(job);
this.drain();
});
this.jobs.add(job);
}
}
private async execute(task: AgentTask) {
const controller = new AbortController();
this.controllers.set(task.id, controller);
const turnId = crypto.randomUUID();
let timeout: Timer | undefined;
let timeoutError: Error | undefined;
try {
if (!await this.store.claim(task.id)) return;
await this.lifecycle?.({ task, turnId, status: "running" });
const taskTimeoutMs = this.configured().taskTimeoutMs;
timeoutError = new Error(`Background task timed out after ${taskTimeoutMs}ms`);
const interrupted = Promise.withResolvers<never>();
controller.signal.addEventListener(
"abort",
() => interrupted.reject(controller.signal.reason),
{ once: true },
);
timeout = setTimeout(() => controller.abort(timeoutError), taskTimeoutMs);
const output = await Promise.race([
this.executeTask(
task,
controller.signal,
turnId,
(progress) => this.store.updateProgress(task.id, progress),
),
interrupted.promise,
]);
if (controller.signal.aborted) throw controller.signal.reason;
const completion = await this.store.complete(task.id, output);
if (completion === "cancelled") {
await this.lifecycle?.({ task, turnId, status: "cancelled" });
} else if (completion === "completed") {
await this.lifecycle?.({ task, turnId, status: "completed" });
}
} catch (error) {
const message = error instanceof Error ? error.message : String(error);
if (timeoutError && controller.signal.reason === timeoutError) {
await this.store.fail(task.id, timeoutError.message);
await this.lifecycle?.({ task, turnId, status: "failed", error: timeoutError.message });
} else if (controller.signal.aborted) {
const cancelled = message === "Background task cancelled";
if (cancelled) await this.store.finishCancellation(task.id);
await this.lifecycle?.({
task,
turnId,
status: cancelled ? "cancelled" : "interrupted",
error: message,
});
} else {
await this.store.fail(task.id, message);
await this.lifecycle?.({ task, turnId, status: "failed", error: message });
}
} finally {
clearTimeout(timeout);
this.controllers.delete(task.id);
}
}
private poll() {
if (this.polling || !this.started) return;
this.polling = true;
void this.tick()
.catch((error) => {
appLogger().error("task.scheduler_failed", { error: error instanceof Error ? error.message : String(error) });
})
.finally(() => {
this.polling = false;
});
}
private async tick() {
for (const task of await this.store.materializeDueTasks()) this.enqueue(task);
}
}