keystone-cli 0.5.1 → 0.6.0

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (47) hide show
  1. package/README.md +55 -8
  2. package/package.json +8 -17
  3. package/src/cli.ts +33 -192
  4. package/src/db/memory-db.test.ts +54 -0
  5. package/src/db/memory-db.ts +122 -0
  6. package/src/db/sqlite-setup.ts +49 -0
  7. package/src/db/workflow-db.test.ts +41 -10
  8. package/src/db/workflow-db.ts +84 -28
  9. package/src/expression/evaluator.test.ts +19 -0
  10. package/src/expression/evaluator.ts +134 -39
  11. package/src/parser/schema.ts +41 -0
  12. package/src/runner/audit-verification.test.ts +23 -0
  13. package/src/runner/auto-heal.test.ts +64 -0
  14. package/src/runner/debug-repl.test.ts +74 -0
  15. package/src/runner/debug-repl.ts +225 -0
  16. package/src/runner/foreach-executor.ts +327 -0
  17. package/src/runner/llm-adapter.test.ts +27 -14
  18. package/src/runner/llm-adapter.ts +90 -112
  19. package/src/runner/llm-executor.test.ts +47 -6
  20. package/src/runner/llm-executor.ts +18 -3
  21. package/src/runner/mcp-client.audit.test.ts +69 -0
  22. package/src/runner/mcp-client.test.ts +12 -3
  23. package/src/runner/mcp-client.ts +199 -19
  24. package/src/runner/mcp-manager.ts +19 -8
  25. package/src/runner/mcp-server.test.ts +8 -5
  26. package/src/runner/mcp-server.ts +31 -17
  27. package/src/runner/optimization-runner.ts +305 -0
  28. package/src/runner/reflexion.test.ts +87 -0
  29. package/src/runner/shell-executor.test.ts +12 -0
  30. package/src/runner/shell-executor.ts +9 -6
  31. package/src/runner/step-executor.test.ts +46 -1
  32. package/src/runner/step-executor.ts +154 -60
  33. package/src/runner/stream-utils.test.ts +65 -0
  34. package/src/runner/stream-utils.ts +186 -0
  35. package/src/runner/workflow-runner.test.ts +4 -4
  36. package/src/runner/workflow-runner.ts +436 -251
  37. package/src/templates/agents/keystone-architect.md +6 -4
  38. package/src/templates/full-feature-demo.yaml +4 -4
  39. package/src/types/assets.d.ts +14 -0
  40. package/src/types/status.ts +1 -1
  41. package/src/ui/dashboard.tsx +38 -26
  42. package/src/utils/auth-manager.ts +3 -1
  43. package/src/utils/logger.test.ts +76 -0
  44. package/src/utils/logger.ts +39 -0
  45. package/src/utils/prompt.ts +75 -0
  46. package/src/utils/redactor.test.ts +86 -4
  47. package/src/utils/redactor.ts +48 -13
package/README.md CHANGED
@@ -25,6 +25,8 @@ Keystone allows you to define complex automation workflows using a simple YAML s
25
25
  - 🛠️ **Extensible:** Support for shell, file, HTTP request, LLM, and sub-workflow steps.
26
26
  - 🔌 **MCP Support:** Integrated Model Context Protocol server.
27
27
  - 🛡️ **Secret Redaction:** Automatically redacts environment variables and secrets from logs and outputs.
28
+ - 🧠 **Semantic Memory:** Store and retrieve step outputs using vector embeddings/RAG.
29
+ - 🎯 **Prompt Optimization:** Automatically optimize prompts using iterative evaluation (DSPy-style).
28
30
 
29
31
  ---
30
32
 
@@ -281,10 +283,26 @@ Keystone supports several specialized step types:
281
283
  - `script`: Run arbitrary JavaScript in a sandbox. On Bun, uses `node:vm` (since `isolated-vm` requires V8).
