diff --git a/.agents/skills/agent-tracing/SKILL.md b/.agents/skills/agent-tracing/SKILL.md new file mode 100644 index 0000000000..d890514d95 --- /dev/null +++ b/.agents/skills/agent-tracing/SKILL.md @@ -0,0 +1,167 @@ +--- +name: agent-tracing +description: "Agent tracing CLI for inspecting agent execution snapshots. Use when user mentions 'agent-tracing', 'trace', 'snapshot', wants to debug agent execution, inspect LLM calls, view context engine data, or analyze agent steps. Triggers on agent debugging, trace inspection, or execution analysis tasks." +user-invocable: false +--- + +# Agent Tracing CLI Guide + +`@lobechat/agent-tracing` is a zero-config local dev tool that records agent execution snapshots to disk and provides a CLI to inspect them. + +## How It Works + +In `NODE_ENV=development`, `AgentRuntimeService.executeStep()` automatically records each step to `.agent-tracing/` as partial snapshots. When the operation completes, the partial is finalized into a complete `ExecutionSnapshot` JSON file. + +**Data flow**: executeStep loop -> build `StepPresentationData` -> write partial snapshot to disk -> on completion, finalize to `.agent-tracing/{timestamp}_{traceId}.json` + +**Context engine capture**: In `RuntimeExecutors.ts`, the `call_llm` executor emits a `context_engine_result` event after `serverMessagesEngine()` processes messages. This event carries the full `contextEngineInput` (DB messages, systemRole, model, knowledge, tools, userMemory, etc.) and the processed `output` messages (the final LLM payload). + +## Package Location + +``` +packages/agent-tracing/ + src/ + types.ts # ExecutionSnapshot, StepSnapshot, SnapshotSummary + store/ + types.ts # ISnapshotStore interface + file-store.ts # FileSnapshotStore (.agent-tracing/*.json) + recorder/ + index.ts # appendStepToPartial(), finalizeSnapshot() + viewer/ + index.ts # Terminal rendering: renderSnapshot, renderStepDetail, renderMessageDetail, renderSummaryTable + cli/ + index.ts # CLI entry point (#!/usr/bin/env bun) + index.ts # Barrel exports +``` + +## Data Storage + +- Completed snapshots: `.agent-tracing/{ISO-timestamp}_{traceId-short}.json` +- Latest symlink: `.agent-tracing/latest.json` +- In-progress partials: `.agent-tracing/_partial/{operationId}.json` +- `FileSnapshotStore` resolves from `process.cwd()` — **run CLI from the repo root** + +## CLI Commands + +All commands run from the **repo root**: + +```bash +# View latest trace (tree overview) +agent-tracing trace + +# View specific trace +agent-tracing trace + +# List recent snapshots +agent-tracing list +agent-tracing list -l 20 + +# Inspect trace detail (overview) +agent-tracing inspect + +# Inspect specific step (-s is short for --step) +agent-tracing inspect -s 0 + +# View messages (-m is short for --messages) +agent-tracing inspect -s 0 -m + +# View full content of a specific message (by index shown in -m output) +agent-tracing inspect -s 0 --msg 2 +agent-tracing inspect -s 0 --msg-input 1 + +# View tool call/result details (-t is short for --tools) +agent-tracing inspect -s 1 -t + +# View raw events (-e is short for --events) +agent-tracing inspect -s 0 -e + +# View runtime context (-c is short for --context) +agent-tracing inspect -s 0 -c + +# Raw JSON output (-j is short for --json) +agent-tracing inspect -j +agent-tracing inspect -s 0 -j +``` + +## Typical Debug Workflow + +```bash +# 1. Trigger an agent operation in the dev UI + +# 2. See the overview +agent-tracing trace + +# 3. List all traces, get traceId +agent-tracing list + +# 4. Inspect a specific step's messages to see what was sent to the LLM +agent-tracing inspect TRACE_ID -s 0 -m + +# 5. Drill into a truncated message for full content +agent-tracing inspect TRACE_ID -s 0 --msg 2 + +# 6. Check tool calls and results +agent-tracing inspect 1 -t TRACE_ID -s +``` + +## Key Types + +```typescript +interface ExecutionSnapshot { + traceId: string; + operationId: string; + model?: string; + provider?: string; + startedAt: number; + completedAt?: number; + completionReason?: + | 'done' + | 'error' + | 'interrupted' + | 'max_steps' + | 'cost_limit' + | 'waiting_for_human'; + totalSteps: number; + totalTokens: number; + totalCost: number; + error?: { type: string; message: string }; + steps: StepSnapshot[]; +} + +interface StepSnapshot { + stepIndex: number; + stepType: 'call_llm' | 'call_tool'; + executionTimeMs: number; + content?: string; // LLM output + reasoning?: string; // Reasoning/thinking + inputTokens?: number; + outputTokens?: number; + toolsCalling?: Array<{ apiName: string; identifier: string; arguments?: string }>; + toolsResult?: Array<{ + apiName: string; + identifier: string; + isSuccess?: boolean; + output?: string; + }>; + messages?: any[]; // DB messages before step + context?: { phase: string; payload?: unknown; stepContext?: unknown }; + events?: Array<{ type: string; [key: string]: unknown }>; + // context_engine_result event contains: + // input: full contextEngineInput (messages, systemRole, model, knowledge, tools, userMemory, ...) + // output: processed messages array (final LLM payload) +} +``` + +## --messages Output Structure + +When using `--messages`, the output shows three sections (if context engine data is available): + +1. **Context Engine Input** — DB messages passed to the engine, with `[0]`, `[1]`, ... indices. Use `--msg-input N` to view full content. +2. **Context Engine Params** — systemRole, model, provider, knowledge, tools, userMemory, etc. +3. **Final LLM Payload** — Processed messages after context engine (system date injection, user memory, history truncation, etc.), with `[0]`, `[1]`, ... indices. Use `--msg N` to view full content. + +## Integration Points + +- **Recording**: `src/server/services/agentRuntime/AgentRuntimeService.ts` — in the `executeStep()` method, after building `stepPresentationData`, writes partial snapshot in dev mode +- **Context engine event**: `src/server/modules/AgentRuntime/RuntimeExecutors.ts` — in `call_llm` executor, after `serverMessagesEngine()` returns, emits `context_engine_result` event +- **Store**: `FileSnapshotStore` reads/writes to `.agent-tracing/` relative to `process.cwd()` diff --git a/.github/workflows/deploy-device-gateway.yml b/.github/workflows/deploy-device-gateway.yml new file mode 100644 index 0000000000..8b1e1f29f5 --- /dev/null +++ b/.github/workflows/deploy-device-gateway.yml @@ -0,0 +1,39 @@ +name: Deploy Device Gateway + +permissions: + contents: read + +on: + push: + branches: [canary] + paths: + - 'apps/device-gateway/**' + workflow_dispatch: + +concurrency: + group: ${{ github.workflow }}-${{ github.ref }} + cancel-in-progress: true + +jobs: + deploy: + name: Deploy to Cloudflare + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v6 + + - uses: actions/setup-node@v6 + with: + node-version: 22 + + - uses: oven-sh/setup-bun@v2 + + - run: cd apps/device-gateway && bun i + + - name: Type check + run: cd apps/device-gateway && npx tsc --noEmit + + - name: Deploy + run: cd apps/device-gateway && npx wrangler deploy + env: + CLOUDFLARE_API_TOKEN: ${{ secrets.CLOUDFLARE_API_TOKEN }} + CLOUDFLARE_ACCOUNT_ID: ${{ secrets.CLOUDFLARE_ACCOUNT_ID }} diff --git a/.gitignore b/.gitignore index e3ca9b89e9..7b1a917be1 100644 --- a/.gitignore +++ b/.gitignore @@ -100,6 +100,9 @@ robots.txt # Cloud service keys vertex-ai-key.json +# Agent tracing snapshots +.agent-tracing/ + # AI coding tools .local/ .claude/ diff --git a/apps/device-gateway/package.json b/apps/device-gateway/package.json new file mode 100644 index 0000000000..d8dfad632e --- /dev/null +++ b/apps/device-gateway/package.json @@ -0,0 +1,18 @@ +{ + "name": "@lobechat/device-gateway", + "version": "0.1.0", + "private": true, + "scripts": { + "dev": "wrangler dev", + "deploy": "wrangler deploy", + "type-check": "tsc --noEmit" + }, + "dependencies": { + "jose": "^6.1.3" + }, + "devDependencies": { + "@cloudflare/workers-types": "^4.20250214.0", + "typescript": "^5.9.3", + "wrangler": "^4.14.4" + } +} diff --git a/apps/device-gateway/scripts/extract-public-key.mjs b/apps/device-gateway/scripts/extract-public-key.mjs new file mode 100755 index 0000000000..042dbe74d4 --- /dev/null +++ b/apps/device-gateway/scripts/extract-public-key.mjs @@ -0,0 +1,48 @@ +#!/usr/bin/env node + +/** + * Extract RS256 public key from JWKS_KEY environment variable. + * Output is the JSON string to use with `wrangler secret put JWKS_PUBLIC_KEY`. + * + * Usage: + * JWKS_KEY='{"keys":[...]}' node scripts/extract-public-key.mjs + * # or load from .env + * node --env-file=../../.env scripts/extract-public-key.mjs + */ + +const jwksString = process.env.JWKS_KEY; + +if (!jwksString) { + console.error('Error: JWKS_KEY environment variable is not set.'); + process.exit(1); +} + +const jwks = JSON.parse(jwksString); +const privateKey = jwks.keys?.find((k) => k.alg === 'RS256' && k.kty === 'RSA'); + +if (!privateKey) { + console.error('Error: No RS256 RSA key found in JWKS_KEY.'); + process.exit(1); +} + +const publicJwks = { + keys: [ + { + alg: privateKey.alg, + e: privateKey.e, + kid: privateKey.kid, + kty: privateKey.kty, + n: privateKey.n, + use: privateKey.use, + }, + ], +}; + +// Remove undefined fields +for (const key of publicJwks.keys) { + for (const [k, v] of Object.entries(key)) { + if (v === undefined) delete key[k]; + } +} + +console.log(JSON.stringify(publicJwks)); diff --git a/apps/device-gateway/src/DeviceGatewayDO.ts b/apps/device-gateway/src/DeviceGatewayDO.ts new file mode 100644 index 0000000000..3e909366f6 --- /dev/null +++ b/apps/device-gateway/src/DeviceGatewayDO.ts @@ -0,0 +1,147 @@ +import { DurableObject } from 'cloudflare:workers'; + +import { DeviceAttachment, Env } from './types'; + +export class DeviceGatewayDO extends DurableObject { + private pendingRequests = new Map< + string, + { + resolve: (result: any) => void; + timer: ReturnType; + } + >(); + + async fetch(request: Request): Promise { + const url = new URL(request.url); + + // ─── WebSocket upgrade (from Desktop) ─── + if (request.headers.get('Upgrade') === 'websocket') { + const pair = new WebSocketPair(); + const [client, server] = Object.values(pair); + + this.ctx.acceptWebSocket(server); + + const deviceId = url.searchParams.get('deviceId') || 'unknown'; + const hostname = url.searchParams.get('hostname') || ''; + const platform = url.searchParams.get('platform') || ''; + + server.serializeAttachment({ + connectedAt: Date.now(), + deviceId, + hostname, + platform, + } satisfies DeviceAttachment); + + return new Response(null, { status: 101, webSocket: client }); + } + + // ─── HTTP API (from Vercel Agent) ─── + if (url.pathname === '/api/device/status') { + const sockets = this.ctx.getWebSockets(); + return Response.json({ + deviceCount: sockets.length, + online: sockets.length > 0, + }); + } + + if (url.pathname === '/api/device/tool-call') { + return this.handleToolCall(request); + } + + if (url.pathname === '/api/device/devices') { + const sockets = this.ctx.getWebSockets(); + const devices = sockets.map((ws) => ws.deserializeAttachment() as DeviceAttachment); + return Response.json({ devices }); + } + + return new Response('Not Found', { status: 404 }); + } + + // ─── Hibernation Handlers ─── + + async webSocketMessage(ws: WebSocket, message: string | ArrayBuffer) { + const data = JSON.parse(message as string); + + if (data.type === 'tool_call_response') { + const pending = this.pendingRequests.get(data.requestId); + if (pending) { + clearTimeout(pending.timer); + pending.resolve(data.result); + this.pendingRequests.delete(data.requestId); + } + } + + if (data.type === 'heartbeat') { + ws.send(JSON.stringify({ type: 'heartbeat_ack' })); + } + } + + async webSocketClose(_ws: WebSocket, _code: number) { + // Hibernation API handles connection cleanup automatically + } + + async webSocketError(ws: WebSocket, _error: unknown) { + ws.close(1011, 'Internal error'); + } + + // ─── Tool Call RPC ─── + + private async handleToolCall(request: Request): Promise { + const sockets = this.ctx.getWebSockets(); + if (sockets.length === 0) { + return Response.json( + { content: '桌面设备不在线', error: 'DEVICE_OFFLINE', success: false }, + { status: 503 }, + ); + } + + const { deviceId, timeout = 30_000, toolCall } = (await request.json()) as { + deviceId?: string; + timeout?: number; + toolCall: unknown; + }; + const requestId = crypto.randomUUID(); + + // Select target device (specified > first available) + const targetWs = deviceId + ? sockets.find((ws) => { + const att = ws.deserializeAttachment() as DeviceAttachment; + return att.deviceId === deviceId; + }) + : sockets[0]; + + if (!targetWs) { + return Response.json({ error: 'DEVICE_NOT_FOUND', success: false }, { status: 503 }); + } + + try { + const result = await new Promise((resolve, reject) => { + const timer = setTimeout(() => { + this.pendingRequests.delete(requestId); + reject(new Error('TIMEOUT')); + }, timeout); + + this.pendingRequests.set(requestId, { resolve, timer }); + + targetWs.send( + JSON.stringify({ + requestId, + toolCall, + type: 'tool_call_request', + }), + ); + }); + + return Response.json({ success: true, ...(result as object) }); + } catch (err) { + return Response.json( + { + content: `工具调用超时（${timeout / 1000}s）`, + error: (err as Error).message, + success: false, + }, + { status: 504 }, + ); + } + } +} diff --git a/apps/device-gateway/src/auth.ts b/apps/device-gateway/src/auth.ts new file mode 100644 index 0000000000..61752ee1e0 --- /dev/null +++ b/apps/device-gateway/src/auth.ts @@ -0,0 +1,36 @@ +import { importJWK, jwtVerify } from 'jose'; + +import { Env } from './types'; + +let cachedKey: CryptoKey | null = null; + +async function getPublicKey(env: Env): Promise { + if (cachedKey) return cachedKey; + + const jwks = JSON.parse(env.JWKS_PUBLIC_KEY); + const rsaKey = jwks.keys.find((k: any) => k.alg === 'RS256'); + + if (!rsaKey) { + throw new Error('No RS256 key found in JWKS_PUBLIC_KEY'); + } + + cachedKey = (await importJWK(rsaKey, 'RS256')) as CryptoKey; + return cachedKey; +} + +export async function verifyDesktopToken( + env: Env, + token: string, +): Promise<{ clientId: string; userId: string }> { + const publicKey = await getPublicKey(env); + const { payload } = await jwtVerify(token, publicKey, { + algorithms: ['RS256'], + }); + + if (!payload.sub) throw new Error('Missing sub claim'); + + return { + clientId: payload.client_id as string, + userId: payload.sub, + }; +} diff --git a/apps/device-gateway/src/index.ts b/apps/device-gateway/src/index.ts new file mode 100644 index 0000000000..c766c0a871 --- /dev/null +++ b/apps/device-gateway/src/index.ts @@ -0,0 +1,51 @@ +import { verifyDesktopToken } from './auth'; +import { DeviceGatewayDO } from './DeviceGatewayDO'; +import { Env } from './types'; + +export { DeviceGatewayDO }; + +export default { + async fetch(request: Request, env: Env): Promise { + const url = new URL(request.url); + + // ─── Health check ─── + if (url.pathname === '/health') { + return new Response('OK', { status: 200 }); + } + + // ─── Desktop WebSocket connection ─── + if (url.pathname === '/ws') { + const token = url.searchParams.get('token'); + if (!token) return new Response('Missing token', { status: 401 }); + + try { + const { userId } = await verifyDesktopToken(env, token); + + const id = env.DEVICE_GATEWAY.idFromName(`user:${userId}`); + const stub = env.DEVICE_GATEWAY.get(id); + + // Forward WebSocket upgrade to DO + const headers = new Headers(request.headers); + headers.set('X-User-Id', userId); + return stub.fetch(new Request(request, { headers })); + } catch { + return new Response('Invalid token', { status: 401 }); + } + } + + // ─── Vercel Agent HTTP API ─── + if (url.pathname.startsWith('/api/device/')) { + const authHeader = request.headers.get('Authorization'); + if (authHeader !== `Bearer ${env.SERVICE_TOKEN}`) { + return new Response('Unauthorized', { status: 401 }); + } + + const body = (await request.clone().json()) as { userId: string }; + const id = env.DEVICE_GATEWAY.idFromName(`user:${body.userId}`); + const stub = env.DEVICE_GATEWAY.get(id); + return stub.fetch(request); + } + + return new Response('Not Found', { status: 404 }); + }, +}; diff --git a/apps/device-gateway/src/types.ts b/apps/device-gateway/src/types.ts new file mode 100644 index 0000000000..4b0eb051a4 --- /dev/null +++ b/apps/device-gateway/src/types.ts @@ -0,0 +1,53 @@ +export interface Env { + DEVICE_GATEWAY: DurableObjectNamespace; + JWKS_PUBLIC_KEY: string; + SERVICE_TOKEN: string; +} + +// ─── Device Info ─── + +export interface DeviceAttachment { + connectedAt: number; + deviceId: string; + hostname: string; + platform: string; +} + +// ─── WebSocket Protocol Messages ─── + +// Desktop → CF +export interface HeartbeatMessage { + type: 'heartbeat'; +} + +export interface ToolCallResponseMessage { + requestId: string; + result: { + content: string; + error?: string; + success: boolean; + }; + type: 'tool_call_response'; +} + +// CF → Desktop +export interface HeartbeatAckMessage { + type: 'heartbeat_ack'; +} + +export interface AuthExpiredMessage { + type: 'auth_expired'; +} + +export interface ToolCallRequestMessage { + requestId: string; + toolCall: { + apiName: string; + arguments: string; + identifier: string; + }; + type: 'tool_call_request'; +} + +export type ClientMessage = HeartbeatMessage | ToolCallResponseMessage; +export type ServerMessage = AuthExpiredMessage | HeartbeatAckMessage | ToolCallRequestMessage; diff --git a/apps/device-gateway/tsconfig.json b/apps/device-gateway/tsconfig.json new file mode 100644 index 0000000000..d63ef5e826 --- /dev/null +++ b/apps/device-gateway/tsconfig.json @@ -0,0 +1,17 @@ +{ + "compilerOptions": { + "target": "ESNext", + "module": "ESNext", + "moduleResolution": "bundler", + "lib": ["ESNext"], + "types": ["@cloudflare/workers-types"], + "strict": true, + "noEmit": true, + "skipLibCheck": true, + "esModuleInterop": true, + "forceConsistentCasingInFileNames": true, + "resolveJsonModule": true, + "isolatedModules": true + }, + "include": ["src"] +} diff --git a/apps/device-gateway/wrangler.toml b/apps/device-gateway/wrangler.toml new file mode 100644 index 0000000000..32863b74df --- /dev/null +++ b/apps/device-gateway/wrangler.toml @@ -0,0 +1,16 @@ +name = "device-gateway" +main = "src/index.ts" +compatibility_date = "2025-01-01" + +[durable_objects] +bindings = [ + { name = "DEVICE_GATEWAY", class_name = "DeviceGatewayDO" } +] + +[[migrations]] +tag = "v1" +new_classes = ["DeviceGatewayDO"] + +# Secrets (injected via `wrangler secret put`): +# - JWKS_PUBLIC_KEY: RS256 public key JSON (extracted from JWKS_KEY) +# - SERVICE_TOKEN: Vercel → CF service-to-service auth secret diff --git a/eslint.config.mjs b/eslint.config.mjs index 55ea9a87d6..f798785625 100644 --- a/eslint.config.mjs +++ b/eslint.config.mjs @@ -133,4 +133,11 @@ export default eslint( 'no-console': 0, }, }, + // agent-tracing CLI - console output is the primary interface + { + files: ['packages/agent-tracing/**/*'], + rules: { + 'no-console': 0, + }, + }, ); diff --git a/package.json b/package.json index 2a26d5d127..04eac49459 100644 --- a/package.json +++ b/package.json @@ -399,6 +399,7 @@ "@edge-runtime/vm": "^5.0.0", "@huggingface/tasks": "^0.19.80", "@inquirer/prompts": "^8.2.0", + "@lobechat/agent-tracing": "workspace:*", "@lobechat/types": "workspace:*", "@lobehub/i18n-cli": "^1.26.0", "@lobehub/lint": "2.1.5", diff --git a/packages/agent-tracing/package.json b/packages/agent-tracing/package.json new file mode 100644 index 0000000000..4aea841625 --- /dev/null +++ b/packages/agent-tracing/package.json @@ -0,0 +1,17 @@ +{ + "name": "@lobechat/agent-tracing", + "version": "1.0.0", + "private": true, + "exports": { + ".": "./src/index.ts", + "./store": "./src/store/types.ts", + "./recorder": "./src/recorder/index.ts" + }, + "main": "./src/index.ts", + "bin": { + "agent-tracing": "./src/cli/index.ts" + }, + "dependencies": { + "commander": "^13.1.0" + } +} diff --git a/packages/agent-tracing/src/cli/index.ts b/packages/agent-tracing/src/cli/index.ts new file mode 100755 index 0000000000..f25dcd1928 --- /dev/null +++ b/packages/agent-tracing/src/cli/index.ts @@ -0,0 +1,17 @@ +#!/usr/bin/env bun + +import { Command } from 'commander'; + +import { registerInspectCommand } from './inspect'; +import { registerListCommand } from './list'; +import { registerTraceCommand } from './trace'; + +const program = new Command(); + +program.name('agent-tracing').description('Local agent execution snapshot viewer').version('1.0.0'); + +registerTraceCommand(program); +registerListCommand(program); +registerInspectCommand(program); + +program.parse(); diff --git a/packages/agent-tracing/src/cli/inspect.ts b/packages/agent-tracing/src/cli/inspect.ts new file mode 100644 index 0000000000..a421fedb56 --- /dev/null +++ b/packages/agent-tracing/src/cli/inspect.ts @@ -0,0 +1,101 @@ +import type { Command } from 'commander'; + +import { FileSnapshotStore } from '../store/file-store'; +import { renderMessageDetail, renderSnapshot, renderStepDetail } from '../viewer'; + +export function registerInspectCommand(program: Command) { + program + .command('inspect') + .description('Inspect trace details') + .argument('', 'Trace ID to inspect') + .option('-s, --step ', 'View specific step') + .option('-m, --messages', 'Show messages context') + .option('-t, --tools', 'Show tool call details') + .option('-e, --events', 'Show raw events (llm_start, llm_result, etc.)') + .option('-c, --context', 'Show runtime context & payload') + .option( + '--msg ', + 'Show full content of message [N] from Final LLM Payload (use with --step)', + ) + .option( + '--msg-input ', + 'Show full content of message [N] from Context Engine Input (use with --step)', + ) + .option('-j, --json', 'Output as JSON') + .action( + async ( + traceId: string, + opts: { + context?: boolean; + events?: boolean; + json?: boolean; + messages?: boolean; + msg?: string; + msgInput?: string; + step?: string; + tools?: boolean; + }, + ) => { + const store = new FileSnapshotStore(); + const snapshot = await store.get(traceId); + if (!snapshot) { + console.error(`Snapshot not found: ${traceId}`); + process.exit(1); + } + + const stepIndex = opts.step !== undefined ? Number.parseInt(opts.step, 10) : undefined; + + if (opts.json) { + if (stepIndex !== undefined) { + const step = snapshot.steps.find((s) => s.stepIndex === stepIndex); + console.log(JSON.stringify(step ?? null, null, 2)); + } else { + console.log(JSON.stringify(snapshot, null, 2)); + } + return; + } + + // --msg or --msg-input: show full message detail + const msgIndex = + opts.msg !== undefined + ? Number.parseInt(opts.msg, 10) + : opts.msgInput !== undefined + ? Number.parseInt(opts.msgInput, 10) + : undefined; + const msgSource: 'input' | 'output' = opts.msgInput !== undefined ? 'input' : 'output'; + + if (msgIndex !== undefined && stepIndex !== undefined) { + const step = snapshot.steps.find((s) => s.stepIndex === stepIndex); + if (!step) { + console.error( + `Step ${stepIndex} not found. Available: ${snapshot.steps.map((s) => s.stepIndex).join(', ')}`, + ); + process.exit(1); + } + console.log(renderMessageDetail(step, msgIndex, msgSource)); + return; + } + + if (stepIndex !== undefined) { + const step = snapshot.steps.find((s) => s.stepIndex === stepIndex); + if (!step) { + console.error( + `Step ${stepIndex} not found. Available: ${snapshot.steps.map((s) => s.stepIndex).join(', ')}`, + ); + process.exit(1); + } + console.log( + renderStepDetail(step, { + context: opts.context, + events: opts.events, + messages: opts.messages, + tools: opts.tools, + }), + ); + return; + } + + console.log(renderSnapshot(snapshot)); + }, + ); +} diff --git a/packages/agent-tracing/src/cli/list.ts b/packages/agent-tracing/src/cli/list.ts new file mode 100644 index 0000000000..1da0fb30bb --- /dev/null +++ b/packages/agent-tracing/src/cli/list.ts @@ -0,0 +1,17 @@ +import type { Command } from 'commander'; + +import { FileSnapshotStore } from '../store/file-store'; +import { renderSummaryTable } from '../viewer'; + +export function registerListCommand(program: Command) { + program + .command('list') + .description('List recent snapshots') + .option('-l, --limit ', 'Max number of snapshots to show', '10') + .action(async (opts: { limit: string }) => { + const store = new FileSnapshotStore(); + let limit = Number.parseInt(opts.limit, 10); + if (Number.isNaN(limit) || limit < 1) limit = 10; + console.log(renderSummaryTable(await store.list({ limit }))); + }); +} diff --git a/packages/agent-tracing/src/cli/trace.ts b/packages/agent-tracing/src/cli/trace.ts new file mode 100644 index 0000000000..b8c7a71c6a --- /dev/null +++ b/packages/agent-tracing/src/cli/trace.ts @@ -0,0 +1,24 @@ +import type { Command } from 'commander'; + +import { FileSnapshotStore } from '../store/file-store'; +import { renderSnapshot } from '../viewer'; + +export function registerTraceCommand(program: Command) { + program + .command('trace') + .description('View latest or specific trace') + .argument('[traceId]', 'Trace ID to view (defaults to latest)') + .action(async (traceId?: string) => { + const store = new FileSnapshotStore(); + const snapshot = traceId ? await store.get(traceId) : await store.getLatest(); + if (!snapshot) { + console.error( + traceId + ? `Snapshot not found: ${traceId}` + : 'No snapshots found. Run an agent operation first.', + ); + process.exit(1); + } + console.log(renderSnapshot(snapshot)); + }); +} diff --git a/packages/agent-tracing/src/index.ts b/packages/agent-tracing/src/index.ts new file mode 100644 index 0000000000..7d57482aad --- /dev/null +++ b/packages/agent-tracing/src/index.ts @@ -0,0 +1,10 @@ +export { appendStepToPartial, finalizeSnapshot } from './recorder'; +export { FileSnapshotStore } from './store/file-store'; +export type { ISnapshotStore } from './store/types'; +export type { ExecutionSnapshot, SnapshotSummary, StepSnapshot } from './types'; +export { + renderMessageDetail, + renderSnapshot, + renderStepDetail, + renderSummaryTable, +} from './viewer'; diff --git a/packages/agent-tracing/src/recorder/index.ts b/packages/agent-tracing/src/recorder/index.ts new file mode 100644 index 0000000000..ba5748509f --- /dev/null +++ b/packages/agent-tracing/src/recorder/index.ts @@ -0,0 +1,63 @@ +import type { ISnapshotStore } from '../store/types'; +import type { ExecutionSnapshot, StepSnapshot } from '../types'; + +/** + * Append a step to a partial snapshot on disk. + * Called from the executeStep loop on each step completion. + */ +export async function appendStepToPartial( + store: ISnapshotStore, + operationId: string, + step: StepSnapshot, + metadata?: { model?: string; provider?: string }, +): Promise { + const partial = (await store.loadPartial(operationId)) ?? { steps: [] }; + + if (!partial.startedAt) { + partial.startedAt = Date.now(); + partial.model = metadata?.model; + partial.provider = metadata?.provider; + } + + if (!partial.steps) partial.steps = []; + partial.steps.push(step); + + await store.savePartial(operationId, partial); +} + +/** + * Finalize a partial snapshot into a completed ExecutionSnapshot. + * Called from the executeStep loop when the operation completes. + */ +export async function finalizeSnapshot( + store: ISnapshotStore, + operationId: string, + completion: { + error?: { message: string; type: string }; + reason: string; + totalCost: number; + totalSteps: number; + totalTokens: number; + }, +): Promise { + const partial = await store.loadPartial(operationId); + if (!partial) return; + + const snapshot: ExecutionSnapshot = { + completedAt: Date.now(), + completionReason: completion.reason as ExecutionSnapshot['completionReason'], + error: completion.error, + model: partial.model, + operationId, + provider: partial.provider, + startedAt: partial.startedAt ?? Date.now(), + steps: (partial.steps ?? []).sort((a, b) => a.stepIndex - b.stepIndex), + totalCost: completion.totalCost, + totalSteps: completion.totalSteps, + totalTokens: completion.totalTokens, + traceId: operationId, + }; + + await store.save(snapshot); + await store.removePartial(operationId); +} diff --git a/packages/agent-tracing/src/store/file-store.ts b/packages/agent-tracing/src/store/file-store.ts new file mode 100644 index 0000000000..f980b705f7 --- /dev/null +++ b/packages/agent-tracing/src/store/file-store.ts @@ -0,0 +1,144 @@ +import fs from 'node:fs/promises'; +import path from 'node:path'; + +import type { ExecutionSnapshot, SnapshotSummary } from '../types'; +import type { ISnapshotStore } from './types'; + +const DEFAULT_DIR = '.agent-tracing'; +const PARTIAL_DIR = '_partial'; + +export class FileSnapshotStore implements ISnapshotStore { + private dir: string; + + constructor(rootDir?: string) { + this.dir = path.resolve(rootDir ?? process.cwd(), DEFAULT_DIR); + } + + // ==================== Completed snapshots ==================== + + async save(snapshot: ExecutionSnapshot): Promise { + await fs.mkdir(this.dir, { recursive: true }); + + const ts = new Date(snapshot.startedAt).toISOString().replaceAll(':', '-'); + const shortId = snapshot.traceId.slice(0, 12); + const filename = `${ts}_${shortId}.json`; + const filePath = path.join(this.dir, filename); + + await fs.writeFile(filePath, JSON.stringify(snapshot, null, 2), 'utf8'); + + // Update latest symlink + const latestPath = path.join(this.dir, 'latest.json'); + try { + await fs.unlink(latestPath); + } catch { + // ignore if doesn't exist + } + await fs.symlink(filename, latestPath); + } + + async get(traceId: string): Promise { + const files = await this.listFiles(); + const match = files.find((f) => f.includes(traceId.slice(0, 12))); + if (!match) return null; + + const content = await fs.readFile(path.join(this.dir, match), 'utf8'); + return JSON.parse(content) as ExecutionSnapshot; + } + + async list(options?: { limit?: number }): Promise { + const files = await this.listFiles(); + const limit = options?.limit ?? 10; + const recent = files.slice(0, limit); + + const summaries: SnapshotSummary[] = []; + + for (const file of recent) { + try { + const content = await fs.readFile(path.join(this.dir, file), 'utf8'); + const snapshot = JSON.parse(content) as ExecutionSnapshot; + summaries.push(toSummary(snapshot)); + } catch { + // skip corrupted files + } + } + + return summaries; + } + + async getLatest(): Promise { + const latestPath = path.join(this.dir, 'latest.json'); + try { + const realPath = await fs.realpath(latestPath); + const content = await fs.readFile(realPath, 'utf8'); + return JSON.parse(content) as ExecutionSnapshot; + } catch { + // No latest symlink — fall back to most recent by filename + const files = await this.listFiles(); + if (files.length === 0) return null; + + const content = await fs.readFile(path.join(this.dir, files[0]), 'utf8'); + return JSON.parse(content) as ExecutionSnapshot; + } + } + + // ==================== Partial snapshots ==================== + + private partialDir(): string { + return path.join(this.dir, PARTIAL_DIR); + } + + private partialPath(operationId: string): string { + const safe = operationId.replaceAll('/', '_'); + return path.join(this.partialDir(), `${safe}.json`); + } + + async loadPartial(operationId: string): Promise | null> { + try { + const content = await fs.readFile(this.partialPath(operationId), 'utf8'); + return JSON.parse(content) as Partial; + } catch { + return null; + } + } + + async savePartial(operationId: string, partial: Partial): Promise { + await fs.mkdir(this.partialDir(), { recursive: true }); + await fs.writeFile(this.partialPath(operationId), JSON.stringify(partial), 'utf8'); + } + + async removePartial(operationId: string): Promise { + try { + await fs.unlink(this.partialPath(operationId)); + } catch { + // ignore + } + } + + // ==================== Internal ==================== + + private async listFiles(): Promise { + try { + const entries = await fs.readdir(this.dir); + return entries + .filter((f) => f.endsWith('.json') && f !== 'latest.json') + .sort() + .reverse(); // newest first (ISO timestamp prefix) + } catch { + return []; + } + } +} + +function toSummary(snapshot: ExecutionSnapshot): SnapshotSummary { + return { + completionReason: snapshot.completionReason, + createdAt: snapshot.startedAt, + durationMs: (snapshot.completedAt ?? Date.now()) - snapshot.startedAt, + hasError: !!snapshot.error, + model: snapshot.model, + operationId: snapshot.operationId, + totalSteps: snapshot.totalSteps, + totalTokens: snapshot.totalTokens, + traceId: snapshot.traceId, + }; +} diff --git a/packages/agent-tracing/src/store/types.ts b/packages/agent-tracing/src/store/types.ts new file mode 100644 index 0000000000..4ccc0b723b --- /dev/null +++ b/packages/agent-tracing/src/store/types.ts @@ -0,0 +1,15 @@ +import type { ExecutionSnapshot, SnapshotSummary } from '../types'; + +export interface ISnapshotStore { + get: (traceId: string) => Promise; + getLatest: () => Promise; + list: (options?: { limit?: number }) => Promise; + /** Load in-progress partial snapshot */ + loadPartial: (operationId: string) => Promise | null>; + + /** Remove partial snapshot (after finalizing) */ + removePartial: (operationId: string) => Promise; + save: (snapshot: ExecutionSnapshot) => Promise; + /** Save in-progress partial snapshot */ + savePartial: (operationId: string, partial: Partial) => Promise; +} diff --git a/packages/agent-tracing/src/types.ts b/packages/agent-tracing/src/types.ts new file mode 100644 index 0000000000..2324e176b5 --- /dev/null +++ b/packages/agent-tracing/src/types.ts @@ -0,0 +1,73 @@ +export interface ExecutionSnapshot { + completedAt?: number; + completionReason?: + | 'done' + | 'error' + | 'interrupted' + | 'max_steps' + | 'cost_limit' + | 'waiting_for_human'; + error?: { type: string; message: string }; + model?: string; + operationId: string; + provider?: string; + startedAt: number; + steps: StepSnapshot[]; + totalCost: number; + totalSteps: number; + totalTokens: number; + traceId: string; +} + +export interface StepSnapshot { + completedAt: number; + // LLM data + content?: string; + context?: { + phase: string; + payload?: unknown; + stepContext?: unknown; + }; + events?: Array<{ type: string; [key: string]: unknown }>; + executionTimeMs: number; + + inputTokens?: number; + // Detailed data (for inspect --step N) + messages?: any[]; + messagesAfter?: any[]; + outputTokens?: number; + + reasoning?: string; + startedAt: number; + + stepIndex: number; + stepType: 'call_llm' | 'call_tool'; + + // Tool data + toolsCalling?: Array<{ + apiName: string; + identifier: string; + arguments?: string; + }>; + toolsResult?: Array<{ + apiName: string; + identifier: string; + isSuccess?: boolean; + output?: string; + }>; + totalCost: number; + // Cumulative + totalTokens: number; +} + +export interface SnapshotSummary { + completionReason?: string; + createdAt: number; + durationMs: number; + hasError: boolean; + model?: string; + operationId: string; + totalSteps: number; + totalTokens: number; + traceId: string; +} diff --git a/packages/agent-tracing/src/viewer/index.ts b/packages/agent-tracing/src/viewer/index.ts new file mode 100644 index 0000000000..73ef989cf7 --- /dev/null +++ b/packages/agent-tracing/src/viewer/index.ts @@ -0,0 +1,433 @@ +import type { ExecutionSnapshot, SnapshotSummary, StepSnapshot } from '../types'; + +// ANSI color helpers +const dim = (s: string) => `\x1B[2m${s}\x1B[22m`; +const bold = (s: string) => `\x1B[1m${s}\x1B[22m`; +const green = (s: string) => `\x1B[32m${s}\x1B[39m`; +const red = (s: string) => `\x1B[31m${s}\x1B[39m`; +const yellow = (s: string) => `\x1B[33m${s}\x1B[39m`; +const cyan = (s: string) => `\x1B[36m${s}\x1B[39m`; +const magenta = (s: string) => `\x1B[35m${s}\x1B[39m`; + +function formatMs(ms: number): string { + if (ms < 1000) return `${ms}ms`; + return `${(ms / 1000).toFixed(1)}s`; +} + +function formatTokens(n: number): string { + if (n >= 1000) return `${(n / 1000).toFixed(1)}k`; + return String(n); +} + +function formatCost(cost: number): string { + if (cost === 0) return '$0'; + if (cost < 0.001) return `$${cost.toFixed(6)}`; + return `$${cost.toFixed(4)}`; +} + +function truncate(s: string, maxLen: number): string { + const single = s.replaceAll('\n', ' '); + if (single.length <= maxLen) return single; + return single.slice(0, maxLen - 3) + '...'; +} + +function padEnd(s: string, len: number): string { + if (s.length >= len) return s; + return s + ' '.repeat(len - s.length); +} + +export function renderSnapshot(snapshot: ExecutionSnapshot): string { + const lines: string[] = []; + const durationMs = (snapshot.completedAt ?? Date.now()) - snapshot.startedAt; + const shortId = snapshot.traceId.slice(0, 12); + + // Header + lines.push( + bold('Agent Operation') + + ` ${cyan(shortId)}` + + (snapshot.model ? ` ${magenta(snapshot.model)}` : '') + + ` ${snapshot.totalSteps} steps` + + ` ${formatMs(durationMs)}`, + ); + + // Steps + const lastIdx = snapshot.steps.length - 1; + for (let i = 0; i <= lastIdx; i++) { + const step = snapshot.steps[i]; + const isLast = i === lastIdx; + const prefix = isLast ? '└─' : '├─'; + const childPrefix = isLast ? ' ' : '│ '; + + lines.push( + `${prefix} Step ${step.stepIndex} ${dim(`[${step.stepType}]`)} ${formatMs(step.executionTimeMs)}`, + ); + + if (step.stepType === 'call_llm') { + renderLlmStep(lines, step, childPrefix); + } else { + renderToolStep(lines, step, childPrefix); + } + } + + // Footer + const reasonColor = snapshot.completionReason === 'done' ? green : snapshot.error ? red : yellow; + lines.push( + `${dim('└─')} ${reasonColor(snapshot.completionReason ?? 'unknown')}` + + ` tokens=${formatTokens(snapshot.totalTokens)}` + + ` cost=${formatCost(snapshot.totalCost)}`, + ); + + if (snapshot.error) { + lines.push(` ${red('Error:')} ${snapshot.error.type} — ${snapshot.error.message}`); + } + + return lines.join('\n'); +} + +function renderLlmStep(lines: string[], step: StepSnapshot, prefix: string): void { + const tokenInfo: string[] = []; + if (step.inputTokens) tokenInfo.push(`in:${formatTokens(step.inputTokens)}`); + if (step.outputTokens) tokenInfo.push(`out:${formatTokens(step.outputTokens)}`); + + if (tokenInfo.length > 0) { + lines.push(`${prefix}${dim('├─')} LLM ${tokenInfo.join(' ')} tokens`); + } + + if (step.toolsCalling && step.toolsCalling.length > 0) { + const names = step.toolsCalling.map((t) => t.identifier || t.apiName); + lines.push( + `${prefix}${dim('├─')} ${yellow('→')} ${step.toolsCalling.length} tool_calls: [${names.join(', ')}]`, + ); + } + + if (step.content) { + lines.push(`${prefix}${dim('└─')} Output ${dim(truncate(step.content, 80))}`); + } + + if (step.reasoning) { + lines.push(`${prefix}${dim('└─')} Reason ${dim(truncate(step.reasoning, 80))}`); + } +} + +function renderMessageList(lines: string[], messages: any[], maxContentLen: number): void { + for (let i = 0; i < messages.length; i++) { + const msg = messages[i]; + const role = msg.role ?? 'unknown'; + const roleColor = + role === 'user' ? green : role === 'assistant' ? cyan : role === 'system' ? magenta : yellow; + const rawContent = typeof msg.content === 'string' ? msg.content : JSON.stringify(msg.content); + const preview = + rawContent && rawContent.length > maxContentLen + ? rawContent.slice(0, maxContentLen) + '...' + : rawContent; + lines.push( + ` ${dim(`[${i}]`)} ${roleColor(role)}${msg.tool_call_id ? dim(` [${msg.tool_call_id}]`) : ''}`, + ); + if (preview) lines.push(` ${dim(preview)}`); + if (msg.tool_calls) { + for (const tc of msg.tool_calls) { + lines.push(` ${yellow('→')} ${tc.function?.name ?? tc.id}`); + } + } + } +} + +function renderToolStep(lines: string[], step: StepSnapshot, prefix: string): void { + if (step.toolsResult) { + for (let i = 0; i < step.toolsResult.length; i++) { + const tool = step.toolsResult[i]; + const isLast = i === step.toolsResult.length - 1; + const connector = isLast ? '└─' : '├─'; + const status = tool.isSuccess === false ? red('✗') : green('✓'); + const name = tool.identifier || tool.apiName; + lines.push(`${prefix}${dim(connector)} Tool ${name} ${status}`); + } + } +} + +export function renderSummaryTable(summaries: SnapshotSummary[]): string { + if (summaries.length === 0) return dim('No snapshots found.'); + + const lines: string[] = [ + bold( + padEnd('Trace ID', 14) + + padEnd('Model', 20) + + padEnd('Steps', 7) + + padEnd('Tokens', 10) + + padEnd('Duration', 10) + + padEnd('Reason', 12) + + 'Time', + ), + dim('─'.repeat(90)), + ]; + + for (const s of summaries) { + const reasonColor = s.completionReason === 'done' ? green : s.hasError ? red : yellow; + const time = new Date(s.createdAt).toLocaleString(); + + lines.push( + cyan(padEnd(s.traceId.slice(0, 12), 14)) + + padEnd(s.model ?? '-', 20) + + padEnd(String(s.totalSteps), 7) + + padEnd(formatTokens(s.totalTokens), 10) + + padEnd(formatMs(s.durationMs), 10) + + reasonColor(padEnd(s.completionReason ?? '-', 12)) + + dim(time), + ); + } + + return lines.join('\n'); +} + +export function renderMessageDetail( + step: StepSnapshot, + msgIndex: number, + source: 'input' | 'output' = 'output', +): string { + const ceEvent = step.events?.find((e) => e.type === 'context_engine_result') as any; + let messages: any[] | undefined; + let label: string; + + if (source === 'input') { + messages = ceEvent?.input?.messages ?? step.messages; + label = ceEvent ? 'Context Engine Input' : 'Messages (before step)'; + } else { + messages = ceEvent?.output ?? step.messages; + label = ceEvent ? 'Final LLM Payload' : 'Messages (before step)'; + } + + if (!messages || messages.length === 0) { + return red('No messages available.'); + } + if (msgIndex < 0 || msgIndex >= messages.length) { + return red(`Message index ${msgIndex} out of range. Available: 0-${messages.length - 1}`); + } + + const msg = messages[msgIndex]; + const lines: string[] = []; + const role = msg.role ?? 'unknown'; + const roleColor = + role === 'user' ? green : role === 'assistant' ? cyan : role === 'system' ? magenta : yellow; + + lines.push(bold(`Message [${msgIndex}]`) + ` from ${label} (${messages.length} total)`); + lines.push( + `Role: ${roleColor(role)}${msg.tool_call_id ? ` tool_call_id: ${msg.tool_call_id}` : ''}`, + ); + if (msg.name) lines.push(`Name: ${msg.name}`); + lines.push(dim('─'.repeat(60))); + + const rawContent = + typeof msg.content === 'string' ? msg.content : JSON.stringify(msg.content, null, 2); + if (rawContent) lines.push(rawContent); + + if (msg.tool_calls && msg.tool_calls.length > 0) { + lines.push(''); + lines.push(bold('Tool Calls:')); + for (const tc of msg.tool_calls) { + lines.push(` ${yellow('→')} ${tc.function?.name ?? tc.id}`); + if (tc.function?.arguments) { + lines.push(` ${dim(tc.function.arguments)}`); + } + } + } + + return lines.join('\n'); +} + +export function renderStepDetail( + step: StepSnapshot, + options?: { context?: boolean; events?: boolean; messages?: boolean; tools?: boolean }, +): string { + const lines: string[] = [ + bold(`Step ${step.stepIndex}`) + ` [${step.stepType}] ${formatMs(step.executionTimeMs)}`, + ]; + if (step.context?.phase) { + lines.push(`Phase: ${cyan(step.context.phase)}`); + } + lines.push(''); + + if (step.inputTokens || step.outputTokens) { + lines.push(`Tokens: in=${step.inputTokens ?? 0} out=${step.outputTokens ?? 0}`); + } + + // Default view: show content & reasoning (unless specific flags are set) + const hasSpecificFlag = + options?.messages || options?.tools || options?.events || options?.context; + if (!hasSpecificFlag || options?.messages) { + if (step.content) { + lines.push(''); + lines.push(bold('Content:')); + lines.push(step.content); + } + if (step.reasoning) { + lines.push(''); + lines.push(bold('Reasoning:')); + lines.push(step.reasoning); + } + } + + if (options?.tools) { + if (step.toolsCalling && step.toolsCalling.length > 0) { + lines.push(''); + lines.push(bold('Tool Calls:')); + for (const tc of step.toolsCalling) { + lines.push(` ${cyan(tc.identifier || tc.apiName)}`); + if (tc.arguments) { + lines.push(` args: ${tc.arguments}`); + } + } + } + if (step.toolsResult && step.toolsResult.length > 0) { + lines.push(''); + lines.push(bold('Tool Results:')); + for (const tr of step.toolsResult) { + const status = tr.isSuccess === false ? red('✗') : green('✓'); + lines.push(` ${status} ${cyan(tr.identifier || tr.apiName)}`); + if (tr.output) { + const output = tr.output.length > 500 ? tr.output.slice(0, 500) + '...' : tr.output; + lines.push(` output: ${output}`); + } + } + } + } + + if (options?.messages) { + // Show context engine input/output from events if available + const ceEvent = step.events?.find((e) => e.type === 'context_engine_result') as any; + + if (ceEvent) { + // Context engine input messages (DB messages passed to engine) + const inputMsgs = ceEvent.input?.messages; + if (inputMsgs) { + lines.push(''); + lines.push(bold(`Context Engine Input: ${inputMsgs.length} messages`) + dim(' (from DB)')); + lines.push(dim('─'.repeat(60))); + renderMessageList(lines, inputMsgs, 200); + } + + // Context engine params + lines.push(''); + lines.push(bold('Context Engine Params:')); + if (ceEvent.input?.systemRole) { + const sr = ceEvent.input.systemRole; + lines.push(` systemRole: ${dim(sr.length > 100 ? sr.slice(0, 100) + '...' : sr)}`); + } + if (ceEvent.input?.model) lines.push(` model: ${cyan(ceEvent.input.model)}`); + if (ceEvent.input?.provider) lines.push(` provider: ${ceEvent.input.provider}`); + if (ceEvent.input?.enableHistoryCount != null) + lines.push(` enableHistoryCount: ${ceEvent.input.enableHistoryCount}`); + if (ceEvent.input?.historyCount != null) + lines.push(` historyCount: ${ceEvent.input.historyCount}`); + if (ceEvent.input?.forceFinish) lines.push(` forceFinish: ${ceEvent.input.forceFinish}`); + if (ceEvent.input?.knowledge) { + const k = ceEvent.input.knowledge; + const fileCount = k.fileContents?.length ?? 0; + const kbCount = k.knowledgeBases?.length ?? 0; + if (fileCount > 0 || kbCount > 0) { + lines.push(` knowledge: ${fileCount} files, ${kbCount} knowledge bases`); + } + } + if (ceEvent.input?.toolsConfig?.tools?.length) { + lines.push(` tools: ${ceEvent.input.toolsConfig.tools.length} plugins`); + } + if (ceEvent.input?.userMemory) lines.push(` userMemory: ${dim('present')}`); + + // Final messages sent to LLM + const outputMsgs = ceEvent.output; + if (outputMsgs) { + lines.push(''); + lines.push( + bold(`Final LLM Payload: ${outputMsgs.length} messages`) + dim(' (after context engine)'), + ); + lines.push(dim('─'.repeat(60))); + renderMessageList(lines, outputMsgs, 300); + } + } else if (step.messages) { + // Fallback: show raw DB messages if no context engine event + lines.push(''); + lines.push(bold(`Messages (before step): ${step.messages.length} messages`)); + lines.push(dim('─'.repeat(60))); + renderMessageList(lines, step.messages, 200); + } + } + + if (options?.events && step.events) { + lines.push(''); + lines.push(bold(`Events: ${step.events.length}`)); + lines.push(dim('─'.repeat(60))); + for (const event of step.events) { + const typeColor = + event.type === 'llm_result' + ? cyan + : event.type === 'llm_start' + ? magenta + : event.type === 'done' + ? green + : event.type === 'error' + ? red + : yellow; + lines.push(` ${typeColor(event.type)}`); + + if (event.type === 'llm_result') { + const result = event.result as any; + if (result?.content) { + const preview = + result.content.length > 300 ? result.content.slice(0, 300) + '...' : result.content; + lines.push(` content: ${dim(preview)}`); + } + if (result?.tool_calls) { + for (const tc of result.tool_calls) { + lines.push(` ${yellow('→')} ${tc.function?.name ?? tc.id}`); + } + } + } else if (event.type === 'llm_start') { + const payload = event.payload as any; + if (payload?.messages) { + lines.push(` ${dim(`${payload.messages.length} messages`)}`); + } + if (payload?.model) { + lines.push(` model: ${payload.model}`); + } + } else if (event.type === 'done') { + if (event.reason) lines.push(` reason: ${event.reason}`); + } else if (event.type === 'tool_result') { + const result = + typeof event.result === 'string' + ? event.result.length > 200 + ? event.result.slice(0, 200) + '...' + : event.result + : JSON.stringify(event.result)?.slice(0, 200); + lines.push(` id: ${event.id ?? '-'}`); + if (result) lines.push(` result: ${dim(result)}`); + } + } + } + + if (options?.context && step.context) { + lines.push(''); + lines.push(bold('Runtime Context:')); + lines.push(dim('─'.repeat(60))); + lines.push(` phase: ${cyan(step.context.phase)}`); + if (step.context.payload) { + lines.push(''); + lines.push(bold(' Payload:')); + const payloadStr = JSON.stringify(step.context.payload, null, 2); + const payloadLines = payloadStr.split('\n'); + for (const line of payloadLines.slice(0, 50)) { + lines.push(` ${dim(line)}`); + } + if (payloadLines.length > 50) { + lines.push(` ${dim(`... (${payloadLines.length - 50} more lines)`)}`); + } + } + if (step.context.stepContext) { + lines.push(''); + lines.push(bold(' Step Context:')); + lines.push( + ` ${dim(JSON.stringify(step.context.stepContext, null, 2).split('\n').join('\n '))}`, + ); + } + } + + return lines.join('\n'); +} diff --git a/packages/agent-tracing/tsconfig.json b/packages/agent-tracing/tsconfig.json new file mode 100644 index 0000000000..6ddc33fbfb --- /dev/null +++ b/packages/agent-tracing/tsconfig.json @@ -0,0 +1,7 @@ +{ + "extends": "../../tsconfig.json", + "compilerOptions": { + "outDir": "./dist" + }, + "include": ["src/**/*.ts"] +} diff --git a/packages/context-engine/src/engine/messages/types.ts b/packages/context-engine/src/engine/messages/types.ts index 30a9c615ac..31446ddb98 100644 --- a/packages/context-engine/src/engine/messages/types.ts +++ b/packages/context-engine/src/engine/messages/types.ts @@ -116,6 +116,7 @@ export interface UserMemoryActivityItem { } export interface UserMemoryIdentityItem { + capturedAt?: string | Date | null; description?: string | null; id?: string; role?: string | null; @@ -124,6 +125,12 @@ export interface UserMemoryIdentityItem { [key: string]: unknown; } +export interface UserMemoryPersonaItem { + narrative?: string | null; + tagline?: string | null; + [key: string]: unknown; +} + /** * User memory data structure * Compatible with SearchMemoryResult from @lobechat/types @@ -133,6 +140,7 @@ export interface UserMemoryData { contexts: UserMemoryContextItem[]; experiences: UserMemoryExperienceItem[]; identities?: UserMemoryIdentityItem[]; + persona?: UserMemoryPersonaItem; preferences: UserMemoryPreferenceItem[]; } diff --git a/packages/context-engine/src/providers/UserMemoryInjector.ts b/packages/context-engine/src/providers/UserMemoryInjector.ts index b8aebe6ec8..f070f40b75 100644 --- a/packages/context-engine/src/providers/UserMemoryInjector.ts +++ b/packages/context-engine/src/providers/UserMemoryInjector.ts @@ -42,13 +42,14 @@ export class UserMemoryInjector extends BaseFirstUserContentProvider { return null; } + const hasPersona = !!(memories.persona?.narrative || memories.persona?.tagline); const identitiesCount = memories.identities?.length || 0; const contextsCount = memories.contexts?.length || 0; const experiencesCount = memories.experiences?.length || 0; const preferencesCount = memories.preferences?.length || 0; log( - `User memories prepared: ${identitiesCount} identity(ies), ${contextsCount} context(s), ${experiencesCount} experience(s), ${preferencesCount} preference(s)`, + `User memories prepared: persona=${hasPersona}, ${identitiesCount} identity(ies), ${contextsCount} context(s), ${experiencesCount} experience(s), ${preferencesCount} preference(s)`, ); return content; diff --git a/packages/context-engine/src/providers/__tests__/UserMemoryInjector.test.ts b/packages/context-engine/src/providers/__tests__/UserMemoryInjector.test.ts index dc2d4f1fc7..81c28fdaf3 100644 --- a/packages/context-engine/src/providers/__tests__/UserMemoryInjector.test.ts +++ b/packages/context-engine/src/providers/__tests__/UserMemoryInjector.test.ts @@ -165,6 +165,73 @@ describe('UserMemoryInjector', () => { expect(result.messages[0].content).toMatchSnapshot(); }); + + it('should inject identities with capturedAt', async () => { + const provider = new UserMemoryInjector({ + memories: { + contexts: [], + experiences: [], + identities: [ + { + capturedAt: '2025-02-23T10:30:00.000Z', + description: 'User is a senior engineer', + id: 'id-1', + role: 'Engineer', + type: 'professional', + }, + ], + preferences: [], + }, + }); + + const context = createContext([{ content: 'Hello', id: 'user-1', role: 'user' }]); + + const result = await provider.process(context); + + expect(result.messages[0].content).toMatchSnapshot(); + }); + }); + + describe('persona injection', () => { + it('should inject persona via UserMemoryInjector', async () => { + const provider = new UserMemoryInjector({ + memories: { + contexts: [], + experiences: [], + persona: { + narrative: 'A senior engineer who loves TypeScript and open-source development.', + tagline: 'Senior OSS engineer', + }, + preferences: [], + }, + }); + + const context = createContext([{ content: 'Hello', id: 'user-1', role: 'user' }]); + + const result = await provider.process(context); + + expect(result.messages[0].content).toMatchSnapshot(); + }); + + it('should inject persona combined with other memory types', async () => { + const provider = new UserMemoryInjector({ + memories: { + contexts: [{ description: 'Context desc', id: 'ctx-1', title: 'Context' }], + experiences: [{ id: 'exp-1', keyLearning: 'Learning', situation: 'Situation' }], + persona: { + narrative: 'A product designer turned engineer.', + tagline: 'Design-minded engineer', + }, + preferences: [{ conclusionDirectives: 'Preference', id: 'pref-1' }], + }, + }); + + const context = createContext([{ content: 'Hello', id: 'user-1', role: 'user' }]); + + const result = await provider.process(context); + + expect(result.messages[0].content).toMatchSnapshot(); + }); }); describe('XML format', () => { diff --git a/packages/context-engine/src/providers/__tests__/__snapshots__/UserMemoryInjector.test.ts.snap b/packages/context-engine/src/providers/__tests__/__snapshots__/UserMemoryInjector.test.ts.snap index b14f142f1a..e56ad6b6ec 100644 --- a/packages/context-engine/src/providers/__tests__/__snapshots__/UserMemoryInjector.test.ts.snap +++ b/packages/context-engine/src/providers/__tests__/__snapshots__/UserMemoryInjector.test.ts.snap @@ -75,6 +75,15 @@ exports[`UserMemoryInjector > memory types > should inject experiences correctly " `; +exports[`UserMemoryInjector > memory types > should inject identities with capturedAt 1`] = ` +" +The following are memories about this user retrieved from previous conversations. Use this information to personalize your responses and maintain continuity. + + User is a senior engineer + +" +`; + exports[`UserMemoryInjector > memory types > should inject preferences correctly 1`] = ` " The following are memories about this user retrieved from previous conversations. Use this information to personalize your responses and maintain continuity. @@ -84,3 +93,33 @@ exports[`UserMemoryInjector > memory types > should inject preferences correctly " `; + +exports[`UserMemoryInjector > persona injection > should inject persona combined with other memory types 1`] = ` +" +The following are memories about this user retrieved from previous conversations. Use this information to personalize your responses and maintain continuity. + +A product designer turned engineer. + + + Context desc + + + + Situation + Learning + + + + Preference + +" +`; + +exports[`UserMemoryInjector > persona injection > should inject persona via UserMemoryInjector 1`] = ` +" +The following are memories about this user retrieved from previous conversations. Use this information to personalize your responses and maintain continuity. + +A senior engineer who loves TypeScript and open-source development. + +" +`; diff --git a/packages/database/src/models/__tests__/agent.test.ts b/packages/database/src/models/__tests__/agent.test.ts index 2443670cf5..de5780f023 100644 --- a/packages/database/src/models/__tests__/agent.test.ts +++ b/packages/database/src/models/__tests__/agent.test.ts @@ -1694,4 +1694,238 @@ describe('AgentModel', () => { expect(offsetResults.length).toBe(2); }); }); + + describe('checkByMarketIdentifier', () => { + it('should return true when agent with marketIdentifier exists', async () => { + await serverDB.insert(agents).values({ + userId, + title: 'Market Agent', + marketIdentifier: 'market-test-123', + }); + + const result = await agentModel.checkByMarketIdentifier('market-test-123'); + expect(result).toBe(true); + }); + + it('should return false when no agent with marketIdentifier exists', async () => { + const result = await agentModel.checkByMarketIdentifier('non-existent-market-id'); + expect(result).toBe(false); + }); + + it('should not find agents belonging to other users', async () => { + await serverDB.insert(agents).values({ + userId: userId2, + title: 'Other User Agent', + marketIdentifier: 'other-user-market', + }); + + const result = await agentModel.checkByMarketIdentifier('other-user-market'); + expect(result).toBe(false); + }); + }); + + describe('getAgentByMarketIdentifier', () => { + it('should return agent id when found', async () => { + const [agent] = await serverDB + .insert(agents) + .values({ userId, title: 'Market Agent', marketIdentifier: 'market-get-123' }) + .returning(); + + const result = await agentModel.getAgentByMarketIdentifier('market-get-123'); + expect(result).toBe(agent.id); + }); + + it('should return null when not found', async () => { + const result = await agentModel.getAgentByMarketIdentifier('nonexistent'); + expect(result).toBeNull(); + }); + + it('should return the most recently updated agent when multiple match', async () => { + const [older] = await serverDB + .insert(agents) + .values({ userId, title: 'Older', marketIdentifier: 'dup-market' }) + .returning(); + + await new Promise((resolve) => setTimeout(resolve, 10)); + + const [newer] = await serverDB + .insert(agents) + .values({ userId, title: 'Newer', marketIdentifier: 'dup-market' }) + .returning(); + + const result = await agentModel.getAgentByMarketIdentifier('dup-market'); + expect(result).toBe(newer.id); + }); + + it('should not return agents from other users', async () => { + await serverDB.insert(agents).values({ + userId: userId2, + title: 'Other', + marketIdentifier: 'other-market-get', + }); + + const result = await agentModel.getAgentByMarketIdentifier('other-market-get'); + expect(result).toBeNull(); + }); + }); + + describe('getAgentByForkedFromIdentifier', () => { + it('should return agent id when forkedFromIdentifier matches', async () => { + const [agent] = await serverDB + .insert(agents) + .values({ + userId, + title: 'Forked Agent', + params: { forkedFromIdentifier: 'source-market-id' }, + }) + .returning(); + + const result = await agentModel.getAgentByForkedFromIdentifier('source-market-id'); + expect(result).toBe(agent.id); + }); + + it('should return null when no match', async () => { + const result = await agentModel.getAgentByForkedFromIdentifier('no-match'); + expect(result).toBeNull(); + }); + + it('should not return agents from other users', async () => { + await serverDB.insert(agents).values({ + userId: userId2, + title: 'Other Forked', + params: { forkedFromIdentifier: 'other-fork-id' }, + }); + + const result = await agentModel.getAgentByForkedFromIdentifier('other-fork-id'); + expect(result).toBeNull(); + }); + }); + + describe('updateSessionGroupId', () => { + it('should update agent sessionGroupId', async () => { + const [group] = await serverDB + .insert(sessionGroups) + .values({ userId, name: 'Test Group' }) + .returning(); + + const [agent] = await serverDB.insert(agents).values({ userId, title: 'Agent' }).returning(); + + const result = await agentModel.updateSessionGroupId(agent.id, group.id); + + expect(result).toBeDefined(); + expect(result.sessionGroupId).toBe(group.id); + }); + + it('should set sessionGroupId to null', async () => { + const [group] = await serverDB + .insert(sessionGroups) + .values({ userId, name: 'Group' }) + .returning(); + + const [agent] = await serverDB + .insert(agents) + .values({ userId, title: 'Agent', sessionGroupId: group.id }) + .returning(); + + const result = await agentModel.updateSessionGroupId(agent.id, null); + expect(result.sessionGroupId).toBeNull(); + }); + + it('should not update agents from other users', async () => { + const [agent] = await serverDB + .insert(agents) + .values({ userId, title: 'User1 Agent' }) + .returning(); + + const result = await agentModel2.updateSessionGroupId(agent.id, null); + expect(result).toBeUndefined(); + }); + }); + + describe('updateConfig edge cases', () => { + it('should return early for null data', async () => { + const [agent] = await serverDB + .insert(agents) + .values({ userId, title: 'Original' }) + .returning(); + + const result = await agentModel.updateConfig(agent.id, null); + expect(result).toBeUndefined(); + + const dbAgent = await serverDB.query.agents.findFirst({ + where: eq(agents.id, agent.id), + }); + expect(dbAgent?.title).toBe('Original'); + }); + + it('should return early for undefined data', async () => { + const [agent] = await serverDB + .insert(agents) + .values({ userId, title: 'Original' }) + .returning(); + + const result = await agentModel.updateConfig(agent.id, undefined); + expect(result).toBeUndefined(); + }); + + it('should return early for empty object', async () => { + const [agent] = await serverDB + .insert(agents) + .values({ userId, title: 'Original' }) + .returning(); + + const result = await agentModel.updateConfig(agent.id, {}); + expect(result).toBeUndefined(); + }); + + it('should return early for non-existent agent', async () => { + const result = await agentModel.updateConfig('non-existent-id', { title: 'New' }); + expect(result).toBeUndefined(); + }); + + it('should delete params field when value is undefined', async () => { + const [agent] = await serverDB + .insert(agents) + .values({ + userId, + title: 'Params Agent', + params: { temperature: 0.7, topP: 0.9 }, + }) + .returning(); + + await agentModel.updateConfig(agent.id, { + params: { temperature: undefined } as any, + }); + + const result = await serverDB.query.agents.findFirst({ + where: eq(agents.id, agent.id), + }); + + // temperature should be deleted, topP should remain + expect((result?.params as any)?.temperature).toBeUndefined(); + expect((result?.params as any)?.topP).toBe(0.9); + }); + + it('should handle null param values (disable flag)', async () => { + const [agent] = await serverDB + .insert(agents) + .values({ + userId, + title: 'Params Agent', + params: { temperature: 0.7 }, + }) + .returning(); + + await agentModel.updateConfig(agent.id, { + params: { temperature: null, topP: 0.5 } as any, + }); + + const result = await serverDB.query.agents.findFirst({ + where: eq(agents.id, agent.id), + }); + + expect((result?.params as any)?.temperature).toBeNull(); + expect((result?.params as any)?.topP).toBe(0.5); + }); + }); }); diff --git a/packages/database/src/models/__tests__/agentSkill.test.ts b/packages/database/src/models/__tests__/agentSkill.test.ts index d7757763f9..1bd47fdecf 100644 --- a/packages/database/src/models/__tests__/agentSkill.test.ts +++ b/packages/database/src/models/__tests__/agentSkill.test.ts @@ -299,4 +299,42 @@ describe('AgentSkillModel', () => { expect(results.total).toBe(1); }); }); + + describe('findByName', () => { + it('should find a skill by name', async () => { + await serverDB.insert(agentSkills).values({ + name: 'Unique Skill Name', + description: 'A unique skill', + identifier: 'unique-skill', + source: 'user', + manifest: createManifest(), + userId, + }); + + const result = await agentSkillModel.findByName('Unique Skill Name'); + expect(result).toBeDefined(); + expect(result?.name).toBe('Unique Skill Name'); + }); + + it('should return undefined for non-existent name', async () => { + const result = await agentSkillModel.findByName('Non Existent'); + expect(result).toBeUndefined(); + }); + + it('should not find skills from other users', async () => { + const otherUserId = 'other-skill-user'; + await serverDB.insert(users).values({ id: otherUserId }); + await serverDB.insert(agentSkills).values({ + name: 'Other Skill', + description: 'Other skill desc', + identifier: 'other-skill', + source: 'user', + manifest: createManifest(), + userId: otherUserId, + }); + + const result = await agentSkillModel.findByName('Other Skill'); + expect(result).toBeUndefined(); + }); + }); }); diff --git a/packages/database/src/models/__tests__/asyncTask.test.ts b/packages/database/src/models/__tests__/asyncTask.test.ts index 3196fc4952..031a82ae7b 100644 --- a/packages/database/src/models/__tests__/asyncTask.test.ts +++ b/packages/database/src/models/__tests__/asyncTask.test.ts @@ -381,3 +381,29 @@ describe('initUserMemoryExtractionMetadata', () => { expect(result.source).toBe('chat_topic'); }); }); + +describe('AsyncTaskModel.findByInferenceId', () => { + it('should find a task by inferenceId', async () => { + const [task] = await serverDB + .insert(asyncTasks) + .values({ + status: AsyncTaskStatus.Processing, + userId, + inferenceId: 'inference-123', + type: AsyncTaskType.UserMemoryExtractionWithChatTopic, + }) + .returning(); + + const result = await AsyncTaskModel.findByInferenceId(serverDB, 'inference-123'); + expect(result).toBeDefined(); + expect(result?.id).toBe(task.id); + }); + + it('should return undefined for non-existent inferenceId', async () => { + const result = await AsyncTaskModel.findByInferenceId( + serverDB, + '00000000-0000-0000-0000-000000000000', + ); + expect(result).toBeUndefined(); + }); +}); diff --git a/packages/database/src/models/__tests__/document.test.ts b/packages/database/src/models/__tests__/document.test.ts index ea53947337..6a4a1f4e65 100644 --- a/packages/database/src/models/__tests__/document.test.ts +++ b/packages/database/src/models/__tests__/document.test.ts @@ -223,6 +223,55 @@ describe('DocumentModel', () => { expect(allResult.total).toBe(3); }); + it('should filter documents by fileTypes', async () => { + const { id: fileId1 } = await fileModel.create({ + fileType: 'text/plain', + name: 'test1.txt', + size: 100, + url: 'https://example.com/test1.txt', + }); + const file1 = await fileModel.findById(fileId1); + if (!file1) throw new Error('File not found'); + + const { id: fileId2 } = await fileModel.create({ + fileType: 'application/pdf', + name: 'test2.pdf', + size: 200, + url: 'https://example.com/test2.pdf', + }); + const file2 = await fileModel.findById(fileId2); + if (!file2) throw new Error('File not found'); + + await documentModel.create({ + content: 'Text document', + fileId: file1.id, + fileType: 'text/plain', + source: file1.url, + sourceType: 'file', + totalCharCount: 13, + totalLineCount: 1, + }); + + await documentModel.create({ + content: 'PDF document', + fileId: file2.id, + fileType: 'application/pdf', + source: file2.url, + sourceType: 'file', + totalCharCount: 12, + totalLineCount: 1, + }); + + // Filter by fileTypes + const textResult = await documentModel.query({ fileTypes: ['text/plain'] }); + expect(textResult.items).toHaveLength(1); + expect(textResult.total).toBe(1); + + // Without filter returns all + const allResult = await documentModel.query(); + expect(allResult.items).toHaveLength(2); + }); + it('should return documents ordered by updatedAt desc', async () => { const { documentId: doc1Id } = await createTestDocument( documentModel, diff --git a/packages/database/src/models/__tests__/file.test.ts b/packages/database/src/models/__tests__/file.test.ts index 9685d7b62b..ef3ac56910 100644 --- a/packages/database/src/models/__tests__/file.test.ts +++ b/packages/database/src/models/__tests__/file.test.ts @@ -1335,4 +1335,72 @@ describe('FileModel', () => { expect(remainingFileChunks).toHaveLength(0); }); }); + + describe('static getFileById', () => { + it('should return a file by id', async () => { + const [file] = await serverDB + .insert(files) + .values({ + id: 'static-file-id', + userId, + name: 'static-file.txt', + url: 'https://example.com/file.txt', + fileType: 'text/plain', + size: 100, + }) + .returning(); + + const result = await FileModel.getFileById(serverDB, file.id); + expect(result).toBeDefined(); + expect(result?.id).toBe(file.id); + expect(result?.name).toBe('static-file.txt'); + }); + + it('should return undefined for non-existent file', async () => { + const result = await FileModel.getFileById(serverDB, 'non-existent'); + expect(result).toBeUndefined(); + }); + }); + + describe('findByIds', () => { + it('should find multiple files by ids', async () => { + await serverDB.insert(files).values([ + { + id: 'find-id-1', + userId, + name: 'file1.txt', + url: 'url1', + fileType: 'text/plain', + size: 100, + }, + { + id: 'find-id-2', + userId, + name: 'file2.txt', + url: 'url2', + fileType: 'text/plain', + size: 200, + }, + ]); + + const result = await fileModel.findByIds(['find-id-1', 'find-id-2']); + expect(result).toHaveLength(2); + }); + + it('should only return files belonging to current user', async () => { + const otherUserId = 'other-file-user'; + await serverDB.insert(users).values({ id: otherUserId }); + await serverDB.insert(files).values({ + id: 'other-file-id', + userId: otherUserId, + name: 'other.txt', + url: 'url', + fileType: 'text/plain', + size: 100, + }); + + const result = await fileModel.findByIds(['other-file-id']); + expect(result).toHaveLength(0); + }); + }); }); diff --git a/packages/database/src/models/__tests__/generation.test.ts b/packages/database/src/models/__tests__/generation.test.ts index ddd3ee36ed..697f2c4556 100644 --- a/packages/database/src/models/__tests__/generation.test.ts +++ b/packages/database/src/models/__tests__/generation.test.ts @@ -919,4 +919,30 @@ describe('GenerationModel', () => { ); }); }); + + describe('findByAsyncTaskId', () => { + it('should find generation by asyncTaskId', async () => { + const [task] = await serverDB + .insert(asyncTasks) + .values({ status: 'processing', userId }) + .returning(); + + await serverDB.insert(generations).values({ + ...testGeneration, + userId, + asyncTaskId: task.id, + }); + + const result = await generationModel.findByAsyncTaskId(task.id); + expect(result).toBeDefined(); + expect(result?.asyncTaskId).toBe(task.id); + }); + + it('should return undefined for non-existent asyncTaskId', async () => { + const result = await generationModel.findByAsyncTaskId( + '00000000-0000-0000-0000-000000000000', + ); + expect(result).toBeUndefined(); + }); + }); }); diff --git a/packages/database/src/models/__tests__/messages/message.stats.test.ts b/packages/database/src/models/__tests__/messages/message.stats.test.ts index a048e7aae6..6a078f3ac1 100644 --- a/packages/database/src/models/__tests__/messages/message.stats.test.ts +++ b/packages/database/src/models/__tests__/messages/message.stats.test.ts @@ -630,4 +630,35 @@ describe('MessageModel Statistics Tests', () => { expect(result3).toBe(true); }); }); + + describe('countUpTo', () => { + it('should count messages up to a limit', async () => { + await serverDB.insert(messages).values([ + { id: 'count-1', userId, role: 'user', content: 'msg 1' }, + { id: 'count-2', userId, role: 'user', content: 'msg 2' }, + { id: 'count-3', userId, role: 'user', content: 'msg 3' }, + ]); + + const result = await messageModel.countUpTo(5); + expect(result).toBe(3); + }); + + it('should return at most n', async () => { + await serverDB.insert(messages).values([ + { id: 'count-a', userId, role: 'user', content: 'msg a' }, + { id: 'count-b', userId, role: 'user', content: 'msg b' }, + { id: 'count-c', userId, role: 'user', content: 'msg c' }, + ]); + + const result = await messageModel.countUpTo(2); + expect(result).toBe(2); + }); + + it('should return 0 for empty user', async () => { + const otherModel = new MessageModel(serverDB, 'empty-count-user'); + await serverDB.insert(users).values({ id: 'empty-count-user' }); + const result = await otherModel.countUpTo(10); + expect(result).toBe(0); + }); + }); }); diff --git a/packages/database/src/models/__tests__/messages/message.update.test.ts b/packages/database/src/models/__tests__/messages/message.update.test.ts index 7f7a0d3ec9..bc02127dc7 100644 --- a/packages/database/src/models/__tests__/messages/message.update.test.ts +++ b/packages/database/src/models/__tests__/messages/message.update.test.ts @@ -1301,4 +1301,105 @@ describe('MessageModel Update Tests', () => { expect(result[0].voice).toBe('updated voice1'); }); }); + + describe('addFiles', () => { + it('should add file associations to a message', async () => { + await serverDB.insert(messages).values({ + id: 'msg-add-files', + userId, + role: 'user', + content: 'test message', + }); + + const result = await messageModel.addFiles('msg-add-files', ['f1']); + expect(result.success).toBe(true); + + const messageFiles = await serverDB + .select() + .from(messagesFiles) + .where(eq(messagesFiles.messageId, 'msg-add-files')); + expect(messageFiles).toHaveLength(1); + expect(messageFiles[0].fileId).toBe('f1'); + }); + + it('should return success true for empty fileIds array', async () => { + const result = await messageModel.addFiles('msg-any', []); + expect(result.success).toBe(true); + }); + + it('should return success false on database error', async () => { + // Try to add a file with a non-existent fileId (FK constraint violation) + await serverDB.insert(messages).values({ + id: 'msg-add-files-err', + userId, + role: 'user', + content: 'test', + }); + + const result = await messageModel.addFiles('msg-add-files-err', ['non-existent-file-id']); + expect(result.success).toBe(false); + }); + + it('should add multiple files at once', async () => { + await serverDB + .insert(files) + .values([ + { id: 'f2', userId, url: 'url2', name: 'file-2', fileType: 'image/jpeg', size: 500 }, + ]); + await serverDB.insert(messages).values({ + id: 'msg-multi-files', + userId, + role: 'user', + content: 'test', + }); + + const result = await messageModel.addFiles('msg-multi-files', ['f1', 'f2']); + expect(result.success).toBe(true); + + const messageFiles = await serverDB + .select() + .from(messagesFiles) + .where(eq(messagesFiles.messageId, 'msg-multi-files')); + expect(messageFiles).toHaveLength(2); + }); + }); + + describe('updateToolArguments - parent message without tools', () => { + it('should return success false when parent message has no tools', async () => { + // Create assistant message WITHOUT tools + await serverDB.insert(messages).values({ + id: 'assistant-no-tools', + userId, + role: 'assistant', + content: 'No tools here', + tools: null, + }); + + // Create tool message pointing to the assistant + await serverDB.insert(messages).values({ + id: 'tool-msg-orphan', + userId, + role: 'tool', + content: 'tool result', + parentId: 'assistant-no-tools', + tool_call_id: 'orphan-tool-call', + }); + + // Create plugin record + await serverDB.insert(messagePlugins).values({ + id: 'tool-msg-orphan', + toolCallId: 'orphan-tool-call', + identifier: 'test-plugin', + arguments: '{"key":"val"}', + userId, + }); + + // Should fail because parent message has no tools + const result = await messageModel.updateToolArguments( + 'orphan-tool-call', + '{"key":"updated"}', + ); + expect(result.success).toBe(false); + }); + }); }); diff --git a/packages/database/src/models/__tests__/session.test.ts b/packages/database/src/models/__tests__/session.test.ts index 29933531c7..81b8a9b97a 100644 --- a/packages/database/src/models/__tests__/session.test.ts +++ b/packages/database/src/models/__tests__/session.test.ts @@ -262,6 +262,25 @@ describe('SessionModel', () => { // 断言结果 expect(result).toBe(0); }); + + it('should count sessions with date range filter', async () => { + await serverDB.insert(sessions).values([ + { id: 's1', userId, createdAt: new Date('2024-01-01') }, + { id: 's2', userId, createdAt: new Date('2024-06-01') }, + { id: 's3', userId, createdAt: new Date('2024-12-01') }, + ]); + + const rangeResult = await sessionModel.count({ + range: ['2024-03-01', '2024-09-01'], + }); + expect(rangeResult).toBe(1); + + const startResult = await sessionModel.count({ startDate: '2024-05-01' }); + expect(startResult).toBe(2); + + const endResult = await sessionModel.count({ endDate: '2024-07-01' }); + expect(endResult).toBe(2); + }); }); describe('queryByKeyword', () => { @@ -1386,6 +1405,37 @@ describe('SessionModel', () => { expect(result[1].id).toBe('2'); // Second most topics (1) }); + it('should include inbox topics in ranking when topics have no sessionId', async () => { + await serverDB.transaction(async (trx) => { + await trx.insert(sessions).values([{ id: '1', userId }]); + await trx + .insert(agents) + .values([{ id: 'a1', userId, title: 'Agent 1', avatar: 'av1', backgroundColor: 'bg1' }]); + await trx.insert(agentsToSessions).values([{ sessionId: '1', agentId: 'a1', userId }]); + + // Create topics: 1 for session, 3 for inbox (no sessionId) + await trx.insert(topics).values([ + { id: 'inbox-t1', userId, sessionId: null }, + { id: 'inbox-t2', userId, sessionId: null }, + { id: 'inbox-t3', userId, sessionId: null }, + { id: 'session-t1', sessionId: '1', userId }, + ]); + }); + + const result = await sessionModel.rank(); + + // Should include both inbox and session entries + expect(result.length).toBeGreaterThanOrEqual(2); + // Inbox should have 3 topics and be ranked first + const inboxEntry = result.find((r) => r.id === 'inbox'); + expect(inboxEntry).toBeDefined(); + expect(inboxEntry?.count).toBe(3); + // Session should have 1 topic + const sessionEntry = result.find((r) => r.id === '1'); + expect(sessionEntry).toBeDefined(); + expect(sessionEntry?.count).toBe(1); + }); + it('should handle sessions with no topics', async () => { // Create test data await serverDB.transaction(async (trx) => { diff --git a/packages/database/src/models/__tests__/topicShare.test.ts b/packages/database/src/models/__tests__/topicShare.test.ts index febd9419fa..c54ec2b2ac 100644 --- a/packages/database/src/models/__tests__/topicShare.test.ts +++ b/packages/database/src/models/__tests__/topicShare.test.ts @@ -3,7 +3,15 @@ import { TRPCError } from '@trpc/server'; import { afterEach, beforeEach, describe, expect, it } from 'vitest'; import { getTestDB } from '../../core/getTestDB'; -import { agents, sessions, topics, topicShares, users } from '../../schemas'; +import { + agents, + chatGroups, + chatGroupsAgents, + sessions, + topics, + topicShares, + users, +} from '../../schemas'; import type { LobeChatDatabase } from '../../type'; import { TopicShareModel } from '../topicShare'; @@ -75,6 +83,15 @@ describe('TopicShareModel', () => { 'Topic not found or not owned by user', ); }); + + it('should return existing share on conflict (duplicate topic)', async () => { + const first = await topicShareModel.create(topicId); + const second = await topicShareModel.create(topicId); + + expect(second).toBeDefined(); + expect(second.topicId).toBe(topicId); + expect(second.id).toBe(first.id); + }); }); describe('updateVisibility', () => { @@ -315,4 +332,41 @@ describe('TopicShareModel', () => { expect(stillExists).not.toBeNull(); }); }); + + describe('findByShareId with group topic', () => { + it('should return group members for a group topic share', async () => { + // Create a chat group with agents + const [group] = await serverDB + .insert(chatGroups) + .values({ userId, title: 'Test Group' }) + .returning(); + + const agent2Id = 'group-member-agent'; + await serverDB.insert(agents).values({ id: agent2Id, userId, title: 'Group Agent' }); + await serverDB + .insert(chatGroupsAgents) + .values({ chatGroupId: group.id, agentId: agent2Id, userId, order: 0 }); + + // Create a topic with groupId + const groupTopicId = 'group-topic-id'; + await serverDB.insert(topics).values({ + id: groupTopicId, + sessionId, + userId, + title: 'Group Topic', + groupId: group.id, + }); + + // Create a share + const share = await topicShareModel.create(groupTopicId); + + // Find by share ID + const result = await TopicShareModel.findByShareId(serverDB, share.id); + expect(result).toBeDefined(); + expect(result?.groupId).toBe(group.id); + expect(result?.groupMembers).toBeDefined(); + expect(result?.groupMembers).toHaveLength(1); + expect(result?.groupMembers?.[0].id).toBe(agent2Id); + }); + }); }); diff --git a/packages/database/src/models/__tests__/topics/topic.create.test.ts b/packages/database/src/models/__tests__/topics/topic.create.test.ts index 4d78a4508f..4a60f47e5b 100644 --- a/packages/database/src/models/__tests__/topics/topic.create.test.ts +++ b/packages/database/src/models/__tests__/topics/topic.create.test.ts @@ -410,5 +410,22 @@ describe('TopicModel - Create', () => { `Topic with id ${topicId} not found`, ); }); + + it('should duplicate a topic with no messages (empty messageIds)', async () => { + const topicId = 'topic-no-messages'; + + await serverDB + .insert(topics) + .values({ id: topicId, sessionId, userId, title: 'Empty Topic' }); + + const { topic: duplicated, messages: duplicatedMessages } = await topicModel.duplicate( + topicId, + 'Duplicated Empty', + ); + + expect(duplicated.id).not.toBe(topicId); + expect(duplicated.title).toBe('Duplicated Empty'); + expect(duplicatedMessages).toHaveLength(0); + }); }); }); diff --git a/packages/database/src/models/__tests__/topics/topic.query.test.ts b/packages/database/src/models/__tests__/topics/topic.query.test.ts index 05a24ca718..04c2626fcd 100644 --- a/packages/database/src/models/__tests__/topics/topic.query.test.ts +++ b/packages/database/src/models/__tests__/topics/topic.query.test.ts @@ -187,6 +187,26 @@ describe('TopicModel - Query', () => { expect(result2).toHaveLength(1); expect(result2[0].id).toBe('topic1'); }); + + it('should exclude topics with specified triggers via excludeTriggers', async () => { + await serverDB.insert(topics).values([ + { id: 'normal-topic', sessionId, userId, title: 'Normal' }, + { id: 'cron-topic', sessionId, userId, title: 'Cron', trigger: 'cron' }, + { id: 'null-trigger', sessionId, userId, title: 'Null Trigger' }, + ]); + + const result = await topicModel.query({ + containerId: sessionId, + excludeTriggers: ['cron'], + }); + + // Should return topics with null trigger or triggers not in the exclude list + expect(result.items).toHaveLength(2); + const ids = result.items.map((t) => t.id); + expect(ids).toContain('normal-topic'); + expect(ids).toContain('null-trigger'); + expect(ids).not.toContain('cron-topic'); + }); }); describe('query with agentId filter', () => { @@ -1361,4 +1381,85 @@ describe('TopicModel - Query', () => { expect(rows.map((t) => t.id)).toEqual(['cursor-topic-z', 'after-1', 'after-2']); }); }); + + describe('getCronTopicsGroupedByCronJob', () => { + it('should return cron topics grouped by cronJobId', async () => { + const agentId = 'cron-agent'; + await serverDB.insert(agents).values({ id: agentId, userId }); + const [session] = await serverDB + .insert(sessions) + .values({ userId, type: 'agent' }) + .returning(); + await serverDB.insert(agentsToSessions).values({ agentId, sessionId: session.id, userId }); + + await serverDB.insert(topics).values([ + { + id: 'cron-topic-1', + userId, + sessionId: session.id, + agentId, + trigger: 'cron', + title: 'Cron Topic 1', + metadata: { cronJobId: 'job-a' }, + }, + { + id: 'cron-topic-2', + userId, + sessionId: session.id, + agentId, + trigger: 'cron', + title: 'Cron Topic 2', + metadata: { cronJobId: 'job-a' }, + }, + { + id: 'cron-topic-3', + userId, + sessionId: session.id, + agentId, + trigger: 'cron', + title: 'Cron Topic 3', + metadata: { cronJobId: 'job-b' }, + }, + ]); + + const result = await topicModel.getCronTopicsGroupedByCronJob(agentId); + + expect(result).toHaveLength(2); + const jobA = result.find((g) => g.cronJobId === 'job-a'); + const jobB = result.find((g) => g.cronJobId === 'job-b'); + expect(jobA?.topics).toHaveLength(2); + expect(jobB?.topics).toHaveLength(1); + }); + + it('should return empty array when no cron topics exist', async () => { + const agentId = 'no-cron-agent'; + await serverDB.insert(agents).values({ id: agentId, userId }); + + const result = await topicModel.getCronTopicsGroupedByCronJob(agentId); + expect(result).toEqual([]); + }); + + it('should not return topics without cronJobId in metadata', async () => { + const agentId = 'cron-agent-no-meta'; + await serverDB.insert(agents).values({ id: agentId, userId }); + const [session] = await serverDB + .insert(sessions) + .values({ userId, type: 'agent' }) + .returning(); + await serverDB.insert(agentsToSessions).values({ agentId, sessionId: session.id, userId }); + + await serverDB.insert(topics).values({ + id: 'cron-no-meta', + userId, + sessionId: session.id, + agentId, + trigger: 'cron', + title: 'No Meta', + metadata: {}, + }); + + const result = await topicModel.getCronTopicsGroupedByCronJob(agentId); + expect(result).toEqual([]); + }); + }); }); diff --git a/packages/database/src/models/__tests__/user.test.ts b/packages/database/src/models/__tests__/user.test.ts index b6583f3700..ee80a9b0b2 100644 --- a/packages/database/src/models/__tests__/user.test.ts +++ b/packages/database/src/models/__tests__/user.test.ts @@ -672,5 +672,29 @@ describe('UserModel', () => { expect(result.responseLanguage).toBe('en-US'); }); }); + + describe('getUserPreference', () => { + it('should return user preference after update', async () => { + await serverDB + .update(users) + .set({ preference: { telemetry: true, useCmdEnterKey: false } }) + .where(eq(users.id, userId)); + + const result = await userModel.getUserPreference(); + expect(result).toBeDefined(); + expect(result).toMatchObject({ telemetry: true, useCmdEnterKey: false }); + }); + + it('should return default preference for existing user', async () => { + const result = await userModel.getUserPreference(); + expect(result).toBeDefined(); + }); + + it('should return undefined for non-existent user', async () => { + const nonExistentModel = new UserModel(serverDB, 'non-existent-user'); + const result = await nonExistentModel.getUserPreference(); + expect(result).toBeUndefined(); + }); + }); }); }); diff --git a/packages/database/src/models/ragEval/__tests__/dataset.test.ts b/packages/database/src/models/ragEval/__tests__/dataset.test.ts new file mode 100644 index 0000000000..da6903b835 --- /dev/null +++ b/packages/database/src/models/ragEval/__tests__/dataset.test.ts @@ -0,0 +1,224 @@ +import { eq } from 'drizzle-orm'; +import { afterEach, beforeEach, describe, expect, it } from 'vitest'; + +import { getTestDB } from '../../../core/getTestDB'; +import { evalDatasets, knowledgeBases, users } from '../../../schemas'; +import { EvalDatasetModel } from '../dataset'; + +const serverDB = await getTestDB(); + +const userId = 'dataset-test-user'; +const userId2 = 'dataset-test-user-2'; +const datasetModel = new EvalDatasetModel(serverDB, userId); + +let knowledgeBaseId: string; + +beforeEach(async () => { + await serverDB.delete(evalDatasets); + await serverDB.delete(knowledgeBases); + await serverDB.delete(users); + + await serverDB.insert(users).values([{ id: userId }, { id: userId2 }]); + + const [kb] = await serverDB + .insert(knowledgeBases) + .values({ name: 'Test KB', userId }) + .returning(); + knowledgeBaseId = kb.id; +}); + +afterEach(async () => { + await serverDB.delete(evalDatasets); + await serverDB.delete(knowledgeBases); + await serverDB.delete(users); +}); + +describe('EvalDatasetModel', () => { + describe('create', () => { + it('should create a new dataset with userId', async () => { + const result = await datasetModel.create({ + knowledgeBaseId, + name: 'Test Dataset', + }); + + expect(result).toBeDefined(); + expect(result.name).toBe('Test Dataset'); + expect(result.knowledgeBaseId).toBe(knowledgeBaseId); + expect(result.userId).toBe(userId); + }); + + it('should create dataset with description', async () => { + const result = await datasetModel.create({ + knowledgeBaseId, + name: 'Dataset with desc', + description: 'A test dataset description', + }); + + expect(result.description).toBe('A test dataset description'); + }); + }); + + describe('delete', () => { + it('should delete a dataset owned by the user', async () => { + const [dataset] = await serverDB + .insert(evalDatasets) + .values({ knowledgeBaseId, name: 'Delete me', userId }) + .returning(); + + await datasetModel.delete(dataset.id); + + const deleted = await serverDB.query.evalDatasets.findFirst({ + where: eq(evalDatasets.id, dataset.id), + }); + expect(deleted).toBeUndefined(); + }); + + it('should not delete a dataset owned by another user', async () => { + const [dataset] = await serverDB + .insert(evalDatasets) + .values({ knowledgeBaseId, name: 'Other dataset', userId: userId2 }) + .returning(); + + await datasetModel.delete(dataset.id); + + const stillExists = await serverDB.query.evalDatasets.findFirst({ + where: eq(evalDatasets.id, dataset.id), + }); + expect(stillExists).toBeDefined(); + }); + }); + + describe('query', () => { + it('should query datasets by knowledgeBaseId for current user', async () => { + await serverDB.insert(evalDatasets).values([ + { knowledgeBaseId, name: 'Dataset 1', userId }, + { knowledgeBaseId, name: 'Dataset 2', userId }, + ]); + + const results = await datasetModel.query(knowledgeBaseId); + + expect(results).toHaveLength(2); + expect(results[0]).toHaveProperty('id'); + expect(results[0]).toHaveProperty('name'); + expect(results[0]).toHaveProperty('description'); + expect(results[0]).toHaveProperty('createdAt'); + expect(results[0]).toHaveProperty('updatedAt'); + }); + + it('should not return datasets from other users', async () => { + await serverDB.insert(evalDatasets).values([ + { knowledgeBaseId, name: 'My dataset', userId }, + { knowledgeBaseId, name: 'Other dataset', userId: userId2 }, + ]); + + const results = await datasetModel.query(knowledgeBaseId); + + expect(results).toHaveLength(1); + expect(results[0].name).toBe('My dataset'); + }); + + it('should return empty array for non-existent knowledge base', async () => { + const results = await datasetModel.query('non-existent'); + expect(results).toHaveLength(0); + }); + + it('should order results by createdAt desc', async () => { + const [d1] = await serverDB + .insert(evalDatasets) + .values({ knowledgeBaseId, name: 'First', userId }) + .returning(); + + const [d2] = await serverDB + .insert(evalDatasets) + .values({ knowledgeBaseId, name: 'Second', userId }) + .returning(); + + const results = await datasetModel.query(knowledgeBaseId); + + // Second should come first (desc order) + expect(results).toHaveLength(2); + }); + }); + + describe('findById', () => { + it('should find a dataset by id', async () => { + const [dataset] = await serverDB + .insert(evalDatasets) + .values({ knowledgeBaseId, name: 'Find me', userId }) + .returning(); + + const result = await datasetModel.findById(dataset.id); + + expect(result).toBeDefined(); + expect(result?.name).toBe('Find me'); + }); + + it('should not find dataset owned by another user', async () => { + const [dataset] = await serverDB + .insert(evalDatasets) + .values({ knowledgeBaseId, name: 'Other', userId: userId2 }) + .returning(); + + const result = await datasetModel.findById(dataset.id); + expect(result).toBeUndefined(); + }); + + it('should return undefined for non-existent id', async () => { + const result = await datasetModel.findById('non-existent'); + expect(result).toBeUndefined(); + }); + }); + + describe('update', () => { + it('should update a dataset owned by the user', async () => { + const [dataset] = await serverDB + .insert(evalDatasets) + .values({ knowledgeBaseId, name: 'Original', userId }) + .returning(); + + await datasetModel.update(dataset.id, { + name: 'Updated', + description: 'New description', + }); + + const updated = await serverDB.query.evalDatasets.findFirst({ + where: eq(evalDatasets.id, dataset.id), + }); + expect(updated?.name).toBe('Updated'); + expect(updated?.description).toBe('New description'); + }); + + it('should not update a dataset owned by another user', async () => { + const [dataset] = await serverDB + .insert(evalDatasets) + .values({ knowledgeBaseId, name: 'Other', userId: userId2 }) + .returning(); + + await datasetModel.update(dataset.id, { name: 'Hacked' }); + + const unchanged = await serverDB.query.evalDatasets.findFirst({ + where: eq(evalDatasets.id, dataset.id), + }); + expect(unchanged?.name).toBe('Other'); + }); + + it('should update the updatedAt timestamp', async () => { + const [dataset] = await serverDB + .insert(evalDatasets) + .values({ knowledgeBaseId, name: 'Timestamp test', userId }) + .returning(); + + const originalUpdatedAt = dataset.updatedAt; + + // Small delay to ensure different timestamp + await new Promise((resolve) => setTimeout(resolve, 10)); + + await datasetModel.update(dataset.id, { name: 'Updated' }); + + const updated = await serverDB.query.evalDatasets.findFirst({ + where: eq(evalDatasets.id, dataset.id), + }); + expect(updated?.updatedAt).not.toEqual(originalUpdatedAt); + }); + }); +}); diff --git a/packages/database/src/models/ragEval/__tests__/datasetRecord.test.ts b/packages/database/src/models/ragEval/__tests__/datasetRecord.test.ts new file mode 100644 index 0000000000..8151a473e4 --- /dev/null +++ b/packages/database/src/models/ragEval/__tests__/datasetRecord.test.ts @@ -0,0 +1,259 @@ +import { eq } from 'drizzle-orm'; +import { afterEach, beforeEach, describe, expect, it } from 'vitest'; + +import { getTestDB } from '../../../core/getTestDB'; +import { evalDatasetRecords, evalDatasets, files, knowledgeBases, users } from '../../../schemas'; +import { EvalDatasetRecordModel } from '../datasetRecord'; + +const serverDB = await getTestDB(); + +const userId = 'dataset-record-test-user'; +const userId2 = 'dataset-record-test-user-2'; +const recordModel = new EvalDatasetRecordModel(serverDB, userId); + +let datasetId: string; +let knowledgeBaseId: string; + +beforeEach(async () => { + await serverDB.delete(evalDatasetRecords); + await serverDB.delete(evalDatasets); + await serverDB.delete(files); + await serverDB.delete(knowledgeBases); + await serverDB.delete(users); + + await serverDB.insert(users).values([{ id: userId }, { id: userId2 }]); + + const [kb] = await serverDB + .insert(knowledgeBases) + .values({ name: 'Test KB', userId }) + .returning(); + knowledgeBaseId = kb.id; + + const [dataset] = await serverDB + .insert(evalDatasets) + .values({ knowledgeBaseId, name: 'Test Dataset', userId }) + .returning(); + datasetId = dataset.id; +}); + +afterEach(async () => { + await serverDB.delete(evalDatasetRecords); + await serverDB.delete(evalDatasets); + await serverDB.delete(files); + await serverDB.delete(knowledgeBases); + await serverDB.delete(users); +}); + +describe('EvalDatasetRecordModel', () => { + describe('create', () => { + it('should create a new record with userId', async () => { + const result = await recordModel.create({ + datasetId, + question: 'What is AI?', + ideal: 'Artificial Intelligence', + }); + + expect(result).toBeDefined(); + expect(result.datasetId).toBe(datasetId); + expect(result.question).toBe('What is AI?'); + expect(result.ideal).toBe('Artificial Intelligence'); + expect(result.userId).toBe(userId); + }); + + it('should create a record with referenceFiles', async () => { + const result = await recordModel.create({ + datasetId, + question: 'Test question', + referenceFiles: ['file-1', 'file-2'], + }); + + expect(result.referenceFiles).toEqual(['file-1', 'file-2']); + }); + }); + + describe('batchCreate', () => { + it('should batch create records', async () => { + const result = await recordModel.batchCreate([ + { datasetId, question: 'Q1', ideal: 'A1' }, + { datasetId, question: 'Q2', ideal: 'A2' }, + ]); + + expect(result).toBeDefined(); + expect(result.userId).toBe(userId); + + const allRecords = await serverDB.query.evalDatasetRecords.findMany({ + where: eq(evalDatasetRecords.datasetId, datasetId), + }); + expect(allRecords).toHaveLength(2); + }); + }); + + describe('delete', () => { + it('should delete a record owned by the user', async () => { + const [record] = await serverDB + .insert(evalDatasetRecords) + .values({ datasetId, question: 'Delete me', userId }) + .returning(); + + await recordModel.delete(record.id); + + const deleted = await serverDB.query.evalDatasetRecords.findFirst({ + where: eq(evalDatasetRecords.id, record.id), + }); + expect(deleted).toBeUndefined(); + }); + + it('should not delete a record owned by another user', async () => { + const [record] = await serverDB + .insert(evalDatasetRecords) + .values({ datasetId, question: 'Other user record', userId: userId2 }) + .returning(); + + await recordModel.delete(record.id); + + const stillExists = await serverDB.query.evalDatasetRecords.findFirst({ + where: eq(evalDatasetRecords.id, record.id), + }); + expect(stillExists).toBeDefined(); + }); + }); + + describe('query', () => { + it('should query records by datasetId with resolved reference files', async () => { + const [file1] = await serverDB + .insert(files) + .values({ + fileType: 'application/pdf', + name: 'doc.pdf', + size: 1024, + url: 'https://example.com/doc.pdf', + userId, + }) + .returning(); + + await serverDB.insert(evalDatasetRecords).values({ + datasetId, + question: 'Q1', + referenceFiles: [file1.id], + userId, + }); + + const results = await recordModel.query(datasetId); + + expect(results).toHaveLength(1); + expect(results[0].referenceFiles).toHaveLength(1); + expect(results[0].referenceFiles[0]).toMatchObject({ + id: file1.id, + name: 'doc.pdf', + fileType: 'application/pdf', + }); + }); + + it('should return records without reference files', async () => { + await serverDB.insert(evalDatasetRecords).values({ + datasetId, + question: 'Q1', + userId, + }); + + const results = await recordModel.query(datasetId); + + expect(results).toHaveLength(1); + expect(results[0].referenceFiles).toEqual([]); + }); + + it('should only return records for current user', async () => { + await serverDB.insert(evalDatasetRecords).values([ + { datasetId, question: 'User Q', userId }, + { datasetId, question: 'Other Q', userId: userId2 }, + ]); + + const results = await recordModel.query(datasetId); + + expect(results).toHaveLength(1); + expect(results[0].question).toBe('User Q'); + }); + }); + + describe('findByDatasetId', () => { + it('should find all records by datasetId for current user', async () => { + await serverDB.insert(evalDatasetRecords).values([ + { datasetId, question: 'Q1', userId }, + { datasetId, question: 'Q2', userId }, + { datasetId, question: 'Q3', userId: userId2 }, + ]); + + const results = await recordModel.findByDatasetId(datasetId); + + expect(results).toHaveLength(2); + }); + + it('should return empty array when no records found', async () => { + const results = await recordModel.findByDatasetId('non-existent'); + expect(results).toHaveLength(0); + }); + }); + + describe('findById', () => { + it('should find a record by id', async () => { + const [record] = await serverDB + .insert(evalDatasetRecords) + .values({ datasetId, question: 'Find me', userId }) + .returning(); + + const result = await recordModel.findById(record.id); + + expect(result).toBeDefined(); + expect(result?.question).toBe('Find me'); + }); + + it('should not find a record owned by another user', async () => { + const [record] = await serverDB + .insert(evalDatasetRecords) + .values({ datasetId, question: 'Other', userId: userId2 }) + .returning(); + + const result = await recordModel.findById(record.id); + expect(result).toBeUndefined(); + }); + + it('should return undefined for non-existent id', async () => { + const result = await recordModel.findById('non-existent'); + expect(result).toBeUndefined(); + }); + }); + + describe('update', () => { + it('should update a record owned by the user', async () => { + const [record] = await serverDB + .insert(evalDatasetRecords) + .values({ datasetId, question: 'Original', ideal: 'Original', userId }) + .returning(); + + await recordModel.update(record.id, { + question: 'Updated', + ideal: 'Updated ideal', + }); + + const updated = await serverDB.query.evalDatasetRecords.findFirst({ + where: eq(evalDatasetRecords.id, record.id), + }); + expect(updated?.question).toBe('Updated'); + expect(updated?.ideal).toBe('Updated ideal'); + }); + + it('should not update a record owned by another user', async () => { + const [record] = await serverDB + .insert(evalDatasetRecords) + .values({ datasetId, question: 'Other', userId: userId2 }) + .returning(); + + await recordModel.update(record.id, { question: 'Hacked' }); + + const unchanged = await serverDB.query.evalDatasetRecords.findFirst({ + where: eq(evalDatasetRecords.id, record.id), + }); + expect(unchanged?.question).toBe('Other'); + }); + }); +}); diff --git a/packages/database/src/models/ragEval/__tests__/evaluation.test.ts b/packages/database/src/models/ragEval/__tests__/evaluation.test.ts new file mode 100644 index 0000000000..fa54bbe099 --- /dev/null +++ b/packages/database/src/models/ragEval/__tests__/evaluation.test.ts @@ -0,0 +1,258 @@ +import { EvalEvaluationStatus } from '@lobechat/types'; +import { eq } from 'drizzle-orm'; +import { afterEach, beforeEach, describe, expect, it } from 'vitest'; + +import { getTestDB } from '../../../core/getTestDB'; +import { + evalDatasets, + evalEvaluation, + evaluationRecords, + knowledgeBases, + users, +} from '../../../schemas'; +import { EvalEvaluationModel } from '../evaluation'; + +const serverDB = await getTestDB(); + +const userId = 'eval-test-user'; +const userId2 = 'eval-test-user-2'; +const evalModel = new EvalEvaluationModel(serverDB, userId); + +let datasetId: string; +let knowledgeBaseId: string; + +beforeEach(async () => { + await serverDB.delete(evaluationRecords); + await serverDB.delete(evalEvaluation); + await serverDB.delete(evalDatasets); + await serverDB.delete(knowledgeBases); + await serverDB.delete(users); + + await serverDB.insert(users).values([{ id: userId }, { id: userId2 }]); + + const [kb] = await serverDB + .insert(knowledgeBases) + .values({ name: 'Test KB', userId }) + .returning(); + knowledgeBaseId = kb.id; + + const [dataset] = await serverDB + .insert(evalDatasets) + .values({ knowledgeBaseId, name: 'Test Dataset', userId }) + .returning(); + datasetId = dataset.id; +}); + +afterEach(async () => { + await serverDB.delete(evaluationRecords); + await serverDB.delete(evalEvaluation); + await serverDB.delete(evalDatasets); + await serverDB.delete(knowledgeBases); + await serverDB.delete(users); +}); + +describe('EvalEvaluationModel', () => { + describe('create', () => { + it('should create a new evaluation with userId', async () => { + const result = await evalModel.create({ + datasetId, + knowledgeBaseId, + name: 'Test Evaluation', + }); + + expect(result).toBeDefined(); + expect(result.name).toBe('Test Evaluation'); + expect(result.datasetId).toBe(datasetId); + expect(result.knowledgeBaseId).toBe(knowledgeBaseId); + expect(result.userId).toBe(userId); + }); + + it('should create evaluation with description', async () => { + const result = await evalModel.create({ + datasetId, + knowledgeBaseId, + name: 'Eval with desc', + description: 'A test evaluation', + }); + + expect(result.description).toBe('A test evaluation'); + }); + }); + + describe('delete', () => { + it('should delete an evaluation owned by the user', async () => { + const [evaluation] = await serverDB + .insert(evalEvaluation) + .values({ datasetId, knowledgeBaseId, name: 'Delete me', userId }) + .returning(); + + await evalModel.delete(evaluation.id); + + const deleted = await serverDB.query.evalEvaluation.findFirst({ + where: eq(evalEvaluation.id, evaluation.id), + }); + expect(deleted).toBeUndefined(); + }); + + it('should not delete an evaluation owned by another user', async () => { + const [evaluation] = await serverDB + .insert(evalEvaluation) + .values({ datasetId, knowledgeBaseId, name: 'Other eval', userId: userId2 }) + .returning(); + + await evalModel.delete(evaluation.id); + + const stillExists = await serverDB.query.evalEvaluation.findFirst({ + where: eq(evalEvaluation.id, evaluation.id), + }); + expect(stillExists).toBeDefined(); + }); + }); + + describe('queryByKnowledgeBaseId', () => { + it('should query evaluations with dataset info and record stats', async () => { + const [evaluation] = await serverDB + .insert(evalEvaluation) + .values({ datasetId, knowledgeBaseId, name: 'Eval 1', userId }) + .returning(); + + // Create a dataset record for the evaluation records to reference + const { evalDatasetRecords } = await import('../../../schemas'); + const [datasetRecord] = await serverDB + .insert(evalDatasetRecords) + .values({ datasetId, question: 'test q', userId }) + .returning(); + + await serverDB.insert(evaluationRecords).values([ + { + evaluationId: evaluation.id, + question: 'Q1', + status: EvalEvaluationStatus.Success, + datasetRecordId: datasetRecord.id, + userId, + }, + { + evaluationId: evaluation.id, + question: 'Q2', + status: EvalEvaluationStatus.Success, + datasetRecordId: datasetRecord.id, + userId, + }, + ]); + + const results = await evalModel.queryByKnowledgeBaseId(knowledgeBaseId); + + expect(results).toHaveLength(1); + expect(results[0].name).toBe('Eval 1'); + expect(results[0].dataset).toMatchObject({ id: datasetId, name: 'Test Dataset' }); + expect(results[0].recordsStats.total).toBe(2); + expect(results[0].recordsStats.success).toBe(2); + }); + + it('should return empty stats when no records exist', async () => { + await serverDB + .insert(evalEvaluation) + .values({ datasetId, knowledgeBaseId, name: 'Empty eval', userId }) + .returning(); + + const results = await evalModel.queryByKnowledgeBaseId(knowledgeBaseId); + + expect(results).toHaveLength(1); + expect(results[0].recordsStats).toEqual({ success: 0, total: 0 }); + }); + + it('should only return evaluations for current user', async () => { + await serverDB.insert(evalEvaluation).values([ + { datasetId, knowledgeBaseId, name: 'My eval', userId }, + { datasetId, knowledgeBaseId, name: 'Other eval', userId: userId2 }, + ]); + + const results = await evalModel.queryByKnowledgeBaseId(knowledgeBaseId); + + expect(results).toHaveLength(1); + expect(results[0].name).toBe('My eval'); + }); + + it('should return empty array for non-existent knowledge base', async () => { + const results = await evalModel.queryByKnowledgeBaseId('non-existent'); + expect(results).toHaveLength(0); + }); + }); + + describe('findById', () => { + it('should find an evaluation by id', async () => { + const [evaluation] = await serverDB + .insert(evalEvaluation) + .values({ datasetId, knowledgeBaseId, name: 'Find me', userId }) + .returning(); + + const result = await evalModel.findById(evaluation.id); + + expect(result).toBeDefined(); + expect(result?.name).toBe('Find me'); + }); + + it('should not find evaluation owned by another user', async () => { + const [evaluation] = await serverDB + .insert(evalEvaluation) + .values({ datasetId, knowledgeBaseId, name: 'Other', userId: userId2 }) + .returning(); + + const result = await evalModel.findById(evaluation.id); + expect(result).toBeUndefined(); + }); + + it('should return undefined for non-existent id', async () => { + const result = await evalModel.findById('non-existent'); + expect(result).toBeUndefined(); + }); + }); + + describe('update', () => { + it('should update an evaluation owned by the user', async () => { + const [evaluation] = await serverDB + .insert(evalEvaluation) + .values({ datasetId, knowledgeBaseId, name: 'Original', userId }) + .returning(); + + await evalModel.update(evaluation.id, { + name: 'Updated', + description: 'New description', + }); + + const updated = await serverDB.query.evalEvaluation.findFirst({ + where: eq(evalEvaluation.id, evaluation.id), + }); + expect(updated?.name).toBe('Updated'); + expect(updated?.description).toBe('New description'); + }); + + it('should not update an evaluation owned by another user', async () => { + const [evaluation] = await serverDB + .insert(evalEvaluation) + .values({ datasetId, knowledgeBaseId, name: 'Other', userId: userId2 }) + .returning(); + + await evalModel.update(evaluation.id, { name: 'Hacked' }); + + const unchanged = await serverDB.query.evalEvaluation.findFirst({ + where: eq(evalEvaluation.id, evaluation.id), + }); + expect(unchanged?.name).toBe('Other'); + }); + + it('should update status field', async () => { + const [evaluation] = await serverDB + .insert(evalEvaluation) + .values({ datasetId, knowledgeBaseId, name: 'Status test', userId }) + .returning(); + + await evalModel.update(evaluation.id, { status: EvalEvaluationStatus.Success }); + + const updated = await serverDB.query.evalEvaluation.findFirst({ + where: eq(evalEvaluation.id, evaluation.id), + }); + expect(updated?.status).toBe(EvalEvaluationStatus.Success); + }); + }); +}); diff --git a/packages/database/src/models/ragEval/__tests__/evaluationRecord.test.ts b/packages/database/src/models/ragEval/__tests__/evaluationRecord.test.ts new file mode 100644 index 0000000000..3386e387af --- /dev/null +++ b/packages/database/src/models/ragEval/__tests__/evaluationRecord.test.ts @@ -0,0 +1,276 @@ +import { EvalEvaluationStatus } from '@lobechat/types'; +import { eq } from 'drizzle-orm'; +import { afterEach, beforeEach, describe, expect, it } from 'vitest'; + +import { getTestDB } from '../../../core/getTestDB'; +import { + evalDatasetRecords, + evalDatasets, + evalEvaluation, + evaluationRecords, + knowledgeBases, + users, +} from '../../../schemas'; +import { EvaluationRecordModel } from '../evaluationRecord'; + +const serverDB = await getTestDB(); + +const userId = 'eval-record-test-user'; +const userId2 = 'eval-record-test-user-2'; +const recordModel = new EvaluationRecordModel(serverDB, userId); + +let datasetId: string; +let evaluationId: string; +let datasetRecordId: string; +let knowledgeBaseId: string; + +beforeEach(async () => { + await serverDB.delete(evaluationRecords); + await serverDB.delete(evalDatasetRecords); + await serverDB.delete(evalEvaluation); + await serverDB.delete(evalDatasets); + await serverDB.delete(knowledgeBases); + await serverDB.delete(users); + + await serverDB.insert(users).values([{ id: userId }, { id: userId2 }]); + + const [kb] = await serverDB + .insert(knowledgeBases) + .values({ name: 'Test KB', userId }) + .returning(); + knowledgeBaseId = kb.id; + + const [dataset] = await serverDB + .insert(evalDatasets) + .values({ knowledgeBaseId, name: 'Test Dataset', userId }) + .returning(); + datasetId = dataset.id; + + const [evaluation] = await serverDB + .insert(evalEvaluation) + .values({ datasetId, knowledgeBaseId, name: 'Test Evaluation', userId }) + .returning(); + evaluationId = evaluation.id; + + const [dsRecord] = await serverDB + .insert(evalDatasetRecords) + .values({ datasetId, question: 'Test Q', userId }) + .returning(); + datasetRecordId = dsRecord.id; +}); + +afterEach(async () => { + await serverDB.delete(evaluationRecords); + await serverDB.delete(evalDatasetRecords); + await serverDB.delete(evalEvaluation); + await serverDB.delete(evalDatasets); + await serverDB.delete(knowledgeBases); + await serverDB.delete(users); +}); + +describe('EvaluationRecordModel', () => { + describe('create', () => { + it('should create a new evaluation record with userId', async () => { + const result = await recordModel.create({ + evaluationId, + datasetRecordId, + question: 'What is AI?', + answer: 'Artificial Intelligence', + }); + + expect(result).toBeDefined(); + expect(result.evaluationId).toBe(evaluationId); + expect(result.datasetRecordId).toBe(datasetRecordId); + expect(result.question).toBe('What is AI?'); + expect(result.answer).toBe('Artificial Intelligence'); + expect(result.userId).toBe(userId); + }); + + it('should create a record with status and context', async () => { + const result = await recordModel.create({ + evaluationId, + datasetRecordId, + question: 'Test', + status: EvalEvaluationStatus.Success, + context: ['ctx1', 'ctx2'], + }); + + expect(result.status).toBe(EvalEvaluationStatus.Success); + expect(result.context).toEqual(['ctx1', 'ctx2']); + }); + }); + + describe('batchCreate', () => { + it('should batch create evaluation records', async () => { + const results = await recordModel.batchCreate([ + { evaluationId, datasetRecordId, question: 'Q1' }, + { evaluationId, datasetRecordId, question: 'Q2' }, + ]); + + expect(results).toHaveLength(2); + expect(results.every((r) => r.userId === userId)).toBe(true); + }); + }); + + describe('delete', () => { + it('should delete a record owned by the user', async () => { + const [record] = await serverDB + .insert(evaluationRecords) + .values({ + evaluationId, + datasetRecordId, + question: 'Delete me', + userId, + }) + .returning(); + + await recordModel.delete(record.id); + + const deleted = await serverDB.query.evaluationRecords.findFirst({ + where: eq(evaluationRecords.id, record.id), + }); + expect(deleted).toBeUndefined(); + }); + + it('should not delete a record owned by another user', async () => { + const [record] = await serverDB + .insert(evaluationRecords) + .values({ + evaluationId, + datasetRecordId, + question: 'Other user record', + userId: userId2, + }) + .returning(); + + await recordModel.delete(record.id); + + const stillExists = await serverDB.query.evaluationRecords.findFirst({ + where: eq(evaluationRecords.id, record.id), + }); + expect(stillExists).toBeDefined(); + }); + }); + + describe('query', () => { + it('should query records by evaluationId (reportId)', async () => { + await serverDB.insert(evaluationRecords).values([ + { evaluationId, datasetRecordId, question: 'Q1', userId }, + { evaluationId, datasetRecordId, question: 'Q2', userId }, + ]); + + const results = await recordModel.query(evaluationId); + + expect(results).toHaveLength(2); + }); + + it('should only return records for current user', async () => { + await serverDB.insert(evaluationRecords).values([ + { evaluationId, datasetRecordId, question: 'My Q', userId }, + { evaluationId, datasetRecordId, question: 'Other Q', userId: userId2 }, + ]); + + const results = await recordModel.query(evaluationId); + + expect(results).toHaveLength(1); + expect(results[0].question).toBe('My Q'); + }); + + it('should return empty array for non-existent evaluationId', async () => { + const results = await recordModel.query('non-existent'); + expect(results).toHaveLength(0); + }); + }); + + describe('findById', () => { + it('should find a record by id', async () => { + const [record] = await serverDB + .insert(evaluationRecords) + .values({ evaluationId, datasetRecordId, question: 'Find me', userId }) + .returning(); + + const result = await recordModel.findById(record.id); + + expect(result).toBeDefined(); + expect(result?.question).toBe('Find me'); + }); + + it('should not find a record owned by another user', async () => { + const [record] = await serverDB + .insert(evaluationRecords) + .values({ evaluationId, datasetRecordId, question: 'Other', userId: userId2 }) + .returning(); + + const result = await recordModel.findById(record.id); + expect(result).toBeUndefined(); + }); + + it('should return undefined for non-existent id', async () => { + const result = await recordModel.findById('non-existent'); + expect(result).toBeUndefined(); + }); + }); + + describe('findByEvaluationId', () => { + it('should find all records by evaluationId', async () => { + await serverDB.insert(evaluationRecords).values([ + { evaluationId, datasetRecordId, question: 'Q1', userId }, + { evaluationId, datasetRecordId, question: 'Q2', userId }, + ]); + + const results = await recordModel.findByEvaluationId(evaluationId); + + expect(results).toHaveLength(2); + }); + + it('should only return records for current user', async () => { + await serverDB.insert(evaluationRecords).values([ + { evaluationId, datasetRecordId, question: 'My Q', userId }, + { evaluationId, datasetRecordId, question: 'Other Q', userId: userId2 }, + ]); + + const results = await recordModel.findByEvaluationId(evaluationId); + + expect(results).toHaveLength(1); + }); + + it('should return empty array when no records found', async () => { + const results = await recordModel.findByEvaluationId('non-existent'); + expect(results).toHaveLength(0); + }); + }); + + describe('update', () => { + it('should update a record owned by the user', async () => { + const [record] = await serverDB + .insert(evaluationRecords) + .values({ evaluationId, datasetRecordId, question: 'Original', userId }) + .returning(); + + await recordModel.update(record.id, { + answer: 'New answer', + status: EvalEvaluationStatus.Success, + }); + + const updated = await serverDB.query.evaluationRecords.findFirst({ + where: eq(evaluationRecords.id, record.id), + }); + expect(updated?.answer).toBe('New answer'); + expect(updated?.status).toBe(EvalEvaluationStatus.Success); + }); + + it('should not update a record owned by another user', async () => { + const [record] = await serverDB + .insert(evaluationRecords) + .values({ evaluationId, datasetRecordId, question: 'Other', userId: userId2 }) + .returning(); + + await recordModel.update(record.id, { answer: 'Hacked' }); + + const unchanged = await serverDB.query.evaluationRecords.findFirst({ + where: eq(evaluationRecords.id, record.id), + }); + expect(unchanged?.answer).toBeNull(); + }); + }); +}); diff --git a/packages/database/src/models/userMemory/__tests__/model.test.ts b/packages/database/src/models/userMemory/__tests__/model.test.ts new file mode 100644 index 0000000000..5ae00f2f1a --- /dev/null +++ b/packages/database/src/models/userMemory/__tests__/model.test.ts @@ -0,0 +1,1598 @@ +// @vitest-environment node +import { + ActivityTypeEnum, + IdentityTypeEnum, + LayersEnum, + RelationshipEnum, + UserMemoryContextObjectType, +} from '@lobechat/types'; +import { eq } from 'drizzle-orm'; +import { beforeEach, describe, expect, it } from 'vitest'; + +import { getTestDB } from '../../../core/getTestDB'; +import { + userMemories, + userMemoriesActivities, + userMemoriesContexts, + userMemoriesExperiences, + userMemoriesIdentities, + userMemoriesPreferences, + users, +} from '../../../schemas'; +import type { LobeChatDatabase } from '../../../type'; +import { UserMemoryModel } from '../model'; + +const userId = 'memory-model-test-user'; +const otherUserId = 'other-memory-model-user'; + +let memoryModel: UserMemoryModel; +const serverDB: LobeChatDatabase = await getTestDB(); + +beforeEach(async () => { + await serverDB.delete(userMemoriesActivities); + await serverDB.delete(userMemoriesContexts); + await serverDB.delete(userMemoriesExperiences); + await serverDB.delete(userMemoriesIdentities); + await serverDB.delete(userMemoriesPreferences); + await serverDB.delete(userMemories); + await serverDB.delete(users); + + await serverDB.insert(users).values([{ id: userId }, { id: otherUserId }]); + memoryModel = new UserMemoryModel(serverDB, userId); +}); + +// Helper to create a base memory + identity pair +async function createIdentityPair(opts: { + baseTitle?: string; + description?: string; + relationship?: string; + role?: string; + tags?: string[]; + type?: string; + user?: string; +}) { + const uid = opts.user ?? userId; + const [mem] = await serverDB + .insert(userMemories) + .values({ + details: 'details', + lastAccessedAt: new Date(), + memoryLayer: 'identity', + memoryType: 'identity', + summary: 'summary', + tags: opts.tags, + title: opts.baseTitle ?? 'Identity memory', + userId: uid, + }) + .returning(); + + const [id] = await serverDB + .insert(userMemoriesIdentities) + .values({ + description: opts.description ?? 'A test identity', + relationship: opts.relationship ?? RelationshipEnum.Self, + role: opts.role, + tags: opts.tags, + type: opts.type ?? 'personal', + userId: uid, + userMemoryId: mem.id, + }) + .returning(); + + return { identity: id, memory: mem }; +} + +// Helper to create a base memory + experience pair +async function createExperiencePair(opts?: { + action?: string; + tags?: string[]; + type?: string; + user?: string; +}) { + const uid = opts?.user ?? userId; + const [mem] = await serverDB + .insert(userMemories) + .values({ + details: 'exp details', + lastAccessedAt: new Date(), + memoryLayer: 'experience', + memoryType: 'experience', + summary: 'exp summary', + tags: opts?.tags, + title: 'Experience memory', + userId: uid, + }) + .returning(); + + const [exp] = await serverDB + .insert(userMemoriesExperiences) + .values({ + action: opts?.action ?? 'did something', + keyLearning: 'learned stuff', + situation: 'a situation', + tags: opts?.tags, + type: opts?.type ?? 'learning', + userId: uid, + userMemoryId: mem.id, + }) + .returning(); + + return { experience: exp, memory: mem }; +} + +// Helper to create a base memory + preference pair +async function createPreferencePair(opts?: { + conclusionDirectives?: string; + tags?: string[]; + type?: string; + user?: string; +}) { + const uid = opts?.user ?? userId; + const [mem] = await serverDB + .insert(userMemories) + .values({ + details: 'pref details', + lastAccessedAt: new Date(), + memoryLayer: 'preference', + memoryType: 'preference', + summary: 'pref summary', + tags: opts?.tags, + title: 'Preference memory', + userId: uid, + }) + .returning(); + + const [pref] = await serverDB + .insert(userMemoriesPreferences) + .values({ + conclusionDirectives: opts?.conclusionDirectives ?? 'use dark mode', + tags: opts?.tags, + type: opts?.type ?? 'ui', + userId: uid, + userMemoryId: mem.id, + }) + .returning(); + + return { memory: mem, preference: pref }; +} + +// Helper to create a base memory + activity pair +async function createActivityPair(opts?: { + status?: string; + tags?: string[]; + type?: string; + user?: string; +}) { + const uid = opts?.user ?? userId; + const [mem] = await serverDB + .insert(userMemories) + .values({ + details: 'activity details', + lastAccessedAt: new Date(), + memoryLayer: 'activity', + memoryType: 'activity', + summary: 'activity summary', + tags: opts?.tags, + title: 'Activity memory', + userId: uid, + }) + .returning(); + + const [act] = await serverDB + .insert(userMemoriesActivities) + .values({ + narrative: 'did a thing', + status: opts?.status ?? 'completed', + tags: opts?.tags, + type: opts?.type ?? 'task', + userId: uid, + userMemoryId: mem.id, + }) + .returning(); + + return { activity: act, memory: mem }; +} + +// Helper to create a base memory + context pair +async function createContextPair(opts?: { + description?: string; + tags?: string[]; + title?: string; + type?: string; + user?: string; +}) { + const uid = opts?.user ?? userId; + const [mem] = await serverDB + .insert(userMemories) + .values({ + details: 'context details', + lastAccessedAt: new Date(), + memoryLayer: 'context', + memoryType: 'context', + summary: 'context summary', + tags: opts?.tags, + title: opts?.title ?? 'Context memory', + userId: uid, + }) + .returning(); + + const [ctx] = await serverDB + .insert(userMemoriesContexts) + .values({ + description: opts?.description ?? 'A context description', + tags: opts?.tags, + title: opts?.title ?? 'A context', + type: opts?.type ?? 'project', + userId: uid, + userMemoryIds: [mem.id], + }) + .returning(); + + return { context: ctx, memory: mem }; +} + +describe('UserMemoryModel', () => { + // ========== Static Methods ========== + describe('parseAssociatedObjects', () => { + it('should return empty array for non-array input', () => { + expect(UserMemoryModel.parseAssociatedObjects(undefined)).toEqual([]); + expect(UserMemoryModel.parseAssociatedObjects('string')).toEqual([]); + expect(UserMemoryModel.parseAssociatedObjects(null)).toEqual([]); + }); + + it('should parse items with name field', () => { + const result = UserMemoryModel.parseAssociatedObjects([{ name: 'test' }]); + expect(result).toHaveLength(1); + expect(result[0]).toMatchObject({ name: 'test' }); + }); + + it('should skip invalid items', () => { + const result = UserMemoryModel.parseAssociatedObjects([null, 42, { noName: true }]); + expect(result).toEqual([]); + }); + }); + + describe('parseAssociatedSubjects', () => { + it('should return empty array for non-array input', () => { + expect(UserMemoryModel.parseAssociatedSubjects(undefined)).toEqual([]); + }); + + it('should parse items with name field', () => { + const result = UserMemoryModel.parseAssociatedSubjects([{ name: 'subject' }]); + expect(result).toHaveLength(1); + }); + }); + + describe('parseAssociatedLocations', () => { + it('should return empty array for null/undefined', () => { + expect(UserMemoryModel.parseAssociatedLocations(null)).toEqual([]); + expect(UserMemoryModel.parseAssociatedLocations(undefined)).toEqual([]); + }); + + it('should parse array of locations', () => { + const result = UserMemoryModel.parseAssociatedLocations([ + { address: '123 Main St', name: 'Home', type: 'residential' }, + ]); + expect(result).toHaveLength(1); + expect(result[0]).toEqual({ + address: '123 Main St', + name: 'Home', + tags: undefined, + type: 'residential', + }); + }); + + it('should handle object input (wraps in array)', () => { + const result = UserMemoryModel.parseAssociatedLocations({ name: 'Office' } as any); + expect(result).toHaveLength(1); + expect(result[0].name).toBe('Office'); + }); + + it('should handle tags array', () => { + const result = UserMemoryModel.parseAssociatedLocations([ + { name: 'Place', tags: ['tag1', 'tag2'] }, + ]); + expect(result[0].tags).toEqual(['tag1', 'tag2']); + }); + + it('should skip items with no valid fields', () => { + const result = UserMemoryModel.parseAssociatedLocations([{ invalid: true } as any]); + expect(result).toEqual([]); + }); + }); + + describe('parseDateFromString', () => { + it('should return null for falsy input', () => { + expect(UserMemoryModel.parseDateFromString(null)).toBeNull(); + expect(UserMemoryModel.parseDateFromString(undefined)).toBeNull(); + expect(UserMemoryModel.parseDateFromString('')).toBeNull(); + }); + + it('should parse valid date string', () => { + const result = UserMemoryModel.parseDateFromString('2024-01-01T00:00:00Z'); + expect(result).toBeInstanceOf(Date); + }); + + it('should return Date as-is if valid', () => { + const date = new Date('2024-01-01'); + expect(UserMemoryModel.parseDateFromString(date)).toBe(date); + }); + + it('should return null for invalid Date', () => { + expect(UserMemoryModel.parseDateFromString(new Date('invalid'))).toBeNull(); + }); + + it('should return null for non-string input', () => { + expect(UserMemoryModel.parseDateFromString(42 as any)).toBeNull(); + }); + }); + + // ========== queryTags ========== + describe('queryTags', () => { + it('should return grouped tags with counts', async () => { + await serverDB.insert(userMemories).values([ + { + lastAccessedAt: new Date(), + memoryLayer: 'context', + tags: ['work', 'coding'], + title: 'M1', + userId, + }, + { + lastAccessedAt: new Date(), + memoryLayer: 'context', + tags: ['work', 'design'], + title: 'M2', + userId, + }, + ]); + + const result = await memoryModel.queryTags(); + + expect(result.length).toBeGreaterThanOrEqual(2); + const workTag = result.find((r) => r.tag === 'work'); + expect(workTag?.count).toBe(2); + }); + + it('should filter by layers', async () => { + await serverDB.insert(userMemories).values([ + { + lastAccessedAt: new Date(), + memoryLayer: 'context', + tags: ['ctx-tag'], + title: 'M1', + userId, + }, + { + lastAccessedAt: new Date(), + memoryLayer: 'experience', + tags: ['exp-tag'], + title: 'M2', + userId, + }, + ]); + + const result = await memoryModel.queryTags({ layers: [LayersEnum.Context] }); + const tags = result.map((r) => r.tag); + expect(tags).toContain('ctx-tag'); + expect(tags).not.toContain('exp-tag'); + }); + + it('should respect pagination', async () => { + await serverDB.insert(userMemories).values( + Array.from({ length: 15 }, (_, i) => ({ + lastAccessedAt: new Date(), + memoryLayer: 'context', + tags: [`tag-${i}`], + title: `M${i}`, + userId, + })), + ); + + const page1 = await memoryModel.queryTags({ page: 1, size: 5 }); + expect(page1).toHaveLength(5); + + const page2 = await memoryModel.queryTags({ page: 2, size: 5 }); + expect(page2).toHaveLength(5); + }); + + it('should not include other user tags', async () => { + await serverDB.insert(userMemories).values([ + { + lastAccessedAt: new Date(), + memoryLayer: 'context', + tags: ['my-tag'], + title: 'M1', + userId, + }, + { + lastAccessedAt: new Date(), + memoryLayer: 'context', + tags: ['other-tag'], + title: 'M2', + userId: otherUserId, + }, + ]); + + const result = await memoryModel.queryTags(); + const tags = result.map((r) => r.tag); + expect(tags).toContain('my-tag'); + expect(tags).not.toContain('other-tag'); + }); + + it('should return empty array when no tags exist', async () => { + const result = await memoryModel.queryTags(); + expect(result).toEqual([]); + }); + }); + + // ========== queryIdentityRoles ========== + describe('queryIdentityRoles', () => { + it('should return tags and roles from self-relationship identities', async () => { + await createIdentityPair({ + role: 'developer', + tags: ['tech'], + }); + await createIdentityPair({ + role: 'developer', + tags: ['tech', 'senior'], + }); + + const result = await memoryModel.queryIdentityRoles(); + + expect(result.tags.length).toBeGreaterThanOrEqual(1); + const techTag = result.tags.find((t) => t.tag === 'tech'); + expect(techTag?.count).toBe(2); + + expect(result.roles.length).toBeGreaterThanOrEqual(1); + const devRole = result.roles.find((r) => r.role === 'developer'); + expect(devRole?.count).toBe(2); + }); + + it('should not include other user identity roles', async () => { + await createIdentityPair({ role: 'my-role' }); + await createIdentityPair({ role: 'other-role', user: otherUserId }); + + const result = await memoryModel.queryIdentityRoles(); + const roles = result.roles.map((r) => r.role); + expect(roles).toContain('my-role'); + expect(roles).not.toContain('other-role'); + }); + + it('should return empty when no identities', async () => { + const result = await memoryModel.queryIdentityRoles(); + expect(result).toEqual({ roles: [], tags: [] }); + }); + + it('should respect pagination', async () => { + for (let i = 0; i < 12; i++) { + await createIdentityPair({ role: `role-${i}` }); + } + const result = await memoryModel.queryIdentityRoles({ size: 5 }); + expect(result.roles.length).toBeLessThanOrEqual(5); + }); + }); + + // ========== queryMemories ========== + describe('queryMemories', () => { + describe('context layer', () => { + it('should return context memories with pagination info', async () => { + await createContextPair({}); + await createContextPair({ title: 'Second context' }); + + const result = await memoryModel.queryMemories({ layer: LayersEnum.Context }); + + expect(result.items.length).toBe(2); + expect(result.page).toBe(1); + expect(result.pageSize).toBe(20); + expect(result.total).toBe(2); + }); + + it('should filter by text query', async () => { + await createContextPair({ title: 'Apple project' }); + await createContextPair({ title: 'Banana project' }); + + const result = await memoryModel.queryMemories({ + layer: LayersEnum.Context, + q: 'Apple', + }); + + expect(result.items.length).toBe(1); + }); + + it('should not return other user data', async () => { + await createContextPair({}); + await createContextPair({ user: otherUserId }); + + const result = await memoryModel.queryMemories({ layer: LayersEnum.Context }); + expect(result.total).toBe(1); + }); + }); + + describe('activity layer', () => { + it('should return activity memories', async () => { + await createActivityPair({}); + + const result = await memoryModel.queryMemories({ layer: LayersEnum.Activity }); + + expect(result.items.length).toBe(1); + expect(result.total).toBe(1); + }); + }); + + describe('experience layer', () => { + it('should return experience memories', async () => { + await createExperiencePair({}); + + const result = await memoryModel.queryMemories({ layer: LayersEnum.Experience }); + + expect(result.items.length).toBe(1); + expect(result.total).toBe(1); + }); + }); + + describe('identity layer', () => { + it('should return identity memories', async () => { + await createIdentityPair({}); + + const result = await memoryModel.queryMemories({ layer: LayersEnum.Identity }); + + expect(result.items.length).toBe(1); + expect(result.total).toBe(1); + }); + }); + + describe('preference layer', () => { + it('should return preference memories', async () => { + await createPreferencePair({}); + + const result = await memoryModel.queryMemories({ layer: LayersEnum.Preference }); + + expect(result.items.length).toBe(1); + expect(result.total).toBe(1); + }); + }); + + describe('pagination', () => { + it('should normalize negative page to 1', async () => { + await createActivityPair({}); + + const result = await memoryModel.queryMemories({ + layer: LayersEnum.Activity, + page: -1, + }); + + expect(result.page).toBe(1); + }); + + it('should cap pageSize at 100', async () => { + const result = await memoryModel.queryMemories({ + layer: LayersEnum.Activity, + pageSize: 200, + }); + + expect(result.pageSize).toBe(100); + }); + + it('should return empty items for unknown layer', async () => { + const result = await memoryModel.queryMemories({ + layer: 'unknown-layer' as any, + }); + + expect(result.items).toEqual([]); + expect(result.total).toBe(0); + }); + }); + }); + + // ========== listMemories ========== + describe('listMemories', () => { + it('should list experience memories', async () => { + await createExperiencePair({}); + + const result = await memoryModel.listMemories({ layer: LayersEnum.Experience }); + + expect(result).toHaveLength(1); + }); + + it('should list identity memories', async () => { + await createIdentityPair({}); + + const result = await memoryModel.listMemories({ layer: LayersEnum.Identity }); + + expect(result).toHaveLength(1); + }); + + it('should list preference memories', async () => { + await createPreferencePair({}); + + const result = await memoryModel.listMemories({ layer: LayersEnum.Preference }); + + expect(result).toHaveLength(1); + }); + + it('should list context memories', async () => { + await createContextPair({}); + + const result = await memoryModel.listMemories({ layer: LayersEnum.Context }); + + expect(result).toHaveLength(1); + }); + + it('should respect pagination parameters', async () => { + for (let i = 0; i < 5; i++) { + await createExperiencePair({}); + } + + const result = await memoryModel.listMemories({ + layer: LayersEnum.Experience, + pageSize: 2, + }); + + expect(result).toHaveLength(2); + }); + + it('should not return other user memories', async () => { + await createExperiencePair({}); + await createExperiencePair({ user: otherUserId }); + + const result = await memoryModel.listMemories({ layer: LayersEnum.Experience }); + + expect(result).toHaveLength(1); + }); + }); + + // ========== getMemoryDetail ========== + describe('getMemoryDetail', () => { + it('should get context detail', async () => { + const { context } = await createContextPair({}); + + const result = await memoryModel.getMemoryDetail({ + id: context.id, + layer: LayersEnum.Context, + }); + + expect(result).toBeDefined(); + expect((result as any)?.context).toBeDefined(); + expect(result?.memory).toBeDefined(); + expect(result?.layer).toBe(LayersEnum.Context); + }); + + it('should get activity detail', async () => { + const { activity } = await createActivityPair({}); + + const result = await memoryModel.getMemoryDetail({ + id: activity.id, + layer: LayersEnum.Activity, + }); + + expect(result).toBeDefined(); + expect((result as any)?.activity).toBeDefined(); + expect(result?.layer).toBe(LayersEnum.Activity); + }); + + it('should get experience detail', async () => { + const { experience } = await createExperiencePair({}); + + const result = await memoryModel.getMemoryDetail({ + id: experience.id, + layer: LayersEnum.Experience, + }); + + expect(result).toBeDefined(); + expect((result as any)?.experience).toBeDefined(); + expect(result?.layer).toBe(LayersEnum.Experience); + }); + + it('should get identity detail', async () => { + const { identity } = await createIdentityPair({}); + + const result = await memoryModel.getMemoryDetail({ + id: identity.id, + layer: LayersEnum.Identity, + }); + + expect(result).toBeDefined(); + expect((result as any)?.identity).toBeDefined(); + expect(result?.layer).toBe(LayersEnum.Identity); + }); + + it('should get preference detail', async () => { + const { preference } = await createPreferencePair({}); + + const result = await memoryModel.getMemoryDetail({ + id: preference.id, + layer: LayersEnum.Preference, + }); + + expect(result).toBeDefined(); + expect((result as any)?.preference).toBeDefined(); + expect(result?.layer).toBe(LayersEnum.Preference); + }); + + it('should return undefined for non-existent id', async () => { + const result = await memoryModel.getMemoryDetail({ + id: 'non-existent', + layer: LayersEnum.Context, + }); + + expect(result).toBeUndefined(); + }); + + it('should not return other user detail', async () => { + const { identity } = await createIdentityPair({ user: otherUserId }); + + const result = await memoryModel.getMemoryDetail({ + id: identity.id, + layer: LayersEnum.Identity, + }); + + expect(result).toBeUndefined(); + }); + }); + + // ========== searchActivities ========== + describe('searchActivities', () => { + it('should return activities for current user (no embedding)', async () => { + await createActivityPair({ type: 'task' }); + await createActivityPair({ type: 'event' }); + + const result = await memoryModel.searchActivities({}); + + expect(result).toHaveLength(2); + }); + + it('should filter by type', async () => { + await createActivityPair({ type: 'task' }); + await createActivityPair({ type: 'event' }); + + const result = await memoryModel.searchActivities({ type: 'task' }); + + expect(result).toHaveLength(1); + }); + + it('should respect limit', async () => { + for (let i = 0; i < 10; i++) { + await createActivityPair({ type: 'task' }); + } + + const result = await memoryModel.searchActivities({ limit: 3 }); + + expect(result).toHaveLength(3); + }); + + it('should return empty array for limit <= 0', async () => { + await createActivityPair({}); + + const result = await memoryModel.searchActivities({ limit: 0 }); + + expect(result).toEqual([]); + }); + + it('should not return other user activities', async () => { + await createActivityPair({}); + await createActivityPair({ user: otherUserId }); + + const result = await memoryModel.searchActivities({}); + + expect(result).toHaveLength(1); + }); + }); + + // ========== searchContexts ========== + describe('searchContexts', () => { + it('should return contexts for current user (no embedding)', async () => { + await createContextPair({}); + + const result = await memoryModel.searchContexts({}); + + expect(result).toHaveLength(1); + }); + + it('should filter by type', async () => { + await createContextPair({ type: 'project' }); + await createContextPair({ type: 'meeting' }); + + const result = await memoryModel.searchContexts({ type: 'project' }); + + expect(result).toHaveLength(1); + }); + + it('should return empty for limit <= 0', async () => { + await createContextPair({}); + expect(await memoryModel.searchContexts({ limit: 0 })).toEqual([]); + }); + }); + + // ========== searchExperiences ========== + describe('searchExperiences', () => { + it('should return experiences for current user', async () => { + await createExperiencePair({}); + + const result = await memoryModel.searchExperiences({}); + + expect(result).toHaveLength(1); + }); + + it('should filter by type', async () => { + await createExperiencePair({ type: 'learning' }); + await createExperiencePair({ type: 'failure' }); + + const result = await memoryModel.searchExperiences({ type: 'learning' }); + + expect(result).toHaveLength(1); + }); + + it('should return empty for limit <= 0', async () => { + await createExperiencePair({}); + expect(await memoryModel.searchExperiences({ limit: 0 })).toEqual([]); + }); + }); + + // ========== searchPreferences ========== + describe('searchPreferences', () => { + it('should return preferences for current user', async () => { + await createPreferencePair({}); + + const result = await memoryModel.searchPreferences({}); + + expect(result).toHaveLength(1); + }); + + it('should filter by type', async () => { + await createPreferencePair({ type: 'ui' }); + await createPreferencePair({ type: 'language' }); + + const result = await memoryModel.searchPreferences({ type: 'ui' }); + + expect(result).toHaveLength(1); + }); + + it('should return empty for limit <= 0', async () => { + await createPreferencePair({}); + expect(await memoryModel.searchPreferences({ limit: 0 })).toEqual([]); + }); + }); + + // ========== updateUserMemoryVectors ========== + describe('updateUserMemoryVectors', () => { + it('should update vectors on base memory', async () => { + const [mem] = await serverDB + .insert(userMemories) + .values({ + lastAccessedAt: new Date(), + memoryLayer: 'context', + title: 'Vector test', + userId, + }) + .returning(); + + const vector1024 = Array.from({ length: 1024 }, () => Math.random()); + await memoryModel.updateUserMemoryVectors(mem.id, { + detailsVector1024: vector1024, + summaryVector1024: vector1024, + }); + + const updated = await serverDB.query.userMemories.findFirst({ + where: eq(userMemories.id, mem.id), + }); + expect(updated?.detailsVector1024).toBeDefined(); + expect(updated?.summaryVector1024).toBeDefined(); + }); + + it('should skip update when no vectors provided', async () => { + const [mem] = await serverDB + .insert(userMemories) + .values({ + lastAccessedAt: new Date(), + memoryLayer: 'context', + title: 'No vector test', + userId, + }) + .returning(); + + const beforeUpdate = await serverDB.query.userMemories.findFirst({ + where: eq(userMemories.id, mem.id), + }); + + await memoryModel.updateUserMemoryVectors(mem.id, {}); + + const afterUpdate = await serverDB.query.userMemories.findFirst({ + where: eq(userMemories.id, mem.id), + }); + expect(afterUpdate?.updatedAt.getTime()).toBe(beforeUpdate?.updatedAt.getTime()); + }); + }); + + // ========== updateContextVectors ========== + describe('updateContextVectors', () => { + it('should update description vector', async () => { + const { context } = await createContextPair({}); + const vector1024 = Array.from({ length: 1024 }, () => Math.random()); + + await memoryModel.updateContextVectors(context.id, { descriptionVector: vector1024 }); + + const updated = await serverDB.query.userMemoriesContexts.findFirst({ + where: eq(userMemoriesContexts.id, context.id), + }); + expect(updated?.descriptionVector).toBeDefined(); + }); + + it('should skip when no vectors provided', async () => { + const { context } = await createContextPair({}); + + await memoryModel.updateContextVectors(context.id, {}); + // No error means success + }); + }); + + // ========== updatePreferenceVectors ========== + describe('updatePreferenceVectors', () => { + it('should update conclusion directives vector', async () => { + const { preference } = await createPreferencePair({}); + const vector1024 = Array.from({ length: 1024 }, () => Math.random()); + + await memoryModel.updatePreferenceVectors(preference.id, { + conclusionDirectivesVector: vector1024, + }); + + const updated = await serverDB.query.userMemoriesPreferences.findFirst({ + where: eq(userMemoriesPreferences.id, preference.id), + }); + expect(updated?.conclusionDirectivesVector).toBeDefined(); + }); + + it('should skip when no vectors provided', async () => { + const { preference } = await createPreferencePair({}); + await memoryModel.updatePreferenceVectors(preference.id, {}); + }); + }); + + // ========== updateIdentityVectors ========== + describe('updateIdentityVectors', () => { + it('should update description vector', async () => { + const { identity } = await createIdentityPair({}); + const vector1024 = Array.from({ length: 1024 }, () => Math.random()); + + await memoryModel.updateIdentityVectors(identity.id, { descriptionVector: vector1024 }); + + const updated = await serverDB.query.userMemoriesIdentities.findFirst({ + where: eq(userMemoriesIdentities.id, identity.id), + }); + expect(updated?.descriptionVector).toBeDefined(); + }); + + it('should skip when no vectors provided', async () => { + const { identity } = await createIdentityPair({}); + await memoryModel.updateIdentityVectors(identity.id, {}); + }); + }); + + // ========== updateExperienceVectors ========== + describe('updateExperienceVectors', () => { + it('should update multiple vectors', async () => { + const { experience } = await createExperiencePair({}); + const vector1024 = Array.from({ length: 1024 }, () => Math.random()); + + await memoryModel.updateExperienceVectors(experience.id, { + actionVector: vector1024, + keyLearningVector: vector1024, + situationVector: vector1024, + }); + + const updated = await serverDB.query.userMemoriesExperiences.findFirst({ + where: eq(userMemoriesExperiences.id, experience.id), + }); + expect(updated?.actionVector).toBeDefined(); + expect(updated?.keyLearningVector).toBeDefined(); + expect(updated?.situationVector).toBeDefined(); + }); + + it('should skip when no vectors provided', async () => { + const { experience } = await createExperiencePair({}); + await memoryModel.updateExperienceVectors(experience.id, {}); + }); + }); + + // ========== updateActivityVectors ========== + describe('updateActivityVectors', () => { + it('should update narrative and feedback vectors', async () => { + const { activity } = await createActivityPair({}); + const vector1024 = Array.from({ length: 1024 }, () => Math.random()); + + await memoryModel.updateActivityVectors(activity.id, { + feedbackVector: vector1024, + narrativeVector: vector1024, + }); + + const updated = await serverDB.query.userMemoriesActivities.findFirst({ + where: eq(userMemoriesActivities.id, activity.id), + }); + expect(updated?.narrativeVector).toBeDefined(); + expect(updated?.feedbackVector).toBeDefined(); + }); + + it('should skip when no vectors provided', async () => { + const { activity } = await createActivityPair({}); + await memoryModel.updateActivityVectors(activity.id, {}); + }); + }); + + // ========== addIdentityEntry ========== + describe('addIdentityEntry', () => { + it('should create both base memory and identity in transaction', async () => { + const result = await memoryModel.addIdentityEntry({ + base: { + details: 'I am a developer', + summary: 'Developer identity', + title: 'Developer', + }, + identity: { + description: 'Software developer', + relationship: RelationshipEnum.Self, + role: 'developer', + type: IdentityTypeEnum.Personal, + }, + }); + + expect(result.identityId).toBeDefined(); + expect(result.userMemoryId).toBeDefined(); + + const mem = await serverDB.query.userMemories.findFirst({ + where: eq(userMemories.id, result.userMemoryId), + }); + expect(mem?.userId).toBe(userId); + expect(mem?.memoryLayer).toBe('identity'); + + const identity = await serverDB.query.userMemoriesIdentities.findFirst({ + where: eq(userMemoriesIdentities.id, result.identityId), + }); + expect(identity?.role).toBe('developer'); + expect(identity?.userMemoryId).toBe(result.userMemoryId); + }); + + it('should handle empty params with defaults', async () => { + const result = await memoryModel.addIdentityEntry({ + base: {}, + identity: {}, + }); + + expect(result.identityId).toBeDefined(); + expect(result.userMemoryId).toBeDefined(); + }); + + it('should normalize relationship and type values', async () => { + const result = await memoryModel.addIdentityEntry({ + base: {}, + identity: { + relationship: ' Self ', + type: ' Personal ', + }, + }); + + const identity = await serverDB.query.userMemoriesIdentities.findFirst({ + where: eq(userMemoriesIdentities.id, result.identityId), + }); + expect(identity?.relationship).toBe(RelationshipEnum.Self); + expect(identity?.type).toBe(IdentityTypeEnum.Personal); + }); + }); + + // ========== updateIdentityEntry ========== + describe('updateIdentityEntry', () => { + it('should update identity and base memory', async () => { + const { identityId, userMemoryId } = await memoryModel.addIdentityEntry({ + base: { title: 'Original' }, + identity: { role: 'original' }, + }); + + const success = await memoryModel.updateIdentityEntry({ + base: { title: 'Updated' }, + identity: { role: 'updated' }, + identityId, + }); + + expect(success).toBe(true); + + const mem = await serverDB.query.userMemories.findFirst({ + where: eq(userMemories.id, userMemoryId), + }); + expect(mem?.title).toBe('Updated'); + + const identity = await serverDB.query.userMemoriesIdentities.findFirst({ + where: eq(userMemoriesIdentities.id, identityId), + }); + expect(identity?.role).toBe('updated'); + }); + + it('should return false for non-existent identity', async () => { + const success = await memoryModel.updateIdentityEntry({ + identity: { role: 'test' }, + identityId: 'non-existent', + }); + + expect(success).toBe(false); + }); + + it('should support replace merge strategy', async () => { + const { identityId } = await memoryModel.addIdentityEntry({ + base: {}, + identity: { + description: 'original desc', + role: 'original role', + type: IdentityTypeEnum.Personal, + }, + }); + + const success = await memoryModel.updateIdentityEntry({ + identity: { description: 'replaced desc' }, + identityId, + mergeStrategy: 'replace' as any, + }); + + expect(success).toBe(true); + const updated = await serverDB.query.userMemoriesIdentities.findFirst({ + where: eq(userMemoriesIdentities.id, identityId), + }); + expect(updated?.description).toBe('replaced desc'); + // In replace mode, unspecified fields become null + expect(updated?.role).toBeNull(); + }); + + it('should not update other user identity', async () => { + const otherModel = new UserMemoryModel(serverDB, otherUserId); + const { identityId } = await otherModel.addIdentityEntry({ + base: {}, + identity: { role: 'other-role' }, + }); + + const success = await memoryModel.updateIdentityEntry({ + identity: { role: 'hacked' }, + identityId, + }); + + expect(success).toBe(false); + }); + }); + + // ========== removeIdentityEntry ========== + describe('removeIdentityEntry', () => { + it('should delete identity and associated base memory', async () => { + const { identityId, userMemoryId } = await memoryModel.addIdentityEntry({ + base: { title: 'To delete' }, + identity: { role: 'disposable' }, + }); + + const success = await memoryModel.removeIdentityEntry(identityId); + + expect(success).toBe(true); + + const identity = await serverDB.query.userMemoriesIdentities.findFirst({ + where: eq(userMemoriesIdentities.id, identityId), + }); + expect(identity).toBeUndefined(); + + const mem = await serverDB.query.userMemories.findFirst({ + where: eq(userMemories.id, userMemoryId), + }); + expect(mem).toBeUndefined(); + }); + + it('should return false for non-existent identity', async () => { + const success = await memoryModel.removeIdentityEntry('non-existent'); + expect(success).toBe(false); + }); + + it('should not delete other user identity', async () => { + const otherModel = new UserMemoryModel(serverDB, otherUserId); + const { identityId } = await otherModel.addIdentityEntry({ + base: {}, + identity: { role: 'other' }, + }); + + const success = await memoryModel.removeIdentityEntry(identityId); + expect(success).toBe(false); + }); + }); + + // ========== getAllIdentities ========== + describe('getAllIdentities', () => { + it('should return all identities for current user', async () => { + await createIdentityPair({ type: 'personal' }); + await createIdentityPair({ type: 'professional' }); + await createIdentityPair({ type: 'other', user: otherUserId }); + + const result = await memoryModel.getAllIdentities(); + + expect(result).toHaveLength(2); + expect(result.every((r) => r.userId === userId)).toBe(true); + }); + + it('should order by capturedAt descending', async () => { + await createIdentityPair({ type: 'first' }); + // Small delay to ensure different timestamps + await createIdentityPair({ type: 'second' }); + + const result = await memoryModel.getAllIdentities(); + + expect(result).toHaveLength(2); + // Most recent first + expect((result[0] as any).capturedAt.getTime()).toBeGreaterThanOrEqual( + (result[1] as any).capturedAt.getTime(), + ); + }); + + it('should return empty array when no identities', async () => { + const result = await memoryModel.getAllIdentities(); + expect(result).toEqual([]); + }); + }); + + // ========== getAllIdentitiesWithMemory ========== + describe('getAllIdentitiesWithMemory', () => { + it('should return identities joined with base memories', async () => { + await createIdentityPair({ baseTitle: 'My identity memory' }); + + const result = await memoryModel.getAllIdentitiesWithMemory(); + + expect(result).toHaveLength(1); + expect(result[0].identity).toBeDefined(); + expect(result[0].memory).toBeDefined(); + expect(result[0].memory.title).toBe('My identity memory'); + }); + + it('should not return other user data', async () => { + await createIdentityPair({}); + await createIdentityPair({ user: otherUserId }); + + const result = await memoryModel.getAllIdentitiesWithMemory(); + + expect(result).toHaveLength(1); + }); + }); + + // ========== getIdentitiesByType ========== + describe('getIdentitiesByType', () => { + it('should filter identities by type', async () => { + await createIdentityPair({ type: 'personal' }); + await createIdentityPair({ type: 'professional' }); + await createIdentityPair({ type: 'personal' }); + + const result = await memoryModel.getIdentitiesByType('personal'); + + expect(result).toHaveLength(2); + expect(result.every((r) => r.type === 'personal')).toBe(true); + }); + + it('should return empty array for non-matching type', async () => { + await createIdentityPair({ type: 'personal' }); + + const result = await memoryModel.getIdentitiesByType('unknown'); + + expect(result).toEqual([]); + }); + + it('should not return other user identities', async () => { + await createIdentityPair({ type: 'personal' }); + await createIdentityPair({ type: 'personal', user: otherUserId }); + + const result = await memoryModel.getIdentitiesByType('personal'); + + expect(result).toHaveLength(1); + }); + }); + + // ========== removeContextEntry ========== + describe('removeContextEntry', () => { + it('should delete context and associated memories', async () => { + const { context, memory } = await createContextPair({}); + + const success = await memoryModel.removeContextEntry(context.id); + + expect(success).toBe(true); + + const ctx = await serverDB.query.userMemoriesContexts.findFirst({ + where: eq(userMemoriesContexts.id, context.id), + }); + expect(ctx).toBeUndefined(); + + const mem = await serverDB.query.userMemories.findFirst({ + where: eq(userMemories.id, memory.id), + }); + expect(mem).toBeUndefined(); + }); + + it('should return false for non-existent context', async () => { + const success = await memoryModel.removeContextEntry('non-existent'); + expect(success).toBe(false); + }); + }); + + // ========== removeExperienceEntry ========== + describe('removeExperienceEntry', () => { + it('should delete experience and associated base memory', async () => { + const { experience, memory } = await createExperiencePair({}); + + const success = await memoryModel.removeExperienceEntry(experience.id); + + expect(success).toBe(true); + + const exp = await serverDB.query.userMemoriesExperiences.findFirst({ + where: eq(userMemoriesExperiences.id, experience.id), + }); + expect(exp).toBeUndefined(); + }); + + it('should return false for non-existent experience', async () => { + const success = await memoryModel.removeExperienceEntry('non-existent'); + expect(success).toBe(false); + }); + }); + + // ========== removePreferenceEntry ========== + describe('removePreferenceEntry', () => { + it('should delete preference and associated base memory', async () => { + const { preference } = await createPreferencePair({}); + + const success = await memoryModel.removePreferenceEntry(preference.id); + + expect(success).toBe(true); + + const pref = await serverDB.query.userMemoriesPreferences.findFirst({ + where: eq(userMemoriesPreferences.id, preference.id), + }); + expect(pref).toBeUndefined(); + }); + + it('should return false for non-existent preference', async () => { + const success = await memoryModel.removePreferenceEntry('non-existent'); + expect(success).toBe(false); + }); + }); + + // ========== Create Memory Methods (model-level) ========== + describe('createActivityMemory', () => { + it('should create activity memory via model', async () => { + const result = await memoryModel.createActivityMemory({ + details: 'Activity details', + memoryLayer: LayersEnum.Activity, + memoryType: 'activity' as any, + summary: 'Activity summary', + title: 'Activity test', + activity: { + narrative: 'Did a thing', + status: 'completed', + type: ActivityTypeEnum.Other, + tags: ['test-tag'], + startsAt: new Date('2025-01-01'), + endsAt: new Date('2025-01-02'), + } as any, + }); + + expect(result.memory).toBeDefined(); + expect(result.memory.memoryLayer).toBe(LayersEnum.Activity); + expect(result.activity).toBeDefined(); + expect(result.activity.narrative).toBe('Did a thing'); + expect(result.activity.userMemoryId).toBe(result.memory.id); + }); + }); + + describe('createExperienceMemory', () => { + it('should create experience memory via model', async () => { + const result = await memoryModel.createExperienceMemory({ + details: 'Experience details', + memoryLayer: LayersEnum.Experience, + memoryType: 'experience' as any, + summary: 'Experience summary', + title: 'Experience test', + experience: { + action: 'learned something', + keyLearning: 'important lesson', + situation: 'at work', + type: 'learning', + tags: ['learn'], + reasoning: 'because reasons', + possibleOutcome: 'better outcomes', + } as any, + }); + + expect(result.memory).toBeDefined(); + expect(result.memory.memoryLayer).toBe(LayersEnum.Experience); + expect(result.experience).toBeDefined(); + expect(result.experience.action).toBe('learned something'); + expect(result.experience.userMemoryId).toBe(result.memory.id); + }); + }); + + describe('createContextMemory', () => { + it('should create context memory via model', async () => { + const result = await memoryModel.createContextMemory({ + details: 'Context details', + memoryLayer: LayersEnum.Context, + memoryType: 'context' as any, + summary: 'Context summary', + title: 'Context test', + context: { + description: 'A test context', + title: 'Test Context', + type: 'project', + tags: ['ctx-tag'], + associatedObjects: [], + associatedSubjects: [], + } as any, + }); + + expect(result.memory).toBeDefined(); + expect(result.memory.memoryLayer).toBe(LayersEnum.Context); + expect(result.context).toBeDefined(); + expect(result.context.description).toBe('A test context'); + }); + }); + + // ========== parseAssociatedObjects with valid schema ========== + describe('parseAssociatedObjects - valid AssociatedObjectSchema', () => { + it('should parse valid associated objects with extra', () => { + const result = UserMemoryModel.parseAssociatedObjects([ + { + name: 'TestObj', + type: UserMemoryContextObjectType.Application, + extra: '{"key":"value"}', + }, + ]); + + expect(result).toHaveLength(1); + expect(result[0]).toMatchObject({ + name: 'TestObj', + type: UserMemoryContextObjectType.Application, + extra: { key: 'value' }, + }); + }); + + it('should parse valid objects with null extra', () => { + const result = UserMemoryModel.parseAssociatedObjects([ + { + name: 'TestObj', + type: UserMemoryContextObjectType.Person, + extra: null, + }, + ]); + + expect(result).toHaveLength(1); + expect(result[0]).toMatchObject({ + name: 'TestObj', + type: UserMemoryContextObjectType.Person, + }); + }); + }); + + describe('parseAssociatedSubjects - valid schema', () => { + it('should parse valid associated subjects with extra', () => { + // UserMemoryContextSubjectType has: Item, Other, Person, Pet + const result = UserMemoryModel.parseAssociatedSubjects([ + { + name: 'TestSubject', + type: 'person', // lowercase matches nativeEnum + extra: '{"role":"admin"}', + }, + ]); + + expect(result).toHaveLength(1); + expect(result[0]).toMatchObject({ + name: 'TestSubject', + extra: { role: 'admin' }, + }); + }); + }); + + // ========== getMemoryDetail edge cases ========== + describe('getMemoryDetail - edge cases', () => { + it('should return undefined for context with no userMemoryIds', async () => { + // Create a context with empty userMemoryIds + const [ctx] = await serverDB + .insert(userMemoriesContexts) + .values({ + description: 'No memory', + title: 'Empty context', + type: 'project', + userId, + userMemoryIds: [], + }) + .returning(); + + const result = await memoryModel.getMemoryDetail({ + id: ctx.id, + layer: LayersEnum.Context, + }); + + expect(result).toBeUndefined(); + }); + + it('should return undefined for activity with no userMemoryId', async () => { + // Create an activity without a linked memory + const [act] = await serverDB + .insert(userMemoriesActivities) + .values({ + narrative: 'orphan', + status: 'pending', + type: ActivityTypeEnum.Other, + userId, + }) + .returning(); + + const result = await memoryModel.getMemoryDetail({ + id: act.id, + layer: LayersEnum.Activity, + }); + + expect(result).toBeUndefined(); + }); + + it('should return undefined for experience with no userMemoryId', async () => { + const [exp] = await serverDB + .insert(userMemoriesExperiences) + .values({ + action: 'orphan', + situation: 'test', + type: 'learning', + userId, + }) + .returning(); + + const result = await memoryModel.getMemoryDetail({ + id: exp.id, + layer: LayersEnum.Experience, + }); + + expect(result).toBeUndefined(); + }); + }); + + // ========== queryMemories with tags filter ========== + describe('queryMemories - tags filter', () => { + it('should filter activities by tags', async () => { + await createActivityPair({ tags: ['urgent'] }); + await createActivityPair({ tags: ['low-priority'] }); + + const result = await memoryModel.queryMemories({ + layer: LayersEnum.Activity, + tags: ['urgent'], + }); + + expect(result.items.length).toBeGreaterThanOrEqual(1); + }); + + it('should filter contexts by tags', async () => { + await createContextPair({ tags: ['work'] }); + await createContextPair({ tags: ['personal'] }); + + const result = await memoryModel.queryMemories({ + layer: LayersEnum.Context, + tags: ['work'], + }); + + expect(result.items.length).toBeGreaterThanOrEqual(1); + }); + }); + + // ========== updateIdentityEntry with capturedAt ========== + describe('updateIdentityEntry - capturedAt', () => { + it('should update capturedAt on identity', async () => { + const { identity, memory } = await createIdentityPair({}); + const capturedDate = new Date('2025-06-15T12:00:00Z'); + + const result = await memoryModel.updateIdentityEntry({ + identityId: identity.id, + identity: { + capturedAt: capturedDate, + }, + }); + + expect(result).toBe(true); + + const updated = await serverDB.query.userMemoriesIdentities.findFirst({ + where: eq(userMemoriesIdentities.id, identity.id), + }); + expect(updated?.capturedAt).toEqual(capturedDate); + }); + }); +}); diff --git a/packages/database/src/models/userMemory/__tests__/persona.test.ts b/packages/database/src/models/userMemory/__tests__/persona.test.ts index 997cc4696a..87eb062fd9 100644 --- a/packages/database/src/models/userMemory/__tests__/persona.test.ts +++ b/packages/database/src/models/userMemory/__tests__/persona.test.ts @@ -88,6 +88,33 @@ describe('UserPersonaModel', () => { expect(latest?.version).toBe(2); }); + describe('appendDiff', () => { + it('should insert a diff record directly', async () => { + // First create a persona document to reference + const { document } = await personaModel.upsertPersona({ persona: '# v1' }); + + const diff = await personaModel.appendDiff({ + diffPersona: '- manual change', + memoryIds: ['mem-manual'], + nextVersion: 2, + personaId: document.id, + previousVersion: 1, + reasoning: 'Manual diff', + snapshot: '# v2', + sourceIds: ['src-manual'], + }); + + expect(diff).toBeDefined(); + expect(diff.personaId).toBe(document.id); + expect(diff.userId).toBe(userId); + expect(diff.diffPersona).toBe('- manual change'); + expect(diff.previousVersion).toBe(1); + expect(diff.nextVersion).toBe(2); + expect(diff.memoryIds).toEqual(['mem-manual']); + expect(diff.sourceIds).toEqual(['src-manual']); + }); + }); + it('lists diffs ordered by createdAt desc', async () => { await personaModel.upsertPersona({ diffPersona: '- change', diff --git a/packages/database/src/repositories/agentMigration/__tests__/agentMigrationRepo.test.ts b/packages/database/src/repositories/agentMigration/__tests__/agentMigrationRepo.test.ts index 9e1d75a609..bf8e63b44b 100644 --- a/packages/database/src/repositories/agentMigration/__tests__/agentMigrationRepo.test.ts +++ b/packages/database/src/repositories/agentMigration/__tests__/agentMigrationRepo.test.ts @@ -2,7 +2,15 @@ import { eq, inArray } from 'drizzle-orm'; import { afterEach, beforeEach, describe, expect, it } from 'vitest'; import { getTestDB } from '../../../core/getTestDB'; -import { agents, messages, sessions, topics, users } from '../../../schemas'; +import { + agents, + agentsToSessions, + messages, + sessionGroups, + sessions, + topics, + users, +} from '../../../schemas'; import type { LobeChatDatabase } from '../../../type'; import { AgentMigrationRepo } from '../index'; @@ -759,4 +767,133 @@ describe('AgentMigrationRepo', () => { expect(result).toBeNull(); }); }); + + describe('migrateSessionGroupId', () => { + it('should backfill sessionGroupId from session groupId', async () => { + // Create session group + const [group] = await serverDB + .insert(sessionGroups) + .values({ id: 'test-group', name: 'Test Group', userId }) + .returning(); + + // Create a session with groupId + const groupedSessionId = 'grouped-session'; + await serverDB.insert(sessions).values({ id: groupedSessionId, userId, groupId: group.id }); + + // Create an agent with null sessionGroupId + const agentId = 'migrate-group-agent'; + await serverDB.insert(agents).values({ id: agentId, userId, sessionGroupId: null }); + + // Link agent to session + await serverDB + .insert(agentsToSessions) + .values({ agentId, sessionId: groupedSessionId, userId }); + + // Run migration + await agentMigrationRepo.migrateSessionGroupId(); + + // Verify agent's sessionGroupId was updated + const updatedAgent = await serverDB.query.agents.findFirst({ + where: eq(agents.id, agentId), + }); + expect(updatedAgent?.sessionGroupId).toBe(group.id); + }); + + it('should not modify agents that already have sessionGroupId', async () => { + const [group1] = await serverDB + .insert(sessionGroups) + .values({ id: 'group-1', name: 'Group 1', userId }) + .returning(); + + const [group2] = await serverDB + .insert(sessionGroups) + .values({ id: 'group-2', name: 'Group 2', userId }) + .returning(); + + const groupedSessionId = 'session-with-group2'; + await serverDB.insert(sessions).values({ id: groupedSessionId, userId, groupId: group2.id }); + + // Agent already has a sessionGroupId set + const agentId = 'already-grouped-agent'; + await serverDB.insert(agents).values({ id: agentId, userId, sessionGroupId: group1.id }); + + await serverDB + .insert(agentsToSessions) + .values({ agentId, sessionId: groupedSessionId, userId }); + + await agentMigrationRepo.migrateSessionGroupId(); + + // Should still have the original group + const agent = await serverDB.query.agents.findFirst({ + where: eq(agents.id, agentId), + }); + expect(agent?.sessionGroupId).toBe(group1.id); + }); + + it('should do nothing when no agents need migration', async () => { + // No agents at all - should not throw + await agentMigrationRepo.migrateSessionGroupId(); + }); + + it('should not modify other user agents', async () => { + const [group] = await serverDB + .insert(sessionGroups) + .values({ id: 'user2-group', name: 'User2 Group', userId: userId2 }) + .returning(); + + const user2Session = 'user2-session-grouped'; + await serverDB + .insert(sessions) + .values({ id: user2Session, userId: userId2, groupId: group.id }); + + const agentId = 'user2-agent'; + await serverDB.insert(agents).values({ id: agentId, userId: userId2, sessionGroupId: null }); + + await serverDB + .insert(agentsToSessions) + .values({ agentId, sessionId: user2Session, userId: userId2 }); + + // Run migration for user1 + await agentMigrationRepo.migrateSessionGroupId(); + + // User2's agent should not be affected + const agent = await serverDB.query.agents.findFirst({ + where: eq(agents.id, agentId), + }); + expect(agent?.sessionGroupId).toBeNull(); + }); + + it('should handle multiple agents needing migration', async () => { + const [group] = await serverDB + .insert(sessionGroups) + .values({ id: 'multi-group', name: 'Multi Group', userId }) + .returning(); + + const session1 = 'multi-session-1'; + const session2 = 'multi-session-2'; + await serverDB.insert(sessions).values([ + { id: session1, userId, groupId: group.id }, + { id: session2, userId, groupId: group.id }, + ]); + + const agent1 = 'multi-agent-1'; + const agent2 = 'multi-agent-2'; + await serverDB.insert(agents).values([ + { id: agent1, userId, sessionGroupId: null }, + { id: agent2, userId, sessionGroupId: null }, + ]); + + await serverDB.insert(agentsToSessions).values([ + { agentId: agent1, sessionId: session1, userId }, + { agentId: agent2, sessionId: session2, userId }, + ]); + + await agentMigrationRepo.migrateSessionGroupId(); + + const updatedAgents = await serverDB.query.agents.findMany({ + where: inArray(agents.id, [agent1, agent2]), + }); + expect(updatedAgents.every((a) => a.sessionGroupId === group.id)).toBe(true); + }); + }); }); diff --git a/packages/database/src/repositories/compression/index.test.ts b/packages/database/src/repositories/compression/index.test.ts index 518b20b273..739fec1e24 100644 --- a/packages/database/src/repositories/compression/index.test.ts +++ b/packages/database/src/repositories/compression/index.test.ts @@ -176,6 +176,88 @@ describe('CompressionRepository', () => { }); }); + describe('updateMetadata', () => { + it('should update metadata jsonb column on a compression group', async () => { + const groupId = await compressionRepo.createCompressionGroup({ + content: 'Summary', + messageIds: [], + metadata: {}, + topicId, + }); + + await compressionRepo.updateMetadata(groupId, { expanded: true } as any); + + // Query metadata column directly since getCompressionGroups returns description-based metadata + const { eq, and } = await import('drizzle-orm'); + const [group] = await serverDB + .select({ metadata: messageGroups.metadata }) + .from(messageGroups) + .where(and(eq(messageGroups.id, groupId), eq(messageGroups.userId, userId))); + expect((group.metadata as any)?.expanded).toBe(true); + }); + + it('should merge with existing metadata', async () => { + const groupId = await compressionRepo.createCompressionGroup({ + content: 'Summary', + messageIds: [], + metadata: {}, + topicId, + }); + + await compressionRepo.updateMetadata(groupId, { expanded: true } as any); + await compressionRepo.updateMetadata(groupId, { collapsed: false } as any); + + const { eq, and } = await import('drizzle-orm'); + const [group] = await serverDB + .select({ metadata: messageGroups.metadata }) + .from(messageGroups) + .where(and(eq(messageGroups.id, groupId), eq(messageGroups.userId, userId))); + const meta = group.metadata as any; + expect(meta?.expanded).toBe(true); + expect(meta?.collapsed).toBe(false); + }); + + it('should handle updating metadata on non-existent group gracefully', async () => { + // Should not throw + await compressionRepo.updateMetadata('non-existent-group-id', { expanded: true } as any); + }); + }); + + describe('updateCompressionContent with metadata merge', () => { + it('should merge description metadata when updating content', async () => { + const groupId = await compressionRepo.createCompressionGroup({ + content: 'Original summary', + messageIds: [], + metadata: { originalTokenCount: 1000 }, + topicId, + }); + + await compressionRepo.updateCompressionContent(groupId, 'Updated summary', { + compressedTokenCount: 100, + }); + + const groups = await compressionRepo.getCompressionGroups(topicId); + expect(groups[0].content).toBe('Updated summary'); + // Verify the description-based metadata was merged + expect((groups[0].metadata as any)?.originalTokenCount).toBe(1000); + expect((groups[0].metadata as any)?.compressedTokenCount).toBe(100); + }); + + it('should update content without metadata', async () => { + const groupId = await compressionRepo.createCompressionGroup({ + content: 'Original', + messageIds: [], + metadata: {}, + topicId, + }); + + await compressionRepo.updateCompressionContent(groupId, 'New content'); + + const groups = await compressionRepo.getCompressionGroups(topicId); + expect(groups[0].content).toBe('New content'); + }); + }); + describe('toggleMessagePin', () => { it('should pin a message', async () => { await serverDB.insert(messages).values({ diff --git a/packages/database/src/repositories/dataExporter/index.test.ts b/packages/database/src/repositories/dataExporter/index.test.ts index 5a19b5452c..f6dfb991a3 100644 --- a/packages/database/src/repositories/dataExporter/index.test.ts +++ b/packages/database/src/repositories/dataExporter/index.test.ts @@ -276,23 +276,56 @@ describe('DataExporterRepos', () => { expect(result.sessions).toHaveLength(1); }); - it.skip('should skip relation tables when source tables have no data', async () => { - // 删除文件数据，这将导致 globalFiles 表被跳过 - await db.delete(files); + it('should skip relation tables when source tables have no data', async () => { + // Delete agents and sessions, so agentsToSessions source tables have no data + await db.delete(agentsToSessions); + await db.delete(agents); + await db.delete(messages); + await db.delete(topics); + await db.delete(sessions); - // 创建导出器实例 const dataExporter = new DataExporterRepos(db, userId); - - // 执行导出 const result = await dataExporter.export(); - // 验证文件表为空 - // expect(result).toHaveProperty('files'); - // expect(result.files).toEqual([]); + // agentsToSessions should be empty because both source tables have no data + expect(result).toHaveProperty('agentsToSessions'); + expect(result.agentsToSessions).toEqual([]); + }); - // 验证关联表也为空 - // expect(result).toHaveProperty('globalFiles'); - // expect(result.globalFiles).toEqual([]); + it('should handle base table query error gracefully', async () => { + // Mock a specific base table to throw an error + // @ts-ignore + vi.spyOn(db.query.userSettings, 'findMany').mockRejectedValueOnce( + new Error('DB connection failed'), + ); + + const dataExporter = new DataExporterRepos(db, userId); + const result = await dataExporter.export(); + + // userSettings should return empty array due to error handling + expect(result).toHaveProperty('userSettings'); + expect(result.userSettings).toEqual([]); + + // Other tables should still export successfully + expect(result.sessions).toHaveLength(1); + }); + + it('should handle relation table query error gracefully', async () => { + // Mock agentsToSessions query to throw an error + // @ts-ignore + vi.spyOn(db.query.agentsToSessions, 'findMany').mockRejectedValueOnce( + new Error('Relation query failed'), + ); + + const dataExporter = new DataExporterRepos(db, userId); + const result = await dataExporter.export(); + + // agentsToSessions should return empty array due to error handling + expect(result).toHaveProperty('agentsToSessions'); + expect(result.agentsToSessions).toEqual([]); + + // Base tables should still export successfully + expect(result.sessions).toHaveLength(1); }); it('should export data for a different user', async () => { diff --git a/packages/database/src/repositories/dataImporter/__tests__/index.test.ts b/packages/database/src/repositories/dataImporter/__tests__/index.test.ts index ae21dc9507..c61a00dca7 100644 --- a/packages/database/src/repositories/dataImporter/__tests__/index.test.ts +++ b/packages/database/src/repositories/dataImporter/__tests__/index.test.ts @@ -137,6 +137,255 @@ describe('DataImporter', () => { }); }); + describe('import with empty tables', () => { + it('should skip tables with empty data', async () => { + const data: ImportPgDataStructure = { + data: { + agents: [], + sessions: [], + sessionGroups: [], + }, + mode: 'pglite', + schemaHash: 'test', + } as any; + + const result = await importer.importPgData(data); + + expect(result.success).toBe(true); + // No results should be returned for empty tables + expect(Object.keys(result.results)).toHaveLength(0); + }); + }); + + describe('import with sessionGroups (relation mapping)', () => { + it('should import sessions with sessionGroup relations', async () => { + const data: ImportPgDataStructure = { + data: { + agents: [ + { + id: 'agt_rel_test', + slug: 'rel-test-agent', + model: 'gpt-4', + provider: 'openai', + systemRole: '', + createdAt: '2025-01-01T00:00:00Z', + updatedAt: '2025-01-01T00:00:00Z', + }, + ], + agentsToSessions: [{ agentId: 'agt_rel_test', sessionId: 'ssn_rel_test' }], + sessionGroups: [ + { + id: 'sg_test1', + name: 'Test Group', + sort: 0, + createdAt: '2025-01-01T00:00:00Z', + updatedAt: '2025-01-01T00:00:00Z', + }, + ], + sessions: [ + { + id: 'ssn_rel_test', + slug: 'rel-test-session', + type: 'agent', + groupId: 'sg_test1', + createdAt: '2025-01-01T00:00:00Z', + updatedAt: '2025-01-01T00:00:00Z', + }, + ], + }, + mode: 'pglite', + schemaHash: 'test', + } as any; + + const result = await importer.importPgData(data); + + expect(result.success).toBe(true); + expect(result.results.sessionGroups).toMatchObject({ added: 1, errors: 0 }); + expect(result.results.sessions).toMatchObject({ added: 1, errors: 0 }); + + // Verify the session's groupId was mapped to the new sessionGroup ID + const sessions = await clientDB.query.sessions.findMany({ + where: eq(Schema.sessions.userId, userId), + }); + expect(sessions).toHaveLength(1); + expect(sessions[0].groupId).not.toBeNull(); + }); + + it('should set relation to null when group mapping not found', async () => { + // Import a session with a groupId that has no corresponding sessionGroup in the import data + // The code handles this by: if idMaps[sessionGroups] exists but mappedId is undefined → set to null + // However, if sessionGroups was never imported, idMaps[sessionGroups] won't exist and groupId stays as-is + // Let's import sessionGroups first (empty) to ensure the map exists, then sessions with unmapped groupId + const data: ImportPgDataStructure = { + data: { + agents: [ + { + id: 'agt_nomap', + slug: 'nomap-agent', + model: 'gpt-4', + provider: 'openai', + systemRole: '', + createdAt: '2025-01-01T00:00:00Z', + updatedAt: '2025-01-01T00:00:00Z', + }, + ], + agentsToSessions: [{ agentId: 'agt_nomap', sessionId: 'ssn_nomap' }], + sessionGroups: [ + { + id: 'sg_exists', + name: 'Exists Group', + sort: 0, + createdAt: '2025-01-01T00:00:00Z', + updatedAt: '2025-01-01T00:00:00Z', + }, + ], + sessions: [ + { + id: 'ssn_nomap', + slug: 'nomap-session', + type: 'agent', + groupId: 'non_existent_group', + createdAt: '2025-01-01T00:00:00Z', + updatedAt: '2025-01-01T00:00:00Z', + }, + ], + }, + mode: 'pglite', + schemaHash: 'test', + } as any; + + const result = await importer.importPgData(data); + + expect(result.success).toBe(true); + expect(result.results.sessions).toMatchObject({ added: 1, errors: 0 }); + + // Session should be imported but groupId should be null (unmapped) + const sessions = await clientDB.query.sessions.findMany({ + where: eq(Schema.sessions.userId, userId), + }); + expect(sessions).toHaveLength(1); + expect(sessions[0].groupId).toBeNull(); + }); + }); + + describe('import with self-references', () => { + it('should nullify self-reference fields (parentId, quotaId)', async () => { + const data: ImportPgDataStructure = { + data: { + agents: [ + { + id: 'agt_selfref', + slug: 'selfref-agent', + model: 'gpt-4', + provider: 'openai', + systemRole: '', + createdAt: '2025-01-01T00:00:00Z', + updatedAt: '2025-01-01T00:00:00Z', + }, + ], + agentsToSessions: [{ agentId: 'agt_selfref', sessionId: 'ssn_selfref' }], + sessions: [ + { + id: 'ssn_selfref', + slug: 'selfref-session', + type: 'agent', + createdAt: '2025-01-01T00:00:00Z', + updatedAt: '2025-01-01T00:00:00Z', + }, + ], + messages: [ + { + id: 'msg_selfref_1', + role: 'user', + content: 'Hello', + sessionId: 'ssn_selfref', + parentId: 'msg_selfref_parent', + quotaId: 'msg_selfref_quota', + createdAt: '2025-01-01T00:00:00Z', + updatedAt: '2025-01-01T00:00:00Z', + }, + ], + }, + mode: 'pglite', + schemaHash: 'test', + } as any; + + const result = await importer.importPgData(data); + + expect(result.success).toBe(true); + expect(result.results.messages).toMatchObject({ added: 1, errors: 0 }); + + // Verify self-reference fields are set to null + const messages = await clientDB.query.messages.findMany({ + where: eq(Schema.messages.userId, userId), + }); + expect(messages).toHaveLength(1); + expect(messages[0].parentId).toBeNull(); + expect(messages[0].quotaId).toBeNull(); + }); + }); + + describe('import with override conflict strategy', () => { + it('should apply field processor when overriding duplicate slug', async () => { + // First, create an agent with a specific slug + const firstData: ImportPgDataStructure = { + data: { + agents: [ + { + id: 'agt_override1', + slug: 'override-test-slug', + model: 'gpt-4', + provider: 'openai', + systemRole: '', + createdAt: '2025-01-01T00:00:00Z', + updatedAt: '2025-01-01T00:00:00Z', + }, + ], + agentsToSessions: [], + sessions: [], + }, + mode: 'pglite', + schemaHash: 'test', + } as any; + + await importer.importPgData(firstData); + + // Now create a new importer and import a DIFFERENT agent with the same slug + const importer2 = new DataImporterRepos(clientDB, userId); + const secondData: ImportPgDataStructure = { + data: { + agents: [ + { + id: 'agt_override2', + slug: 'override-test-slug', + model: 'gpt-4', + provider: 'openai', + systemRole: '', + createdAt: '2025-01-02T00:00:00Z', + updatedAt: '2025-01-02T00:00:00Z', + }, + ], + agentsToSessions: [], + sessions: [], + }, + mode: 'pglite', + schemaHash: 'test', + } as any; + + // Default conflictStrategy for agents is 'override' (no conflictStrategy in config = default 'override') + const result = await importer2.importPgData(secondData); + + expect(result.success).toBe(true); + // The override strategy should apply the field processor (appends UUID suffix to slug) + expect(result.results.agents).toMatchObject({ added: 1, errors: 0 }); + + const agents = await clientDB.query.agents.findMany({ + where: eq(Schema.agents.userId, userId), + }); + expect(agents).toHaveLength(2); + }); + }); + describe('import message and topic', () => { it('should import return correct result', async () => { const exportData = topicsData as ImportPgDataStructure; diff --git a/packages/database/src/repositories/knowledge/__tests__/index.test.ts b/packages/database/src/repositories/knowledge/__tests__/index.test.ts index 0894981159..e1442fc40e 100644 --- a/packages/database/src/repositories/knowledge/__tests__/index.test.ts +++ b/packages/database/src/repositories/knowledge/__tests__/index.test.ts @@ -597,6 +597,218 @@ describe('KnowledgeRepo', () => { }); }); + describe('query with knowledgeBaseId + filters', () => { + beforeEach(async () => { + await serverDB + .insert(knowledgeBases) + .values([{ id: 'kb-filter', name: 'Filter KB', userId }]); + + // Create a folder doc in KB + await serverDB.insert(documents).values([ + { + id: 'kb-folder-doc', + userId, + title: 'KB Folder', + fileType: 'custom/folder', + sourceType: 'topic', + source: 'internal://folder/kb-folder-doc', + slug: 'kb-folder', + knowledgeBaseId: 'kb-filter', + totalCharCount: 0, + totalLineCount: 0, + }, + { + id: 'kb-standalone-doc', + userId, + title: 'KB Standalone Note', + fileType: 'custom/note', + sourceType: 'topic', + source: 'internal://note/kb-standalone-doc', + knowledgeBaseId: 'kb-filter', + totalCharCount: 200, + totalLineCount: 4, + }, + { + id: 'kb-standalone-doc-searchable', + userId, + title: 'Searchable KB Note', + fileType: 'custom/note', + sourceType: 'topic', + source: 'internal://note/kb-standalone-doc-searchable', + knowledgeBaseId: 'kb-filter', + totalCharCount: 100, + totalLineCount: 2, + }, + { + id: 'kb-app-doc', + userId, + title: 'KB App Doc', + fileType: 'application/pdf', + sourceType: 'topic', + source: 'internal://doc/kb-app-doc', + knowledgeBaseId: 'kb-filter', + totalCharCount: 300, + totalLineCount: 6, + }, + ]); + + // Create files in KB + await serverDB.insert(files).values([ + { + id: 'kb-f-image', + userId, + name: 'kb-image.png', + fileType: 'image/png', + size: 1000, + url: 'https://example.com/kb-image.png', + }, + { + id: 'kb-f-pdf', + userId, + name: 'kb-doc.pdf', + fileType: 'application/pdf', + size: 2000, + url: 'https://example.com/kb-doc.pdf', + parentId: 'kb-folder-doc', + }, + { + id: 'kb-f-searchable', + userId, + name: 'searchable-file.txt', + fileType: 'text/plain', + size: 100, + url: 'https://example.com/searchable.txt', + }, + ]); + + await serverDB.insert(knowledgeBaseFiles).values([ + { fileId: 'kb-f-image', knowledgeBaseId: 'kb-filter', userId }, + { fileId: 'kb-f-pdf', knowledgeBaseId: 'kb-filter', userId }, + { fileId: 'kb-f-searchable', knowledgeBaseId: 'kb-filter', userId }, + ]); + }); + + it('should filter KB files by parentId', async () => { + const result = await knowledgeRepo.query({ + knowledgeBaseId: 'kb-filter', + parentId: 'kb-folder-doc', + }); + + expect(result.some((item) => item.id === 'kb-f-pdf')).toBe(true); + expect(result.every((item) => item.id !== 'kb-f-image')).toBe(true); + }); + + it('should filter KB files by null parentId', async () => { + const result = await knowledgeRepo.query({ + knowledgeBaseId: 'kb-filter', + parentId: null, + }); + + // Files without parentId should be returned + expect(result.some((item) => item.id === 'kb-f-image')).toBe(true); + }); + + it('should filter KB files by search query', async () => { + const result = await knowledgeRepo.query({ + knowledgeBaseId: 'kb-filter', + q: 'searchable', + }); + + expect(result.some((item) => item.id === 'kb-f-searchable')).toBe(true); + }); + + it('should filter KB files by category (Images)', async () => { + const result = await knowledgeRepo.query({ + knowledgeBaseId: 'kb-filter', + category: FilesTabs.Images, + }); + + // Images category returns only image files, document query returns empty set + expect(result.some((item) => item.id === 'kb-f-image')).toBe(true); + expect(result.every((item) => item.fileType.startsWith('image'))).toBe(true); + }); + + it('should filter KB files by category (Documents) and exclude custom/document', async () => { + const result = await knowledgeRepo.query({ + knowledgeBaseId: 'kb-filter', + category: FilesTabs.Documents, + }); + + // Should include application/* files and custom/* docs + expect( + result.every( + (item) => + item.fileType.startsWith('application') || + (item.fileType.startsWith('custom') && item.fileType !== 'custom/document'), + ), + ).toBe(true); + }); + + it('should return KB standalone documents (no fileId) with search', async () => { + const result = await knowledgeRepo.query({ + knowledgeBaseId: 'kb-filter', + q: 'Searchable KB', + }); + + expect(result.some((item) => item.id === 'kb-standalone-doc-searchable')).toBe(true); + }); + + it('should handle KB with parentId for documents', async () => { + // Add a child document under the KB folder + await serverDB.insert(documents).values([ + { + id: 'kb-child-doc', + userId, + title: 'KB Child Doc', + fileType: 'custom/note', + sourceType: 'topic', + source: 'internal://note/kb-child-doc', + knowledgeBaseId: 'kb-filter', + parentId: 'kb-folder-doc', + totalCharCount: 50, + totalLineCount: 1, + }, + ]); + + const result = await knowledgeRepo.query({ + knowledgeBaseId: 'kb-filter', + parentId: 'kb-folder-doc', + }); + + expect(result.some((item) => item.id === 'kb-child-doc')).toBe(true); + }); + + it('should handle KB with null parentId for documents', async () => { + const result = await knowledgeRepo.query({ + knowledgeBaseId: 'kb-filter', + parentId: null, + }); + + // Standalone docs without parentId should be returned + expect(result.some((item) => item.id === 'kb-standalone-doc')).toBe(true); + }); + }); + + describe('query with non-matching categories (empty document results)', () => { + it('should return empty document set for Images category', async () => { + // Images only match files, documents should return empty + const result = await knowledgeRepo.query({ category: FilesTabs.Images }); + + // All results should be files with image/* type + result.forEach((item) => { + expect(item.fileType.startsWith('image')).toBe(true); + }); + }); + + it('should return empty document set for Videos category', async () => { + const result = await knowledgeRepo.query({ category: FilesTabs.Videos }); + + result.forEach((item) => { + expect(item.fileType.startsWith('video')).toBe(true); + }); + }); + }); + describe('query with website category', () => { beforeEach(async () => { await serverDB.insert(files).values([ @@ -668,4 +880,86 @@ describe('KnowledgeRepo', () => { expect(result.length).toBeGreaterThanOrEqual(3); }); }); + + describe('query with unknown category (default getFileTypePrefix)', () => { + beforeEach(async () => { + await serverDB.insert(files).values([ + { + id: 'default-cat-file', + userId, + name: 'file.txt', + fileType: 'text/plain', + size: 100, + url: 'https://example.com/default.txt', + }, + ]); + }); + + it('should handle Home category (default prefix returns empty string)', async () => { + // FilesTabs.Home triggers getFileTypePrefix default case, returning '' + // This means all file types match since ILIKE '%' matches everything + const result = await knowledgeRepo.query({ category: FilesTabs.Home as any }); + + expect(result.length).toBeGreaterThanOrEqual(0); + }); + }); + + describe('query sort edge cases', () => { + beforeEach(async () => { + await serverDB.insert(files).values([ + { + id: 'sort-file-1', + userId, + name: 'a-file.txt', + fileType: 'text/plain', + size: 300, + url: 'https://example.com/a.txt', + createdAt: new Date('2024-01-01T10:00:00Z'), + updatedAt: new Date('2024-01-05T10:00:00Z'), + }, + { + id: 'sort-file-2', + userId, + name: 'z-file.txt', + fileType: 'text/plain', + size: 100, + url: 'https://example.com/z.txt', + createdAt: new Date('2024-01-03T10:00:00Z'), + updatedAt: new Date('2024-01-03T10:00:00Z'), + }, + ]); + }); + + it('should sort by updatedAt asc', async () => { + const result = await knowledgeRepo.query({ + sorter: 'updatedAt', + sortType: SortType.Asc, + }); + + if (result.length >= 2) { + expect(result[0].updatedAt.getTime()).toBeLessThanOrEqual(result[1].updatedAt.getTime()); + } + }); + + it('should sort by createdAt desc', async () => { + const result = await knowledgeRepo.query({ + sorter: 'createdAt', + sortType: SortType.Desc, + }); + + if (result.length >= 2) { + expect(result[0].createdAt.getTime()).toBeGreaterThanOrEqual(result[1].createdAt.getTime()); + } + }); + + it('should fallback to default sort when invalid sorter is given', async () => { + const result = await knowledgeRepo.query({ + sorter: 'invalidField', + sortType: SortType.Asc, + }); + + // Should still return results (falls back to created_at DESC) + expect(result.length).toBeGreaterThan(0); + }); + }); }); diff --git a/packages/database/vitest.config.mts b/packages/database/vitest.config.mts index 97a04f905a..616cd1d740 100644 --- a/packages/database/vitest.config.mts +++ b/packages/database/vitest.config.mts @@ -17,6 +17,28 @@ export default defineConfig({ '@': resolve(__dirname, '../../src'), }, + coverage: { + exclude: [ + 'src/server/**', + 'src/repositories/dataImporter/deprecated/**', + 'src/types/**', + 'src/models/userMemory/sources/index.ts', + 'src/models/userMemory/sources/shared.ts', + 'src/models/ragEval/index.ts', + 'src/models/agentEval/index.ts', + 'src/repositories/userMemory/index.ts', + 'src/models/_template.ts', + 'src/models/__tests__/_test_template.ts', + 'src/models/web-server.ts', + 'src/core/web-server.ts', + 'src/core/db-adaptor.ts', + 'src/core/getTestDB.ts', + 'src/index.ts', + 'tests/**', + 'vitest.config*.mts', + ], + reporter: ['text', 'json'], + }, environment: 'happy-dom', exclude: [ 'node_modules/**/**', diff --git a/packages/prompts/src/prompts/userMemory/__snapshots__/index.test.ts.snap b/packages/prompts/src/prompts/userMemory/__snapshots__/index.test.ts.snap index ed502880a6..d707514e52 100644 --- a/packages/prompts/src/prompts/userMemory/__snapshots__/index.test.ts.snap +++ b/packages/prompts/src/prompts/userMemory/__snapshots__/index.test.ts.snap @@ -127,6 +127,24 @@ exports[`promptUserMemory > identities only > should format identities grouped b " `; +exports[`promptUserMemory > identities only > should format identity with capturedAt Date object 1`] = ` +" +The following are memories about this user retrieved from previous conversations. Use this information to personalize your responses and maintain continuity. + + User lives in Tokyo + +" +`; + +exports[`promptUserMemory > identities only > should format identity with capturedAt string 1`] = ` +" +The following are memories about this user retrieved from previous conversations. Use this information to personalize your responses and maintain continuity. + + User is a software engineer + +" +`; + exports[`promptUserMemory > identities only > should format single type identities (demographic only) 1`] = ` " The following are memories about this user retrieved from previous conversations. Use this information to personalize your responses and maintain continuity. @@ -174,6 +192,15 @@ exports[`promptUserMemory > identities only > should handle identity without rol " `; +exports[`promptUserMemory > identities only > should skip capturedAt when null 1`] = ` +" +The following are memories about this user retrieved from previous conversations. Use this information to personalize your responses and maintain continuity. + + User is a runner + +" +`; + exports[`promptUserMemory > mixed memory types > should format all memory types together 1`] = ` " The following are memories about this user retrieved from previous conversations. Use this information to personalize your responses and maintain continuity. @@ -260,6 +287,42 @@ exports[`promptUserMemory > mixed memory types > should format experiences and p " `; +exports[`promptUserMemory > mixed memory types > should format persona combined with other memory types 1`] = ` +" +The following are memories about this user retrieved from previous conversations. Use this information to personalize your responses and maintain continuity. + +A tech lead who focuses on AI-powered developer tools and has a strong preference for TypeScript. + + + User is a senior engineer + + + Working on AI products + + + Prefer concise responses with code examples + +" +`; + +exports[`promptUserMemory > persona only > should format persona with narrative and tagline 1`] = ` +" +The following are memories about this user retrieved from previous conversations. Use this information to personalize your responses and maintain continuity. + +A senior software engineer based in Shanghai who specializes in frontend development with React and TypeScript. Passionate about open source and building developer tools. + +" +`; + +exports[`promptUserMemory > persona only > should format persona with only narrative (no tagline) 1`] = ` +" +The following are memories about this user retrieved from previous conversations. Use this information to personalize your responses and maintain continuity. + +A product designer who transitioned from engineering. Enjoys bridging the gap between design and development. + +" +`; + exports[`promptUserMemory > preferences only > should filter out empty preferences and keep valid ones 1`] = ` " The following are memories about this user retrieved from previous conversations. Use this information to personalize your responses and maintain continuity. diff --git a/packages/prompts/src/prompts/userMemory/index.test.ts b/packages/prompts/src/prompts/userMemory/index.test.ts index a2eace10bb..658351c214 100644 --- a/packages/prompts/src/prompts/userMemory/index.test.ts +++ b/packages/prompts/src/prompts/userMemory/index.test.ts @@ -348,6 +348,106 @@ describe('promptUserMemory', () => { }); expect(result).toMatchSnapshot(); }); + + it('should format identity with capturedAt string', () => { + const result = promptUserMemory({ + memories: { + identities: [ + { + capturedAt: '2025-02-23T10:30:00.000Z', + description: 'User is a software engineer', + id: 'id-1', + role: 'Engineer', + type: 'professional', + }, + ], + }, + }); + expect(result).toMatchSnapshot(); + }); + + it('should format identity with capturedAt Date object', () => { + const result = promptUserMemory({ + memories: { + identities: [ + { + capturedAt: new Date('2025-06-15T00:00:00.000Z'), + description: 'User lives in Tokyo', + id: 'id-2', + type: 'demographic', + }, + ], + }, + }); + expect(result).toMatchSnapshot(); + }); + + it('should skip capturedAt when null', () => { + const result = promptUserMemory({ + memories: { + identities: [ + { + capturedAt: null, + description: 'User is a runner', + id: 'id-3', + type: 'personal', + }, + ], + }, + }); + expect(result).toMatchSnapshot(); + }); + }); + + describe('persona only', () => { + it('should format persona with narrative and tagline', () => { + const result = promptUserMemory({ + memories: { + persona: { + narrative: + 'A senior software engineer based in Shanghai who specializes in frontend development with React and TypeScript. Passionate about open source and building developer tools.', + tagline: 'Senior frontend engineer & OSS contributor', + }, + }, + }); + expect(result).toMatchSnapshot(); + }); + + it('should format persona with only narrative (no tagline)', () => { + const result = promptUserMemory({ + memories: { + persona: { + narrative: + 'A product designer who transitioned from engineering. Enjoys bridging the gap between design and development.', + }, + }, + }); + expect(result).toMatchSnapshot(); + }); + + it('should skip persona with null values', () => { + const result = promptUserMemory({ + memories: { + persona: { + narrative: null, + tagline: null, + }, + }, + }); + expect(result).toBe(''); + }); + + it('should skip persona with empty string values', () => { + const result = promptUserMemory({ + memories: { + persona: { + narrative: '', + tagline: '', + }, + }, + }); + expect(result).toBe(''); + }); }); describe('mixed memory types', () => { @@ -477,6 +577,40 @@ describe('promptUserMemory', () => { expect(result).toMatchSnapshot(); }); + it('should format persona combined with other memory types', () => { + const result = promptUserMemory({ + memories: { + contexts: [ + { + description: 'Working on AI products', + id: 'ctx-1', + title: 'Current Work', + }, + ], + identities: [ + { + description: 'User is a senior engineer', + id: 'id-1', + role: 'Senior Engineer', + type: 'professional', + }, + ], + persona: { + narrative: + 'A tech lead who focuses on AI-powered developer tools and has a strong preference for TypeScript.', + tagline: 'AI-focused tech lead', + }, + preferences: [ + { + conclusionDirectives: 'Prefer concise responses with code examples', + id: 'pref-1', + }, + ], + }, + }); + expect(result).toMatchSnapshot(); + }); + it('should format experiences and preferences without contexts', () => { const result = promptUserMemory({ memories: { diff --git a/packages/prompts/src/prompts/userMemory/index.ts b/packages/prompts/src/prompts/userMemory/index.ts index 203aa0ca1d..37561371a8 100644 --- a/packages/prompts/src/prompts/userMemory/index.ts +++ b/packages/prompts/src/prompts/userMemory/index.ts @@ -23,16 +23,23 @@ export interface UserMemoryPreferenceItem { export type IdentityType = 'demographic' | 'personal' | 'professional'; export interface UserMemoryIdentityItem { + capturedAt?: string | Date | null; description?: string | null; id?: string; role?: string | null; type?: IdentityType | string | null; } +export interface UserMemoryPersonaItem { + narrative?: string | null; + tagline?: string | null; +} + export interface UserMemoryData { contexts?: UserMemoryContextItem[]; experiences?: UserMemoryExperienceItem[]; identities?: UserMemoryIdentityItem[]; + persona?: UserMemoryPersonaItem; preferences?: UserMemoryPreferenceItem[]; } @@ -97,11 +104,34 @@ const isValidIdentityItem = (item: UserMemoryIdentityItem): boolean => { /** * Formats a single identity memory item */ +const formatDateOnly = (value: string | Date): string => { + const d = typeof value === 'string' ? new Date(value) : value; + if (Number.isNaN(d.getTime())) return ''; + return d.toISOString().slice(0, 10); // "2025-02-23" +}; + const formatIdentityItem = (item: UserMemoryIdentityItem): string => { const typeAttr = item.type ? ` type="${item.type}"` : ''; const roleAttr = item.role ? ` role="${item.role}"` : ''; const idAttr = item.id ? ` id="${item.id}"` : ''; - return ` ${item.description || ''}`; + const capturedAtAttr = item.capturedAt ? ` capturedAt="${formatDateOnly(item.capturedAt)}"` : ''; + return ` ${item.description || ''}`; +}; + +/** + * Check if a persona item has meaningful content + */ +const isValidPersonaItem = (item?: UserMemoryPersonaItem | null): item is UserMemoryPersonaItem => { + if (!item) return false; + return !!(item.narrative || item.tagline); +}; + +/** + * Formats a persona memory item as XML + */ +const formatPersonaItem = (item: UserMemoryPersonaItem): string => { + const taglineAttr = item.tagline ? ` tagline="${item.tagline}"` : ''; + return `\n${item.narrative || ''}\n`; }; /** @@ -115,6 +145,7 @@ const formatIdentityItem = (item: UserMemoryIdentityItem): string => { */ export const promptUserMemory = ({ memories }: PromptUserMemoryOptions): string => { // Filter out empty/invalid items + const hasPersona = isValidPersonaItem(memories.persona); const identities = (memories.identities || []).filter(isValidIdentityItem); const contexts = (memories.contexts || []).filter(isValidContextItem); const experiences = (memories.experiences || []).filter(isValidExperienceItem); @@ -126,7 +157,7 @@ export const promptUserMemory = ({ memories }: PromptUserMemoryOptions): string const hasPreferences = preferences.length > 0; // If no memories at all, return empty - if (!hasIdentities && !hasContexts && !hasExperiences && !hasPreferences) { + if (!hasPersona && !hasIdentities && !hasContexts && !hasExperiences && !hasPreferences) { return ''; } @@ -134,7 +165,10 @@ export const promptUserMemory = ({ memories }: PromptUserMemoryOptions): string 'The following are memories about this user retrieved from previous conversations. Use this information to personalize your responses and maintain continuity.', ]; - // Add instruction + // Add persona section (highest-level user context) + if (hasPersona) { + contentParts.push(formatPersonaItem(memories.persona!)); + } // Add identities section (user's identity information) if (hasIdentities) { diff --git a/src/app/(backend)/api/agent/webhooks/bot-callback/route.ts b/src/app/(backend)/api/agent/webhooks/bot-callback/route.ts index 64aec0343c..680ae34acd 100644 --- a/src/app/(backend)/api/agent/webhooks/bot-callback/route.ts +++ b/src/app/(backend)/api/agent/webhooks/bot-callback/route.ts @@ -3,6 +3,7 @@ import { NextResponse } from 'next/server'; import { getServerDB } from '@/database/core/db-adaptor'; import { AgentBotProviderModel } from '@/database/models/agentBotProvider'; +import { TopicModel } from '@/database/models/topic'; import { verifyQStashSignature } from '@/libs/qstash'; import { KeyVaultsGateKeeper } from '@/server/modules/KeyVaultsEncrypt'; import { DiscordRestApi } from '@/server/services/bot/discordRestApi'; @@ -12,6 +13,7 @@ import { renderStepProgress, splitMessage, } from '@/server/services/bot/replyTemplate'; +import { SystemAgentService } from '@/server/services/systemAgent'; const log = debug('api-route:agent:bot-callback'); @@ -112,6 +114,39 @@ export async function POST(request: Request): Promise { log('bot-callback: failed to remove eyes reaction: %O', error); } } + + // Fire-and-forget: summarize topic title and update Discord thread name + const { reason, topicId, userId, userPrompt, lastAssistantContent } = body; + if (reason !== 'error' && topicId && userId && userPrompt && lastAssistantContent) { + const topicModel = new TopicModel(serverDB, userId); + topicModel + .findById(topicId) + .then(async (topic) => { + // Only generate when topic has an empty title + if (topic?.title) return; + + const systemAgent = new SystemAgentService(serverDB, userId); + const title = await systemAgent.generateTopicTitle({ + lastAssistantContent, + userPrompt, + }); + if (!title) return; + + await topicModel.update(topicId, { title }); + + // Update Discord thread name if there's a thread ID + const parts = platformThreadId.split(':'); + const threadId = parts[3]; + if (threadId) { + discord.updateChannelName(threadId, title).catch((error) => { + log('bot-callback: failed to update Discord thread name: %O', error); + }); + } + }) + .catch((error) => { + log('bot-callback: topic title summarization failed: %O', error); + }); + } } else { return NextResponse.json({ error: `Unknown callback type: ${type}` }, { status: 400 }); } diff --git a/src/server/modules/AgentRuntime/RuntimeExecutors.ts b/src/server/modules/AgentRuntime/RuntimeExecutors.ts index 89917290d0..cefd4920f6 100644 --- a/src/server/modules/AgentRuntime/RuntimeExecutors.ts +++ b/src/server/modules/AgentRuntime/RuntimeExecutors.ts @@ -4,8 +4,8 @@ import { type CallLLMPayload, type GeneralAgentCallLLMResultPayload, type InstructionExecutor, + UsageCounter, } from '@lobechat/agent-runtime'; -import { UsageCounter } from '@lobechat/agent-runtime'; import { ToolNameResolver } from '@lobechat/context-engine'; import { parse } from '@lobechat/conversation-flow'; import { consumeStreamUntilDone } from '@lobechat/model-runtime'; @@ -135,7 +135,8 @@ export const createRuntimeExecutors = ( let processedMessages; if (agentConfig) { const { LOBE_DEFAULT_MODEL_LIST } = await import('model-bank'); - const processedResult = await serverMessagesEngine({ + + const contextEngineInput = { capabilities: { isCanUseFC: (m: string, p: string) => { const info = LOBE_DEFAULT_MODEL_LIST.find( @@ -183,8 +184,17 @@ export const createRuntimeExecutors = ( toolsConfig: { tools: agentConfig.plugins ?? [], }, - }); - processedMessages = processedResult; + userMemory: state.metadata?.userMemory, + }; + + processedMessages = await serverMessagesEngine(contextEngineInput); + + // Emit context engine event for tracing (captures input params and final LLM messages) + events.push({ + input: contextEngineInput, + output: processedMessages, + type: 'context_engine_result', + } as any); } else { processedMessages = llmPayload.messages; } diff --git a/src/server/services/agentRuntime/AgentRuntimeService.ts b/src/server/services/agentRuntime/AgentRuntimeService.ts index 2704fad137..154ac088a7 100644 --- a/src/server/services/agentRuntime/AgentRuntimeService.ts +++ b/src/server/services/agentRuntime/AgentRuntimeService.ts @@ -266,10 +266,24 @@ export class AgentRuntimeService { discordContext, evalContext, maxSteps, + userMemory, } = params; try { - log('[%s] Creating new operation (autoStart: %s)', operationId, autoStart); + const memories = userMemory?.memories; + log( + '[%s] Creating new operation (autoStart: %s) with params: model=%s, provider=%s, tools=%d, messages=%d, manifests=%d, memory=%s', + operationId, + autoStart, + agentConfig?.model, + agentConfig?.provider, + tools?.length ?? 0, + initialMessages.length, + toolManifestMap ? Object.keys(toolManifestMap).length : 0, + memories + ? `{contexts:${memories.contexts?.length ?? 0},experiences:${memories.experiences?.length ?? 0},preferences:${memories.preferences?.length ?? 0},identities:${memories.identities?.length ?? 0},activities:${memories.activities?.length ?? 0},persona:${memories.persona ? 'yes' : 'no'}}` + : 'none', + ); // Initialize operation state - create state before saving const initialState = { @@ -289,6 +303,7 @@ export class AgentRuntimeService { stepWebhook, stream, userId, + userMemory, webhookDelivery, workingDirectory: agentConfig?.chatConfig?.localSystem?.workingDirectory, ...appContext, @@ -358,7 +373,6 @@ export class AgentRuntimeService { const { operationId, stepIndex, context, humanInput, approvedToolCall, rejectionReason } = params; - // Get registered callbacks const callbacks = this.getStepCallbacks(operationId); // ===== Distributed lock: prevent duplicate execution from QStash retries ===== @@ -701,6 +715,55 @@ export class AgentRuntimeService { } } + // Dev mode: record step snapshot to disk for agent-tracing CLI + if (process.env.NODE_ENV === 'development') { + try { + const { FileSnapshotStore } = await import('@lobechat/agent-tracing'); + const store = new FileSnapshotStore(); + + const partial = (await store.loadPartial(operationId)) ?? { steps: [] }; + + if (!partial.startedAt) { + partial.startedAt = Date.now(); + partial.model = + (agentState?.metadata as any)?.agentConfig?.model ?? + agentState?.modelRuntimeConfig?.model; + partial.provider = + (agentState?.metadata as any)?.agentConfig?.provider ?? + agentState?.modelRuntimeConfig?.provider; + } + + if (!partial.steps) partial.steps = []; + partial.steps.push({ + completedAt: Date.now(), + content: stepPresentationData.content, + context: { + payload: currentContext?.payload, + phase: currentContext?.phase ?? 'unknown', + stepContext: currentContext?.stepContext, + }, + events: stepResult.events as any, + executionTimeMs: stepPresentationData.executionTimeMs, + inputTokens: stepPresentationData.stepInputTokens, + messages: agentState?.messages, + messagesAfter: stepResult.newState.messages, + outputTokens: stepPresentationData.stepOutputTokens, + reasoning: stepPresentationData.reasoning, + startedAt: startAt, + stepIndex, + stepType: stepPresentationData.stepType, + toolsCalling: stepPresentationData.toolsCalling, + toolsResult: stepPresentationData.toolsResult, + totalCost: stepPresentationData.totalCost, + totalTokens: stepPresentationData.totalTokens, + }); + + await store.savePartial(operationId, partial); + } catch { + // agent-tracing not available, skip silently + } + } + // Update step tracking in state metadata and trigger step webhook if (stepResult.newState.metadata?.stepWebhook) { const prevTracking = stepResult.newState.metadata._stepTracking || {}; @@ -770,6 +833,44 @@ export class AgentRuntimeService { log('[%s] onComplete callback error: %O', operationId, callbackError); } } + + // Dev mode: finalize tracing snapshot + if (process.env.NODE_ENV === 'development') { + try { + const { FileSnapshotStore } = await import('@lobechat/agent-tracing'); + const store = new FileSnapshotStore(); + const partial = await store.loadPartial(operationId); + + if (partial) { + const snapshot = { + completedAt: Date.now(), + completionReason: reason, + error: stepResult.newState.error + ? { + message: String( + stepResult.newState.error.message ?? stepResult.newState.error, + ), + type: String(stepResult.newState.error.type ?? 'unknown'), + } + : undefined, + model: partial.model, + operationId, + provider: partial.provider, + startedAt: partial.startedAt ?? Date.now(), + steps: (partial.steps ?? []).sort((a, b) => a.stepIndex - b.stepIndex), + totalCost: stepResult.newState.cost?.total ?? 0, + totalSteps: stepResult.newState.stepCount, + totalTokens: stepResult.newState.usage?.llm?.tokens?.total ?? 0, + traceId: operationId, + }; + + await store.save(snapshot as any); + await store.removePartial(operationId); + } + } catch { + // agent-tracing not available, skip silently + } + } } return { @@ -1280,6 +1381,11 @@ export class AgentRuntimeService { (m: { content?: string; role: string }) => m.role === 'assistant' && m.content, )?.content; + // Extract first user prompt for downstream consumers (e.g., topic title summarization) + const userPrompt = state.messages?.find( + (m: { content?: string; role: string }) => m.role === 'user', + )?.content; + const delivery = state.metadata?.webhookDelivery || 'fetch'; await this.deliverWebhook( @@ -1297,8 +1403,11 @@ export class AgentRuntimeService { status: state.status, steps: state.stepCount, toolCalls: state.usage?.tools?.totalCalls, + topicId: state.metadata?.topicId, totalTokens: state.usage?.llm?.tokens?.total, type: 'completion', + userId: state.metadata?.userId, + userPrompt, }, delivery, operationId, diff --git a/src/server/services/agentRuntime/types.ts b/src/server/services/agentRuntime/types.ts index c89cdc2187..fb6edb318b 100644 --- a/src/server/services/agentRuntime/types.ts +++ b/src/server/services/agentRuntime/types.ts @@ -2,6 +2,8 @@ import { type AgentRuntimeContext, type AgentState } from '@lobechat/agent-runti import { type LobeToolManifest } from '@lobechat/context-engine'; import { type UserInterventionConfig } from '@lobechat/types'; +import { type ServerUserMemoryConfig } from '@/server/modules/Mecha/ContextEngineering/types'; + // ==================== Step Lifecycle Callbacks ==================== /** @@ -171,6 +173,8 @@ export interface OperationCreationParams { * Use { approvalMode: 'headless' } for async tasks that should never wait for human approval */ userInterventionConfig?: UserInterventionConfig; + /** User memory (persona) for injection into LLM context */ + userMemory?: ServerUserMemoryConfig; /** * Webhook delivery method. * - 'fetch': plain HTTP POST (default) diff --git a/src/server/services/aiAgent/index.ts b/src/server/services/aiAgent/index.ts index bb893e0889..c848052d8b 100644 --- a/src/server/services/aiAgent/index.ts +++ b/src/server/services/aiAgent/index.ts @@ -23,11 +23,14 @@ import { MessageModel } from '@/database/models/message'; import { PluginModel } from '@/database/models/plugin'; import { ThreadModel } from '@/database/models/thread'; import { TopicModel } from '@/database/models/topic'; +import { UserModel } from '@/database/models/user'; +import { UserPersonaModel } from '@/database/models/userMemory/persona'; import { createServerAgentToolsEngine, type EvalContext, type ServerAgentToolsContext, } from '@/server/modules/Mecha'; +import { type ServerUserMemoryConfig } from '@/server/modules/Mecha/ContextEngineering/types'; import { AgentService } from '@/server/services/agent'; import { AgentRuntimeService } from '@/server/services/agentRuntime'; import { type StepLifecycleCallbacks } from '@/server/services/agentRuntime/types'; @@ -106,6 +109,12 @@ interface InternalExecAgentParams extends ExecAgentParams { * Defaults to true. Set to false for non-streaming scenarios (e.g., bot integrations). */ stream?: boolean; + /** + * Custom title for the topic. + * When provided (including empty string), overrides the default prompt-based title. + * When undefined, falls back to prompt.slice(0, 50). + */ + title?: string; /** Topic creation trigger source ('cron' | 'chat' | 'api') */ trigger?: string; /** @@ -182,6 +191,7 @@ export class AiAgentService { files, stepCallbacks, stream, + title, trigger, cronJobId, evalContext, @@ -231,7 +241,8 @@ export class AiAgentService { const newTopic = await this.topicModel.create({ agentId: resolvedAgentId, metadata, - title: prompt.slice(0, 50) + (prompt.length > 50 ? '...' : ''), + title: + title !== undefined ? title : prompt.slice(0, 50) + (prompt.length > 50 ? '...' : ''), trigger, }); topicId = newTopic.id; @@ -432,7 +443,48 @@ export class AiAgentService { ); } - // 8. Get existing messages if provided + // 8. Fetch user persona for memory injection + // Persona is user-level global memory, only depends on user's global memory setting + let userMemory: ServerUserMemoryConfig | undefined; + + let globalMemoryEnabled = true; // default: enabled (matches DEFAULT_MEMORY_SETTINGS) + try { + const userModel = new UserModel(this.db, this.userId); + const settings = await userModel.getUserSettings(); + const memorySettings = settings?.memory as { enabled?: boolean } | undefined; + globalMemoryEnabled = memorySettings?.enabled !== false; + } catch (error) { + log('execAgent: failed to fetch user memory settings: %O', error); + } + + log('execAgent: memory check — globalMemoryEnabled=%s', globalMemoryEnabled); + + if (globalMemoryEnabled) { + try { + const personaModel = new UserPersonaModel(this.db, this.userId); + const persona = await personaModel.getLatestPersonaDocument(); + + if (persona?.persona) { + userMemory = { + fetchedAt: Date.now(), + memories: { + contexts: [], + experiences: [], + persona: { + narrative: persona.persona, + tagline: persona.tagline, + }, + preferences: [], + }, + }; + log('execAgent: fetched user persona (version: %d)', persona.version); + } + } catch (error) { + log('execAgent: failed to fetch user persona: %O', error); + } + } + + // 9. Get existing messages if provided let historyMessages: any[] = []; if (existingMessageIds.length > 0) { historyMessages = await this.messageModel.query({ @@ -581,6 +633,7 @@ export class AiAgentService { tools, userId: this.userId, userInterventionConfig, + userMemory, webhookDelivery, }); diff --git a/src/server/services/bot/AgentBridgeService.ts b/src/server/services/bot/AgentBridgeService.ts index 3c013041b2..868402cf01 100644 --- a/src/server/services/bot/AgentBridgeService.ts +++ b/src/server/services/bot/AgentBridgeService.ts @@ -4,11 +4,13 @@ import { emoji } from 'chat'; import debug from 'debug'; import urlJoin from 'url-join'; +import { TopicModel } from '@/database/models/topic'; import { UserModel } from '@/database/models/user'; import type { LobeChatDatabase } from '@/database/type'; import { appEnv } from '@/envs/app'; import { AiAgentService } from '@/server/services/aiAgent'; import { isQueueAgentRuntimeEnabled } from '@/server/services/queue/impls'; +import { SystemAgentService } from '@/server/services/systemAgent'; import { formatPrompt as formatPromptUtil } from './formatPrompt'; import { @@ -334,6 +336,7 @@ export class AgentBridgeService { files, prompt, stepWebhook: { body: webhookBody, url: callbackUrl }, + title: '', trigger, userInterventionConfig: { approvalMode: 'headless' }, webhookDelivery: 'qstash', @@ -415,6 +418,7 @@ export class AgentBridgeService { : undefined, files, prompt, + title: '', stepCallbacks: { onAfterStep: async (stepData) => { const { content, shouldContinue, toolsCalling } = stepData; @@ -499,6 +503,32 @@ export class AgentBridgeService { chunks.length, ); resolve({ reply: lastAssistantContent, topicId: resolvedTopicId }); + + // Fire-and-forget: summarize topic title in DB (no Discord rename in local mode) + if (resolvedTopicId && prompt) { + const topicModel = new TopicModel(this.db, this.userId); + topicModel + .findById(resolvedTopicId) + .then(async (topic) => { + if (topic?.title) return; + + const systemAgent = new SystemAgentService(this.db, this.userId); + const title = await systemAgent.generateTopicTitle({ + lastAssistantContent, + userPrompt: prompt, + }); + if (!title) return; + + await topicModel.update(resolvedTopicId, { title }); + }) + .catch((error) => { + log( + 'executeWithInMemoryCallbacks: topic title summarization failed: %O', + error, + ); + }); + } + return; } diff --git a/src/server/services/bot/BotMessageRouter.ts b/src/server/services/bot/BotMessageRouter.ts index 05d0ff873a..a76fa2f233 100644 --- a/src/server/services/bot/BotMessageRouter.ts +++ b/src/server/services/bot/BotMessageRouter.ts @@ -73,25 +73,19 @@ export class BotMessageRouter { // Check for forwarded Gateway event (from Gateway worker) const gatewayToken = req.headers.get('x-discord-gateway-token'); if (gatewayToken) { - log('Gateway forwarded event, token=%s...', gatewayToken.slice(0, 10)); - log( - 'Known tokens: %o', - [...this.botInstancesByToken.keys()].map((t) => t.slice(0, 10)), - ); - - // Log forwarded event details for debugging + // Log forwarded event details try { const bodyText = new TextDecoder().decode(bodyBuffer); const event = JSON.parse(bodyText); if (event.type === 'GATEWAY_MESSAGE_CREATE') { const d = event.data; + const mentions = d?.mentions?.map((m: any) => m.username).join(', '); log( - 'MESSAGE_CREATE: author=%s (id=%s, bot=%s), mentions=%o, content=%s', + 'Gateway MESSAGE_CREATE: author=%s (bot=%s), mentions=[%s], content=%s', d?.author?.username, - d?.author?.id, d?.author?.bot, - d?.mentions?.map((m: any) => ({ id: m.id, username: m.username })), + mentions || '', d?.content?.slice(0, 100), ); } @@ -101,7 +95,6 @@ export class BotMessageRouter { const bot = this.botInstancesByToken.get(gatewayToken); if (bot?.webhooks && 'discord' in bot.webhooks) { - log('Matched bot by token'); return bot.webhooks.discord(this.cloneRequest(req, bodyBuffer)); } diff --git a/src/server/services/bot/discordRestApi.ts b/src/server/services/bot/discordRestApi.ts index 3a9d43cfbf..c5806d8bda 100644 --- a/src/server/services/bot/discordRestApi.ts +++ b/src/server/services/bot/discordRestApi.ts @@ -28,6 +28,12 @@ export class DiscordRestApi { ); } + async updateChannelName(channelId: string, name: string): Promise { + const truncatedName = name.slice(0, 100); // Discord thread name limit + log('updateChannelName: channel=%s, name=%s', channelId, truncatedName); + await this.rest.patch(Routes.channel(channelId), { body: { name: truncatedName } }); + } + async createMessage(channelId: string, content: string): Promise<{ id: string }> { log('createMessage: channel=%s', channelId); const data = (await this.rest.post(Routes.channelMessages(channelId), { diff --git a/src/server/services/systemAgent/index.ts b/src/server/services/systemAgent/index.ts new file mode 100644 index 0000000000..5c5a1661b7 --- /dev/null +++ b/src/server/services/systemAgent/index.ts @@ -0,0 +1,117 @@ +import { DEFAULT_SYSTEM_AGENT_CONFIG } from '@lobechat/const'; +import { chainSummaryTitle } from '@lobechat/prompts'; +import { type UserSystemAgentConfig, type UserSystemAgentConfigKey } from '@lobechat/types'; +import debug from 'debug'; + +import { UserModel } from '@/database/models/user'; +import { type LobeChatDatabase } from '@/database/type'; +import { initModelRuntimeFromDB } from '@/server/modules/ModelRuntime'; + +const log = debug('lobe-server:system-agent-service'); + +const TOPIC_TITLE_SCHEMA = { + name: 'topic_title', + schema: { + additionalProperties: false, + properties: { + title: { description: 'A concise topic title', type: 'string' }, + }, + required: ['title'], + type: 'object' as const, + }, + strict: true, +}; + +/** + * Server-side service for SystemAgent automated tasks. + * + * Encapsulates the common pattern: read user's systemAgent config → build chain prompt + * → call LLM via generateObject → return structured result. + * + * Each public method corresponds to a `UserSystemAgentConfigKey` task type + * (topic, translation, agentMeta, etc.). + */ +export class SystemAgentService { + private readonly db: LobeChatDatabase; + private readonly userId: string; + + constructor(db: LobeChatDatabase, userId: string) { + this.db = db; + this.userId = userId; + } + + /** + * Generate a concise topic title from user prompt + assistant reply. + * + * @returns The generated title string, or null on failure + */ + async generateTopicTitle(params: { + lastAssistantContent: string; + userPrompt: string; + }): Promise { + const { userPrompt, lastAssistantContent } = params; + + try { + const { model, provider } = await this.getTaskModelConfig('topic'); + const locale = await this.getUserLocale(); + + log('generateTopicTitle: locale=%s, model=%s, provider=%s', locale, model, provider); + + const messages = [ + { content: userPrompt, role: 'user' as const }, + { content: lastAssistantContent, role: 'assistant' as const }, + ]; + + const payload = chainSummaryTitle(messages, locale); + + const modelRuntime = await initModelRuntimeFromDB(this.db, this.userId, provider); + const result = await modelRuntime.generateObject({ + messages: payload.messages as any[], + model, + schema: TOPIC_TITLE_SCHEMA, + }); + + const title = (result as { title?: string })?.title?.trim(); + if (!title) { + log('generateTopicTitle: LLM returned empty title'); + return null; + } + + log('generateTopicTitle: generated title="%s"', title); + return title; + } catch (error) { + console.error('SystemAgentService.generateTopicTitle failed:', error); + return null; + } + } + + // ============== Private Helpers ============== // + + /** + * Get the model/provider config for a specific systemAgent task type. + * Falls back to DEFAULT_SYSTEM_AGENT_CONFIG when user has no custom settings. + */ + private async getTaskModelConfig( + taskKey: UserSystemAgentConfigKey, + ): Promise<{ model: string; provider: string }> { + const userModel = new UserModel(this.db, this.userId); + const settings = await userModel.getUserSettings(); + const systemAgent = settings?.systemAgent as Partial | undefined; + + const taskConfig = systemAgent?.[taskKey]; + const defaults = DEFAULT_SYSTEM_AGENT_CONFIG[taskKey]; + + return { + model: taskConfig?.model || defaults.model, + provider: taskConfig?.provider || defaults.provider, + }; + } + + /** + * Get the user's preferred response language (locale). + */ + private async getUserLocale(): Promise { + const userInfo = await UserModel.getInfoForAIGeneration(this.db, this.userId); + return userInfo.responseLanguage || 'en-US'; + } +}