cherry-studio/src/main/services/agents/services/claudecode/index.ts

// src/main/services/agents/services/claudecode/index.ts
import { EventEmitter } from 'node:events'
import { createRequire } from 'node:module'

import { McpHttpServerConfig, Options, query, SDKMessage } from '@anthropic-ai/claude-code'
import { loggerService } from '@logger'
import { config as apiConfigService } from '@main/apiServer/config'
import { validateModelId } from '@main/apiServer/utils'

import { GetAgentSessionResponse } from '../..'
import { AgentServiceInterface, AgentStream, AgentStreamEvent } from '../../interfaces/AgentStreamInterface'
import { transformSDKMessageToStreamParts } from './transform'

const require_ = createRequire(import.meta.url)
const logger = loggerService.withContext('ClaudeCodeService')

class ClaudeCodeStream extends EventEmitter implements AgentStream {
  declare emit: (event: 'data', data: AgentStreamEvent) => boolean
  declare on: (event: 'data', listener: (data: AgentStreamEvent) => void) => this
  declare once: (event: 'data', listener: (data: AgentStreamEvent) => void) => this
}

class ClaudeCodeService implements AgentServiceInterface {
  private claudeExecutablePath: string

  constructor() {
    // Resolve Claude Code CLI robustly (works in dev and in asar)
    this.claudeExecutablePath = require_.resolve('@anthropic-ai/claude-code/cli.js')
  }

  async invoke(
    prompt: string,
    session: GetAgentSessionResponse,
    abortController: AbortController,
    lastAgentSessionId?: string
  ): Promise<AgentStream> {
    const aiStream = new ClaudeCodeStream()

    // Validate session accessible paths and make sure it exists as a directory
    const cwd = session.accessible_paths[0]
    if (!cwd) {
      aiStream.emit('data', {
        type: 'error',
        error: new Error('No accessible paths defined for the agent session')
      })
      return aiStream
    }

    // Validate model info
    const modelInfo = await validateModelId(session.model)
    if (!modelInfo.valid) {
      aiStream.emit('data', {
        type: 'error',
        error: new Error(`Invalid model ID '${session.model}': ${JSON.stringify(modelInfo.error)}`)
      })
      return aiStream
    }
    if (modelInfo.provider?.type !== 'anthropic' || modelInfo.provider.apiKey === '') {
      aiStream.emit('data', {
        type: 'error',
        error: new Error(`Invalid provider type '${modelInfo.provider?.type}'. Expected 'anthropic' provider type.`)
      })
      return aiStream
    }

    // TODO: use cherry studio api server config instead of direct provider config to provide more flexibility (e.g. custom headers, proxy, statistics, etc).
    const apiConfig = await apiConfigService.get()
    // process.env.ANTHROPIC_AUTH_TOKEN = apiConfig.apiKey
    // process.env.ANTHROPIC_BASE_URL = `http://${apiConfig.host}:${apiConfig.port}`
    const env = {
      ...process.env,
      ANTHROPIC_API_KEY: apiConfig.apiKey,
      ANTHROPIC_BASE_URL: `http://${apiConfig.host}:${apiConfig.port}/${modelInfo.provider.id}`,
      ELECTRON_RUN_AS_NODE: '1'
    }

    // Build SDK options from parameters
    const options: Options = {
      abortController,
      cwd,
      env,
      model: modelInfo.modelId,
      pathToClaudeCodeExecutable: this.claudeExecutablePath,
      stderr: (chunk: string) => {
        logger.info('claude stderr', { chunk })
      },
      appendSystemPrompt: session.instructions,
      permissionMode: session.configuration?.permission_mode,
      maxTurns: session.configuration?.max_turns
    }

    if (session.accessible_paths.length > 1) {
      options.additionalDirectories = session.accessible_paths.slice(1)
    }

    if (session.mcps && session.mcps.length > 0) {
      // mcp configs
      const mcpList: Record<string, McpHttpServerConfig> = {}
      for (const mcpId of session.mcps) {
        mcpList[mcpId] = {
          type: 'http',
          url: `http://${apiConfig.host}:${apiConfig.port}/v1/mcps/${mcpId}/mcp`,
          headers: {
            Authorization: `Bearer ${apiConfig.apiKey}`
          }
        }
      }
      options.mcpServers = mcpList
      options.strictMcpConfig = true
    }

    if (lastAgentSessionId) {
      options.resume = lastAgentSessionId
    }

    logger.silly('Starting Claude Code SDK query', {
      prompt,
      options
    })

    // Start async processing
    this.processSDKQuery(prompt, options, aiStream)

    return aiStream
  }

  private async *userMessages(prompt: string) {
    {
      yield {
        type: 'user' as const,
        parent_tool_use_id: null,
        session_id: '',
        message: {
          role: 'user' as const,
          content: prompt
        }
      }
    }
  }

  /**
   * Process SDK query and emit stream events
   */
  private async processSDKQuery(prompt: string, options: Options, stream: ClaudeCodeStream): Promise<void> {
    const jsonOutput: SDKMessage[] = []
    let hasCompleted = false
    const startTime = Date.now()

    try {
      // Process streaming responses using SDK query
      for await (const message of query({
        prompt: this.userMessages(prompt),
        options
      })) {
        if (hasCompleted) break

        jsonOutput.push(message)
        logger.silly('claude response', { message })
        if (message.type === 'assistant' || message.type === 'user') {
          logger.silly('message content', {
            message: JSON.stringify({ role: message.message.role, content: message.message.content })
          })
        }

        // Transform SDKMessage to UIMessageChunks
        const chunks = transformSDKMessageToStreamParts(message)
        for (const chunk of chunks) {
          stream.emit('data', {
            type: 'chunk',
            chunk
          })
        }
      }

      // Successfully completed
      hasCompleted = true
      const duration = Date.now() - startTime

      logger.debug('SDK query completed successfully', {
        duration,
        messageCount: jsonOutput.length
      })

      // Emit completion event
      stream.emit('data', {
        type: 'complete'
      })
    } catch (error) {
      if (hasCompleted) return
      hasCompleted = true

      const duration = Date.now() - startTime

      // Check if this is an abort error
      const errorObj = error as any
      const isAborted =
        errorObj?.name === 'AbortError' ||
        errorObj?.message?.includes('aborted') ||
        options.abortController?.signal.aborted

      if (isAborted) {
        logger.info('SDK query aborted by client disconnect', { duration })
        // Simply cleanup and return - don't emit error events
        stream.emit('data', {
          type: 'cancelled',
          error: new Error('Request aborted by client')
        })
        return
      }

      // Original error handling for non-abort errors
      logger.error('SDK query error:', {
        error: errorObj instanceof Error ? errorObj.message : String(errorObj),
        duration,
        messageCount: jsonOutput.length
      })

      // Emit error event
      stream.emit('data', {
        type: 'error',
        error: errorObj instanceof Error ? errorObj : new Error(String(errorObj))
      })
    }
  }
}

export default ClaudeCodeService