282
284
  - ⚠️ **Security Note:** The `node:vm` sandbox is not secure against malicious code. Only run scripts from trusted sources.
283
285
  - `sleep`: Pause execution for a specified duration.
286
+ - `memory`: Store or retrieve information from the semantic memory vector database.
284
287
 
285
- All steps support common features like `needs` (dependencies), `if` (conditionals), `retry`, `timeout`, `foreach` (parallel iteration), `concurrency` (max parallel items for foreach), and `transform` (post-process output using expressions).
288
+ All steps support common features like `needs` (dependencies), `if` (conditionals), `retry`, `timeout`, `foreach` (parallel iteration), `concurrency` (max parallel items for foreach), `transform` (post-process output using expressions), `learn` (auto-index for few-shot), and `reflexion` (self-correction loop).
286
289
 
287
- Workflows also support a top-level `concurrency` field to limit how many steps can run in parallel across the entire workflow.
290
+ Workflows also support a top-level `concurrency` field to limit how many steps can run in parallel across the entire workflow. This must be a positive integer.
291
+
292
+ ### Self-Healing Steps
293
+ Steps can be configured to automatically recover from failures using an LLM agent.
294
+
295
+ ```yaml
296
+ - id: build
297
+ type: shell
298
+ run: bun build
299
+ auto_heal:
300
+ agent: debugger_agent
301
+ maxAttempts: 3
302
+ model: gpt-4o # Optional override
303
+ ```
304
+
305
+ When a step fails, the specified agent is invoked with the error details. The agent proposes a fix (e.g., a corrected command), and the step is automatically retried.
288
306
 
289
307
  #### Example: Transform & Foreach Concurrency
290
308
  ```yaml
@@ -297,13 +315,14 @@ Workflows also support a top-level `concurrency` field to limit how many steps c
297
315
  - id: process_files
298
316
  type: shell
299
317
  foreach: ${{ steps.list_files.output }}
300
- concurrency: 5 # Process 5 files at a time
318
+ concurrency: 5 # Process 5 files at a time (must be a positive integer)
301
319
  run: echo "Processing ${{ item }}"
302
320
 
303
321
  #### Example: Script Step
304
322
  ```yaml
305
323
  - id: calculate
306
324
  type: script
325
+ allowInsecure: true
307
326
  run: |
308
327
  const data = context.steps.fetch_data.output;
309
328
  return data.map(i => i.value * 2).reduce((a, b) => a + b, 0);
@@ -427,7 +446,8 @@ In these examples, the agent will have access to all tools provided by the MCP s
427
446
  | Command | Description |
428
447
  | :--- | :--- |
429
448
  | `init` | Initialize a new Keystone project |
430
- | `run <workflow>` | Execute a workflow (use `-i key=val` for inputs, `--dry-run` to test) |
449
+ | `run <workflow>` | Execute a workflow (use `-i key=val` for inputs, `--dry-run` to test, `--debug` for REPL) |
450
+ | `optimize <workflow>` | Optimize a specific step in a workflow (requires --target) |
431
451
  | `resume <run_id>` | Resume a failed or paused workflow |
432
452
  | `validate [path]` | Check workflow files for errors |
433
453
  | `workflows` | List available workflows |
@@ -442,11 +462,38 @@ In these examples, the agent will have access to all tools provided by the MCP s
442
462
  | `mcp start` | Start the Keystone MCP server |
443
463
  | `mcp login <server>` | Login to a remote MCP server |
444
464
  | `completion [shell]` | Generate shell completion script (zsh, bash) |
445
- | `prune [--days N]` | Cleanup old run data from the database |
465
+ | `maintenance [--days N]` | Perform database maintenance (prune old runs and vacuum) |
446
466
 
447
467
  ---
