diff --git a/.codex/AGENTS.md b/.codex/AGENTS.md new file mode 100644 index 00000000..291ed7ca --- /dev/null +++ b/.codex/AGENTS.md @@ -0,0 +1,37 @@ +# Project Instructions for Claude Code + +## Development Guidelines + +- We use **pnpm** as our package manager +- At the beginning of any conversation, familiarize yourself with the repo structure +- Prefer method names with linux style, lowercase short words +- Never promote yourself in commits (no "Generated with Claude" messages) +- If we refactor something and are no longer using an old file, clean up after yourself + +`pnpm build` is the build command - not `pnpm build:_` or anything else. + +--- + +# Rule: Read the Docs First + +**You must do this before proposing designs, writing code, or editing files:** + +## 1) Read the canonical specs in `/docs/` +Review the relevant files and design constraints before you suggest anything. + +## 2) Follow established codebase patterns +- **Examine existing code** to understand established patterns, conventions, and architectural decisions +- **Maintain consistency** with existing naming conventions, file structures, and coding patterns +- **Reuse existing patterns** rather than inventing new ones unless the spec explicitly requires deviation +- If you must deviate from existing patterns, explicitly justify why based on the `/docs/` specifications + +## 3) If the spec is missing or ambiguous +- **Stop and propose a spec addition** first (succinct ADR-style note). Do not implement until the spec gap is resolved. +- Offer a default, but mark it as **"Proposed addition to /docs"**. + +## Checklists + +### Before coding +- [ ] I looked up the relevant `/docs/*` file(s). +- [ ] I examined existing codebase patterns for similar functionality. +- [ ] I verified my approach is consistent with established conventions. diff --git a/microprojects/playground/server/package.json b/microprojects/playground/server/package.json index eda79426..68f2757d 100644 --- a/microprojects/playground/server/package.json +++ b/microprojects/playground/server/package.json @@ -13,7 +13,7 @@ "@kernl-sdk/pg": "workspace:*", "@kernl-sdk/server": "workspace:*", "@kernl-sdk/turbopuffer": "workspace:*", - "kernl": "*" + "kernl": "workspace:*" }, "devDependencies": { "@types/node": "^24.10.0", diff --git a/microprojects/playground/server/src/agents/sleeper.ts b/microprojects/playground/server/src/agents/sleeper.ts new file mode 100644 index 00000000..aa4ec562 --- /dev/null +++ b/microprojects/playground/server/src/agents/sleeper.ts @@ -0,0 +1,16 @@ +// import "@kernl-sdk/ai/openai"; +import { openai } from "@kernl-sdk/ai/openai"; +import { Agent } from "kernl"; +// import { anthropic } from "@kernl-sdk/ai/anthropic"; + + +export const sleeper = new Agent({ + id: "sleeper", + name: "Sleeper", + description: "An agent that demonstrates the sleep/wakeup system tool", + instructions: `You are a helpful assistant that can pause and wait. +When asked to wait or sleep, use the wait_until tool with an appropriate delay. +Always explain what you're doing before sleeping.`, + model: openai("gpt-5.1"), + memory: { enabled: true }, +}); diff --git a/microprojects/playground/server/src/app.ts b/microprojects/playground/server/src/app.ts index ed17886b..025bfee0 100644 --- a/microprojects/playground/server/src/app.ts +++ b/microprojects/playground/server/src/app.ts @@ -1,8 +1,11 @@ import { Kernl } from "kernl"; import { pgvector, postgres } from "@kernl-sdk/pg"; -import { turbopuffer } from "@kernl-sdk/turbopuffer"; +import "@kernl-sdk/ai/openai"; + +// import { turbopuffer } from "@kernl-sdk/turbopuffer"; import { echo } from "./agents/echo"; +import { sleeper } from "./agents/sleeper"; import { titler } from "./agents/titler"; import { watson } from "./agents/watson"; @@ -12,13 +15,21 @@ export function build(): Kernl { db: postgres({ connstr: process.env.DATABASE_URL! }), vector: pgvector({ connstr: process.env.DATABASE_URL! }), }, + // scheduler: true, + scheduler: { + autoStart: true + } }); // --- agents --- kernl.register(echo); + kernl.register(sleeper); kernl.register(titler); kernl.register(watson); + // start wakeup scheduler + // kernl.schedule?.start(); + return kernl; } diff --git a/packages/kernl/src/agent.ts b/packages/kernl/src/agent.ts index 3aebec08..158a7f1f 100644 --- a/packages/kernl/src/agent.ts +++ b/packages/kernl/src/agent.ts @@ -1,3 +1,7 @@ +/** + * /packages/kernl/src/agent.ts + */ + import { message, LanguageModel, @@ -15,7 +19,7 @@ import type { RThreadUpdateParams, } from "@/api/resources/threads/types"; import type { Context, UnknownContext } from "./context"; -import { Tool, memory } from "./tool"; +import { Tool, memory, sleep } from "./tool"; import { BaseToolkit } from "./tool/toolkit"; import { InputGuardrail, @@ -112,14 +116,23 @@ export class Agent< this.kernl = kernl; // initialize system toolkits + + // Memory System Tool if (this.memory.enabled) { // safety: system tools only rely on ctx.agent, not ctx.context const toolkit = memory as unknown as BaseToolkit; this.systools.push(toolkit); toolkit.bind(this); } + // Sleep System Tool + { + const sleepToolkit = sleep as unknown as BaseToolkit; + this.systools.push(sleepToolkit); + sleepToolkit.bind(this); + } } + /** * Blocking execution - spawns or resumes thread and waits for completion * @@ -263,6 +276,21 @@ export class Agent< return undefined; } + /** + * @internal + * + * Check if a tool ID belongs to a system toolkit. + * + * @param id The tool ID to check + * @returns true if the tool is a system tool + */ + isSysTool(id: string): boolean { + for (const toolkit of this.systools) { + if (toolkit.get(id)) return true; + } + return false; + } + /** * @internal * diff --git a/packages/kernl/src/agent/__tests__/systools.test.ts b/packages/kernl/src/agent/__tests__/systools.test.ts index ca2f776d..2ee458f5 100644 --- a/packages/kernl/src/agent/__tests__/systools.test.ts +++ b/packages/kernl/src/agent/__tests__/systools.test.ts @@ -26,11 +26,13 @@ describe("Agent systools", () => { const kernl = new Kernl(); kernl.register(agent); - expect(agent.systools.length).toBe(1); + // memory + sleep toolkits + expect(agent.systools.length).toBe(2); expect(agent.systools[0].id).toBe("sys.memory"); + expect(agent.systools[1].id).toBe("sys.sleep"); }); - it("has no systools when memory not configured", () => { + it("has only sleep toolkit when memory not configured", () => { const agent = new Agent({ id: "test-agent", name: "Test", @@ -41,10 +43,12 @@ describe("Agent systools", () => { const kernl = new Kernl(); kernl.register(agent); - expect(agent.systools.length).toBe(0); + // sleep is always registered + expect(agent.systools.length).toBe(1); + expect(agent.systools[0].id).toBe("sys.sleep"); }); - it("has no systools when memory.enabled is false", () => { + it("has only sleep toolkit when memory.enabled is false", () => { const agent = new Agent({ id: "test-agent", name: "Test", @@ -56,7 +60,9 @@ describe("Agent systools", () => { const kernl = new Kernl(); kernl.register(agent); - expect(agent.systools.length).toBe(0); + // sleep is always registered + expect(agent.systools.length).toBe(1); + expect(agent.systools[0].id).toBe("sys.sleep"); }); it("can retrieve memory tools via agent.tool()", () => { @@ -114,12 +120,13 @@ describe("Agent systools", () => { const ctx = new Context("test"); const tools = await agent.tools(ctx); - // Memory tools should be first (from systools) - // Order: list, create, update, search + // Memory tools should be first (from systools), then sleep tools + // Order: memory tools (list, create, update, search), then sleep tools (wait_until) expect(tools[0].id).toBe("list_memories"); expect(tools[1].id).toBe("create_memory"); expect(tools[2].id).toBe("update_memory"); expect(tools[3].id).toBe("search_memories"); + expect(tools[4].id).toBe("wait_until"); }); }); @@ -147,4 +154,69 @@ describe("Agent systools", () => { expect(agent.memory).toEqual({ enabled: true }); }); }); + + describe("sleep toolkit", () => { + it("sleep toolkit is always registered", () => { + const agent = new Agent({ + id: "test-agent", + name: "Test", + instructions: "Test", + model, + }); + + const kernl = new Kernl(); + kernl.register(agent); + + const sleepToolkit = agent.systools.find((t) => t.id === "sys.sleep"); + expect(sleepToolkit).toBeDefined(); + }); + + it("can retrieve wait_until tool via agent.tool()", () => { + const agent = new Agent({ + id: "test-agent", + name: "Test", + instructions: "Test", + model, + }); + + const kernl = new Kernl(); + kernl.register(agent); + + expect(agent.tool("wait_until")).toBeDefined(); + }); + + it("includes wait_until in agent.tools() output", async () => { + const agent = new Agent({ + id: "test-agent", + name: "Test", + instructions: "Test", + model, + }); + + const kernl = new Kernl(); + kernl.register(agent); + + const ctx = new Context("test"); + const tools = await agent.tools(ctx); + const ids = tools.map((t) => t.id); + + expect(ids).toContain("wait_until"); + }); + + it("wait_until tool has correct description", () => { + const agent = new Agent({ + id: "test-agent", + name: "Test", + instructions: "Test", + model, + }); + + const kernl = new Kernl(); + kernl.register(agent); + + const tool = agent.tool("wait_until"); + expect(tool).toBeDefined(); + expect(tool!.id).toBe("wait_until"); + }); + }); }); diff --git a/packages/kernl/src/agent/types.ts b/packages/kernl/src/agent/types.ts index fb392939..8221a73e 100644 --- a/packages/kernl/src/agent/types.ts +++ b/packages/kernl/src/agent/types.ts @@ -1,3 +1,7 @@ +/** + * /packages/kernl/src/agent/types.ts + */ + import { type ZodType } from "zod"; import { diff --git a/packages/kernl/src/context.ts b/packages/kernl/src/context.ts index edecab7d..0da468f7 100644 --- a/packages/kernl/src/context.ts +++ b/packages/kernl/src/context.ts @@ -1,3 +1,7 @@ +/** + * /packages/kernl/src/context.ts + */ + import type { Agent } from "@/agent"; /** @@ -26,6 +30,12 @@ export class Context { */ agent?: Agent; + /** + * The thread ID for the current execution. + * Set by the thread during tool execution. + */ + threadId?: string; + // ---------------------- // TEMPORARY: Tool approval tracking until actions system is refined // ---------------------- diff --git a/packages/kernl/src/index.ts b/packages/kernl/src/index.ts index 109df52f..cb8a31b9 100644 --- a/packages/kernl/src/index.ts +++ b/packages/kernl/src/index.ts @@ -1,3 +1,6 @@ +/** + * /package/kernl/src/index.ts +*/ export { Kernl } from "./kernl"; export type { KernlOptions, @@ -79,3 +82,15 @@ export type { MemoryByte, MemoryByteCodec, } from "./memory"; + +// --- wakeups --- +export type { + WakeupStore, + NewScheduledWakeup, + ScheduledWakeup, + ScheduledWakeupUpdate, +} from "./wakeup"; + +// --- scheduler --- +export { WakeupScheduler } from "./scheduler"; +export type { WakeupSchedulerOptions, WakeupSchedulerState } from "./scheduler"; diff --git a/packages/kernl/src/kernl/index.ts b/packages/kernl/src/kernl/index.ts index cd050863..8f2bfa1e 100644 --- a/packages/kernl/src/kernl/index.ts +++ b/packages/kernl/src/kernl/index.ts @@ -1,3 +1,6 @@ +/** + * /packages/kernl/src/kernl/index.ts + */ export { Kernl } from "./kernl"; export type { KernlOptions, diff --git a/packages/kernl/src/kernl/kernl.ts b/packages/kernl/src/kernl/kernl.ts index cc00600f..b55affb2 100644 --- a/packages/kernl/src/kernl/kernl.ts +++ b/packages/kernl/src/kernl/kernl.ts @@ -1,3 +1,6 @@ +/** + * /packages/kernl/src/kernl/kernl.ts + */ import type { LanguageModel } from "@kernl-sdk/protocol"; import { resolveEmbeddingModel } from "@kernl-sdk/retrieval"; @@ -15,6 +18,7 @@ import { MemoryIndexHandle, buildMemoryIndexSchema, } from "@/memory"; +import { WakeupScheduler } from "@/scheduler"; import type { ThreadExecuteResult, ThreadStreamEvent } from "@/thread/types"; import type { AgentOutputType } from "@/agent/types"; @@ -38,6 +42,7 @@ export class Kernl extends KernlHooks { readonly threads: RThreads; readonly agents: RAgents; readonly memories: Memory; + readonly scheduler: WakeupScheduler | null; constructor(options: KernlOptions = {}) { super(); @@ -73,6 +78,15 @@ export class Kernl extends KernlHooks { : undefined, encoder, }); + + // initialize scheduler + if (options.scheduler) { + const schedulerOpts = + typeof options.scheduler === "boolean" ? {} : options.scheduler; + this.scheduler = new WakeupScheduler(this, schedulerOpts); + } else { + this.scheduler = null; + } } /** diff --git a/packages/kernl/src/kernl/types.ts b/packages/kernl/src/kernl/types.ts index 44972cc3..3bf252d0 100644 --- a/packages/kernl/src/kernl/types.ts +++ b/packages/kernl/src/kernl/types.ts @@ -1,8 +1,12 @@ +/** + * /packages/kernl/src/kernl/types.ts + */ import { LanguageModel, EmbeddingModel } from "@kernl-sdk/protocol"; import { SearchIndex } from "@kernl-sdk/retrieval"; import { Agent } from "@/agent"; import { KernlStorage } from "@/storage"; +import type { WakeupSchedulerOptions } from "@/scheduler/types"; /** * Storage configuration for Kernl. @@ -81,6 +85,18 @@ export interface KernlOptions { * Memory system configuration. */ memory?: MemoryOptions; + + /** + * Scheduler configuration for wakeup polling. + * + * - If `true`, creates a scheduler with default options (30s interval, batch size 10) + * - If an object, creates a scheduler with the provided options + * - If `false` or omitted, no scheduler is created + * + * The scheduler polls for due wakeups and resumes sleeping threads. + * Call `kernl.scheduler.start()` to begin polling. + */ + scheduler?: boolean | WakeupSchedulerOptions; } /** diff --git a/packages/kernl/src/memory/codecs/domain.ts b/packages/kernl/src/memory/codecs/domain.ts index 24207ebc..0caf7a40 100644 --- a/packages/kernl/src/memory/codecs/domain.ts +++ b/packages/kernl/src/memory/codecs/domain.ts @@ -1,5 +1,6 @@ /** * Domain-level memory codecs. + * /packages/kernl/src/memory/codecs/domain.ts * * Codecs for transforming between memory domain types and search/index formats. */ diff --git a/packages/kernl/src/memory/codecs/identity.ts b/packages/kernl/src/memory/codecs/identity.ts index d9c6be66..2e4ba738 100644 --- a/packages/kernl/src/memory/codecs/identity.ts +++ b/packages/kernl/src/memory/codecs/identity.ts @@ -1,5 +1,6 @@ /** * Identity codecs - pass through unchanged. + * /packages/kernl/src/memory/codecs/identity.ts * * Used for backends that support the full IndexMemoryRecord schema natively. */ diff --git a/packages/kernl/src/memory/codecs/index.ts b/packages/kernl/src/memory/codecs/index.ts index df4bd6a1..02792688 100644 --- a/packages/kernl/src/memory/codecs/index.ts +++ b/packages/kernl/src/memory/codecs/index.ts @@ -1,5 +1,6 @@ /** * Memory codecs. + * /packages/kernl/src/memory/codecs/index.ts * * Re-exports all memory codecs: * - Domain codecs (MEMORY_FILTER, PATCH_CODEC, recordCodec) diff --git a/packages/kernl/src/memory/codecs/tpuf.ts b/packages/kernl/src/memory/codecs/tpuf.ts index fb7d4c0b..d3975286 100644 --- a/packages/kernl/src/memory/codecs/tpuf.ts +++ b/packages/kernl/src/memory/codecs/tpuf.ts @@ -1,5 +1,6 @@ /** * Turbopuffer backend codecs. + * /packages/kernl/src/memory/codecs/tpuf.ts * * Turbopuffer constraints: * - Exactly one ANN vector field named "vector" per namespace. diff --git a/packages/kernl/src/memory/encoder.ts b/packages/kernl/src/memory/encoder.ts index 35398933..736b170e 100644 --- a/packages/kernl/src/memory/encoder.ts +++ b/packages/kernl/src/memory/encoder.ts @@ -1,5 +1,6 @@ /** * MemoryByte encoder - converts MemoryByte to IndexableByte with embeddings. + * /packages/kernl/src/memory/encoder.ts */ import type { EmbeddingModel, JSONObject } from "@kernl-sdk/protocol"; diff --git a/packages/kernl/src/memory/handle.ts b/packages/kernl/src/memory/handle.ts index 765ba3aa..d08e9b60 100644 --- a/packages/kernl/src/memory/handle.ts +++ b/packages/kernl/src/memory/handle.ts @@ -1,5 +1,6 @@ /** * Memory index handle with lazy initialization. + * /packages/kernl/src/memory/handle.ts */ import { diff --git a/packages/kernl/src/memory/index.ts b/packages/kernl/src/memory/index.ts index f2ec7707..48c75509 100644 --- a/packages/kernl/src/memory/index.ts +++ b/packages/kernl/src/memory/index.ts @@ -1,5 +1,6 @@ /** * Memory module. + * /packages/kernl/src/memory/index.ts */ export { Memory } from "./memory"; diff --git a/packages/kernl/src/memory/indexes.ts b/packages/kernl/src/memory/indexes.ts index aafd0503..0004a765 100644 --- a/packages/kernl/src/memory/indexes.ts +++ b/packages/kernl/src/memory/indexes.ts @@ -1,5 +1,6 @@ /** * Memory index interfaces. + * /packages/kernl/src/memory/indexes.ts * * Indexes are projections of the primary store (DB) that enable * specialized query patterns (vector search, graph traversal, archival). diff --git a/packages/kernl/src/memory/memory.ts b/packages/kernl/src/memory/memory.ts index d29f2783..a4ab996b 100644 --- a/packages/kernl/src/memory/memory.ts +++ b/packages/kernl/src/memory/memory.ts @@ -20,6 +20,7 @@ import { MEMORY_FILTER, PATCH_CODEC, recordCodec } from "./codecs"; /** * Memory is the primary memory abstraction for agents. + * /packages/kernl/src/memory/memory.ts * * Sits above storage/index layers + owns cognitive policy, eviction/TTL, consolidation. * diff --git a/packages/kernl/src/memory/schema.ts b/packages/kernl/src/memory/schema.ts index 4308589d..e451d9f9 100644 --- a/packages/kernl/src/memory/schema.ts +++ b/packages/kernl/src/memory/schema.ts @@ -1,5 +1,6 @@ /** * Memory index schema builder. + * /packages/kernl/src/memory/schema.ts */ import type { FieldSchema } from "@kernl-sdk/retrieval"; diff --git a/packages/kernl/src/memory/store.ts b/packages/kernl/src/memory/store.ts index 62466ec8..daf6583c 100644 --- a/packages/kernl/src/memory/store.ts +++ b/packages/kernl/src/memory/store.ts @@ -1,5 +1,6 @@ /** * Memory store interface. + * /packages/kernl/src/memory/store.ts */ import type { diff --git a/packages/kernl/src/memory/types.ts b/packages/kernl/src/memory/types.ts index 3f573afa..137101ed 100644 --- a/packages/kernl/src/memory/types.ts +++ b/packages/kernl/src/memory/types.ts @@ -1,5 +1,6 @@ /** * Memory types. + * /packages/kernl/src/memory/types.ts */ import type { JSONObject } from "@kernl-sdk/protocol"; diff --git a/packages/kernl/src/scheduler/index.ts b/packages/kernl/src/scheduler/index.ts new file mode 100644 index 00000000..90a10254 --- /dev/null +++ b/packages/kernl/src/scheduler/index.ts @@ -0,0 +1,7 @@ +/** + * Scheduler Module + * /packages/kernl/src/scheduler/index.ts + */ + +export { WakeupScheduler } from "./scheduler"; +export type { WakeupSchedulerOptions, WakeupSchedulerState } from "./types"; diff --git a/packages/kernl/src/scheduler/scheduler.ts b/packages/kernl/src/scheduler/scheduler.ts new file mode 100644 index 00000000..aa647ea7 --- /dev/null +++ b/packages/kernl/src/scheduler/scheduler.ts @@ -0,0 +1,196 @@ +/** + * Wakeup Scheduler + * /packages/kernl/src/scheduler/scheduler.ts + * + * Polls the wakeup store for due wakeups and resumes threads. + * Can be attached to a Kernl instance or run standalone. + */ + +import type { Kernl } from "@/kernl"; +import type { ScheduledWakeup } from "@/wakeup"; +import { getLogger } from "@/lib/logger"; +import { message } from "@kernl-sdk/protocol"; + +import type { WakeupSchedulerOptions, WakeupSchedulerState } from "./types"; + +const logger = getLogger("kernl:scheduler"); + +const DEFAULT_INTERVAL_MS = 30_000; // 30 seconds +const DEFAULT_BATCH_SIZE = 10; + +export class WakeupScheduler { + private readonly kernl: Kernl; + private readonly intervalMs: number; + private readonly batchSize: number; + + private timer: ReturnType | null = null; + private polling = false; // guard against overlapping polls + + private _state: WakeupSchedulerState = { + running: false, + processed: 0, + failed: 0, + lastPollAt: null, + }; + + constructor(kernl: Kernl, options: WakeupSchedulerOptions = {}) { + this.kernl = kernl; + this.intervalMs = options.intervalMs ?? DEFAULT_INTERVAL_MS; + this.batchSize = options.batchSize ?? DEFAULT_BATCH_SIZE; + + if (options.autoStart) { + this.start(); + } + } + + /** + * Current scheduler state (read-only snapshot). + */ + get state(): Readonly { + return { ...this._state }; + } + + /** + * Start the polling loop. + * No-op if already running. + */ + start(): void { + if (this._state.running) { + logger.warn("scheduler already running"); + return; + } + + logger.info({ intervalMs: this.intervalMs }, "scheduler starting"); + this._state.running = true; + + // Run first poll immediately, then on interval + this.poll(); + this.timer = setInterval(() => this.poll(), this.intervalMs); + } + + /** + * Stop the polling loop. + * In-flight poll will complete but no new polls will start. + */ + stop(): void { + if (!this._state.running) { + logger.warn("scheduler not running"); + return; + } + + logger.info("scheduler stopping"); + this._state.running = false; + + if (this.timer) { + clearInterval(this.timer); + this.timer = null; + } + } + + /** + * Execute a single poll iteration. + * Claims due wakeups and resumes their threads. + * + * Can be called manually for testing or one-off processing. + */ + async poll(): Promise { + // Guard against overlapping polls (previous poll still running) + if (this.polling) { + logger.debug("skipping poll, previous poll still running"); + return; + } + + this.polling = true; + this._state.lastPollAt = Date.now(); + + try { + const wakeupStore = this.kernl.storage?.wakeups; + if (!wakeupStore) { + logger.warn("no wakeup store configured, skipping poll"); + return; + } + + const nowMs = Date.now(); + const claimed = await wakeupStore.claimDue(nowMs, this.batchSize); + + if (claimed.length === 0) { + logger.debug("no due wakeups found"); + return; + } + + logger.info({ count: claimed.length }, "claimed due wakeups"); + + // Process each wakeup + await Promise.all(claimed.map((wakeup) => this.processWakeup(wakeup))); + } catch (err) { + logger.error({ err }, "poll failed"); + } finally { + this.polling = false; + } + } + + /** + * Process a single wakeup: resume the thread and mark wakeup as woken. + */ + private async processWakeup(wakeup: ScheduledWakeup): Promise { + const { id, threadId } = wakeup; + const wakeupStore = this.kernl.storage?.wakeups; + const threadStore = this.kernl.storage?.threads; + + if (!wakeupStore || !threadStore) { + logger.error({ wakeupId: id }, "storage not configured"); + return; + } + + try { + logger.debug({ wakeupId: id, threadId }, "processing wakeup"); + + // Load thread with history + const thread = await threadStore.get(threadId, { history: true }); + + if (!thread) { + throw new Error(`Thread ${threadId} not found`); + } + + // Append a user message to signal wakeup and continuation + const wakeupMessage = wakeup.reason + ? `[Sleep period complete - reason: ${wakeup.reason}] You have woken up. Continue with your task.` + : "[Sleep period complete] You have woken up. Continue with your task."; + + thread.append( + message({ + role: "user", + text: wakeupMessage, + }), + ); + + // Resume thread execution (thread.execute continues from where it left off) + await this.kernl.schedule(thread); + + // Mark wakeup as complete + await wakeupStore.update(id, { + id, + woken: true, + updatedAt: Date.now(), + }); + + this._state.processed++; + logger.info({ wakeupId: id, threadId }, "wakeup processed successfully"); + } catch (err) { + this._state.failed++; + const errorMsg = err instanceof Error ? err.message : String(err); + logger.error({ wakeupId: id, threadId, err: errorMsg }, "wakeup processing failed"); + + // Record error on the wakeup record + try { + await wakeupStore.update(id, { + id, + error: errorMsg, + updatedAt: Date.now(), + }); + } catch (updateErr) { + logger.error({ wakeupId: id, err: updateErr }, "failed to update wakeup with error"); + } + } + } +} diff --git a/packages/kernl/src/scheduler/types.ts b/packages/kernl/src/scheduler/types.ts new file mode 100644 index 00000000..68edf835 --- /dev/null +++ b/packages/kernl/src/scheduler/types.ts @@ -0,0 +1,38 @@ +/** + * Scheduler Types + * /packages/kernl/src/scheduler/types.ts + */ + +export interface WakeupSchedulerOptions { + /** + * Polling interval in milliseconds. + * @default 30000 (30 seconds) + */ + intervalMs?: number; + + /** + * Maximum number of wakeups to claim per poll. + * @default 10 + */ + batchSize?: number; + + /** + * Whether to start polling automatically when the scheduler is created. + * @default false + */ + autoStart?: boolean; +} + +export interface WakeupSchedulerState { + /** Whether the scheduler is currently polling */ + running: boolean; + + /** Number of wakeups processed since start */ + processed: number; + + /** Number of wakeups that failed to process */ + failed: number; + + /** Timestamp of last poll (epoch ms), null if never polled */ + lastPollAt: number | null; +} diff --git a/packages/kernl/src/storage/base.ts b/packages/kernl/src/storage/base.ts index cf42ad63..90e1736a 100644 --- a/packages/kernl/src/storage/base.ts +++ b/packages/kernl/src/storage/base.ts @@ -1,10 +1,12 @@ /** * Core storage contracts. + * /packages/kernl/src/storage/base.ts */ import type { AgentRegistry, ModelRegistry } from "@/kernl/types"; import type { ThreadStore } from "./thread"; -import type { MemoryStore } from "@/memory/store"; +import type { MemoryStore } from "@/memory" +import type { WakeupStore } from "@/wakeup"; /** * The main storage interface for Kernl. @@ -22,6 +24,11 @@ export interface KernlStorage { */ memories: MemoryStore; + /** + * Wakeup store - manages wakeup records for agents. + */ + wakeups: WakeupStore; + // tasks: TaskStore; // traces: TraceStore; diff --git a/packages/kernl/src/storage/in-memory.ts b/packages/kernl/src/storage/in-memory.ts index 6cc0b3ae..a5970f33 100644 --- a/packages/kernl/src/storage/in-memory.ts +++ b/packages/kernl/src/storage/in-memory.ts @@ -1,5 +1,6 @@ /** * In-memory storage implementation for Kernl. + * /packages/kernl/src/storage/in-memory.ts * * Pure domain-level - no codecs, schemas, or DB records. * @@ -33,6 +34,13 @@ import type { MemoryListOptions, MemoryFilter, } from "@/memory"; +import type { + NewScheduledWakeup, + ScheduledWakeup, + ScheduledWakeupUpdate, + WakeupStore, +} from "@/wakeup"; + /** * In-memory storage implementation. @@ -40,10 +48,12 @@ import type { export class InMemoryStorage implements KernlStorage { threads: InMemoryThreadStore; memories: InMemoryMemoryStore; + wakeups: InMemoryWakeupStore; constructor() { this.threads = new InMemoryThreadStore(); this.memories = new InMemoryMemoryStore(); + this.wakeups = new InMemoryWakeupStore(); } bind(registries: { agents: AgentRegistry; models: ModelRegistry }): void { @@ -542,3 +552,108 @@ export class InMemoryMemoryStore implements MemoryStore { }); } } + + +/** + * In-memory wakeup store implementation + */ +export class InMemoryWakeupStore implements WakeupStore { + private wakeups = new Map(); + + async create(input: NewScheduledWakeup): Promise { + const now = Date.now(); + const nowS = Math.floor(now / 1000); + + // Support both new (sleepFor) and legacy (runAt) fields. + const sleepForS = + typeof input.sleepFor === "number" + ? Math.max(0, Math.floor(input.sleepFor)) + : typeof input.runAt === "number" + ? Math.max(0, Math.floor(input.runAt / 1000) - nowS) + : 0; + + const runAt = nowS * 1000 + sleepForS * 1000; + + const record: ScheduledWakeup = { + id: input.id, + threadId: input.threadId, + runAt, + sleepFor: sleepForS, + reason: input.reason ?? null, + woken: false, + claimedAt: null, + createdAt: now, + updatedAt: now, + error: null, + }; + + this.wakeups.set(record.id, record); + return record; + } + + async get(id: string): Promise { + return this.wakeups.get(id) ?? null; + } + + async update( + id: string, + patch: ScheduledWakeupUpdate, + ): Promise { + const existing = this.wakeups.get(id); + if (!existing) { + throw new Error(`Wakeup with id ${id} not found`); + } + + const updated: ScheduledWakeup = { + ...existing, + runAt: patch.runAt ?? existing.runAt, + reason: patch.reason ?? existing.reason, + woken: patch.woken ?? existing.woken, + claimedAt: + patch.claimedAt !== undefined ? patch.claimedAt : existing.claimedAt, + error: patch.error ?? existing.error, + updatedAt: patch.updatedAt ?? Date.now(), + }; + + this.wakeups.set(id, updated); + return updated; + } + + async delete(id: string): Promise { + this.wakeups.delete(id); + } + + async claimDue( + nowMs: number | bigint, + limit: number, + ): Promise { + const now = typeof nowMs === "bigint" ? Number(nowMs) : nowMs; + + const candidates: ScheduledWakeup[] = []; + + for (const w of this.wakeups.values()) { + if (!w.woken && w.claimedAt == null && w.runAt <= now) { + candidates.push(w); + } + } + + candidates.sort((a, b) => a.runAt - b.runAt); + + const selected = candidates.slice(0, limit); + const updatedSelected: ScheduledWakeup[] = []; + + const claimTime = Date.now(); + + for (const w of selected) { + const updated: ScheduledWakeup = { + ...w, + claimedAt: claimTime, + updatedAt: claimTime, + }; + this.wakeups.set(updated.id, updated); + updatedSelected.push(updated); + } + + return updatedSelected; + } +} diff --git a/packages/kernl/src/storage/index.ts b/packages/kernl/src/storage/index.ts index e4aecdba..2481f1ec 100644 --- a/packages/kernl/src/storage/index.ts +++ b/packages/kernl/src/storage/index.ts @@ -1,5 +1,6 @@ /** * Storage contracts for Kernl. + * /packages/kernl/src/storage/index.ts * * Core owns these interfaces; storage packages implement them. * (must be defined here to avoid circular deps) diff --git a/packages/kernl/src/storage/thread.ts b/packages/kernl/src/storage/thread.ts index 566738db..fa94b77c 100644 --- a/packages/kernl/src/storage/thread.ts +++ b/packages/kernl/src/storage/thread.ts @@ -1,5 +1,6 @@ /** * Thread storage contracts. + * /packages/kernl/src/storage/thread.ts */ import type { Thread } from "@/thread"; diff --git a/packages/kernl/src/thread/index.ts b/packages/kernl/src/thread/index.ts index 45388ec2..9ac6996f 100644 --- a/packages/kernl/src/thread/index.ts +++ b/packages/kernl/src/thread/index.ts @@ -1 +1,5 @@ +/** + * /packages/kernl/src/thread/index.ts + */ + export { Thread } from "./thread"; diff --git a/packages/kernl/src/thread/thread.ts b/packages/kernl/src/thread/thread.ts index 21226e75..041e0607 100644 --- a/packages/kernl/src/thread/thread.ts +++ b/packages/kernl/src/thread/thread.ts @@ -1,3 +1,6 @@ +/** + * /packages/kernl/src/thread/thread.ts + */ import assert from "assert"; import { ZodType } from "zod"; import * as z from "zod"; @@ -19,6 +22,7 @@ import { LanguageModel, LanguageModelItem, LanguageModelRequest, + INTERRUPTIBLE } from "@kernl-sdk/protocol"; import { randomID, filter } from "@kernl-sdk/shared/lib"; @@ -141,6 +145,7 @@ export class Thread< } } + // MARK: Execute /** * Blocking execution - runs until terminal state or interruption */ @@ -186,12 +191,17 @@ export class Thread< } catch (err) { throw err; } finally { - this.state = STOPPED; + // Preserve INTERRUPTIBLE state (set by sleep/approval flows). + // Only transition to STOPPED if we're still RUNNING. + if (this.state === RUNNING) { + this.state = STOPPED; + } this.abort = undefined; - await this.checkpoint(); /* c4: final checkpoint - persist STOPPED state */ + await this.checkpoint(); /* c4: final checkpoint - persist final state */ } } + // MARK: _execute main loop /** * Main execution loop - always yields events, callers can propagate or discard. * @@ -251,13 +261,11 @@ export class Thread< await this.checkpoint(); /* c3: tick complete */ if (pendingApprovals.length > 0) { - // publish a batch approval request containing all of them - // - // const reqid = randomID(); - // this.kernl.publish(channel, approvalRequest); - // - // const filter = { reqid } - // await wait_event(Action.ApprovalResponse, filter); + // Thread halts when actions require approval (e.g., sleep). + // External scheduler will resume the thread later. + // Set state to INTERRUPTIBLE so callers know the thread is sleeping, not stopped. + this.state = INTERRUPTIBLE; + return; } } } @@ -402,19 +410,47 @@ export class Thread< const actions: ThreadEventInner[] = []; const pendingApprovals: ToolCall[] = []; - // (TODO): clean this - approval tracking should be handled differently for (const e of toolEvents) { - if ( - e.kind === "tool-result" && - (e.state as any) === "requires_approval" // (TODO): fix this - ) { - // find the original tool call for this pending approval - const call = intentions.toolCalls.find((c) => c.callId === e.callId); - call && pendingApprovals.push(call); - } else { - actions.push(e); + actions.push(e); // always record the tool call in history + + // Only system tools are allowed to affect control flow. + if (e.kind === "tool-result" && this.agent.isSysTool(e.toolId)) { + // 1) Future-proof: generic approval-based interrupts (other sys tools) + // If you later add real approval flows that still use INTERRUPTIBLE, + // this branch continues to work. + if (e.state === INTERRUPTIBLE) { + const call = intentions.toolCalls.find((c) => c.callId === e.callId); + if (call) pendingApprovals.push(call); + continue; + } + + // 2) Sleep tool: state is COMPLETED, but semantically it means + // "alarm set, thread should pause until wakeup." + if (e.toolId === "wait_until") { + const call = intentions.toolCalls.find((c) => c.callId === e.callId); + if (call) pendingApprovals.push(call); + continue; + } } } + // const actions: ThreadEventInner[] = []; + // const pendingApprovals: ToolCall[] = []; + + // // (TODO): clean this - approval tracking should be handled differently + // for (const e of toolEvents) { + // // actions.push(e); + // if ( + // e.kind === "tool-result" && + // e.state === INTERRUPTIBLE && + // this.agent.isSysTool(e.toolId) + // ) { + // // find the original tool call for this pending approval + // const call = intentions.toolCalls.find((c) => c.callId === e.callId); + // call && pendingApprovals.push(call); + // } else { + // actions.push(e); + // } + // } return { actions: actions, @@ -446,7 +482,14 @@ export class Thread< // is refined const ctx = new Context(this.namespace, this.context.context); ctx.agent = this.agent; - ctx.approve(call.callId); // mark this call as approved + ctx.threadId = this.tid; + + // Don't pre-approve system tools - they use requiresApproval to signal + // INTERRUPTIBLE state (e.g., sleep tool halts the execution loop) + if (!this.agent.isSysTool(call.toolId)) { + ctx.approve(call.callId); + } + const res = await tool.invoke(ctx, call.arguments, call.callId); return { diff --git a/packages/kernl/src/thread/types.ts b/packages/kernl/src/thread/types.ts index 2f898c51..faeedd33 100644 --- a/packages/kernl/src/thread/types.ts +++ b/packages/kernl/src/thread/types.ts @@ -1,3 +1,6 @@ +/** + * /packages/kernl/src/thread/types.ts + */ import { ToolCall, LanguageModel, diff --git a/packages/kernl/src/thread/utils.ts b/packages/kernl/src/thread/utils.ts index 45b53e5a..707c6e4c 100644 --- a/packages/kernl/src/thread/utils.ts +++ b/packages/kernl/src/thread/utils.ts @@ -1,3 +1,6 @@ +/** + * /packages/kernl/src/thread/utils.ts + */ import { ZodType } from "zod"; import type { ResolvedAgentResponse } from "@/guardrail"; diff --git a/packages/kernl/src/tool/index.ts b/packages/kernl/src/tool/index.ts index 2dba7f93..a2844fde 100644 --- a/packages/kernl/src/tool/index.ts +++ b/packages/kernl/src/tool/index.ts @@ -1,3 +1,7 @@ +/** + * /packages/kernl/src/tool/index.ts + */ + export { BaseTool, FunctionTool, HostedTool, tool } from "./tool"; export { BaseToolkit, Toolkit, FunctionToolkit, MCPToolkit } from "./toolkit"; export type { @@ -10,4 +14,4 @@ export type { } from "./types"; // --- system toolkits --- -export { memory } from "./sys"; +export { memory, sleep } from "./sys"; diff --git a/packages/kernl/src/tool/sys/index.ts b/packages/kernl/src/tool/sys/index.ts index a2da55a2..32421576 100644 --- a/packages/kernl/src/tool/sys/index.ts +++ b/packages/kernl/src/tool/sys/index.ts @@ -1,7 +1,9 @@ /** * System toolkits. + * /packages/kernl/src/tool/sys/index.ts * * These are internal toolkits that can be enabled via agent config flags. */ export { memory } from "./memory"; +export { sleep } from "./sleep"; \ No newline at end of file diff --git a/packages/kernl/src/tool/sys/memory.ts b/packages/kernl/src/tool/sys/memory.ts index 80b7c991..9693cb16 100644 --- a/packages/kernl/src/tool/sys/memory.ts +++ b/packages/kernl/src/tool/sys/memory.ts @@ -1,5 +1,6 @@ /** * Memory system toolkit. + * /packages/kernl/src/tool/sys/memory.ts * * Provides tools for agents to store and retrieve memories. * Enabled via `memory: true` in agent config. diff --git a/packages/kernl/src/tool/sys/sleep.ts b/packages/kernl/src/tool/sys/sleep.ts new file mode 100644 index 00000000..4381bad3 --- /dev/null +++ b/packages/kernl/src/tool/sys/sleep.ts @@ -0,0 +1,90 @@ +/** + * Sleep system toolkit. + * /packages/kernl/src/tool/sys/sleep.ts + */ + +import assert from "assert"; +import { z } from "zod"; + +import { randomID } from "@kernl-sdk/shared/lib"; + +import { tool } from "../tool"; +import { Toolkit } from "../toolkit"; + +const WaitParamsSchema = z.object({ + delay_s: z + .number() + .int() + .nonnegative() + .describe( + "How many seconds from now to wait before resuming this thread.", + ), + reason: z + .string() + .optional() + .describe("Optional human-readable reason for the sleep."), +}); + +const wait = tool({ + id: "wait_until", + description: + "Pause this agent until a future time. The thread will be resumed automatically.", + mode: "async" as const, + parameters: WaitParamsSchema, + + // NOTE: we rely on execute(), not requiresApproval, so the tool + // ends up with state=COMPLETED (green) in the UI. + // + // This function: + // 1) Schedules the wakeup in storage. + // 2) Returns metadata so the UI can render a 'Sleeping' success state. + execute: async (ctx, params) => { + assert(ctx.agent, "ctx.agent is required for sleep tools"); + assert(ctx.threadId, "ctx.threadId is required for sleep tools"); + + const { delay_s, reason } = params as z.infer; + + const agent: any = ctx.agent; + + if (!agent.kernl || !agent.kernl.storage) { + throw new Error("Agent is not bound to Kernl storage."); + } + + const wakeupStore: any = agent.kernl.storage.wakeups; + if (!wakeupStore || typeof wakeupStore.create !== "function") { + throw new Error("Wakeup store is not configured on Kernl storage."); + } + + const now_s = Math.floor(Date.now() / 1000); + + // For UI display purposes + const targetRunAt_s = now_s + delay_s; + const targetRunAt_ms = targetRunAt_s * 1000; + + await wakeupStore.create({ + id: `wkp_${randomID()}`, + threadId: ctx.threadId, + sleepFor: delay_s, + reason: reason ?? null, + }); + + // This object becomes the `result` in the tool-result event. + // The UI can key off `status: "sleeping"` to show a green "Sleeping" badge. + return { + status: "sleeping", + scheduled: true, + run_at_s: targetRunAt_s, + run_at_ms: targetRunAt_ms, + message: + "Wakeup scheduled; the thread will resume when the wakeup is processed.", + }; + }, +}); + +// --- Toolkit --- + +export const sleep = new Toolkit({ + id: "sys.sleep", + description: "Tools for pausing agents and scheduling wakeups.", + tools: [wait], +}); diff --git a/packages/kernl/src/tool/tool.ts b/packages/kernl/src/tool/tool.ts index dfd5bf23..83909224 100644 --- a/packages/kernl/src/tool/tool.ts +++ b/packages/kernl/src/tool/tool.ts @@ -1,3 +1,7 @@ +/** + * /packages/kernl/src/tool/tool.ts + */ + import { z } from "zod"; import { Context, UnknownContext } from "@/context"; diff --git a/packages/kernl/src/tool/toolkit.ts b/packages/kernl/src/tool/toolkit.ts index 6eed0c81..de1c5854 100644 --- a/packages/kernl/src/tool/toolkit.ts +++ b/packages/kernl/src/tool/toolkit.ts @@ -1,3 +1,7 @@ +/** + * /packages/kernl/src/tool/toolkit.ts + */ + import type { Agent } from "@/agent"; import type { Context, UnknownContext } from "@/context"; diff --git a/packages/kernl/src/tool/types.ts b/packages/kernl/src/tool/types.ts index 475b87e7..bbf99d40 100644 --- a/packages/kernl/src/tool/types.ts +++ b/packages/kernl/src/tool/types.ts @@ -1,3 +1,7 @@ +/** + * /packages/kernl/src/tool/types.ts + */ + import { z, type ZodType } from "zod"; import { Agent } from "@/agent"; diff --git a/packages/kernl/src/wakeup/index.ts b/packages/kernl/src/wakeup/index.ts new file mode 100644 index 00000000..b4de287f --- /dev/null +++ b/packages/kernl/src/wakeup/index.ts @@ -0,0 +1,12 @@ +/** + * Wakeup public API. + * /packages/kernl/src/wakeup/index.ts + */ + +export type { + NewScheduledWakeup, + ScheduledWakeup, + ScheduledWakeupUpdate, +} from "./types"; + +export type { WakeupStore } from "./store"; diff --git a/packages/kernl/src/wakeup/store.ts b/packages/kernl/src/wakeup/store.ts new file mode 100644 index 00000000..9cecf20f --- /dev/null +++ b/packages/kernl/src/wakeup/store.ts @@ -0,0 +1,51 @@ +/** + * Wakeup store contract. + * /packages/kernl/src/wakeup/store.ts + */ + +import type { + NewScheduledWakeup, + ScheduledWakeup, + ScheduledWakeupUpdate, +} from "./types"; + +/** + * Persistence contract for scheduled wakeups. + * + * Implementations live in provider packages (e.g. @kernl-sdk/storage/pg). + */ +export interface WakeupStore { + /** + * Create a new scheduled wakeup. + */ + create(wakeup: NewScheduledWakeup): Promise; + + /** + * Get a wakeup by ID. + */ + get(id: string): Promise; + + /** + * Update an existing wakeup. + */ + update(id: string, patch: ScheduledWakeupUpdate): Promise; + + /** + * Delete a wakeup (admin/cleanup). Normal flow should mark `woken = true` + * instead of deleting. + */ + delete(id: string): Promise; + + /** + * Atomically claim up to `limit` due wakeups. + * + * A wakeup is "due" when: + * - woken = false + * - claimedAt is null + * - runAt <= nowMs + * + * Implementations must be safe under multiple pollers + * (e.g. using SKIP LOCKED semantics). + */ + claimDue(nowMs: number | bigint, limit: number): Promise; +} diff --git a/packages/kernl/src/wakeup/types.ts b/packages/kernl/src/wakeup/types.ts new file mode 100644 index 00000000..5be32875 --- /dev/null +++ b/packages/kernl/src/wakeup/types.ts @@ -0,0 +1,98 @@ +/** + * Wakeup types. + * /packages/kernl/src/wakeup/types.ts + */ + +export interface NewScheduledWakeup { + /** + * Unique wakeup ID (e.g. ulid/uuid). Generated by Kernl, not by the agent. + */ + id: string; + + /** + * Thread to resume when the wakeup becomes due. + */ + threadId: string; + + /** + * Duration to sleep (seconds). Preferred over `runAt`. + * + * The system will calculate the absolute wakeup time as `created_at + sleepFor`. + */ + sleepFor?: number; + + /** + * (Legacy) When this wakeup becomes due (epoch milliseconds). + * + * For backward compatibility. Prefer using `sleepFor` instead. + * Storage providers can store this as seconds (`run_at_s`) or ms, + * codecs handle conversion. + */ + runAt?: number; + + /** + * Optional free-form reason for observability. + */ + reason?: string | null; +} + +/** + * Persisted wakeup record in domain form. + */ +export interface ScheduledWakeup { + id: string; + threadId: string; + + /** + * Due time in epoch milliseconds (domain-level). + */ + runAt: number; + + /** + * Duration to sleep (seconds). + */ + sleepFor: number; + + reason: string | null; + + /** + * True once the wakeup has been fully consumed (thread resumed / handled). + */ + woken: boolean; + + /** + * When a poller claimed this wakeup (epoch ms). `null` if unclaimed. + */ + claimedAt: number | null; + + /** + * Creation and update timestamps (epoch ms). + */ + createdAt: number; + updatedAt: number; + + /** + * Optional error message if processing/resume failed. + */ + error: string | null; +} + +/** + * Patch payload for updating a wakeup. + */ +export interface ScheduledWakeupUpdate { + id: string; + + runAt?: number; + reason?: string | null; + + woken?: boolean; + claimedAt?: number | null; + error?: string | null; + + /** + * Optional explicit updatedAt (ms). Most PG logic will just set + * this to `Date.now()` if omitted. + */ + updatedAt?: number; +} diff --git a/packages/shared/src/lib/codec.ts b/packages/shared/src/lib/codec.ts index c0b0e8f8..b7367096 100644 --- a/packages/shared/src/lib/codec.ts +++ b/packages/shared/src/lib/codec.ts @@ -2,6 +2,7 @@ import { z, type ZodType } from "zod"; /** * Bidirectional codec for converting between types. + * /packages/shared/src/lib/codec.ts * * @example * ```typescript diff --git a/packages/storage/core/src/index.ts b/packages/storage/core/src/index.ts index f6cc7b88..095cf49c 100644 --- a/packages/storage/core/src/index.ts +++ b/packages/storage/core/src/index.ts @@ -1,5 +1,6 @@ /** * @kernl-sdk/storage - Generic storage abstractions for Kernl + * /packages/storage/core/src/index.ts */ export * from "./base"; @@ -8,3 +9,4 @@ export * from "./memory"; export * from "./serde/thread"; export * from "./serde/memory"; export * from "./table"; +export * from "./wakeup"; diff --git a/packages/storage/core/src/memory/index.ts b/packages/storage/core/src/memory/index.ts index 294ff8ff..a12e4367 100644 --- a/packages/storage/core/src/memory/index.ts +++ b/packages/storage/core/src/memory/index.ts @@ -1,5 +1,6 @@ /** * Memory storage types and schema. + * /packages/storage/core/src/memory/index.ts */ export { diff --git a/packages/storage/core/src/memory/schema.ts b/packages/storage/core/src/memory/schema.ts index 8da61f0e..210f8ca9 100644 --- a/packages/storage/core/src/memory/schema.ts +++ b/packages/storage/core/src/memory/schema.ts @@ -1,5 +1,6 @@ /** * Memory table definition and record schema. + * /packages/storage/core/src/memory/schema.ts */ import { z } from "zod"; diff --git a/packages/storage/core/src/table.ts b/packages/storage/core/src/table.ts index 437fba3d..805088ec 100644 --- a/packages/storage/core/src/table.ts +++ b/packages/storage/core/src/table.ts @@ -1,5 +1,6 @@ /** * Schema name for all kernl tables. + * /packages/storage/core/table.ts */ export const KERNL_SCHEMA_NAME = "kernl"; diff --git a/packages/storage/core/src/wakeup/index.ts b/packages/storage/core/src/wakeup/index.ts new file mode 100644 index 00000000..68c9fcaf --- /dev/null +++ b/packages/storage/core/src/wakeup/index.ts @@ -0,0 +1,10 @@ +/** + * Wakeup storage types and utilities. + * /packages/storage/core/src/wakeup/index.ts + */ + +export { + TABLE_SCHEDULED_WAKEUPS, + ScheduledWakeupRecordSchema, + type ScheduledWakeupRecord, +} from "./schema"; diff --git a/packages/storage/core/src/wakeup/schema.ts b/packages/storage/core/src/wakeup/schema.ts new file mode 100644 index 00000000..6f5d601f --- /dev/null +++ b/packages/storage/core/src/wakeup/schema.ts @@ -0,0 +1,70 @@ +/** + * /packages/storage/core/src/wakeup/schema.ts + * + * First implementation: + * - wakeup_at: epoch seconds when the wakeup becomes due + * - sleep_for: seconds to sleep (requested duration) + * - woken: consumed/completed + * - claimed_at_s: set when a poller claims it to avoid double-processing + */ + +import { z } from "zod"; + +import { text, bigint, boolean, timestamps, defineTable } from "@/table"; +import { TABLE_THREADS } from "@/thread/schema"; + +export const TABLE_SCHEDULED_WAKEUPS = defineTable( + "scheduled_wakeups", + { + id: text().primaryKey(), + + thread_id: text().references(() => TABLE_THREADS.columns.id, { + onDelete: "CASCADE", + }), + + // Requested duration (seconds) + sleep_for: bigint(), + + // Due time (epoch seconds) + wakeup_at: bigint(), + + reason: text().nullable(), + + // Consumed/completed + woken: boolean().default(false), + + // Claimed by a poller (epoch seconds); nullable means unclaimed + claimed_at_s: bigint().nullable(), + + ...timestamps, + + error: text().nullable(), + }, + [ + // Polling query: woken=false AND claimed_at_s IS NULL AND wakeup_at <= now + { kind: "index", columns: ["woken", "wakeup_at"] }, + { kind: "index", columns: ["thread_id"] }, + { kind: "index", columns: ["claimed_at_s"] }, + ], +); + +const epochSeconds = z.coerce.number().int().nonnegative(); + +export const ScheduledWakeupRecordSchema = z.object({ + id: z.string(), + thread_id: z.string(), + + sleep_for: epochSeconds, + wakeup_at: epochSeconds, + reason: z.string().nullable(), + + woken: z.boolean(), + claimed_at_s: epochSeconds.nullable(), + + created_at: epochSeconds, + updated_at: epochSeconds, + + error: z.string().nullable(), +}); + +export type ScheduledWakeupRecord = z.infer; diff --git a/packages/storage/pg/src/index.ts b/packages/storage/pg/src/index.ts index 33cffc89..05ade447 100644 --- a/packages/storage/pg/src/index.ts +++ b/packages/storage/pg/src/index.ts @@ -1,5 +1,6 @@ /** * @kernl/pg - PostgreSQL storage adapter for Kernl + * /packages/storage/pg/src/index.ts */ export { @@ -12,6 +13,7 @@ export { resolveVectorConfig, } from "./storage"; export { PGMemoryStore } from "./memory/store"; +export { PGWakeupStore } from "./wakeup/store"; export { postgres, pgvector, type PostgresConfig } from "./postgres"; export { MIGRATIONS, REQUIRED_SCHEMA_VERSION } from "./migrations"; export { diff --git a/packages/storage/pg/src/memory/sql.ts b/packages/storage/pg/src/memory/sql.ts index 666887eb..08a00e57 100644 --- a/packages/storage/pg/src/memory/sql.ts +++ b/packages/storage/pg/src/memory/sql.ts @@ -1,6 +1,7 @@ /** * Memory SQL conversion codecs. - * + * /packages/storage/pg/src/memory/sql.ts + * * TODO: generalize object -> SQL conversion into a shared utility */ diff --git a/packages/storage/pg/src/memory/store.ts b/packages/storage/pg/src/memory/store.ts index b99bf909..70975b6e 100644 --- a/packages/storage/pg/src/memory/store.ts +++ b/packages/storage/pg/src/memory/store.ts @@ -1,5 +1,6 @@ /** * PostgreSQL Memory store implementation. + * /packages/storage/pg/src/memory/store.ts */ import type { Pool, PoolClient } from "pg"; diff --git a/packages/storage/pg/src/migrations.ts b/packages/storage/pg/src/migrations.ts index 77416325..4e064a26 100644 --- a/packages/storage/pg/src/migrations.ts +++ b/packages/storage/pg/src/migrations.ts @@ -8,6 +8,7 @@ import { TABLE_THREADS, TABLE_THREAD_EVENTS, TABLE_MEMORIES, + TABLE_SCHEDULED_WAKEUPS, } from "@kernl-sdk/storage"; /** @@ -27,6 +28,12 @@ export interface Migration { * List of all migrations in order. */ export const MIGRATIONS: Migration[] = [ + { + id: "000_enable_vector", + async up(ctx) { + await ctx.client.query("CREATE EXTENSION IF NOT EXISTS vector"); + }, + }, { id: "001_threads", async up(ctx) { @@ -40,6 +47,98 @@ export const MIGRATIONS: Migration[] = [ await ctx.createTable(TABLE_MEMORIES); }, }, + { + id: "003_scheduled_wakeups", + async up(ctx) { + await ctx.createTable(TABLE_SCHEDULED_WAKEUPS); + }, + }, + { + id: "004_scheduled_wakeups_v2", + async up(ctx) { + // v1 -> v2: + // - run_at_s -> wakeup_at + // - add sleep_for + // - convert created_at/updated_at from ms -> seconds (if needed) + // - backfill sleep_for for existing rows + const schema = "kernl"; + + const colsRes = await ctx.client.query<{ column_name: string }>( + ` + SELECT column_name + FROM information_schema.columns + WHERE table_schema = $1 AND table_name = 'scheduled_wakeups' + `, + [schema], + ); + + if (colsRes.rows.length === 0) return; + + const cols = new Set(colsRes.rows.map((r) => r.column_name)); + + if (cols.has("run_at_s") && !cols.has("wakeup_at")) { + await ctx.client.query( + `ALTER TABLE ${schema}.scheduled_wakeups RENAME COLUMN run_at_s TO wakeup_at`, + ); + cols.delete("run_at_s"); + cols.add("wakeup_at"); + } + + // If we somehow ended up with both columns, keep wakeup_at and drop the old name. + if (cols.has("run_at_s") && cols.has("wakeup_at")) { + await ctx.client.query( + `ALTER TABLE ${schema}.scheduled_wakeups DROP COLUMN run_at_s`, + ); + cols.delete("run_at_s"); + } + + if (!cols.has("sleep_for")) { + await ctx.client.query( + `ALTER TABLE ${schema}.scheduled_wakeups ADD COLUMN sleep_for bigint`, + ); + cols.add("sleep_for"); + } + + // Convert ms -> seconds only when the values are clearly ms. + await ctx.client.query( + ` + UPDATE ${schema}.scheduled_wakeups + SET created_at = created_at / 1000, + updated_at = updated_at / 1000 + WHERE created_at > 100000000000 OR updated_at > 100000000000; + `, + ); + + // Backfill sleep_for for existing rows. + await ctx.client.query( + ` + UPDATE ${schema}.scheduled_wakeups + SET sleep_for = GREATEST(0, wakeup_at - created_at) + WHERE sleep_for IS NULL; + `, + ); + + // Enforce non-null/default after backfill. + await ctx.client.query( + `ALTER TABLE ${schema}.scheduled_wakeups ALTER COLUMN sleep_for SET DEFAULT 0`, + ); + await ctx.client.query( + `ALTER TABLE ${schema}.scheduled_wakeups ALTER COLUMN sleep_for SET NOT NULL`, + ); + + // Ensure the polling indexes exist (safe even if duplicates already exist under other names). + await ctx.client.query( + `CREATE INDEX IF NOT EXISTS scheduled_wakeups_woken_wakeup_at_idx ON ${schema}.scheduled_wakeups (woken, wakeup_at)`, + ); + await ctx.client.query( + `CREATE INDEX IF NOT EXISTS scheduled_wakeups_thread_id_idx ON ${schema}.scheduled_wakeups (thread_id)`, + ); + await ctx.client.query( + `CREATE INDEX IF NOT EXISTS scheduled_wakeups_claimed_at_s_idx ON ${schema}.scheduled_wakeups (claimed_at_s)`, + ); + }, + }, + ]; /** diff --git a/packages/storage/pg/src/storage.ts b/packages/storage/pg/src/storage.ts index b7cc21d7..9292e78c 100644 --- a/packages/storage/pg/src/storage.ts +++ b/packages/storage/pg/src/storage.ts @@ -15,6 +15,7 @@ import { UnimplementedError } from "@kernl-sdk/shared/lib"; /* pg */ import { PGThreadStore } from "./thread/store"; import { PGMemoryStore } from "./memory/store"; +import { PGWakeupStore } from "./wakeup/store"; import { MIGRATIONS } from "./migrations"; import { SQL_IDENTIFIER_REGEX } from "./sql"; @@ -109,11 +110,13 @@ export class PGStorage implements KernlStorage { threads: PGThreadStore; memories: PGMemoryStore; + wakeups: PGWakeupStore; constructor(config: PGStorageConfig) { this.pool = config.pool; this.threads = new PGThreadStore(this.pool, () => this.ensureInit()); this.memories = new PGMemoryStore(this.pool, () => this.ensureInit()); + this.wakeups = new PGWakeupStore(this.pool, () => this.ensureInit()); } /** diff --git a/packages/storage/pg/src/wakeup/codec.ts b/packages/storage/pg/src/wakeup/codec.ts new file mode 100644 index 00000000..803ee751 --- /dev/null +++ b/packages/storage/pg/src/wakeup/codec.ts @@ -0,0 +1,119 @@ +/** + * Wakeup codecs. + * /packages/storage/pg/src/wakeup/codec.ts + */ + +import type { Codec } from "@kernl-sdk/shared/lib"; + +import type { + ScheduledWakeupRecord, +} from "@kernl-sdk/storage"; + +import type { + NewScheduledWakeup, + ScheduledWakeup, +} from "kernl"; + +/** + * Convert NewScheduledWakeup (domain) -> ScheduledWakeupRecord (DB). + */ +export const NewScheduledWakeupCodec: Codec< + NewScheduledWakeup, + ScheduledWakeupRecord +> = { + encode(input) { + // Store timestamps in epoch *seconds* to keep the DB representation stable. + const nowS = Math.floor(Date.now() / 1000); + + // Back-compat: allow either `sleepFor` (seconds) or legacy `runAt` (ms epoch). + const anyInput = input as any; + const sleepForS = + typeof anyInput.sleepFor === "number" + ? Math.max(0, Math.floor(anyInput.sleepFor)) + : Math.max(0, Math.floor(anyInput.runAt / 1000) - nowS); + + const wakeupAtS = nowS + sleepForS; + + return { + id: input.id, + thread_id: input.threadId, + sleep_for: sleepForS, + wakeup_at: wakeupAtS, + reason: input.reason ?? null, + woken: false, + claimed_at_s: null, + created_at: nowS, + updated_at: nowS, + error: null, + }; + }, + + decode() { + throw new Error("NewScheduledWakeupCodec.decode not implemented"); + }, +}; + +/** + * Convert between ScheduledWakeup (domain) and ScheduledWakeupRecord (DB). + */ +export const ScheduledWakeupCodec: Codec< + ScheduledWakeup, + ScheduledWakeupRecord +> = { + encode(wakeup) { + const anyWakeup = wakeup as any; + + const wakeupAtMs = anyWakeup.wakeupAt ?? anyWakeup.runAt; + const wakeupAtS = Math.floor(wakeupAtMs / 1000); + + // Domain timestamps are typically ms; DB stores seconds. + const createdAtS = + typeof anyWakeup.createdAt === "number" + ? Math.floor(anyWakeup.createdAt / 1000) + : wakeupAtS; + + const updatedAtS = + typeof anyWakeup.updatedAt === "number" + ? Math.floor(anyWakeup.updatedAt / 1000) + : createdAtS; + + const sleepForS = + typeof anyWakeup.sleepFor === "number" + ? Math.max(0, Math.floor(anyWakeup.sleepFor)) + : Math.max(0, wakeupAtS - createdAtS); + + const claimedAtS = + anyWakeup.claimedAt != null ? Math.floor(anyWakeup.claimedAt / 1000) : null; + + return { + id: wakeup.id, + thread_id: wakeup.threadId, + sleep_for: sleepForS, + wakeup_at: wakeupAtS, + reason: wakeup.reason, + woken: wakeup.woken, + claimed_at_s: claimedAtS, + created_at: createdAtS, + updated_at: updatedAtS, + error: wakeup.error, + }; + }, + + decode(record) { + return { + id: record.id, + threadId: record.thread_id, + // Back-compat: expose legacy `runAt` in ms while the DB stores seconds. + runAt: record.wakeup_at * 1000, + // New: also expose the stored duration (seconds). + sleepFor: record.sleep_for, + reason: record.reason, + woken: record.woken, + claimedAt: + record.claimed_at_s != null ? record.claimed_at_s * 1000 : null, + createdAt: record.created_at * 1000, + updatedAt: record.updated_at * 1000, + error: record.error, + }; + }, +}; diff --git a/packages/storage/pg/src/wakeup/sql.ts b/packages/storage/pg/src/wakeup/sql.ts new file mode 100644 index 00000000..d8d0cbe8 --- /dev/null +++ b/packages/storage/pg/src/wakeup/sql.ts @@ -0,0 +1,90 @@ +/** + * Wakeup SQL conversion codecs. + * /packages/storage/pg/src/wakeup/sql.ts + */ + +import type { Codec } from "@kernl-sdk/shared/lib"; +import type { ScheduledWakeupUpdate } from "kernl"; + +export interface SQLClause { + sql: string; + params: unknown[]; +} + +/** + * Input for building a partial UPDATE ... SET clause. + */ +export interface PatchInput { + patch: ScheduledWakeupUpdate; + startIdx: number; +} + +/** + * Build `SET ...` clause for scheduled_wakeups updates. + * + * Mirrors the style of /packages/storage/pg/src/memory/sql.ts::PATCH + */ +export const PATCH: Codec = { + encode({ patch, startIdx }) { + const sets: string[] = []; + const params: unknown[] = []; + let idx = startIdx; + + // Back-compat: allow either `wakeupAt` (ms), legacy `runAt` (ms), and/or `sleepFor` (seconds). + const anyPatch = patch as any; + const wakeupAtMs = anyPatch.wakeupAt ?? anyPatch.runAt; + + if (wakeupAtMs !== undefined) { + const wakeupAtS = Math.floor(wakeupAtMs / 1000); + const wakeupIdx = idx++; + sets.push(`wakeup_at = $${wakeupIdx}`); + params.push(wakeupAtS); + + // If caller did not explicitly set sleepFor, keep it consistent with wakeup_at and created_at. + if (anyPatch.sleepFor === undefined) { + sets.push(`sleep_for = GREATEST(0, $${wakeupIdx} - created_at)`); + } + } + + if (anyPatch.sleepFor !== undefined) { + sets.push(`sleep_for = $${idx++}`); + params.push(Math.max(0, Math.floor(anyPatch.sleepFor))); + } + + if (patch.reason !== undefined) { + sets.push(`reason = $${idx++}`); + params.push(patch.reason); + } + + if (patch.woken !== undefined) { + sets.push(`woken = $${idx++}`); + params.push(patch.woken); + } + + if (patch.claimedAt !== undefined) { + sets.push(`claimed_at_s = $${idx++}`); + params.push( + patch.claimedAt != null ? Math.floor(patch.claimedAt / 1000) : null, + ); + } + + if (patch.error !== undefined) { + sets.push(`error = $${idx++}`); + params.push(patch.error); + } + + // Always bump updated_at in epoch seconds. + const nowS = Math.floor(((patch as any).updatedAt ?? Date.now()) / 1000); + sets.push(`updated_at = $${idx++}`); + params.push(nowS); + + return { + sql: sets.join(", "), + params, + }; + }, + + decode() { + throw new Error("PATCH.decode not implemented"); + }, +}; diff --git a/packages/storage/pg/src/wakeup/store.ts b/packages/storage/pg/src/wakeup/store.ts new file mode 100644 index 00000000..bb722716 --- /dev/null +++ b/packages/storage/pg/src/wakeup/store.ts @@ -0,0 +1,169 @@ +/** + * PostgreSQL Wakeup store implementation. + * /packages/storage/pg/src/wakeup/store.ts + */ + +import type { Pool, PoolClient } from "pg"; + +import type { + WakeupStore, + NewScheduledWakeup, + ScheduledWakeup, + ScheduledWakeupUpdate, +} from "kernl"; + +import { + KERNL_SCHEMA_NAME, + ScheduledWakeupRecordSchema, + type ScheduledWakeupRecord, +} from "@kernl-sdk/storage"; + +import { NewScheduledWakeupCodec, ScheduledWakeupCodec } from "./codec"; +import { PATCH } from "./sql"; + +/** + * Convert ms/bigint ms to a JS number in ms. + */ +const toMs = (value: number | bigint): number => + typeof value === "bigint" ? Number(value) : value; + +/** + * Convert ms to epoch seconds (integer). + */ +const toSeconds = (ms: number | bigint): number => + Math.floor(toMs(ms) / 1000); + +/** + * PostgreSQL implementation of WakeupStore. + * + * Follows the same pattern as PGMemoryStore: + * - depends on ensureInit() to create tables + * - always validates with Zod Schemas + * - converts between domain types and DB records via codecs + */ +export class PGWakeupStore implements WakeupStore { + private db: Pool | PoolClient; + private ensureInit: () => Promise; + + constructor(db: Pool | PoolClient, ensureInit: () => Promise) { + this.db = db; + this.ensureInit = ensureInit; + } + + async create(input: NewScheduledWakeup): Promise { + await this.ensureInit(); + + const row = NewScheduledWakeupCodec.encode(input); + + const result = await this.db.query( + `INSERT INTO ${KERNL_SCHEMA_NAME}.scheduled_wakeups + (id, thread_id, sleep_for, wakeup_at, reason, woken, claimed_at_s, created_at, updated_at, error) + VALUES ($1,$2,$3,$4,$5,$6,$7,$8,$9,$10) + RETURNING *`, + [ + row.id, + row.thread_id, + row.sleep_for, + row.wakeup_at, + row.reason, + row.woken, + row.claimed_at_s, + row.created_at, + row.updated_at, + row.error, + ], + ); + + const record = ScheduledWakeupRecordSchema.parse(result.rows[0]); + return ScheduledWakeupCodec.decode(record); + } + + async get(id: string): Promise { + await this.ensureInit(); + + const result = await this.db.query( + `SELECT * FROM ${KERNL_SCHEMA_NAME}.scheduled_wakeups WHERE id = $1`, + [id], + ); + + if (result.rows.length === 0) return null; + + const record = ScheduledWakeupRecordSchema.parse(result.rows[0]); + return ScheduledWakeupCodec.decode(record); + } + + async update( + id: string, + patch: ScheduledWakeupUpdate, + ): Promise { + await this.ensureInit(); + + const { sql, params } = PATCH.encode({ patch, startIdx: 2 }); + + const result = await this.db.query( + `UPDATE ${KERNL_SCHEMA_NAME}.scheduled_wakeups + SET ${sql} + WHERE id = $1 + RETURNING *`, + [id, ...params], + ); + + if (result.rows.length === 0) { + throw new Error(`Wakeup with id ${id} not found`); + } + + const record = ScheduledWakeupRecordSchema.parse(result.rows[0]); + return ScheduledWakeupCodec.decode(record); + } + + async delete(id: string): Promise { + await this.ensureInit(); + + await this.db.query( + `DELETE FROM ${KERNL_SCHEMA_NAME}.scheduled_wakeups WHERE id = $1`, + [id], + ); + } + + /** + * Atomically claim up to `limit` due wakeups. + * + * Uses SELECT ... FOR UPDATE SKIP LOCKED pattern. + */ + async claimDue( + nowMs: number | bigint, + limit: number, + ): Promise { + await this.ensureInit(); + + const nowS = toSeconds(nowMs); + + const result = await this.db.query( + ` + WITH due AS ( + SELECT id + FROM ${KERNL_SCHEMA_NAME}.scheduled_wakeups + WHERE woken = FALSE + AND claimed_at_s IS NULL + AND wakeup_at <= $1 + ORDER BY wakeup_at ASC + LIMIT $2 + FOR UPDATE SKIP LOCKED + ) + UPDATE ${KERNL_SCHEMA_NAME}.scheduled_wakeups AS sw + SET claimed_at_s = $1, + updated_at = $1 + FROM due + WHERE sw.id = due.id + RETURNING sw.*; + `, + [nowS, limit], + ); + + return result.rows.map((row) => + ScheduledWakeupCodec.decode( + ScheduledWakeupRecordSchema.parse(row), + ), + ); + } +}