Просмотр исходного кода

fix(bot): 用量按文字自计,环境手册按需带上

回复结束后不再向 Hermes 查用量。消息和上文都没提到环境时,只留开放名单。
Samien 1 месяц назад
Родитель
Сommit
0ba7963749
1 измененных файлов с 46 добавлено и 43 удалено
  1. 46 43
      src/server/workspaceRuntime.ts

+ 46 - 43
src/server/workspaceRuntime.ts

@@ -19,12 +19,26 @@ import type {
   WorkspaceAgentUsageSummary,
 } from '../shared/workspace.js'
 import { personaSection } from '../shared/workspace.js'
-import { deniedHostToolsets, readHostTools } from './hostToolSettings.js'
+import { readHostTools } from './hostToolSettings.js'
 function globalComputers(home: string) {
   try { return readHostTools(home) }
   catch { return { denyServerTools: true, scriptMachine: false, serverComputer: false, vm: false, cloud: false } }
 }
 
+function estimateTokens(text: string): number {
+  let units = 0
+  for (const ch of text) units += (ch.codePointAt(0) ?? 0) > 0xff ? 1 : 0.25
+  return Math.ceil(units)
+}
+
+function environmentMentions(text: string, hostNames: string[]) {
+  const cloud = /云虚拟机|cloud_computer|Grok Bot 虚拟机/.test(text)
+  const vm = /虚拟环境/.test(text) || /computer_(desktop|shell|action|read_file|write_file)/.test(text)
+  const named = hostNames.some(name => name.length >= 2 && name !== '本机' && text.includes(name))
+  const desktop = /脚本机|本机|desktop_/.test(text) || named
+  return { desktop, vm, cloud }
+}
+
 import { WorkspaceScheduler, type Work, NO_REPLY, HOST_FALLBACK, WORKSPACE_CONCURRENCY_LIMIT, RESOURCE_WAIT_CODES } from './workspaceScheduler.js'
 import type { WorkspacePlugins, OpenBotPlugins } from './botPlugins/workspacePlugins.js'
 import { mentionedAgents } from './workspaceMentions.js'
