返回 AiToEarn
agent-runtime.service.ts
根目录 / project / aitoearn-backend / apps / aitoearn-ai / src / core / agent / services / agent-runtime.service.ts
1 import { spawn } from 'node:child_process'
2 import * as fs from 'node:fs'
3 import { join } from 'node:path'
4 import {
5 AbortError,
6 createSdkMcpServer,
7 McpServerConfig,
8 Options,
9 OutputFormat,
10 query,
11 SDKMessage,
12 SpawnedProcess,
13 SpawnOptions,
14 } from '@anthropic-ai/claude-agent-sdk'
15 import { ContentBlockParam } from '@anthropic-ai/sdk/resources'
16 import { Injectable, Logger, Optional } from '@nestjs/common'
17 import { StorageProvider } from '@yikart/assets'
18 import {
19 AppException,
20 getExceptionPayload,
21 ResponseCode,
22 UserType,
23 WithLoggerContext,
24 } from '@yikart/common'
25 import {
26 AiLogChannel,
27 AiLogRepository,
28 AiLogStatus,
29 AiLogType,
30 ContentGenerationTask,
31 ContentGenerationTaskRepository,
32 ContentGenerationTaskStatus,
33 } from '@yikart/mongodb'
34 import { Redlock } from '@yikart/redlock'
35 import { Request, Response } from 'express'
36 import { firstValueFrom, from, interval, merge, Observable, of, throwError } from 'rxjs'
37 import { catchError, concatMap, filter, finalize, first, map, mergeMap, share, skip, takeUntil, tap, timeout } from 'rxjs/operators'
38 import { z } from 'zod'
39 import { RedlockKey } from '../../../common/enums'
40 import { config } from '../../../config'
41 import { AiAvailabilityService } from '../../ai-availability'
42 import { RelayMediaResolverService } from '../../ai/relay-media'
43 import { ChannelsToolName, CLAUDE_CODE_ROUTER_PROVIDER_NAME, McpServerName, POLLING_TASK_AGENT_PROMPT, SKILL_ANALYZER_AGENT_PROMPT, SYSTEM_PROMPT } from '../agent.constants'
44 import { ContentBlock, CreateContentGenerationTaskDto } from '../agent.dto'
45 import { enhancePrompt, filterHeaders, normalizePrompt, sanitizeMessage, shouldFilterSyntheticMessage } from '../agent.utils'
46 import {
47 AgentMessageType,
48 AgentMessageVoHelper,
49 ContentGenerationTaskAgentChunkVo,
50 ContentGenerationTaskChunkVo,
51 ContentGenerationTaskErrorChunkVo,
52 ContentGenerationTaskInitChunkVo,
53 ContentGenerationTaskKeepAliveChunkVo,
54 ContentGenerationTaskResultSchema,
55 ContentGenerationTaskResultUnionSchema,
56 } from '../agent.vo'
57 import { ImageEditMcp, ImageEditToolName } from '../mcp/image-edit.mcp'
58 import { successResult, wrapTool } from '../mcp/mcp.utils'
59 import { MediaMcp, MediaToolName } from '../mcp/media.mcp'
60 import { SubtitleMcp, SubtitleToolName } from '../mcp/subtitle.mcp'
61 import { UtilMcp, UtilToolName } from '../mcp/util.mcp'
62 import { VideoUtilsMcp, VideoUtilsToolName } from '../mcp/video-utils.mcp'
63 import { AideoMcp, AideoToolName } from '../mcp/volcengine/aideo.mcp'
64 import { DramaRecapMcp, DramaRecapToolName } from '../mcp/volcengine/drama-recap.mcp'
65 import { StyleTransferMcp, StyleTransferToolName } from '../mcp/volcengine/style-transfer.mcp'
66 import { VideoEditMcp, VideoEditToolName } from '../mcp/volcengine/video-edit.mcp'
67
68 export interface ClaudeQueryOptions {
69 includePartialMessages?: boolean
70 sessionId?: string
71 model?: string
72 taskId?: string
73 outputFormat?: OutputFormat
74 persistSession?: boolean
75 availabilityOperation?: string
76 }
77
78 type TaskResult = z.infer<typeof ContentGenerationTaskResultUnionSchema>
79
80 export interface RuntimeRunningTaskInfo {
81 taskId: string
82 userId: string
83 abortController: AbortController
84 claudeCodeProcess?: ReturnType<typeof spawn>
85 completionPromise: Promise<void>
86 sessionId?: string
87 }
88
89 @Injectable()
90 export class AgentRuntimeService {
91 private readonly logger = new Logger(AgentRuntimeService.name)
92 private readonly sessionDir = join(process.cwd(), '.claude-session')
93 private readonly projectDir = this.sessionDir.replace(/\//g, '-').replace(/\./g, '-')
94 private readonly runningTasks = new Map<string, RuntimeRunningTaskInfo>()
95
96 constructor(
97 private readonly mediaMcp: MediaMcp,
98 private readonly aideoMcp: AideoMcp,
99 private readonly utilMcp: UtilMcp,
100 private readonly videoEditMcp: VideoEditMcp,
101 private readonly aiLogRepo: AiLogRepository,
102 private readonly contentGenerateRepository: ContentGenerationTaskRepository,
103 private readonly aiAvailability: AiAvailabilityService,
104 private readonly dramaRecapMcp: DramaRecapMcp,
105 private readonly videoUtilsMcp: VideoUtilsMcp,
106 private readonly styleTransferMcp: StyleTransferMcp,
107 private readonly imageEditMcp: ImageEditMcp,
108 private readonly subtitleMcp: SubtitleMcp,
109 private readonly storageProvider: StorageProvider,
110 @Optional() private readonly relayMediaResolver?: RelayMediaResolverService,
111 ) {}
112
113 private getTaskCwd(taskId: string): string {
114 return join(this.sessionDir, 'tasks', taskId)
115 }
116
117 private getTaskProjectDir(taskId: string): string {
118 const taskCwd = this.getTaskCwd(taskId)
119 return taskCwd.replace(/\//g, '-').replace(/\./g, '-')
120 }
121
122 private getMessageType(chunk: SDKMessage): AgentMessageType.Assistant | AgentMessageType.User | AgentMessageType.Result | AgentMessageType.System | AgentMessageType.StreamEvent | AgentMessageType.ToolProgress | AgentMessageType.AuthStatus {
123 switch (chunk.type) {
124 case 'assistant':
125 return AgentMessageType.Assistant
126 case 'user':
127 return AgentMessageType.User
128 case 'result':
129 return AgentMessageType.Result
130 case 'system':
131 return AgentMessageType.System
132 case 'stream_event':
133 return AgentMessageType.StreamEvent
134 case 'tool_progress':
135 return AgentMessageType.ToolProgress
136 case 'auth_status':
137 return AgentMessageType.AuthStatus
138 default:
139 return AgentMessageType.System
140 }
141 }
142
143 private generateAllowedTools(mcpServers: Record<string, McpServerConfig>): string[] {
144 const allowedTools: string[] = []
145
146 for (const serverName of Object.keys(mcpServers)) {
147 let toolNames: string[] = []
148
149 switch (serverName) {
150 case McpServerName.MediaGeneration:
151 toolNames = Object.values(MediaToolName)
152 break
153 case McpServerName.Util:
154 toolNames = [UtilToolName.GetCurrentTime, UtilToolName.Wait]
155 break
156 case McpServerName.SessionTools:
157 toolNames = [UtilToolName.OutputTaskResult, UtilToolName.SetTitle]
158 break
159 case McpServerName.Aideo:
160 toolNames = Object.values(AideoToolName)
161 break
162 case McpServerName.VideoEdit:
163 toolNames = Object.values(VideoEditToolName)
164 break
165 case McpServerName.DramaRecap:
166 toolNames = Object.values(DramaRecapToolName)
167 break
168 case McpServerName.VideoUtils:
169 toolNames = Object.values(VideoUtilsToolName)
170 break
171 case McpServerName.StyleTransfer:
172 toolNames = Object.values(StyleTransferToolName)
173 break
174 case McpServerName.ImageEdit:
175 toolNames = Object.values(ImageEditToolName)
176 break
177 case McpServerName.Subtitle:
178 toolNames = Object.values(SubtitleToolName)
179 break
180 case McpServerName.Channels:
181 toolNames = Object.values(ChannelsToolName)
182 break
183 default:
184 continue
185 }
186
187 for (const toolName of toolNames) {
188 allowedTools.push(`mcp__${serverName}__${toolName}`)
189 }
190 }
191
192 return allowedTools
193 }
194
195 public claudeQuery(
196 systemPromptContent: ContentBlockParam[],
197 enhancedContent: ContentBlockParam[],
198 abortController: AbortController,
199 options: ClaudeQueryOptions,
200 mcpServers?: Record<string, McpServerConfig>,
201 spawnClaudeCodeProcess?: (options: SpawnOptions) => SpawnedProcess,
202 ) {
203 const taskCwd = options.taskId
204 ? this.getTaskCwd(options.taskId)
205 : this.sessionDir
206
207 if (options.taskId && !fs.existsSync(taskCwd)) {
208 fs.mkdirSync(taskCwd, { recursive: true })
209 }
210
211 const queryModel = options.model || config.agent.defaultModel
212 const ccrSubAgentModelPrompt = `<CCR-SUBAGENT-MODEL>${CLAUDE_CODE_ROUTER_PROVIDER_NAME},${config.agent.backgroundModel}</CCR-SUBAGENT-MODEL>`
213
214 const queryOptions: Options = {
215 permissionMode: 'default',
216 includePartialMessages: options.includePartialMessages ?? false,
217 model: queryModel,
218 cwd: taskCwd,
219 settingSources: ['project', 'user'],
220 env: {
221 ...process.env,
222 DEBUG_CLAUDE_AGENT_SDK: '1',
223 ANTHROPIC_AUTH_TOKEN: 'ccr',
224 ANTHROPIC_API_KEY: '',
225 ANTHROPIC_BASE_URL: 'http://127.0.0.1:3456',
226 NO_PROXY: '127.0.0.1',
227 DISABLE_TELEMETRY: 'true',
228 DISABLE_COST_WARNINGS: 'true',
229 HOME: this.sessionDir,
230 },
231 mcpServers,
232 allowedTools: this.generateAllowedTools(mcpServers ?? {}),
233 tools: [
234 'Task',
235 'TaskOutput',
236 'Read',
237 'WebFetch',
238 'TodoWrite',
239 'TaskStop',
240 'Skill',
241 'ListMcpResourcesTool',
242 'ReadMcpResourceTool',
243 ],
244 agents: {
245 'polling-task': {
246 description: 'AI task polling specialist for monitoring asynchronous video/media generation task status. Use when polling task status, checking completion, or handling timeouts.',
247 model: 'inherit',
248 mcpServers: [
249 {
250 ...(mcpServers?.[McpServerName.MediaGeneration] && { [McpServerName.MediaGeneration]: mcpServers[McpServerName.MediaGeneration] }),
251 ...(mcpServers?.[McpServerName.Util] && { [McpServerName.Util]: mcpServers[McpServerName.Util] }),
252 ...(mcpServers?.[McpServerName.Aideo] && { [McpServerName.Aideo]: mcpServers[McpServerName.Aideo] }),
253 ...(mcpServers?.[McpServerName.VideoEdit] && { [McpServerName.VideoEdit]: mcpServers[McpServerName.VideoEdit] }),
254 ...(mcpServers?.[McpServerName.DramaRecap] && { [McpServerName.DramaRecap]: mcpServers[McpServerName.DramaRecap] }),
255 ...(mcpServers?.[McpServerName.VideoUtils] && { [McpServerName.VideoUtils]: mcpServers[McpServerName.VideoUtils] }),
256 ...(mcpServers?.[McpServerName.StyleTransfer] && { [McpServerName.StyleTransfer]: mcpServers[McpServerName.StyleTransfer] }),
257 },
258 ],
259 tools: [
260 'Task',
261 'TaskOutput',
262 'Read',
263 'NotebookEdit',
264 'WebFetch',
265 'TodoWrite',
266 'TaskStop',
267 'Skill',
268 'ToolSearch',
269 'ListMcpResourcesTool',
270 'ReadMcpResourceTool',
271 `mcp__${McpServerName.MediaGeneration}__getVideoStatus`,
272 `mcp__${McpServerName.MediaGeneration}__getSoraCharacter`,
273 `mcp__${McpServerName.Aideo}__getAideoTaskStatus`,
274 `mcp__${McpServerName.VideoEdit}__getVideoEditTaskStatus`,
275 `mcp__${McpServerName.StyleTransfer}__getVideoStyleTransferStatus`,
276 `mcp__${McpServerName.DramaRecap}__getDramaRecapTaskStatus`,
277 `mcp__${McpServerName.Util}__wait`,
278 `mcp__${McpServerName.Util}__getCurrentTime`,
279 ],
280 prompt: `${ccrSubAgentModelPrompt}\n${POLLING_TASK_AGENT_PROMPT}`,
281 skills: [],
282 },
283 'skill-analyzer': {
284 description: 'Analyzes user requests to determine which skills are needed for content generation tasks. Call this agent FIRST before any generation to identify required skills.',
285 model: 'inherit',
286 mcpServers: [],
287 tools: [],
288 prompt: `${ccrSubAgentModelPrompt}\n${SKILL_ANALYZER_AGENT_PROMPT}`,
289 skills: [],
290 },
291 },
292 canUseTool: async (name, input, options) => {
293 this.logger.debug({ options, name, input }, 'Received tool request')
294 return { behavior: 'allow', updatedInput: input }
295 },
296 hooks: {
297 PostToolUse: [
298 {
299 hooks: [
300 async (input) => {
301 if (input.hook_event_name === 'PostToolUse') {
302 this.logger.debug({ input }, 'Received PostToolUse hook')
303 const response = input.tool_response
304 if (Array.isArray(response)) {
305 const mappedResponse = response.map<ContentBlockParam>((block) => {
306 if (block.type === 'text' && typeof block.text === 'string' && block.text.startsWith('[Resource link: Image')) {
307 const url = block.text.split('] ')[1]
308 return {
309 type: 'image',
310 source: {
311 type: 'url',
312 url,
313 },
314 }
315 }
316 return block
317 })
318 const updatedResponse = await this.resolveRelayJson(mappedResponse)
319 return {
320 hookSpecificOutput: {
321 hookEventName: 'PostToolUse',
322 updatedMCPToolOutput: updatedResponse,
323 },
324 }
325 }
326 }
327 return {}
328 },
329 ],
330 },
331 ],
332 },
333 outputFormat: options.outputFormat,
334 persistSession: options.persistSession,
335 spawnClaudeCodeProcess,
336 abortController,
337 }
338
339 if (options.sessionId) {
340 queryOptions.resume = options.sessionId
341 }
342
343 const content = [
344 ...systemPromptContent,
345 ...enhancedContent,
346 ]
347
348 this.logger.debug({ content }, 'Content')
349
350 const req = query({
351 prompt: (async function* () {
352 yield {
353 session_id: options.sessionId || '',
354 type: 'user',
355 message: {
356 role: 'user',
357 content,
358 },
359 parent_tool_use_id: null,
360 }
361 })(),
362 options: queryOptions,
363 })
364
365 const startedAt = Date.now()
366 const aiAvailability = this.aiAvailability
367 const availabilityContext = {
368 provider: 'agent',
369 operation: options.availabilityOperation ?? 'claudeQuery',
370 model: queryModel,
371 module: 'agent',
372 }
373
374 return (async function* () {
375 try {
376 for await (const chunk of req) {
377 yield chunk
378 }
379
380 await aiAvailability.recordSuccess(availabilityContext, Date.now() - startedAt)
381 }
382 catch (error) {
383 await aiAvailability.recordFailure(availabilityContext, error, Date.now() - startedAt)
384 throw error
385 }
386 })()
387 }
388
389 // @Cron(CronExpression.EVERY_HOUR)
390 @Redlock(RedlockKey.AgentHealthCheck, 600, { throwOnFailure: false })
391 @WithLoggerContext()
392 async runHealthCheck(): Promise<void> {
393 const abortController = new AbortController()
394 const res = this.claudeQuery(
395 [{ type: 'text', text: '' }],
396 [{ type: 'text', text: 'Hello, how are you?' }],
397 abortController,
398 {
399 includePartialMessages: true,
400 availabilityOperation: 'healthCheck',
401 },
402 )
403
404 const timeoutMs = 5 * 60 * 1000
405 const chunks: unknown[] = []
406
407 await firstValueFrom(
408 from(res).pipe(
409 timeout({
410 each: timeoutMs,
411 with: () => throwError(() => new Error('Timeout')),
412 }),
413 tap(chunk => chunks.push(chunk)),
414 first(chunk => chunk.type === 'stream_event'),
415 map(() => true),
416 catchError((error) => {
417 this.logger.error({ error, chunks }, '健康检查失败')
418 return of(false)
419 }),
420 ),
421 )
422 abortController.abort()
423 }
424
425 createContentGenerationTask(params: {
426 userId: string
427 userType: UserType
428 dto: CreateContentGenerationTaskDto
429 abortController: AbortController
430 req: Request
431 res: Response
432 }): Observable<ContentGenerationTaskChunkVo> {
433 const {
434 userId,
435 userType,
436 dto,
437 abortController,
438 req,
439 res,
440 } = params
441
442 return from(this.initializeTask(userId, userType, dto, abortController, req)).pipe(
443 mergeMap(({ taskId, sessionId: initialSessionId, abortController, mcpServers }) => {
444 let sessionId = initialSessionId
445 let completionResolver: () => void
446 const completionPromise = new Promise<void>((resolve) => {
447 completionResolver = resolve
448 })
449
450 const taskInfo: RuntimeRunningTaskInfo = {
451 taskId,
452 userId,
453 abortController,
454 claudeCodeProcess: undefined,
455 completionPromise,
456 sessionId,
457 }
458 this.runningTasks.set(taskId, taskInfo)
459
460 let completeTitleUpdate: (() => void) | undefined
461
462 abortController.signal.addEventListener('abort', () => {
463 this.logger.warn({
464 taskId,
465 sessionId: taskInfo.sessionId,
466 userId,
467 }, `Task ${taskId} was aborted`)
468 void this.contentGenerateRepository.updateStatus(taskId, ContentGenerationTaskStatus.Aborted)
469 })
470
471 return from(this.prepareAgentPrompt(dto)).pipe(
472 mergeMap(({ normalizedContent, enhancedContent, systemPromptContent }) => {
473 const userMessage = {
474 type: 'user',
475 content: normalizedContent,
476 }
477
478 void this.contentGenerateRepository.updateMessage(taskId, userMessage)
479
480 let taskResult: TaskResult | undefined
481
482 const outputTaskResultTool = wrapTool(
483 this.logger,
484 UtilToolName.OutputTaskResult,
485 'output task result JSON, the result will be included in the completion message.',
486 ContentGenerationTaskResultSchema.shape,
487 async (args) => {
488 taskResult = args.result
489 return successResult('Task result submitted successfully')
490 },
491 this.aiAvailability,
492 )
493
494 const [setTitleTool, titleUpdate$, completeTitleUpdateFn] = this.utilMcp.createSetTitleTool(taskId)
495 completeTitleUpdate = completeTitleUpdateFn
496
497 const sessionToolsMcp = createSdkMcpServer({
498 name: McpServerName.SessionTools,
499 version: '1.0.0',
500 tools: [outputTaskResultTool, setTitleTool],
501 })
502
503 const queryGenerator = this.claudeQuery(
504 systemPromptContent,
505 enhancedContent,
506 abortController,
507 {
508 includePartialMessages: dto.includePartialMessages ?? false,
509 sessionId,
510 model: dto.model,
511 taskId,
512 },
513 {
514 [McpServerName.SessionTools]: sessionToolsMcp,
515 ...mcpServers,
516 },
517 (options: SpawnOptions): SpawnedProcess => {
518 const childProcess = spawn(options.command, options.args, {
519 cwd: options.cwd,
520 env: options.env,
521 signal: abortController.signal,
522 stdio: ['pipe', 'pipe', 'pipe'],
523 windowsHide: true,
524 })
525 childProcess.stderr.on('data', (data) => {
526 this.logger.debug({ taskId, sessionId }, `Received Claude DEBUG Log: ${data}`)
527 })
528 childProcess.on('exit', (code) => {
529 this.logger.debug({ taskId, sessionId, code }, `Claude Code process exited with code ${code}`)
530 })
531
532 taskInfo.claudeCodeProcess = childProcess
533
534 return childProcess
535 },
536 )
537
538 const messageStream$ = this.createMessageStream(queryGenerator).pipe(share())
539
540 const firstMessage$ = messageStream$.pipe(
541 first(),
542 concatMap(async (chunk) => {
543 const extractedSessionId = 'session_id' in chunk ? chunk.session_id : undefined
544
545 this.logger.debug({ taskId, chunk }, `Task ${taskId} first message received`)
546
547 if (extractedSessionId) {
548 await this.contentGenerateRepository.updateById(taskId, {
549 sessionId: extractedSessionId,
550 })
551 sessionId = extractedSessionId
552 taskInfo.sessionId = extractedSessionId
553
554 taskInfo.completionPromise
555 .then(() => {
556 this.logger.log({ taskId: taskInfo.taskId, sessionId }, `Task ${taskInfo.taskId} completed during shutdown`)
557 })
558 this.logger.debug({ taskId, sessionId: extractedSessionId }, `Task ${taskId} sessionId saved: ${extractedSessionId}`)
559 }
560
561 return ContentGenerationTaskInitChunkVo.create({
562 type: AgentMessageType.Init,
563 taskId,
564 messages: undefined,
565 })
566 }),
567 )
568
569 const restMessages$ = messageStream$.pipe(
570 tap((chunk) => {
571 if (chunk.type !== 'stream_event') {
572 this.logger.debug({ taskId, sessionId, chunk: sanitizeMessage(chunk) }, `Received message for task ${taskId}`)
573 }
574 }),
575 skip(1),
576 filter(chunk => !shouldFilterSyntheticMessage(chunk)),
577 concatMap(async (chunk) => {
578 return this.transformMessage(chunk, taskId, userId, userType, taskResult, sessionId)
579 }),
580 )
581
582 const completionSignal$ = new Observable<null>((subscriber) => {
583 messageStream$.subscribe({
584 complete: () => subscriber.next(null),
585 error: () => subscriber.next(null),
586 })
587 })
588
589 const keepAlive$ = interval(5000).pipe(
590 map(() => ContentGenerationTaskKeepAliveChunkVo.create({
591 type: AgentMessageType.KeepAlive,
592 })),
593 takeUntil(completionSignal$),
594 )
595
596 const titleUpdateStream$ = titleUpdate$.pipe(
597 takeUntil(completionSignal$),
598 )
599
600 return merge(
601 firstMessage$,
602 restMessages$,
603 keepAlive$,
604 titleUpdateStream$,
605 )
606 }),
607 catchError((error) => {
608 this.logger.error(Object.assign(error, { taskId, sessionId }), `Error in task ${taskId}, session ${sessionId}`)
609 if (error instanceof AbortError) {
610 this.logger.debug({ taskId, sessionId }, `Task ${taskId} aborted`)
611 return of()
612 }
613
614 void this.contentGenerateRepository.updateStatus(taskId, ContentGenerationTaskStatus.Error)
615
616 const payload = getExceptionPayload(error)
617 const errorChunk = ContentGenerationTaskErrorChunkVo.create({
618 type: AgentMessageType.Error,
619 ...payload,
620 timestamp: Date.now(),
621 })
622
623 void this.contentGenerateRepository.updateMessage(taskId, errorChunk)
624
625 return of(errorChunk)
626 }),
627 finalize(async () => {
628 if (completeTitleUpdate) {
629 completeTitleUpdate()
630 }
631
632 if (!res.closed) {
633 res.end()
634 }
635
636 if (taskInfo.claudeCodeProcess) {
637 this.logger.debug({ taskId, sessionId }, `正在关闭任务 ${taskId} 的子进程`)
638 taskInfo.claudeCodeProcess.kill()
639 }
640
641 if (sessionId) {
642 await this.uploadAgentSession(sessionId, taskId)
643 }
644
645 completionResolver()
646 this.runningTasks.delete(taskId)
647
648 this.logger.log({
649 taskId,
650 sessionId,
651 userId,
652 }, `Task ${taskId} finished`)
653 }),
654 )
655 }),
656 catchError((error) => {
657 this.logger.error(error, 'Error in createContentGenerationTask')
658
659 const payload = getExceptionPayload(error)
660 const errorChunk = ContentGenerationTaskErrorChunkVo.create({
661 type: AgentMessageType.Error,
662 ...payload,
663 timestamp: Date.now(),
664 })
665
666 return of(errorChunk)
667 }),
668 )
669 }
670
671 abortTask(taskId: string): void {
672 const taskInfo = this.runningTasks.get(taskId)
673 if (!taskInfo || taskInfo.abortController.signal.aborted) {
674 return
675 }
676
677 this.logger.debug({ taskId, sessionId: taskInfo.sessionId }, `Aborting task ${taskId} via Redis broadcast`)
678 taskInfo.abortController.abort()
679 }
680
681 async waitForRunningTasks(): Promise<void> {
682 const memoryTasks = Array.from(this.runningTasks.values())
683
684 if (memoryTasks.length === 0) {
685 return
686 }
687
688 const waitPromises = memoryTasks.map(taskInfo =>
689 taskInfo.completionPromise
690 .then(() => {
691 this.logger.log({ taskId: taskInfo.taskId }, `Task ${taskInfo.taskId} completed during shutdown`)
692 })
693 .catch((error) => {
694 this.logger.error({ error, taskId: taskInfo.taskId }, `Task ${taskInfo.taskId} failed during shutdown`)
695 }),
696 )
697 await Promise.all(waitPromises)
698 }
699
700 private async initializeTask(
701 userId: string,
702 userType: UserType,
703 dto: CreateContentGenerationTaskDto,
704 abortController: AbortController,
705 req: Request,
706 ): Promise<{
707 taskId: string
708 sessionId: string | undefined
709 historicalMessages: Array<Record<string, unknown>>
710 abortController: AbortController
711 mcpServers: Record<string, McpServerConfig>
712 }> {
713 let task
714 let originalTask
715 let sessionId: string | undefined
716 let historicalMessages: Array<Record<string, unknown>> = []
717
718 if (dto.taskId) {
719 const originalTaskId = dto.taskId
720 this.logger.debug({ taskId: originalTaskId }, `Resuming conversation for task ${originalTaskId}`)
721
722 originalTask = await this.contentGenerateRepository.getByUserIdAndId(userId, originalTaskId)
723 if (!originalTask) {
724 this.logger.warn({ taskId: originalTaskId }, `Task ${originalTaskId} not found for user ${userId}`)
725 throw new AppException(ResponseCode.AgentTaskNotFound)
726 }
727
728 sessionId = originalTask.sessionId
729 if (!sessionId) {
730 this.logger.warn({ taskId: originalTaskId }, `Task ${originalTaskId} has no sessionId, cannot resume`)
731 throw new AppException(ResponseCode.AgentTaskNotFound)
732 }
733
734 this.logger.debug({ taskId: originalTaskId, sessionId }, `Task ${originalTaskId} resuming with sessionId: ${sessionId}`)
735 task = originalTask
736 historicalMessages = originalTask.messages || []
737
738 await this.downloadAgentSession(originalTask)
739 }
740 else {
741 task = await this.contentGenerateRepository.create({
742 userId,
743 })
744 this.logger.debug({ taskId: task.id }, `Created new task ${task.id} for user ${userId}`)
745 }
746
747 await this.contentGenerateRepository.updateStatus(task.id, ContentGenerationTaskStatus.Running)
748
749 const headers = filterHeaders(req.headers)
750 this.logger.debug({ headers }, 'mcp headers')
751 const mcpServers: Record<string, McpServerConfig> = {
752 [McpServerName.MediaGeneration]: this.mediaMcp.createServer(userId, userType),
753 [McpServerName.Util]: this.utilMcp.server,
754 [McpServerName.Aideo]: this.aideoMcp.createServer(userId, userType),
755 [McpServerName.VideoEdit]: this.videoEditMcp.createServer(userId, userType),
756 [McpServerName.DramaRecap]: this.dramaRecapMcp.createServer(userId, userType),
757 [McpServerName.VideoUtils]: this.videoUtilsMcp.createServer(userId, userType),
758 [McpServerName.StyleTransfer]: this.styleTransferMcp.createServer(userId, userType),
759 [McpServerName.ImageEdit]: this.imageEditMcp.createServer(userId, userType),
760 // [McpServerName.Subtitle]: this.subtitleMcp.createServer(userId, userType),
761 [McpServerName.Account]: {
762 type: 'http',
763 url: `${config.serverClient.baseUrl}/account/mcp`,
764 headers,
765 },
766 [McpServerName.Content]: {
767 type: 'http',
768 url: `${config.serverClient.baseUrl}/content/mcp`,
769 headers,
770 },
771 [McpServerName.Statistics]: {
772 type: 'http',
773 url: `${config.serverClient.baseUrl}/statistics/mcp`,
774 headers,
775 },
776 [McpServerName.Channels]: {
777 type: 'http',
778 url: `${config.serverClient.baseUrl}/channels/mcp`,
779 headers,
780 },
781 }
782
783 return {
784 taskId: task.id,
785 sessionId,
786 historicalMessages,
787 abortController,
788 mcpServers,
789 }
790 }
791
792 private createMessageStream(req: AsyncGenerator<SDKMessage, void>): Observable<SDKMessage> {
793 return from(req)
794 }
795
796 private async transformMessage(
797 chunk: SDKMessage,
798 taskId: string,
799 userId: string,
800 userType: UserType,
801 taskResult?: TaskResult,
802 sessionId?: string,
803 ): Promise<ContentGenerationTaskAgentChunkVo | ContentGenerationTaskErrorChunkVo> {
804 if (chunk.type === 'result') {
805 const { session_id, modelUsage, ...restChunk } = chunk
806 const currentSessionId = sessionId || session_id || undefined
807 if ('total_cost_usd' in chunk) {
808 await this.aiLogRepo.create({
809 userId,
810 userType,
811 taskId,
812 model: 'claude-agent',
813 channel: AiLogChannel.ClaudeAgent,
814 startedAt: new Date(),
815 type: AiLogType.Agent,
816 request: {},
817 response: { modelUsage, taskResult },
818 status: AiLogStatus.Success,
819 })
820 }
821
822 this.logger.debug({
823 taskId,
824 sessionId: currentSessionId,
825 modelUsage,
826 })
827
828 switch (chunk.subtype) {
829 case 'success': {
830 const messageToSave = {
831 ...restChunk,
832 message: chunk.result,
833 result: taskResult,
834 } as Omit<typeof restChunk, 'session_id'> & { message?: string, result?: TaskResult }
835
836 await this.contentGenerateRepository.updateMessage(taskId, messageToSave as unknown as Record<string, unknown>)
837
838 const requiresActionTypes = ['createChannel', 'updateChannel', 'loginChannel']
839 const resultArray = taskResult ? (Array.isArray(taskResult) ? taskResult : [taskResult]) : []
840 const hasRequiresAction = resultArray.some(item => item && 'action' in item && requiresActionTypes.includes(item.action as string))
841 const finalStatus = hasRequiresAction ? ContentGenerationTaskStatus.RequiresAction : ContentGenerationTaskStatus.Completed
842
843 void this.contentGenerateRepository.updateStatus(taskId, finalStatus)
844
845 return ContentGenerationTaskAgentChunkVo.create({
846 type: this.getMessageType(chunk),
847 message: AgentMessageVoHelper.create(messageToSave),
848 })
849 }
850
851 default: {
852 void this.contentGenerateRepository.updateStatus(taskId, ContentGenerationTaskStatus.Error)
853
854 const errorCodeMap: Record<string, ResponseCode> = {
855 error_max_budget_usd: ResponseCode.AgentTaskFailed,
856 error_during_execution: ResponseCode.AgentTaskFailed,
857 error_max_turns: ResponseCode.AgentTaskFailed,
858 error_max_structured_output_retries: ResponseCode.AgentTaskFailed,
859 }
860 const responseCode = errorCodeMap[chunk.subtype] || ResponseCode.AgentTaskFailed
861
862 const payload = getExceptionPayload(new AppException(responseCode))
863 const errorResult = ContentGenerationTaskErrorChunkVo.create({
864 type: AgentMessageType.Error,
865 ...payload,
866 timestamp: Date.now(),
867 })
868
869 await this.contentGenerateRepository.updateMessage(taskId, errorResult)
870 return errorResult
871 }
872 }
873 }
874
875 const messageToSave = 'session_id' in chunk
876 ? sanitizeMessage(chunk)
877 : chunk
878
879 if (messageToSave.type !== 'stream_event') {
880 await this.contentGenerateRepository.updateMessage(taskId, messageToSave as unknown as Record<string, unknown>)
881 }
882
883 const messageVo = AgentMessageVoHelper.create(messageToSave)
884
885 return ContentGenerationTaskAgentChunkVo.create({
886 type: this.getMessageType(chunk),
887 message: messageVo,
888 })
889 }
890
891 private buildSystemPromptContent(): ContentBlockParam[] {
892 return [{ type: 'text', text: SYSTEM_PROMPT }]
893 }
894
895 private async prepareAgentPrompt(dto: CreateContentGenerationTaskDto): Promise<{
896 normalizedContent: ContentBlock[]
897 enhancedContent: ContentBlockParam[]
898 systemPromptContent: ContentBlockParam[]
899 }> {
900 const normalizedContent = normalizePrompt(dto.prompt)
901 const relayContent = await this.resolveRelayJson(normalizedContent)
902 return {
903 normalizedContent,
904 enhancedContent: enhancePrompt(relayContent),
905 systemPromptContent: await this.buildSystemPromptContent(),
906 }
907 }
908
909 private async resolveRelayJson<T>(value: T): Promise<T> {
910 if (!this.relayMediaResolver) {
911 return value
912 }
913 return await this.relayMediaResolver.resolveJson(value)
914 }
915
916 private async uploadAgentSession(sessionId: string, taskId: string): Promise<void> {
917 try {
918 this.logger.debug(`Uploading session ${sessionId} for task ${taskId}`)
919 const taskProjectDir = this.getTaskProjectDir(taskId)
920 const projectsDir = join(this.sessionDir, '.claude/projects', taskProjectDir)
921 const sessionFile = join(projectsDir, `${sessionId}.jsonl`)
922
923 if (!fs.existsSync(sessionFile)) {
924 this.logger.fatal(`Session file not found locally: ${sessionFile}`)
925 return
926 }
927
928 const s3Key = `claude-session/.claude/projects/${this.projectDir}/${sessionId}.jsonl`
929 await this.storageProvider.putObject(s3Key, fs.readFileSync(sessionFile), 'application/jsonl')
930 this.logger.debug({ taskId, sessionId }, `Uploaded session file to S3: ${s3Key}`)
931
932 const agentIds = await this.parseAgentIds(sessionFile)
933 const uploadedTodoFiles = await this.uploadAgentSessionFiles(sessionId, agentIds, projectsDir, taskId)
934
935 await this.contentGenerateRepository.updateById(taskId, {
936 subAgentIds: agentIds,
937 todos: uploadedTodoFiles,
938 })
939 }
940 catch (error) {
941 this.logger.fatal({ error, sessionId, taskId }, 'Failed to upload session and agent files to S3')
942 }
943 }
944
945 private async downloadAgentSession(task: ContentGenerationTask): Promise<void> {
946 if (!task.sessionId) {
947 return
948 }
949 const taskId = task.id
950 const sessionId = task.sessionId
951
952 try {
953 this.logger.debug({ taskId, sessionId }, `Downloading session ${sessionId} for task ${taskId}`)
954 const s3Key = `claude-session/.claude/projects/${this.projectDir}/${task.sessionId}.jsonl`
955 const response = await this.storageProvider.getObject(s3Key)
956
957 const taskProjectDir = this.getTaskProjectDir(task.id)
958 const projectsDir = join(this.sessionDir, '.claude/projects', taskProjectDir)
959
960 if (response.buffer) {
961 fs.mkdirSync(projectsDir, { recursive: true })
962 fs.writeFileSync(join(projectsDir, `${task.sessionId}.jsonl`), response.buffer)
963 }
964 else {
965 this.logger.fatal({ taskId, sessionId }, `Session file not found in S3: ${s3Key}`)
966 }
967
968 if (task.subAgentIds?.length) {
969 const subagentsDir = join(projectsDir, task.sessionId, 'subagents')
970 fs.mkdirSync(subagentsDir, { recursive: true })
971
972 for (const agentId of task.subAgentIds) {
973 const agentS3Key = `claude-session/.claude/projects/${this.projectDir}/agent-${agentId}.jsonl`
974 const agentResponse = await this.storageProvider.getObject(agentS3Key)
975 if (agentResponse.buffer) {
976 fs.writeFileSync(join(subagentsDir, `agent-${agentId}.jsonl`), agentResponse.buffer)
977 this.logger.debug({ taskId, sessionId, agentId }, `Downloaded agent ${agentId} file from S3: ${agentS3Key}`)
978 }
979 else {
980 this.logger.fatal({ taskId, sessionId, agentId }, `Agent file not found in S3: ${agentS3Key}`)
981 }
982 }
983 }
984
985 if (task.todos?.length) {
986 const todosDir = join(this.sessionDir, '.claude/todos')
987 fs.mkdirSync(todosDir, { recursive: true })
988
989 for (const todoS3Key of task.todos) {
990 const fileResponse = await this.storageProvider.getObject(todoS3Key)
991 if (fileResponse.buffer) {
992 const fileName = todoS3Key.replace('claude-session/.claude/todos/', '')
993 fs.writeFileSync(join(todosDir, fileName), fileResponse.buffer)
994 this.logger.debug({ taskId, sessionId }, `Downloaded todo file ${fileName} from S3: ${todoS3Key}`)
995 }
996 else {
997 this.logger.fatal({ taskId, sessionId }, `Todo file not found in S3: ${todoS3Key}`)
998 }
999 }
1000 }
1001 }
1002 catch (error) {
1003 this.logger.fatal({ error, taskId, sessionId }, 'Failed to download session and agent files from S3')
1004 throw new AppException(ResponseCode.AgentSessionRecoveryFailed)
1005 }
1006 }
1007
1008 private async parseAgentIds(filePath: string): Promise<string[]> {
1009 const content = fs.readFileSync(filePath, 'utf-8')
1010 const agentIdPattern = /"agentId"\\s*:\\s*"([^"]+)"/g
1011 const agentIds = new Set<string>()
1012
1013 for (const match of content.matchAll(agentIdPattern)) {
1014 if (match[1]) {
1015 agentIds.add(match[1])
1016 }
1017 }
1018
1019 return Array.from(agentIds)
1020 }
1021
1022 private async uploadAgentSessionFiles(sessionId: string, agentIds: string[], projectsDir: string, taskId: string): Promise<string[]> {
1023 const uploadedTodoFiles: string[] = []
1024
1025 const subagentsDir = join(projectsDir, sessionId, 'subagents')
1026 for (const agentId of agentIds) {
1027 const agentProjectFile = join(subagentsDir, `agent-${agentId}.jsonl`)
1028 if (fs.existsSync(agentProjectFile)) {
1029 const s3Key = `claude-session/.claude/projects/${this.projectDir}/agent-${agentId}.jsonl`
1030 await this.storageProvider.putObject(s3Key, fs.readFileSync(agentProjectFile), 'application/jsonl')
1031 this.logger.debug({ taskId, sessionId, agentId }, `Uploaded agent ${agentId} file to S3: ${s3Key}`)
1032 }
1033 }
1034
1035 const todosDir = join(this.sessionDir, '.claude/todos')
1036 if (fs.existsSync(todosDir)) {
1037 const files = fs.readdirSync(todosDir).filter(file => file.startsWith(sessionId))
1038 for (const file of files) {
1039 const todoFilePath = join(todosDir, file)
1040 const s3Key = `claude-session/.claude/todos/${file}`
1041 await this.storageProvider.putObject(s3Key, fs.readFileSync(todoFilePath), 'application/json')
1042 this.logger.debug({ taskId, sessionId }, `Uploaded todo file ${file} to S3: ${s3Key}`)
1043 uploadedTodoFiles.push(s3Key)
1044 }
1045 }
1046
1047 return uploadedTodoFiles
1048 }
1049 }
1050
1050 lines TYPESCRIPT