@lakitu/sdk 0.1.0

package/convex/sandbox/state/checkpoints.ts

@@ -0,0 +1,341 @@
+/**
+ * Checkpoints - State Persistence for Chained Agent Runs
+ *
+ * Save and restore agent state for long-running tasks that
+ * exceed sandbox timeout limits.
+ */
+
+import { mutation, query, internalMutation, internalQuery } from "../_generated/server";
+import { v } from "convex/values";
+
+// ============================================
+// Types
+// ============================================
+
+const fileSnapshotValidator = v.object({
+  path: v.string(),
+  contentHash: v.string(),
+  size: v.number(),
+  lastModified: v.number(),
+});
+
+const beadSnapshotValidator = v.object({
+  id: v.string(),
+  title: v.string(),
+  status: v.string(),
+  type: v.string(),
+  priority: v.number(),
+});
+
+const messageValidator = v.object({
+  role: v.string(),
+  content: v.string(),
+  timestamp: v.optional(v.number()),
+});
+
+// ============================================
+// Mutations
+// ============================================
+
+/**
+ * Create a checkpoint
+ */
+export const create = mutation({
+  args: {
+    sessionId: v.string(),
+    threadId: v.string(),
+    iteration: v.number(),
+    // Compressed conversation history
+    messageHistory: v.array(messageValidator),
+    // Files touched so far
+    fileState: v.array(fileSnapshotValidator),
+    // Task tracking state
+    beadsState: v.array(beadSnapshotValidator),
+    // Artifacts saved (references)
+    artifactsProduced: v.array(v.string()),
+    // What to do next
+    nextTask: v.string(),
+    // Why we're checkpointing
+    reason: v.union(
+      v.literal("timeout"),
+      v.literal("token_limit"),
+      v.literal("manual"),
+      v.literal("error_recovery")
+    ),
+    // Additional context
+    metadata: v.optional(v.any()),
+  },
+  handler: async (ctx, args) => {
+    // Mark previous checkpoints as superseded
+    const previous = await ctx.db
+      .query("checkpoints")
+      .withIndex("by_session", (q) => q.eq("sessionId", args.sessionId))
+      .filter((q) => q.eq(q.field("status"), "active"))
+      .collect();
+
+    for (const checkpoint of previous) {
+      await ctx.db.patch(checkpoint._id, { status: "superseded" });
+    }
+
+    // Create new checkpoint
+    return await ctx.db.insert("checkpoints", {
+      sessionId: args.sessionId,
+      threadId: args.threadId,
+      iteration: args.iteration,
+      messageHistory: args.messageHistory,
+      fileState: args.fileState,
+      beadsState: args.beadsState,
+      artifactsProduced: args.artifactsProduced,
+      nextTask: args.nextTask,
+      reason: args.reason,
+      status: "active",
+      metadata: args.metadata,
+      createdAt: Date.now(),
+    });
+  },
+});
+
+/**
+ * Mark checkpoint as restored (used)
+ */
+export const markRestored = internalMutation({
+  args: {
+    id: v.id("checkpoints"),
+    newThreadId: v.optional(v.string()),
+  },
+  handler: async (ctx, args) => {
+    await ctx.db.patch(args.id, {
+      status: "restored",
+      restoredAt: Date.now(),
+      restoredToThread: args.newThreadId,
+    });
+  },
+});
+
+/**
+ * Mark checkpoint as completed (task finished)
+ */
+export const markCompleted = mutation({
+  args: {
+    sessionId: v.string(),
+    finalResult: v.optional(v.any()),
+  },
+  handler: async (ctx, args) => {
+    const checkpoints = await ctx.db
+      .query("checkpoints")
+      .withIndex("by_session", (q) => q.eq("sessionId", args.sessionId))
+      .collect();
+
+    for (const checkpoint of checkpoints) {
+      await ctx.db.patch(checkpoint._id, {
+        status: "completed",
+        completedAt: Date.now(),
+        finalResult: args.finalResult,
+      });
+    }
+
+    return { updated: checkpoints.length };
+  },
+});
+
+/**
+ * Mark checkpoint as failed
+ */
+export const markFailed = mutation({
+  args: {
+    id: v.id("checkpoints"),
+    error: v.string(),
+  },
+  handler: async (ctx, args) => {
+    await ctx.db.patch(args.id, {
+      status: "failed",
+      failedAt: Date.now(),
+      error: args.error,
+    });
+  },
+});
+
+// ============================================
+// Queries
+// ============================================
+
+/**
+ * Get the latest active checkpoint for a session
+ */
+export const getLatest = query({
+  args: { sessionId: v.string() },
+  handler: async (ctx, args) => {
+    return await ctx.db
+      .query("checkpoints")
+      .withIndex("by_session", (q) => q.eq("sessionId", args.sessionId))
+      .filter((q) => q.eq(q.field("status"), "active"))
+      .order("desc")
+      .first();
+  },
+});
+
+/**
+ * Get checkpoint by ID
+ */
+export const get = query({
+  args: { id: v.id("checkpoints") },
+  handler: async (ctx, args) => {
+    return await ctx.db.get(args.id);
+  },
+});
+
+/**
+ * Internal: Get checkpoint by ID (for internal actions)
+ */
+export const internalGet = internalQuery({
+  args: { id: v.id("checkpoints") },
+  handler: async (ctx, args) => {
+    return await ctx.db.get(args.id);
+  },
+});
+
+/**
+ * Get all checkpoints for a session (history)
+ */
+export const getHistory = query({
+  args: { sessionId: v.string() },
+  handler: async (ctx, args) => {
+    return await ctx.db
+      .query("checkpoints")
+      .withIndex("by_session", (q) => q.eq("sessionId", args.sessionId))
+      .order("desc")
+      .collect();
+  },
+});
+
+/**
+ * Get checkpoint stats
+ */
+export const getStats = query({
+  args: {},
+  handler: async (ctx) => {
+    const all = await ctx.db.query("checkpoints").collect();
+
+    const byStatus = {
+      active: 0,
+      restored: 0,
+      completed: 0,
+      failed: 0,
+      superseded: 0,
+    };
+
+    const byReason = {
+      timeout: 0,
+      token_limit: 0,
+      manual: 0,
+      error_recovery: 0,
+    };
+
+    for (const cp of all) {
+      byStatus[cp.status as keyof typeof byStatus]++;
+      byReason[cp.reason as keyof typeof byReason]++;
+    }
+
+    return {
+      total: all.length,
+      byStatus,
+      byReason,
+      averageIteration:
+        all.length > 0
+          ? all.reduce((sum, cp) => sum + cp.iteration, 0) / all.length
+          : 0,
+    };
+  },
+});
+
+// ============================================
+// Internal Mutations
+// ============================================
+
+/**
+ * Create checkpoint from current state
+ * Called internally during timeout handling
+ */
+export const createFromCurrentState = internalMutation({
+  args: {
+    threadId: v.string(),
+    nextTask: v.string(),
+    iteration: v.number(),
+  },
+  handler: async (ctx, args) => {
+    // Get current file state
+    const fileStates = await ctx.db.query("fileState").collect();
+    const fileState = fileStates.map((f) => ({
+      path: f.path,
+      contentHash: f.contentHash || "",
+      size: f.size || 0,
+      lastModified: f.lastAccessAt,
+    }));
+
+    // Get current beads state
+    const beads = await ctx.db.query("beads").collect();
+    const beadsState = beads.map((b) => ({
+      id: b._id,
+      title: b.title,
+      status: b.status,
+      type: b.type,
+      priority: b.priority,
+    }));
+
+    // Get artifacts produced
+    const artifacts = await ctx.db.query("artifacts").collect();
+    const artifactsProduced = artifacts.map((a) => a.name);
+
+    // Create session ID based on thread
+    const sessionId = `session_${args.threadId}`;
+
+    // Create checkpoint
+    return await ctx.db.insert("checkpoints", {
+      sessionId,
+      threadId: args.threadId,
+      iteration: args.iteration,
+      messageHistory: [], // Will be populated by agent SDK
+      fileState,
+      beadsState,
+      artifactsProduced,
+      nextTask: args.nextTask,
+      reason: "timeout",
+      status: "active",
+      createdAt: Date.now(),
+    });
+  },
+});
+
+// ============================================
+// Cleanup
+// ============================================
+
+/**
+ * Clean up old checkpoints
+ */
+export const cleanup = mutation({
+  args: {
+    olderThanMs: v.optional(v.number()),
+    keepCompleted: v.optional(v.boolean()),
+  },
+  handler: async (ctx, args) => {
+    const cutoff = Date.now() - (args.olderThanMs ?? 7 * 24 * 3600000); // 7 days default
+
+    const old = await ctx.db
+      .query("checkpoints")
+      .filter((q) => q.lt(q.field("createdAt"), cutoff))
+      .collect();
+
+    let deleted = 0;
+    for (const checkpoint of old) {
+      // Skip completed if requested
+      if (args.keepCompleted && checkpoint.status === "completed") {
+        continue;
+      }
+      await ctx.db.delete(checkpoint._id);
+      deleted++;
+    }
+
+    return { deleted };
+  },
+});