@@ -359,6 +373,7 @@ export class WorkspaceRuntime extends WorkspaceScheduler {
       lastFlush = 0
     let flushTimer: ReturnType<typeof setTimeout> | undefined
     let usageCompletion: Promise<void> | undefined
+    let submittedText = ''
     const toolController = new AbortController()
     let toolLease: WorkspaceToolLease | undefined
     let pluginLease: OpenBotPlugins | undefined
@@ -525,42 +540,8 @@ export class WorkspaceRuntime extends WorkspaceScheduler {
         } else if (!completing) {
           completedEvidence = true
           completing = true
-          let timeout: ReturnType<typeof setTimeout>
-          usageCompletion = Promise.race([
-            gateway.rpc('session.usage', { session_id: runtimeId }),
-            new Promise<undefined>((resolve) => {
-              timeout = setTimeout(() => resolve(undefined), 1000)
-            }),
-          ])
-            .then((usage) => {
-              if (usage && !this.closing && this.store.get<{ taskId: string }>(owner, 'binding', key)?.taskId === run.id) {
-                this.recordAgentUsage(owner, key, agent.id, usage)
-                const used = usage.context_used ?? usage.used_tokens,
-                  limit = usage.context_max ?? usage.context_limit
-                const contextKey = run.conversationTaskId ?? c.id,
-                  context = {
-                  ...usage,
-                  conversationTaskId: run.conversationTaskId,
-                  usedTokens: used,
-                  limitTokens: limit,
-                  percent:
-                    typeof used === 'number' && typeof limit === 'number' && limit > 0
-                      ? Math.round((used / limit) * 1000) / 10
-                      : undefined,
-                  }
-                this.store.put(owner, 'context', contextKey, context)
-                this.store.event(
-                  owner,
-                  'context.changed',
-                  context,
-                  c.id,
-                )
-              }
-            })
-            .catch(() => {})
-            .finally(() => {
-              clearTimeout(timeout)
-            })
+          const outputText = [resultMessage.content, resultMessage.reasoning, resultMessage.tools.length ? JSON.stringify(resultMessage.tools) : ''].filter(Boolean).join('\n')
+          this.recordTurnTokens(owner, agent.id, estimateTokens(submittedText), estimateTokens(outputText))
           finish()
         }
         return
@@ -785,10 +766,15 @@ export class WorkspaceRuntime extends WorkspaceScheduler {
         const goal = run.conversationTaskId ? this.store.get<import('../shared/agentTasks.js').AgentGoal>(owner, 'goal', run.conversationTaskId) : undefined
         const assignmentId = this.store.require<Run>(owner, 'run', run.runId).assignmentId
         const globals = globalComputers(this.store.home)
-        const deniedToolsets = deniedHostToolsets(this.store.home)
         const eligible = !agent.remoteAgentId && !agent.temporaryGoalId && !agent.archived
         const open = [eligible && (globals.scriptMachine || globals.serverComputer) ? '脚本机' : '', eligible && globals.vm && agent.envs ? '虚拟环境' : '', eligible && globals.cloud ? '云虚拟机' : ''].filter(Boolean)
-        const envDirective = `同一条对话里,只使用 Bot 设置里全局开放的环境:${open.join('、') || '无'}。脚本机里,本机是夭夭正在运行的这台电脑,省略 host 或写「本机」就看它,其他脚本机按名称点名。虚拟环境是服务端上的隔离虚拟机,云虚拟机是 Grok Bot 虚拟机。没开放的不要用。${deniedToolsets.length ? '服务端本机的终端和文件已全局禁用,不要用服务端本机顶上。' : '全局设置已允许使用服务端本机的终端和文件。'}`
+        const hostNames = this.desktopEnvironments?.hostStates(owner).map(host => host.name).filter(Boolean) ?? []
+        const prior = this.store.messages(owner, c.id, Number.MAX_SAFE_INTEGER, Number.MAX_SAFE_INTEGER, true, run.conversationTaskId)
+          .filter(message => message.id !== resultMessage.id && message.seq <= run.triggerSeq)
+          .map(message => `${message.content}\n${message.tools.map(tool => String(tool.name ?? tool.tool_name ?? '')).join(' ')}`)
+          .join('\n')
+        const hit = environmentMentions(`${trigger.content}\n${prior}`, hostNames)
+        const envDirective = `同一条对话里,只使用 Bot 设置里全局开放的环境:${open.join('、') || '无'}。没开放的不要用。`
         const rules = [
           envDirective,
           `你是 ${agent.name}。${agent.description ? `描述:${agent.description}\n\n` : ''}${personaSection(agent) ? `用户为它设定的工作规范:\n${personaSection(agent)}\n\n` : ''}以下是用户为这个独立机器人配置的角色与规则(版本 ${agent.revision}):\n${agent.remoteAgentId ? '配置由远端机器人管理。' : agent.instructions}`,
@@ -816,9 +802,9 @@ export class WorkspaceRuntime extends WorkspaceScheduler {
           root.projectId ? `当前项目 ID:${root.projectId}。只能使用当前项目记忆。` : '',
           `本次来源消息 ID:${run.messageId};当前会话 ID:${c.id}。`,
           plugins ? '本轮已挂载用户为当前 Bot 授权的插件工具,工具名以 plugin_ 开头,说明中包含实际服务和操作。仅按用户当前任务使用;连接或重新授权应用请让用户打开 Bot 模式的工具 → 已连接应用。不要索取 API Key 或在回复中展示凭据。' : '',
-          desktop ? DESKTOP_ENVIRONMENT_RULES+(this.desktopEnvironments?.onlineHostsLine(owner)??'') : '',
-          cloud ? grokComputerRules() : '',
-          dispatchedVm ? VM_COMPUTER_RULES : '',
+          desktop && hit.desktop ? DESKTOP_ENVIRONMENT_RULES+(this.desktopEnvironments?.onlineHostsLine(owner)??'') : '',
+          cloud && hit.cloud ? grokComputerRules() : '',
+          dispatchedVm && hit.vm ? VM_COMPUTER_RULES : '',
           `[yaoyao-run:${run.runId}:${resultMessage.id}]`,
         ]
           .filter(Boolean)
@@ -838,10 +824,11 @@ export class WorkspaceRuntime extends WorkspaceScheduler {
         this.saveWork(owner, admission)
         submitted = true
         try {
+          submittedText = `${rules}\n\n${text}\n${attachmentRefs.join('\n')}`
           await gateway.rpc('prompt.submit', {
             session_id: runtimeId,
             ...(target.runner?.computer?{workMarker:`[yaoyao-run:${run.runId}:${resultMessage.id}]`}:{}),
-            text: `${rules}\n\n${text}\n${attachmentRefs.join('\n')}`,
+            text: submittedText,
           })
           this.store.remove(owner,'binding-reset',key)
           binding.contextSeq = Math.max(binding.contextSeq ?? 0, admission.contextThroughSeq)
@@ -983,6 +970,22 @@ export class WorkspaceRuntime extends WorkspaceScheduler {
       if (task.conversationId === conversationId) this.tasks.cancel(owner, task.id)
     await Promise.all([super.stopConversation(owner, conversationId), this.tasks.cancelOrigin(owner, conversationId)])
   }
+  private recordTurnTokens(owner: string, agentId: string, input: number, output: number): void {
+    if (!input && !output) return
+    this.store.atomic(() => {
+      const now = new Date(),
+        date = `${now.getFullYear()}-${String(now.getMonth() + 1).padStart(2, '0')}-${String(now.getDate()).padStart(2, '0')}`,
+        id = `${agentId}:${date}`
+      const day = this.store.get<WorkspaceAgentUsageSummary & { agentId: string; date: string; updatedAt: number }>(owner, 'agent-token-day', id)
+        ?? { agentId, date, input: 0, output: 0, total: 0, updatedAt: 0 }
+      day.input += input
+      day.output += output
+      day.total += input + output
+      day.updatedAt = now.getTime()
+      this.store.put(owner, 'agent-token-day', id, day)
+    })
+  }
+
   /**
    * Hermes reports cumulative per-session token counters (input/output/total) alongside
    * context occupancy. Attribute the growth since the last snapshot to the agent's daily