nolo-cli 0.1.40 → 0.1.42

package/chat/dialog/actions/updateTokensAction.ts

@@ -0,0 +1,233 @@
+import { TokenUsageData } from "../../../ai/token/types";
+import { DataType } from "../../../create/types";
+import { createTokenStatsKey } from "../../../database/keys";
+import { ulid } from "ulid";
+import { format } from "date-fns";
+import { patch, read, selectById, write } from "../../../database/dbSlice";
+import { toast } from "../../../app/utils/toast";
+import {
+  createTokenRecord,
+  saveTokenRecord,
+  ModelStats,
+} from "../../../ai/token/saveTokenRecord";
+import { pino } from "pino";
+import { deductBalance } from "../../../auth/authSlice"; // <--- 1. 导入新的 deductBalance action
+import { prepareTokenUsageData } from "../../../ai/token/prepareTokenUsageData";
+
+const logger = pino({ name: "token-usage", level: "info" });
+const dialogTokenPatchQueue = new Map<string, Promise<void>>();
+
+interface DayStats {
+  userId: string;
+  period: "day";
+  timeKey: string;
+  total: ModelStats;
+  models: Record<string, ModelStats>;
+  providers: Record<string, ModelStats>;
+}
+
+const queueDialogTokenPatch = async <T>(
+  dialogKey: string,
+  task: () => Promise<T>
+): Promise<T> => {
+  const previousTask = dialogTokenPatchQueue.get(dialogKey) ?? Promise.resolve();
+  const nextTask = previousTask.catch(() => undefined).then(task);
+  const queueEntry = nextTask.then(
+    () => undefined,
+    () => undefined
+  );
+
+  dialogTokenPatchQueue.set(dialogKey, queueEntry);
+
+  try {
+    return await nextTask;
+  } finally {
+    if (dialogTokenPatchQueue.get(dialogKey) === queueEntry) {
+      dialogTokenPatchQueue.delete(dialogKey);
+    }
+  }
+};
+
+const updateStatsCounter = (
+  data: TokenUsageData,
+  stats: ModelStats = { count: 0, tokens: { input: 0, output: 0 }, cost: 0 }
+): ModelStats => ({
+  count: stats.count + 1,
+  tokens: {
+    input: stats.tokens.input + data.input_tokens,
+    output: stats.tokens.output + data.output_tokens,
+  },
+  cost: stats.cost + data.cost,
+});
+
+const updateStats = async (
+  data: TokenUsageData,
+  existingStats: DayStats | null,
+  key: string,
+  thunkApi
+) => {
+  try {
+    const stats = existingStats ?? {
+      userId: data.userId,
+      period: "day",
+      timeKey: format(Date.now(), "yyyy-MM-dd"),
+      total: { count: 0, tokens: { input: 0, output: 0 }, cost: 0 },
+      models: {},
+      providers: {},
+    };
+
+    const modelName = data.model || "unknown";
+    const providerName = data.provider || "unknown";
+
+    const cleanModels = Object.fromEntries(
+      Object.entries(stats.models).filter(
+        ([key]) => !["unknown", "undefined"].includes(key)
+      )
+    );
+
+    const updatedStats = {
+      ...stats,
+      total: updateStatsCounter(data, stats.total),
+      models: {
+        ...cleanModels,
+        [modelName]: updateStatsCounter(data, cleanModels[modelName]),
+      },
+      providers: {
+        ...stats.providers,
+        [providerName]: updateStatsCounter(data, stats.providers[providerName]),
+      },
+    };
+
+    await thunkApi.dispatch(
+      write({
+        data: { ...updatedStats, id: key, type: DataType.TOKEN },
+        customKey: key,
+      })
+    );
+
+    return updatedStats;
+  } catch (error) {
+    logger.error(
+      { key, userId: data.userId, error: error.message },
+      "Failed to update token stats"
+    );
+    toast.error("Failed to update token stats");
+    throw error;
+  }
+};
+
+export const saveTokenUsage = async (data: TokenUsageData, thunkApi) => {
+  const dateKey = format(Date.now(), "yyyy-MM-dd");
+  const tokenDayStatsKey = createTokenStatsKey(data.userId, dateKey);
+  try {
+    let currentStats = null;
+    try {
+      currentStats = await thunkApi.dispatch(read({
+        dbKey: tokenDayStatsKey
+      })).unwrap();
+    } catch (err) {
+      logger.warn({ tokenDayStatsKey }, "No existing stats found");
+    }
+
+    const updatedStats = await updateStats(
+      data,
+      currentStats,
+      tokenDayStatsKey,
+      thunkApi
+    );
+
+    return {
+      success: true,
+      id: ulid(Date.now()),
+      record: updatedStats,
+    };
+  } catch (error) {
+    logger.error(
+      {
+        key,
+        userId: data.userId,
+        error: error.message,
+        tokenData: {
+          input: data.input_tokens,
+          output: data.output_tokens,
+          model: data.model,
+        },
+      },
+      "Failed to process token usage"
+    );
+
+    toast.error("Failed to process token usage");
+    throw error;
+  }
+};
+
+export const updateTokensAction = async (
+  { dialogId, dialogKey, usage: usageRaw, agentConfig },
+  thunkApi
+) => {
+  const { currentUser } = thunkApi.getState().auth;
+  const timestamp = Date.now();
+  const prepared = prepareTokenUsageData({
+    rawUsage: usageRaw,
+    agentConfig,
+    userId: currentUser?.userId,
+    username: currentUser?.username,
+    cybotId: agentConfig.id,
+    dialogId,
+    timestamp,
+  });
+  const { usage, tokenData } = prepared;
+  const result = { cost: tokenData.cost, pay: tokenData.pay };
+
+  const persistedTokenData: TokenUsageData = {
+    ...tokenData,
+    type: DataType.TOKEN,
+    id: ulid(timestamp),
+    dateKey: format(timestamp, "yyyy-MM-dd"),
+  };
+
+  const record = createTokenRecord(persistedTokenData, {
+    cost: result.cost,
+    inputPrice: agentConfig.inputPrice,
+    outputPrice: agentConfig.outputPrice,
+  });
+
+
+  await saveTokenRecord(persistedTokenData, record, thunkApi);
+  await saveTokenUsage(persistedTokenData, thunkApi);
+
+  if (result.cost > 0) {
+    thunkApi.dispatch(deductBalance(result.cost));
+  }
+
+  if (dialogKey) {
+    await queueDialogTokenPatch(dialogKey, async () => {
+      const latestState = thunkApi.getState();
+      const dialogConfig =
+        selectById(latestState, dialogKey) ??
+        await thunkApi.dispatch(read({ dbKey: dialogKey })).unwrap();
+
+      if (!dialogConfig) {
+        throw new Error(`Dialog not found for token update: ${dialogKey}`);
+      }
+
+      await thunkApi.dispatch(
+        patch({
+          dbKey: dialogKey,
+          changes: {
+            inputTokens: (dialogConfig.inputTokens ?? 0) + usage.input_tokens,
+            outputTokens: (dialogConfig.outputTokens ?? 0) + usage.output_tokens,
+            totalCost: (dialogConfig.totalCost ?? 0) + result.cost,
+          },
+        })
+      ).unwrap();
+    });
+  }
+
+  // 返回本轮最终记账结果；TopBar 展示由持久化 dialog 统计 + runtime live 增量合并得到
+  return {
+    input_tokens: usage.input_tokens,
+    output_tokens: usage.output_tokens,
+    cost: result.cost,
+  };
+};

