diff --git a/apps/server/__tests__/services/session-create.spec.ts b/apps/server/__tests__/services/session-create.spec.ts index c84473a25..5568e998b 100644 --- a/apps/server/__tests__/services/session-create.spec.ts +++ b/apps/server/__tests__/services/session-create.spec.ts @@ -6,6 +6,7 @@ import { cancelSessionCreation, resetSessionCreationCancellationState } from '#~/services/session/creation-cancellation.js' +import { resetSessionCreationLifecycleState, waitForSessionCreation } from '#~/services/session/creation-lifecycle.js' const mocks = vi.hoisted(() => ({ getWorkspaceFolder: vi.fn(), @@ -62,6 +63,7 @@ describe('createSessionWithInitialMessage', () => { const createSession = vi.fn() const updateSession = vi.fn() const updateSessionRuntimeState = vi.fn() + const getSessionRuntimeState = vi.fn() const getSession = vi.fn() const getSessionWorkspace = vi.fn() const saveMessage = vi.fn() @@ -74,6 +76,7 @@ describe('createSessionWithInitialMessage', () => { beforeEach(() => { vi.clearAllMocks() resetSessionCreationCancellationState() + resetSessionCreationLifecycleState() createSession.mockImplementation((title?: string, id?: string) => ({ id: id ?? 'sess-1', @@ -84,11 +87,16 @@ describe('createSessionWithInitialMessage', () => { id, createdAt: Date.now() })) + getSessionRuntimeState.mockReturnValue({ + runtimeKind: 'interactive', + historySeedPending: false + }) saveMessage.mockReturnValue(true) vi.mocked(getDb).mockReturnValue({ createSession, updateSession, updateSessionRuntimeState, + getSessionRuntimeState, getSession, getSessionWorkspace, saveMessage, @@ -158,6 +166,72 @@ describe('createSessionWithInitialMessage', () => { ) }) + it('wakes early websocket waiters only after workspace provisioning completes', async () => { + let sessionPersisted = false + let finishProvisioning: (() => void) | undefined + createSession.mockImplementationOnce((title?: string, id?: string) => { + sessionPersisted = true + return { + id: id ?? 'sess-creation-race', + title, + createdAt: Date.now() + } + }) + getSession.mockImplementation((id: string) => + sessionPersisted + ? { id, createdAt: Date.now() } + : undefined + ) + mocks.provisionSessionWorkspace.mockImplementationOnce(() => + new Promise((resolve) => { + finishProvisioning = () => + resolve({ + sessionId: 'sess-creation-race', + workspaceFolder: '/workspace/root' + }) + }) + ) + + const waitPromise = waitForSessionCreation('sess-creation-race') + let waiterSettled = false + void waitPromise.then(() => { + waiterSettled = true + }) + const creationPromise = createSessionWithInitialMessage({ + id: 'sess-creation-race', + shouldStart: false + }) + + await vi.waitFor(() => { + expect(mocks.provisionSessionWorkspace).toHaveBeenCalledOnce() + }) + expect(waiterSettled).toBe(false) + + finishProvisioning?.() + await creationPromise + await expect(waitPromise).resolves.toBeUndefined() + }) + + it('does not discard an existing non-shell session on duplicate creation', async () => { + getSession.mockReturnValue({ + id: 'sess-existing', + createdAt: Date.now(), + messageCount: 0, + title: 'Existing session' + }) + createSession.mockImplementationOnce(() => { + throw new Error('UNIQUE constraint failed: sessions.id') + }) + + await expect(createSessionWithInitialMessage({ + id: 'sess-existing', + shouldStart: false + })).rejects.toThrow('UNIQUE constraint failed') + + expect(mocks.deleteRuntimeSessionStores).not.toHaveBeenCalled() + expect(deleteSession).not.toHaveBeenCalled() + }) + it('cancels creation before a pending session is created', async () => { cancelSessionCreation('sess-cancel-before-create') @@ -199,6 +273,46 @@ describe('createSessionWithInitialMessage', () => { expect(deleteSession).toHaveBeenCalledWith('sess-cancel-during-workspace') }) + it('keeps the first creation cancellable after an overlapping duplicate request fails', async () => { + let rejectProvisioning: ((error: unknown) => void) | undefined + createSession + .mockImplementationOnce((title?: string, id?: string) => ({ + id: id ?? 'sess-overlap-cancel', + title, + createdAt: Date.now() + })) + .mockImplementationOnce(() => { + throw new Error('UNIQUE constraint failed: sessions.id') + }) + mocks.provisionSessionWorkspace.mockImplementationOnce(async (_sessionId, options) => + new Promise((_resolve, reject) => { + rejectProvisioning = reject + options.signal?.addEventListener('abort', () => { + reject(options.signal?.reason) + }, { once: true }) + }) + ) + + const firstCreation = createSessionWithInitialMessage({ + id: 'sess-overlap-cancel', + initialMessage: 'hello' + }) + await vi.waitFor(() => { + expect(rejectProvisioning).toBeTypeOf('function') + }) + + await expect(createSessionWithInitialMessage({ + id: 'sess-overlap-cancel', + initialMessage: 'duplicate' + })).rejects.toThrow('UNIQUE constraint failed') + + expect(cancelSessionCreation('sess-overlap-cancel')).toBe('active') + await expect(firstCreation).rejects.toMatchObject({ + code: 'session_creation_cancelled' + }) + expect(deleteSession).toHaveBeenCalledWith('sess-overlap-cancel') + }) + it('uses the shared workspace by default when the project config is not set', async () => { await createSessionWithInitialMessage({ title: 'Demo', diff --git a/apps/server/__tests__/services/session-creation-lifecycle.spec.ts b/apps/server/__tests__/services/session-creation-lifecycle.spec.ts new file mode 100644 index 000000000..a3e377999 --- /dev/null +++ b/apps/server/__tests__/services/session-creation-lifecycle.spec.ts @@ -0,0 +1,121 @@ +import { beforeEach, describe, expect, it, vi } from 'vitest' + +import { getDb } from '#~/db/index.js' +import { + SessionCreationWaitError, + beginSessionCreation, + resetSessionCreationLifecycleState, + waitForSessionCreation +} from '#~/services/session/creation-lifecycle.js' + +vi.mock('#~/db/index.js', () => ({ + getDb: vi.fn() +})) + +describe('session creation lifecycle', () => { + const getSession = vi.fn() + + beforeEach(() => { + vi.useRealTimers() + resetSessionCreationLifecycleState() + vi.clearAllMocks() + getSession.mockReturnValue(undefined) + vi.mocked(getDb).mockReturnValue({ getSession } as any) + }) + + it('keeps an early waiter pending until HTTP creation completes', async () => { + const waitPromise = waitForSessionCreation('sess-early') + const attempt = beginSessionCreation('sess-early') + let settled = false + void waitPromise.then( + () => { + settled = true + }, + () => { + settled = true + } + ) + + await Promise.resolve() + expect(settled).toBe(false) + + attempt.complete() + + await expect(waitPromise).resolves.toBeUndefined() + }) + + it('does not fail waiters while another creation attempt is still active', async () => { + const firstAttempt = beginSessionCreation('sess-overlap') + const secondAttempt = beginSessionCreation('sess-overlap') + const waitPromise = waitForSessionCreation('sess-overlap') + let settled = false + void waitPromise.then( + () => { + settled = true + }, + () => { + settled = true + } + ) + + firstAttempt.fail(new Error('first failed')) + await Promise.resolve() + expect(settled).toBe(false) + + secondAttempt.fail(new Error('second failed')) + + await expect(waitPromise).rejects.toBeInstanceOf(SessionCreationWaitError) + }) + + it('resolves all waiters as soon as one overlapping attempt succeeds', async () => { + const firstAttempt = beginSessionCreation('sess-success') + const secondAttempt = beginSessionCreation('sess-success') + const firstWaiter = waitForSessionCreation('sess-success') + const secondWaiter = waitForSessionCreation('sess-success') + + firstAttempt.complete() + secondAttempt.fail(new Error('late duplicate failed')) + + await expect(Promise.all([firstWaiter, secondWaiter])).resolves.toEqual([undefined, undefined]) + }) + + it('returns immediately for an already-created session', async () => { + getSession.mockReturnValue({ id: 'sess-existing' }) + + await expect(waitForSessionCreation('sess-existing')).resolves.toBeUndefined() + }) + + it('times out when neither a session nor a creation attempt appears', async () => { + vi.useFakeTimers() + const waitPromise = waitForSessionCreation('sess-missing', { timeoutMs: 100 }) + const rejection = expect(waitPromise).rejects.toMatchObject({ + code: 'session_creation_wait_failed', + message: 'Session creation timed out' + }) + + await vi.advanceTimersByTimeAsync(100) + + await rejection + }) + + it('stops the wait timeout once HTTP creation begins', async () => { + vi.useFakeTimers() + const waitPromise = waitForSessionCreation('sess-slow', { timeoutMs: 100 }) + const attempt = beginSessionCreation('sess-slow') + let settled = false + void waitPromise.then( + () => { + settled = true + }, + () => { + settled = true + } + ) + + await vi.advanceTimersByTimeAsync(1_000) + expect(settled).toBe(false) + + attempt.complete() + await expect(waitPromise).resolves.toBeUndefined() + }) +}) diff --git a/apps/server/__tests__/services/session-start.spec.ts b/apps/server/__tests__/services/session-start.spec.ts index 84f7fbb2f..5a7ce28d9 100644 --- a/apps/server/__tests__/services/session-start.spec.ts +++ b/apps/server/__tests__/services/session-start.spec.ts @@ -259,6 +259,18 @@ describe('startAdapterSession', () => { expect(runtime.session.kill).not.toHaveBeenCalled() }) + it('rejects a missing session without creating database or workspace state', async () => { + currentSession = undefined + + await expect(startAdapterSession('sess-missing')).rejects.toThrow( + 'Session sess-missing does not exist' + ) + + expect(createSession).not.toHaveBeenCalled() + expect(mocks.provisionSessionWorkspace).not.toHaveBeenCalled() + expect(mocks.run).not.toHaveBeenCalled() + }) + it('deduplicates concurrent start requests for the same session', async () => { let resolveRun: | ((value: { session: { emit: ReturnType; kill: ReturnType } }) => void) diff --git a/apps/server/__tests__/websocket/server.spec.ts b/apps/server/__tests__/websocket/server.spec.ts index 2d0f6670a..b3c0d8ded 100644 --- a/apps/server/__tests__/websocket/server.spec.ts +++ b/apps/server/__tests__/websocket/server.spec.ts @@ -4,6 +4,12 @@ import type { Server } from 'node:http' import { beforeEach, describe, expect, it, vi } from 'vitest' +import { SessionCreationWaitError } from '#~/services/session/creation-lifecycle.js' + +const lifecycleMocks = vi.hoisted(() => ({ + waitForSessionCreation: vi.fn() +})) + const addSessionSubscriberSocket = vi.fn() const removeSessionSubscriberSocket = vi.fn() const notifySessionUpdated = vi.fn() @@ -69,6 +75,16 @@ vi.mock('#~/services/session/interaction.js', () => ({ handleInteractionResponse })) +vi.mock('#~/services/session/creation-lifecycle.js', async () => { + const actual = await vi.importActual( + '#~/services/session/creation-lifecycle.js' + ) + return { + ...actual, + waitForSessionCreation: lifecycleMocks.waitForSessionCreation + } +}) + vi.mock('#~/services/session/runtime.js', () => ({ addSessionSubscriberSocket, removeSessionSubscriberSocket, @@ -100,6 +116,7 @@ describe('setupWebSocket', () => { runtimeKind: 'interactive', historySeedPending: false }) + lifecycleMocks.waitForSessionCreation.mockResolvedValue(undefined) const { setupWebSocket } = await import('#~/websocket/server.js') setupWebSocket(new EventEmitter() as Server, { @@ -216,6 +233,159 @@ describe('setupWebSocket', () => { expect(attachSocketToSession).toHaveBeenCalledWith('sess-1', ws, 'adapter') }) + it('waits for HTTP session creation before starting the adapter', async () => { + let completeCreation: (() => void) | undefined + lifecycleMocks.waitForSessionCreation.mockImplementationOnce(() => + new Promise((resolve) => { + completeCreation = resolve + }) + ) + startAdapterSession.mockResolvedValue({ sockets: new Set(), session: {} }) + + const ws = { + close: vi.fn(), + on: vi.fn(), + readyState: 1, + send: vi.fn() + } + const connection = connectionHandler?.(ws, { + url: '/ws?sessionId=sess-early', + headers: { host: 'localhost' } + }) + + await vi.waitFor(() => { + expect(lifecycleMocks.waitForSessionCreation).toHaveBeenCalledWith('sess-early', { + signal: expect.any(AbortSignal) + }) + }) + expect(startAdapterSession).not.toHaveBeenCalled() + expect(attachSocketToSession).not.toHaveBeenCalled() + + getSession.mockReturnValue({ id: 'sess-early', status: 'running' }) + completeCreation?.() + await connection + + expect(startAdapterSession).toHaveBeenCalledWith('sess-early', expect.any(Object)) + expect(attachSocketToSession).toHaveBeenCalledWith('sess-early', ws, 'adapter') + }) + + it('ignores a socket that closed before the lazy session handler loaded', async () => { + const ws = { + close: vi.fn(), + on: vi.fn(), + readyState: 3, + send: vi.fn() + } + + await connectionHandler?.(ws, { + url: '/ws?sessionId=sess-already-closed', + headers: { host: 'localhost' } + }) + + expect(lifecycleMocks.waitForSessionCreation).not.toHaveBeenCalled() + expect(startAdapterSession).not.toHaveBeenCalled() + expect(attachSocketToSession).not.toHaveBeenCalled() + expect(ws.send).not.toHaveBeenCalled() + }) + + it('closes an early websocket when HTTP session creation fails', async () => { + lifecycleMocks.waitForSessionCreation.mockRejectedValueOnce( + new SessionCreationWaitError('sess-failed', 'Session creation failed') + ) + + const ws = { + close: vi.fn(), + on: vi.fn(), + readyState: 1, + send: vi.fn() + } + + await connectionHandler?.(ws, { + url: '/ws?sessionId=sess-failed', + headers: { host: 'localhost' } + }) + + expect(startAdapterSession).not.toHaveBeenCalled() + expect(ws.send).toHaveBeenCalledOnce() + expect(JSON.parse(String(ws.send.mock.calls[0]?.[0]))).toMatchObject({ + type: 'error', + data: { + message: 'Session creation failed', + fatal: true + } + }) + expect(ws.close).toHaveBeenCalledWith(1008, 'Session creation failed') + }) + + it('cancels the creation wait when the websocket closes early', async () => { + lifecycleMocks.waitForSessionCreation.mockImplementationOnce((_sessionId, options) => + new Promise((_resolve, reject) => { + options.signal?.addEventListener('abort', () => { + reject(new SessionCreationWaitError('sess-closed', 'Session creation wait cancelled')) + }, { once: true }) + }) + ) + + let closeHandler: (() => void) | undefined + const ws = { + close: vi.fn(), + on: vi.fn((event: string, handler: () => void) => { + if (event === 'close') closeHandler = handler + }), + readyState: 1, + send: vi.fn() + } + const connection = connectionHandler?.(ws, { + url: '/ws?sessionId=sess-closed', + headers: { host: 'localhost' } + }) + + await vi.waitFor(() => { + expect(closeHandler).toBeTypeOf('function') + }) + closeHandler?.() + await connection + + expect(startAdapterSession).not.toHaveBeenCalled() + expect(attachSocketToSession).not.toHaveBeenCalled() + expect(ws.send).not.toHaveBeenCalled() + }) + + it('does not attach a socket that closes while the adapter is starting', async () => { + let finishAdapterStart: ((runtime: { sockets: Set; session: object }) => void) | undefined + startAdapterSession.mockImplementationOnce(() => + new Promise((resolve) => { + finishAdapterStart = resolve + }) + ) + getSession.mockReturnValue({ id: 'sess-starting', status: 'running' }) + + let closeHandler: (() => void) | undefined + const ws = { + close: vi.fn(), + on: vi.fn((event: string, handler: () => void) => { + if (event === 'close') closeHandler = handler + }), + readyState: 1, + send: vi.fn() + } + const connection = connectionHandler?.(ws, { + url: '/ws?sessionId=sess-starting', + headers: { host: 'localhost' } + }) + + await vi.waitFor(() => { + expect(startAdapterSession).toHaveBeenCalledOnce() + expect(closeHandler).toBeTypeOf('function') + }) + closeHandler?.() + finishAdapterStart?.({ sockets: new Set(), session: {} }) + await connection + + expect(attachSocketToSession).not.toHaveBeenCalled() + expect(ws.send).not.toHaveBeenCalled() + }) + it('keeps completed sessions in passive mode when opening the page', async () => { getSession.mockReturnValue({ id: 'sess-1', diff --git a/apps/server/src/services/session/AGENTS.md b/apps/server/src/services/session/AGENTS.md index ee2308a45..dd8f83121 100644 --- a/apps/server/src/services/session/AGENTS.md +++ b/apps/server/src/services/session/AGENTS.md @@ -2,6 +2,7 @@ - index.ts：会话生命周期主入口，负责 Adapter 会话启动、用户消息注入、状态更新、中断与终止；`getWorkspaceActivitySnapshot` 汇总 workspace 是否仍有运行中 / 等待输入的会话 - create.ts：会话创建编排，负责初始消息注入、标签初始化与按需启动会话 +- creation-lifecycle.ts：协调 HTTP 创建与提前到达的 WebSocket；HTTP 是唯一持久化写入方，WebSocket 只等待完整创建结果 - events.ts：会话事件落库与状态派生，负责从消息中提取摘要并更新会话元信息 - interaction.ts：交互请求/响应服务，负责等待用户输入、超时处理与 external session 交互闭环 - notification.ts：会话状态通知，基于 config 子域装载的统一配置决定是否发送系统通知 diff --git a/apps/server/src/services/session/create.ts b/apps/server/src/services/session/create.ts index 23df1a67a..406eda804 100644 --- a/apps/server/src/services/session/create.ts +++ b/apps/server/src/services/session/create.ts @@ -16,6 +16,7 @@ import { registerSessionCreationCancellation, throwIfSessionCreationCancelled } from '#~/services/session/creation-cancellation.js' +import { beginSessionCreation, isSessionCreationActive } from '#~/services/session/creation-lifecycle.js' import { notifySessionUpdated } from '#~/services/session/runtime.js' import { deleteSessionWorkspace, @@ -139,6 +140,12 @@ export async function createSessionWithInitialMessage(options: { } = options const db = getDb() const creationCancellationId = id ?? `session-creation-${Date.now()}-${Math.random().toString(16).slice(2)}` + const creationAlreadyActive = id != null && isSessionCreationActive(id) + const existingRequestedSession = id == null ? undefined : db.getSession(id) + const creationAttempt = id != null && + (existingRequestedSession == null || creationAlreadyActive) + ? beginSessionCreation(id) + : undefined const creationCancellation = registerSessionCreationCancellation(creationCancellationId) let session: Session | undefined @@ -271,14 +278,19 @@ export async function createSessionWithInitialMessage(options: { } } + creationAttempt?.complete() return session } catch (err) { - if (session != null) { - await deleteSessionWorkspace(session.id, { force: true }) - if (db.getSessionWorkspace(session.id)?.state !== 'deleting') { - db.deleteSessionWorkspace(session.id) - db.deleteSession(session.id) + try { + if (session != null) { + await deleteSessionWorkspace(session.id, { force: true }) + if (db.getSessionWorkspace(session.id)?.state !== 'deleting') { + db.deleteSessionWorkspace(session.id) + db.deleteSession(session.id) + } } + } finally { + creationAttempt?.fail(err) } throw err } finally { diff --git a/apps/server/src/services/session/creation-cancellation.ts b/apps/server/src/services/session/creation-cancellation.ts index 97d47f7ec..e16e98854 100644 --- a/apps/server/src/services/session/creation-cancellation.ts +++ b/apps/server/src/services/session/creation-cancellation.ts @@ -1,6 +1,11 @@ const PENDING_CANCELLATION_TTL_MS = 5 * 60 * 1000 -const activeCreationControllers = new Map() +interface ActiveSessionCreationCancellation { + controller: AbortController + registrations: number +} + +const activeCreationControllers = new Map() const pendingCreationCancellations = new Map() export class SessionCreationCancelledError extends Error { @@ -26,16 +31,27 @@ const prunePendingCancellations = () => { export const registerSessionCreationCancellation = (sessionId: string) => { prunePendingCancellations() - const controller = new AbortController() - activeCreationControllers.set(sessionId, controller) - if (pendingCreationCancellations.delete(sessionId)) { - controller.abort(new SessionCreationCancelledError(sessionId)) + let state = activeCreationControllers.get(sessionId) + if (state == null) { + state = { + controller: new AbortController(), + registrations: 0 + } + activeCreationControllers.set(sessionId, state) + if (pendingCreationCancellations.delete(sessionId)) { + state.controller.abort(new SessionCreationCancelledError(sessionId)) + } } + state.registrations += 1 + let unregistered = false return { - signal: controller.signal, + signal: state.controller.signal, unregister: () => { - if (activeCreationControllers.get(sessionId) === controller) { + if (unregistered) return + unregistered = true + state.registrations = Math.max(0, state.registrations - 1) + if (state.registrations === 0 && activeCreationControllers.get(sessionId) === state) { activeCreationControllers.delete(sessionId) } } @@ -50,10 +66,10 @@ export const cancelSessionCreation = ( ) => { prunePendingCancellations() - const controller = activeCreationControllers.get(sessionId) - if (controller != null) { - if (!controller.signal.aborted) { - controller.abort(new SessionCreationCancelledError(sessionId)) + const state = activeCreationControllers.get(sessionId) + if (state != null) { + if (!state.controller.signal.aborted) { + state.controller.abort(new SessionCreationCancelledError(sessionId)) } return 'active' as const } diff --git a/apps/server/src/services/session/creation-lifecycle.ts b/apps/server/src/services/session/creation-lifecycle.ts new file mode 100644 index 000000000..e0795726e --- /dev/null +++ b/apps/server/src/services/session/creation-lifecycle.ts @@ -0,0 +1,180 @@ +import { getDb } from '#~/db/index.js' + +const DEFAULT_SESSION_CREATION_WAIT_TIMEOUT_MS = 30_000 + +interface SessionCreationWaiter { + resolve: () => void + reject: (error: Error) => void + creationStarted: () => void +} + +interface SessionCreationState { + activeAttempts: number + waiters: Set +} + +const sessionCreationStates = new Map() + +export class SessionCreationWaitError extends Error { + code = 'session_creation_wait_failed' + sessionId: string + + constructor(sessionId: string, message: string) { + super(message) + this.name = 'SessionCreationWaitError' + this.sessionId = sessionId + } +} + +const getOrCreateSessionCreationState = (sessionId: string) => { + const existing = sessionCreationStates.get(sessionId) + if (existing != null) { + return existing + } + + const created: SessionCreationState = { + activeAttempts: 0, + waiters: new Set() + } + sessionCreationStates.set(sessionId, created) + return created +} + +const settleWaiters = ( + sessionId: string, + state: SessionCreationState, + error?: SessionCreationWaitError +) => { + for (const waiter of [...state.waiters]) { + if (error == null) { + waiter.resolve() + } else { + waiter.reject(error) + } + } + + if (sessionCreationStates.get(sessionId) === state) { + sessionCreationStates.delete(sessionId) + } +} + +export const beginSessionCreation = (sessionId: string) => { + const state = getOrCreateSessionCreationState(sessionId) + state.activeAttempts += 1 + for (const waiter of state.waiters) { + waiter.creationStarted() + } + let settled = false + + return { + complete: () => { + if (settled) return + settled = true + state.activeAttempts = Math.max(0, state.activeAttempts - 1) + settleWaiters(sessionId, state) + }, + fail: (_error: unknown) => { + if (settled) return + settled = true + state.activeAttempts = Math.max(0, state.activeAttempts - 1) + if (state.activeAttempts === 0) { + settleWaiters( + sessionId, + state, + new SessionCreationWaitError(sessionId, 'Session creation failed') + ) + } + } + } +} + +export const isSessionCreationActive = (sessionId: string) => + (sessionCreationStates.get(sessionId)?.activeAttempts ?? 0) > 0 + +export const waitForSessionCreation = ( + sessionId: string, + options: { + signal?: AbortSignal + timeoutMs?: number + } = {} +) => { + const existingState = sessionCreationStates.get(sessionId) + if (existingState == null && getDb().getSession(sessionId) != null) { + return Promise.resolve() + } + + const state = existingState ?? getOrCreateSessionCreationState(sessionId) + const timeoutMs = options.timeoutMs ?? DEFAULT_SESSION_CREATION_WAIT_TIMEOUT_MS + + return new Promise((resolve, reject) => { + let settled = false + let timeout: ReturnType | undefined + + const clearWaitTimeout = () => { + if (timeout != null) { + clearTimeout(timeout) + timeout = undefined + } + } + const cleanup = () => { + clearWaitTimeout() + options.signal?.removeEventListener('abort', handleAbort) + state.waiters.delete(waiter) + if (state.activeAttempts === 0 && state.waiters.size === 0 && sessionCreationStates.get(sessionId) === state) { + sessionCreationStates.delete(sessionId) + } + } + const finish = (error?: Error) => { + if (settled) return + settled = true + cleanup() + if (error == null) { + resolve() + } else { + reject(error) + } + } + const waiter: SessionCreationWaiter = { + resolve: () => finish(), + reject: error => finish(error), + creationStarted: clearWaitTimeout + } + const handleAbort = () => { + finish(new SessionCreationWaitError(sessionId, 'Session creation wait cancelled')) + } + + state.waiters.add(waiter) + if (state.activeAttempts === 0) { + timeout = setTimeout(() => { + finish(new SessionCreationWaitError(sessionId, 'Session creation timed out')) + }, timeoutMs) + } + + if (options.signal?.aborted === true) { + handleAbort() + } else { + options.signal?.addEventListener('abort', handleAbort, { once: true }) + } + }) +} + +export const isSessionCreationWaitError = (error: unknown): error is SessionCreationWaitError => { + if (error instanceof SessionCreationWaitError) { + return true + } + if (!(error instanceof Error) || !('code' in error)) { + return false + } + return (error as Error & { code?: unknown }).code === 'session_creation_wait_failed' +} + +export const resetSessionCreationLifecycleState = () => { + for (const [sessionId, state] of sessionCreationStates) { + settleWaiters( + sessionId, + state, + new SessionCreationWaitError(sessionId, 'Session creation lifecycle reset') + ) + } + sessionCreationStates.clear() +} diff --git a/apps/server/src/services/session/index.ts b/apps/server/src/services/session/index.ts index ea02360dc..8be4d5ef6 100644 --- a/apps/server/src/services/session/index.ts +++ b/apps/server/src/services/session/index.ts @@ -91,7 +91,7 @@ import { clearSessionWorkspaceChangeTracking, finalizeSessionWorkspaceChangeTracking } from '#~/services/session/workspace-changes.js' -import { provisionSessionWorkspace, resolveSessionWorkspace } from '#~/services/session/workspace.js' +import { resolveSessionWorkspace } from '#~/services/session/workspace.js' import { runConfiguredWorktreeEnvironmentScripts } from '#~/services/worktree-environments.js' import { getSessionLogger } from '#~/utils/logger.js' @@ -361,12 +361,6 @@ const buildAgentRoomRuntimeContextPrompt = ( ].join('\n') } -const resolveInteractiveSessionWorktreeDefault = async () => { - const { mergedConfig } = await loadConfigState() - .catch(() => ({ mergedConfig: {} as { conversation?: { createSessionWorktree?: boolean } } })) - return mergedConfig.conversation?.createSessionWorktree ?? false -} - const uniqueStrings = (values: string[]) => [...new Set(values)] const getPermissionToolUseCache = (sessionId: string) => { @@ -663,10 +657,13 @@ export async function startAdapterSession( const startPromise = (async () => { const db = getDb() + const existing = db.getSession(sessionId) + if (existing == null) { + throw new Error(`Session ${sessionId} does not exist`) + } const historyMessages = db.getMessages(sessionId) as WSEvent[] const hasHistory = historyMessages.some(event => event.type === 'message') const serverLogger = getSessionLogger(sessionId, 'server') - const existing = db.getSession(sessionId) const runtimeState = db.getSessionRuntimeState(sessionId) const resolvedModel = options.model ?? existing?.model const resolvedAdapter = options.adapter ?? existing?.adapter @@ -750,16 +747,6 @@ export async function startAdapterSession( resolvedPromptName }, '[server] Starting new adapter process') - if (existing == null) { - serverLogger.info({ sessionId }, '[server] Session not found in DB, creating new entry') - db.createSession(undefined, sessionId, undefined, undefined, { - runtimeKind: 'interactive' - }) - await provisionSessionWorkspace(sessionId, { - createWorktree: await resolveInteractiveSessionWorktreeDefault() - }) - } - const withResolvedPermissionMode = ( updates: Partial> ) => { diff --git a/apps/server/src/websocket/session-connection.ts b/apps/server/src/websocket/session-connection.ts index b7e0ec6b0..63ae8cdae 100644 --- a/apps/server/src/websocket/session-connection.ts +++ b/apps/server/src/websocket/session-connection.ts @@ -5,6 +5,11 @@ import { v4 as uuidv4 } from 'uuid' import type { WebSocket } from 'ws' import { getDb } from '#~/db/index.js' +import { + SessionCreationWaitError, + isSessionCreationWaitError, + waitForSessionCreation +} from '#~/services/session/creation-lifecycle.js' import { interruptSession, killSession, processUserMessage, startAdapterSession } from '#~/services/session/index.js' import { handleInteractionResponse } from '#~/services/session/interaction.js' import { @@ -49,15 +54,33 @@ export const handleSessionWebSocketConnection = async ( const serverLogger = getSessionLogger(sessionId, 'server') serverLogger.info({ sessionId }, '[server] Connection established') + const creationWaitController = new AbortController() + let socketClosed = ws.readyState !== WEBSOCKET_OPEN + const isSocketOpen = () => !socketClosed && ws.readyState === WEBSOCKET_OPEN + ws.on('close', () => { + socketClosed = true + creationWaitController.abort() + }) + if (!isSocketOpen()) { + creationWaitController.abort() + return + } try { + await waitForSessionCreation(sessionId, { + signal: creationWaitController.signal + }) + if (!isSocketOpen()) return + const db = getDb() const sessionData = db.getSession(sessionId) + if (sessionData == null) { + throw new SessionCreationWaitError(sessionId, 'Session creation completed without a session record') + } const sessionRuntimeState = db.getSessionRuntimeState(sessionId) const isExternalSession = sessionRuntimeState?.runtimeKind === 'external' const cachedRuntime = getAdapterSessionRuntime(sessionId) - const shouldAutoStartAdapter = sessionData == null || - sessionData.status === 'running' || + const shouldAutoStartAdapter = sessionData.status === 'running' || sessionData.status === 'waiting_input' if (isExternalSession) { @@ -83,15 +106,20 @@ export const handleSessionWebSocketConnection = async ( adapter, account }) - attachSocketToSession(sessionId, ws, 'adapter') if (cached == null) { throw new Error(`Failed to initialize session runtime for ${sessionId}`) } + if (!isSocketOpen()) return + attachSocketToSession(sessionId, ws, 'adapter') } else { attachSocketToSession(sessionId, ws, 'external') } } catch (err) { + if (!isSocketOpen()) return sendSocketError(ws, err) + if (isSessionCreationWaitError(err)) { + ws.close(1008, 'Session creation failed') + } return } diff --git a/changelog/1.0.0-rc.4/server.md b/changelog/1.0.0-rc.4/server.md new file mode 100644 index 000000000..385dbd540 --- /dev/null +++ b/changelog/1.0.0-rc.4/server.md @@ -0,0 +1,3 @@ +# @oneworks/server 1.0.0-rc.4 + +- Let session WebSockets connect before HTTP creation finishes without racing to create duplicate database records, and keep slow or cancelled creation attempts coordinated through completion.