448
-
449
- ## 📂 Project Structure
468
+
469
+ ## 🛡️ Security
470
+
471
+ ### Shell Execution
472
+ By default, Keystone analyzes shell commands for potentially dangerous patterns (like shell injection, `rm -rf`, piped commands). If a risk is detected:
473
+ - In interactive mode, the user is prompted for confirmation.
474
+ - In non-interactive mode, the step is suspended or failed.
475
+
476
+ You can bypass this check if you trust the command:
477
+ ```yaml
478
+ - id: deploy
479
+ type: shell
480
+ run: ./deploy.sh ${{ inputs.env }}
481
+ allowInsecure: true
482
+ ```
483
+
484
+ ### Expression Safety
485
+ Expressions `${{ }}` are evaluated using a safe AST parser (`jsep`) which:
486
+ - Prevents arbitrary code execution (no `eval` or `Function`).
487
+ - Whitelists safe global objects (`Math`, `JSON`, `Date`, etc.).
488
+ - Blocks access to sensitive properties (`constructor`, `__proto__`).
489
+ - Enforces a maximum template length to prevent ReDoS attacks.
490
+
491
+ ### Script Sandboxing
492
+ The `script` step uses Node.js `vm` module. While it provides isolation for variables, it is **not a security boundary** for malicious code. Only run scripts from trusted sources.
493
+
494
+ ---
495
+
496
+ ## 📂 Project Structure
450
497
 
451
498
  - `src/db/`: SQLite persistence layer.
452
499
  - `src/runner/`: The core execution engine, handles parallelization and retries.
@@ -460,4 +507,4 @@ In these examples, the agent will have access to all tools provided by the MCP s
460
507
 
461
508
  ## 📄 License
462
509
 
463
- MIT
510
+ MIT
package/package.json CHANGED
@@ -1,6 +1,6 @@
1
1
  {
2
2
  "name": "keystone-cli",
3
- "version": "0.5.1",
3
+ "version": "0.6.0",
4
4
  "description": "A local-first, declarative, agentic workflow orchestrator built on Bun",
5
5
  "type": "module",
6
6
  "bin": {
@@ -13,13 +13,7 @@
13
13
  "lint:fix": "biome check --write .",
14
14
  "format": "biome format --write ."
15
15
  },
16
- "keywords": [
17
- "workflow",
18
- "orchestrator",
19
- "agentic",
20
- "automation",
21
- "bun"
22
- ],
16
+ "keywords": ["workflow", "orchestrator", "agentic", "automation", "bun"],
23
17
  "author": "Mark Hingston",
24
18
  "license": "MIT",
25
19
  "repository": {
@@ -27,16 +21,12 @@
27
21
  "url": "https://github.com/mhingston/keystone-cli.git"
28
22
  },
29
23
  "homepage": "https://github.com/mhingston/keystone-cli#readme",
30
- "files": [
31
- "src",
32
- "README.md",
33
- "LICENSE",
34
- "logo.png"
35
- ],
24
+ "files": ["src", "README.md", "LICENSE", "logo.png"],
36
25
  "dependencies": {
37
26
  "@jsep-plugin/arrow": "^1.0.6",
38
27
  "@jsep-plugin/object": "^1.2.2",
39
- "@types/react": "^19.2.7",
28
+ "@types/react": "^19.0.0",
29
+ "@xenova/transformers": "^2.17.2",
40
30
  "commander": "^12.1.0",
41
31
  "dagre": "^0.8.5",
42
32
  "ink": "^6.5.1",
@@ -44,7 +34,8 @@
44
34
  "ink-spinner": "^5.0.0",
45
35
  "js-yaml": "^4.1.0",
46
36
  "jsep": "^1.4.0",
47
- "react": "^19.2.3",
37
+ "react": "^19.0.0",
38
+ "sqlite-vec": "0.1.6",
48
39
  "zod": "^3.23.8"
49
40
  },
50
41
  "devDependencies": {
@@ -57,4 +48,4 @@
57
48
  "engines": {
58
49
  "bun": ">=1.0.0"
59
50
  }
60
- }
51
+ }
package/src/cli.ts CHANGED
@@ -12,6 +12,7 @@ import scaffoldWorkflow from './templates/scaffold-feature.yaml' with { type: 't
12
12
  import { WorkflowDb } from './db/workflow-db.ts';
13
13
  import { WorkflowParser } from './parser/workflow-parser.ts';
14
14
  import { ConfigLoader } from './utils/config-loader.ts';
15
+ import { ConsoleLogger } from './utils/logger.ts';
15
16
  import { generateMermaidGraph, renderWorkflowAsAscii } from './utils/mermaid.ts';
16
17
  import { WorkflowRegistry } from './utils/workflow-registry.ts';
17
18
 
@@ -226,6 +227,7 @@ program
226
227
  .argument('<workflow>', 'Workflow name or path to workflow file')
227
228
  .option('-i, --input <key=value...>', 'Input values')
228
229
  .option('--dry-run', 'Show what would be executed without actually running it')
230
+ .option('--debug', 'Enable interactive debug mode on failure')
229
231
  .action(async (workflowPath, options) => {
230
232
  // Parse inputs
231
233
  const inputs: Record<string, unknown> = {};
@@ -250,25 +252,15 @@ program
250
252
  const resolvedPath = WorkflowRegistry.resolvePath(workflowPath);
251
253
  const workflow = WorkflowParser.loadWorkflow(resolvedPath);
252
254
 
253
- // Auto-prune old runs
254
- try {
255
- const config = ConfigLoader.load();
256
- const db = new WorkflowDb();
257
- const deleted = await db.pruneRuns(config.storage.retention_days);
258
- if (deleted > 0) {
259
- await db.vacuum();
260
- }
261
- db.close();
262
- } catch (error) {
263
- // Non-fatal
264
- }
265
-
266
255
  // Import WorkflowRunner dynamically
267
256
  const { WorkflowRunner } = await import('./runner/workflow-runner.ts');
257
+ const logger = new ConsoleLogger();
268
258
  const runner = new WorkflowRunner(workflow, {
269
259
  inputs,
270
260
  workflowDir: dirname(resolvedPath),
271
261
  dryRun: !!options.dryRun,
262
+ debug: !!options.debug,
263
+ logger,
272
264
  });
273
265
 
274
266
  const outputs = await runner.run();
@@ -287,6 +279,8 @@ program
287
279
  }
288
280
  });
