diff --git a/src/main/core/conversations/impl/local-conversation.ts b/src/main/core/conversations/impl/local-conversation.ts index 32bd135ec..412354352 100644 --- a/src/main/core/conversations/impl/local-conversation.ts +++ b/src/main/core/conversations/impl/local-conversation.ts @@ -27,6 +27,7 @@ const MAX_RESPAWNS = 2; export class LocalConversationProvider implements ConversationProvider { private sessions = new Map(); + private knownSessionIds = new Set(); private respawnCounts = new Map(); private readonly projectId: string; private readonly taskPath: string; @@ -73,6 +74,7 @@ export class LocalConversationProvider implements ConversationProvider { conversation.taskId, conversation.id ); + this.knownSessionIds.add(sessionId); if (this.sessions.has(sessionId)) return; await claudeTrustService.maybeAutoTrustLocal({ @@ -188,28 +190,31 @@ export class LocalConversationProvider implements ConversationProvider { async stopSession(conversationId: string): Promise { const sessionId = makePtySessionId(this.projectId, this.taskId, conversationId); + this.knownSessionIds.delete(sessionId); const pty = this.sessions.get(sessionId); - if (!pty) return; - try { - pty.kill(); - } catch (e) { - log.warn('LocalAgentProvider: error killing PTY', { sessionId, error: String(e) }); + if (pty) { + try { + pty.kill(); + } catch (e) { + log.warn('LocalAgentProvider: error killing PTY', { sessionId, error: String(e) }); + } + this.sessions.delete(sessionId); + ptySessionRegistry.unregister(sessionId); } - this.sessions.delete(sessionId); - ptySessionRegistry.unregister(sessionId); if (this.tmux) { await killTmuxSession(this.exec, makeTmuxSessionName(sessionId)); } } async destroyAll(): Promise { - const sessionIds = Array.from(this.sessions.keys()); + const sessionIds = Array.from(this.knownSessionIds); await this.detachAll(); if (this.tmux) { await Promise.all( sessionIds.map((id) => killTmuxSession(this.exec, makeTmuxSessionName(id))) ); } + this.knownSessionIds.clear(); } async detachAll(): Promise { diff --git a/src/main/core/conversations/impl/ssh-conversation.ts b/src/main/core/conversations/impl/ssh-conversation.ts index 8e8a70727..a54bb868c 100644 --- a/src/main/core/conversations/impl/ssh-conversation.ts +++ b/src/main/core/conversations/impl/ssh-conversation.ts @@ -24,6 +24,7 @@ const MAX_RESPAWNS = 2; export class SshConversationProvider implements ConversationProvider { private sessions = new Map(); + private knownSessionIds = new Set(); private respawnCounts = new Map(); private readonly projectId: string; private readonly taskPath: string; @@ -74,6 +75,7 @@ export class SshConversationProvider implements ConversationProvider { conversation.taskId, conversation.id ); + this.knownSessionIds.add(sessionId); if (this.sessions.has(sessionId)) return; @@ -185,28 +187,31 @@ export class SshConversationProvider implements ConversationProvider { async stopSession(conversationId: string): Promise { const sessionId = makePtySessionId(this.projectId, this.taskId, conversationId); + this.knownSessionIds.delete(sessionId); const pty = this.sessions.get(sessionId); - if (!pty) return; - try { - pty.kill(); - } catch (e) { - log.warn('SshAgentProvider: error killing PTY', { sessionId, error: String(e) }); + if (pty) { + try { + pty.kill(); + } catch (e) { + log.warn('SshAgentProvider: error killing PTY', { sessionId, error: String(e) }); + } + this.sessions.delete(sessionId); + ptySessionRegistry.unregister(sessionId); } - this.sessions.delete(sessionId); - ptySessionRegistry.unregister(sessionId); if (this.tmux) { await killTmuxSession(this.exec, makeTmuxSessionName(sessionId)); } } async destroyAll(): Promise { - const sessionIds = Array.from(this.sessions.keys()); + const sessionIds = Array.from(this.knownSessionIds); await this.detachAll(); if (this.tmux) { await Promise.all( sessionIds.map((id) => killTmuxSession(this.exec, makeTmuxSessionName(id))) ); } + this.knownSessionIds.clear(); } async detachAll(): Promise { diff --git a/src/main/core/projects/impl/local-project-provider.ts b/src/main/core/projects/impl/local-project-provider.ts index 4918d5f05..0a61f7d54 100644 --- a/src/main/core/projects/impl/local-project-provider.ts +++ b/src/main/core/projects/impl/local-project-provider.ts @@ -2,6 +2,7 @@ import fs from 'node:fs'; import path from 'node:path'; import { Conversation } from '@shared/conversations'; import { LocalProject } from '@shared/projects'; +import { makePtySessionId } from '@shared/ptySessionId'; import { err, ok, type Result } from '@shared/result'; import { getTaskEnvVars } from '@shared/task/envVars'; import { Task, type TaskBootstrapStatus } from '@shared/tasks'; @@ -15,9 +16,11 @@ import { GitService } from '@main/core/git/impl/git-service'; import { bareRefName } from '@main/core/git/impl/git-utils'; import type { GitProvider } from '@main/core/git/types'; import { githubAuthService } from '@main/core/github/services/github-auth-service'; +import { killTmuxSession, makeTmuxSessionName } from '@main/core/pty/tmux-session-name'; import { appSettingsService } from '@main/core/settings/settings-service'; +import { getTaskSessionLeafIds } from '@main/core/tasks/session-targets'; import { LocalTerminalProvider } from '@main/core/terminals/impl/local-terminal-provider'; -import { getGitLocalExec } from '@main/core/utils/exec'; +import { getGitLocalExec, getLocalExec } from '@main/core/utils/exec'; import type { Workspace } from '@main/core/workspaces/workspace'; import { WorkspaceLifecycleService } from '@main/core/workspaces/workspace-lifecycle-service'; import { WorkspaceRegistry } from '@main/core/workspaces/workspace-registry'; @@ -36,6 +39,7 @@ import { TimeoutSignal, withTimeout } from '../utils'; import { WorktreeService } from '../worktrees/worktree-service'; const TASK_TIMEOUT_MS = 60_000; +const TEARDOWN_SCRIPT_WAIT_MS = 10_000; function toProvisionError(e: unknown): ProvisionTaskError { if (e instanceof TimeoutSignal) return { type: 'timeout', message: e.message, timeout: e.ms }; @@ -72,6 +76,7 @@ export class LocalProjectProvider implements ProjectProvider { private bootstrapErrors = new Map(); private worktreeService: WorktreeService; private workspaceRegistry = new WorkspaceRegistry(); + private readonly localExec = getLocalExec(); constructor( private readonly project: LocalProject, @@ -307,15 +312,24 @@ export class LocalProjectProvider implements ProjectProvider { async teardownTask(taskId: string): Promise> { if (this.tearingDownTasks.has(taskId)) return this.tearingDownTasks.get(taskId)!; const task = this.tasks.get(taskId); - if (!task) return ok(); + if (!task) { + await this.cleanupDetachedTmuxSessions(taskId); + return ok(); + } const promise = withTimeout(this.doTeardownTask(task), TASK_TIMEOUT_MS) .then(() => ok()) - .catch((e) => { + .catch(async (e) => { log.error('LocalProjectProvider: failed to teardown task', { taskId, error: String(e), }); + await this.cleanupDetachedTmuxSessions(taskId).catch((cleanupError) => { + log.warn('LocalProjectProvider: fallback tmux cleanup failed', { + taskId, + error: String(cleanupError), + }); + }); return err(toTeardownError(e)); }) .finally(() => { @@ -345,10 +359,25 @@ export class LocalProjectProvider implements ProjectProvider { const scripts = settings.scripts; if (scripts?.teardown && this.workspaceRegistry.refCount(wsId) === 1) { - await workspace.lifecycleService.runLifecycleScript( - { type: 'teardown', script: scripts.teardown }, - { waitForExit: true, exit: true } - ); + try { + const runTeardown = workspace.lifecycleService.runLifecycleScript( + { type: 'teardown', script: scripts.teardown }, + { waitForExit: true, exit: true } + ); + await withTimeout(runTeardown, TEARDOWN_SCRIPT_WAIT_MS); + } catch (error) { + if (error instanceof TimeoutSignal) { + log.debug('LocalProjectProvider: teardown script wait timed out', { + taskId: task.taskId, + timeoutMs: TEARDOWN_SCRIPT_WAIT_MS, + }); + } else { + log.warn('LocalProjectProvider: teardown script failed (continuing cleanup)', { + taskId: task.taskId, + error: String(error), + }); + } + } } } @@ -357,6 +386,16 @@ export class LocalProjectProvider implements ProjectProvider { await this.workspaceRegistry.release(wsId); } + private async cleanupDetachedTmuxSessions(taskId: string): Promise { + const { conversationIds, terminalIds } = await getTaskSessionLeafIds(this.project.id, taskId); + const sessionIds = [...conversationIds, ...terminalIds].map((leafId) => + makePtySessionId(this.project.id, taskId, leafId) + ); + await Promise.all( + sessionIds.map((sessionId) => killTmuxSession(this.localExec, makeTmuxSessionName(sessionId))) + ); + } + async removeTaskWorktree(taskBranch: string): Promise { const worktreePath = await this.worktreeService.getWorktree(taskBranch); if (worktreePath) { diff --git a/src/main/core/projects/impl/ssh-project-provider.ts b/src/main/core/projects/impl/ssh-project-provider.ts index ad5cc3c49..6d22e32bd 100644 --- a/src/main/core/projects/impl/ssh-project-provider.ts +++ b/src/main/core/projects/impl/ssh-project-provider.ts @@ -3,6 +3,7 @@ import path from 'node:path'; import type { SFTPWrapper } from 'ssh2'; import { Conversation } from '@shared/conversations'; import type { SshProject } from '@shared/projects'; +import { makePtySessionId } from '@shared/ptySessionId'; import { err, ok, type Result } from '@shared/result'; import { getTaskEnvVars } from '@shared/task/envVars'; import { Task, type TaskBootstrapStatus } from '@shared/tasks'; @@ -15,8 +16,10 @@ import { GitService } from '@main/core/git/impl/git-service'; import { bareRefName } from '@main/core/git/impl/git-utils'; import type { GitProvider } from '@main/core/git/types'; import { githubAuthService } from '@main/core/github/services/github-auth-service'; +import { killTmuxSession, makeTmuxSessionName } from '@main/core/pty/tmux-session-name'; import { SshClientProxy } from '@main/core/ssh/ssh-client-proxy'; import { SshConnectionEvent, sshConnectionManager } from '@main/core/ssh/ssh-connection-manager'; +import { getTaskSessionLeafIds } from '@main/core/tasks/session-targets'; import { SshTerminalProvider } from '@main/core/terminals/impl/ssh-terminal-provider'; import { getGitSshExec, getSshExec } from '@main/core/utils/exec'; import type { Workspace } from '@main/core/workspaces/workspace'; @@ -37,6 +40,7 @@ import { TimeoutSignal, withTimeout } from '../utils'; import { WorktreeService } from '../worktrees/worktree-service'; const TASK_TIMEOUT_MS = 60_000; +const TEARDOWN_SCRIPT_WAIT_MS = 10_000; function toProvisionError(e: unknown): ProvisionTaskError { if (e instanceof TimeoutSignal) return { type: 'timeout', message: e.message, timeout: e.ms }; @@ -347,15 +351,24 @@ export class SshProjectProvider implements ProjectProvider { async teardownTask(taskId: string): Promise> { if (this.tearingDownTasks.has(taskId)) return this.tearingDownTasks.get(taskId)!; const task = this.tasks.get(taskId); - if (!task) return ok(); + if (!task) { + await this.cleanupDetachedTmuxSessions(taskId); + return ok(); + } const promise = withTimeout(this.doTeardownTask(task), TASK_TIMEOUT_MS) .then(() => ok()) - .catch((e) => { + .catch(async (e) => { log.error('SshProjectProvider: failed to teardown task', { taskId, error: String(e), }); + await this.cleanupDetachedTmuxSessions(taskId).catch((cleanupError) => { + log.warn('SshProjectProvider: fallback tmux cleanup failed', { + taskId, + error: String(cleanupError), + }); + }); return err(toTeardownError(e)); }) .finally(() => { @@ -387,10 +400,25 @@ export class SshProjectProvider implements ProjectProvider { const scripts = settings.scripts; if (scripts?.teardown && this.workspaceRegistry.refCount(wsId) === 1) { - await workspace.lifecycleService.runLifecycleScript( - { type: 'teardown', script: scripts.teardown }, - { waitForExit: true, exit: true } - ); + try { + const runTeardown = workspace.lifecycleService.runLifecycleScript( + { type: 'teardown', script: scripts.teardown }, + { waitForExit: true, exit: true } + ); + await withTimeout(runTeardown, TEARDOWN_SCRIPT_WAIT_MS); + } catch (error) { + if (error instanceof TimeoutSignal) { + log.debug('SshProjectProvider: teardown script wait timed out', { + taskId: task.taskId, + timeoutMs: TEARDOWN_SCRIPT_WAIT_MS, + }); + } else { + log.warn('SshProjectProvider: teardown script failed (continuing cleanup)', { + taskId: task.taskId, + error: String(error), + }); + } + } } } @@ -399,6 +427,17 @@ export class SshProjectProvider implements ProjectProvider { await this.workspaceRegistry.release(wsId); } + private async cleanupDetachedTmuxSessions(taskId: string): Promise { + const { conversationIds, terminalIds } = await getTaskSessionLeafIds(this.project.id, taskId); + const sessionIds = [...conversationIds, ...terminalIds].map((leafId) => + makePtySessionId(this.project.id, taskId, leafId) + ); + const exec = getSshExec(this.proxy); + await Promise.all( + sessionIds.map((sessionId) => killTmuxSession(exec, makeTmuxSessionName(sessionId))) + ); + } + async removeTaskWorktree(taskBranch: string): Promise { const worktreePath = await this.worktreeService.getWorktree(taskBranch); if (worktreePath) { diff --git a/src/main/core/tasks/deleteTask.ts b/src/main/core/tasks/deleteTask.ts index b3549f60f..276fd2922 100644 --- a/src/main/core/tasks/deleteTask.ts +++ b/src/main/core/tasks/deleteTask.ts @@ -12,22 +12,22 @@ export async function deleteTask(projectId: string, taskId: string): Promise { + log.warn('deleteTask: teardown failed', { taskId, error: String(e) }); + return null; + }); + + if (teardownResult && !teardownResult.success) { + log.warn('deleteTask: teardown failed', { taskId, error: teardownResult.error.message }); + } + } + await db.delete(tasks).where(eq(tasks.id, taskId)); void viewStateService.del(`task:${taskId}`); capture('task_deleted'); if (project) { - void project - .teardownTask(taskId) - .then((teardownResult) => { - if (!teardownResult.success) { - log.warn('deleteTask: teardown failed', { taskId, error: teardownResult.error.message }); - } - }) - .catch((e) => { - log.warn('deleteTask: teardown failed', { taskId, error: String(e) }); - }); - if (task.taskBranch) { const siblings = await db .select({ id: tasks.id }) diff --git a/src/main/core/tasks/session-targets.ts b/src/main/core/tasks/session-targets.ts new file mode 100644 index 000000000..beab72edc --- /dev/null +++ b/src/main/core/tasks/session-targets.ts @@ -0,0 +1,29 @@ +import { and, eq } from 'drizzle-orm'; +import { db } from '@main/db/client'; +import { conversations, terminals } from '@main/db/schema'; + +export type TaskSessionLeafIds = { + conversationIds: string[]; + terminalIds: string[]; +}; + +export async function getTaskSessionLeafIds( + projectId: string, + taskId: string +): Promise { + const [conversationRows, terminalRows] = await Promise.all([ + db + .select({ id: conversations.id }) + .from(conversations) + .where(and(eq(conversations.projectId, projectId), eq(conversations.taskId, taskId))), + db + .select({ id: terminals.id }) + .from(terminals) + .where(and(eq(terminals.projectId, projectId), eq(terminals.taskId, taskId))), + ]); + + return { + conversationIds: conversationRows.map((row) => row.id), + terminalIds: terminalRows.map((row) => row.id), + }; +} diff --git a/src/main/core/terminals/impl/local-terminal-provider.ts b/src/main/core/terminals/impl/local-terminal-provider.ts index 5238ad732..028af2904 100644 --- a/src/main/core/terminals/impl/local-terminal-provider.ts +++ b/src/main/core/terminals/impl/local-terminal-provider.ts @@ -24,6 +24,7 @@ type SpawnPolicy = { export class LocalTerminalProvider implements TerminalProvider { private sessions = new Map(); + private knownSessionIds = new Set(); private respawnCounts = new Map(); private readonly projectId: string; private readonly scopeId: string; @@ -98,6 +99,7 @@ export class LocalTerminalProvider implements TerminalProvider { policy: SpawnPolicy ): Promise { const sessionId = makePtySessionId(terminal.projectId, terminal.taskId, terminal.id); + this.knownSessionIds.add(sessionId); if (this.sessions.has(sessionId)) return; const cfg: GeneralSessionConfig = { @@ -162,6 +164,7 @@ export class LocalTerminalProvider implements TerminalProvider { async killTerminal(terminalId: string): Promise { const sessionId = makePtySessionId(this.projectId, this.scopeId, terminalId); + this.knownSessionIds.delete(sessionId); const pty = this.sessions.get(sessionId); if (pty) { try { @@ -176,13 +179,14 @@ export class LocalTerminalProvider implements TerminalProvider { } async destroyAll(): Promise { - const sessionIds = Array.from(this.sessions.keys()); + const sessionIds = Array.from(this.knownSessionIds); await this.detachAll(); if (this.tmux) { await Promise.all( sessionIds.map((id) => killTmuxSession(this.exec, makeTmuxSessionName(id))) ); } + this.knownSessionIds.clear(); } async detachAll(): Promise { diff --git a/src/main/core/terminals/impl/ssh-terminal-provider.ts b/src/main/core/terminals/impl/ssh-terminal-provider.ts index 9565eaf77..a40ad0675 100644 --- a/src/main/core/terminals/impl/ssh-terminal-provider.ts +++ b/src/main/core/terminals/impl/ssh-terminal-provider.ts @@ -28,6 +28,7 @@ type SpawnPolicy = { export class SshTerminalProvider implements TerminalProvider { private sessions = new Map(); + private knownSessionIds = new Set(); private respawnCounts = new Map(); private terminals = new Map(); private readonly projectId: string; @@ -109,6 +110,7 @@ export class SshTerminalProvider implements TerminalProvider { policy: SpawnPolicy ): Promise { const sessionId = makePtySessionId(terminal.projectId, terminal.taskId, terminal.id); + this.knownSessionIds.add(sessionId); if (this.sessions.has(sessionId)) return; if (policy.trackForRehydrate) { this.terminals.set(terminal.id, terminal); @@ -209,6 +211,7 @@ export class SshTerminalProvider implements TerminalProvider { async killTerminal(terminalId: string): Promise { const sessionId = makePtySessionId(this.projectId, this.scopeId, terminalId); + this.knownSessionIds.delete(sessionId); const pty = this.sessions.get(sessionId); if (pty) { try { @@ -224,13 +227,14 @@ export class SshTerminalProvider implements TerminalProvider { } async destroyAll(): Promise { - const sessionIds = Array.from(this.sessions.keys()); + const sessionIds = Array.from(this.knownSessionIds); await this.detachAll(); if (this.tmux) { await Promise.all( sessionIds.map((id) => killTmuxSession(this.exec, makeTmuxSessionName(id))) ); } + this.knownSessionIds.clear(); this.terminals.clear(); }