package/convex/sandbox/state/files.ts

@@ -0,0 +1,383 @@
+/**
+ * State - File State Tracking
+ *
+ * Explicit state management with diff-driven architecture.
+ * Tracks file changes, generates diffs, and supports rollback.
+ */
+
+import { mutation, query, internalMutation, internalQuery } from "../_generated/server";
+import { v } from "convex/values";
+
+// ============================================
+// Mutations
+// ============================================
+
+/**
+ * Track a file access (read or write)
+ */
+export const trackFileAccess = mutation({
+  args: {
+    path: v.string(),
+    operation: v.union(v.literal("read"), v.literal("write"), v.literal("edit")),
+    contentHash: v.optional(v.string()),
+    size: v.optional(v.number()),
+    threadId: v.optional(v.string()),
+  },
+  handler: async (ctx, args) => {
+    // Get existing file state
+    const existing = await ctx.db
+      .query("fileState")
+      .withIndex("by_path", (q) => q.eq("path", args.path))
+      .first();
+
+    if (existing) {
+      // Update existing
+      await ctx.db.patch(existing._id, {
+        lastOperation: args.operation,
+        lastAccessAt: Date.now(),
+        contentHash: args.contentHash ?? existing.contentHash,
+        size: args.size ?? existing.size,
+        accessCount: existing.accessCount + 1,
+      });
+      return existing._id;
+    }
+
+    // Create new file state entry
+    return await ctx.db.insert("fileState", {
+      path: args.path,
+      lastOperation: args.operation,
+      lastAccessAt: Date.now(),
+      contentHash: args.contentHash,
+      size: args.size,
+      threadId: args.threadId,
+      accessCount: 1,
+      createdAt: Date.now(),
+    });
+  },
+});
+
+/**
+ * Record a file edit with diff
+ */
+export const recordEdit = mutation({
+  args: {
+    path: v.string(),
+    oldContent: v.string(),
+    newContent: v.string(),
+    diff: v.string(),
+    verified: v.boolean(),
+    threadId: v.optional(v.string()),
+  },
+  handler: async (ctx, args) => {
+    // Get file state
+    const fileState = await ctx.db
+      .query("fileState")
+      .withIndex("by_path", (q) => q.eq("path", args.path))
+      .first();
+
+    // Insert edit history
+    const editId = await ctx.db.insert("editHistory", {
+      path: args.path,
+      fileStateId: fileState?._id,
+      oldContentHash: simpleHash(args.oldContent),
+      newContentHash: simpleHash(args.newContent),
+      diff: args.diff,
+      verified: args.verified,
+      threadId: args.threadId,
+      createdAt: Date.now(),
+    });
+
+    // Update file state
+    if (fileState) {
+      await ctx.db.patch(fileState._id, {
+        lastOperation: "edit",
+        lastAccessAt: Date.now(),
+        contentHash: simpleHash(args.newContent),
+        size: args.newContent.length,
+        lastEditId: editId,
+      });
+    }
+
+    return editId;
+  },
+});
+
+/**
+ * Mark a file as rolled back
+ */
+export const rollback = mutation({
+  args: {
+    editId: v.id("editHistory"),
+    reason: v.string(),
+  },
+  handler: async (ctx, args) => {
+    const edit = await ctx.db.get(args.editId);
+    if (!edit) {
+      throw new Error("Edit not found");
+    }
+
+    await ctx.db.patch(args.editId, {
+      rolledBack: true,
+      rollbackReason: args.reason,
+      rolledBackAt: Date.now(),
+    });
+
+    return { success: true };
+  },
+});
+
+// ============================================
+// Queries
+// ============================================
+
+/**
+ * Get file state by path
+ */
+export const getByPath = query({
+  args: { path: v.string() },
+  handler: async (ctx, args) => {
+    return await ctx.db
+      .query("fileState")
+      .withIndex("by_path", (q) => q.eq("path", args.path))
+      .first();
+  },
+});
+
+/**
+ * Get all files accessed in a thread
+ */
+export const getByThread = query({
+  args: { threadId: v.string() },
+  handler: async (ctx, args) => {
+    return await ctx.db
+      .query("fileState")
+      .withIndex("by_thread", (q) => q.eq("threadId", args.threadId))
+      .collect();
+  },
+});
+
+/**
+ * Get edit history for a file
+ */
+export const getEditHistory = query({
+  args: {
+    path: v.string(),
+    limit: v.optional(v.number()),
+  },
+  handler: async (ctx, args) => {
+    return await ctx.db
+      .query("editHistory")
+      .withIndex("by_path", (q) => q.eq("path", args.path))
+      .order("desc")
+      .take(args.limit ?? 20);
+  },
+});
+
+/**
+ * Get recent edits across all files
+ */
+export const getRecentEdits = query({
+  args: { limit: v.optional(v.number()) },
+  handler: async (ctx, args) => {
+    return await ctx.db
+      .query("editHistory")
+      .order("desc")
+      .take(args.limit ?? 50);
+  },
+});
+
+/**
+ * Check if a file has changed since a given hash
+ */
+export const hasChanged = query({
+  args: {
+    path: v.string(),
+    expectedHash: v.string(),
+  },
+  handler: async (ctx, args) => {
+    const fileState = await ctx.db
+      .query("fileState")
+      .withIndex("by_path", (q) => q.eq("path", args.path))
+      .first();
+
+    if (!fileState) {
+      return { changed: true, reason: "File not tracked" };
+    }
+
+    if (fileState.contentHash !== args.expectedHash) {
+      return {
+        changed: true,
+        reason: "Content hash mismatch",
+        currentHash: fileState.contentHash,
+      };
+    }
+
+    return { changed: false };
+  },
+});
+
+// ============================================
+// Internal Mutations
+// ============================================
+
+/**
+ * Restore file state from a checkpoint
+ */
+export const restoreFromCheckpoint = internalMutation({
+  args: {
+    checkpointId: v.id("checkpoints"),
+  },
+  handler: async (ctx, args) => {
+    const checkpoint = await ctx.db.get(args.checkpointId);
+    if (!checkpoint) {
+      throw new Error(`Checkpoint ${args.checkpointId} not found`);
+    }
+
+    // Clear existing file state
+    const existingStates = await ctx.db.query("fileState").collect();
+    for (const state of existingStates) {
+      await ctx.db.delete(state._id);
+    }
+
+    // Restore file states from checkpoint
+    for (const file of checkpoint.fileState) {
+      await ctx.db.insert("fileState", {
+        path: file.path,
+        lastOperation: "read",
+        lastAccessAt: file.lastModified,
+        contentHash: file.contentHash,
+        size: file.size,
+        accessCount: 1,
+        createdAt: Date.now(),
+      });
+    }
+
+    // Restore beads state
+    // Note: This creates new beads, doesn't update existing ones
+    // In a real implementation, you'd want to sync with existing beads
+    for (const bead of checkpoint.beadsState) {
+      const existing = await ctx.db
+        .query("beads")
+        .filter((q) => q.eq(q.field("title"), bead.title))
+        .first();
+
+      if (!existing) {
+        await ctx.db.insert("beads", {
+          title: bead.title,
+          type: bead.type as any,
+          status: bead.status as any,
+          priority: bead.priority,
+          createdAt: Date.now(),
+          updatedAt: Date.now(),
+        });
+      }
+    }
+
+    return { restored: true, checkpointId: args.checkpointId };
+  },
+});
+
+// ============================================
+// Internal Mutation Wrappers (for tools)
+// ============================================
+
+/**
+ * Internal version of trackFileAccess for tool use
+ */
+export const internalTrackFileAccess = internalMutation({
+  args: {
+    path: v.string(),
+    operation: v.union(v.literal("read"), v.literal("write"), v.literal("edit")),
+    contentHash: v.optional(v.string()),
+    size: v.optional(v.number()),
+    threadId: v.optional(v.string()),
+  },
+  handler: async (ctx, args) => {
+    const existing = await ctx.db
+      .query("fileState")
+      .withIndex("by_path", (q) => q.eq("path", args.path))
+      .first();
+
+    if (existing) {
+      await ctx.db.patch(existing._id, {
+        lastOperation: args.operation,
+        lastAccessAt: Date.now(),
+        contentHash: args.contentHash ?? existing.contentHash,
+        size: args.size ?? existing.size,
+        accessCount: existing.accessCount + 1,
+      });
+      return existing._id;
+    }
+
+    return await ctx.db.insert("fileState", {
+      path: args.path,
+      lastOperation: args.operation,
+      lastAccessAt: Date.now(),
+      contentHash: args.contentHash,
+      size: args.size,
+      threadId: args.threadId,
+      accessCount: 1,
+      createdAt: Date.now(),
+    });
+  },
+});
+
+/**
+ * Internal version of recordEdit for tool use
+ */
+export const internalRecordEdit = internalMutation({
+  args: {
+    path: v.string(),
+    oldContent: v.string(),
+    newContent: v.string(),
+    diff: v.string(),
+    verified: v.boolean(),
+    threadId: v.optional(v.string()),
+  },
+  handler: async (ctx, args) => {
+    const fileState = await ctx.db
+      .query("fileState")
+      .withIndex("by_path", (q) => q.eq("path", args.path))
+      .first();
+
+    const editId = await ctx.db.insert("editHistory", {
+      path: args.path,
+      fileStateId: fileState?._id,
+      oldContentHash: simpleHash(args.oldContent),
+      newContentHash: simpleHash(args.newContent),
+      diff: args.diff,
+      verified: args.verified,
+      threadId: args.threadId,
+      createdAt: Date.now(),
+    });
+
+    if (fileState) {
+      await ctx.db.patch(fileState._id, {
+        lastOperation: "edit",
+        lastAccessAt: Date.now(),
+        contentHash: simpleHash(args.newContent),
+        size: args.newContent.length,
+        lastEditId: editId,
+      });
+    }
+
+    return editId;
+  },
+});
+
+// ============================================
+// Helpers
+// ============================================
+
+/**
+ * Simple hash function for content comparison
+ */
+function simpleHash(str: string): string {
+  let hash = 0;
+  for (let i = 0; i < str.length; i++) {
+    const char = str.charCodeAt(i);
+    hash = (hash << 5) - hash + char;
+    hash = hash & hash; // Convert to 32bit integer
+  }
+  return hash.toString(16);
+}