@lobehub/lobehub 2.0.0-next.38 → 2.0.0-next.39

package/src/libs/trpc/middleware/openTelemetry.ts

@@ -0,0 +1,141 @@
+import type { Attributes, Span } from '@lobechat/observability-otel/api';
+import { SpanKind, SpanStatusCode, diag, trace } from '@lobechat/observability-otel/api';
+import {
+  ATTR_ERROR_TYPE,
+  ATTR_EXCEPTION_MESSAGE,
+  ATTR_EXCEPTION_STACKTRACE,
+  DEFAULT_ERROR_CODE,
+  DEFAULT_SUCCESS_STATUS,
+  TRPCAttribute,
+  createAttributesForMetrics,
+  getPayloadSize,
+  serverDurationHistogram,
+  serverRequestSizeHistogram,
+  serverRequestsPerRpcHistogram,
+  serverResponseSizeHistogram,
+  serverResponsesPerRpcHistogram,
+  tRPCConventionFromPathAndType,
+} from '@lobechat/observability-otel/trpc';
+import { TRPCError } from '@trpc/server';
+import { env } from 'node:process';
+
+import { name } from '../../../../package.json';
+import { trpc } from '../lambda/init';
+
+const tracer = trace.getTracer('trpc-server');
+
+const recordRpcServerMetrics = ({
+  attributes,
+  durationMs,
+  requestSize,
+  responseSize,
+}: {
+  attributes: Attributes;
+  durationMs: number;
+  requestSize?: number;
+  responseSize?: number;
+}) => {
+  serverDurationHistogram.record(durationMs, attributes);
+  serverRequestsPerRpcHistogram.record(1, attributes);
+  serverResponsesPerRpcHistogram.record(1, attributes);
+
+  if (typeof requestSize === 'number') {
+    serverRequestSizeHistogram.record(requestSize, attributes);
+  }
+
+  if (typeof responseSize === 'number') {
+    serverResponseSizeHistogram.record(responseSize, attributes);
+  }
+};
+
+const finalizeSpanWithError = (span: Span, error: unknown) => {
+  span.setStatus({
+    code: SpanStatusCode.ERROR,
+    message: error instanceof Error ? error.message : 'Unknown error',
+  });
+
+  if (error instanceof Error) {
+    span.recordException(error);
+    span.setAttribute(ATTR_ERROR_TYPE, error.constructor.name);
+    span.setAttribute(ATTR_EXCEPTION_MESSAGE, error.message);
+    span.setAttribute(ATTR_EXCEPTION_STACKTRACE, error.stack || '');
+  }
+};
+
+export const openTelemetry = trpc.middleware(async ({ path, type, next, getRawInput }) => {
+  if (!env.ENABLE_TELEMETRY) {
+    diag.debug(name, 'telemetry disabled', env.ENABLE_TELEMETRY);
+
+    return next();
+  }
+
+  diag.debug(name, 'tRPC instrumentation', 'incomingRequest');
+
+  const spanName = `tRPC ${type.toUpperCase()} ${path}`;
+  const baseAttributes = tRPCConventionFromPathAndType(path, type);
+  const input = getRawInput();
+  const requestSize = getPayloadSize(input);
+
+  const span = tracer.startSpan(spanName, {
+    attributes: baseAttributes,
+    kind: SpanKind.SERVER,
+  });
+
+  const startTimestamp = Date.now();
+
+  try {
+    const result = await next();
+    diag.debug(name, 'tRPC instrumentation', 'requestHandled');
+
+    const responseSize = getPayloadSize(result.ok ? result.data : result.error);
+
+    const durationMs = Date.now() - startTimestamp;
+    const statusCode = result.ok ? DEFAULT_SUCCESS_STATUS : result.error.code;
+    span.setAttribute(TRPCAttribute.RPC_TRPC_STATUS_CODE, statusCode);
+
+    if (result.ok) {
+      span.setStatus({ code: SpanStatusCode.OK });
+    } else {
+      finalizeSpanWithError(span, result.error);
+    }
+
+    recordRpcServerMetrics({
+      attributes: createAttributesForMetrics(baseAttributes, statusCode, {
+        [TRPCAttribute.RPC_TRPC_SUCCESS]: result.ok,
+        ...(result.ok ? undefined : { [ATTR_ERROR_TYPE]: result.error.code }),
+      }),
+      durationMs,
+      requestSize,
+      responseSize,
+    });
+
+    diag.debug(name, 'tRPC instrumentation', 'metrics recorded');
+
+    return result;
+  } catch (error) {
+    diag.error(name, 'tRPC instrumentation', 'requestError', error);
+
+    const durationMs = Date.now() - startTimestamp;
+    const trpcError = error instanceof TRPCError ? error : undefined;
+    const statusCode = trpcError ? trpcError.code : DEFAULT_ERROR_CODE;
+
+    span.setAttribute(TRPCAttribute.RPC_TRPC_STATUS_CODE, statusCode);
+    finalizeSpanWithError(span, error);
+
+    recordRpcServerMetrics({
+      attributes: createAttributesForMetrics(baseAttributes, statusCode, {
+        [TRPCAttribute.RPC_TRPC_SUCCESS]: false,
+        ...(trpcError ? { [ATTR_ERROR_TYPE]: trpcError.code } : undefined),
+      }),
+      durationMs,
+      requestSize,
+      responseSize: getPayloadSize(trpcError ? trpcError : error),
+    });
+
+    diag.error(name, 'tRPC instrumentation', 'metrics recorded with error', error);
+
+    throw error;
+  } finally {
+    span.end();
+  }
+});

