send finish event from backend

Author: charleslien

backend/src/websockets/websocket-action.ts

@@ -1,12 +1,12 @@
 import { calculateUsageAndBalance } from '@codebuff/billing'
 import { trackEvent } from '@codebuff/common/analytics'
+import { AnalyticsEvent } from '@codebuff/common/constants/analytics-events'
+import db from '@codebuff/common/db/index'
+import * as schema from '@codebuff/common/db/schema'
 import {
   ASYNC_AGENTS_ENABLED,
   toOptionalFile,
 } from '@codebuff/common/old-constants'
-import { AnalyticsEvent } from '@codebuff/common/constants/analytics-events'
-import db from '@codebuff/common/db/index'
-import * as schema from '@codebuff/common/db/schema'
 import { getErrorObject } from '@codebuff/common/util/error'
 import { ensureEndsWithNewline } from '@codebuff/common/util/file'
 import { generateCompactId } from '@codebuff/common/util/string'
@@ -234,6 +234,16 @@ export const callMainPrompt = async (
 
   const { sessionState, output } = result
 
+  sendAction(ws, {
+    type: 'response-chunk',
+    userInputId: promptId,
+    chunk: {
+      type: 'finish',
+      agentId: sessionState.mainAgentState.agentType ?? undefined,
+      totalCost: sessionState.mainAgentState.creditsUsed,
+    },
+  })
+
   // Send prompt data back
   sendAction(ws, {
     type: 'prompt-response',

npm-app/src/cli.ts

@@ -4,11 +4,11 @@ import { homedir } from 'os'
 import path, { basename, dirname, isAbsolute, parse } from 'path'
 import * as readline from 'readline'
 
+import { AnalyticsEvent } from '@codebuff/common/constants/analytics-events'
 import {
   API_KEY_ENV_VAR,
   ASYNC_AGENTS_ENABLED,
 } from '@codebuff/common/old-constants'
-import { AnalyticsEvent } from '@codebuff/common/constants/analytics-events'
 import {
   getAllAgents,
   getAgentDisplayName,
@@ -109,7 +109,6 @@ import { withHangDetection } from './utils/with-hang-detection'
 import type { CliOptions, GitCommand } from './types'
 import type { ApiKeyType } from '@codebuff/common/api-keys/constants'
 import type { CostMode } from '@codebuff/common/old-constants'
-import type { PrintModeFinish } from '@codebuff/common/types/print-mode'
 import type { ProjectFileContext } from '@codebuff/common/util/file'
 
 // Cache for local agent info to avoid async issues in sync methods
@@ -1372,17 +1371,6 @@ export class CLI {
       .flat()
       .reduce((sum, credits) => sum + credits, 0)
 
-    if (printModeIsEnabled()) {
-      const finishObj: PrintModeFinish = {
-        type: 'finish',
-        totalCost: totalCreditsUsedThisSession,
-      }
-      const agentId = CLI.getInstance().agent
-      if (agentId) {
-        finishObj.agentId = agentId
-      }
-      printModeLog(finishObj)
-    }
     let exitUsageMessage = `${pluralize(totalCreditsUsedThisSession, 'credit')} used this session`
     if (client.usageData.remainingBalance !== null) {
       exitUsageMessage += `, ${client.usageData.remainingBalance.toLocaleString()} credits left.`