Skip to content

Commit af15dd5

Browse files
waterdrop86claude
andcommitted
refactor: replace raw session SQL in agent runtimes, task-queue, and detach tests with repository functions
Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
1 parent bdf3898 commit af15dd5

6 files changed

Lines changed: 128 additions & 152 deletions

File tree

src/main/session/agent-runtimes/codex-runtime.ts

Lines changed: 7 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
import { basename } from 'node:path';
2-
import { getDb } from '../../db';
2+
import { getSessionRecord, getClaimedAgentIds, setAgentIdIfNull } from '../session-repository';
33
import { logger } from '../../logger';
44
import { findCodexThreadMatch } from '../codex-session-store';
55
import { hasPermissionPrompt } from '../prompt-detect';
@@ -30,19 +30,10 @@ export function scheduleCodexThreadCapture(
3030

3131
const deadline = Date.now() + 15_000;
3232
const poll = async (): Promise<void> => {
33-
const db = getDb();
34-
const row = db.prepare(
35-
'SELECT session_type, codex_thread_id FROM sessions WHERE session_id = ?',
36-
).get(input.sessionId) as { session_type: string; codex_thread_id: string | null } | undefined;
37-
if (!row || row.session_type !== 'codex' || row.codex_thread_id) return;
38-
39-
const claimedThreadIds = new Set(
40-
(
41-
db.prepare(
42-
'SELECT codex_thread_id FROM sessions WHERE codex_thread_id IS NOT NULL AND session_id != ?',
43-
).all(input.sessionId) as { codex_thread_id: string }[]
44-
).map((entry) => entry.codex_thread_id),
45-
);
33+
const record = getSessionRecord(input.sessionId);
34+
if (!record || record.session_type !== 'codex' || record.codex_thread_id) return;
35+
36+
const claimedThreadIds = getClaimedAgentIds('codex_thread_id', input.sessionId);
4637

4738
const match = findCodexThreadMatch({
4839
cwd: input.cwd,
@@ -52,10 +43,8 @@ export function scheduleCodexThreadCapture(
5243
claimedThreadIds,
5344
});
5445
if (match) {
55-
const result = db.prepare(
56-
'UPDATE sessions SET codex_thread_id = ? WHERE session_id = ? AND codex_thread_id IS NULL',
57-
).run(match.id, input.sessionId);
58-
if (result.changes > 0) {
46+
const claimed = setAgentIdIfNull(input.sessionId, 'codex_thread_id', match.id);
47+
if (claimed) {
5948
logger.info('session', 'Captured Codex thread ID', {
6049
sessionId: input.sessionId,
6150
codexThreadId: match.id,

src/main/session/agent-runtimes/copilot-runtime.ts

Lines changed: 7 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,5 @@
11
import { basename } from 'node:path';
2-
import { getDb } from '../../db';
2+
import { getSessionRecord, getClaimedAgentIds, setAgentIdIfNull } from '../session-repository';
33
import { logger } from '../../logger';
44
import { findCopilotSessionId } from '../copilot-session-store';
55
import { hasPermissionPrompt } from '../prompt-detect';
@@ -80,19 +80,10 @@ export function scheduleCopilotSessionCapture(
8080
const deadline = Date.now() + 15_000;
8181

8282
const poll = async (): Promise<void> => {
83-
const db = getDb();
84-
const row = db.prepare(
85-
'SELECT session_type, copilot_session_id FROM sessions WHERE session_id = ?',
86-
).get(input.sessionId) as { session_type: string; copilot_session_id: string | null } | undefined;
87-
if (!row || row.session_type !== 'copilot' || row.copilot_session_id) return;
88-
89-
const claimedSessionIds = new Set(
90-
(
91-
db.prepare(
92-
'SELECT copilot_session_id FROM sessions WHERE copilot_session_id IS NOT NULL AND session_id != ?',
93-
).all(input.sessionId) as { copilot_session_id: string }[]
94-
).map((entry) => entry.copilot_session_id),
95-
);
83+
const record = getSessionRecord(input.sessionId);
84+
if (!record || record.session_type !== 'copilot' || record.copilot_session_id) return;
85+
86+
const claimedSessionIds = getClaimedAgentIds('copilot_session_id', input.sessionId);
9687

9788
const match = findCopilotSessionId({
9889
cwd: input.cwd,
@@ -102,10 +93,8 @@ export function scheduleCopilotSessionCapture(
10293
});
10394

10495
if (match) {
105-
const result = db.prepare(
106-
'UPDATE sessions SET copilot_session_id = ? WHERE session_id = ? AND copilot_session_id IS NULL',
107-
).run(match, input.sessionId);
108-
if (result.changes > 0) {
96+
const claimed = setAgentIdIfNull(input.sessionId, 'copilot_session_id', match);
97+
if (claimed) {
10998
logger.info('session', 'Captured Copilot session ID', {
11099
sessionId: input.sessionId,
111100
copilotSessionId: match,

src/main/session/agent-runtimes/gemini-runtime.ts

Lines changed: 6 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
import { execFileSync } from 'node:child_process';
22
import { basename } from 'node:path';
3-
import { getDb } from '../../db';
3+
import { getSessionRecord, getClaimedAgentIds, setAgentIdIfNull } from '../session-repository';
44
import { logger } from '../../logger';
55
import {
66
parseGeminiSessionList,
@@ -45,20 +45,11 @@ export function scheduleGeminiSessionCapture(
4545
): void {
4646
const deadline = Date.now() + 15_000;
4747
const poll = async (): Promise<void> => {
48-
const db = getDb();
49-
const row = db.prepare(
50-
'SELECT session_type, gemini_session_id FROM sessions WHERE session_id = ?',
51-
).get(input.sessionId) as { session_type: string; gemini_session_id: string | null } | undefined;
52-
if (!row || row.session_type !== 'gemini' || row.gemini_session_id) return;
48+
const record = getSessionRecord(input.sessionId);
49+
if (!record || record.session_type !== 'gemini' || record.gemini_session_id) return;
5350

5451
try {
55-
const claimedSessionIds = new Set(
56-
(
57-
db.prepare(
58-
'SELECT gemini_session_id FROM sessions WHERE gemini_session_id IS NOT NULL AND session_id != ?',
59-
).all(input.sessionId) as { gemini_session_id: string }[]
60-
).map((entry) => entry.gemini_session_id),
61-
);
52+
const claimedSessionIds = getClaimedAgentIds('gemini_session_id', input.sessionId);
6253

6354
const entries = listGeminiSessions(input.command, input.cwd);
6455
const match = selectGeminiSessionCandidate(entries, {
@@ -67,10 +58,8 @@ export function scheduleGeminiSessionCapture(
6758
});
6859

6960
if (match) {
70-
const result = db.prepare(
71-
'UPDATE sessions SET gemini_session_id = ? WHERE session_id = ? AND gemini_session_id IS NULL',
72-
).run(match.geminiSessionId, input.sessionId);
73-
if (result.changes > 0) {
61+
const claimed = setAgentIdIfNull(input.sessionId, 'gemini_session_id', match.geminiSessionId);
62+
if (claimed) {
7463
logger.info('session', 'Captured Gemini session ID', {
7564
sessionId: input.sessionId,
7665
geminiSessionId: match.geminiSessionId,

src/main/session/session-repository.ts

Lines changed: 13 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -395,14 +395,26 @@ export function updateAutoLabel(id: string, label: string): boolean {
395395

396396
export function setAgentIdIfNull(
397397
id: string,
398-
column: 'claude_session_id' | 'gemini_session_id' | 'copilot_session_id',
398+
column: 'codex_thread_id' | 'gemini_session_id' | 'copilot_session_id',
399399
value: string,
400400
): boolean {
401401
return getDb()
402402
.prepare(`UPDATE sessions SET ${column} = ? WHERE session_id = ? AND ${column} IS NULL`)
403403
.run(value, id).changes > 0;
404404
}
405405

406+
export function getClaimedAgentIds(
407+
column: 'codex_thread_id' | 'gemini_session_id' | 'copilot_session_id',
408+
excludeSessionId: string,
409+
): Set<string> {
410+
return new Set(
411+
(getDb()
412+
.prepare(`SELECT ${column} FROM sessions WHERE ${column} IS NOT NULL AND session_id != ?`)
413+
.all(excludeSessionId) as Record<string, string>[])
414+
.map((r) => r[column]),
415+
);
416+
}
417+
406418
export function deleteSessionWithEvents(id: string): void {
407419
const db = getDb();
408420
db.transaction(() => {

src/main/task-queue.ts

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@ import type { WebContents } from 'electron';
33
import type { IPtyManager } from '../shared/pty-manager-interface';
44
import type { SessionManager } from './session/session-manager';
55
import { getDb } from './db';
6+
import { updateSession } from './session/session-repository';
67
import { logger } from './logger';
78
import { isAtClaudePrompt, isAtUserChoice, parseUserChoices } from './session/prompt-detect';
89
import { getTranscriptPath } from './session/transcript-path';
@@ -599,8 +600,7 @@ export class TaskQueue {
599600
// Update DB permission_mode so future tasks have accurate state
600601
if (task.permissionMode) {
601602
const modeForDb = task.permissionMode === 'default' ? null : task.permissionMode;
602-
getDb().prepare('UPDATE sessions SET permission_mode = ? WHERE session_id = ?')
603-
.run(modeForDb, task.targetSessionId);
603+
updateSession(task.targetSessionId!, { permissionMode: modeForDb });
604604
}
605605

606606
logger.info('task', 'Dispatched task after permission mode cycling', {

0 commit comments

Comments
 (0)