Account for full context tokens in compression (#908)

* Account for full context tokens in compression * Fix group chat final context updates --------- Co-authored-by: Codex <codex@openai.com>
2026-05-21 19:40:52 +08:00
parent b2ec321990
commit 39ead94352
16 changed files with 730 additions and 35 deletions
@@ -190,9 +190,9 @@ class AgentClient {
        this.socket!.emit('stop_typing', { roomId })
    }

-    emitContextStatus(roomId: string, status: 'compressing' | 'replying' | 'ready'): void {
+    emitContextStatus(roomId: string, status: 'compressing' | 'replying' | 'ready', extra?: Record<string, unknown>): void {
        this.ensureConnected()
-        this.socket!.emit('context_status', { roomId, agentName: this.name, status })
+        this.socket!.emit('context_status', { roomId, agentName: this.name, status, ...extra })
    }

    emitApprovalRequested(roomId: string, payload: Record<string, unknown>): void {
@@ -261,7 +261,7 @@ class AgentClient {
    async replyToMention(
        roomId: string,
        msg: MentionMessage,
-        onStatus?: (status: 'compressing' | 'replying' | 'ready') => void,
+        onStatus?: (status: 'compressing' | 'replying' | 'ready', extra?: Record<string, unknown>) => void,
    ): Promise<void> {
        logger.debug(`[AgentClients] ${this.name} mentioned by ${msg.senderName}: "${msg.content.slice(0, 50)}"`)
        const runMessageId = groupMessageId(roomId, this.profile, this.name)
@@ -278,6 +278,9 @@ class AgentClient {
            // Build compressed context if context engine is available
            let conversationHistory: Array<{ role: string; content: string }> = []
            let instructions: string | undefined
+            const bridge = new AgentBridgeClient()
+            const sessionSeed = String(this.storage?.getRoom?.(roomId)?.sessionSeed || '0')
+            const sessionId = groupBridgeSessionId(roomId, this.profile, this.name, sessionSeed)

            if (this.contextEngine && this.storage) {
                try {
@@ -310,9 +313,32 @@ class AgentClient {
                        currentMessage: msg,
                        compression,
                        profile: this.profile,
+                        contextTokenEstimator: async (history: Array<{ role: 'user' | 'assistant'; content: string }>, estimateInstructions: string) => {
+                            const estimate = await bridge.contextEstimate(
+                                sessionId,
+                                history,
+                                estimateInstructions,
+                                this.profile,
+                            )
+                            logger.info({
+                                roomId,
+                                agentName: this.name,
+                                profile: this.profile,
+                                sessionId,
+                                messages: estimate.message_count,
+                                toolCount: estimate.tool_count,
+                                systemPromptChars: estimate.system_prompt_chars,
+                                fullContextTokens: estimate.token_count,
+                            }, '[GroupChat] full context estimate')
+                            return estimate.token_count
+                        },
                    })
                    conversationHistory = ctx.conversationHistory
                    instructions = ctx.instructions
+                    if (typeof ctx.meta.contextTokenEstimate === 'number' && Number.isFinite(ctx.meta.contextTokenEstimate)) {
+                        this.storage.updateRoomTotalTokens?.(roomId, ctx.meta.contextTokenEstimate)
+                        onStatus?.('replying', { totalTokens: ctx.meta.contextTokenEstimate })
+                    }
                    logger.debug(`[AgentClients] ${this.name}: context built — historyLen=${conversationHistory.length}, meta=%j`, ctx.meta)
                    onStatus?.('replying')
                } catch (err: any) {
@@ -339,9 +365,6 @@ class AgentClient {
            const bridgeInput: AgentBridgeMessage = isContentBlockArray(input)
                ? await convertContentBlocksForAgent(input)
                : input
-            const bridge = new AgentBridgeClient()
-            const sessionSeed = String(this.storage?.getRoom?.(roomId)?.sessionSeed || '0')
-            const sessionId = groupBridgeSessionId(roomId, this.profile, this.name, sessionSeed)
            const flushedAssistantParts = new Set<string>()
            let lastChunk: AgentBridgeOutput | null = null
            const started = await bridge.chat(
@@ -409,6 +432,7 @@ class AgentClient {
                    reasoning_content: reasoningContent || null,
                })
                this.emitMessageStreamEnd(roomId, streamMessageId)
+                await this.refreshRoomFullContextEstimate(roomId, sessionId, bridge, instructions)
                onStatus?.('ready')
                return
            }
@@ -429,6 +453,94 @@ class AgentClient {
        }
    }

+    private async refreshRoomFullContextEstimate(
+        roomId: string,
+        sessionId: string,
+        bridge: AgentBridgeClient,
+        instructions?: string,
+    ): Promise<void> {
+        if (!this.storage?.getMessages) return
+        try {
+            const history = this.buildRoomEstimateHistory(roomId)
+            const estimate = await bridge.contextEstimate(
+                sessionId,
+                history,
+                instructions,
+                this.profile,
+            )
+            const totalTokens = Number(estimate.token_count || 0)
+            if (!Number.isFinite(totalTokens) || totalTokens <= 0) return
+            const rounded = Math.floor(totalTokens)
+            this.storage.updateRoomTotalTokens?.(roomId, rounded)
+            this.emitContextStatus(roomId, 'replying', { totalTokens: rounded })
+            logger.info({
+                roomId,
+                agentName: this.name,
+                profile: this.profile,
+                sessionId,
+                messages: estimate.message_count,
+                toolCount: estimate.tool_count,
+                systemPromptChars: estimate.system_prompt_chars,
+                fullContextTokens: rounded,
+                phase: 'final',
+            }, '[GroupChat] full context estimate')
+        } catch (err: any) {
+            logger.warn(`[GroupChat] failed to refresh final context estimate room=${roomId} agent=${this.name}: ${err.message}`)
+        }
+    }
+
+    private buildRoomEstimateHistory(roomId: string): Array<{ role: 'user' | 'assistant'; content: string }> {
+        const messages = this.storage?.getMessages?.(roomId) || []
+        return messages.map((message: any) => this.mapRoomMessageForEstimate(message))
+    }
+
+    private mapRoomMessageForEstimate(message: any): { role: 'user' | 'assistant'; content: string } {
+        const senderName = String(message?.senderName || 'unknown')
+        const role = String(message?.role || 'user')
+        const isOwnAgent = message?.senderId === this.socket?.id || senderName === this.name
+
+        if (role === 'tool') {
+            const label = message?.tool_name ? `Tool result: ${message.tool_name}` : 'Tool result'
+            return { role: 'user', content: `[${senderName}] [${label}]\n${message?.content || ''}` }
+        }
+
+        if (role === 'assistant' && Array.isArray(message?.tool_calls) && message.tool_calls.length > 0) {
+            const toolsInfo = message.tool_calls.map((toolCall: any) => {
+                const name = toolCall?.function?.name || 'unknown'
+                let args = String(toolCall?.function?.arguments || '{}')
+                if (args.length > 4000) args = `${args.slice(0, 4000)}...`
+                return `[Calling tool: ${name} with arguments: ${args}]`
+            }).join('\n')
+            const content = String(message?.content || '').trim()
+            return {
+                role: isOwnAgent ? 'assistant' : 'user',
+                content: content
+                    ? `${this.formatAttributedContent(senderName, content)}\n${this.formatAttributionPrefix(senderName)}${toolsInfo}`
+                    : `${this.formatAttributionPrefix(senderName)}${toolsInfo}`,
+            }
+        }
+
+        return {
+            role: isOwnAgent ? 'assistant' : 'user',
+            content: this.formatAttributedContent(senderName, String(message?.content || '')),
+        }
+    }
+
+    private formatAttributedContent(senderName: string, content: string): string {
+        return `${this.formatAttributionPrefix(senderName)}${this.stripMentions(content)}`
+    }
+
+    private formatAttributionPrefix(senderName: string): string {
+        return `[${senderName}]: `
+    }
+
+    private stripMentions(content: string): string {
+        return String(content || '')
+            .replace(/@([^\s@]+)/g, '')
+            .replace(/[ \t]{2,}/g, ' ')
+            .replace(/^\s+/, '')
+    }
+
    private async sendAgentErrorMessage(
        roomId: string,
        messageId: string,
@@ -897,8 +1009,8 @@ export class AgentClients {
        }

        this._processingRooms.add(agentKey)
-        const onStatus = (status: 'compressing' | 'replying' | 'ready') => {
-            agent.emitContextStatus(roomId, status)
+        const onStatus = (status: 'compressing' | 'replying' | 'ready', extra?: Record<string, unknown>) => {
+            agent.emitContextStatus(roomId, status, extra)
            logger.debug(`[AgentClients] room ${roomId} agent ${agent.name} status: ${status}`)
        }