diff --git a/packages/khala/package.json b/packages/khala/package.json index bf57d90..4272ff1 100644 --- a/packages/khala/package.json +++ b/packages/khala/package.json @@ -10,6 +10,7 @@ "check": "biome check ." }, "dependencies": { + "@uncaged/nerve-core": "workspace:*", "hono": "^4.7.0", "jsonata": "^2.0.5", "ulidx": "^2.4.1" diff --git a/packages/khala/src/__tests__/result.test.ts b/packages/khala/src/__tests__/result.test.ts index b183df5..47aae71 100644 --- a/packages/khala/src/__tests__/result.test.ts +++ b/packages/khala/src/__tests__/result.test.ts @@ -1,5 +1,5 @@ -import { describe, it, expect } from "vitest"; -import { ok, err } from "../result.js"; +import { err, ok } from "@uncaged/nerve-core"; +import { describe, expect, it } from "vitest"; describe("Result", () => { it("ok wraps a value", () => { diff --git a/packages/khala/src/db.ts b/packages/khala/src/db.ts index c49db59..b0c877f 100644 --- a/packages/khala/src/db.ts +++ b/packages/khala/src/db.ts @@ -1,6 +1,6 @@ +import type { Result } from "@uncaged/nerve-core"; +import { err, ok } from "@uncaged/nerve-core"; import { type ULID, monotonicFactory } from "ulidx"; -import type { Result } from "./result.js"; -import { err, ok } from "./result.js"; import type { Agent, GetThreadMessagesOpts, Message, Task, Thread } from "./types.js"; const generateUlid = monotonicFactory((): number => Date.now()); @@ -88,19 +88,16 @@ export async function appendMessage( step: number, agentId: string | null, ): Promise { - await db + const row = await db .prepare( `INSERT INTO messages (thread_id, role, content, meta, step, agent_id, created_at) - VALUES (?, ?, ?, ?, ?, ?, datetime('now'))`, + VALUES (?, ?, ?, ?, ?, ?, datetime('now')) + RETURNING *`, ) .bind(threadId, role, content, meta, step, agentId) - .run(); - const row = await db - .prepare("SELECT * FROM messages WHERE thread_id = ? AND step = ? ORDER BY id DESC LIMIT 1") - .bind(threadId, step) .first(); if (!row) { - throw new Error("appendMessage: row not found after insert"); + throw new Error("appendMessage: INSERT RETURNING returned no row"); } return rowToMessage(row); } diff --git a/packages/khala/src/index.ts b/packages/khala/src/index.ts index 02d3383..8967f32 100644 --- a/packages/khala/src/index.ts +++ b/packages/khala/src/index.ts @@ -97,6 +97,7 @@ app.get("/tasks", agentAuth, async (c) => { }); // Claim a task +// TODO(#132): Add max_concurrent_claims per agent to prevent greedy claiming app.post("/tasks/:id/claim", agentAuth, async (c) => { const taskId = c.req.param("id"); const agentId = c.get("agentId"); diff --git a/packages/khala/src/moderator.ts b/packages/khala/src/moderator.ts index 697956b..9f121a4 100644 --- a/packages/khala/src/moderator.ts +++ b/packages/khala/src/moderator.ts @@ -1,6 +1,6 @@ +import type { Result } from "@uncaged/nerve-core"; +import { err, ok } from "@uncaged/nerve-core"; import jsonata from "jsonata"; -import type { Result } from "./result.js"; -import { err, ok } from "./result.js"; export type ModeratorDecision = { role: string }; diff --git a/packages/khala/src/result.ts b/packages/khala/src/result.ts deleted file mode 100644 index a166c35..0000000 --- a/packages/khala/src/result.ts +++ /dev/null @@ -1,9 +0,0 @@ -export type Result = { ok: true; value: T } | { ok: false; error: E }; - -export function ok(value: T): Result { - return { ok: true, value }; -} - -export function err(error: E): Result { - return { ok: false, error }; -} diff --git a/packages/khala/src/routes/admin.ts b/packages/khala/src/routes/admin.ts index 1e231ca..1341129 100644 --- a/packages/khala/src/routes/admin.ts +++ b/packages/khala/src/routes/admin.ts @@ -1,8 +1,8 @@ +import type { Result } from "@uncaged/nerve-core"; import { Hono } from "hono"; import { requireAdmin, sha256Hex } from "../auth.js"; import { deleteAgent, listAgents, registerAgent } from "../db.js"; import type { KhalaBindings } from "../env.js"; -import type { Result } from "../result.js"; const admin = new Hono<{ Bindings: KhalaBindings }>(); @@ -21,6 +21,7 @@ admin.get("/agents", async (c) => { return c.json({ agents: [...agents] }); }); +// TODO(#132): Add rate limiting for agent registration endpoint admin.post("/agents", async (c) => { if (!requireAdmin(c)) { return c.json({ error: "unauthorized" }, 401); diff --git a/packages/khala/src/thread-do.ts b/packages/khala/src/thread-do.ts index 9f5f187..4376b01 100644 --- a/packages/khala/src/thread-do.ts +++ b/packages/khala/src/thread-do.ts @@ -87,7 +87,7 @@ export class ThreadDO { ); return { ok: true }; } - await createTask(db, threadId, role, r.prompt, 300); + await createTask(db, threadId, role, r.prompt, r.timeoutSeconds ?? 300); return { ok: true }; } diff --git a/packages/khala/src/workflows.ts b/packages/khala/src/workflows.ts index c9f1e7b..966f60a 100644 --- a/packages/khala/src/workflows.ts +++ b/packages/khala/src/workflows.ts @@ -1,5 +1,7 @@ export type CloudRole = { prompt: string; + /** Turn timeout in seconds. Defaults to 300. */ + timeoutSeconds: number | null; }; export type CloudWorkflowDef = { diff --git a/packages/khala/src/workflows/ping-pong.ts b/packages/khala/src/workflows/ping-pong.ts index 11890e8..24be37c 100644 --- a/packages/khala/src/workflows/ping-pong.ts +++ b/packages/khala/src/workflows/ping-pong.ts @@ -3,8 +3,8 @@ import { type CloudWorkflowDef, registerWorkflow } from "../workflows.js"; const pingPong: CloudWorkflowDef = { name: "ping-pong", roles: { - pinger: { prompt: "Send the literal word ping" }, - ponger: { prompt: "Send the literal word pong" }, + pinger: { prompt: "Send the literal word ping", timeoutSeconds: null }, + ponger: { prompt: "Send the literal word pong", timeoutSeconds: null }, }, moderator: `$count(steps) >= 6 ? { "role": "__end__" } : $count(steps) % 2 = 0 ? { "role": "pinger" } : { "role": "ponger" }`, }; diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index e1349aa..2619ee5 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -86,6 +86,9 @@ importers: packages/khala: dependencies: + '@uncaged/nerve-core': + specifier: workspace:* + version: link:../core hono: specifier: ^4.7.0 version: 4.12.15