package/chat/dialog/dialogAgents.ts

@@ -0,0 +1,72 @@
+import type { DialogConfig } from "../../app/types";
+
+type LegacyDialogConfig = {
+  llmId?: string;
+  primaryAgentKey?: string;
+};
+
+const normalizeAgentId = (value: unknown): string | null =>
+  typeof value === "string" && value.trim() ? value.trim() : null;
+
+const dedupeAgentIds = (agentIds: Array<string | null | undefined>): string[] => {
+  const seen = new Set<string>();
+  const result: string[] = [];
+  for (const agentId of agentIds) {
+    const normalized = normalizeAgentId(agentId);
+    if (!normalized || seen.has(normalized)) continue;
+    seen.add(normalized);
+    result.push(normalized);
+  }
+  return result;
+};
+
+export const getDialogAgentIds = (
+  dialogConfig:
+    | DialogConfig
+    | (LegacyDialogConfig & { cybots?: string[] })
+    | null
+    | undefined,
+): string[] => {
+  if (!dialogConfig) return [];
+  const primaryAgentKey = normalizeAgentId(dialogConfig.primaryAgentKey);
+  if (Array.isArray(dialogConfig.cybots)) {
+    return dedupeAgentIds([primaryAgentKey, ...dialogConfig.cybots]);
+  }
+  if (primaryAgentKey) {
+    return [primaryAgentKey];
+  }
+  const legacyAgentId = normalizeAgentId(dialogConfig.llmId);
+  if (legacyAgentId) {
+    return [legacyAgentId];
+  }
+  return [];
+};
+
+export const getPrimaryDialogAgentId = (
+  dialogConfig:
+    | DialogConfig
+    | (LegacyDialogConfig & { cybots?: string[] })
+    | null
+    | undefined,
+): string | null => getDialogAgentIds(dialogConfig)[0] ?? null;
+
+export const addDialogAgentIds = (
+  existingAgentIds: string[],
+  nextAgentIds: string[],
+): string[] => {
+  const merged = [...existingAgentIds, ...nextAgentIds];
+  return merged.filter((id, index) => merged.indexOf(id) === index);
+};
+
+export const replacePrimaryDialogAgentId = (
+  existingAgentIds: string[],
+  nextPrimaryAgentId: string,
+): string[] => [
+  nextPrimaryAgentId,
+  ...existingAgentIds.filter((id) => id !== nextPrimaryAgentId),
+];
+
+export const removeDialogAgentId = (
+  existingAgentIds: string[],
+  agentIdToRemove: string,
+): string[] => existingAgentIds.filter((id) => id !== agentIdToRemove);

