@ai-setting/roy-agent-core 1.0.0

package/src/env/log-trace/span-storage.test.ts

@@ -0,0 +1,145 @@
+/**
+ * @fileoverview SQLiteSpanStorage WAL 模式测试
+ */
+
+import { describe, test, expect, beforeEach, afterEach } from 'bun:test';
+import { SQLiteSpanStorage } from './span-storage';
+import { join } from 'path';
+import { mkdtempSync, rmSync } from 'fs';
+import os from 'os';
+
+describe('SQLiteSpanStorage with WAL', () => {
+  let storage: SQLiteSpanStorage;
+  let tempDir: string;
+
+  beforeEach(() => {
+    tempDir = mkdtempSync(join(os.tmpdir(), 'span-storage-test-'));
+    storage = new SQLiteSpanStorage(join(tempDir, 'traces.db'));
+  });
+
+  afterEach(() => {
+    storage.close();
+    rmSync(tempDir, { recursive: true, force: true });
+  });
+
+  describe('WAL mode', () => {
+    test('should enable WAL journal mode', async () => {
+      await storage.initialize();
+      
+      // WAL mode should be enabled
+      // We can verify by checking the database journal mode
+      const spans = storage.findBySpanId('non-existent');
+      expect(spans).toBeUndefined();
+    });
+
+    test('should save and retrieve spans', async () => {
+      await storage.initialize();
+      
+      const span = {
+        traceId: 'trace-123',
+        spanId: 'span-456',
+        parentSpanId: undefined,
+        name: 'test-span',
+        kind: 'internal' as const,
+        status: 'ok' as const,
+        startTime: Date.now(),
+        endTime: Date.now() + 100,
+        attributes: { key: 'value' },
+      };
+      
+      storage.save(span);
+      
+      const retrieved = storage.findBySpanId('span-456');
+      expect(retrieved).toBeDefined();
+      expect(retrieved!.name).toBe('test-span');
+      expect(retrieved!.traceId).toBe('trace-123');
+    });
+
+    test('should handle concurrent writes', async () => {
+      await storage.initialize();
+      
+      // Simulate concurrent writes
+      const promises: Promise<void>[] = [];
+      for (let i = 0; i < 10; i++) {
+        const span = {
+          traceId: 'trace-concurrent',
+          spanId: `span-concurrent-${i}`,
+          parentSpanId: undefined,
+          name: `concurrent-span-${i}`,
+          kind: 'internal' as const,
+          status: 'ok' as const,
+          startTime: Date.now(),
+          attributes: {},
+        };
+        promises.push(Promise.resolve(storage.save(span)));
+      }
+      
+      await Promise.all(promises);
+      
+      // All spans should be saved
+      const traces = storage.findByTraceId('trace-concurrent');
+      expect(traces.length).toBeGreaterThanOrEqual(10);
+    });
+
+    test('should save batch of spans', async () => {
+      await storage.initialize();
+      
+      const spans = Array.from({ length: 5 }, (_, i) => ({
+        traceId: 'trace-batch',
+        spanId: `span-batch-${i}`,
+        parentSpanId: undefined,
+        name: `batch-span-${i}`,
+        kind: 'internal' as const,
+        status: 'ok' as const,
+        startTime: Date.now(),
+        attributes: {},
+      }));
+      
+      storage.saveBatch(spans);
+      
+      const traces = storage.findByTraceId('trace-batch');
+      expect(traces.length).toBe(5);
+    });
+
+    test('should list traces', async () => {
+      await storage.initialize();
+      
+      // Save some spans
+      for (let i = 0; i < 5; i++) {
+        storage.save({
+          traceId: `trace-list-${i}`,
+          spanId: `span-list-${i}`,
+          parentSpanId: undefined,
+          name: `list-span-${i}`,
+          kind: 'internal' as const,
+          status: 'ok' as const,
+          startTime: Date.now(),
+          attributes: {},
+        });
+      }
+      
+      const traces = storage.listTraces(10);
+      expect(traces.length).toBeGreaterThanOrEqual(5);
+    });
+
+    test('should delete by traceId', async () => {
+      await storage.initialize();
+      
+      storage.save({
+        traceId: 'trace-to-delete',
+        spanId: 'span-to-delete',
+        parentSpanId: undefined,
+        name: 'delete-span',
+        kind: 'internal' as const,
+        status: 'ok' as const,
+        startTime: Date.now(),
+        attributes: {},
+      });
+      
+      storage.deleteByTraceId('trace-to-delete');
+      
+      const retrieved = storage.findBySpanId('span-to-delete');
+      expect(retrieved).toBeUndefined();
+    });
+  });
+});