package/src/locales/default/auth.ts

@@ -147,5 +147,49 @@ export default {
     profile: '个人资料',
     security: '安全',
     stats: '数据统计',
+    usage: '用量统计',
+  },
+  usage: {
+    activeModels: {
+      modelTable: '模型列表',
+      models: '活跃模型',
+      providerTable: '提供商列表',
+      providers: '活跃提供商',
+      table: {
+        calls: '调用次数',
+        model: '模型',
+        provider: '提供商',
+        spend: '花费',
+      },
+    },
+    cards: {
+      month: {
+        modelCalls: '模型调用',
+        title: '本月花费',
+      },
+      today: {
+        title: '今日花费',
+        yesterday: '昨日',
+      },
+    },
+    table: {
+      actions: '操作',
+      createdAt: '使用时间',
+      inputTokens: '输入 Token',
+      model: '模型',
+      outputTokens: '输出 Token',
+      spend: '花费',
+      tps: 'TPS',
+      ttft: 'TTFT',
+      type: '调用类型',
+    },
+    trends: {
+      spend: '金额',
+      tokens: 'Token',
+    },
+    welcome: {
+      model: '模型',
+      provider: '提供商',
+    },
   },
 };

package/src/locales/default/chat.ts

@@ -188,6 +188,7 @@ export default {
   },
 
   messageAction: {
+    continueGeneration: '继续生成',
     delAndRegenerate: '删除并重新生成',
     deleteDisabledByThreads: '存在子话题，不能删除',
     regenerate: '重新生成',

package/src/server/routers/desktop/mcp.ts

@@ -87,13 +87,11 @@ export const mcpRouter = router({
     )
     .mutation(async ({ input }) => {
       // Pass the validated params, toolName, and args to the service
-      const data = await mcpService.callTool(
+      return await mcpService.callTool(
         { ...input.params, env: input.env },
         input.toolName,
         input.args,
       );
-
-      return JSON.stringify(data);
     }),
 
   validMcpServerInstallable: mcpProcedure

package/src/server/routers/lambda/index.ts

@@ -30,6 +30,7 @@ import { sessionGroupRouter } from './sessionGroup';
 import { threadRouter } from './thread';
 import { topicRouter } from './topic';
 import { uploadRouter } from './upload';
+import { usageRouter } from './usage';
 import { userRouter } from './user';
 
 export const lambdaRouter = router({
@@ -61,6 +62,7 @@ export const lambdaRouter = router({
   thread: threadRouter,
   topic: topicRouter,
   upload: uploadRouter,
+  usage: usageRouter,
   user: userRouter,
 });
 

package/src/server/routers/lambda/usage.ts

@@ -0,0 +1,36 @@
+import { z } from 'zod';
+
+import { authedProcedure, router } from '@/libs/trpc/lambda';
+import { serverDatabase } from '@/libs/trpc/lambda/middleware';
+import { UsageRecordService } from '@/server/services/usage';
+
+const usageProcedure = authedProcedure.use(serverDatabase).use(async (opts) => {
+  const { ctx } = opts;
+  return opts.next({
+    ctx: {
+      usageRecordService: new UsageRecordService(ctx.serverDB, ctx.userId),
+    },
+  });
+});
+
+export const usageRouter = router({
+  findAndGroupByDay: usageProcedure
+    .input(
+      z.object({
+        mo: z.string().optional(),
+      }),
+    )
+    .query(async ({ ctx, input }) => {
+      return await ctx.usageRecordService.findAndGroupByDay(input.mo);
+    }),
+
+  findByMonth: usageProcedure
+    .input(
+      z.object({
+        mo: z.string().optional(),
+      }),
+    )
+    .query(async ({ ctx, input }) => {
+      return await ctx.usageRecordService.findByMonth(input.mo);
+    }),
+});

package/src/server/routers/tools/mcp.ts

@@ -100,8 +100,6 @@ export const mcpRouter = router({
       checkStdioEnvironment(input.params);
 
       // Pass the validated params, toolName, and args to the service
-      const data = await mcpService.callTool(input.params, input.toolName, input.args);
-
-      return JSON.stringify(data);
+      return await mcpService.callTool(input.params, input.toolName, input.args);
     }),
 });

package/src/server/services/mcp/index.test.ts

@@ -33,7 +33,7 @@ describe('MCPService', () => {
       args: ['--test'],
     };
 
-    it('should return original data when content array is empty', async () => {
+    it('should return MCPToolCallResult when content array is empty', async () => {
       mockClient.callTool.mockResolvedValue({
         content: [],
         isError: false,
@@ -41,10 +41,12 @@ describe('MCPService', () => {
 
       const result = await mcpService.callTool(mockParams, 'testTool', '{}');
 
-      expect(result).toEqual([]);
+      expect(result.content).toBe('');
+      expect(result.success).toBe(true);
+      expect(result.state).toEqual({ content: [], isError: false });
     });
 
-    it('should return original data when content is null or undefined', async () => {
+    it('should handle null content', async () => {
       mockClient.callTool.mockResolvedValue({
         content: null,
         isError: false,
@@ -52,10 +54,11 @@ describe('MCPService', () => {
 
       const result = await mcpService.callTool(mockParams, 'testTool', '{}');
 
-      expect(result).toBeNull();
+      expect(result.content).toBe('');
+      expect(result.success).toBe(true);
     });
 
-    it('should return parsed JSON when single element contains valid JSON', async () => {
+    it('should return JSON string when single element contains valid JSON', async () => {
       const jsonData = { message: 'Hello World', status: 'success' };
       mockClient.callTool.mockResolvedValue({
         content: [{ type: 'text', text: JSON.stringify(jsonData) }],
@@ -64,7 +67,8 @@ describe('MCPService', () => {
 
       const result = await mcpService.callTool(mockParams, 'testTool', '{}');
 
-      expect(result).toEqual(jsonData);
+      expect(result.content).toBe(JSON.stringify(jsonData));
+      expect(result.success).toBe(true);
     });
 
     it('should return plain text when single element contains non-JSON text', async () => {
@@ -76,10 +80,12 @@ describe('MCPService', () => {
 
       const result = await mcpService.callTool(mockParams, 'testTool', '{}');
 
-      expect(result).toBe(textData);
+      expect(result.content).toBe(textData);
+      expect(result.success).toBe(true);
+      expect(result.state.content).toEqual([{ type: 'text', text: textData }]);
     });
 
-    it('should return original data when single element has no text', async () => {
+    it('should return empty content when single element has no text', async () => {
       const contentData = [{ type: 'text', text: '' }];
       mockClient.callTool.mockResolvedValue({
         content: contentData,
@@ -88,10 +94,12 @@ describe('MCPService', () => {
 
       const result = await mcpService.callTool(mockParams, 'testTool', '{}');
 
-      expect(result).toEqual(contentData);
+      expect(result.content).toBe('');
+      expect(result.success).toBe(true);
+      expect(result.state.content).toEqual(contentData);
     });
 
-    it('should return complete array when content has multiple elements', async () => {
+    it('should join multiple text elements with double newlines', async () => {
       const multipleContent = [
         { type: 'text', text: 'First message' },
         { type: 'text', text: 'Second message' },
@@ -105,11 +113,12 @@ describe('MCPService', () => {
 
       const result = await mcpService.callTool(mockParams, 'testTool', '{}');
 
-      // 应该直接返回完整的数组，不进行任何处理
-      expect(result).toEqual(multipleContent);
+      expect(result.content).toBe('First message\n\nSecond message\n\n{"json": "data"}');
+      expect(result.success).toBe(true);
+      expect(result.state.content).toEqual(multipleContent);
     });
 
-    it('should return complete array when content has two elements', async () => {
+    it('should join two text elements with double newline', async () => {
       const twoContent = [
         { type: 'text', text: 'First message' },
         { type: 'text', text: 'Second message' },
@@ -122,7 +131,9 @@ describe('MCPService', () => {
 
       const result = await mcpService.callTool(mockParams, 'testTool', '{}');
 
-      expect(result).toEqual(twoContent);
+      expect(result.content).toBe('First message\n\nSecond message');
+      expect(result.success).toBe(true);
+      expect(result.state.content).toEqual(twoContent);
     });
 
     it('should return error result when isError is true', async () => {
@@ -135,7 +146,9 @@ describe('MCPService', () => {
 
       const result = await mcpService.callTool(mockParams, 'testTool', '{}');
 
-      expect(result).toEqual(errorResult);
+      expect(result.content).toBe('Error occurred');
+      expect(result.success).toBe(true);
+      expect(result.state).toEqual(errorResult);
     });
 
     it('should throw TRPCError when client throws error', async () => {

package/src/server/services/mcp/index.ts

@@ -175,29 +175,40 @@ export class MCPService {
         loggableParams,
         result,
       );
-      const { content, isError } = result;
 
-      if (isError) return result;
-
-      const data = content as { text: string; type: 'text' }[];
-
-      if (!data || data.length === 0) return data;
-
-      if (data.length > 1) return data;
-
-      const text = data[0]?.text;
-      if (!text) return data;
-
-      // try to get json object, which will be stringify in the client
-      const json = safeParseJSON(text);
-      if (json) return json;
-
-      return text;
+      // TODO: map more type
+      const content = result.content
+        ? result.content
+            .map((item) => {
+              switch (item.type) {
+                case 'text': {
+                  return item.text;
+                }
+                default: {
+                  return '';
+                }
+              }
+            })
+            .filter(Boolean)
+            .join('\n\n')
+        : '';
+
+      if (result.isError) return { content, state: result, success: true };
+
+      return { content, state: result, success: true };
     } catch (error) {
       if (error instanceof McpError) {
         const mcpError = error as McpError;
 
-        return mcpError.message;
+        return {
+          content: mcpError.message,
+          error: error,
+          state: {
+            content: [{ text: mcpError.message, type: 'text' }],
+            isError: true,
+          },
+          success: false,
+        };
       }
 
       console.error(