package/chat/dialog/dialogAttachmentCleanup.ts

@@ -0,0 +1,135 @@
+const FILE_CONTENT_RE = /(?:https?:\/\/[^/"'\s]+)?\/api\/v1\/db\/file\/content\/([^?#"'\s)]+)/g;
+
+export type DialogAttachmentMessage = {
+  id?: string;
+  dbKey?: string;
+  content?: unknown;
+  [key: string]: unknown;
+};
+
+export type DialogAttachmentCandidate = {
+  fileId: string;
+  fileDbKey: string | null;
+  size: number | null;
+  ownerType: string | null;
+  ownerId: string | null;
+  ownerDbKey: string | null;
+  source: string | null;
+  status: "delete" | "retain";
+  reason: string;
+};
+
+export type DialogAttachmentPlan = {
+  dialogId: string;
+  messageCount: number;
+  referencedFileIds: string[];
+  candidates: DialogAttachmentCandidate[];
+  deleteCandidates: DialogAttachmentCandidate[];
+  retainedCandidates: DialogAttachmentCandidate[];
+  bytesToDelete: number;
+  metadataReadFailures: Array<{ fileId: string; error: string }>;
+};
+
+export function extractFileContentIds(value: unknown): string[] {
+  const ids = new Set<string>();
+  const visit = (input: unknown) => {
+    if (typeof input === "string") {
+      for (const match of input.matchAll(FILE_CONTENT_RE)) {
+        const fileId = decodeURIComponent(match[1] ?? "").trim();
+        if (fileId) ids.add(fileId);
+      }
+      return;
+    }
+    if (Array.isArray(input)) {
+      for (const item of input) visit(item);
+      return;
+    }
+    if (input && typeof input === "object") {
+      for (const item of Object.values(input as Record<string, unknown>)) visit(item);
+    }
+  };
+  visit(value);
+  return [...ids];
+}
+
+function messageDbKey(dialogId: string, message: DialogAttachmentMessage) {
+  if (typeof message.dbKey === "string" && message.dbKey.trim()) {
+    return message.dbKey.trim();
+  }
+  if (typeof message.id === "string" && message.id.trim()) {
+    return `dialog-${dialogId}-msg-${message.id.trim()}`;
+  }
+  return null;
+}
+
+function asText(value: unknown) {
+  return typeof value === "string" && value.trim() ? value.trim() : null;
+}
+
+export function buildDialogAttachmentPlan(args: {
+  dialogId: string;
+  ownerId?: string | null;
+  includeUserOwnedReferenced?: boolean;
+  messages: DialogAttachmentMessage[];
+  metadataByFileId: Record<string, any>;
+  metadataReadFailures?: Array<{ fileId: string; error: string }>;
+}): DialogAttachmentPlan {
+  const messageKeys = new Set(
+    args.messages
+      .map((message) => messageDbKey(args.dialogId, message))
+      .filter((key): key is string => Boolean(key))
+  );
+  const referencedFileIds = [
+    ...new Set(args.messages.flatMap((message) => extractFileContentIds(message))),
+  ];
+
+  const candidates = referencedFileIds.map<DialogAttachmentCandidate>((fileId) => {
+    const metadata = args.metadataByFileId[fileId];
+    const ownerType = asText(metadata?.ownerType);
+    const ownerId = asText(metadata?.ownerId);
+    const ownerDbKey = asText(metadata?.ownerDbKey);
+    const fileDbKey = asText(metadata?.dbKey);
+    const source = asText(metadata?.source);
+    const ownedByDialog = ownerType === "dialog" && ownerId === args.dialogId;
+    const ownedByMessage = Boolean(ownerDbKey && messageKeys.has(ownerDbKey));
+    const userOwnedReferenced =
+      args.includeUserOwnedReferenced === true &&
+      Boolean(args.ownerId) &&
+      ownerType === "user" &&
+      ownerId === args.ownerId;
+    const canDelete = Boolean(fileDbKey && (ownedByDialog || ownedByMessage || userOwnedReferenced));
+
+    return {
+      fileId,
+      fileDbKey,
+      size: typeof metadata?.size === "number" && Number.isFinite(metadata.size) ? metadata.size : null,
+      ownerType,
+      ownerId,
+      ownerDbKey,
+      source,
+      status: canDelete ? "delete" : "retain",
+      reason: canDelete
+        ? ownedByMessage
+          ? "ownerDbKey matches a message in this dialog"
+          : userOwnedReferenced
+            ? "user-owned file is referenced by this dialog and explicit referenced-attachment deletion is enabled"
+            : "ownerType/ownerId matches this dialog"
+        : fileDbKey
+          ? "file ownership is not exclusive to this dialog"
+          : "file metadata not found",
+    };
+  });
+
+  const deleteCandidates = candidates.filter((candidate) => candidate.status === "delete");
+  const retainedCandidates = candidates.filter((candidate) => candidate.status === "retain");
+  return {
+    dialogId: args.dialogId,
+    messageCount: args.messages.length,
+    referencedFileIds,
+    candidates,
+    deleteCandidates,
+    retainedCandidates,
+    bytesToDelete: deleteCandidates.reduce((sum, candidate) => sum + (candidate.size ?? 0), 0),
+    metadataReadFailures: args.metadataReadFailures ?? [],
+  };
+}