package/src/env/log-trace/span-storage.ts

@@ -0,0 +1,230 @@
+/**
+ * @fileoverview Span Storage 实现
+ * 
+ * 提供 SQLiteSpanStorage 一种实现，用于持久化存储 span 数据
+ */
+
+import {
+  type Span,
+  type SpanStorage,
+  type TraceInfo,
+  SpanKind,
+  SpanStatus,
+} from "./types";
+
+// 重新导出类型
+export type { SpanStorage, TraceInfo } from "./types";
+
+// ============================================================================
+// SQLiteSpanStorage
+// ============================================================================
+
+/**
+ * SQLite 存储实现
+ * 
+ * 使用 Bun 内置的 bun:sqlite 模块进行持久化存储
+ * 
+ * 注意：此实现仅支持 SQLite 文件存储，不支持 :memory: 模式
+ * WAL 模式默认启用以支持跨进程并发访问
+ */
+export class SQLiteSpanStorage implements SpanStorage {
+  private db: any = null;
+  private dbPath: string;
+  private initialized = false;
+
+  /**
+   * 创建 SQLiteSpanStorage 实例
+   * 
+   * @param dbPath SQLite 数据库文件路径（必填）
+   * @throws Error 如果 dbPath 为空或为 :memory:
+   */
+  constructor(dbPath: string) {
+    // 验证：dbPath 必填，且不能是 :memory:
+    if (!dbPath) {
+      throw new Error("SQLiteSpanStorage requires a valid dbPath parameter");
+    }
+    if (dbPath === ":memory:") {
+      throw new Error("SQLiteSpanStorage does not support :memory: mode. Use a file path instead.");
+    }
+    this.dbPath = dbPath;
+  }
+
+  async initialize(): Promise<void> {
+    if (this.initialized) return;
+
+    // 使用 Bun 内置的 bun:sqlite
+    const { Database } = require("bun:sqlite");
+
+    // 确保目录存在
+    const fs = require("fs");
+    const path = require("path");
+    const dir = path.dirname(this.dbPath);
+    if (!fs.existsSync(dir)) {
+      fs.mkdirSync(dir, { recursive: true });
+    }
+
+    this.db = new Database(this.dbPath);
+    
+    // 启用 WAL 模式以支持更好的并发写入
+    // WAL (Write-Ahead Logging) 允许读操作与写操作并发执行
+    // WAL 模式支持跨进程访问，是本实现的核心特性
+    this.db.run("PRAGMA journal_mode=WAL");
+    // 设置 busy_timeout 以处理锁冲突（5秒）
+    this.db.run("PRAGMA busy_timeout=5000");
+    // 使用 NORMAL 同步模式，平衡性能和数据安全
+    this.db.run("PRAGMA synchronous=NORMAL");
+    
+    // 创建表
+    this.db.run(`
+      CREATE TABLE IF NOT EXISTS span (
+        span_id TEXT PRIMARY KEY,
+        trace_id TEXT NOT NULL,
+        parent_span_id TEXT,
+        name TEXT NOT NULL,
+        kind TEXT NOT NULL,
+        status TEXT NOT NULL,
+        start_time INTEGER NOT NULL,
+        end_time INTEGER,
+        attributes TEXT,
+        result TEXT,
+        error TEXT,
+        time_created INTEGER NOT NULL
+      )
+    `);
+    this.db.run("CREATE INDEX IF NOT EXISTS idx_span_trace ON span(trace_id)");
+    this.db.run("CREATE INDEX IF NOT EXISTS idx_span_parent ON span(parent_span_id)");
+
+    this.initialized = true;
+  }
+
+  save(span: Span): void {
+    if (!this.db) return;
+
+    const stmt = this.db.prepare(`
+      INSERT OR REPLACE INTO span 
+      (span_id, trace_id, parent_span_id, name, kind, status, start_time, end_time, attributes, result, error, time_created)
+      VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
+    `);
+    
+    // result 处理：如果已经是字符串，直接存储；否则 stringify
+    const resultValue = span.result !== undefined && span.result !== null
+      ? (typeof span.result === "string" ? span.result : JSON.stringify(span.result))
+      : null;
+    
+    stmt.run(
+      span.spanId,
+      span.traceId,
+      span.parentSpanId || null,
+      span.name,
+      span.kind,
+      span.status,
+      span.startTime,
+      span.endTime || null,
+      JSON.stringify(span.attributes),
+      resultValue,
+      span.error || null,
+      Date.now()
+    );
+  }
+
+  saveBatch(spans: Span[]): void {
+    for (const span of spans) {
+      this.save(span);
+    }
+  }
+
+  findByTraceId(traceId: string): Span[] {
+    if (!this.db) return [];
+
+    const rows = this.db.prepare("SELECT * FROM span WHERE trace_id = ?").all(traceId);
+    return this.buildTree(rows.map(this.rowToSpan));
+  }
+
+  findBySpanId(spanId: string): Span | undefined {
+    if (!this.db) return undefined;
+
+    const row = this.db.prepare("SELECT * FROM span WHERE span_id = ?").get(spanId);
+    return row ? this.rowToSpan(row) : undefined;
+  }
+
+  listTraces(limit: number = 10): TraceInfo[] {
+    if (!this.db) return [];
+
+    const rows = this.db.prepare(`
+      SELECT trace_id, 
+             MIN(start_time) as start_time,
+             MAX(end_time) as end_time,
+             COUNT(*) as span_count
+      FROM span 
+      GROUP BY trace_id 
+      ORDER BY start_time DESC 
+      LIMIT ?
+    `).all(limit);
+
+    return rows.map((row: any) => ({
+      traceId: row.trace_id,
+      rootSpanName: "unknown",
+      startTime: row.start_time,
+      endTime: row.end_time,
+      duration: row.end_time - row.start_time,
+      spanCount: row.span_count,
+      status: "ok" as const,
+    }));
+  }
+
+  deleteByTraceId(traceId: string): void {
+    if (this.db) {
+      this.db.prepare("DELETE FROM span WHERE trace_id = ?").run(traceId);
+    }
+  }
+
+  close(): void {
+    if (this.db) {
+      this.db.close();
+      this.db = null;
+    }
+  }
+
+  private rowToSpan(row: any): Span {
+    return {
+      traceId: row.trace_id,
+      spanId: row.span_id,
+      parentSpanId: row.parent_span_id,
+      name: row.name,
+      kind: row.kind as SpanKind,
+      status: row.status as SpanStatus,
+      startTime: row.start_time,
+      endTime: row.end_time,
+      attributes: row.attributes ? JSON.parse(row.attributes) : {},
+      // result 保持原样存储，不进行额外的 parse
+      // 如果原来是对象，会被序列化为字符串存储
+      // 如果原来就是字符串（已经是序列化的），也会保持原样
+      result: row.result || undefined,
+      error: row.error || undefined,
+    };
+  }
+
+  private buildTree(spans: Span[]): Span[] {
+    const spanMap = new Map<string, Span>();
+    const roots: Span[] = [];
+
+    for (const span of spans) {
+      spanMap.set(span.spanId, { ...span, children: [] });
+    }
+
+    for (const span of spanMap.values()) {
+      if (span.parentSpanId) {
+        const parent = spanMap.get(span.parentSpanId);
+        if (parent) {
+          parent.children!.push(span);
+        } else {
+          roots.push(span);
+        }
+      } else {
+        roots.push(span);
+      }
+    }
+
+    return roots;
+  }
+}