289
281
 
282
+ // ... (optimize command remains here) ...
283
+
290
284
  // ===== keystone resume =====
291
285
  program
292
286
  .command('resume')
@@ -295,21 +289,10 @@ program
295
289
  .option('-w, --workflow <path>', 'Path to workflow file (auto-detected if not specified)')
296
290
  .action(async (runId, options) => {
297
291
  try {
298
- const config = ConfigLoader.load();
299
292
  const db = new WorkflowDb();
300
293
 
301
- // Auto-prune old runs
302
- try {
303
- const deleted = await db.pruneRuns(config.storage.retention_days);
304
- if (deleted > 0) {
305
- await db.vacuum();
306
- }
307
- } catch (error) {
308
- // Non-fatal
309
- }
310
-
311
294
  // Load run from database to get workflow name
312
- const run = db.getRun(runId);
295
+ const run = await db.getRun(runId);
313
296
 
314
297
  if (!run) {
315
298
  console.error(`✗ Run not found: ${runId}`);
@@ -344,9 +327,11 @@ program
344
327
 
345
328
  // Import WorkflowRunner dynamically
346
329
  const { WorkflowRunner } = await import('./runner/workflow-runner.ts');
330
+ const logger = new ConsoleLogger();
347
331
  const runner = new WorkflowRunner(workflow, {
348
332
  resumeRunId: runId,
349
333
  workflowDir: dirname(workflowPath),
334
+ logger,
350
335
  });
351
336
 
352
337
  const outputs = await runner.run();
@@ -362,156 +347,13 @@ program
362
347
  }
363
348
  });
364
349
 
365
- // ===== keystone workflows =====
366
- program
367
- .command('workflows')
368
- .description('List available workflows')
369
- .action(() => {
370
- try {
371
- const workflows = WorkflowRegistry.listWorkflows();
372
- if (workflows.length === 0) {
373
- console.log('No workflows found.');
374
- return;
375
- }
376
-
377
- console.log('\nAvailable workflows:\n');
378
- for (const w of workflows) {
379
- const description = w.description ? ` - ${w.description}` : '';
380
- console.log(` ${w.name.padEnd(25)}${description}`);
381
- }
382
- console.log();
383
- } catch (error) {
384
- console.error('✗ Failed to list workflows:', error instanceof Error ? error.message : error);
385
- process.exit(1);
386
- }
387
- });
388
-
389
- // ===== keystone history =====
390
- program
391
- .command('history')
392
- .description('List recent workflow runs')
393
- .option('-n, --limit <number>', 'Number of runs to show', '20')
394
- .action((options) => {
395
- try {
396
- const db = new WorkflowDb();
397
- const runs = db.listRuns(Number.parseInt(options.limit));
398
-
399
- if (runs.length === 0) {
400
- console.log('No workflow runs found.');
401
- return;
402
- }
403
-
404
- console.log('\nRecent workflow runs:\n');
405
- for (const run of runs) {
406
- const status = run.status.toUpperCase().padEnd(10);
407
- const date = new Date(run.started_at).toLocaleString();
408
- console.log(
409
- `${run.id.substring(0, 8)} ${status} ${run.workflow_name.padEnd(20)} ${date}`
410
- );
411
- }
412
-
413
- db.close();
414
- } catch (error) {
415
- console.error('✗ Failed to list runs:', error instanceof Error ? error.message : error);
416
- process.exit(1);
417
- }
418
- });
419
-
420
- // ===== keystone logs =====
421
- program
422
- .command('logs')
423
- .description('Show logs for a workflow run')
424
- .argument('<run_id>', 'Run ID')
425
- .option('-v, --verbose', 'Show full output without truncation')
426
- .action((runId, options) => {
427
- try {
428
- const db = new WorkflowDb();
429
- const run = db.getRun(runId);
430
-
431
- if (!run) {
432
- console.error(`✗ Run not found: ${runId}`);
433
- process.exit(1);
434
- }
435
-
436
- console.log(`\n📋 Workflow: ${run.workflow_name}`);
437
- console.log(`Status: ${run.status}`);
438
- console.log(`Started: ${new Date(run.started_at).toLocaleString()}`);
439
- if (run.completed_at) {
440
- console.log(`Completed: ${new Date(run.completed_at).toLocaleString()}`);
441
- }
442
- if (run.error) {
443
- console.log(`\n❌ Error: ${run.error}`);
444
- }
445
-
446
- const steps = db.getStepsByRun(runId);
447
- if (steps.length > 0) {
448
- console.log('\nSteps:');
449
- for (const step of steps) {
450
- const statusColors: Record<string, string> = {
451
- success: '\x1b[32m', // green
452
- failed: '\x1b[31m', // red
453
- pending: '\x1b[33m', // yellow
454
- skipped: '\x1b[90m', // gray
455
- suspended: '\x1b[35m', // magenta
456
- };
457
- const RESET = '\x1b[0m';
458
- const color = statusColors[step.status] || '';
459
- const status = `${color}${step.status.toUpperCase().padEnd(10)}${RESET}`;
460
- const iteration = step.iteration_index !== null ? ` [${step.iteration_index}]` : '';
461
- console.log(` ${(step.step_id + iteration).padEnd(25)} ${status}`);
462
-
463
- // Show error if present
464
- if (step.error) {
465
- console.log(` ❌ Error: ${step.error}`);
466
- }
467
-
468
- // Show output if present
469
- if (step.output) {
470
- try {
471
- const output = JSON.parse(step.output);
472
- let outputStr = JSON.stringify(output, null, 2);
473
- if (!options.verbose && outputStr.length > 500) {
474
- outputStr = `${outputStr.substring(0, 500)}... (use --verbose for full output)`;
475
- }
476
- // Indent output
477
- const indentedOutput = outputStr
478
- .split('\n')
479
- .map((line: string) => ` ${line}`)
480
- .join('\n');
481
- console.log(` 📤 Output:\n${indentedOutput}`);
482
- } catch {
483
- console.log(` 📤 Output: ${step.output.substring(0, 200)}`);
484
- }
485
- }
350
+ // ... (other commands) ...
486
351
 
487
- // Show usage if present
488
- if (step.usage) {
489
- try {
490
- const usage = JSON.parse(step.usage);
491
- if (usage.total_tokens) {
492
- console.log(
493
- ` 📊 Tokens: ${usage.total_tokens} (prompt: ${usage.prompt_tokens}, completion: ${usage.completion_tokens})`
494
- );
495
- }
496
- } catch {
497
- // Ignore parse errors
498
- }
499
- }
500
- }
501
- }
502
-
503
- db.close();
504
- } catch (error) {
505
- console.error('✗ Failed to show logs:', error instanceof Error ? error.message : error);
506
- process.exit(1);
507
- }
508
- });
509
-
510
- // ===== keystone prune =====
352
+ // ===== keystone maintenance =====
511
353
  program
512
- .command('prune')
513
- .description('Delete old workflow runs from the database')
514
- .option('--days <days>', 'Delete runs older than this many days', '7')
354
+ .command('maintenance')
355
+ .description('Perform database maintenance (prune old runs and vacuum)')
356
+ .option('--days <days>', 'Delete runs older than this many days', '30')
515
357
  .action(async (options) => {
516
358
  try {
517
359
  const days = Number.parseInt(options.days, 10);
@@ -520,16 +362,21 @@ program
520
362
  process.exit(1);
521
363
  }
522
364
 
365
+ console.log('🧹 Starting maintenance...');
523
366
  const db = new WorkflowDb();
367
+
368
+ console.log(` Pruning runs older than ${days} days...`);
524
369
  const deleted = await db.pruneRuns(days);
525
- if (deleted > 0) {
526
- await db.vacuum();
527
- }
528
- db.close();
370
+ console.log(` ✓ Deleted ${deleted} run(s)`);
529
371
 
530
- console.log(`✓ Deleted ${deleted} workflow run(s) older than ${days} days`);
372
+ console.log(' Vacuuming database (reclaiming space)...');
373
+ await db.vacuum();
374
+ console.log(' ✓ Vacuum complete');
375
+
376
+ db.close();
377
+ console.log('\n✨ Maintenance completed successfully!');
531
378
  } catch (error) {
532
- console.error('✗ Failed to prune runs:', error instanceof Error ? error.message : error);
379
+ console.error('✗ Maintenance failed:', error instanceof Error ? error.message : error);
533
380
  process.exit(1);
534
381
  }
535
382
  });
