ouyang/cherry-studio
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

refactor: Rename message stream handler and update session creation logic

Browse Source
This commit is contained in:
Vaayne
2025-09-17 14:10:10 +08:00
parent 219d162e1a
commit d1ff8591a6
20 changed files with 147 additions and 269 deletions
Download Patch File Download Diff File Expand all files Collapse all files
src/main/services/agents/BaseService.ts
+4 -4
View File
@@ -1,11 +1,11 @@
import { type Client, createClient } from '@libsql/client'
import { loggerService } from '@logger'
import { drizzle } from 'drizzle-orm/libsql'
import { drizzle, type LibSQLDatabase } from 'drizzle-orm/libsql'
import fs from 'fs'
import path from 'path'
import * as schema from './database/schema'
import { MigrationService } from './database/MigrationService'
import * as schema from './database/schema'
import { dbPath } from './drizzle.config'
const logger = loggerService.withContext('BaseService')
@@ -24,7 +24,7 @@ const logger = loggerService.withContext('BaseService')
*/
export abstract class BaseService {
protected static client: Client | null = null
protected static db: ReturnType<typeof drizzle> | null = null
protected static db: LibSQLDatabase<typeof schema> | null = null
protected static isInitialized = false
protected static initializationPromise: Promise<void> | null = null
protected jsonFields: string[] = ['built_in_tools', 'mcps', 'configuration', 'accessible_paths']
@@ -110,7 +110,7 @@ export abstract class BaseService {
}
}
protected get database(): ReturnType<typeof drizzle> {
protected get database(): LibSQLDatabase<typeof schema> {
this.ensureInitialized()
return BaseService.db!
}
src/main/services/agents/database/schema/index.ts
+1 -1
View File
@@ -3,6 +3,6 @@
*/
export * from './agents.schema'
export * from './sessions.schema'
export * from './messages.schema'
export * from './migrations.schema'
export * from './sessions.schema'
src/main/services/agents/database/schema/messages.schema.ts
+1
View File
@@ -1,4 +1,5 @@
import { foreignKey, index, integer, sqliteTable, text } from 'drizzle-orm/sqlite-core'
import { sessionsTable } from './sessions.schema'
// session_messages table to log all messages, thoughts, actions, observations in a session
src/main/services/agents/database/schema/sessions.schema.ts
+2
View File
@@ -3,10 +3,12 @@
*/
import { foreignKey, index, sqliteTable, text } from 'drizzle-orm/sqlite-core'
import { agentsTable } from './agents.schema'
export const sessionsTable = sqliteTable('sessions', {
id: text('id').primaryKey(),
agent_type: text('agent_type').notNull(),
agent_id: text('agent_id').notNull(), // Primary agent ID for the session
name: text('name').notNull(),
description: text('description'),
src/main/services/agents/drizzle.config.ts
+5 -1
View File
@@ -16,12 +16,16 @@ function getDbPath() {
return path.join(app.getPath('userData'), 'agents.db')
}
const resolvedDbPath = getDbPath()
export const dbPath = resolvedDbPath
export default defineConfig({
dialect: 'sqlite',
schema: './src/main/services/agents/database/schema/index.ts',
out: './src/main/services/agents/database/drizzle',
dbCredentials: {
url: `file:${getDbPath()}`
url: `file:${resolvedDbPath}`
},
verbose: true,
strict: true
src/main/services/agents/services/SessionMessageService.ts
+14 -6
View File
@@ -11,7 +11,7 @@ import { UIMessageChunk } from 'ai'
import { count, eq } from 'drizzle-orm'
import { BaseService } from '../BaseService'
import { type InsertSessionMessageRow, sessionMessagesTable } from '../database/schema'
import { sessionMessagesTable } from '../database/schema'
import ClaudeCodeService from './claudecode'
const logger = loggerService.withContext('SessionMessageService')
@@ -76,19 +76,19 @@ export class SessionMessageService extends BaseService {
return { messages, total }
}
createSessionMessageStream(session: GetAgentSessionResponse, messageData: CreateSessionMessageRequest): EventEmitter {
createSessionMessage(session: GetAgentSessionResponse, messageData: CreateSessionMessageRequest): EventEmitter {
this.ensureInitialized()
// Create a new EventEmitter to manage the session message lifecycle
const sessionStream = new EventEmitter()
// No parent validation needed, start immediately
this.startClaudeCodeStream(session, messageData, sessionStream)
this.startSessionMessageStream(session, messageData, sessionStream)
return sessionStream
}
private startClaudeCodeStream(
private startSessionMessageStream(
session: GetAgentSessionResponse,
req: CreateSessionMessageRequest,
sessionStream: EventEmitter
@@ -99,7 +99,12 @@ export class SessionMessageService extends BaseService {
session_id = previousMessages[0].session_id
}
logger.debug('Claude Code stream message data:', { message: req, session_id })
logger.debug('Session Message stream message data:', { message: req, session_id })
if (session.agent_type !== 'claude-code') {
logger.error('Unsupported agent type for streaming:', { agent_type: session.agent_type })
throw new Error('Unsupported agent type for streaming')
}
// Create the streaming agent invocation (using invokeStream for streaming)
const claudeStream = this.cc.invoke(req.content, session.accessible_paths[0], session_id, {
@@ -107,7 +112,6 @@ export class SessionMessageService extends BaseService {
maxTurns: session.configuration?.maxTurns || 10
})
let sessionMessage: AgentSessionMessageEntity | null = null
const streamedChunks: UIMessageChunk[] = []
const rawAgentMessages: any[] = [] // Generic agent messages storage
@@ -202,6 +206,10 @@ export class SessionMessageService extends BaseService {
// error: new Error('Failed to save session message to database')
// })
// }
sessionStream.emit('data', {
type: 'complete',
result: structuredContent
})
break
}
src/main/services/agents/services/SessionService.ts
+8 -5
View File
@@ -1,4 +1,5 @@
import type {
AgentEntity,
AgentSessionEntity,
CreateSessionRequest,
GetAgentSessionResponse,
@@ -25,18 +26,18 @@ export class SessionService extends BaseService {
await BaseService.initialize()
}
async createSession(req: CreateSessionRequest): Promise<AgentSessionEntity> {
async createSession(agentId: string, req: CreateSessionRequest): Promise<AgentSessionEntity> {
this.ensureInitialized()
// Validate agent exists - we'll need to import AgentService for this check
// For now, we'll skip this validation to avoid circular dependencies
// The database foreign key constraint will handle this
const agents = await this.database.select().from(agentsTable).where(eq(agentsTable.id, req.agent_id)).limit(1)
const agents = await this.database.select().from(agentsTable).where(eq(agentsTable.id, agentId)).limit(1)
if (!agents[0]) {
throw new Error('Agent not found')
}
const agent = this.deserializeJsonFields(agents[0]) as AgentSessionEntity
const agent = this.deserializeJsonFields(agents[0]) as AgentEntity
const id = `session_${Date.now()}_${Math.random().toString(36).substring(2, 11)}`
const now = new Date().toISOString()
@@ -51,15 +52,17 @@ export class SessionService extends BaseService {
const insertData: InsertSessionRow = {
id,
agent_id: agentId,
agent_type: agent.type,
name: serializedData.name || null,
agent_id: serializedData.agent_id,
description: serializedData.description || null,
accessible_paths: serializedData.accessible_paths || null,
instructions: serializedData.instructions || null,
model: serializedData.model || null,
plan_model: serializedData.plan_model || null,
small_model: serializedData.small_model || null,
mcps: serializedData.mcps || null,
configuration: serializedData.configuration || null,
accessible_paths: serializedData.accessible_paths || null,
created_at: now,
updated_at: now
}
src/main/services/agents/services/claudecode/index.ts
-149
View File
@@ -242,155 +242,6 @@ class ClaudeCodeService implements AgentServiceInterface {
process.on('error', () => clearTimeout(timeout))
}
/**
* Set up process event handlers and return a promise that resolves with complete output
*/
private setupProcessHandlers(process: ChildProcess): Promise<ClaudeCodeResult> {
return new Promise((resolve, reject) => {
let stdoutData = ''
let stderrData = ''
const jsonOutput: any[] = []
let hasResolved = false
const startTime = Date.now()
// Handle stdout with proper encoding and buffering
if (process.stdout) {
process.stdout.setEncoding('utf8')
process.stdout.on('data', (data: string) => {
stdoutData += data
logger.debug('Agent stdout chunk:', { length: data.length })
// Parse JSON stream output line by line
const lines = data.split('\n')
for (const line of lines) {
if (line.trim()) {
try {
const parsed = JSON.parse(line.trim())
jsonOutput.push(parsed)
logger.silly('Parsed JSON output:', parsed)
} catch (e) {
// Not JSON, might be plain text output
logger.debug('Non-JSON stdout line:', { line: line.trim() })
}
}
}
})
process.stdout.on('end', () => {
logger.debug('Agent stdout stream ended')
})
}
// Handle stderr with proper encoding
if (process.stderr) {
process.stderr.setEncoding('utf8')
process.stderr.on('data', (data: string) => {
stderrData += data
logger.warn('Agent stderr chunk:', { data: data.trim() })
})
process.stderr.on('end', () => {
logger.debug('Agent stderr stream ended')
})
}
// Handle process exit
process.on('exit', (code, signal) => {
const duration = Date.now() - startTime
const success = code === 0
const status = success ? 'completed' : 'failed'
logger.info('Agent process exited', {
code,
signal,
success,
status,
duration,
stdoutLength: stdoutData.length,
stderrLength: stderrData.length,
jsonItems: jsonOutput.length
})
if (!hasResolved) {
hasResolved = true
resolve({
success,
stdout: stdoutData,
stderr: stderrData,
jsonOutput,
exitCode: code || undefined
})
}
})
// Handle process errors
process.on('error', (error) => {
const duration = Date.now() - startTime
logger.error('Agent process error:', {
error: error.message,
duration,
stdoutLength: stdoutData.length,
stderrLength: stderrData.length
})
if (!hasResolved) {
hasResolved = true
reject({
success: false,
stdout: stdoutData,
stderr: stderrData,
jsonOutput,
error
})
}
})
// Handle close event as a fallback
process.on('close', (code, signal) => {
const duration = Date.now() - startTime
logger.debug('Agent process closed', { code, signal, duration })
// Only resolve here if exit event hasn't fired
if (!hasResolved) {
hasResolved = true
const success = code === 0
resolve({
success,
stdout: stdoutData,
stderr: stderrData,
jsonOutput,
exitCode: code || undefined
})
}
})
// Set a timeout to prevent hanging indefinitely (reduced for debugging)
const timeout = setTimeout(() => {
if (!hasResolved) {
hasResolved = true
logger.error('Agent process timeout after 30 seconds', {
pid: process.pid,
stdoutLength: stdoutData.length,
stderrLength: stderrData.length,
jsonItems: jsonOutput.length
})
process.kill('SIGTERM')
reject({
success: false,
stdout: stdoutData,
stderr: stderrData,
jsonOutput,
error: new Error('Process timeout after 30 seconds')
})
}
}, 30 * 1000) // 30 seconds timeout for debugging
// Clear timeout when process ends
process.on('exit', () => clearTimeout(timeout))
process.on('error', () => clearTimeout(timeout))
})
}
}
export default ClaudeCodeService
src/main/services/agents/services/index.ts
+7 -2
View File
@@ -16,5 +16,10 @@ export { sessionMessageService } from './SessionMessageService'
export { sessionService } from './SessionService'
// Type definitions for service requests and responses
export type { CreateSessionRequest, ListSessionsOptions, UpdateSessionRequest } from './SessionService'
export type { AgentEntity, AgentSessionEntity,CreateAgentRequest, UpdateAgentRequest } from '@types'
export type {
AgentSessionMessageEntity,
CreateSessionRequest,
GetAgentSessionResponse,
ListOptions as SessionListOptions,
UpdateSessionRequest} from '@types'