package/src/env/log-trace/trace-context.test.ts

@@ -0,0 +1,187 @@
+import { describe, test, expect, beforeEach, afterEach } from "bun:test";
+import {
+  getTraceContext,
+  initTraceContext,
+  type ITraceContext,
+} from "./trace-context";
+import { contextStore } from "./trace-context";
+
+describe("TraceContext", () => {
+  let traceCtx: ITraceContext;
+
+  beforeEach(() => {
+    // 重置为新的实例
+    traceCtx = initTraceContext();
+  });
+
+  afterEach(() => {
+    // 清理 AsyncLocalStorage 状态
+    contextStore.disable();
+  });
+
+  describe("generateRequestId", () => {
+    test("should generate unique request IDs", () => {
+      const id1 = traceCtx.generateRequestId();
+      const id2 = traceCtx.generateRequestId();
+      expect(id1).toBeTruthy();
+      expect(id2).toBeTruthy();
+      expect(id1).not.toBe(id2);
+    });
+
+    test("should start with req_ prefix", () => {
+      const id = traceCtx.generateRequestId();
+      expect(id.startsWith("req_")).toBe(true);
+    });
+  });
+
+  describe("initContext", () => {
+    test("should set requestId", () => {
+      const requestId = "test-request-123";
+      traceCtx.initContext(requestId);
+      expect(traceCtx.getRequestId()).toBe(requestId);
+    });
+
+    test("should set sessionId", () => {
+      const requestId = "test-request-123";
+      const sessionId = "session-456";
+      traceCtx.initContext(requestId, sessionId);
+      expect(traceCtx.getRequestId()).toBe(requestId);
+      expect(traceCtx.getSpanContext()?.traceId).toBe(requestId);
+    });
+  });
+
+  describe("runWithContext", () => {
+    test("should run function with current context", () => {
+      const requestId = "test-request-123";
+      traceCtx.initContext(requestId);
+
+      let capturedId: string | undefined;
+      traceCtx.runWithContext(() => {
+        capturedId = traceCtx.getRequestId();
+      });
+
+      expect(capturedId).toBe(requestId);
+    });
+
+    test("should preserve context across nested calls", () => {
+      const requestId = "test-request-123";
+      traceCtx.initContext(requestId);
+
+      let outerId: string | undefined;
+      let innerId: string | undefined;
+
+      traceCtx.runWithContext(() => {
+        outerId = traceCtx.getRequestId();
+        traceCtx.runWithContext(() => {
+          innerId = traceCtx.getRequestId();
+        });
+      });
+
+      expect(outerId).toBe(requestId);
+      expect(innerId).toBe(requestId);
+    });
+  });
+
+  describe("runWithNewContext", () => {
+    test("should run function in new context", () => {
+      const requestId1 = "request-1";
+      const requestId2 = "request-2";
+      traceCtx.initContext(requestId1);
+
+      let capturedId: string | undefined;
+      traceCtx.runWithNewContext(requestId2, undefined, () => {
+        capturedId = traceCtx.getRequestId();
+      });
+
+      expect(capturedId).toBe(requestId2);
+    });
+
+    test("should restore previous context after new context", () => {
+      const requestId1 = "request-1";
+      const requestId2 = "request-2";
+      traceCtx.initContext(requestId1);
+
+      let beforeId: string | undefined;
+      let duringId: string | undefined;
+      let afterId: string | undefined;
+
+      beforeId = traceCtx.getRequestId();
+      traceCtx.runWithNewContext(requestId2, undefined, () => {
+        duringId = traceCtx.getRequestId();
+      });
+      afterId = traceCtx.getRequestId();
+
+      expect(beforeId).toBe(requestId1);
+      expect(duringId).toBe(requestId2);
+      expect(afterId).toBe(requestId1);
+    });
+  });
+
+  describe("runWithContextAsync", () => {
+    test("should preserve context in async function", async () => {
+      const requestId = "async-request";
+      traceCtx.initContext(requestId);
+
+      let capturedId: string | undefined;
+      await traceCtx.runWithContextAsync(async () => {
+        await Promise.resolve();
+        capturedId = traceCtx.getRequestId();
+      });
+
+      expect(capturedId).toBe(requestId);
+    });
+  });
+
+  describe("runWithNewContextAsync", () => {
+    test("should run async function in new context", async () => {
+      const requestId1 = "request-1";
+      const requestId2 = "request-2";
+      traceCtx.initContext(requestId1);
+
+      let capturedId: string | undefined;
+      await traceCtx.runWithNewContextAsync(requestId2, undefined, async () => {
+        await Promise.resolve();
+        capturedId = traceCtx.getRequestId();
+      });
+
+      expect(capturedId).toBe(requestId2);
+    });
+  });
+
+  describe("startNewContext", () => {
+    test("should enter new context and return cleanup function", () => {
+      const requestId = "new-context";
+      
+      const cleanup = traceCtx.startNewContext(requestId);
+      expect(traceCtx.getRequestId()).toBe(requestId);
+      
+      cleanup();
+      // After cleanup, context should be undefined (no previous context was set)
+      // Note: AsyncLocalStorage.enterWith persists, so cleanup to previous state
+      expect(traceCtx.getRequestId()).toBeUndefined();
+    });
+
+    test("should restore previous context on cleanup", () => {
+      const requestId1 = "context-1";
+      const requestId2 = "context-2";
+      
+      traceCtx.initContext(requestId1);
+      const cleanup = traceCtx.startNewContext(requestId2);
+      cleanup();
+      
+      expect(traceCtx.getRequestId()).toBe(requestId1);
+    });
+  });
+
+  describe("getSpanContext", () => {
+    test("should return span context with traceId and spanId", () => {
+      const requestId = "span-context-test";
+      traceCtx.initContext(requestId);
+      
+      const spanCtx = traceCtx.getSpanContext();
+      expect(spanCtx).toBeDefined();
+      expect(spanCtx?.traceId).toBe(requestId);
+      expect(spanCtx?.spanId).toBeDefined();
+    });
+  });
+});