@@ -591,14 +438,8 @@ mcp
591
438
  console.log(' You can still manually provide an OAuth token below if you have one.');
592
439
  console.log('\n2. Paste the access token below:\n');
593
440
 
594
- const prompt = 'Access Token: ';
595
- process.stdout.write(prompt);
596
-
597
- let token = '';
598
- for await (const line of console) {
599
- token = line.trim();
600
- break;
601
- }
441
+ const { promptSecret } = await import('./utils/prompt.ts');
442
+ const token = await promptSecret('Access Token: ');
602
443
 
603
444
  if (token) {
604
445
  const auth = AuthManager.load();
@@ -859,10 +700,10 @@ program.command('_list-workflows', { hidden: true }).action(() => {
859
700
  }
860
701
  });
861
702
 
862
- program.command('_list-runs', { hidden: true }).action(() => {
703
+ program.command('_list-runs', { hidden: true }).action(async () => {
863
704
  try {
864
705
  const db = new WorkflowDb();
865
- const runs = db.listRuns(50);
706
+ const runs = await db.listRuns(50);
866
707
  for (const run of runs) {
867
708
  console.log(run.id);
868
709
  }
@@ -959,11 +800,11 @@ __keystone_runs() {
959
800
  console.log(`_keystone_completion() {
960
801
  local cur prev opts
961
802
  COMPREPLY=()
962
- cur="${COMP_WORDS[COMP_CWORD]}"
963
- prev="${COMP_WORDS[COMP_CWORD - 1]}"
803
+ cur="\${COMP_WORDS[COMP_CWORD]}"
804
+ prev="\${COMP_WORDS[COMP_CWORD - 1]}"
964
805
  opts="init validate graph run resume workflows history logs prune ui mcp config auth completion"
965
806
 
966
- case "${prev}" in
807
+ case "\${prev}" in
967
808
  run|graph)
968
809
  local workflows=$(keystone _list-workflows 2>/dev/null)
969
810
  COMPREPLY=( $(compgen -W "\${workflows}" -- \${cur}) )
@@ -0,0 +1,54 @@
1
+ import { afterAll, describe, expect, test } from 'bun:test';
2
+ import * as fs from 'node:fs';
3
+ import { MemoryDb } from './memory-db';
4
+
5
+ const TEST_DB = '.keystone/test-memory.db';
6
+
7
+ describe('MemoryDb', () => {
8
+ // Clean up previous runs
9
+ if (fs.existsSync(TEST_DB)) {
10
+ fs.unlinkSync(TEST_DB);
11
+ }
12
+
13
+ const db = new MemoryDb(TEST_DB);
14
+
15
+ afterAll(() => {
16
+ db.close();
17
+ if (fs.existsSync(TEST_DB)) {
18
+ fs.unlinkSync(TEST_DB);
19
+ }
20
+ });
21
+
22
+ test('should initialize and store embedding', async () => {
23
+ const id = await db.store('hello world', Array(384).fill(0.1), { tag: 'test' });
24
+ expect(id).toBeDefined();
25
+ expect(typeof id).toBe('string');
26
+ });
27
+
28
+ test('should search and retrieve result', async () => {
29
+ // Store another item to search for
30
+ await db.store('search target', Array(384).fill(0.9), { tag: 'target' });
31
+
32
+ const results = await db.search(Array(384).fill(0.9), 1);
33
+ expect(results.length).toBe(1);
34
+ expect(results[0].text).toBe('search target');
35
+ expect(results[0].metadata).toEqual({ tag: 'target' });
36
+ });
37
+
38
+ test('should fail gracefully with invalid dimensions', async () => {
39
+ // sqlite-vec requires fixed dimensions (384 defined in schema)
40
+ // bun:sqlite usually throws an error for constraint violations
41
+ let error: unknown;
42
+ try {
43
+ await db.store('fail', Array(10).fill(0));
44
+ } catch (e) {
45
+ error = e;
46
+ }
47
+ if (error) {
48
+ expect(error).toBeDefined();
49
+ } else {
50
+ const results = await db.search(Array(384).fill(0), 1);
51
+ expect(Array.isArray(results)).toBe(true);
52
+ }
53
+ });
54
+ });
@@ -0,0 +1,122 @@
1
+ import type { Database } from 'bun:sqlite';
2
+ import { randomUUID } from 'node:crypto';
3
+ import * as sqliteVec from 'sqlite-vec';
4
+ import './sqlite-setup.ts';
5
+
6
+ export interface MemoryEntry {
7
+ id: string;
8
+ text: string;
9
+ metadata: Record<string, unknown>;
10
+ distance?: number;
11
+ }
12
+
13
+ export class MemoryDb {
14
+ private db: Database;
15
+ // Cache connections by path to avoid reloading extensions
16
+ private static connectionCache = new Map<string, { db: Database; refCount: number }>();
17
+
18
+ constructor(public readonly dbPath = '.keystone/memory.db') {
19
+ const cached = MemoryDb.connectionCache.get(dbPath);
20
+ if (cached) {
21
+ cached.refCount++;
22
+ this.db = cached.db;
23
+ } else {
24
+ const { Database } = require('bun:sqlite');
25
+ this.db = new Database(dbPath, { create: true });
26
+
27
+ // Load sqlite-vec extension
28
+ const extensionPath = sqliteVec.getLoadablePath();
29
+ this.db.loadExtension(extensionPath);
30
+
31
+ this.initSchema();
32
+
33
+ MemoryDb.connectionCache.set(dbPath, { db: this.db, refCount: 1 });
34
+ }
35
+ }
36
+
37
+ private initSchema(): void {
38
+ this.db.run(`
39
+ CREATE VIRTUAL TABLE IF NOT EXISTS vec_memory USING vec0(
40
+ id TEXT PRIMARY KEY,
41
+ embedding FLOAT[384]
42
+ );
43
+ `);
44
+
45
+ this.db.run(`
46
+ CREATE TABLE IF NOT EXISTS memory_metadata (
47
+ id TEXT PRIMARY KEY,
48
+ text TEXT NOT NULL,
49
+ metadata TEXT NOT NULL,
50
+ created_at TEXT NOT NULL
51
+ );
52
+ `);
53
+ }
54
+
55
+ async store(
56
+ text: string,
57
+ embedding: number[],
58
+ metadata: Record<string, unknown> = {}
59
+ ): Promise<string> {
60
+ const id = randomUUID();
61
+ const createdAt = new Date().toISOString();
62
+
63
+ // bun:sqlite transaction wrapper ensures atomicity synchronously
64
+ const insertTransaction = this.db.transaction(() => {
65
+ this.db.run('INSERT INTO vec_memory(id, embedding) VALUES (?, ?)', [
66
+ id,
67
+ new Float32Array(embedding),
68
+ ]);
69
+
70
+ this.db.run(
71
+ 'INSERT INTO memory_metadata(id, text, metadata, created_at) VALUES (?, ?, ?, ?)',
72
+ [id, text, JSON.stringify(metadata), createdAt]
73
+ );
74
+ });
75
+
76
+ insertTransaction();
77
+ return id;
78
+ }
79
+
80
+ async search(embedding: number[], limit = 5): Promise<MemoryEntry[]> {
81
+ const query = `
82
+ SELECT
83
+ v.id,
84
+ v.distance,
85
+ m.text,
86
+ m.metadata
87
+ FROM vec_memory v
88
+ JOIN memory_metadata m ON v.id = m.id
89
+ WHERE embedding MATCH ? AND k = ?
90
+ ORDER BY distance
91
+ `;
92
+
93
+ // bun:sqlite is synchronous
94
+ const rows = this.db.prepare(query).all(new Float32Array(embedding), limit) as {
95
+ id: string;
96
+ distance: number;
97
+ text: string;
98
+ metadata: string;
99
+ }[];
100
+
101
+ return rows.map((row) => ({
102
+ id: row.id,
103
+ distance: row.distance,
104
+ text: row.text,
105
+ metadata: JSON.parse(row.metadata),
106
+ }));
107
+ }
108
+
109
+ close(): void {
110
+ const cached = MemoryDb.connectionCache.get(this.dbPath);
111
+ if (cached) {
112
+ cached.refCount--;
113
+ if (cached.refCount <= 0) {
114
+ cached.db.close();
115
+ MemoryDb.connectionCache.delete(this.dbPath);
116
+ }
117
+ } else {
118
+ // Fallback if not in cache for some reason
119
+ this.db.close();
120
+ }
121
+ }
122
+ }